base.py

# Copyright (c) ModelScope Contributors. All rights reserved.
import dataclasses
import logging
import megatron.core
import operator
import os
import shutil
import torch
import torch.nn
from abc import ABC, abstractmethod
from contextlib import contextmanager, nullcontext
from functools import partial
from mcore_bridge import LoraParallelLinear
from megatron.core import mpu
from megatron.core.distributed import DistributedDataParallel as DDP
from megatron.core.distributed import finalize_model_grads
from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
from megatron.core.pipeline_parallel import get_forward_backward_func
from megatron.core.transformer.module import MegatronModule
from megatron.core.transformer.moe.moe_utils import track_moe_metrics
from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper
from modelscope import check_local_model_is_latest
from packaging import version
from pathlib import Path
from transformers.utils import is_torch_npu_available
from typing import Callable, Dict, List, Optional

from swift.dataset import RowPreprocessor
from swift.megatron.callbacks import megatron_callbacks_map
from swift.megatron.model import get_mcore_model
from swift.megatron.utils import (apply_router_replay_patch, disable_forward_pre_hook, enable_forward_pre_hook,
                                  get_optimizer_param_scheduler, get_padding_to, init_persistent_async_worker,
                                  initialize_tp_communicators, load_mcore_checkpoint,
                                  logical_and_across_model_parallel_group, maybe_finalize_async_save,
                                  prepare_mcore_model, reconstruct_tensor_cp,
                                  reduce_max_stat_across_model_parallel_group, save_mcore_checkpoint,
                                  should_disable_forward_pre_hook, warmup_jit_function, wrap_model)
from swift.template import Template
from swift.trainers import dynamic_gradient_checkpointing
from swift.trainers.utils import patch_modelscope_hub_timeout
from swift.utils import (deep_getattr, gc_collect, get_current_device, get_last_valid_indices, get_logger, is_last_rank,
                         is_master, ms_logger_context)
from .batch_sampler import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
from .utils import TrainerState, build_streaming_dataloader, prepare_batch

try:
    from megatron.core.optimizer import param_group_identifier_keys
    from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
except ImportError:
    param_group_identifier_keys = None
    RouterReplay = None
    RouterReplayAction = None

mcore_016 = version.parse(megatron.core.__version__) >= version.parse('0.16.0rc0')

logger = get_logger()


