wehub-resource-sync
526 lines
22 KiB
Python
526 lines
22 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import gc
|
|
import torch
|
|
from accelerate.utils import gather as hf_gather
|
|
from accelerate.utils import gather_object as hf_gather_object
|
|
from contextlib import nullcontext
|
|
from dataclasses import dataclass
|
|
from megatron.core import mpu
|
|
from megatron.core.distributed import DistributedDataParallel as DDP
|
|
from megatron.core.optimizer import ChainedOptimizer
|
|
from typing import Any, Optional
|
|
|
|
from swift.dataloader import DataLoaderDispatcher
|
|
from swift.megatron.utils import get_batch_on_this_cp_rank, get_packed_seq_params
|
|
from swift.utils import empty_cache, get_current_device, get_logger, to_device
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def get_batch_on_this_pp_rank(args, data, vp_stage=None):
|
|
if args.task_type == 'causal_lm':
|
|
data['labels'] = torch.roll(data['labels'], -1, dims=-1)
|
|
if 'loss_scale' in data:
|
|
data['loss_scale'] = torch.roll(data['loss_scale'], -1, dims=-1)
|
|
batch = to_device(data, get_current_device(), non_blocking=True)
|
|
if args.pipeline_model_parallel_size == 1:
|
|
return batch
|
|
is_pp_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage)
|
|
if not is_pp_last_stage:
|
|
batch['labels'] = None
|
|
if 'loss_scale' in batch:
|
|
batch['loss_scale'] = None
|
|
|
|
return batch
|
|
|
|
|
|
def gather(tensor, group: Optional[torch.distributed.ProcessGroup] = None):
|
|
if group is None:
|
|
return hf_gather(tensor)
|
|
size = torch.distributed.get_world_size(group=group)
|
|
output = [torch.empty_like(tensor) for _ in range(size)]
|
|
torch.distributed.all_gather(output, tensor, group=group, async_op=False)
|
|
|
|
return torch.cat(output, dim=0)
|
|
|
|
|
|
def gather_object(object: Any, group: Optional[torch.distributed.ProcessGroup] = None):
|
|
if group is None:
|
|
return hf_gather_object(object)
|
|
size = torch.distributed.get_world_size(group=group)
|
|
output_objects = [None for _ in range(size)]
|
|
torch.distributed.all_gather_object(output_objects, object, group=group)
|
|
return [x for y in output_objects for x in y]
|
|
|
|
|
|
# code borrowed from verl
|
|
@torch.no_grad()
|
|
def load_megatron_model_to_gpu(models, load_grad=True, load_frozen_params=True):
|
|
for model_chunk in models:
|
|
if isinstance(model_chunk, DDP):
|
|
model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]
|
|
for buffers in model_chunk_all_buffers:
|
|
for buffer in buffers:
|
|
# sometimes, we don't want to load grad for pure inference
|
|
if load_grad and hasattr(buffer, 'grad_data_size'):
|
|
current_storage_size = buffer.grad_data.storage().size()
|
|
if current_storage_size == 0 or current_storage_size == buffer.grad_data_size:
|
|
buffer.grad_data.storage().resize_(buffer.grad_data_size)
|
|
buffer.grad_data.zero_()
|
|
else:
|
|
# Non-standard layers (e.g. GatedDeltaNet) may have grad
|
|
# buffers with mismatched storage size; skip resize and
|
|
# zero in-place with current storage.
|
|
buffer.grad_data.zero_()
|
|
|
|
if buffer.param_data.storage().size() == 0:
|
|
buffer.param_data.storage().resize_(buffer.param_data_size)
|
|
# copy data from cpu to cuda
|
|
buffer.param_data.copy_(buffer.param_data.cpu_data, non_blocking=True)
|
|
|
|
if load_frozen_params:
|
|
device_id = get_current_device()
|
|
for param in model_chunk.module.parameters():
|
|
if not param.requires_grad and param.device.type == 'cpu':
|
|
param.data = param.data.to(device_id, non_blocking=True)
|
|
else:
|
|
# we need this for ref module
|
|
device_id = get_current_device()
|
|
for _, param in model_chunk.named_parameters():
|
|
param.data = param.data.to(device_id, non_blocking=True)
|
|
if param.grad is not None:
|
|
param.grad = param.grad.to(device_id, non_blocking=True)
|
|
gc.collect()
|
|
empty_cache()
|
|
|
|
|
|
@torch.no_grad()
|
|
def offload_megatron_model_to_cpu(models):
|
|
"""
|
|
In megatron, the model and optimizer storage are:
|
|
- bf16 parameter data chunked in model parallel group
|
|
- fp32 grad chunked in model parallel group
|
|
- fp32 main_parameter chunked in model and dp group
|
|
- fp32 optimizer state chunked in model and dp group
|
|
|
|
When using LoRA, frozen base model parameters are NOT managed by DDP buffers.
|
|
They must be offloaded separately via direct param iteration.
|
|
"""
|
|
for model_chunk in models:
|
|
if isinstance(model_chunk, DDP):
|
|
model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]
|
|
for buffers in model_chunk_all_buffers:
|
|
for buffer in buffers:
|
|
# offload parameters
|
|
if buffer.param_data.storage().size() > 0:
|
|
existing = getattr(buffer.param_data, 'cpu_data', None)
|
|
if existing is None:
|
|
buffer.param_data.cpu_data = torch.empty(
|
|
buffer.param_data.size(),
|
|
dtype=buffer.param_data.dtype,
|
|
device='cpu',
|
|
pin_memory=True,
|
|
)
|
|
buffer.param_data_size = buffer.param_data.storage().size()
|
|
else:
|
|
assert existing.shape == buffer.param_data.shape, (
|
|
f'cpu_data shape {tuple(existing.shape)} != '
|
|
f'param_data shape {tuple(buffer.param_data.shape)}; '
|
|
'reallocating would reintroduce the 2x peak.')
|
|
assert existing.dtype == buffer.param_data.dtype, (
|
|
f'cpu_data dtype {existing.dtype} != '
|
|
f'param_data dtype {buffer.param_data.dtype}; '
|
|
'reallocating would reintroduce the 2x peak.')
|
|
# Synchronous D2H copy into the preexisting pinned
|
|
# buffer; must complete before resize_(0) frees the
|
|
# GPU storage.
|
|
buffer.param_data.cpu_data.copy_(buffer.param_data.data, non_blocking=False)
|
|
buffer.param_data.storage().resize_(0)
|
|
|
|
assert buffer.param_data_size == buffer.param_data.cpu_data.storage().size()
|
|
|
|
if buffer.grad_data.storage().size() > 0:
|
|
# if the grad_data size is already zero, we assume that it is already offloaded
|
|
buffer.grad_data_size = buffer.grad_data.storage().size()
|
|
buffer.grad_data.storage().resize_(0)
|
|
|
|
for param in model_chunk.module.parameters():
|
|
if not param.requires_grad and param.device.type != 'cpu':
|
|
param.data = param.data.to('cpu', non_blocking=True)
|
|
else:
|
|
# we need this for ref module
|
|
for _, param in model_chunk.named_parameters():
|
|
param.data = param.data.to('cpu', non_blocking=True)
|
|
if param.grad is not None:
|
|
param.grad = param.grad.to('cpu', non_blocking=True)
|
|
gc.collect()
|
|
empty_cache()
|
|
|
|
|
|
@torch.no_grad()
|
|
def load_megatron_copy_params(optimizers):
|
|
"""
|
|
Load optimizer parameters back to GPU. Handles ChainedOptimizer.
|
|
|
|
Args:
|
|
optimizers: Optimizer or ChainedOptimizer instance.
|
|
"""
|
|
|
|
def _iter_opts(opt):
|
|
if isinstance(opt, ChainedOptimizer):
|
|
return opt.chained_optimizers
|
|
return [opt]
|
|
|
|
def load_tensor_to_gpu(tensor):
|
|
if tensor is None:
|
|
return
|
|
device_id = get_current_device()
|
|
tensor.data = tensor.data.to(device_id, non_blocking=True)
|
|
|
|
def load_group_to_gpu(group):
|
|
if group is None:
|
|
return
|
|
|
|
if isinstance(group, list):
|
|
for param_group in group:
|
|
if isinstance(param_group, list):
|
|
for param in param_group:
|
|
load_tensor_to_gpu(param)
|
|
else:
|
|
load_tensor_to_gpu(param_group)
|
|
else:
|
|
load_tensor_to_gpu(group)
|
|
|
|
# Load all parameter groups to GPU for each underlying optimizer
|
|
|
|
for _opt in _iter_opts(optimizers):
|
|
if hasattr(_opt, 'shard_fp32_from_float16_groups'):
|
|
load_group_to_gpu(_opt.shard_fp32_from_float16_groups)
|
|
|
|
|
|
@torch.no_grad()
|
|
def offload_megatron_copy_params(optimizers):
|
|
"""
|
|
Offload optimizer parameters to CPU. Supports both Megatron optimizers
|
|
and `ChainedOptimizer`, which wraps a list of underlying optimizers.
|
|
|
|
Args:
|
|
optimizers: The optimizer or ChainedOptimizer instance.
|
|
"""
|
|
|
|
def _iter_opts(opt):
|
|
if isinstance(opt, ChainedOptimizer):
|
|
return opt.chained_optimizers
|
|
return [opt]
|
|
|
|
def offload_tensor_to_cpu(tensor):
|
|
if tensor is None:
|
|
return
|
|
tensor.data = tensor.data.to('cpu', non_blocking=True)
|
|
|
|
def offload_group_to_cpu(group):
|
|
if group is None:
|
|
return
|
|
|
|
if isinstance(group, list):
|
|
for param_group in group:
|
|
if isinstance(param_group, list):
|
|
for param in param_group:
|
|
offload_tensor_to_cpu(param)
|
|
else:
|
|
offload_tensor_to_cpu(param_group)
|
|
else:
|
|
offload_tensor_to_cpu(group)
|
|
|
|
# Offload all parameter groups to CPU for each underlying optimizer
|
|
|
|
for _opt in _iter_opts(optimizers):
|
|
if hasattr(_opt, 'shard_fp32_from_float16_groups'):
|
|
offload_group_to_cpu(_opt.shard_fp32_from_float16_groups)
|
|
|
|
|
|
@torch.no_grad()
|
|
def load_megatron_optimizer(optimizers):
|
|
|
|
def _iter_opts(opt):
|
|
if isinstance(opt, ChainedOptimizer):
|
|
return opt.chained_optimizers
|
|
return [opt]
|
|
|
|
for _opt in _iter_opts(optimizers):
|
|
load_megatron_copy_params(_opt)
|
|
if _opt.optimizer is not None:
|
|
# if we are using HybridDeviceOptimizer, we need to only move gpu optimizer state to gpu
|
|
if hasattr(_opt.optimizer, '_move_new_state_to_right_device'):
|
|
_opt.optimizer._move_new_state_to_right_device()
|
|
else:
|
|
opt_state_dict_values = _opt.optimizer.state.values()
|
|
for v in opt_state_dict_values:
|
|
if 'exp_avg' in v:
|
|
v['exp_avg'] = v['exp_avg'].to(get_current_device(), non_blocking=True)
|
|
if 'exp_avg_sq' in v:
|
|
v['exp_avg_sq'] = v['exp_avg_sq'].to(get_current_device(), non_blocking=True)
|
|
gc.collect()
|
|
empty_cache()
|
|
|
|
|
|
@torch.no_grad()
|
|
def offload_megatron_optimizer(optimizers):
|
|
|
|
def _iter_opts(opt):
|
|
if isinstance(opt, ChainedOptimizer):
|
|
return opt.chained_optimizers
|
|
return [opt]
|
|
|
|
for _opt in _iter_opts(optimizers):
|
|
offload_megatron_copy_params(_opt)
|
|
# worker may hold zero parameter when enabling custom pipeline layout
|
|
if _opt.optimizer is not None:
|
|
# HybridDeviceOptimizer: offload all sub-optimizer states to CPU
|
|
hdo = _opt.optimizer
|
|
if all(hasattr(hdo, attr) for attr in ('sub_optimizers', 'inner_param_to_orig_param', 'state')):
|
|
for optimizer in hdo.sub_optimizers:
|
|
for param, state in optimizer.state.items():
|
|
for k, v in state.items():
|
|
if not isinstance(v, torch.Tensor):
|
|
continue
|
|
orig_param = hdo.inner_param_to_orig_param.get(param, param)
|
|
hdo.state[orig_param][k] = state[k] = v.to('cpu')
|
|
else:
|
|
opt_state_dict_values = _opt.optimizer.state.values()
|
|
for v in opt_state_dict_values:
|
|
if 'exp_avg' in v:
|
|
v['exp_avg'] = v['exp_avg'].to('cpu', non_blocking=True)
|
|
if 'exp_avg_sq' in v:
|
|
v['exp_avg_sq'] = v['exp_avg_sq'].to('cpu', non_blocking=True)
|
|
gc.collect()
|
|
empty_cache()
|
|
|
|
|
|
def log_gpu_memory(prefix: str = '', info_once: bool = False):
|
|
log_msg = (f'{prefix} GPU memory: {torch.cuda.memory_allocated() / 1024**3:.2f}GB allocated, '
|
|
f'{torch.cuda.memory_reserved() / 1024**3:.2f}GB reserved')
|
|
if info_once:
|
|
logger.info_once(log_msg, hash_id=prefix)
|
|
else:
|
|
logger.info(log_msg)
|
|
|
|
|
|
@dataclass
|
|
class TrainerState:
|
|
should_save: bool = False
|
|
should_eval: bool = False
|
|
should_log: bool = False
|
|
|
|
iteration: int = 0
|
|
consumed_train_samples: int = 0
|
|
# compat transformers
|
|
max_steps: Optional[int] = None
|
|
|
|
best_metric: Optional[float] = None
|
|
best_global_step: Optional[int] = None
|
|
last_model_checkpoint: Optional[str] = None
|
|
best_model_checkpoint: Optional[str] = None
|
|
|
|
@property
|
|
def global_step(self) -> int:
|
|
return self.iteration
|
|
|
|
|
|
class MegatronDataLoaderDispatcher(DataLoaderDispatcher):
|
|
|
|
@property
|
|
def group(self):
|
|
return mpu.get_data_parallel_group()
|
|
|
|
|
|
def build_streaming_dataloader(args, dataset, collate_fn):
|
|
base_dataloader = torch.utils.data.DataLoader(
|
|
dataset,
|
|
num_workers=args.dataloader_num_workers,
|
|
pin_memory=args.dataloader_pin_memory,
|
|
collate_fn=collate_fn,
|
|
batch_size=args.micro_batch_size,
|
|
prefetch_factor=args.dataloader_prefetch_factor if args.dataloader_num_workers > 0 else None,
|
|
persistent_workers=args.dataloader_persistent_workers if args.dataloader_num_workers > 0 else False,
|
|
)
|
|
return MegatronDataLoaderDispatcher(base_dataloader)
|
|
|
|
|
|
_NPU_ATTENTION_MASK_2D_MODEL_TYPES = {'qwen3_5', 'qwen3_5_moe'}
|
|
|
|
|
|
def _should_use_npu_generated_attention_mask(args) -> bool:
|
|
from transformers.utils import is_torch_npu_available
|
|
if not is_torch_npu_available():
|
|
return False
|
|
if args.task_type != 'causal_lm' or args.padding_free:
|
|
return False
|
|
if getattr(args, 'attention_backend', None) == 'local':
|
|
return False
|
|
return bool(getattr(args, 'use_flash_attn', False))
|
|
|
|
|
|
def _prepare_npu_generated_attention_mask(batch, *, keep_attention_mask_2d: bool) -> None:
|
|
if keep_attention_mask_2d:
|
|
attention_mask = batch.get('attention_mask')
|
|
if 'attention_mask_2d' not in batch and attention_mask is not None:
|
|
batch['attention_mask_2d'] = (attention_mask == 0).sum(dim=(1, 2)) > 0
|
|
else:
|
|
batch.pop('attention_mask_2d', None)
|
|
batch['attention_mask'] = None
|
|
|
|
|
|
def prepare_batch(args, data, vp_stage=None):
|
|
"""Prepare a micro-batch for Megatron forward: PP slicing, packed_seq_params, CP slicing.
|
|
|
|
Extracted from BaseMegatronTrainer._prepare_batch for reuse in ray workers.
|
|
"""
|
|
batch = get_batch_on_this_pp_rank(args, data, vp_stage=vp_stage)
|
|
seq_lens = batch.pop('seq_lens', None)
|
|
# Consider compatibility and security.
|
|
num_samples = batch.pop('num_samples', None)
|
|
if seq_lens is not None:
|
|
if num_samples is not None:
|
|
assert num_samples == len(seq_lens), (
|
|
f"'num_samples' ({num_samples}) is inconsistent with len(seq_lens) ({len(seq_lens)}).")
|
|
num_samples = len(seq_lens)
|
|
text_position_ids = batch.pop('text_position_ids', None)
|
|
if text_position_ids is None:
|
|
text_position_ids = batch.get('position_ids')
|
|
if _should_use_npu_generated_attention_mask(args):
|
|
_prepare_npu_generated_attention_mask(
|
|
batch, keep_attention_mask_2d=getattr(args, 'model_type', None) in _NPU_ATTENTION_MASK_2D_MODEL_TYPES)
|
|
else:
|
|
batch.pop('attention_mask_2d', None)
|
|
if args.padding_free and text_position_ids is not None:
|
|
batch['packed_seq_params'] = get_packed_seq_params(args, text_position_ids)
|
|
if seq_lens is not None:
|
|
batch['packed_seq_params'].seq_lens = torch.tensor(seq_lens, device=text_position_ids.device)
|
|
if num_samples is not None:
|
|
batch['packed_seq_params'].num_samples = num_samples
|
|
batch.setdefault('attention_mask', None)
|
|
batch = get_batch_on_this_cp_rank(args, batch)
|
|
return batch
|
|
|
|
|
|
def compute_per_token_logps_fn(model, args, data_iterator, temperature=1.0, no_grad=True, enable_routing_replay=False):
|
|
"""Forward pass → logits → temperature-scaled per-token logps.
|
|
|
|
Returns:
|
|
(per_token_logps, routing_topk_idx) — either may be None on non-last PP stages.
|
|
"""
|
|
from swift.megatron.utils import (RouterReplayHelper, forward_step_helper, get_local_topk_idx_for_current_rank,
|
|
get_router_replay_data, set_router_replay_data)
|
|
from .vocab_parallel_utils import compute_logps_and_entropy_from_logits
|
|
|
|
data = prepare_batch(args, next(data_iterator))
|
|
data.pop('loss_scale', None)
|
|
labels = data.get('labels')
|
|
|
|
routing_topk_idx = None
|
|
global_topk_idx = data.pop('routed_experts', None)
|
|
if enable_routing_replay and RouterReplayHelper.is_replay_forward_action(model.config):
|
|
assert global_topk_idx is not None, 'When router_replay_mode = R3, routed_experts must be in data'
|
|
routing_topk_idx = get_local_topk_idx_for_current_rank(global_topk_idx, model.config,
|
|
data.get('packed_seq_params'))
|
|
set_router_replay_data(routing_topk_idx, model.config)
|
|
|
|
data_for_forward = {k: v for k, v in data.items() if k != 'labels'}
|
|
context = torch.no_grad() if no_grad else nullcontext()
|
|
is_training = model.training
|
|
if is_training:
|
|
model.eval()
|
|
try:
|
|
with context:
|
|
output_tensor = forward_step_helper(model, data_for_forward)
|
|
finally:
|
|
if is_training:
|
|
model.train()
|
|
|
|
if enable_routing_replay and RouterReplayHelper.is_r2_record_action(model.config):
|
|
routing_topk_idx = get_router_replay_data(model.config)
|
|
|
|
if labels is None or output_tensor is None:
|
|
return None, routing_topk_idx
|
|
|
|
if temperature != 1.0:
|
|
output_tensor.div_(temperature)
|
|
per_token_logps, _ = compute_logps_and_entropy_from_logits(output_tensor, labels)
|
|
|
|
packed_seq_params = data.get('packed_seq_params')
|
|
if packed_seq_params is not None:
|
|
num_samples = packed_seq_params.seq_lens.shape[0]
|
|
else:
|
|
input_ids = data.get('input_ids')
|
|
num_samples = input_ids.shape[0] if input_ids is not None else labels.shape[0]
|
|
|
|
if args.context_parallel_size > 1:
|
|
per_token_logps = reconstruct_tensor_cp(args.context_parallel_size, per_token_logps, packed_seq_params,
|
|
num_samples)
|
|
return per_token_logps, routing_topk_idx
|
|
|
|
|
|
def reconstruct_tensor_cp(cp_size, tensor, packed_seq_params, num_samples):
|
|
"""In CP mode, all_gather and reconstruct full tensor sequences."""
|
|
cp_rank = mpu.get_context_parallel_rank()
|
|
|
|
# All-gather across CP ranks
|
|
output_list = [torch.empty_like(tensor) for _ in range(cp_size)]
|
|
torch.distributed.all_gather(output_list, tensor.contiguous(), group=mpu.get_context_parallel_group())
|
|
output_list[cp_rank] = tensor
|
|
|
|
if packed_seq_params is not None:
|
|
cu_seqlens_full = packed_seq_params.cu_seqlens_q
|
|
cu_seqlens_cp = cu_seqlens_full // cp_size
|
|
|
|
# Calculate total packed length
|
|
total_packed_len = cu_seqlens_full[num_samples].item()
|
|
output_full = tensor.new_zeros(1, total_packed_len)
|
|
|
|
# Reconstruct each sequence
|
|
for i in range(num_samples):
|
|
start_full = cu_seqlens_full[i].item()
|
|
end_full = cu_seqlens_full[i + 1].item()
|
|
seq_len = end_full - start_full
|
|
|
|
# Length of each chunk after CP split
|
|
chunk_len = seq_len // cp_size
|
|
half_chunk = chunk_len // 2
|
|
|
|
# Concatenate from each CP rank's output (load-balanced split)
|
|
for j in range(cp_size):
|
|
o = output_list[j][0]
|
|
start_cp = cu_seqlens_cp[i].item()
|
|
o0 = o[start_cp:start_cp + half_chunk]
|
|
o1 = o[start_cp + half_chunk:start_cp + chunk_len]
|
|
|
|
# Place back to full sequence
|
|
output_full[0, start_full + j * half_chunk:start_full + (j + 1) * half_chunk] = o0
|
|
output_full[0, end_full - (j + 1) * half_chunk:end_full - j * half_chunk] = o1
|
|
else:
|
|
# non-padding_free mode: [batch_size, seq_len/cp_size] -> [batch_size, seq_len]
|
|
# Each CP rank has chunks split with load-balanced pattern (2*cp_size chunks)
|
|
batch_size = tensor.shape[0]
|
|
seq_len_per_cp = tensor.shape[1]
|
|
full_seq_len = seq_len_per_cp * cp_size
|
|
output_full = tensor.new_zeros(batch_size, full_seq_len)
|
|
|
|
# Each CP rank j holds chunks j and (2*cp_size - j - 1) from the original 2*cp_size split
|
|
# Reconstruct the full sequence by placing chunks back in correct positions
|
|
chunk_len = full_seq_len // (2 * cp_size)
|
|
|
|
for j in range(cp_size):
|
|
o = output_list[j]
|
|
# This rank holds 2 chunks: chunk j and chunk (2*cp_size - j - 1)
|
|
half_len = seq_len_per_cp // 2
|
|
o0 = o[:, :half_len]
|
|
o1 = o[:, half_len:]
|
|
|
|
# Place chunk j at position j * chunk_len
|
|
output_full[:, j * chunk_len:(j + 1) * chunk_len] = o0
|
|
reverse_idx = 2 * cp_size - j - 1
|
|
output_full[:, reverse_idx * chunk_len:(reverse_idx + 1) * chunk_len] = o1
|
|
|
|
return output_full
|