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()

        Args:
            trainer: Trainer that owns the callback lifecycle.
        """
        super().__init__(trainer)
        self.config: ProfilingConfig = trainer.config.profiler
        self.profiler: Any = None
        self.output_path: Optional[Path] = None
        self._profiler_module: Any = None
        self._torch_npu: Any = None
        self._device_type: Optional[str] = None
        self._is_profiler_started = False
        self._metadata_recorded = False
        self._training_step = 0
        self._mstx_range_id: Optional[int] = None

        if not self.config.enabled:
            return

        self._validate_types()
        self.start_step, self.stop_step = self._normalize_steps()
        self.start_on_init = self._normalize_start_on_init()
        self.level = self._normalize_level()

        rank = trainer.global_rank
        if not self._is_profile_required(rank):
            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)

        schedule_config = self._get_schedule()
        profile_kwargs = {
            "activities": self._get_activities(),
            "profile_memory": self.config.memory,
            "with_stack": self.config.with_stack,
            "schedule": schedule_config,

            "with_stack": self.config.with_stack,
            "schedule": schedule_config,
            "on_trace_ready": self._profiler_module.tensorboard_trace_handler(str(self.output_path)),
        }
        if self._device_type == "npu":
            profile_kwargs["experimental_config"] = self._get_npu_experimental_config()
        self.profiler = self._profiler_module.profile(**profile_kwargs)

    def on_train_begin(self, state: TrainerState, **kwargs: Any) -> None:
        """Start collection before the first training step when requested.

        Args:
            state: Current trainer state.
            **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(
        self,
        state: TrainerState,

            state: Current trainer state.
            micro_batches: Optional micro-batches prepared for the current step.
            **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()
        if self._training_step == self.start_step:
            self._record_metadata()

        if self.config.mstx:
            mstx = self._torch_npu.npu.mstx
            step_num = state.global_step + 1
            self._mstx_range_id = mstx.range_start(
                f"step {step_num}",
                self._torch_npu.npu.current_stream(),
            )

            loss_dict: Optional named losses for the completed step.
            grad_norm: Optional gradient norm for the completed step.
            **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:
            logger.info("End of profiling. Analyze the trace under %s", self.output_path)
            self._stop()

    def on_train_end(self, state: TrainerState, **kwargs: Any) -> None:
        """Flush a partial profile when training ends before the 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 = (
            "enabled",
            "start_on_init",
            "memory",
            "pipeline_stage_leaders",

            "with_stack",
            "data_simplification",
            "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})")

    def _normalize_steps(self) -> tuple[int, int]:
        """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"
            )
            start_step, stop_step = 1, 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
        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:
            active_steps = self.stop_step
            skip_first = 1
        else:
            active_steps = self.stop_step - self.start_step + 1
            skip_first = self.start_step
        return self._profiler_module.schedule(
            wait=0,
            warmup=0,
            active=active_steps,
            repeat=1,

        )

    def _is_profile_required(self, rank: int) -> bool:
        """Return whether the current rank is selected for profiling."""
        ranks = self.config.rank_ids or []
        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:
            return True
        return rank in ranks or rank in pipeline_ranks

    def _resolve_profiler_module(self) -> tuple[str, Any]:
        """Resolve the profiler backend for the active device type."""
        device_type = get_device_type()
        if device_type == "npu":
            # torch-npu is optional and must not become a trainer import-time dependency.
            import torch_npu  # pylint: disable=import-outside-toplevel

            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":
            activities.append(self._profiler_module.ProfilerActivity.CUDA)
        elif self._device_type == "npu":
            activities.append(self._profiler_module.ProfilerActivity.NPU)
        return activities

    def _get_npu_experimental_config(self) -> Any:
        """Build the torch-npu experimental profiler configuration."""
        profiler_level = getattr(self._profiler_module.ProfilerLevel, f"Level{self.level}")
        options = {
            "profiler_level": profiler_level,
            "data_simplification": self.config.data_simplification,
            "mstx": self.config.mstx,
        }
        export_type = getattr(getattr(self._profiler_module, "ExportType", None), "Text", None)
        if export_type is not None:
            options["export_type"] = [export_type]
        # torch-npu exposes this version-dependent profiler entry point only under its private name.
        experimental_config = getattr(self._profiler_module, "_ExperimentalConfig")
        return experimental_config(**options)

    def _start(self) -> None:
        self.profiler.start()
        self.profiler.step()
        self._is_profiler_started = True

    def _record_metadata(self) -> None:
        """Record distributed topology metadata once per profiling session."""
        if self._metadata_recorded:
            return
        parallel = self.trainer.config.accelerator
        metadata = {
            "tensor_model_parallel_size": parallel.tp_size,
            "pipeline_model_parallel_size": parallel.pp_size,
            "data_parallel_size": getattr(self.trainer.mesh, "dp_size", 1),
            "expert_model_parallel_size": parallel.ep_size,

            "sequence_parallel": parallel.sequence_parallel,
            "parallel_mode": "distributed" if self.trainer.world_size > 1 else "stand_alone",
            "world_size": self.trainer.world_size,
        }
        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
        self._torch_npu.npu.mstx.range_end(self._mstx_range_id)
        self._mstx_range_id = None

    def _stop(self) -> None:
        self._close_mstx_range()
        profiler = self.profiler
        self.profiler = None
        if profiler is not None and self._is_profiler_started:
            profiler.stop()
        self._is_profiler_started = False