class BaseMegatronTrainer(ABC):

    def __init__(self, args, template: Template):
        # validate mcore version and patch routing_replay
        self.enable_routing_replay = args.router_replay_mode != 'disabled'
        if self.enable_routing_replay:
            apply_router_replay_patch()

        self.args = args
        self.template = template
        self.prepare_model()
        # Sync template.padding_free after prepare_model(), because _check_padding_free
        # may override args.padding_free for certain models (e.g. DSA attention).
        if template.padding_free != args.padding_free:
            logger.warning(f'template.padding_free({template.padding_free}) != args.padding_free({args.padding_free}), '
                           f'syncing template.padding_free to {args.padding_free}.')
            template.padding_free = args.padding_free
        self.optimizer, self.opt_param_scheduler = self.get_optimizer_and_scheduler()
        self.data_collator = self._get_data_collator()

        self.state = TrainerState(max_steps=args.train_iters)
        initialize_embedding = args.new_special_tokens or args.task_type == 'seq_cls'
        if initialize_embedding:
            for m in self.unwrapped_models:
                self._initialize_embedding(m)
        self._load_checkpoint()

        self.eval_metrics = None
        if args.check_model and hasattr(args, 'model_dir'):
            with ms_logger_context(logging.CRITICAL), patch_modelscope_hub_timeout():
                config_info = self._collect_config_info()
                config_info.update({
                    'invoked_by': 'local_trainer',
                    'third_party': 'swift',
                    'trainer_class': self.__class__.__name__,
                    'trainer_backend': 'megatron',
                })
                check_local_model_is_latest(args.model_info.model_dir, user_agent=config_info)

        self.callbacks = []
        for callback in args.callbacks:
            self.callbacks.append(megatron_callbacks_map[callback](self))

        if args.tp_comm_overlap:
            initialize_tp_communicators(args, self.config)

        warmup_jit_function(self.config, args)

        if args.async_save and args.use_persistent_ckpt_worker:
            init_persistent_async_worker()

    def _load_checkpoint(self):
        args = self.args
        if not args.finetune:
            self.state.iteration = self._load_iteration()
        if args.mcore_model is not None:
            self.state.iteration = load_mcore_checkpoint(
                args, self.wrapped_models, self.optimizer, self.opt_param_scheduler, load_arg='mcore_model')
        if args.mcore_adapter is not None:
            self.state.iteration = load_mcore_checkpoint(
                args, self.wrapped_models, self.optimizer, self.opt_param_scheduler, load_arg='mcore_adapter')
        self.state.consumed_train_samples = getattr(args, 'consumed_train_samples', 0)

    def call_event(self, event, **kwargs):
        for callback in self.callbacks:
            getattr(callback, event)(**kwargs)

    def on_log(self, logs, prefix=''):
        n_steps = logs.pop('n_steps')
        self._log_callback(logs, n_steps)
        if prefix:
            logs = {f'{prefix}{k}': v for k, v in logs.items()}
        self._last_logged_metrics = dict(logs)
        self.call_event('on_log', logs=logs)

    def _log_callback(self, logs, n_steps):
        args = self.args
        config = self.config
        if config.num_moe_experts is not None:
            moe_loss_scale = 1 / args.num_microbatches / n_steps
            track_names = []
            if config.moe_router_load_balancing_type == 'aux_loss':
                track_names.append('load_balancing_loss')
            elif config.moe_router_load_balancing_type == 'seq_aux_loss':
                track_names.append('seq_load_balancing_loss')
            elif config.moe_router_load_balancing_type == 'global_aux_loss':
                track_names.append('global_load_balancing_loss')
            if config.moe_z_loss_coeff is not None:
                track_names.append('z_loss')
            track_moe_metrics(
                loss_scale=moe_loss_scale,
                iteration=self.state.iteration,
                writer=None,
                total_loss_dict=logs,
                force_initialize=True,
                track_names=track_names,
                num_layers=config.num_layers,
                moe_layer_freq=config.moe_layer_freq,
                mtp_num_layers=args.mtp_num_layers)
        if args.mtp_num_layers is not None:
            mtp_loss_scale = 1 / args.num_microbatches / n_steps
            mtp_logs = {}
            MTPLossLoggingHelper.track_mtp_metrics(mtp_loss_scale, self.state.iteration, None, None, mtp_logs)
            logs.update({k.replace(' ', '_'): v for k, v in mtp_logs.items()})
        # Track sparse attention indexer loss.
        if args.dsa_indexer_loss_coeff is not None and args.dsa_indexer_loss_coeff > 0:
            from megatron.core.transformer.experimental_attention_variant.dsa import DSAIndexerLossLoggingHelper
            indexer_loss_scale = 1 / args.num_microbatches / n_steps
            idx_logs = {}
            DSAIndexerLossLoggingHelper.track_indexer_metrics(
                loss_scale=indexer_loss_scale,
                iteration=self.state.iteration,
                writer=None,
                wandb_writer=None,
                total_loss_dict=idx_logs,
            )
            logs.update({k.replace(' ', '_'): v for k, v in idx_logs.items()})

        for k, v in logs.items():
            if isinstance(v, torch.Tensor):
                if v.numel() == 2:
                    v = v[0] / v[1]
                v = v.item()
            logs[k] = v

    def prepare_model(self):
        args = self.args
        self.unwrapped_models = get_mcore_model(args, self.template.config)
        self.config = self.unwrapped_models[0].config
        logger.info(f'model_config: {self.config}')
        self.bridge = self.config.bridge
        self.peft_models = self._prepare_peft_model(self.unwrapped_models)
        self.wrapped_models = wrap_model(args, self.unwrapped_models)

    def _prepare_peft_model(self, models):
        args = self.args
        if args.mcore_model is None:
            self.bridge.load_weights(models, args.model_dir)
        peft_models = [prepare_mcore_model(args, model) for model in models]
        if args.tuner_type == 'lora' and args.adapters and args.mcore_adapter is None:
            assert len(args.adapters) == 1, 'Currently only support one adapter.'
            self.bridge.load_weights(models, args.adapters[0], peft_format=True, adapter_name='default')
        return peft_models

    def get_optimizer_and_scheduler(self):
        args = self.args
        if mcore_016:
            from megatron.core.optimizer import AdamOptimizerConfig, SGDOptimizerConfig
            if args.optimizer == 'adam' or 'muon' in args.optimizer:
                # TODO(deyuf): Muon needs both adam + muon but get() only receive one config
                # So for now we keep using adam config that's back compat with old way
                config_cls = AdamOptimizerConfig
            elif args.optimizer == 'sgd':
                config_cls = SGDOptimizerConfig
            else:
                raise ValueError(f'Invalid optimizer type: {args.optimizer}')
        else:
            config_cls = OptimizerConfig

        kwargs = {
            f.name: getattr(args, f.name)
            for f in dataclasses.fields(config_cls) if hasattr(args, f.name) and f.name != 'loss_scale'
        }
        config = config_cls(**kwargs)

        if args.apply_wd_to_qk_layernorm or self.args.vit_lr is not None or self.args.aligner_lr is not None:
            param_groups_context = self._patch_get_param_groups()
        else:
            param_groups_context = nullcontext()
        with param_groups_context:

            if 'muon' not in config.optimizer:
                # If the user is asking for a non-zero embedding init std, skip weight decay for embeddings
                # to avoid embeddings from shrinking to zero as recommended in https://arxiv.org/abs/2312.16903
                # default_skip_embedding_weight_decay=args.embedding_init_method_std is not None,
                optimizer = get_megatron_optimizer(
                    config,
                    self.wrapped_models,
                )
            else:
                from megatron.core.optimizer.muon import get_megatron_muon_optimizer
                optimizer = get_megatron_muon_optimizer(
                    config,
                    self.wrapped_models,
                    layer_wise_distributed_optimizer='dist' in config.optimizer,
                )
        opt_param_scheduler = get_optimizer_param_scheduler(args, optimizer)
        return optimizer, opt_param_scheduler

    def _get_data_collator(self):
        data_collator = self.template.data_collator
        padding_to = get_padding_to(self.args)
        logger.info(f'padding_to: {padding_to}')
        data_collator = partial(data_collator, padding_to=padding_to)
        return data_collator

    def cyclic_iter(self, iterable, use_origin_cyclic: bool = False):
        training = self.unwrapped_models[0].training
        if not training or use_origin_cyclic:
            while True:
                for x in iterable:
                    yield x
            return

        args = self.args
        epoch = 0
        is_finished = False
        while True:
            if not is_finished:
                logger.info(f'The training of Epoch {epoch} starts...')
            for x in iterable:
                yield x
            # streaming
            if training and args.num_train_epochs and epoch >= args.num_train_epochs - 1:
                is_finished = True
            epoch += 1
            if is_finished:
                # Note that this approach will train for one additional step.
                logger.info(f'Training of {epoch} epochs has been completed, the training has finished.')
                args.train_iters = self.state.iteration + 1

    def _get_param_groups_mcore_016(
        self,
        model_chunks: List[MegatronModule],
        config: OptimizerConfig,
        config_overrides,
    ) -> List[Dict]:
        return self._get_param_groups(
            model_chunks,
            no_weight_decay_cond=None,
            scale_lr_cond=None,
            lr_mult=1.,
            lr=config.lr,
            min_lr=config.min_lr,
            decoupled_lr=config.decoupled_lr,
            decoupled_min_lr=config.decoupled_min_lr,
        )

    # Code borrowed from Megatron-LM
    def _get_param_groups(
        self,
        model_chunks: List[MegatronModule],
        no_weight_decay_cond: Optional[Callable],
        scale_lr_cond: Optional[Callable],
        lr_mult: float,
        lr: float,
        min_lr: float,
        decoupled_lr: Optional[float],
        decoupled_min_lr: Optional[float],
        default_skip_embedding_weight_decay: bool = False,
    ) -> List[Dict]:
        """Create parameter groups for optimizer.

        Creates parameter groups based on weight decay condition (regularized vs
        non regularized), learning rate scale condition (lr vs lr_mult * lr),
        and whether it is expert parameters. scale_lr_cond is used during finetuning
        where head of the network requires a scaled version of the base learning rate.

        Args:
            model_chunks (List[MegatronModule]): model chunks to create parameter
                groups for.
            no_weight_decay_cond (func, optional): function to determine whether a
                parameter should not perform weight decay.
            scale_lr_cond (func, optional): function to determine whether a parameter
                should have a scaled learning rate.
            lr_mult (float): learning rate multiplier for parameters that
                satisfy scale_lr_cond.
            lr (float): learning rate.
            min_lr (float): minimum learning rate.
            decoupled_lr (Optional[float]): optional decoupled learning rate.
            decoupled_min_lr (Optional[float]): optional decoupled minimum learning rate.
            default_skip_embedding_weight_decay (bool): whether to skip weight decay for embedding
                parameters by default, if no_weight_decay_cond is not provided.

        Returns:
            List of parameter groups.
        """
        args = self.args
        if decoupled_min_lr is None:
            decoupled_min_lr = min_lr
        is_multimodal = args.megatron_model_meta.is_multimodal
        if args.vit_lr is not None or args.aligner_lr is not None:
            assert is_multimodal, 'vit_lr and aligner_lr are only supported for multimodal models.'
            vit_lr = args.vit_lr if args.vit_lr is not None else args.lr
            aligner_lr = args.aligner_lr if args.aligner_lr is not None else args.lr
            logger.info_once(f'vit_lr: {vit_lr}, aligner_lr: {aligner_lr}, llm_lr: {args.lr}')
        use_decoupled_learning_rate = decoupled_lr is not None

        # Map (wd_mult, lr_mult, is_expert_parallel, is_decoupled_lr) to params.
        params_map = {}
        for model_chunk in model_chunks:
            visual = model_chunk.module.module.visual if is_multimodal else None
            for name, param in model_chunk.named_parameters():
                if not param.requires_grad:
                    continue

                is_expert_parallel = not getattr(param, 'allreduce', True)

                if no_weight_decay_cond is not None:
                    no_wd: bool = no_weight_decay_cond(name, param)
                elif args.apply_wd_to_qk_layernorm and any(
                        name.endswith(k) for k in ['q_layernorm.weight', 'k_layernorm.weight']):
                    no_wd = False
                else:
                    # Do not regularize biases and norm parameters.
                    #  optionally, also skip weight decay for embedding parameters if requested
                    #  (useful if you do not want embeddings to shrink to zero in training
                    #  https://arxiv.org/abs/2312.16903)
                    no_wd = (
                        name.endswith('.bias') or len(param.shape) == 1
                        or (default_skip_embedding_weight_decay and 'embedding' in name))
                _lr_mult = lr_mult
                if scale_lr_cond is not None:
                    scale_lr = scale_lr_cond(name, param)
                else:
                    scale_lr = False
                    # Handling multimodal models: vit_lr, aligner_lr
                    unwrapped_name = name.removeprefix('module.').removeprefix('module.')
                    if visual is not None:
                        is_aligner = any(unwrapped_name.startswith(f'visual.{k}') for k in visual._aligner or [])
                        is_vit = any(unwrapped_name.startswith(f'visual.{k}')
                                     for k in visual._vision_tower) and not is_aligner
                    else:
                        is_aligner, is_vit = False, False
                    if is_vit and args.vit_lr:
                        scale_lr = True
                        _lr_mult = args.vit_lr / lr
                    elif is_aligner and args.aligner_lr:
                        scale_lr = True
                        _lr_mult = args.aligner_lr / lr

                if not no_wd and not scale_lr:
                    wd_mult, _lr_mult = 1.0, 1.0
                elif not no_wd and scale_lr:
                    wd_mult, _lr_mult = 1.0, _lr_mult
                elif no_wd and not scale_lr:
                    wd_mult, _lr_mult = 0.0, 1.0
                else:
                    wd_mult, _lr_mult = 0.0, _lr_mult

                is_decoupled_lr = False
                # For input/embedding and output layer: embedding.word_embeddings.weight /
                # output_layer.weight.
                if use_decoupled_learning_rate and getattr(param, 'is_embedding_or_output_parameter', False):
                    is_decoupled_lr = True

                key = (wd_mult, _lr_mult, is_expert_parallel, is_decoupled_lr)
                if key not in params_map:
                    params_map[key] = []
                params_map[key].append(param)

        # Distributed checkpoint requires all ranks to have the same param groups,
        # so we need to align the param groups across ranks, otherwise we may have
        # runtime error when loading the checkpoint or numerical error when resuming training.
        params_key = list(params_map.keys())
        gathered_params_key = [None for _ in range(torch.distributed.get_world_size())]
        torch.distributed.all_gather_object(gathered_params_key, params_key)
        for keys in gathered_params_key:
            for key in keys:
                if key not in params_key:
                    params_key.append(key)

        param_groups = []
        for key in params_key:
            wd_mult, _lr_mult, is_expert_parallel, is_decoupled_lr = key
            params = params_map[key] if key in params_map else []
            param_group = {
                'params': params,
                'wd_mult': wd_mult,
                'lr_mult': _lr_mult,
                'is_expert_parallel': is_expert_parallel,
                'is_decoupled_lr': is_decoupled_lr,
            }
            # Ensure param_group has required keys for matching when loading optimizer state
            # See MegatronOptimizer._filter_and_reorder_param_groups.
            if param_group_identifier_keys is not None:
                assert set(param_group.keys()) - set(param_group_identifier_keys) == {'params'}
            param_groups.append(param_group)

        # Update min and max lr in param groups
        # These changes are compatible with mcore 0.16.
        for param_group in param_groups:
            if param_group['is_decoupled_lr']:
                assert decoupled_lr is not None
                param_group['max_lr'] = decoupled_lr
                param_group['min_lr'] = decoupled_min_lr
            else:
                param_group['max_lr'] = lr
                param_group['min_lr'] = min_lr
            lr_mult = param_group.pop('lr_mult')
            param_group['max_lr'] *= lr_mult
            param_group['min_lr'] *= lr_mult
        return param_groups

    @contextmanager
    def _patch_get_param_groups(self):
        from megatron.core import optimizer
        _get_param_groups = optimizer._get_param_groups
        if mcore_016:
            optimizer._get_param_groups = self._get_param_groups_mcore_016
        else:
            optimizer._get_param_groups = self._get_param_groups
        try:
            yield
        finally:
            optimizer._get_param_groups = _get_param_groups

    def _load_iteration(self):
        args = self.args
        ckpt_dir = None
        if args.tuner_type == 'full':
            ckpt_dir = args.model
        else:
            ckpt_dir = args.adapters[0] if args.adapters else args.model
        if ckpt_dir is None:
            return 0
        logger.info(f'checkpoint_dir: {ckpt_dir}')
        tracker_path = os.path.join(ckpt_dir, 'latest_checkpointed_iteration.txt')
        if not os.path.exists(tracker_path):
            return 0
        with open(tracker_path, 'r') as f:
            iteration = int(f.read())

        common_path = os.path.join(ckpt_dir, f'iter_{iteration:07d}', 'common.pt')
        if not os.path.exists(common_path):
            return iteration

        state_dict = torch.load(common_path)
        if 'args' not in state_dict:
            return iteration
        self.args.consumed_train_samples = getattr(state_dict['args'], 'consumed_train_samples', 0)

        return iteration

    def _prepare_vit_gradient_checkpointing(self, model):
        visual = model.visual
        if visual is None:
            return
        for vision_tower in visual._vision_tower:
            module = deep_getattr(visual, vision_tower)
            if self.args.vit_gradient_checkpointing:
                dynamic_gradient_checkpointing(module, False)
                try:
                    module.gradient_checkpointing_enable(
                        gradient_checkpointing_kwargs=self.args.vit_gradient_checkpointing_kwargs)
                    module.enable_input_require_grads()
                except AttributeError:
                    pass

    @staticmethod
    def _initialize_embedding(model):
        # compat new_special_tokens
        init_method = model.config.init_method
        if hasattr(model, 'language_model'):
            model = model.language_model
        for key in ['embedding.word_embeddings', 'output_layer']:
            if key == 'output_layer' and model.share_embeddings_and_output_weights:
                continue
            module = deep_getattr(model, key)
            if module is None:
                continue
            initialize_mask = (module.weight == 0).all(dim=-1)
            num_to_initialize = initialize_mask.sum().item()
            if num_to_initialize == 0:
                continue
            logger.info_if(f'num_to_initialize: {num_to_initialize}', cond=mpu.get_data_parallel_rank() == 0)
            tensor = module.weight.new_empty(num_to_initialize, module.weight.shape[1])
            module.weight.data[initialize_mask] = init_method(tensor)
            if getattr(module.weight, 'main_param', None) is not None:
                module.weight.main_param.copy_(module.weight.view(-1))

    def _all_reduce_metric(self,
                           metric: Dict[str, torch.Tensor],
                           reduction=torch.distributed.ReduceOp.AVG,
                           group=None) -> Dict[str, torch.Tensor]:
        if group is None:
            group = mpu.get_data_parallel_group()
        reporting_metric = torch.stack(list(metric.values()), dim=0)
        torch.distributed.all_reduce(reporting_metric, reduction, group=group)
        return {k: reporting_metric[i] for i, k in enumerate(metric.keys())}

    def merge_lora_adapters(self, adapter_name='default'):
        """Merge LoRA adapters into base model weights for vLLM inference."""
        with torch.no_grad():
            for model in self.unwrapped_models:
                for module in model.modules():
                    if isinstance(module, LoraParallelLinear):
                        # Merge all active adapters
                        module.merge(adapter_names=[adapter_name])

    def unmerge_lora_adapters(self):
        """Unmerge LoRA adapters to restore training state."""
        with torch.no_grad():
            for model in self.unwrapped_models:
                for module in model.modules():
                    if isinstance(module, LoraParallelLinear):
                        # Unmerge to restore separate LoRA weights for training
                        module.unmerge()

    @staticmethod
    def copy_path(src_path: str, tgt_path: str):
        if not is_master():
            return
        if not os.path.exists(src_path):
            raise FileNotFoundError(f'Source path does not exist: {src_path}')

        if os.path.isfile(src_path):
            os.makedirs(os.path.dirname(tgt_path), exist_ok=True)
            shutil.copy(src_path, tgt_path)
        elif os.path.isdir(src_path):
            shutil.copytree(src_path, tgt_path, dirs_exist_ok=True)
        else:
            raise ValueError(f'Source path is neither a file nor a directory: {src_path}')

    def _prepare_data_iterator(self, train_dataset, val_dataset=None, use_origin_cyclic: bool = False):
        train_dataloader, val_dataloader = self._prepare_dataloader(train_dataset, val_dataset)
        train_data_iterator = iter(self.cyclic_iter(train_dataloader, use_origin_cyclic=use_origin_cyclic))
        val_data_iterator = None
        if val_dataset is not None:
            val_data_iterator = iter(self.cyclic_iter(val_dataloader, use_origin_cyclic=use_origin_cyclic))
        return train_data_iterator, val_data_iterator

    def setup_model_training(self):
        args = self.args
        config = self.config

        for m in self.wrapped_models:
            m.train()

        if args.is_multimodal:
            for m in self.unwrapped_models:
                self._prepare_vit_gradient_checkpointing(m)

        config.grad_scale_func = self.optimizer.scale_loss
        if isinstance(self.wrapped_models[0], DDP) and args.overlap_grad_reduce:
            assert config.no_sync_func is None, ('When overlap_grad_reduce is True, config.no_sync_func must be None; '
                                                 'a custom no_sync_func is not supported when overlapping grad-reduce')
            config.no_sync_func = [model_chunk.no_sync for model_chunk in self.wrapped_models]
            if len(self.wrapped_models) == 1:
                config.no_sync_func = config.no_sync_func[0]
            if args.align_grad_reduce:
                config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in self.wrapped_models]
                if len(self.wrapped_models) == 1:
                    config.grad_sync_func = config.grad_sync_func[0]
        if args.overlap_param_gather and args.align_param_gather:
            config.param_sync_func = [model_chunk.start_param_sync for model_chunk in self.wrapped_models]
            if len(self.wrapped_models) == 1:
                config.param_sync_func = config.param_sync_func[0]
        config.finalize_model_grads_func = finalize_model_grads

        self._start_iteration = self.state.iteration
        self._pre_hook_enabled = False
        if should_disable_forward_pre_hook(args):
            disable_forward_pre_hook(self.wrapped_models, param_sync=False)
            self._saved_param_sync_func = config.param_sync_func
            config.param_sync_func = None

        self.call_event('on_train_begin')
        self._train_metrics = {}

    def setup_training(self, train_dataset, val_dataset):
        self.setup_model_training()
        args = self.args
        if args.virtual_pipeline_model_parallel_size is not None:
            train_data_iterator, val_data_iterator = [], []
            for _ in range(args.virtual_pipeline_model_parallel_size):
                train_it, val_it = self._prepare_data_iterator(train_dataset, val_dataset)
                train_data_iterator.append(train_it)
                val_data_iterator.append(val_it)
        else:
            train_data_iterator, val_data_iterator = self._prepare_data_iterator(train_dataset, val_dataset)
        return train_data_iterator, val_data_iterator

    def run_train_step(self, train_data_iterator, val_data_iterator):
        """Execute one training iteration including logging / eval / save.

        Returns ``True`` if training should continue, ``False`` if done.
        """
        args = self.args
        config = self.config
        state = self.state

        self.call_event('on_step_begin')
        maybe_finalize_async_save(args, blocking=False)
        metrics, grad_norm, update_successful = self.train_step(train_data_iterator)

        if state.iteration == self._start_iteration:
            if update_successful:
                if should_disable_forward_pre_hook(args):
                    enable_forward_pre_hook(self.wrapped_models)
                    config.param_sync_func = self._saved_param_sync_func
                    self._pre_hook_enabled = True
            else:
                self._start_iteration = state.iteration + 1

        state.iteration += 1
        self.call_event('on_step_end')
        self._aggregated_metrics(metrics, self._train_metrics)
        self._train_metrics['grad_norm'] = grad_norm
        for param_group in self.optimizer.param_groups:
            if len(param_group['params']) == 0:
                continue
            self._train_metrics['learning_rate'] = param_group['lr']
            break
        if state.should_log:
            state.should_log = False
            self.on_log(logs=self._train_metrics)
            self._train_metrics = {}

        eval_metrics = None
        if state.should_eval:
            state.should_eval = False
            if should_disable_forward_pre_hook(args):
                disable_forward_pre_hook(self.wrapped_models)
                self._pre_hook_enabled = False
            eval_metrics = self.evaluate(val_data_iterator)
            for m in self.wrapped_models:
                m.train()
            if should_disable_forward_pre_hook(args):
                enable_forward_pre_hook(self.wrapped_models)
                self._pre_hook_enabled = True

        if state.should_save:
            self._determine_best_metric(eval_metrics)
            if should_disable_forward_pre_hook(args):
                disable_forward_pre_hook(self.wrapped_models)
            state.should_save = False
            self.save_checkpoint()
            self.call_event('on_save', output_dir=self.state.last_model_checkpoint)
            if should_disable_forward_pre_hook(args):
                enable_forward_pre_hook(self.wrapped_models)

        return state.iteration < args.train_iters

    def finalize_training(self):
        """Cleanup after training loop completes."""
        self.call_event('on_train_end')
        if self._pre_hook_enabled:
            disable_forward_pre_hook(self.wrapped_models)
        maybe_finalize_async_save(self.args, blocking=True, terminate=True)

    def train(self, train_dataset, val_dataset):
        train_data_iterator, val_data_iterator = self.setup_training(train_dataset, val_dataset)
        while self.state.iteration < self.args.train_iters:
            self.run_train_step(train_data_iterator, val_data_iterator)
        self.finalize_training()

    def _determine_best_metric(self, metrics) -> bool:
        args = self.args
        state = self.state
        if (args.metric_for_best_model is None or metrics is None or not is_last_rank()
                or args.metric_for_best_model not in metrics):
            if args.metric_for_best_model is None:
                return False
            tensor = torch.zeros((3, ), device=get_current_device(), dtype=torch.float64)
            torch.distributed.all_reduce(tensor)
            is_new_best_metric = bool(tensor[2].item())
            if is_new_best_metric:
                state.best_metric = tensor[0].item()
                state.best_global_step = int(tensor[1].item())
            return is_new_best_metric
        metric_value = metrics[args.metric_for_best_model]
        op = operator.ge if args.greater_is_better else operator.le
        if state.best_metric is None:
            state.best_metric = float('-inf') if args.greater_is_better else float('inf')

        is_new_best_metric = False
        if op(metric_value, state.best_metric):
            state.best_metric = metric_value
            state.best_global_step = state.global_step
            is_new_best_metric = True
        tensor = torch.tensor([state.best_metric, state.best_global_step, is_new_best_metric],
                              device=get_current_device(),
                              dtype=torch.float64)
        torch.distributed.all_reduce(tensor)
        return is_new_best_metric

    def save_checkpoint(self):
        args = self.args
        state = self.state
        args.consumed_train_samples = state.consumed_train_samples
        iteration = state.iteration
        output_dir = os.path.join(args.output_dir, f'checkpoint-{iteration}')
        os.makedirs(output_dir, exist_ok=True)
        args_path = os.path.join(os.path.dirname(output_dir), 'args.json')
        self.copy_path(args_path, os.path.join(output_dir, 'args.json'))
        if args.save_safetensors and args.no_save_optim:
            model = []
        else:
            model = self.wrapped_models
        gc_collect()
        save_mcore_checkpoint(
            args,
            model,
            self.optimizer,
            self.opt_param_scheduler,
            iteration=iteration,
            peft_format=args.tuner_type == 'lora',
            output_dir=output_dir)
        state.last_model_checkpoint = output_dir
        if state.best_global_step is not None:
            best_model_checkpoint = os.path.join(args.output_dir, f'checkpoint-{state.best_global_step}')
            if os.path.exists(best_model_checkpoint):
                state.best_model_checkpoint = best_model_checkpoint
        # safetensors
        if args.save_safetensors:
            skip_saving_adapter = args.tuner_type == 'lora_llm' or (
                args.tuner_type == 'lora' and args.merge_lora and not hasattr(self.bridge, '_support_hf_grouped_lora'))

            if not skip_saving_adapter:
                self.bridge.save_weights(
                    self.unwrapped_models,
                    output_dir,
                    peft_format=args.tuner_type == 'lora',
                    args=args,
                    processor=self.template.processor,
                )
            # merge-lora does not store lora, lora saving may report an error (Qwen3-VL-Moe)
            if args.tuner_type != 'full' and args.merge_lora:
                self.merge_lora_adapters()
                origin_output_dir = output_dir
                output_dir = f'{output_dir}-merged'
                os.makedirs(output_dir, exist_ok=True)
                for fname in ['latest_checkpointed_iteration.txt', 'args.json']:
                    src_path = os.path.join(origin_output_dir, fname)
                    self.copy_path(src_path, os.path.join(output_dir, fname))
                # common.pt
                common_path = os.path.join(origin_output_dir, f'iter_{iteration:07d}', 'common.pt')
                tgt_common_path = os.path.join(output_dir, f'iter_{iteration:07d}', 'common.pt')
                os.makedirs(os.path.dirname(tgt_common_path), exist_ok=True)
                self.copy_path(common_path, tgt_common_path)
                self.bridge.save_weights(
                    self.unwrapped_models,
                    output_dir,
                    peft_format=False,
                    args=args,
                    processor=self.template.processor,
                )
                self.unmerge_lora_adapters()

        if is_master():
            self._rotate_checkpoints(args.output_dir)

    def _rotate_checkpoints(self, output_dir: str):
        # Code borrowed from huggingface/transformers
        args = self.args
        if args.save_total_limit is None or args.save_total_limit <= 0:
            return

        checkpoints_sorted = self._sorted_checkpoints(output_dir)
        if len(checkpoints_sorted) <= args.save_total_limit:
            return

        number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - args.save_total_limit)
        checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
        for checkpoint in checkpoints_to_be_deleted:
            shutil.rmtree(checkpoint, ignore_errors=True)
            if os.path.exists(f'{checkpoint}-merged'):
                shutil.rmtree(f'{checkpoint}-merged', ignore_errors=True)

    def _sorted_checkpoints(self, output_dir: str):
        # Code borrowed from huggingface/transformers
        state = self.state
        glob_checkpoints = [
            str(p) for p in Path(output_dir).glob('checkpoint-*') if p.is_dir() and not p.name.endswith('-merged')
        ]
        # Sort by modification time
        checkpoints_sorted = sorted(glob_checkpoints, key=os.path.getmtime)

        # Make sure we don't delete the best model.
        if state.best_model_checkpoint is not None and state.best_model_checkpoint in checkpoints_sorted:
            best_model_index = checkpoints_sorted.index(state.best_model_checkpoint)
            checkpoints_sorted.pop(best_model_index)
            checkpoints_sorted.append(state.best_model_checkpoint)
        return checkpoints_sorted

    def training_log(self, metrics, grad_norm):
        learning_rate = None
        for param_group in self.optimizer.param_groups:
            if len(param_group['params']) == 0:
                continue
            learning_rate = param_group['lr']
        logger.info(f'metrics: {metrics}, grad_norm: {grad_norm}, learning_rate: {learning_rate}')

    def evaluate(self, val_data_iterator):
        args = self.args
        for m in self.wrapped_models:
            m.eval()
        eval_metrics = {}
        forward_backward_func = get_forward_backward_func()
        self.call_event('on_eval_begin')
        with torch.no_grad():
            for _ in range(args.eval_iters):
                data_iterator = self._replace_data_iterator(val_data_iterator)
                metrics = forward_backward_func(
                    forward_step_func=self.forward_step,
                    data_iterator=data_iterator,
                    model=self.wrapped_models,
                    num_microbatches=self.args.num_microbatches,
                    seq_length=args.seq_length,
                    micro_batch_size=args.micro_batch_size,
                    forward_only=True,
                )
                self.call_event('on_eval_step')
                self._aggregated_metrics(metrics, eval_metrics)
        self.compute_eval_metrics(eval_metrics)
        self.on_log(logs=eval_metrics, prefix='eval_')
        self.call_event('on_eval_end')
        return eval_metrics

    def compute_eval_metrics(self, metrics):
        if self.eval_metrics is not None and mpu.is_pipeline_last_stage():
            metric = self.eval_metrics.compute()
            for k, v in metric.items():
                metrics[k] = v if isinstance(v, torch.Tensor) else torch.tensor(v)
            self.eval_metrics.reset()

    def _replace_data_iterator(self, data_iterator):
        return data_iterator

    def train_step(self, train_data_iterator):
        args = self.args
        forward_backward_func = get_forward_backward_func()
        for m in self.wrapped_models:
            m.zero_grad_buffer()
        self.optimizer.zero_grad()
        # TODO: refactor _replace_data_iterator
        data_iterator = self._replace_data_iterator(train_data_iterator)

        if self.enable_routing_replay:
            RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)

        metrics = forward_backward_func(
            forward_step_func=self.forward_step,
            data_iterator=data_iterator,
            model=self.wrapped_models,
            num_microbatches=args.num_microbatches,
            seq_length=args.seq_length,
            micro_batch_size=args.micro_batch_size,
            forward_only=False,
        )

        update_successful, grad_norm, _ = self.optimizer.step()
        update_successful = logical_and_across_model_parallel_group(update_successful)
        grad_norm = reduce_max_stat_across_model_parallel_group(grad_norm)
        if update_successful:
            self.opt_param_scheduler.step(increment=args.global_batch_size)

        if self.enable_routing_replay:
            RouterReplay.clear_global_router_replay_action()
            RouterReplay.clear_global_indices()

        return metrics, grad_norm, update_successful

    def _aggregated_metrics(self, metrics, total_metrics):
        if 'n_steps' not in total_metrics:
            total_metrics['n_steps'] = 0
        total_metrics['n_steps'] += 1
        if not metrics:
            return
        metrics = RowPreprocessor.rows_to_batched(metrics)
        for key, val in metrics.items():
            val = torch.stack([v for v in val if v is not None], dim=0)
            if val[0].numel() == 2:
                val = val.sum(dim=0)
                if val[1] == 0:
                    continue
            elif val[0].numel() == 1:
                val = val.new_tensor([val.sum(), val.shape[0]])
            else:
                raise ValueError(f'Invalid value shape: {val[0].shape} for key {key}')
            if key not in total_metrics:
                total_metrics[key] = torch.tensor([0.0, 0.0], dtype=torch.float32, device=torch.cuda.current_device())
            total_metrics[key] += val

    def _prepare_dataloader(self, train_dataset, val_dataset=None):
        args = self.args
        val_dataloader = None
        if args.streaming:
            train_dataloader = build_streaming_dataloader(args, train_dataset, self.data_collator)
            if val_dataset is not None:
                val_dataloader = build_streaming_dataloader(args, val_dataset, self.data_collator)
            return train_dataloader, val_dataloader
        train_batch_sampler = MegatronPretrainingRandomSampler(
            train_dataset,
            total_samples=len(train_dataset),
            consumed_samples=self.state.consumed_train_samples,
            micro_batch_size=args.micro_batch_size,
            data_parallel_rank=mpu.get_data_parallel_rank(),
            data_parallel_size=mpu.get_data_parallel_world_size(),
            data_sharding=args.data_sharding,
            shuffle=args.train_dataloader_shuffle,
            group_by_length=args.group_by_length,
        )
        train_dataloader = self._create_dataloader(train_dataset, train_batch_sampler)
        if val_dataset is not None:
            val_batch_sampler = MegatronPretrainingSampler(
                total_samples=len(val_dataset),
                consumed_samples=0,
                micro_batch_size=args.micro_batch_size,
                data_parallel_rank=mpu.get_data_parallel_rank(),
                data_parallel_size=mpu.get_data_parallel_world_size(),
            )
            val_dataloader = self._create_dataloader(val_dataset, val_batch_sampler)
        return train_dataloader, val_dataloader

    def _create_dataloader(self, dataset, batch_sampler):
        args = self.args

        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_sampler=batch_sampler,
            num_workers=args.dataloader_num_workers,
            pin_memory=args.dataloader_pin_memory,
            persistent_workers=args.dataloader_persistent_workers if args.dataloader_num_workers > 0 else False,
            prefetch_factor=args.dataloader_prefetch_factor if args.dataloader_num_workers > 0 else None,
            collate_fn=self.data_collator,
        )
        return dataloader

    @abstractmethod
    def forward_step(self, data_iterator, model):
        pass

    def _should_use_npu_generated_attention_mask(self, args) -> bool:
        return (is_torch_npu_available() and args.task_type == 'causal_lm' and not args.padding_free
                and getattr(args, 'attention_backend', None) != 'local' and getattr(args, 'use_flash_attn', False))

    def _prepare_batch(self, data, vp_stage=None):
        return prepare_batch(self.args, data, vp_stage=vp_stage)

    def get_batch(self, data_iterator, vp_stage=None):
        """Generate a batch."""
        return self._prepare_batch(next(data_iterator), vp_stage)

    def _collect_config_info(self) -> Dict[str, str]:
        """
        Collects trainer-specific configuration details.

        Subclasses can override this method to provide additional configuration
        information for model compatibility verification.

        Returns:
            Dict[str, str]: Configuration parameters as key-value pairs.
        """
        if self.__class__.__name__ == 'MegatronTrainer':
            if not self.template.use_chat_template:
                return {
                    'seq2seq_mode': 'pt',
                }
            else:
                return {
                    'seq2seq_mode': 'sft',
                }
        return {}

    def get_last_tokens(self, output_tensor, packed_seq_params=None, attention_mask=None, num_samples=None):
        if self.args.context_parallel_size > 1:
            output_tensor = reconstruct_tensor_cp(output_tensor, packed_seq_params, dim=1)
        if packed_seq_params is None:
            # Compatible with attention_mask_2d
            if attention_mask.dim() > 2:
                attention_mask = (~attention_mask).sum(dim=(1, 2)) > 0
            last_token_idx = get_last_valid_indices(attention_mask.long())
            last_tokens = output_tensor[torch.arange(output_tensor.shape[0]), last_token_idx]
        else:
            num_samples = num_samples or packed_seq_params.num_samples
            if self.args.context_parallel_size > 1:
                last_token_idx = packed_seq_params.cu_seqlens_q[:num_samples] + packed_seq_params.seq_lens - 1
            else:
                last_token_idx = packed_seq_params.cu_seqlens_q[1:num_samples + 1] - 1
            last_tokens = output_tensor[0, last_token_idx]
        return last_tokens