Megatron-LM pretrain_gpt.py 源码精读

if __name__ == "__main__":

    # Temporary for transition to core datasets
    train_valid_test_datasets_provider.is_distributed = True

    # Optionally enable inprocess restart on pretrain
    pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain)

    pretrain(
        train_valid_test_datasets_provider,          # 回调①：构建数据集
        partial(model_provider, gpt_builder),         # 回调②：构建模型
        ModelType.encoder_or_decoder,                 # 模型类型枚举
        forward_step,                                 # 回调③：单步前向
        args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
        extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None,
        store=store,
        get_embedding_ranks=get_embedding_ranks,
    )

def get_batch(data_iterator, vp_stage: Optional[int] = None):
    """Generate a batch."""
    args = get_args()
    config = core_transformer_config_from_args(args)
    # 如果不在 PP 首尾 stage，且不是 MTP rank → 不需要数据
    if not is_first_or_last_pipeline_stage(vp_stage) and (
    (not mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage))):
        return None, None, None, None, None, None, None

    # 只在 TP rank 0 上真正读数据，其他 TP rank 通过 broadcast 接收
    batch = get_batch_on_this_tp_rank(
        data_iterator,
        mtp_on_this_rank=mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage)
    )

    # Multi-Token Prediction：在 CP 切分之前，预先准备多个预测头的数据
    mtp_data = None
    if config.mtp_num_layers and mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage):
        mtp_data = prepare_mtp_data(
            batch['tokens'], batch['position_ids'],
            batch['labels'], batch['loss_mask'],
            config.mtp_num_layers, cu_seqlens=cu_seqlens,
        )

    # 如果同时有 MTP 和 packed/hybrid CP，需要把 MTP 数据混入 batch
    # 这样后续 CP 切分会一并处理它们
    mtp_needs_batch_split = mtp_data is not None and (cu_seqlens is not None or local_cp_size is not None)
    if mtp_needs_batch_split:
        for k in range(config.mtp_num_layers):
            batch[f'mtp_input_ids_{k}'] = mtp_data['mtp_input_ids'][k]
            batch[f'mtp_position_ids_{k}'] = mtp_data['mtp_position_ids'][k]
            # ...

    if cu_seqlens is None and local_cp_size is None:
        # 普通 CP：沿序列维度均匀切分
        batch = get_batch_on_this_cp_rank(batch)
        packed_seq_params = None
    elif local_cp_size is None:  # Packed THD format
        # 变长序列打包格式（多条不等长序列拼在一起）
        batch, packed_seq_params = get_thd_batch_on_this_cp_rank(
            batch, cu_seqlens, cu_seqlens_padded, max_seqlen
        )
    else: # Hybrid CP format
        batch, packed_seq_params = get_batch_on_this_hybrid_cp_rank(batch, local_cp_size)

    # 最终返回 7 个元素
    return (*batch.values(), packed_seq_params, mtp_data)

def loss_func(
    loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optional[GPTModel] = None
):
    """Loss function.

    Args:
        loss_mask: 用于屏蔽 padding token 的 loss
        output_tensor: 模型输出的 per-token loss
        model: GPT 模型实例（用于 ModelOpt 蒸馏场景）
    """
    args = get_args()

    # 标准 loss 计算
    losses = output_tensor.view(-1).float()     # 展平为一维
    loss_mask = loss_mask.view(-1).float()       # 展平 mask
    loss = torch.sum(losses * loss_mask)         # 只对非 padding token 求和

    num_tokens = loss_mask.sum().clone().detach().to(torch.int)
    report = {'lm loss': torch.cat([loss.clone().detach().view(1), num_tokens.view(1)])}

    # NaN / Inf 检查
    rerun_state_machine = get_rerun_state_machine()
    if args.check_for_nan_in_loss_and_grad:
        rerun_state_machine.validate_result(
            result=loss,
            rejection_func=torch.isnan,                    # 检测 NaN
            message="found NaN in local forward loss calculation",
            tolerance=0.0,       # 前向计算是确定性的，不允许任何容忍度
            fatal=True,          # 致命错误，必须 rerun
        )
        rerun_state_machine.validate_result(
            result=loss,
            rejection_func=torch.isinf,                    # 检测 Inf
            message="found Inf in local forward loss calculation",
            tolerance=0.0,
            fatal=True,
        )

    # Spiky loss 检查（loss 突然飙升）
    if args.check_for_spiky_loss:
        rerun_state_machine.validate_result(
            result=loss,
            rejection_func=partial(
                rerun_state_machine.is_unexpectedly_large,
                threshold=SPIKY_LOSS_FACTOR,               # 阈值 = 10x
                context="loss",
            ),
            message="Spiky loss",
            tolerance=0.0,
            fatal=False,         # 非致命，会 rerun 但不会崩溃
        )

