wehub-resource-sync
1060 lines
48 KiB
Python
1060 lines
48 KiB
Python
# 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
|