# 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 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 = [] load_balancing_type = config.moe_router_load_balancing_type if isinstance(load_balancing_type, str): load_balancing_type = [load_balancing_type] if 'aux_loss' in load_balancing_type: track_names.append('load_balancing_loss') if 'seq_aux_loss' in load_balancing_type: track_names.append('seq_load_balancing_loss') if 'global_aux_loss' in load_balancing_type: 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') # Instead of using lr_mult to control the learning rate, we directly use max_lr/min_lr. param_group['lr_mult'] = 1. 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 if args.nccl_comm_warmup: # Eagerly create NCCL communicators while GPU memory is still free. Lazily-initialized # comms (e.g. the dp/cp loss all-reduce and grad-sync coalescing) otherwise first fire # at the iteration-1 memory peak, where NCCL's internal cudaMalloc can fail with # "Failed to CUDA calloc async N bytes". A 1-element dummy all-reduce per group is # numerically inert and forces the communicator to be created up front. dummy = torch.zeros(1, device=get_current_device()) warmed = 0 for getter, kwargs in ( (mpu.get_data_parallel_group, { 'with_context_parallel': True }), (mpu.get_data_parallel_group, {}), (mpu.get_context_parallel_group, {}), (mpu.get_tensor_model_parallel_group, {}), (mpu.get_pipeline_model_parallel_group, {}), (mpu.get_model_parallel_group, {}), (mpu.get_embedding_group, {}), (mpu.get_position_embedding_group, {}), ): try: group = getter(**kwargs) except (AssertionError, ValueError, TypeError): continue for g in (group if isinstance(group, list) else [group]): if g is not None: torch.distributed.all_reduce(dummy, group=g) warmed += 1 torch.cuda.synchronize() logger.info(f'NCCL communicator warm-up done ({warmed} groups).') 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 _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): 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 = packed_seq_params.seq_lens.shape[0] last_token_idx = packed_seq_params.cu_seqlens_q[:num_samples] + packed_seq_params.seq_lens - 1 last_tokens = output_tensor[0, last_token_idx] return last_tokens