Diff Coverage

            #   - 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_collate_fn()
            self._build_dataloader()
            self._compute_train_iters()

            self._build_optimizer()
            self.model_integration.attach_optimizer(self.optimizer)
            self._build_lr_scheduler()
            self._build_training_context()
            self._init_callbacks()
        except BaseException:
            self.abort_profilers()
            raise

        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,

            try:
                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

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

                memory_profiler.stop()
                self.on_train_end()
            finally:

            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(

            "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
        ):

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

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

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

            self.base._build_optimizer()
            self.base.model_integration.attach_optimizer(self.base.optimizer)
            self.base._build_lr_scheduler()
            self.base._build_training_context()
            self.base._build_graph_compiler()
            self.base._init_callbacks()
        except BaseException:
            self.base.abort_profilers()
            raise

            )

            # 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

            memory_profiler.stop()
            self.on_train_end()

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

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

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

            self.base._build_optimizer()
            self.base.model_integration.attach_optimizer(self.base.optimizer)
            self.base._build_lr_scheduler()
            self.base._build_training_context()
            self.base._init_callbacks()
        except BaseException:
            self.base.abort_profilers()
            raise

            )

            # 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

            memory_profiler.stop()
            self.on_train_end()