Diff Coverage

"""

import logging
import re
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Optional, Union

import torch

logger = logging.getLogger(__name__)


def _apply_parameter_freezing(model: nn.Module, freeze_config: Any) -> None:
    """Freeze parameters selected by fully qualified module prefixes."""
    if not isinstance(freeze_config, Mapping):
        raise ValueError("freeze_config must be a mapping")
    module_prefixes = freeze_config.get("module_prefixes")
    if not isinstance(module_prefixes, list) or not module_prefixes or not all(
            isinstance(prefix, str) and prefix for prefix in module_prefixes
    ):
        raise ValueError("freeze_config.module_prefixes must be a non-empty list of strings")

    frozen_names = []
    for name, parameter in model.named_parameters():
        if any(name == prefix or name.startswith(f"{prefix}.") for prefix in module_prefixes):
            parameter.requires_grad_(False)
            frozen_names.append(name)
    if not frozen_names:
        raise ValueError(f"freeze_config.module_prefixes matched no parameters: {module_prefixes}")
    logger.info("Froze %s parameters under module prefixes %s", len(frozen_names), module_prefixes)


@dataclass(frozen=True)
class _FinalizeTarget:

        context=context,
    )


def _resolve_compile_config(
    compile_config: Optional[Union[CompileConfig, dict]],
    validate_placement: bool,
    fsdp2_manager: Optional[FSDP2Manager],
) -> tuple[Optional[CompileConfig], bool]:

    validate_placement: bool,
    fsdp2_manager: Optional[FSDP2Manager],
) -> tuple[Optional[CompileConfig], bool]:
    """Normalize compile configuration and validate execution constraints."""
    if isinstance(compile_config, dict):
        compile_config = CompileConfig(enabled=True, **compile_config)
    if compile_config is not None and not isinstance(compile_config, CompileConfig):
        raise TypeError("compile_config must be a CompileConfig, mapping, or None")

    compile_for_execution = bool(
        not validate_placement
        and compile_config is not None
        and compile_config.enabled
    )
    if validate_placement and compile_config is not None and compile_config.enabled:
        logger.info("Skipping decoder-layer compile during placement validation")
    if (
        compile_for_execution
        and compile_config.fullgraph
        and isinstance(fsdp2_manager, FSDP2Manager)
    ):
        raise ValueError(
            "compile.fullgraph=True is incompatible with FSDP hooks kept eager "
            "by _dynamo_disable; set compile.fullgraph=False"
        )
    return compile_config, compile_for_execution


def apply_model_infrastructure(
    model: nn.Module,

    """

    distributed_setup = kwargs.get("distributed_setup")

    compile_config, compile_for_execution = _resolve_compile_config(
        compile_config,
        validate_placement,
        fsdp2_manager,
    )

    )

    # Step 6: Parameter freezing (before sharding)
    if freeze_config is not None:
        _apply_parameter_freezing(model, freeze_config)

    # Steps 7-8: plan and apply parameter/activation layouts.
    model, source_shard_info = _plan_and_apply_sharding(
        model,

            )

        return attention_mask, swa_mask, packed_seq_params

    @staticmethod
    def _split_model_and_loss_inputs(
            parallel_batch: Mapping[str, Any],
    ) -> tuple[Mapping[str, Any], Mapping[str, Any]]:
        """Split forward fields from loss and token-accounting fields."""

import torch
from torch.utils.data import default_collate

from hyper_models.components.utils.constants import IGNORE_INDEX

