Diff Coverage

            target_ccfg.b,
            target_ccfg.vp,
        )
        if hasattr(target_ccfg.parser, "config_shard_emb"):
            target_ccfg.parser.config_shard_emb(target_ccfg)
        if hasattr(target_ccfg.parser, "config_shard_recompute"):
            target_ccfg.parser.config_shard_recompute(target_ccfg)
        target_ccfg.parser.config_dp_tp_exp(target_ccfg)
        target_ccfg.parser.config_optimizer_shard(target_ccfg)
        target_ccfg.parser.config_comm_flag(target_ccfg)
        if fr is not None:

        ``config_shard_emb`` call (guarded by ``hasattr``) and the initial
        ``shard_embed`` value computed in ``_init_shard`` is never refreshed,
        producing an embedding-memory mismatch versus the MF parser.
        """
        target_ccfg = self.ccfg if ccfg is None else ccfg
        target_ccfg.shard_embed = (
            target_ccfg.d
            if (target_ccfg.vocab_emb_dp and target_ccfg.p == 1)
            else (target_ccfg.t * target_ccfg.d)
        )

        computed at initial parse time (using the default ``t`` from the
        YAML), causing memory-estimation errors when the search explores
        strategies with different ``t`` values.
        """
        target_ccfg = self.ccfg if ccfg is None else ccfg
        target_ccfg.shard_recompute_input = (
            target_ccfg.t if self._recompute_slice_activation else 1
        )

    def _init_shard(self):

# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""parser child class"""
from typing import Any

from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config
from hyper_parallel.auto_parallel.sapp_nd.nd.common.framework_parsers._cost_model_parser import _CostModelParser
from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory

class CostModelParserMindformers(_CostModelParser):
    """parser class for MindFormers format"""

    def parse(self) -> None:
        """Parse the configured MindFormers YAML fields."""
        self.__config_parse_yaml()

    def __config_parse_yaml_parallelism(self):

        self.ccfg.rec_op.ffAct = int(
            not (self.config.recompute_config.select_recompute and self.ccfg.sp > 1)
        )

    def config_shard_emb(self, ccfg: Any = None) -> None:
        """Configure embedding and output activation sharding.

        Args:
            ccfg: Cost-model config to update. Defaults to the parser's config.

        Args:
            ccfg: Cost-model config to update. Defaults to the parser's config.
        """
        target_ccfg = self.ccfg if ccfg is None else ccfg
        target_ccfg.shard_embed = (
            target_ccfg.d
            if (target_ccfg.vocab_emb_dp and target_ccfg.p == 1)
            else (target_ccfg.t * target_ccfg.d)
        )

        Returns:
            loss: Loss value
        """
        if self.optimizer is None:
            self.compile(input_batch, label_batch)
        return self._compiler.forward_backward(input_batch, label_batch)

    def to(self, device: torch.device) -> "GraphTrainer":
        """Move the model to ``device`` and remember it for batch placement.

        self._local_mapping.update({key: value})

    def __delitem__(self, key: str) -> None:
        """Delete every registration for ``key`` from this registry."""
        if key not in self._local_mapping and key not in self._global_mapping:
            raise KeyError(key)
        self._local_mapping.pop(key, None)
        self._global_mapping.pop(key, None)

    def __iter__(self) -> Iterator[str]:
        """Iterate over all valid keys, local overrides taking precedence."""
        # Ensure we use all keys, with the overwritten ones on top

        return len(self._global_mapping.keys() | self._local_mapping.keys())

    def __contains__(self, key: object) -> bool:
        """Return whether ``key`` has a local or global registration."""
        return key in self._local_mapping or key in self._global_mapping

    def register(
        self,
        key: str,

        """Synchronize all ranks and tear down the distributed process group."""
        if not dist.is_available() or not dist.is_initialized():
            return

        try:
            empty_cache()
            dist.barrier()
            synchronize()
        finally:
            destroy_process_group()

    def train(self) -> None:
        """Run the configured training loop."""
        config: TrainerConfig = self.config
        self.data_iterator = None
        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}")

                self.data_iterator.stop()
                self.data_iterator = None
        finally:
            # ExitStack runs every callback even when an earlier cleanup raises.
            with ExitStack() as cleanup:
                cleanup.callback(self.destroy_distributed)
                cleanup.callback(synchronize)
                cleanup.callback(self.on_train_end)
                if self.data_iterator is not None:
                    cleanup.callback(self.data_iterator.stop)

    """Base callback bound to a Trainer through a weak reference."""

    def __init__(self, trainer: "BaseTrainer") -> None:
        """Bind the callback to its owning trainer."""
        self.trainer = weakref.proxy(trainer)
        self.mesh = trainer.mesh

    def on_step_begin(self, state: TrainerState, micro_batches: List[Dict[str, Any]] = None, **kwargs: Any) -> None:
        """Hook invoked at the start of each training step.