utils.py — 引擎通用工具

文件概述

提供训练引擎的通用工具函数，包括确定性训练、微批次准备和批次后处理。

核心函数

enable_full_determinism

def enable_full_determinism(seed=42):
    """启用完全确定性训练（可复现）

    设置所有随机种子:
    - torch.manual_seed
    - np.random.seed
    - random.seed
    - CUDA 确定性算法
    """

prepare_micro_batches

def prepare_micro_batches(data, dp_group, num_batches_divided_by=None,
                          same_micro_num_in_dp=False, min_num_micro_batch=None):
    """将大批次数据分成微批次

    参数:
    - data: 输入数据
    - dp_group: 数据并行通信组
    - same_micro_num_in_dp: 确保所有 DP rank 有相同数量的微批次
    - min_num_micro_batch: 最少微批次数

    返回:
    - micro_batches: 微批次列表
    - indices: 用于后续还原顺序的索引
    """

postprocess_batch_func

def postprocess_batch_func(output_lst, indices, data):
    """将多个微批次的输出合并还原为完整批次

    功能:
    - 收集所有微批次的 model_output
    - 按原始顺序还原
    - 聚合 loss 和 metrics
    """

与其他模块的关系

• 被所有 Engine 实现（FSDP/Megatron/TorchTitan/VeOmni）使用
• prepare_micro_batches 是控制显存使用的关键

小结

这些工具函数是所有训练引擎共享的基础设施，特别是微批次处理对 GPU 显存管理至关重要。