# Copyright (c) ModelScope Contributors. All rights reserved. # Parts of the functions in this file are code borrowed from NVIDIA/Megatron-LM import copy import dataclasses import megatron.core import numpy as np import os import random import torch from argparse import Namespace from contextlib import contextmanager from datetime import timedelta from mcore_bridge import set_random_seed, split_cp_inputs, unwrap_model from megatron.core import dist_checkpointing, mpu, parallel_state, tensor_parallel from megatron.core.dist_checkpointing.mapping import ShardedObject from megatron.core.dist_checkpointing.serialization import (get_default_load_sharded_strategy, get_default_save_sharded_strategy) from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue, AsyncRequest from megatron.core.dist_checkpointing.strategies.fully_parallel import (FullyParallelLoadStrategyWrapper, FullyParallelSaveStrategyWrapper) from megatron.core.dist_checkpointing.strategies.torch import TorchDistLoadShardedStrategy, TorchDistSaveShardedStrategy from megatron.core.distributed import DistributedDataParallel as DDP from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.fusions.fused_bias_dropout import bias_dropout_add_fused_train from megatron.core.fusions.fused_bias_gelu import bias_gelu from megatron.core.fusions.fused_bias_swiglu import bias_swiglu from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler from megatron.core.transformer.module import Float16Module from megatron.core.utils import get_torch_version, is_te_min_version, is_torch_min_version from packaging import version from transformers.utils import is_torch_npu_available from typing import Any, Dict, Optional from swift.utils import check_json_format, get_logger, init_process_group, is_master, set_device from .patcher import patch_merge_fn logger = get_logger() mcore_017 = version.parse(megatron.core.__version__) >= version.parse('0.17.0rc0') @contextmanager def _patch_megatron_timeout(distributed_timeout_minutes): origin_create_group = parallel_state.create_group def create_group(ranks=None, timeout=None, *_args, **kwargs): if timeout is None: timeout = timedelta(minutes=distributed_timeout_minutes) return origin_create_group(ranks, timeout, *_args, **kwargs) parallel_state.create_group = create_group try: yield finally: parallel_state.create_group = origin_create_group def _initialize_mpu(args): """Initialize torch.distributed and core model parallel.""" if not torch.distributed.is_initialized(): set_device() init_process_group(args.ddp_backend, args.ddp_timeout) args.rank = torch.distributed.get_rank() args.world_size = torch.distributed.get_world_size() if mpu.model_parallel_is_initialized(): logger.info('model parallel is already initialized') else: distributed_timeout_minutes = args.ddp_timeout // 60 with _patch_megatron_timeout(distributed_timeout_minutes): mpu.initialize_model_parallel( args.tensor_model_parallel_size, args.pipeline_model_parallel_size, args.virtual_pipeline_model_parallel_size, context_parallel_size=args.context_parallel_size, expert_model_parallel_size=args.expert_model_parallel_size, expert_tensor_parallel_size=args.expert_tensor_parallel_size, distributed_timeout_minutes=distributed_timeout_minutes, ) if is_master(): logger.info(f'TP: {args.tensor_model_parallel_size}, PP: {args.pipeline_model_parallel_size}, ' f'VPP: {args.virtual_pipeline_model_parallel_size}, CP: {args.context_parallel_size}, ' f'EP: {args.expert_model_parallel_size}, ETP: {args.expert_tensor_parallel_size}') def initialize_megatron(args): # Pytorch distributed. _initialize_mpu(args) # Random seeds for reproducibility. logger.info(f'Setting random seeds to {args.seed}.') set_random_seed(args.seed, args.data_parallel_random_init, args.te_rng_tracker) # Setup MoE aux loss scale value. if args.model_info.is_moe_model: from megatron.core.transformer.moe.router import MoEAuxLossAutoScaler MoEAuxLossAutoScaler.set_loss_scale(torch.ones(1, device=torch.cuda.current_device())) def _get_rng_state(): """Collect rng state across data parallel ranks.""" rng_state = { 'random_rng_state': random.getstate(), 'np_rng_state': np.random.get_state(), 'torch_rng_state': torch.get_rng_state(), 'cuda_rng_state': torch.cuda.get_rng_state(), 'rng_tracker_states': tensor_parallel.get_cuda_rng_tracker().get_states() } # data_parallel_random_init False rng_state_list = [rng_state] pp_rank = mpu.get_pipeline_model_parallel_rank() pp_size = mpu.get_pipeline_model_parallel_world_size() tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() rng_state_list = ShardedObject( 'rng_state', rng_state_list, (pp_size, tp_size), (pp_rank, tp_rank), replica_id=mpu.get_data_parallel_rank(with_context_parallel=True)) return rng_state_list def _generate_state_dict(args, models, optimizer=None, opt_param_scheduler=None, rng_state=None, iteration=None, model_sd_kwargs=None, optim_sd_kwargs=None): model_sd_kwargs = model_sd_kwargs or {} state_dict = { 'args': Namespace(**check_json_format(vars(args))), 'checkpoint_version': 3.0, } if iteration is not None: state_dict['iteration'] = iteration for i, m in enumerate(models): key = 'model' if len(models) > 1: key = f'model{i}' model_sd = models[i].sharded_state_dict(**model_sd_kwargs) state_dict[key] = model_sd if not args.no_save_optim: if optimizer is not None: state_dict['optimizer'] = _optimizer_sharded_state_dict(optimizer, state_dict, optim_sd_kwargs or {}) if opt_param_scheduler is not None: state_dict['opt_param_scheduler'] = opt_param_scheduler.state_dict() if not args.no_save_rng and rng_state is not None: state_dict['rng_state'] = rng_state return state_dict def _optimizer_sharded_state_dict(optimizer, state_dict, optim_sd_kwargs): if is_torch_npu_available(): from swift.model.npu_patch.megatron_checkpoint import optimizer_sharded_state_dict return optimizer_sharded_state_dict(optimizer, state_dict, **optim_sd_kwargs) return optimizer.sharded_state_dict(state_dict, **optim_sd_kwargs) def _load_optimizer_state_dict(optimizer, state_dict): if is_torch_npu_available(): from swift.model.npu_patch.megatron_checkpoint import load_optimizer_state_dict load_optimizer_state_dict(optimizer, state_dict) return optimizer.load_state_dict(state_dict) def _filter_adapter_state_dict(state_dict, peft_format: bool, adapter_name: str = 'default'): """ When peft_format is True, keep only the PEFT format state_dict; when False, remove the PEFT format state_dict. This function ensures it is called when tuner_type != 'full'. """ if 'model' in state_dict: n_models = 1 else: n_models = 0 while f'model{n_models}' in state_dict: n_models += 1 for i in range(n_models): if i == 0 and n_models == 1: model_key = 'model' else: model_key = f'model{i}' new_state_dict = {} state_dict_model = state_dict[model_key] for k, v in state_dict_model.items(): if peft_format: if '.lora_A.' in k or '.lora_B.' in k or '.modules_to_save.' in k: new_state_dict[k] = v else: if '.lora_A.' in k or '.lora_B.' in k or 'original_module.' in k: continue k = k.replace('base_layer.', '') k = k.replace(f'modules_to_save.{adapter_name}.', '') v.key = v.key.replace('base_layer.', '') v.key = v.key.replace(f'modules_to_save.{adapter_name}.', '') new_state_dict[k] = v state_dict[model_key] = new_state_dict def _preprocess_common_before_consistancy_check(common_state_dict): # Convert args key of type namespace to dictionary preprocessed_common_state_dict = copy.deepcopy(common_state_dict) preprocessed_common_state_dict['args'] = vars(preprocessed_common_state_dict['args']) # Remove rank and local rank from state dict if it exists, since they are expected to be different preprocessed_common_state_dict['args'].pop('local_rank', None) preprocessed_common_state_dict['args'].pop('rank', None) return preprocessed_common_state_dict def get_sharded_sd_metadata(args): sharded_sd_metadata = {'singleton_local_shards': False, 'chained_optim_avoid_prefix': True} force_pre_mcore_014 = not is_torch_min_version('2.6a0') if force_pre_mcore_014 and not args.dist_ckpt_save_pre_mcore_014: args.dist_ckpt_save_pre_mcore_014 = True logger.warning(f'PyTorch version {get_torch_version()} below 2.6 detected.' f' Forcing dist_ckpt_save_pre_mcore_014 behavior.') if args.dist_ckpt_save_pre_mcore_014: sharded_sd_metadata['distrib_optim_sharding_type'] = 'fully_sharded_model_space' else: if args.dist_ckpt_optim_fully_reshardable: sharded_sd_metadata['distrib_optim_sharding_type'] = 'fully_reshardable' sharded_sd_metadata[ 'distrib_optim_fully_reshardable_mem_efficient'] = args.distrib_optim_fully_reshardable_mem_efficient else: sharded_sd_metadata['distrib_optim_sharding_type'] = 'dp_reshardable' return sharded_sd_metadata def save_mcore_checkpoint( args, models, optimizer=None, opt_param_scheduler=None, iteration=1, output_dir: Optional[str] = None, peft_format: bool = False, ): if output_dir is None: output_dir = args.output_dir models = unwrap_model(models) rng_state = _get_rng_state() if models else None checkpoint_dir = os.path.join(output_dir, f'iter_{iteration:07d}') sharded_sd_metadata = get_sharded_sd_metadata(args) os.makedirs(checkpoint_dir, exist_ok=True) state_dict = _generate_state_dict( args, models, optimizer, opt_param_scheduler, rng_state, iteration=iteration, model_sd_kwargs={'metadata': sharded_sd_metadata}, optim_sd_kwargs={'metadata': sharded_sd_metadata}, ) _filter_adapter_state_dict(state_dict, peft_format) if mcore_017: save_strategy = TorchDistSaveShardedStrategy() else: save_strategy = get_default_save_sharded_strategy() save_strategy = FullyParallelSaveStrategyWrapper( save_strategy, mpu.get_data_parallel_group(with_context_parallel=True), ) kwargs = {'content_metadata': sharded_sd_metadata} async_save = args.async_save if not models: # save GPU memory assert 'optimizer' not in state_dict async_save = False common_path = os.path.join(checkpoint_dir, 'common.pt') if is_master(): state_dict.update(kwargs) torch.save(state_dict, common_path) async_save_request = None else: async_save_request = dist_checkpointing.save( state_dict, checkpoint_dir, save_strategy, async_sharded_save=async_save, validate_access_integrity=True, preprocess_common_before_consistancy_check=_preprocess_common_before_consistancy_check, **kwargs) if not async_save: assert async_save_request is None # Wait so everyone is done (necessary) if torch.distributed.is_initialized(): torch.distributed.barrier() if is_master(): tracker_path = os.path.join(output_dir, 'latest_checkpointed_iteration.txt') try: from megatron.core.msc_utils import open_file except ImportError: open_file = open with open_file(tracker_path, 'w') as f: f.write(str(iteration)) def iter_finalize_fn(): logger.info(f'Successfully saved Megatron model weights in `{output_dir}`.') if async_save: assert async_save_request is not None async_save_request.add_finalize_fn(iter_finalize_fn) else: iter_finalize_fn() if async_save: schedule_async_save(async_save_request) # Singleton manager of async calls # The default is `TemporalAsyncCaller` _async_calls_queue = AsyncCallsQueue() def init_persistent_async_worker(): global _async_calls_queue # Recreate the async_calls_queue for persistent worker # This duplicate step is for backward compatiblity _async_calls_queue = AsyncCallsQueue(persistent=True) def schedule_async_save(async_request: AsyncRequest): """Schedule the async save request. Args: async_request (AsyncRequest): the async save request. """ _async_calls_queue.schedule_async_request(async_request) def maybe_finalize_async_save(args, blocking: bool = False, terminate=False): """Finalizes active async save calls. Args: blocking (bool, optional): if True, will wait until all active requests are done. Otherwise, finalizes only the async request that already finished. Defaults to False. terminate (bool, optional): if True, the asynchronous queue will be closed as the last action of this function. """ if not args.async_save: return _async_calls_queue.maybe_finalize_async_calls(blocking, no_dist=False) if terminate: _async_calls_queue.close() def is_empty_async_queue() -> bool: """Check if async calls queue is empty. This result is consistent across ranks.""" return _async_calls_queue.get_num_unfinalized_calls() == 0 def _load_iteration(tracker_path: str): if not os.path.exists(tracker_path): return 0 with open(tracker_path, 'r') as f: iteration = int(f.read()) # Get the max iteration retrieved across the ranks. if torch.distributed.is_initialized(): iters_cuda = torch.tensor([iteration], dtype=torch.long, device='cuda') torch.distributed.all_reduce(iters_cuda, op=torch.distributed.ReduceOp.MAX) iteration = iters_cuda[0].item() return iteration def load_mcore_checkpoint(args, ddp_models: list, optimizer=None, opt_param_scheduler=None, load_arg: str = 'mcore_model', adapter_name: str = 'default'): if load_arg in {'mcore_adapter', 'mcore_ref_adapter'}: peft_format = True else: # 'mcore_model', 'mcore_ref_model' peft_format = False load_dir = getattr(args, load_arg) no_load_optim = args.no_load_optim no_load_rng = args.no_load_rng finetune = args.finetune if not peft_format and args.tuner_type != 'full': # When training with LoRA and loading the base model no_load_optim = True no_load_rng = True finetune = True models = unwrap_model(ddp_models) tracker_path = os.path.join(load_dir, 'latest_checkpointed_iteration.txt') iteration = _load_iteration(tracker_path) checkpoint_dir = os.path.join(load_dir, f'iter_{iteration:07d}') state_dict = dist_checkpointing.load_common_state_dict(checkpoint_dir) ckpt_tp_pp = ( state_dict['args'].tensor_model_parallel_size, state_dict['args'].pipeline_model_parallel_size, ) run_tp_pp = ( args.tensor_model_parallel_size, args.pipeline_model_parallel_size, ) mismatch_msg = f'(TP, PP) mismatch after resume ({run_tp_pp} vs {ckpt_tp_pp} from checkpoint)' # Determine if RNG state will be loaded if (ckpt_tp_pp == run_tp_pp and not finetune and not no_load_rng and not getattr(state_dict['args'], 'no_save_rng', False)): gen_sd_rng_state = _get_rng_state() # we can load the rng state else: gen_sd_rng_state = None if ckpt_tp_pp != run_tp_pp: logger.info(f'{mismatch_msg}: RNG state will be ignored') sharded_sd_metadata = state_dict.get('content_metadata') if (not finetune and not no_load_optim and not getattr(state_dict['args'], 'no_save_optim', False)): gen_sd_optim = optimizer gen_sd_opt_param_scheduler = opt_param_scheduler if (args.use_distributed_optimizer and ckpt_tp_pp != run_tp_pp and (sharded_sd_metadata or {}).get('distrib_optim_sharding_type') not in { 'fully_reshardable', 'fully_sharded_model_space', 'fsdp_dtensor', }): raise RuntimeError(f'{mismatch_msg}: not supported for DistributedOptimizer') else: gen_sd_optim, gen_sd_opt_param_scheduler = None, None optim_sd_kwargs = dict(metadata=sharded_sd_metadata, is_loading=True) model_sd_kwargs = dict(metadata=sharded_sd_metadata) # TODO: check no_save_optim sharded_state_dict = _generate_state_dict( args, models, gen_sd_optim, gen_sd_opt_param_scheduler, gen_sd_rng_state, iteration=iteration, model_sd_kwargs=model_sd_kwargs, optim_sd_kwargs=optim_sd_kwargs) _filter_adapter_state_dict(sharded_state_dict, peft_format, adapter_name=adapter_name) model_keys = [k for k in sharded_state_dict.keys() if k.startswith('model')] # compat vpp for k in model_keys: patch_merge_fn(sharded_state_dict[k]) if mcore_017: load_strategy = TorchDistLoadShardedStrategy() else: load_strategy = get_default_load_sharded_strategy(checkpoint_dir) load_strategy = FullyParallelLoadStrategyWrapper(load_strategy, mpu.get_data_parallel_group(with_context_parallel=True)) state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_dir, load_strategy) if finetune: iteration = 0 if 'args' in state_dict and not finetune: args.consumed_train_samples = getattr(state_dict['args'], 'consumed_train_samples', 0) if len(ddp_models) == 1: ddp_models[0].load_state_dict(state_dict['model'], strict=False) else: for i, m in enumerate(ddp_models): if f'model{i}' not in state_dict: continue m.load_state_dict(state_dict[f'model{i}']) if not finetune and not no_load_optim: if optimizer is not None: _load_optimizer_state_dict(optimizer, state_dict['optimizer']) if opt_param_scheduler is not None: opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler']) elif (args.fp16 or args.bf16) and optimizer is not None: optimizer.reload_model_params() if not finetune and not no_load_rng: if 'rng_state' in state_dict: rng_state = state_dict['rng_state'] if args.data_parallel_random_init: rng_state = rng_state[mpu.get_data_parallel_rank()] else: rng_state = rng_state[0] random.setstate(rng_state['random_rng_state']) np.random.set_state(rng_state['np_rng_state']) torch.set_rng_state(rng_state['torch_rng_state']) torch.cuda.set_rng_state(rng_state['cuda_rng_state']) tensor_parallel.get_cuda_rng_tracker().set_states(rng_state['rng_tracker_states']) if torch.distributed.is_initialized(): torch.distributed.barrier() logger.info(f'Successfully loaded Megatron model weights from: {load_dir}') return iteration def wrap_model(args, models, wrap_with_ddp: bool = True): # Set tensor model parallel attributes if not set. # Only parameters that are already tensor model parallel have these # attributes set for them. We should make sure the default attributes # are set for all params so the optimizer can use them. for m in models: for param in m.parameters(): tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param) m.cuda(torch.cuda.current_device()) # Fp16 config = models[0].config if args.fp16 or args.bf16: models = [Float16Module(config, model_module) for model_module in models] # DDP if not wrap_with_ddp: return models kwargs = {} for f in dataclasses.fields(DistributedDataParallelConfig): if hasattr(args, f.name): kwargs[f.name] = getattr(args, f.name) # compat: SWIFT keeps the user-facing Megatron-LM arg name, while MCore # DistributedDataParallelConfig expects grad_reduce_in_fp32. if hasattr(args, 'accumulate_allreduce_grads_in_fp32'): kwargs['grad_reduce_in_fp32'] = args.accumulate_allreduce_grads_in_fp32 kwargs['check_for_nan_in_grad'] = True ddp_config = DistributedDataParallelConfig(**kwargs) # In the Megatron FSDP and DDP use path, we need to initialize the bucket size. # If bucket_size is not provided as an input, use sane default. # If using very large dp_sizes, make buckets larger to ensure that chunks used in NCCL # ring-reduce implementations are large enough to remain bandwidth-bound rather than # latency-bound. if ddp_config.bucket_size is None: ddp_config.bucket_size = max(40000000, 1000000 * mpu.get_data_parallel_world_size(with_context_parallel=True)) # Set bucket_size to infinity if overlap_grad_reduce is False. if not ddp_config.overlap_grad_reduce: ddp_config.bucket_size = None # Setup stream for DDP initialization with proper synchronization. ddp_stream = torch.cuda.Stream() ddp_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(ddp_stream): models = [ DDP( config=config, ddp_config=ddp_config, module=model_chunk, # Turn off bucketing for model_chunk 2 onwards, since communication for these # model chunks is overlapped with compute anyway. disable_bucketing=(model_chunk_idx > 0) or args.overlap_param_gather_with_optimizer_step, ) for (model_chunk_idx, model_chunk) in enumerate(models) ] # Ensure DDP initialization completes before proceeding on the default stream. torch.cuda.current_stream().wait_stream(ddp_stream) # Broadcast params from data parallel src rank to other data parallel ranks. if args.data_parallel_random_init: for m in models: m.broadcast_params() return models def get_optimizer_param_scheduler(args, optimizer): # Iteration-based training. if args.lr_decay_iters is None: args.lr_decay_iters = args.train_iters lr_decay_steps = args.lr_decay_iters * args.global_batch_size wd_incr_steps = args.train_iters * args.global_batch_size wsd_decay_steps = None if args.lr_wsd_decay_iters is not None: wsd_decay_steps = args.lr_wsd_decay_iters * args.global_batch_size if args.lr_warmup_fraction is not None: lr_warmup_steps = args.lr_warmup_fraction * lr_decay_steps else: lr_warmup_steps = args.lr_warmup_iters * args.global_batch_size opt_param_scheduler = OptimizerParamScheduler( optimizer, init_lr=args.lr_warmup_init, max_lr=args.lr, min_lr=args.min_lr, lr_warmup_steps=lr_warmup_steps, lr_decay_steps=lr_decay_steps, lr_decay_style=args.lr_decay_style, start_wd=args.start_weight_decay, end_wd=args.end_weight_decay, wd_incr_steps=wd_incr_steps, wd_incr_style=args.weight_decay_incr_style, wsd_decay_steps=wsd_decay_steps, lr_wsd_decay_style=args.lr_wsd_decay_style, ) return opt_param_scheduler def should_disable_forward_pre_hook(args): """Block forward pre-hook for certain configurations.""" return args.use_distributed_optimizer and args.overlap_param_gather def enable_forward_pre_hook(model_chunks): for model_chunk in model_chunks: assert isinstance(model_chunk, DDP) model_chunk.enable_forward_pre_hook() def disable_forward_pre_hook(model_chunks, param_sync=True): for model_chunk in model_chunks: assert isinstance(model_chunk, DDP) model_chunk.disable_forward_pre_hook(param_sync=param_sync) def initialize_tp_communicators(args, config): """initializing the communicators with user buffers for high-performance tensor-model-parallel communication overlap""" from transformer_engine.pytorch import module as te_module input_shape = [ (args.seq_length * args.micro_batch_size) // args.context_parallel_size, config.hidden_size, ] if is_te_min_version('2.7.0'): UserBufferQuantizationMode = te_module.base.UserBufferQuantizationMode quantization_modes = [UserBufferQuantizationMode.FP8 if args.fp8 else UserBufferQuantizationMode.NONE] if args.fp8 is not None and args.first_last_layers_bf16 and (args.num_layers_at_start_in_bf16 > 0 or args.num_layers_at_end_in_bf16 > 0): quantization_modes.append(UserBufferQuantizationMode.NONE) # The process group with the target bootstrap backend is created in Transformer Engine. te_module.base.initialize_ub( shape=input_shape, tp_size=args.tensor_model_parallel_size, quantization_modes=quantization_modes, ) elif is_te_min_version('1.9.0'): # The process group with the target bootstrap backend is created in Transformer Engine. te_module.base.initialize_ub( shape=input_shape, tp_size=args.tensor_model_parallel_size, use_fp8=(args.fp8 is not None), ) def warmup_jit_function(config, args): if args.bf16: dtype = torch.bfloat16 elif args.fp16: dtype = torch.float16 else: dtype = torch.float32 bias = torch.rand(config.ffn_hidden_size // config.tensor_model_parallel_size, dtype=dtype, device='cuda') input_tensor = torch.rand( ( args.seq_length // config.context_parallel_size, args.micro_batch_size, config.ffn_hidden_size // config.tensor_model_parallel_size, ), dtype=dtype, device='cuda', ) # Warmup JIT fusions with the input_tensor grad_enable state of both forward # prop and recomputation for bias_grad, input_grad in zip([True, True], [False, True]): bias.requires_grad, input_tensor.requires_grad = bias_grad, input_grad for _ in range(5): if config.swiglu: output = bias_swiglu(input_tensor, bias) else: output = bias_gelu(bias, input_tensor) del bias, input_tensor, output # Warmup fused bias+dropout+add if config.sequence_parallel: seq_length = args.seq_length // mpu.get_tensor_model_parallel_world_size() else: seq_length = args.seq_length input_tensor = torch.rand( (seq_length // config.context_parallel_size, args.micro_batch_size, config.hidden_size), dtype=dtype, device='cuda', ) residual = torch.rand( (seq_length // config.context_parallel_size, args.micro_batch_size, config.hidden_size), dtype=dtype, device='cuda', ) bias = torch.rand((config.hidden_size), dtype=dtype, device='cuda').expand_as(residual) dropout_rate = 0.1 # Warmup JIT fusions with the input_tensor grad_enable state of both forward # prop and recomputation for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]): input_tensor.requires_grad = input_grad bias.requires_grad = bias_grad residual.requires_grad = residual_grad for _ in range(5): output = bias_dropout_add_fused_train([input_tensor, bias], residual, dropout_rate) del bias, input_tensor, residual, output torch.cuda.empty_cache() def get_batch_on_this_cp_rank(args, batch: Dict[str, Any]): """Slice batch input along sequence dimension into multiple chunks, which are parallelized across GPUs in a context parallel group. """ # With causal masking, each token only attends to its prior tokens. Simply split # sequence into CP chunks can result in severe load imbalance. That's to say, chunks # at the end of sequence have bigger workload than others. To address this issue, # we split sequence into 2*CP ranks. Assuming CP=2, we then get 4 chunks, chunk_0 # and chunk_3 are assigned to GPU0, chunk_1 and chunk_2 are assigned to GPU1, so # that we can get balanced workload among GPUs in a context parallel group. cp_size = mpu.get_context_parallel_world_size() if cp_size > 1: keys = ['labels', 'position_ids', 'loss_scale'] if not args.is_multimodal: # Multimodal models will handle CP in input_embeds. keys.append('input_ids') packed_seq_params = batch.get('packed_seq_params') cp_partition_mode = getattr(args, 'cp_partition_mode', 'zigzag') kwargs = {} if cp_partition_mode == 'contiguous': kwargs['cp_partition_mode'] = 'contiguous' for key, val in batch.items(): if key not in keys: continue if args.task_type in ('seq_cls', 'embedding', 'generative_reranker') and key == 'labels': continue if val is not None: batch[key] = split_cp_inputs(val, getattr(packed_seq_params, 'cu_seqlens_q', None), -1, **kwargs) attention_mask = batch.get('attention_mask') if is_torch_npu_available() and attention_mask is not None and attention_mask.ndim >= 4: batch['attention_mask'] = split_cp_inputs(attention_mask, getattr(packed_seq_params, 'cu_seqlens_q', None), -2, **kwargs) return batch