utils.py — Megatron 工具函数

文件概述

提供 Megatron 引擎的随机种子设置工具。

核心函数

set_random_seed

def set_random_seed(seed):
    """设置所有随机种子，包括 Megatron 的张量并行种子"""
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    # Megatron 特有: 设置张量并行的 CUDA 随机种子
    tensor_parallel.model_parallel_cuda_manual_seed(seed)

小结

确保 Megatron 分布式训练中所有并行组的随机种子一致，保证可复现性。