Diff Coverage

from hyper_parallel.trainer.runtime import fsdp as fsdp_runtime
from hyper_parallel.trainer.runtime.data_iterator import HyperIter
from hyper_parallel.trainer.runtime.logging import enable_third_party_logging
from hyper_parallel.trainer.runtime.memory import empty_cache, print_device_mem_info
from hyper_parallel.trainer.runtime.memory_profiler import memory_profiler
from hyper_parallel.trainer.runtime.random import enable_high_precision_for_bf16, set_seed
from hyper_parallel.trainer.runtime.device import (  # pylint: disable=syntax-error
    get_device_type,
    get_torch_device,

        # General setup owns distributed initialization as its first stage.
        # ``_build_distributed_setup`` remains a reserved backend hook; the
        # trainer only defines its contract here.
        self._setup()
        try:
            memory_profiler.reset(
                self.config.memory,
                global_rank=self.global_rank,
                tp_rank=self.mesh.tp_rank,
                dp_rank=self.mesh.dp_rank,

            #   - consume the ``DistributedSetup`` created above;
            #   - set ``model`` and the resolved checkpoint ``model_config``;
            #   - return a materialized, weight-loaded, already-parallelized model;
            #   - never call a second trainer-side ``build_parallelize_model``.
            self._build_model()
            self._build_loss()

            # Build trainer-owned data components after the model finalizes parameters and sharding.
            self._build_model_assets()
            self._build_data_transform()
            self._build_dataset()
            self._build_data_batch_adapter()
            self._build_collate_fn()
            self._build_dataloader()
            self._compute_train_iters()

            self._build_optimizer()
            self._build_lr_scheduler()
            self._build_training_context()
            self._init_callbacks()
        except BaseException:
            memory_profiler.abort()
            raise

    def _setup(self):
        """Initialize logging, distributed state, and the local device."""
        # log args

        self.tqdm_callback = TqdmCallback(self)
        self.logging_callback = LoggingCallback(self)
        self.evaluate_callback = EvaluateCallback(self)
        self.garbage_collection_callback = GarbageCollectionCallback(self)
        self.profiler_callback = ProfilerCallback(self)
        self._callbacks = [
            self.environ_meter_callback,
            self.logging_callback,
            self.tqdm_callback,

        Args:
            data_iterator: Iterator providing the next group of micro-batches.
        """
        memory_profiler.step()
        config = self.config

        micro_batches: List[Dict[str, Any]] = next(data_iterator)
        self.state.global_step += 1

    def train(self) -> None:
        """Run the configured training loop."""
        config: TrainerConfig = self.config
        try:
            self.on_train_begin()
            self.data_iterator = HyperIter(
                self.train_dataloader, use_background_prefetcher=config.dataloader.use_background_prefetcher
            )
            logger.info(
                "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
                self.local_rank, self.state.global_step, self.train_iters, self.state.epoch, self.train_epochs,
            )

            start_epoch = self.state.epoch
            for epoch in range(start_epoch, self.train_epochs):
                if epoch != start_epoch:
                    self.train_dataloader.set_epoch(epoch)
                    self.data_iterator = HyperIter(
                        self.train_dataloader, use_background_prefetcher=config.dataloader.use_background_prefetcher
                    )
                self.state.epoch = epoch

                self.on_epoch_begin()

                start_step = self.state.global_step - epoch * self.train_steps
                train_steps = min(self.train_steps, self.train_iters - epoch * self.train_steps)
                for _ in range(start_step, train_steps):
                    try:
                        self.train_step(self.data_iterator)
                    except StopIteration:
                        logger.info(
                            "epoch:%s Dataloader finished with drop_last %s",
                            epoch,
                            config.dataloader.drop_last,
                        )
                        break

                self.on_epoch_end()
                self.state.epoch = epoch + 1

                print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")

                if config.dataloader.use_background_prefetcher:
                    self.data_iterator.stop()
        except BaseException:
            memory_profiler.abort()
            raise

        memory_profiler.stop()
        self.on_train_end()

        if config.dataloader.use_background_prefetcher:
            self.data_iterator.stop()

            return

        self._device_type, self._profiler_module = self._resolve_profiler_module()
        if self.config.mstx and self._device_type != "npu":
            raise ValueError("profiler.mstx is supported only when training on Ascend NPU")

        output_root = Path(self.config.output_path or "./output").expanduser()
        self.output_path = output_root / "profile" / f"rank_{rank}"
        logger.info("Profile save path: %s", self.output_path)

            **kwargs: Additional callback context supplied by the trainer.
        """
        del state, kwargs
        if self.profiler is None or not self.start_on_init:
            return
        self._start()
        self._record_metadata()

    def on_step_begin(

            **kwargs: Additional callback context supplied by the trainer.
        """
        del micro_batches, kwargs
        if self.profiler is None:
            return

        self._training_step += 1
        if not self._is_profiler_started:
            self._start()

            **kwargs: Additional callback context supplied by the trainer.
        """
        del state, loss, loss_dict, grad_norm, kwargs
        if self.profiler is None or not self._is_profiler_started:
            return

        self._close_mstx_range()
        self.profiler.step()
        if self._training_step == self.stop_step:

        Args:
            state: Final trainer state.
            **kwargs: Additional callback context supplied by the trainer.
        """
        del state, kwargs
        self._stop()

    def _validate_types(self) -> None:
        """Validate profiler fields that dataclass construction cannot constrain."""
        boolean_fields = (

            "mstx",
        )
        for name in boolean_fields:
            if not isinstance(getattr(self.config, name), bool):
                raise TypeError(f"profiler.{name} must be a bool")
        for name in ("start_step", "stop_step"):
            value = getattr(self.config, name)
            if isinstance(value, bool) or not isinstance(value, int):
                raise TypeError(f"profiler.{name} must be an int")
        if self.config.output_path is not None and not isinstance(self.config.output_path, str):
            raise TypeError("profiler.output_path must be a string or null")
        ranks = self.config.rank_ids
        if ranks is not None:
            if not isinstance(ranks, list):
                raise TypeError("profiler.rank_ids must be a list or null")
            if any(isinstance(rank, bool) or not isinstance(rank, int) or not 0 <= rank < self.trainer.world_size
                   for rank in ranks):
                raise ValueError(f"profiler.rank_ids must contain only integers in [0, {self.trainer.world_size})")

        """Normalize the inclusive profiling window to valid positive steps."""
        start_step = self.config.start_step
        stop_step = self.config.stop_step
        if start_step < 1:
            logger.warning("profiler.start_step must be greater than 0; reset it to 1")
            start_step = 1
        if stop_step < 1:
            logger.warning("profiler.stop_step must be greater than 0; reset it to 10")
            stop_step = 10
        if start_step > stop_step:
            logger.warning(
                "profiler.stop_step must be greater than or equal to profiler.start_step; reset both to 1 and 10"
            )

        return start_step, stop_step

    def _normalize_start_on_init(self) -> bool:
        if self.start_step != 1 and self.config.start_on_init:
            logger.warning(
                "profiler.start_step and profiler.start_on_init cannot take effect simultaneously; "
                "reset profiler.start_on_init to false"
            )
            return False
        return self.config.start_on_init

    def _normalize_level(self) -> int:
        level = self.config.level

    def _normalize_level(self) -> int:
        level = self.config.level
        if level is None:
            return 0
        if isinstance(level, bool) or not isinstance(level, int):
            raise TypeError("profiler.level must be an int or null")
        if level not in (0, 1, 2):
            logger.warning("Invalid profiler.level %s; reset it to 0", level)
            return 0
        return level

    def _get_schedule(self) -> Any:
        if self.start_on_init:

        pipeline_ranks: list[int] = []
        if self.config.pipeline_stage_leaders:
            pipeline_stages = max(1, self.trainer.config.accelerator.pp_size)
            if self.trainer.world_size % pipeline_stages != 0:
                raise ValueError("device count must be divisible by pipeline stage count")
            devices_per_stage = self.trainer.world_size // pipeline_stages
            pipeline_ranks = [stage * devices_per_stage for stage in range(pipeline_stages)]

        if not ranks and not pipeline_ranks:

            self._torch_npu = torch_npu
            return device_type, torch_npu.profiler
        if device_type in ("cpu", "cuda"):
            return device_type, torch.profiler
        raise RuntimeError(f"profiler does not support device type {device_type!r}")

    def _get_activities(self) -> list[Any]:
        activities = [self._profiler_module.ProfilerActivity.CPU]
        if self._device_type == "cuda":

        }
        try:
            self.profiler.add_metadata_json("distributed_args", json.dumps(metadata))
            self._metadata_recorded = True
        except AttributeError as error:
            logger.warning("Profiler failed to record distributed args: %s", error)

    def _close_mstx_range(self) -> None:
        if self._mstx_range_id is None:
            return

            dp_rank: Rank in the data-parallel group.
        """
        self._discard_previous_session()
        if config is None:
            return

        self._validate(config)
        self.enable = config.enable
        self.mem_info = config.mem_info

        self.enable = False
        self.mem_info = False
        try:
            self._stop_history()
        except Exception:  # pylint: disable=broad-exception-caught
            logger.exception("Failed to stop allocator memory history while aborting training.")
            return

        self._deactivate_session()

    def _discard_previous_session(self) -> None:

    @staticmethod
    def _validate_bool(value: bool, name: str) -> None:
        if not isinstance(value, bool):
            raise TypeError(f"memory.{name} must be a bool")

    @staticmethod
    def _validate_steps(start_step: int, end_step: int) -> None:
        for name, value in (("start_step", start_step), ("end_step", end_step)):

    @staticmethod
    def _validate_steps(start_step: int, end_step: int) -> None:
        for name, value in (("start_step", start_step), ("end_step", end_step)):
            if isinstance(value, bool) or not isinstance(value, int):
                raise TypeError(f"memory.{name} must be an int")
        if start_step < 0:
            raise ValueError("memory.start_step must be non-negative")
        if end_step < start_step:
            raise ValueError("memory.end_step must be greater than or equal to memory.start_step")

    @staticmethod
    def _validate_dump_ranks(dump_ranks: list[int]) -> None:
        if not isinstance(dump_ranks, list):
            raise ValueError("memory.dump_ranks must be a list of non-negative integers")
        if any(
                isinstance(rank, bool) or not isinstance(rank, int) or rank < 0
                for rank in dump_ranks
        ):

from hyper_parallel.data.text import build_chat_template
from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
from hyper_parallel.trainer.runtime.logging import create_logger
from hyper_parallel.trainer.runtime.memory import print_device_mem_info
from hyper_parallel.trainer.runtime.memory_profiler import memory_profiler
from hyper_parallel.trainer.runtime.device import synchronize
from hyper_parallel.trainer.base import BaseTrainer
from hyper_parallel.trainer.config import TrainerConfig

        self.base = BaseTrainer.__new__(BaseTrainer)
        self.base.config = config

        self.base._setup()
        try:
            memory_profiler.reset(
                config.memory,
                global_rank=self.base.global_rank,
                tp_rank=self.base.mesh.tp_rank,
                dp_rank=self.base.mesh.dp_rank,

                global_rank=self.base.global_rank,
                tp_rank=self.base.mesh.tp_rank,
                dp_rank=self.base.mesh.dp_rank,
            )
            self.base._build_model()
            self.base._build_loss()

            # datasets
            self._build_model_assets()
            self._build_data_transform()
            self.base._build_dataset()
            self.base._build_data_batch_adapter()

            # dataloader
            self._build_collate_fn()
            self.base._build_dataloader()

            # get_batch
            self._build_get_batch()
            self.base._compute_train_iters()

            self.base._build_optimizer()
            self.base._build_lr_scheduler()
            self.base._build_training_context()
            self.base._init_callbacks()
        except BaseException:
            memory_profiler.abort()
            raise

    def _build_model_assets(self) -> None:
        """Build tokenizer-backed assets for text training."""
        config: TrainerConfig = self.base.config

        Args:
            data_iterator: Iterator providing raw text batches.
        """
        memory_profiler.step()
        config = self.base.config
        num_micro_steps = self.base.num_micro_batches
        optimizers = self.base.optimizer if isinstance(self.base.optimizer, list) else [self.base.optimizer]

    def train(self) -> None:
        """Run the configured global optimizer steps by Dataset epoch."""
        config = self.base.config
        try:
            self.on_train_begin()
            logger.info(
                "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
                self.base.local_rank,
                self.base.state.global_step,
                self.base.train_iters,

                self.base.train_epochs,
            )

            # Checkpoint resume restores state.global_step, state.epoch, and the DataLoader cursor.
            for epoch in range(self.base.state.epoch, self.base.train_epochs):
                train_dataloader = self.base.train_dataloader
                if hasattr(train_dataloader, "set_epoch"):
                    train_dataloader.set_epoch(epoch)

                self.base.state.epoch = epoch
                self.on_epoch_begin()
                data_iterator = iter(train_dataloader) if train_dataloader is not None else None
                start_step = self.base.state.global_step - epoch * self.base.train_steps
                train_steps = min(self.base.train_steps, self.base.train_iters - epoch * self.base.train_steps)
                for _ in range(start_step, train_steps):
                    try:
                        self.train_step(data_iterator)
                    except StopIteration:
                        logger.info(
                            "epoch:%s Dataloader finished with drop_last %s",
                            epoch,
                            config.dataloader.drop_last,
                        )
                        break

                self.on_epoch_end()
                self.base.state.epoch = epoch + 1
                print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")
                if self.base.state.global_step >= self.base.train_iters:
                    break
        except BaseException:
            memory_profiler.abort()
            raise

        memory_profiler.stop()
        self.on_train_end()

        synchronize()
        self.base.destroy_distributed()

from hyper_parallel.data.vlm import build_processor, build_vlm_get_batch
from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
from hyper_parallel.trainer.runtime.logging import create_logger
from hyper_parallel.trainer.runtime.memory import print_device_mem_info
from hyper_parallel.trainer.runtime.memory_profiler import memory_profiler
from hyper_parallel.trainer.runtime.device import synchronize  # pylint: disable=syntax-error
from hyper_parallel.trainer.base import BaseTrainer
from hyper_parallel.trainer.config import TrainerConfig

        self.base = BaseTrainer.__new__(BaseTrainer)
        self.base.config = config

        self.base._setup()
        try:
            memory_profiler.reset(
                config.memory,
                global_rank=self.base.global_rank,
                tp_rank=self.base.mesh.tp_rank,
                dp_rank=self.base.mesh.dp_rank,

                global_rank=self.base.global_rank,
                tp_rank=self.base.mesh.tp_rank,
                dp_rank=self.base.mesh.dp_rank,
            )
            self.base._build_model()
            self.base._build_loss()

            # datasets
            self._build_model_assets()
            self._build_data_transform()
            self.base._build_dataset()
            self.base._build_data_batch_adapter()

            # dataloader
            self._build_collate_fn()
            self.base._build_dataloader()

            # get_batch
            self._build_get_batch()
            self.base._compute_train_iters()

            self.base._build_optimizer()
            self.base._build_lr_scheduler()
            self.base._build_training_context()
            self.base._init_callbacks()
        except BaseException:
            memory_profiler.abort()
            raise

    def _build_model_assets(self) -> None:
        """Build processor-backed assets for VLM training."""
        config: TrainerConfig = self.base.config

        Args:
            data_iterator: Iterator providing raw multimodal batches.
        """
        memory_profiler.step()
        config = self.base.config
        first_training_batch = self.base.get_batch(data_iterator)
        num_micro_steps = self.base.num_micro_batches
        training_batches = [first_training_batch]

    def train(self) -> None:
        """Run the VLM training loop."""
        config = self.base.config
        try:
            self.on_train_begin()
            logger.info(
                "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
                self.base.local_rank,
                self.base.state.global_step,
                self.base.train_iters,

                self.base.train_epochs,
            )

            # Checkpoint resume restores state.global_step, state.epoch, and the DataLoader cursor.
            for epoch in range(self.base.state.epoch, self.base.train_epochs):
                train_dataloader = self.base.train_dataloader

                if hasattr(train_dataloader, "set_epoch"):
                    train_dataloader.set_epoch(epoch)

                self.on_epoch_begin()
                data_iterator = iter(train_dataloader) if train_dataloader is not None else None

                start_step = self.base.state.global_step - epoch * self.base.train_steps
                train_steps = min(self.base.train_steps, self.base.train_iters - epoch * self.base.train_steps)
                for _ in range(start_step, train_steps):
                    try:
                        self.train_step(data_iterator)
                    except StopIteration:
                        logger.info(
                            "epoch:%s Dataloader finished with drop_last %s",
                            epoch,
                            config.dataloader.drop_last,
                        )
                        break

                self.on_epoch_end()
                self.base.state.epoch = epoch + 1
                print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")

                if self.base.state.global_step >= self.base.train_iters:
                    break
        except BaseException:
            memory_profiler.abort()
            raise

        memory_profiler.stop()
        self.on_train_end()

        synchronize()
        self.base.destroy_distributed()