Diff Coverage

                and not isinstance(x_layout.ragged_shard, RaggedShardInfo)
                and not isinstance(to_layout.ragged_shard, RaggedShardInfo)):
            # Validate the export-only path before any collective or layout rewrite.
            if x_layout.rank_list != to_layout.rank_list:
                raise ValueError("Uneven redistribution requires the same rank list")
            return self._gather_uneven_to_replicate(input_x, to_layout)
        x = input_x
        if input_x.layout.is_partial():
            # Solve partial status first

            to_layout.ragged_shard, RaggedShardInfo
        ):
            return self._redistribute_ragged(x, to_layout)
        if from_layout.has_uneven_shard or to_layout.has_uneven_shard:
            return self._gather_uneven_to_replicate(x, to_layout)
        return self._redistribution_normal(x, to_layout)

    @staticmethod
    @torch.no_grad()

        if len(rank_list) == 1:
            return local_tensor.detach().clone()
        max_size = max(sizes)
        if tensor_dim != 0 and any(size != max_size for size in sizes):
            raise NotImplementedError("Uneven export communication only supports dim=0")
        padded = local_tensor.detach().contiguous()
        if tensor_dim == 0 and local_tensor.shape[0] != max_size:
            padded_shape = (max_size, *local_tensor.shape[1:])
            padded = local_tensor.new_zeros(padded_shape)

            sample: Messages, optional external images, and source path metadata.
        """
        messages = sample.get("messages")
        if not isinstance(messages, list) or not messages:
            raise ValueError("Omni sample must contain non-empty messages")

        prepared_messages = copy.deepcopy(messages)
        self._bind_external_images(prepared_messages, sample.get("images"))
        source_path = sample.get(ONLINE_SOURCE_PATH_KEY)

        source_directory = os.path.dirname(os.path.abspath(source_path))
        for message in prepared_messages:
            content = message.get("content")
            if not isinstance(content, list):
                continue
            for content_block in content:
                self._resolve_media_block(content_block, source_directory)
        return prepared_messages

        """Consume external images only after checking the complete placeholder count."""
        if images is None:
            images = []
        if not isinstance(images, list):
            raise TypeError("Omni sample images must be a list")
        placeholders = []
        for message in messages:
            content = message.get("content")
            if isinstance(content, str) and (images or "<image>" in content):

        """Preserve text around each placeholder and leave explicit media untouched."""
        expanded = []
        for block in content:
            if not isinstance(block, dict):
                expanded.append(block)
                continue
            text_content = block.get("text")
            if (block.get("type") == "text"
                    and isinstance(text_content, str) and "<image>" in text_content):
                for index, text in enumerate(block["text"].split("<image>")):

    @classmethod
    def _resolve_media_block(cls, content_block: Any, source_directory: str) -> None:
        """Resolve path-bearing fields without changing the processor's message format."""
        if not isinstance(content_block, dict):
            return

        block_type = content_block.get("type")
        if block_type == "image_url":
            image_url = content_block.get("image_url")
            if isinstance(image_url, str):
                content_block["image_url"] = cls._resolve_media_path(image_url, source_directory)
            elif isinstance(image_url, dict) and isinstance(image_url.get("url"), str):
                image_url["url"] = cls._resolve_media_path(image_url["url"], source_directory)
            return

        if block_type in ("image", "video", "audio"):
            for field in (block_type, "url", "path"):
                media_field = content_block.get(field)

        for end in range(common_length + 1, len(full_tokens) + 1):
            decoded = tokenizer.decode(full_tokens[common_length:end], **decode_kwargs)
            if full_tail.startswith(decoded) and decoded.startswith(prompt_tail):
                return end
        raise ValueError("Cannot locate assistant supervision boundary in processor tokens")

    def encode_sample(self, sample: dict[str, Any]) -> dict[str, Any]:
        """Build model inputs and assistant-only labels from processor messages."""
        if self.training_mode != "sft":

            value = sample[source_key]
        except KeyError as exc:
            raise ValueError(f"Instruction sample does not contain field {source_key!r}") from exc
        if not isinstance(value, str):
            raise ValueError(f"Instruction field {source_key!r} must be a string")
        if not value:
            raise ValueError(f"Instruction field {source_key!r} must not be empty")
        return value

    ) -> None:
        """Store source configuration and optional training sample targets."""
        super().__init__(data_path, data_config)
        if sampler_type not in ("single", "cyclic"):
            raise ValueError(f"Unsupported Mapping sampler_type: {sampler_type!r}")
        self.sampler_type = sampler_type
        self.train_valid_test_num_samples = self._validate_split_sizes(train_valid_test_num_samples)

    def build_source(

        if load_dataset_path in ("json", "parquet") and "data_files" in load_options and not streaming:
            loaded_dataset = self._load_local_files(load_dataset_path, load_options["data_files"], selected_split)
        else:
            loaded_dataset = load_dataset(
                load_dataset_path,
                **load_options,
                split=selected_split,
                streaming=streaming,

        from datasets import ReadInstruction, concatenate_datasets, load_dataset  # pylint: disable=C0415

        files_by_split = data_files if isinstance(data_files, Mapping) else {"train": data_files}
        if any("://" in path for paths in files_by_split.values() for path in paths):
            return load_dataset(
                file_format, data_files=data_files, split=selected_split,
                cache_dir=self.data_config.get("cache_dir"),
            )
        datasets_by_split = {}

                        file_format, data_files=source_path, split="train",
                        cache_dir=self.data_config.get("cache_dir"),
                    )
                    if ONLINE_SOURCE_PATH_KEY in source.column_names:
                        source = source.remove_columns(ONLINE_SOURCE_PATH_KEY)
                    source = source.add_column(
                        ONLINE_SOURCE_PATH_KEY, [os.path.abspath(source_path)] * len(source),
                    )
                    source_datasets.append(source)

        """Log the first raw conversation and native-template rendering once."""
        if not self.log_first_rendered_template or self._has_logged_first_rendered_template:
            return
        if os.environ.get("RANK", "0") != "0":
            self._has_logged_first_rendered_template = True
            return

        rendered_template = self._apply_chat_template(
            messages,
            tokenize=False,

        )
        logger.info("First chat-template input messages: %r", messages)
        template_kwargs = dict(self.chat_template_kwargs)
        if tools:
            template_kwargs["tools"] = tools
        logger.info("First chat-template kwargs: %r", template_kwargs)
        logger.info("First rendered chat template:\n%s", rendered_template)
        self._has_logged_first_rendered_template = True

        for end, message in enumerate(messages, start=1):
            loss_mask = message.get("loss_mask", 1 if message["role"] == "assistant" else 0)
            if previous_length == 0 and message["role"] == "system":
                if loss_mask == 1:
                    raise ValueError(
                        "A leading system message with loss enabled cannot be grouped with the first user message."
                    )
                has_pending_system_messages = True
                continue

                    )
                has_pending_system_messages = True
                continue
            if has_pending_system_messages and previous_length == 0 and loss_mask == 1:
                raise ValueError(
                    "The first message after leading system messages must have loss disabled so their token "
                    "boundaries can be grouped safely."
                )

                if (
                        len(source_ids) < previous_length
                        or source_ids[:previous_length] != input_ids
                ):
                    raise ValueError(
                        "The tokenizer generation prompt structurally rewrote the conversation prefix; "
                        "assistant-header loss masking requires the inference prompt to extend the training history."
                    )
                # History contains supervised targets from earlier turns: mask only the new source suffix.

            previous_length = current_length
            previous_text = rendered_text

        if has_pending_system_messages and previous_length == 0:
            raise ValueError("A conversation containing only system messages cannot be rendered for training.")

        input_ids = input_ids[-max_seq_len:]
        labels = labels[-max_seq_len:]
        return {

        if not tools:
            return None
        try:
            tools = json.loads(tools)
        except json.JSONDecodeError as exc:
            raise ValueError(f"Tools field {tools_key!r} must contain valid JSON") from exc
        if tools is None or tools is False:
            return None
    if not isinstance(tools, Sequence) or isinstance(tools, (str, bytes)):
        raise ValueError(f"Tools field {tools_key!r} must be a sequence of mappings")

    normalized_tools = []
    for tool in tools:
        if not isinstance(tool, Mapping):
            raise ValueError(f"Every entry in tools field {tools_key!r} must be a mapping")
        normalized_tools.append(dict(tool))
    return normalized_tools or None

            raise ValueError("chat_template is required for conversation data")
        if self.max_seq_len <= 0:
            raise ValueError("max_seq_len must be positive")
        if self.tools_key is not None and (not isinstance(self.tools_key, str) or not self.tools_key):
            raise ValueError("tools_key must be a non-empty string or None")

    def __call__(self, sample: Mapping[str, Any]) -> list[dict[str, Any]]:
        """Encode one conversation record."""
        messages = (self.message_formatter(sample) if self.message_formatter is not None

    if text_keys is None:
        text_keys = {"sharegpt": "conversations"}.get(data_type, "text")
    if data_type == "plaintext":
        if tools_key is not None:
            raise ValueError("tools_key is only supported for sharegpt data")
        data_transform = PlaintextTransform(tokenizer, max_seq_len, text_keys, drop_ragged_tail)
    elif data_type == "sharegpt":
        data_transform = TextConversationTransform(chat_template, max_seq_len, text_keys, tools_key)
    elif data_type == "alpaca":

    elif data_type == "sharegpt":
        data_transform = TextConversationTransform(chat_template, max_seq_len, text_keys, tools_key)
    elif data_type == "alpaca":
        if tools_key is not None:
            raise ValueError("tools_key is only supported for sharegpt data")
        if column_mapping is None:
            raise ValueError("column_mapping is required for alpaca data")
        data_transform = TextConversationTransform(
            chat_template,

            max_seq_len,
            message_formatter=InstructionMessages(column_mapping, input_separator),
        )
    else:
        raise ValueError(f"Unsupported text data format: {data_type!r}")

    logger.debug("Built text transform: data_type=%s, transform=%s", data_type, type(data_transform).__name__)
    return data_transform

                continue
            matched.add(class_name)
            parameters = list(module.parameters())
            if not parameters:
                raise ValueError(f"wrapper_class selected a parameterless module: {fqn}")
            if any(not name.startswith(f"{fqn}.") for param in parameters for name in aliases[param]):
                raise ValueError(f"wrapper_class module shares parameters outside its subtree: {fqn}")
            selected.append(_WrapModuleInfo(fqn, module))
            seen.add(id(module))

            wrapped_module_ids = {id(unit.module) for unit in wrap_modules}
        else:
            declared_wrap_modules = self._find_adapter_declared_wrap_modules(model)
            if declared_wrap_modules:
                logger.info(
                    "Adapter-declared FSDP child units: %s",
                    [wrap_module.fqn for wrap_module in declared_wrap_modules],
                )
            excluded_subtree_module_fqns = {

        original_conv: Callable[..., Any], cp_mesh: Any,
        hidden_states: torch.Tensor, weight: torch.Tensor,
        bias: Any = None, activation: Any = None, **kwargs: Any) -> Any:
    """Run the original causal Conv1d with the previous CP rank's halo."""
    halo_width = weight.shape[-1] - 1
    if halo_width <= 0 or cp_mesh.size() <= 1:
        return original_conv(
            hidden_states, weight, bias, activation, **kwargs)
    if hidden_states.shape[-1] < halo_width:
        raise ValueError(
            "GDN CP requires local sequence length "
            f"({hidden_states.shape[-1]}) to be at least Conv1d halo width "
            f"({halo_width})"
        )
    tail = hidden_states[..., -halo_width:].contiguous()
    halo = _exchange_previous_cp_halo(tail, cp_mesh)
    conv_input = torch.cat((halo, hidden_states), dim=-1)
    output = original_conv(
        conv_input, weight, bias, activation, **kwargs)
    return output[..., halo_width:].contiguous()


def _gdn_cp_conv1d_module(
        original_forward: Callable[..., Any], conv_module: Any,

def _gdn_cp_conv1d_module(
        original_forward: Callable[..., Any], conv_module: Any,
        cp_mesh: Any, hidden_states: torch.Tensor) -> torch.Tensor:
    """Run the Transformers Conv1d fallback with the previous CP halo."""
    kernel_size = conv_module.kernel_size[0]
    dilation = conv_module.dilation[0]
    halo_width = (kernel_size - 1) * dilation
    if halo_width <= 0 or cp_mesh.size() <= 1:
        return original_forward(hidden_states)
    if hidden_states.shape[-1] < halo_width:
        raise ValueError(
            "GDN CP requires local sequence length "
            f"({hidden_states.shape[-1]}) to be at least Conv1d halo width "
            f"({halo_width})"
        )
    tail = hidden_states[..., -halo_width:].contiguous()
    halo = _exchange_previous_cp_halo(tail, cp_mesh)
    conv_input = torch.cat((halo, hidden_states), dim=-1)
    output = original_forward(conv_input)
    local_length = hidden_states.shape[-1]
    return output[..., halo_width:halo_width + local_length].contiguous()


def _gdn_rule_cp_to_hp(
        query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,

        query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
        decay: torch.Tensor, beta: torch.Tensor,
        cp_mesh: Any) -> tuple[torch.Tensor, ...]:
    """Convert local-sequence GDN rule inputs into local-head layouts."""
    tensors = (query, key, value, decay, beta)
    expected_dims = (4, 4, 4, 3, 3)
    if any(
            tensor.dim() != expected
            for tensor, expected in zip(tensors, expected_dims)):
        shapes = [tuple(tensor.shape) for tensor in tensors]
        raise ValueError(
            "GDN CP expects 4-D Q/K/V and 3-D g/beta, got "
            f"{shapes}"
        )
    if any(tensor.shape[1] != query.shape[1] for tensor in tensors[1:]):
        raise ValueError("GDN CP inputs must share the sequence length")
    return tuple(
        ulysses_seq_to_head(
            tensor, seq_dim=1, head_dim=2, cp_mesh=cp_mesh)
        for tensor in tensors
    )

def _gdn_rule_hp_to_cp(
        output: torch.Tensor, cp_mesh: Any) -> torch.Tensor:
    """Restore one GDN rule output to the local-sequence layout."""
    if output.dim() != 4:
        raise ValueError(
            "GDN CP expects a 4-D GDN rule output, got "
            f"shape {tuple(output.shape)}"
        )
    return ulysses_head_to_seq(
        output, seq_dim=1, head_dim=2, cp_mesh=cp_mesh)


def _validate_gdn_cp(target_module: Any, cp_mesh: Any) -> None:

def _validate_gdn_cp(target_module: Any, cp_mesh: Any) -> None:
    """Validate the GDN interface required by CP interception."""
    if cp_mesh is None or cp_mesh.size() <= 1:
        raise ValueError("GDN CP requires an active CP mesh")
    module_name = type(target_module).__name__
    if "GatedDeltaNet" not in module_name:
        raise TypeError(
            "GDN CP wrapper requires a GatedDeltaNet module, got "
            f"{module_name}"
        )
    num_value_heads = getattr(target_module, "num_v_heads", None)

            f"{module_name}"
        )
    num_value_heads = getattr(target_module, "num_v_heads", None)
    if not isinstance(num_value_heads, int) or num_value_heads <= 0:
        raise ValueError(
            "GDN CP requires target_module.num_v_heads to be a "
            "positive integer"
        )
    if num_value_heads % cp_mesh.size():
        raise ValueError(
            f"GDN value heads ({num_value_heads}) must be divisible "
            f"by CP size ({cp_mesh.size()})"
        )

            return implementation
        closure = inspect.getclosurevars(implementation).nonlocals
        wrapped = closure.get("forward_func")
        if not callable(wrapped):
            break
        implementation = inspect.unwrap(wrapped)
    raise RuntimeError(
        "GDN CP could not resolve the Transformers forward "
        "implementation"
    )

        and "conv1d" in referenced_names
        and callable(original_conv_forward)
    )
    if not use_fused_conv and not use_module_conv:
        raise RuntimeError(
            "GDN CP requires a fused or module causal Conv1d path"
        )

    rule_names = tuple(

            or callable(forward_globals.get(name))
        )
    )
    if not rule_names:
        raise RuntimeError(
            "GDN CP could not find a supported chunk Gated Delta "
            "Rule call in the Transformers forward"
        )
    rule_owners = {

            *args: Any,
            **kwargs: Any,
        ) -> Any:
            if args:
                raise TypeError(
                    "GDN CP expects g and beta to be passed by keyword"
                )
            decay = kwargs.get("g")
            beta = kwargs.get("beta")

                )
            decay = kwargs.get("g")
            beta = kwargs.get("beta")
            if not isinstance(decay, torch.Tensor) or not isinstance(beta, torch.Tensor):
                raise TypeError(
                    "GDN CP requires Tensor keyword arguments g and beta"
                )
            if kwargs.get("initial_state") is not None or kwargs.get(
                    "output_final_state", False):
                raise NotImplementedError(
                    "GDN CP does not support recurrent cache state"
                )
            query, key, value, decay, beta = _gdn_rule_cp_to_hp(
                query, key, value, decay, beta, cp_mesh)

            call_kwargs = dict(kwargs)
            call_kwargs.update(g=decay, beta=beta)
            output = original_rule(query, key, value, **call_kwargs)
            if not isinstance(output, tuple) or len(output) != 2:
                raise TypeError(
                    "GDN rule must return (output, recurrent_state)"
                )
            core_output, recurrent_state = output
            if recurrent_state is not None:
                raise NotImplementedError(
                    "GDN CP does not support a returned recurrent state"
                )
            return _gdn_rule_hp_to_cp(core_output, cp_mesh), recurrent_state

    def cp_forward(*args: Any, **kwargs: Any) -> Any:
        """Run the original GDN forward with CP-aware primitive interfaces."""
        cache_params = kwargs.get("cache_params")
        if cache_params is None and len(args) > 1:
            cache_params = args[1]
        if cache_params is not None:
            raise NotImplementedError(
                "GDN CP currently supports training with cache_params=None"
            )
        if kwargs.get("cu_seq_lens_q") is not None:
            raise NotImplementedError(
                "GDN CP does not support packed sequences"
            )

        fired = {"conv": 0, "rule": 0}

                    conv_args = (conv_kwargs.pop("x"),)
                return _gdn_cp_causal_conv1d(
                    original_conv, cp_mesh, *conv_args, **conv_kwargs)
        else:
            @functools.wraps(original_conv_forward)
            def cp_conv(hidden_states: torch.Tensor) -> torch.Tensor:
                fired["conv"] += 1
                return _gdn_cp_conv1d_module(
                    original_conv_forward, conv_module, cp_mesh, hidden_states)

        def count_rule(wrapped_rule):
            @functools.wraps(wrapped_rule)

        if use_fused_conv:
            replace_primitive(conv_owner, conv_name, cp_conv)
        else:
            conv_module.forward = cp_conv
        for name, wrapped_rule in wrapped_rules.items():
            replace_primitive(rule_owners[name], name, wrapped_rule)
        try:
            output = original_forward(*args, **kwargs)

        finally:
            if use_fused_conv:
                replace_primitive(conv_owner, conv_name, original_conv)
            else:
                conv_module.forward = original_conv_forward
            for name, original_rule in original_rules.items():
                replace_primitive(rule_owners[name], name, original_rule)
        if fired["conv"] != 1 or fired["rule"] != 1:
            raise RuntimeError(
                "GDN CP wrapper expected one causal Conv1d and one "
                "Gated Delta Rule call, got "
                f"conv={fired['conv']}, rule={fired['rule']}"
            )

            self.stress_detect_callback = StressDetectCallback(self)
            self._callbacks.append(self.stress_detect_callback)
        # Registered to save, to restore, or both --- the two are independent, so
        # a run that only loads an existing checkpoint still needs the callback.
        if (
            self.config.checkpoint.save_ckpt
            or self.config.checkpoint.save_hf_weights
            or self.config.checkpoint.restore_from is not None
        ):

        self._save_steps = ckpt_cfg.save_steps if self._save_ckpt else 0
        self._save_epochs = ckpt_cfg.save_epochs if self._save_ckpt else 0
        save_time_interval_minutes = float(ckpt_cfg.save_time_interval_minutes)
        if not math.isfinite(save_time_interval_minutes):
            raise ValueError("checkpoint.save_time_interval_minutes must be finite")
        self._save_time_interval_seconds = (
            max(save_time_interval_minutes, 0.0) * 60.0
            if self._save_ckpt
            else 0.0

        local_due = int(time.monotonic() >= self._next_time_save)
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            return bool(all_reduce(local_due, op="max"))
        return bool(local_due)

    def _advance_time_deadline(self) -> None:
        """Move the time deadline forward without accumulating save latency."""
        if self._next_time_save is None:
            return
        now = time.monotonic()
        elapsed = max(now - self._next_time_save, 0.0)
        elapsed_intervals = int(elapsed // self._save_time_interval_seconds)
        self._next_time_save += (elapsed_intervals + 1) * self._save_time_interval_seconds

            # Complete and rotate the preceding async checkpoint before
            # starting a new one. This temporarily allows limit + 1
            # directories instead of deleting the last known-good checkpoint
            # while its replacement is still being written.
            self.wait_for_pending_save()

        save_dir = os.path.join(self._checkpoint_dir, f"{STEP_PREFIX}{state.global_step}")
        save_async = self._is_async and not force_sync

    def _save_hf_checkpoint(self, save_dir: str, global_step: int) -> None:
        """Collect and export one Transformers-compatible model checkpoint."""
        if self._hf_checkpoint_manager is None:
            raise RuntimeError("Hugging Face checkpoint manager is not initialized")

        hf_dir = os.path.join(save_dir, _HF_CHECKPOINT_DIR)
        logger.info(
            "Saving Hugging Face weights: global_step=%s, dir=%s",

        suffix = name[len(STEP_PREFIX):]
        if not suffix or not suffix.isascii() or not suffix.isdigit():
            return None
        if len(suffix) > 1 and suffix.startswith("0"):
            return None
        return int(suffix)

    @staticmethod
    def _is_complete_checkpoint(path: str) -> bool:

                        )
                        for entry in entries
                ):
                    return True
        except OSError:
            return False

        hf_dir = os.path.join(path, _HF_CHECKPOINT_DIR)
        if os.path.islink(hf_dir) or not os.path.isdir(hf_dir):
            return False

        """List completed, non-symlink checkpoint directories from oldest to newest."""
        checkpoint_root = os.path.realpath(self._checkpoint_dir)
        try:
            entries = os.scandir(checkpoint_root)
        except OSError:
            return []

        checkpoints = []
        with entries:
            for entry in entries:

                if step is None or entry.is_symlink() or not entry.is_dir(follow_symlinks=False):
                    continue
                resolved_path = os.path.realpath(entry.path)
                if os.path.dirname(resolved_path) != checkpoint_root:
                    logger.warning("Skipping checkpoint outside configured root: %s", entry.path)
                    continue
                if self._is_complete_checkpoint(resolved_path):
                    checkpoints.append((step, resolved_path))

        checkpoints.sort(key=lambda item: item[0])

        if limit is None or limit <= 0:
            return
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            if torch.distributed.get_rank() != 0:
                return

        checkpoints = self._list_managed_checkpoints()
        delete_count = len(checkpoints) - limit
        if delete_count <= 0:

            if step >= 0
        }
        deletable = [item for item in checkpoints if item[0] not in protected_steps]
        if len(deletable) < delete_count:
            logger.warning(
                "Checkpoint limit %s cannot be met without deleting an active checkpoint; "
                "keeping %s completed checkpoints.",
                limit,
                len(checkpoints) - len(deletable),

                "keeping %s completed checkpoints.",
                limit,
                len(checkpoints) - len(deletable),
            )
            delete_count = len(deletable)

        for step, checkpoint_path in deletable[:delete_count]:
            # Revalidate immediately before deletion so a replaced path or
            # symlink can never redirect rotation outside checkpoint_dir.

                    or not os.path.isdir(checkpoint_path)
                    or os.path.dirname(os.path.realpath(checkpoint_path)) != checkpoint_root
                    or self._checkpoint_step(os.path.basename(checkpoint_path)) != step
            ):
                logger.warning("Skipping unsafe checkpoint deletion target: %s", checkpoint_path)
                continue
            try:
                shutil.rmtree(checkpoint_path)
            except OSError as exc:
                logger.warning("Failed to delete old checkpoint %s: %s", checkpoint_path, exc)
                continue
            logger.info(
                "Deleted old checkpoint due to save_total_limit=%s: %s",
                limit,
                checkpoint_path,

    def _save_hf_assets(self, hf_dir: str) -> None:
        """Write the tokenizer or processor assets next to exported weights."""
        processor = getattr(self.trainer, "processor", None)
        if processor is not None:
            processor.save_pretrained(hf_dir)
            return

        chat_template = getattr(self.trainer, "chat_template", None)
        if chat_template is not None:
            chat_template.save_pretrained(hf_dir)
            return

        tokenizer = getattr(self.trainer, "tokenizer", None)
        if tokenizer is not None:
            save_tokenizer = getattr(tokenizer, "save_pretrained", None)

                    chat_template_kwargs=assets_config.chat_template_kwargs,
                    log_first_rendered_template=config.debug.log_first_rendered_template,
                )
            else:
                if assets_config.chat_template_kwargs or config.debug.log_first_rendered_template:
                    raise ValueError(
                        "dataset.model_assets chat-template options require chat_template='tokenizer'"
                    )
                self.base.chat_template = assets_config.chat_template.build(
                    tokenizer=self.base.tokenizer,

        if config.dataloader.get_batch is None:
            raise ValueError("dataloader.get_batch must define a batching runtime target")
        if getattr(config.dataloader.get_batch, "encoder_parallel", None) is not None:
            raise ValueError("encoder_parallel belongs under accelerator, not dataloader.get_batch")
        get_batch = config.dataloader.get_batch.build(
            mesh_context=self.base.mesh,
            device=self.base.device,
            data_config=getattr(config.dataset, "data_config", {}),
            pp_shared_data=bool(getattr(config.dataloader, "pp_shared_data", False)),

            device=self.base.device,
            data_config=getattr(config.dataset, "data_config", {}),
            pp_shared_data=bool(getattr(config.dataloader, "pp_shared_data", False)),
        )
        bind_model = getattr(get_batch, "bind_model", None)
        if callable(bind_model):
            bind_model(self.base.model)
        encoder_runtime = build_encoder_parallel_runtime(
            config=config.accelerator.encoder_parallel,
            mesh_context=self.base.mesh,
            model=self.base.model,