def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = False):
    """Forward training step."""
    args = get_args()
    timers = get_timers()

    # ① 取数据
    timers('batch-generator', log_level=2).start()
    global stimer
    with stimer(bdata=True):
        vp_stage = get_attr_wrapped_model(model, "vp_stage")
        tokens, labels, loss_mask, attention_mask, position_ids, \
            packed_seq_params, mtp_data = get_batch(data_iterator, vp_stage)
    timers('batch-generator').stop()

    # ② 前向计算
    with stimer:
        output_tensor = model(
            tokens, position_ids, attention_mask,
            labels=labels, loss_mask=loss_mask,
            packed_seq_params=packed_seq_params, mtp_data=mtp_data,
        )

    # ③ 返回 (模型输出, 延迟执行的 loss 函数)
    return output_tensor, partial(loss_func, loss_mask, model=model)

def core_gpt_dataset_config_from_args(args):
    if args.legacy_tokenizer:
        tokenizer = get_tokenizer()
    else:
        tokenizer = build_tokenizer(args)

    # 从命令行参数解析数据混合比例
    blend, blend_per_split = get_blend_and_blend_per_split(args)

    return GPTDatasetConfig(
        random_seed=args.seed,
        sequence_length=args.seq_length,
        blend=blend,                          # 多数据源混合比例
        blend_per_split=blend_per_split,       # train/val/test 各自的混合
        split=args.split,                      # "98,1,1" 格式的 train/val/test 比例
        tokenizer=tokenizer,
        reset_position_ids=args.reset_position_ids,
        reset_attention_mask=args.reset_attention_mask,
        eod_mask_loss=args.eod_mask_loss,      # EOD 后的 token 是否计算 loss
        context_parallel_size=args.context_parallel_size,
        data_parallel_size=args.data_parallel_size,
        # ...
    )

def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None):
    """Build the train test and validation datasets."""
    args = get_args()
    config = core_gpt_dataset_config_from_args(args)

    # 根据训练模式选择数据集类型
    if args.sft:
        dataset_type = SFTDataset           # SFT 指令微调
    else:
        if args.mock_data:
            dataset_type = MockGPTDataset   # 假数据（调试用）
        else:
            dataset_type = GPTDataset       # 标准预训练

    # BlendedMegatronDatasetBuilder 支持多数据集按比例混合
    train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
        dataset_type,
        train_val_test_num_samples,
        partial(is_dataset_built_on_rank, vp_stage=vp_stage),
        config
    ).build()

    return train_ds, valid_ds, test_ds

def model_provider(
    model_builder: Callable,
    pre_process=True,     # 是否包含 embedding（PP 第一个 stage）
    post_process=True,    # 是否包含 output head（PP 最后一个 stage）
    vp_stage=None         # Virtual Pipeline stage 编号
) -> Union[GPTModel, MambaModel]:
    args = get_args()

    # OOM 时自动保存 memory snapshot（调试用）
    if args.record_memory_history:
        torch.cuda.memory._record_memory_history(True, ...)
        torch._C._cuda_attach_out_of_memory_observer(oom_observer)

    return model_builder(args, pre_process, post_process, vp_stage)

def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None):
    # ① 构建 TransformerConfig（从命令行参数转化）
    if config is None:
        config = core_transformer_config_from_args(args)

    # ② 选择 Transformer Layer 的具体实现规格 (spec)
    if args.spec is not None:
        transformer_layer_spec = import_module(args.spec)     # 用户自定义 spec
    else:
        use_te = args.transformer_impl == "transformer_engine"
        if args.num_experts:
            # MoE 模型：使用 decoder block spec（包含 Router + Expert）
            transformer_layer_spec = get_gpt_decoder_block_spec(config, use_te, ...)
        else:
            # 标准 Dense 模型
            if use_te:
                transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(...)
            else:
                transformer_layer_spec = get_gpt_layer_local_spec(...)

    # ③ (可选) Multi-Token Prediction
    mtp_block_spec = None
    if args.mtp_num_layers is not None:
        mtp_block_spec = get_gpt_mtp_block_spec(config, ...)

    # ④ 实例化 GPTModel
    model = GPTModel(
        config=config,
        transformer_layer_spec=transformer_layer_spec,
        vocab_size=args.padded_vocab_size,
        max_sequence_length=args.max_position_embeddings,
        pre_process=pre_process,              # PP: 是否包含 embedding
        post_process=post_process,            # PP: 是否包含 output head
        share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
        position_embedding_type=args.position_embedding_type,
        rotary_base=args.rotary_base,
        mtp_block_spec=mtp_block_spec,
    )
    return model

def is_dataset_built_on_rank(vp_stage=None):
    args = get_args()
    config = core_transformer_config_from_args(args)
    return (
        is_first_or_last_pipeline_stage(vp_stage)
        or mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage)
    ) and parallel_state.get_tensor_model_parallel_rank() == 0

def get_embedding_ranks(pp_ranks: List[int]):
    """确定哪些 PP rank 需要持有 embedding 权重。"""
    embedding_ranks = [pp_ranks[0]]        # 第一个 stage 一定需要
    if len(pp_ranks) > 1:
        args = get_args()
        if not args.untie_embeddings_and_output_weights:
            embedding_ranks.append(pp_ranks[-1])   # 权重共享时，最后一个 stage 也需要
        config = core_transformer_config_from_args(args)
        mtp_ranks = get_mtp_ranks(pp_ranks, config)
        embedding_ranks.extend(mtp_ranks)          # MTP rank 也需要 embedding
    return sorted(list(set(embedding_ranks)))