Megatron-LM training.py 源码精读 · 初始化篇

def pretrain(
    train_valid_test_dataset_provider,  # 回调①：数据怎么来
    model_provider,                     # 回调②：模型怎么建
    model_type,                         # 回调③：模型类型标识
    forward_step_func,                  # 回调④：前向怎么算
    ...
):
    # ---- Step 1: 初始化 ----
    initialize_megatron(...)            # 分布式环境 + 通信组 + 随机种子
    args = get_args()

    # ---- Step 2: 模型 + 优化器 ----
    model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
        model_provider, model_type, ...
    )

    # ---- Step 3: 数据迭代器 ----
    train_data_iterator, valid_data_iterator, test_data_iterator = \
        build_train_valid_test_data_iterators(train_valid_test_dataset_provider)

    # ---- Step 4: 训练 ----
    if not args.skip_train and args.do_train:
        iteration, num_flops = train(
            forward_step_func, model, optimizer, opt_param_scheduler,
            train_data_iterator, valid_data_iterator, config, ...
        )

    # ---- Step 5: 收尾 ----
    if args.do_valid:
        evaluate_and_print_results(...)
    if args.do_test:
        evaluate_and_print_results(...)
    maybe_finalize_async_save(blocking=True, terminate=True)

# 1. 初始化 PyTorch 分布式（NCCL 后端）
torch.distributed.init_process_group(
    backend=args.distributed_backend,   # 通常是 'nccl'
    world_size=args.world_size,
    rank=args.rank,
    timeout=timedelta(minutes=args.distributed_timeout_minutes),
)

# 2. 创建 Megatron 的多维并行通信组
mpu.initialize_model_parallel(
    tensor_model_parallel_size=args.tensor_model_parallel_size,
    pipeline_model_parallel_size=args.pipeline_model_parallel_size,
    virtual_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,
    order='tp-cp-ep-dp-pp',             # 默认 rank 排列顺序
    ...
)

def _set_random_seed(seed_, data_parallel_random_init, ...):
    seed = seed_
    if data_parallel_random_init:
        seed = seed_ + mpu.get_data_parallel_rank()

    # ① 默认 RNG：所有 TP rank 种子相同
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

    # ② TP 专用 RNG：每个 TP rank 种子不同（加上 tp_rank 偏移）
    _CUDA_RNG_STATE_TRACKER.add(
        'model-parallel-rng',
        seed + mpu.get_tensor_model_parallel_rank()
    )

def build_model():
    if pp_size > 1 and virtual_pp_size is not None:
        # VPP 模式：构建 virtual_pp_size 个 model chunk
        model = []
        for i in range(args.virtual_pipeline_model_parallel_size):
            pre_process = mpu.is_pipeline_first_stage(
                ignore_virtual=False, vp_stage=i
            )
            post_process = mpu.is_pipeline_last_stage(
                ignore_virtual=False, vp_stage=i
            )
            this_model = model_provider_func(
                pre_process=pre_process,
                post_process=post_process,
                vp_stage=i          # 告诉 provider 这是第几个虚拟 chunk
            )
            model.append(this_model)
    else:
        # 标准模式：只构建 1 个 model
        pre_process = mpu.is_pipeline_first_stage()
        post_process = mpu.is_pipeline_last_stage()
        model = model_provider_func(
            pre_process=pre_process,
            post_process=post_process
        )

# 1. 移动到 GPU（FSDP2 除外，它在 fully_shard 时分配显存）
for model_module in model:
    model_module.cuda(torch.cuda.current_device())

# 2. FP16/BF16 精度包装
if args.fp16 or args.bf16:
    config = get_model_config(model[0])
    model = [Float16Module(config, model_module) for model_module in model]

# 选择 DP 实现
if args.use_torch_fsdp2:
    DP = torch_FSDP
elif args.use_megatron_fsdp:
    DP = megatron_FSDP
else:
    DP = DDP   # Megatron 自研 DDP

# DDP 配置（仅 Megatron DDP）
ddp_config = DistributedDataParallelConfig(
    grad_reduce_in_fp32=args.accumulate_allreduce_grads_in_fp32,
    overlap_grad_reduce=args.overlap_grad_reduce,   # 梯度通信与计算重叠
    overlap_param_gather=args.overlap_param_gather,  # 参数收集与计算重叠
    bucket_size=args.ddp_bucket_size or auto_bucket_size,
    ...
)

# 包装每个 model chunk
model = [
    DP(
        config=config,
        ddp_config=ddp_config,
        module=model_chunk,
        # VPP 模式下，model_chunk 1+ 的通信已经和计算重叠了
        # 所以关闭它们的 bucketing
        disable_bucketing=(model_chunk_idx > 0),
    )
    for model_chunk_idx, model_chunk in enumerate(model)
]

# 从 args 构建 OptimizerConfig
config = OptimizerConfig(**{
    f.name: getattr(args, f.name)
    for f in dataclasses.fields(OptimizerConfig)
    if hasattr(args, f.name)
})

# 根据 --optimizer 参数选择优化器
if 'muon' not in config.optimizer:
    # 标准路径：Adam / AdamW（分布式优化器）
    optimizer = get_megatron_optimizer(
        config, model,
        no_weight_decay_cond,    # 哪些参数不要 weight decay
        scale_lr_cond,           # 哪些参数用不同学习率
        ...
    )
else:
    # Muon 优化器（实验性）
    optimizer = get_megatron_muon_optimizer(
        config, model, ...
    )

if args.load is not None or args.pretrained_checkpoint is not None:
    # 加载 checkpoint，恢复：
    #   - 模型权重
    #   - 优化器状态（Adam 的 m, v）
    #   - LR 调度器状态
    #   - 训练进度（iteration, consumed_samples）
    args.iteration, args.num_floating_point_operations_so_far = load_checkpoint(
        model,
        optimizer,
        opt_param_scheduler,
        checkpointing_context=checkpointing_context,
    )
else:
    # 从零开始训练
    args.iteration = 0
    args.num_floating_point_operations_so_far = 0

if args.virtual_pipeline_model_parallel_size is not None:
    # VPP 模式：每个 vp_stage 有独立的 data iterator
    train_data_iterator = []
    valid_data_iterator = []
    test_data_iterator = []
    for vp_stage in range(len(model)):
        vp_provider = functools.partial(
            train_valid_test_dataset_provider,
            vp_stage=vp_stage
        )
        iterators = build_train_valid_test_data_iterators(vp_provider)
        train_data_iterator.append(iterators[0])
        valid_data_iterator.append(iterators[1])
        test_data_iterator.append(iterators[2])
else:
    # 标准模式：单一 data iterator
    train_data_iterator, valid_data_iterator, test_data_iterator = \
        build_train_valid_test_data_iterators(
            train_valid_test_dataset_provider
        )