__all__ = ["ProfilerCallback"]

            global_rank: Rank in the global process group.
            tp_rank: Rank in the tensor-parallel group.
            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.current_step = 0
        self.start_step = config.start_step
        self.end_step = config.end_step
        self.save_path = config.save_path
        self.dump_ranks = list(config.dump_ranks)
        self.stacks = config.stacks
        self.max_entries = sys.maxsize if config.max_entries is None else config.max_entries
        self.global_rank = global_rank
        self.tp_rank = tp_rank
        self.dp_rank = dp_rank
        self._session_active = self.enable or self.mem_info

        if self.enable:
            self._resolve_device_apis(require_history=True)
            # Match MindSpeed's session-relative control point: step zero runs
            # after distributed setup and before model construction.
            self.step()

    def step(self) -> None:
        """Advance the state machine before a Trainer training step starts."""
        if not self._session_active:
            return

        try:
            if self.enable:
                if self.current_step == self.start_step:
                    self._record()
                if self.current_step == self.end_step:
                    try:
                        self._dump_and_stop_history()
                    finally:
                        self.enable = False
                        if not self.mem_info and not self._history_started:
                            self._deactivate_session()

            if self.mem_info:
                self._log_memory_info()
        except BaseException:
            self.abort()
            raise
        finally:
            self.current_step += 1

        if not self.enable and not self.mem_info:
            self._deactivate_session()

    def stop(self) -> None:
        """Dump an active partial window and end a normal training session."""
        if not self._session_active and not self._history_started:
            return

        try:
            try:
                if self._history_started:
                    self._dump_and_stop_history()
            finally:
                self.enable = False
                self.mem_info = False
                if not self._history_started:
                    self._deactivate_session()
        except BaseException:
            self.abort()
            raise

    def abort(self) -> None:
        """Stop allocator history without dumping or masking a Trainer error."""
        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:
        if self._history_started:
            self._stop_history()
        self._deactivate_session()

    def _deactivate_session(self) -> None:
        self.enable = False
        self.mem_info = False
        self._session_active = False
        self._device_api = None
        self._memory_api = None

    @staticmethod
    def _validate(config: MemoryConfig) -> None:
        MemoryProfiler._validate_bool(config.enable, "enable")
        MemoryProfiler._validate_bool(config.mem_info, "mem_info")
        MemoryProfiler._validate_steps(config.start_step, config.end_step)
        MemoryProfiler._validate_save_path(config.save_path)
        MemoryProfiler._validate_dump_ranks(config.dump_ranks)
        MemoryProfiler._validate_stacks(config.stacks)
        MemoryProfiler._validate_max_entries(config.max_entries)

    @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)):
            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_save_path(save_path: str) -> None:
        if not isinstance(save_path, str) or not save_path:
            raise ValueError("memory.save_path must be a non-empty string")

    @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
        ):
            raise ValueError("memory.dump_ranks must be a list of non-negative integers")

    @staticmethod
    def _validate_stacks(stacks: str) -> None:
        if stacks not in ("python", "all"):
            raise ValueError("memory.stacks must be 'python' or 'all'")

    @staticmethod
    def _validate_max_entries(max_entries: Optional[int]) -> None:
        if max_entries is not None and (
                isinstance(max_entries, bool)
                or not isinstance(max_entries, int)
                or max_entries <= 0
        ):
            raise ValueError("memory.max_entries must be a positive integer or None")

    def _record(self) -> None:
        """Start allocator memory history recording."""
        self._history_started = True
        try:
            self._memory_api._record_memory_history(
                stacks=self.stacks,
                max_entries=self.max_entries,
            )
        except Exception:
            # The private allocator API does not expose whether a failed call
            # enabled history, so make a best-effort stop before propagating.
            self.abort()
            raise

    def _dump_and_stop_history(self) -> None:
        try:
            self._dump()
        finally:
            self._stop_history()

    def _dump(self) -> None:
        """Dump the allocator snapshot for a configured rank."""
        if self.global_rank not in self.dump_ranks:
            return
        os.makedirs(self.save_path, exist_ok=True)
        timestamp = time.strftime("%Y-%m-%d-%H-%M")
        file_path = os.path.join(
            self.save_path,
            f"snapshot_{timestamp}_{self.global_rank}.pickle",
        )
        self._memory_api._dump_snapshot(file_path)
        logger.info("Memory snapshot dumped to %s", file_path)

    def _stop_history(self) -> None:
        if not self._history_started:
            return
        self._memory_api._record_memory_history(enabled=None)
        self._history_started = False

    def _log_memory_info(self) -> None:
        """Log and reset peak allocator statistics for the current step."""
        self._resolve_device_apis(require_history=False)
        logger.info(
            "Memory usage step=%s global_rank=%s tp_rank=%s dp_rank=%s "
            "max_memory_reserved=%s max_memory_allocated=%s",
            self.current_step,
            self.global_rank,

            self.dp_rank,
            self._device_api.max_memory_reserved(),
            self._device_api.max_memory_allocated(),
        )
        self._device_api.reset_peak_memory_stats()

    def _resolve_device_apis(self, *, require_history: bool) -> None:
        """Resolve and validate the active device memory APIs."""
        if self._device_api is None:
            device_type = get_device_type()
            if device_type not in ("cuda", "npu"):
                raise RuntimeError(
                    "memory profiling requires a CUDA or NPU device, "
                    f"but the active device type is {device_type!r}"
                )
            self._device_api = get_torch_device()

        if not require_history or self._memory_api is not None:
            return
        self._memory_api = getattr(self._device_api, "memory", None)
        required_methods = ("_record_memory_history", "_dump_snapshot")
        missing = [
            method
            for method in required_methods
            if self._memory_api is None
            or not callable(getattr(self._memory_api, method, None))

            for method in required_methods
            if self._memory_api is None
            or not callable(getattr(self._memory_api, method, None))
        ]
        if missing:
            raise RuntimeError(
                "memory profiling backend is missing allocator API(s): "
                + ", ".join(missing)
            )

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()