utils.py

# 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, 'name', args.attention_backend) in {'local', 'unfused'}:
        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