_TEXT_FIELDS = {
    "input_ids",
    "labels",

            {field: value for field, value in sample.items() if field not in _TEXT_FIELDS}
            for sample in samples
        ]

        batch = default_collate(text_samples)
        if any(modal_samples):
            for field in {field for sample in modal_samples for field in sample}:
                values = [sample[field] for sample in modal_samples if field in sample]
                batch[field] = (

    Returns:
        A collator producing one VLM micro-batch dictionary.
    """
    if packing:
        raise NotImplementedError("The temporary VLM collator does not support packing")
    if pad_token_id != 0 or ignore_index != IGNORE_INDEX or pad_to_length is not None:
        raise NotImplementedError("The temporary VLM collator does not support custom text padding")
    return VLMCollator()


__all__ = ["VLMCollator", "build_vlm_collator"]

"""Build the VLM dataset from a LLaVA-style JSON list."""

import json
import os
from collections.abc import Callable
from typing import Any, Optional, TypeAlias

from torch.utils.data import Dataset

from hyper_models.components.utils.constants import IGNORE_INDEX

SampleTransform: TypeAlias = Callable[[Any], Any]


class VLMDataset(Dataset):
    """Load a LLaVA-style JSON list of multimodal conversations.

# limitations under the License.
# ============================================================================
"""Temporary self-contained VLM batch preparation."""

from collections.abc import Mapping
from typing import Any

from hyper_parallel.platform import get_platform

platform = get_platform()

_MODEL_INPUT_FIELDS = {
    "input_ids",
    "labels",

class VLMBatchProcessor:
    """Normalize and classify one VLM batch without shared LLM adapters."""

    @staticmethod
    def normalize_source_batch(source_batch: Mapping[str, Any]) -> dict[str, Any]:
        """Normalize one collated batch into the temporary VLM contract."""
        batch = dict(source_batch)
        if "input_ids" not in batch:
            raise ValueError("VLM batch must contain 'input_ids'")
        if "labels" not in batch:
            raise ValueError("VLM batch must contain 'labels'")
        if "loss_mask" not in batch:
            batch["loss_mask"] = batch["labels"] >= 0
        return batch

    @staticmethod
    def prepare_batch(batch: Mapping[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
        """Split device-resident VLM fields into model and loss inputs."""
        model_inputs = {field: value for field, value in batch.items() if field in _MODEL_INPUT_FIELDS}
        loss_inputs = {field: value for field, value in batch.items() if field in _LOSS_INPUT_FIELDS}
        return model_inputs, loss_inputs


class VLMGetBatch:
    """Prepare VLM batches for the temporary TP=CP=PP=1 training path."""

    def __init__(self, *, mesh_context: Any, device: Any, pp_shared_data: bool = False) -> None:
        """Validate the temporary VLM parallel boundary and store the device.

        Args:
            mesh_context: Trainer mesh exposing TP, CP, and PP sizes.

        Raises:
            NotImplementedError: If model parallelism or pipeline batch sharing is enabled.
        """
        parallel_sizes = {
            "tp_size": int(getattr(mesh_context, "tp_size", 1)),
            "cp_size": int(getattr(mesh_context, "cp_size", 1)),
            "pp_size": int(getattr(mesh_context, "pp_size", 1)),
        }
        unsupported_sizes = {name: size for name, size in parallel_sizes.items() if size != 1}
        if unsupported_sizes:
            raise NotImplementedError(
                "The temporary VLM batch path requires TP=CP=PP=1, but got "
                + ", ".join(f"{name}={size}" for name, size in unsupported_sizes.items())
            )
        if pp_shared_data:
            raise NotImplementedError("The temporary VLM batch path does not support pp_shared_data")
        self.device = device
        self.processor = VLMBatchProcessor()

    def __call__(
            self,
            data_iterator: Any,
            *,
            external_batch: Mapping[str, Any] | None = None,

        Returns:
            Model inputs and loss inputs on the configured device.
        """
        source_batch = external_batch if external_batch is not None else next(data_iterator)
        if not isinstance(source_batch, Mapping):
            raise ValueError("VLM DataLoader must yield a mapping batch")
        normalized_batch = self.processor.normalize_source_batch(source_batch)
        device_batch = {
            field: platform.move_to_device(value, self.device, non_blocking=True)
            if platform.is_tensor(value) else value
            for field, value in normalized_batch.items()
        }
        return self.processor.prepare_batch(device_batch)


def build_vlm_get_batch(
        *,

    Returns:
        Callable VLM batch adapter.
    """
    return VLMGetBatch(mesh_context=mesh_context, device=device, pp_shared_data=pp_shared_data)


__all__ = ["VLMBatchProcessor", "VLMGetBatch", "build_vlm_get_batch"]

from typing import Any, Dict

from hyper_parallel import SkipDTensorDispatch, hsdp_sync_stream
from hyper_parallel.core.utils import clip_grad_norm_
from hyper_models.components.datasets import calculate_num_micro_batches
from hyper_models.components.datasets.vlm import build_processor, build_vlm_get_batch
from hyper_models.components.loss.loss_utils import count_loss_token
from hyper_models.components.utils import helper
from hyper_models.components.utils.device import synchronize  # pylint: disable=syntax-error
from hyper_models.trainer.base import BaseTrainer
from hyper_models.trainer.config import TrainerConfig

logger = helper.create_logger(__name__)