资讯详情

\[\newcommand{\t}{\text} \]

本文考虑各种并行方案在 Megatron-LM 框架下的具体实现。相关 notation 同 前篇。本文大量使用 AI 辅助。

注意：本文基于官方 Megatron-LM v0.19.0 main branch 的某个开发中状态，可能与前一个稳定版本 v0.18.0 有区别，请注意分别。

I. 概述

PretrainConfigContainer from megatron.training.config.container 是训练流程 理论上的顶层配置容器。一部分训练配置会从扁平的 args 被解析为语义丰富的该容器，但并非所有配置。解析后会得到该容器的一份实例，以及它内部持有的那一组子配置对象。

理解这套配置体系的关键在于搞清楚两件事：每个子配置管什么，以及 我想改某个功能时该动哪里。本节围绕这两条线索展开。

I.I. 配置继承树

PretrainConfigContainer 通过 组合 持有所有子配置。以下是它持有的全部子配置及其职责：

• model: HybridModelConfig | GPTModelConfig：

  • 在 上一个稳定版本 (v0.18.0)，尚不支持 GPTModelConfig。目前开发中的 v0.19.0 才接入新选项。
  • 与模型具体架构有关的一切：层数、隐藏维度、head 数、FFN 维度、归一化方式、激活函数、dropout 率、位置编码类型、vocab size、序列长度等。
  • GPTModelConfig from megatron.training.models.gpt 和 HybridModelConfig from megatron.training.models.hybrid 都继承自 ModelConfig from megatron.training.models.base，因此如果要引入全新的模型架构（比如 RWKV、Mamba-2），同样应该继承 ModelConfig 并替换 model 项，但此时需要的适配会很复杂。
  • GPTModelConfig 与 HybridModelConfig 均通过 组合 持有 transformer: TransformerConfig，并通过 __getattr__ 代理可以直接访问其中配置。
    • TransformerConfig from megatron.core.transformer.transformer_config 继承自 ModelParallelConfig from megatron.core.model_parallel_config，后者定义了并行阶段通用的一些配置，如 TP/PP/CP/EP/SP 并行度等。
    • TransformerConfig 在此基础上添加了 Transformer 层相关的配置，如模型具体架构等。
    • 如果要引入 DeepSeek 风格的 Multi-Latent Attention，可进一步继承为 MLATransformerConfig；如果要支持 Nemotron-NAS 的异构层架构（每层可以不同），则使用 HeterogeneousTransformerConfig。

• optimizer: OptimizerConfig from megatron.core.optimizer.optimizer_config：

  • optimizer 相关配置，以及 ZeRO 等分布式 optimizer 方法。

• scheduler: SchedulerConfig from megatron.training.config.training_config：

  • 学习率 schedule、warmup 策略、weight decay schedule 相关配置。

• ddp: DistributedDataParallelConfig from megatron.core.distributed.distributed_data_parallel_config：

  • DDP / FSDP 的分片策略（no_shard = DDP，optim = ZeRO-1，optim_grads = ZeRO-2，optim_grads_params = ZeRO-3 即 FSDP）。
  • 梯度 reduce overlap、参数 all-gather overlap、低精度通讯（FP32 reduce、FP8/FP4 param gather）。
  • 进一步被 TorchFullyShardedDataParallelConfig from megatron.core.distributed.torch_fully_sharded_data_parallel_config 继承。

• train: TrainingConfig from megatron.training.config.training_config：

  • micro/global batch size、train iters/samples。
  • 退出条件（interval/duration/signal handler）、手动 GC 控制等训练循环本身的行为。

• validation: ValidationConfig from megatron.training.config.training_config：

  • 验证频率、eval batch size、skip train、test mode、多数据集验证。

• dist: DistributedInitConfig from megatron.training.config.common_config：

  • 分布式后端选择（nccl/gloo）、通信超时、进程组创建策略。
  • SHARP、flight recorder 等初始化阶段的设置。

• rng: RNGConfig from megatron.training.config.common_config：

  • 随机数种子、TE/inference RNG tracker、DP random init。

• logger: LoggerConfig from megatron.training.config.training_config：

  • TensorBoard、WandB、timing、memory logging、energy monitoring。

• checkpoint: CheckpointConfig from megatron.training.config.training_config：

  • checkpoint save/load 的路径与格式。
  • async save、fully parallel save/load、replication、integrity verification。

• profiling: ProfilingConfig from megatron.training.config.common_config：

  • nsys profiler、PyTorch profiler、memory snapshot。

• tokenizer: TokenizerConfig from megatron.training.config.training_config：

  • vocab size、tokenizer type/model、chat template。

• rerun_state_machine: RerunStateMachineConfig from megatron.training.config.resilience_config：

  • error injection、rerun mode、NaN/spiky loss detection。

• straggler: StragglerDetectionConfig | None from megatron.training.config.resilience_config：

  • 慢节点检测。

除了显式提到的几个类以外，其它 config 均为独立类型，不继承 Megatron 内部的任何基类。

I.II. 具体配置方法

共有三种基本等价的配置方法，均能产出 PretrainConfigContainer 示例，但实际应用时仍大量依赖直接对 args 的读取。

I.II.I. CLI 参数

这是最常见的使用方式。用户通过 --tensor-model-parallel-size 8 这样的命令行参数启动训练，整个流程分为 参数注册 和 config 构造 两个阶段。

参数注册阶段。

• 入口是 parse_and_validate_args() from megatron.training.arguments，它会调用同一文件下的 parse_args() 进行 parser 注册和 args 生成。
• parse_args() 对 parser 注册的方式包括：
  • add_megatron_arguments() from megatron.training.arguments 人工调用了很多内置方法添加 parser，同时也调用了 ArgumentGroupFactory from megatron.training.argument_utils 从 dataclass 字段自动生成，基本上相当于对 parser.parse_args() 的逆向。
  • parse_and_validate_args() 支持输入一个 extra_args_provider，它会被直接传给 parse_args() 以进行外部自定义参数的注册，。
• 注册完后，用 args = parser.parse_args() 解析得到一个 args: argparse.Namespace 对象，并返回。

config 构造阶段。对上一阶段得到的扁平对象 args，进行深度解析。具体解析方式的常见示例如下：

args = parse_and_validate_args(...)
model_cfg = gpt_config_from_args(args)
full_config = pretrain_cfg_container_from_args(args, model_cfg)

其中 gpt_config_from_args() from megatron.training.argument_utils 用于生成 GPTModelConfig；如果要使用 HybridModelConfig 或其它自定义配置，需要重写相关的函数。

而 pretrain_cfg_container_from_args() from megatron.training.argument_utils 在此基础上依次生成若干子配置，并聚合出一个 PretrainConfigContainer 实例。其具体模式如下：

• HybridModelConfig | GPTModelConfig 持有的 TransformerConfig 由 core_transformer_config_from_args() from megatron.training.argument_utils 负责，包含大量字段名映射逻辑。
• 其余大部分子配置（TrainingConfig、SchedulerConfig、DistributedInitConfig 等）由 _default_config_from_args() from megatron.training.argument_utils 通过字段名 1:1 匹配自动构造——它遍历 dataclass 的字段，如果 args 中有同名属性就取出来，否则跳过（使用默认值）。
• OptimizerConfig 和 DistributedDataParallelConfig 各自由专门的函数 get_megatron_optimizer_config() 和 get_megatron_ddp_config() 构造，因为它们的构造逻辑较复杂。
• CheckpointConfig、ProfilingConfig、RerunStateMachineConfig 虽然也用 _default_config_from_args()，但额外做了字段名重映射（例如 args.no_save_optim → save_optim=not args.no_save_optim、args.profile → use_nsys_profiler=True）。

I.II.II. Dataclass 直接构造

在 Python 代码中直接实例化 PretrainConfigContainer 及其子配置，绕过 argparse 和字段名映射。需要人工定义所有持有的配置各一份。

这种方式产出的对象与 CLI 方式 完全相同——都是 PretrainConfigContainer 的实例，都经过 __post_init__ 校验。区别在于绕过了上述工厂函数中的字段名映射和默认值填充，需要用户自行保证字段名正确。

I.II.III. YAML 声明式配置

Megatron-LM 中存在 两套 YAML 配置机制。

旧版 YAML：如果在 CLI 中指定 --yaml-cfg path/to/config.yaml，则会在parse_args() 中调用 load_yaml() from megatron.training.yaml_arguments，将 YAML 文件解析为嵌套的 SimpleNamespace 对象，并 整体替换 argparse 解析出的 args。后续流程与纯 CLI 完全一致。

新版 YAML：PretrainConfigContainer 自带的反序列化能力。它完全绕过 argparse，直接从 YAML 构造出 PretrainConfigContainer 实例。

I.III. 自定义方法

Megatron-LM 的配置体系对扩展是开放的，但扩展的路径取决于你愿意动多少源码。总体上分为两种策略：不改 Megatron-LM 源码（pip install 并 import，适用于外部集成）和 改源码（适用于长期 fork 或深度定制）。

大多数定制需求的起点是自定义一个 MyConfigContainer 继承 PretrainConfigContainer，得到一个扩展容器；大部分配置保持不变，只不过把其中的某个子配置替换成自己扩展过的版本，或者增加一个新配置。

比如你引入了一种新的 MoE 路由策略，需要一个全新的 MyExpertConfig；同时还希望试用一种新 scheduler，对应 MySchedulerConfig，那么：

1. 从零定义 MyExpertConfig dataclass；定义 MySchedulerConfig 继承 SchedulerConfig。
2. 自定义一个 MyConfigContainer 继承 PretrainConfigContainer，添加 my_expert_config: MyExpertConfig 字段，将 scheduler 的类型替换为自定义的 MySchedulerConfig。
3. 然后在 model builder、router、scheduler factory 或训练循环中实际读取新字段。

CLI / YAML 的注册。 如果你扩展了 config 但还想通过 CLI 或 YAML 来配置它，需要额外的注册步骤：

CLI：

• 最简单的方式是用 extra_args_provider 回调。如果你的自定义 config 是完整的 dataclass，也可以直接用 ArgumentGroupFactory from megatron.training.argument_utils 将其全部字段自动注册为 CLI 参数。
• 如果要改源码，则在 megatron.training.arguments 中添加 _add_my_args(parser)，并在 core_transformer_config_from_args() from megatron.training.argument_utils 中添加映射逻辑。

YAML：太复杂不搞了。

II. 单卡场景

本节集中于构建单卡训练需要的相关配置。

II.I. 入口与最小配置

Megatron-LM 提供两种典型的单卡训练入口：

完整训练脚本 pretrain_gpt.py：调用 PretrainConfigContainer 体系，走完整的初始化 → 模型构建 → 数据加载 → 训练循环流程。适用于与生产环境对齐的训练。

最小独立循环 examples/run_simple_mcore_train_loop.py：绕过 PretrainConfigContainer 和 argparse，直接构造 TransformerConfig 和 GPTModel，手动管理 forward-backward-optimizer 循环。适用于理解 MCore API 的最小用法。但是官方实现其实设置了 \(N_\t{TP}=2\)，因此与单卡场景不完全相同。

通过 CLI 使用 pretrain_gpt.py 时，必须手动设置 的参数为：

• 模型架构：--num-layers、--hidden-size、--num-attention-heads。
• 训练流程：--micro-batch-size、--train-iters、--lr。
• 数据：--data-path 和 --vocab-file，或使用 --mock-data 进行无数据测试。

II.II. 模型架构

此部分的配置集中于 TransformerConfig。

II.II.I. 核心必填字段

以下三个字段没有合理默认值，必须由用户指定：

• num_layers: int — Transformer 层数 \(L\)。
• hidden_size: int — 隐藏维度 \(H\)。
• num_attention_heads: int — 注意力头数 \(N_q\)。

II.II.II. 自动推导字段

以下字段在未显式设置时自动推导：

• ffn_hidden_size: Optional[int] = None — FFN 中间维度 \(H_\t{FFN}\)。默认取 4 * hidden_size；如果使用 SwiGLU，一般需要手动对齐到 8 / 3 * hidden_size。
• kv_channels: Optional[int] = None — 单头维度 \(d\)。默认取 hidden_size // num_attention_heads。（为什么可以不等？因为 QKV proj 和 attn proj 承担了衔接桥梁！）
• num_query_groups: Optional[int] = None — KV 头数 \(N_{kv}\)。默认取 num_attention_heads（即标准 MHA）。设为较小值可启用 GQA；设为 1 即 MQA。

II.II.III. Normalization

• normalization: str = "LayerNorm" — 取值 "LayerNorm" 或 "RMSNorm"。LLaMA 系列使用 RMSNorm。
• layernorm_epsilon: float = 1e-5 — 归一化的 \(\epsilon\)。
• layernorm_zero_centered_gamma: bool = False — 启用中心化 gamma 的 LayerNorm，改善数值稳定性。
• qk_layernorm: bool = False — 对 Q 和 K 做归一化。
• qk_l2_norm: bool = False — LLaMA4 风格的 QK L2 归一化。

II.II.IV. Activation

• gated_linear_unit: bool = False — 启用门控线性单元（SwiGLU / GeGLU）。启用时 FFN 从标准的 \(\t{GeLU}(XA)B\) 变为 \(\t{GeLU}(XA_1) \odot (XA_2) B\) 的形式。
• activation_func: Callable = F.gelu — 激活函数。配合 gated_linear_unit=True 使用时，常见设为 F.silu（LLaMA 风格）。

II.II.V. Position Embedding

本小节的相关内容来自 GPTModelConfig 而不是 TransformerConfig。

• position_embedding_type: Literal["learned_absolute", "rope", "mrope", "yarn", "none"] = "learned_absolute" — 位置编码类型。
  • "learned_absolute"：可学习的绝对位置编码，GPT-2/3 风格。
  • "rope"：旋转位置编码，LLaMA/Qwen 风格。进一步由 rotary_base、rotary_percent、rope_scaling 控制。
  • "mrope"：多模态旋转位置编码（Qwen 风格 MRoPE），通过 mrope_section 指定通道分割。
  • "yarn"：Yet Another RoPE extensioN，支持更长上下文的外推。
  • "none"：不使用位置编码。
• rotary_base: float = 10000 — RoPE 基频。LLaMA 3 使用 1000000。
• rotary_percent: float = 1.0 — RoPE 作用于头维度的比例。设为小于 1.0 时部分维度不旋转。
• rope_scaling: bool = False — 启用 RoPE 缩放（LLaMA 3.x 风格的低频/高频平滑）。
• rope_scaling_factor: float = 8.0 — RoPE 缩放因子。
• rotary_interleaved: bool = False — RoPE 风格：False 为 LLaMA 风格（连续配对），True 为 RoFormer 风格（交错配对）。
• no_rope_freq: Optional[Union[int, List[int]]] = None — 每隔 \(N\) 层跳过一次 RoPE。

II.II.VI. Attention

• attention_backend: AttnBackend = AttnBackend.auto — 注意力后端。"auto" 由 TransformerEngine 自动选择；"local" 使用 MCore 原生实现。
• attention_dropout: float = 0.1 — 注意力权重后的 dropout。
• hidden_dropout: float = 0.1 — 隐藏状态的 dropout。
• window_size: Optional[Tuple[int, int]] = None — 滑动窗口注意力的窗口大小。例如 (-1, 0) 表示仅关注左侧无限窗口。

II.II.VII. 初始化

• init_method: Optional[Callable] = None — 权重初始化函数。默认为 init_method_normal(0.02)，即 \(\mathcal{N}(0, 0.02)\)。
• init_method_std: float = 0.02 — 默认正态初始化的标准差。
• output_layer_init_method: Optional[Callable] = None — 输出层初始化，默认按 \(\sqrt{2L}\) 缩放。
• init_model_with_meta_device: bool = False — 在 meta device 上初始化模型，适用于超大模型避免 OOM。

II.II.VIII. Transformer 实现

• transformer_impl: str = "transformer_engine" — 取值 "transformer_engine"（使用 TE fused kernel）、"local"（使用 MCore 原生 PyTorch 实现）、"inference_optimized"（推理优化实现）。

II.II.IX. MTP

MTP 允许模型在每个位置同时预测未来多个 token，是 DeepSeek V3 等模型使用的训练技巧。

• mtp_num_layers: Optional[int] = None — MTP 头的 Transformer 层数。为 None 时不启用 MTP。
• mtp_detach_heads: bool = False — 是否对 MTP 头的输入做 detach（阻断梯度回传到主干）。
• mtp_use_repeated_layer: bool = False — 多个 MTP 头是否共享同一组层权重。
• mtp_hybrid_override_pattern: Optional[str] = None — MTP 层内的 MoE/Dense 模式（如 "*E" 表示全部 MoE 层）。

MTP 的实现位于 megatron.core.transformer.multi_token_prediction，在 GPTModel._postprocess() 中调用 MultiTokenPredictionBlock.forward()。其核心流程为：对每个 MTP 层，将主干隐藏状态与 roll 后的 input embedding 拼接，过一个 Transformer 层产出预测 logits，loss 按 1/num_microbatches 缩放。

II.II.X. Multi-Latent Attention (MLA)

MLA 是 DeepSeek V2/V3 引入的注意力变体，通过低秩压缩 KV 减少推理时的 KV Cache 开销。

• multi_latent_attention: bool = False — 启用 MLA。启用时需配合 MLATransformerConfig from megatron.core.transformer.transformer_config 使用，该类继承 TransformerConfig 并添加 MLA 特有的压缩维度等字段。

II.II.XI. Activation Checkpointing

AC 通过重计算激活值来换取显存节省，是典型的 计算换存储 策略。

• recompute_granularity: Optional[str] = None — 重计算粒度。取值：
  • "selective"：选择性重计算，仅重计算计算量大但存储小的激活值（如 attention 的 softmax 输入）。
  • "full"：完全重计算，把整个层的中间激活值全部丢弃，BP 时重新计算。
• recompute_method: Optional[str] = None — 重计算方法。取值：
  • "uniform"：均匀切分，每 recompute_num_layers 层做一次重计算。
  • "block"：按块切分，每个 block 做一次重计算。
• recompute_num_layers: Optional[int] = None — 每个重计算单元包含的层数。
• recompute_modules: Optional[List[str]] = None — 指定需要重计算的子模块，例如 ["core_attn"] 仅重计算 attention 核心部分。

II.III. 训练配置

训练流程相关的配置分散在多个子配置中。

II.III.I. TrainingConfig

TrainingConfig from megatron.training.config.training_config 控制训练循环本身的行为：

• micro_batch_size: Optional[int] = None — micro-batch size \(b\)。必须设置。
• global_batch_size: Optional[int] = None — global batch size \(B\)。为 None 时自动等于 micro_batch_size（单卡无 DP）。
• train_iters: Optional[int] = None — 总训练迭代数。与 train_samples 二选一。
• train_samples: Optional[int] = None — 总训练样本数。
• empty_unused_memory_level: int = 0 — torch.cuda.empty_cache() 的激进程度。0 = 不调用，1 = 适中，2 = 激进。单卡显存紧张时可设为 1。

II.III.II. SchedulerConfig

SchedulerConfig from megatron.training.config.training_config 控制学习率调度：

• lr_decay_style: str = "linear" — 衰减策略，取值 "constant"、"linear"、"cosine"、"inverse-square-root"、"WSD"（Warmup-Stable-Decay）。
• lr_warmup_iters: int = 0 — warmup 迭代数。
• lr_warmup_init: float = 0.0 — warmup 起始学习率。
• lr_decay_iters: Optional[int] = None — 衰减持续的迭代数。默认取 train_iters。
• start_weight_decay: Optional[float] = None — 初始 weight decay。
• end_weight_decay: Optional[float] = None — 最终 weight decay。

II.III.III. OptimizerConfig

OptimizerConfig from megatron.core.optimizer.optimizer_config 控制优化器行为：

• optimizer: str = "adam" — 优化器类型。取值 "adam"、"sgd"、"muon"、"lion"、"soap"、"adaptive_muon"。
• lr: Optional[float] = None — 初始学习率。必须设置。
• min_lr: Optional[float] = None — 最小学习率。默认取 lr（无衰减）。
• adam_beta1: float = 0.9 — Adam/AdamW 的 \(\beta_1\)。
• adam_beta2: float = 0.999 — Adam/AdamW 的 \(\beta_2\)。
• adam_eps: float = 1e-8 — Adam/AdamW 的 \(\epsilon\)。
• decoupled_weight_decay: bool = True — 是否启用 decoupled WD。
• weight_decay: float = 0.01 — L2 正则 / weight decay 系数。
• clip_grad: float = 1.0 — 全局 L2 梯度裁剪阈值。
• use_distributed_optimizer: bool = False — 是否启用 ZeRO-1。单卡时不需要。
• optimizer_cpu_offload: bool = False — 将优化器状态 offload 到 CPU。

FP16 Loss Scaling：

• loss_scale: Optional[float] = None — 手动设置 loss scale。为 None 时使用动态 loss scaling。
• initial_loss_scale: float = 2**32 — 动态 loss scaling 的初始值。
• min_loss_scale: float = 1.0 — 动态 loss scaling 的最小值。
• loss_scale_window: float = 1000 — 动态 loss scaling 的调整窗口。

II.III.IV. Mixed Precision

混合精度相关配置分布在 ModelParallelConfig 和 OptimizerConfig 中：

• fp16: bool = False（ModelParallelConfig）— 启用 FP16 混合精度训练。
• bf16: bool = False（ModelParallelConfig）— 启用 BF16 混合精度训练。
• params_dtype: torch.dtype = torch.float32（ModelParallelConfig）— 权重初始化的 dtype。

FP8 训练：

• fp8: Optional[str] = None（TransformerConfig）— 取值 "e4m3" 或 "hybrid"，启用 FP8 训练。需要 H100 等支持 FP8 的硬件。
• fp8_recipe: str = "delayed"（TransformerConfig）— FP8 策略，取值 "tensorwise"、"delayed"、"mxfp8"、"blockwise"。
• fp8_param: bool = False（TransformerConfig）— 权重本身也保持 FP8 格式。

II.IV. 日志与检查点

II.IV.I. LoggerConfig

LoggerConfig from megatron.training.config.training_config：

• log_interval: int = 100 — 每隔多少步打印 loss 和 timing 信息。
• tensorboard_dir: Optional[str] = None — TensorBoard 日志目录。
• wandb_project: Optional[str] = None — Weights & Biases 项目名。

II.IV.II. CheckpointConfig

CheckpointConfig from megatron.training.config.training_config：

• save: Optional[str] = None — checkpoint 输出目录。
• save_interval: Optional[int] = None — 每隔多少步保存 checkpoint。
• load: Optional[str] = None — 从指定目录加载 checkpoint。
• finetune: bool = False — 仅加载模型权重（不加载 optimizer 和 RNG 状态），并重置 iteration 计数。用于 finetune 场景。
• ckpt_format: str = "torch_dist" — checkpoint 格式，取值 "torch"、"torch_dist"（分布式保存）、"torch_dcp"、"fsdp_dtensor"。
• async_save: bool = False — 异步保存 checkpoint，不阻塞训练。

II.IV.III. ProfilingConfig

ProfilingConfig from megatron.training.config.common_config：

• use_nsys_profiler: bool = False — 启用 Nsight Systems profiler。
• use_pytorch_profiler: bool = False — 启用 PyTorch profiler。
• profile_step_start: int = 10 — 从第几步开始 profiling。
• profile_step_end: int = 12 — 到第几步停止 profiling。

II.VI. 其它配置

II.VI.I. RNGConfig

RNGConfig from megatron.training.config.common_config：

• seed: int = 1234 — 全局随机种子。
• data_parallel_random_init: bool = False — 跨 DP rank 使用不同随机初始化。单卡时无影响。

II.VI.II. ValidationConfig

ValidationConfig from megatron.training.config.training_config：

• eval_iters: int = 100 — 每次评估运行的迭代数。
• eval_interval: Optional[int] = None — 每隔多少步进行一次评估。
• skip_train: bool = False — 仅评估，跳过训练。

II.VI.III. Deterministic Mode

• deterministic_mode: bool = False（ModelParallelConfig）— 启用确定性模式，保证相同输入产生相同输出。调试时有用，但可能降低性能。
• gradient_accumulation_fusion: bool = False（ModelParallelConfig）— 启用梯度累积与 backward 的融合（需要 APEX）。单卡时可尝试启用以提升性能。

III. 多卡场景

III.I. 总览

Megatron-LM 实现了五维并行空间：

• TP：tensor_model_parallel_size（ModelParallelConfig）。
• PP：pipeline_model_parallel_size（ModelParallelConfig）。
• CP：context_parallel_size（ModelParallelConfig）
• EP：expert_model_parallel_size（ModelParallelConfig）。
• DP：data_parallel_size，会被自动计算。

其中，总 GPU 数目 world_size 通过 torchrun 中的 WORLD_SIZE 设置；而 DP size 通过以下公式计算：

model_size = tensor_model_parallel_size * pipeline_model_parallel_size * context_parallel_size
data_parallel_size: int = world_size // model_size

正如理论分析，MoE 层的 contract 可能不同于 attn 层的 contract：

• expert_model_parallel_size：expert 的 EP size。
• expert_tensor_parallel_size（ModelParallelConfig），衡量了 expert 的 TP size，默认值是 tensor_model_parallel_size。
• expert_data_parallel_size：expert 的 DP size。会通过以下方式自动计算。

expert_tensor_model_pipeline_parallel_size =      expert_tensor_parallel_size * expert_model_parallel_size * pipeline_model_parallel_size
expert_data_parallel_size = world_size // expert_tensor_model_pipeline_parallel_size

III.I.I. Rank 映射顺序

torchrun 默认会把卡按照顺序编号：如果每个节点有 \(8\) 张卡，也即 GPUS_PER_NODE = 8，则 node 0 会对应编号 0-7，node 1 对应编号 8-15，以此类推。

而 Megatron-LM 会把编号分到具体的同步组里。分配方式由 initialize_model_parallel() from megatron.core.parallel_state 的 order 参数控制：

• 默认：Dense 层使用 'tp-cp-dp-pp'，而 Sparse 层使用 etp-ep-edp-pp。
• 可选：启用 --use-tp-pp-dp-mapping 会交换 PP 和 DP 的位置。

以 Dense 的 tp-cp-dp-pp 模式为例，则有以下表格：

维度	步长	含义
tp	1	TP 组内 rank 相邻
cp	tp_size	CP 次之
dp	tp_size \(\times\) cp_size	DP 再次
pp	tp_size \(\times\) cp_size \(\times\) dp_size	PP 最慢变化

而 sparse 阶段的编号可能不同于 dense 阶段。

[!TIP]

这种分配方式从形式和功能上均类似 DRAM 的地址分配：

• 形式上均满足「编号的若干位对应某一组织形式」。
• 功能上均为了利用相邻性质：
  • TP 需要处于同一个节点下。
  • DRAM 需要利用 cached page。

要调整分配方式，可以通过调整 torchrun 为编号分配卡的过程或 megatron 给编号分组的过程实现。

III.I.II. 必要约束

以下约束在 ModelParallelConfig.__post_init__ 和 TransformerConfig.__post_init__ 中强制执行：

• num_attention_heads % tensor_model_parallel_size == 0 — TP 必须整除注意力头数。
• num_query_groups 必须是 tensor_model_parallel_size 的整数倍或约数（GQA 的 KV 头切分要求）。
• sequence_parallel=True 要求 tensor_model_parallel_size > 1。
• pipeline_model_parallel_size > 1 时必须设置 pipeline_dtype。
• expert_model_parallel_size > 1 时 num_moe_experts 不能为 None。
• virtual_pipeline_model_parallel_size 不为 None 时要求 pipeline_model_parallel_size > 1。

III.II. DP

梯度累积步数 num_micro_batches 满足 num_micro_batches = global_batch_size / (micro_batch_size * data_parallel_size)。这部分逻辑位于 ConstantNumMicroBatchesCalculator from megatron.core.num_microbatches_calculator。

Megatron-LM 提供三种独立的 DP 实现，分布在 DistributedDataParallelConfig 中：

III.II.I. 标准 DDP + DistributedOptimizer

对应 ZeRO-1。通过 use_distributed_optimizer: bool = True 启用。

关键字段：

• overlap_grad_reduce: bool = False — 梯度 reduce-scatter 与 backward 计算重叠。生产环境建议启用。
• overlap_param_gather: bool = False — 参数 all-gather 与 forward 计算重叠（也即开启 prefetch）。生产环境建议启用。
• bucket_size: Optional[int] = None — 通讯 bucket 的最大参数量。控制梯度打包粒度。
• grad_reduce_in_fp32: bool = False — 梯度 reduce 使用 FP32 精度。
• check_for_nan_in_grad: bool = False — 通讯前检查梯度 NaN。

III.II.II. Megatron FSDP

对应 ZeRO-1/2/3。通过 use_megatron_fsdp: bool = True 启用。基于 PyTorch DTensor 和 DeviceMesh 实现。

关键字段：

• data_parallel_sharding_strategy: str = "no_shard" — 分片策略：
  • "no_shard"：等价于 DDP（ZeRO-0）。
  • "optim"：ZeRO-1，仅分片优化器状态。
  • "optim_grads"：ZeRO-2，分片优化器状态和梯度。
  • "optim_grads_params"：ZeRO-3/FSDP，分片优化器状态、梯度和参数。
• fsdp_double_buffer: bool = False — FSDP 通讯的双缓冲，重叠 compute 与 comm。
• megatron_fsdp_main_params_dtype: Optional[torch.dtype] = torch.float32 — 主参数精度。
• megatron_fsdp_enable_fine_grained_param_gather: bool = False — 细粒度参数 gather。

Hybrid Sharded DP (HSDP)：通过 num_distributed_optimizer_instances: int > 1 和 outer_dp_sharding_strategy: str = "optim" 启用。在多个 DP 域之间创建层次化的分片结构。

III.II.III. Torch FSDP2

通过 --use-torch-fsdp2 启用。使用 PyTorch 原生 FSDP2 实现。

• reshard_after_forward: bool = True（TorchFullyShardedDataParallelConfig）— forward 后是否 reshard 参数。True = ZeRO-3，False = ZeRO-2。

约束：与 PP、EP、MCore DistributedOptimizer、梯度累积融合不兼容。

III.II.IV. 梯度累积与重叠

• gradient_accumulation_fusion: bool = False（ModelParallelConfig）— 使用 APEX 的 fused kernel 将梯度累积与 backward 融合。

III.III. PP

III.III.I. 基本配置

• pipeline_model_parallel_size: int = 1 — 流水线并行度。
• pipeline_dtype: Optional[torch.dtype] = None — \(N_\t{PP}>1\) 时 P2P 通讯的 dtype。通常设为 params_dtype 的值。
• virtual_pipeline_model_parallel_size: Optional[int] = None — virtual PP 度，也即每张卡负责的 stage 数。如果 \(\t{VPP}>1\) 则意味着开启 interleaved 模式。
• microbatch_group_size_per_vp_stage: Optional[int] = None — 只有开启 interleaved 才有用，刻画 virtual stage 的切换周期。

调度方案由 get_forward_backward_func() from megatron.core.pipeline_parallel.schedules 自动选择：

• PP = 1：forward_backward_no_pipelining — 简单的顺序 micro-batch 处理。
• PP > 1，VPP = None：forward_backward_pipelining_without_interleaving — 标准 1F1B 调度。
• PP > 1，VPP ≠ None：forward_backward_pipelining_with_interleaving — Interleaved 1F1B。

III.III.II. 通讯重叠

• overlap_p2p_comm: bool = False — P2P 通讯与计算重叠。
• batch_p2p_comm: bool = True — 使用 batch_isend_irecv 替代逐个调用。
• deallocate_pipeline_outputs: bool = False — 发送激活值到下一个 stage 后立即释放。
• overlap_p2p_comm_warmup_flush: bool = False — warmup/flush 阶段也重叠通讯与计算。

III.III.III. 不均匀 PP 切分

当模型包含非 Transformer 层（如 embedding、LM head）时，各 stage 的层数可能不均匀。

• num_layers_in_first_pipeline_stage: Optional[int] = None — 第一个 stage 的层数。
• num_layers_in_last_pipeline_stage: Optional[int] = None — 最后一个 stage 的层数。
• pipeline_model_parallel_layout — 自定义 PP 切分模式。支持字符串模式（如 'Et*3|(tt|)*29,m|L'）、列表或 PipelineParallelLayerLayout 对象。其中 E = embedding，t = Transformer 层，m = MoE 层，L = LM head。
• account_for_embedding_in_pipeline_split: bool — 将 embedding 视为标准层参与切分。
• account_for_loss_in_pipeline_split: bool — 将 loss 视为标准层参与切分。

III.III.IV. EP 与 PP 的交互

当 EP > 1 且 PP > 1 时，megatron/core/pipeline_parallel/combined_1f1b.py 提供了 combined 1F1B 调度，让 EP 的 a2a 通讯与 PP 的流水线并行进一步重叠：

• overlap_moe_expert_parallel_comm: bool = False — 启用 EP 通讯重叠。即使 PP = 1 也可生效（此时使用 combined_1f1b_schedule_for_no_pipelining）。
• delay_wgrad_compute: bool = False — 延迟 expert 权重梯度计算以改善重叠效果。

III.IV. TP

III.IV.I. 基本配置

• tensor_model_parallel_size: int = 1 — 张量并行度。

约束：num_attention_heads % tensor_model_parallel_size == 0，num_query_groups 必须是 TP 的整数倍（此时一卡持有多个 group）或约数（此时 KV head 会被赋值多份）。

III.IV.II. Sequence Parallelism

对应 Megatron-SP。

• sequence_parallel: bool = False — 启用 SP。要求 TP > 1。

效果：将 TP 的 all-reduce 拆为 reduce-scatter + all-gather，中间插入 token-wise 的操作（LayerNorm、dropout），避免多卡重复计算。

当 Megatron-SP 启用时，PP stage 间传递的激活值 shape 从 \((b, S, H)\) 变为 \((b, S / (N_\t{CP} \times N_\t{TP}), H)\)。

III.IV.III. TP 通讯重叠

• tp_comm_overlap: bool = False — 允许 TP 通讯与线性层执行重叠。
• tp_comm_bulk_wgrad: bool = True — All-gather 与 backward 激活梯度 GeMM 重叠。
• tp_comm_bulk_dgrad: bool = True — Reduce-scatter 与 backward 权重梯度 GeMM 重叠。
• tp_comm_overlap_ag: bool = True — All-gather 与 forward GeMM 流水线化重叠。
• tp_comm_overlap_rs: bool = True — Reduce-scatter 与 forward GeMM 流水线化重叠。

III.IV.IV. Expert TP

• expert_tensor_parallel_size: Optional[int] = None — MoE expert 层的 TP 度数。默认取 tensor_model_parallel_size。可独立于非 expert 层的 TP 度数设置，允许 MoE 层使用不同的切分策略。

III.V. CP

III.V.I. 基本配置

• context_parallel_size: int = 1 — 上下文并行度。
• cp_comm_type: Optional[Union[str, List[str]]] = None — 逐层可配置的通讯类型：
  • "p2p"：Ring Attention。
  • "a2a"：Ulysses 风格。
  • "all_gather"：All-gather 变体。
  • "a2a+p2p"：层级式，也即 head-context 2D parallelism。先 a2a 分组 head，组内 ring-attn。

支持逐层配置，例如 cp_comm_type=["p2p", "a2a"] 表示第一层用 ring，第二层用 Ulysses。

III.V.II. 层级 CP

• hierarchical_context_parallel_sizes: Optional[list[int]] = None — 两层 CP 度数列表，如 [2, 4]。第一层使用 NVLink 内的 a2a，第二层使用跨 IB 的 p2p。对应 Para.txt III.III 中的层级通讯拓扑。

III.V.III. 混合 CP

• hybrid_context_parallel: bool = False — 对变长 packed samples 启用负载均衡。当多个 prompt 被 pack 到同一个序列中时，不同 CP rank 可能持有不同长度的 context，导致计算不均。启用后框架会自动调整切分边界。
• max_seqlen_per_dp_cp_rank: Optional[int] = None — 每个 DP×CP rank 的最大序列长度。

约束：Hybrid CP 与 Megatron FSDP 不兼容。

III.VI. EP

III.VI.I. 基本配置

• expert_model_parallel_size: int = 1 — 专家并行度。
• num_moe_experts: Optional[int] = None（TransformerConfig）— expert 数量 \(E\)。非 None 时启用 MoE。

III.VI.II. MoE Router

• moe_router_topk: int = 2 — 每个 token 路由到的 expert 数 \(K_\t{MoE}\)。
• moe_router_load_balancing_type: Union[str, List[str]] = "aux_loss" — 负载均衡策略：
  • "aux_loss"：辅助损失（Switch Transformer 风格）。
  • "seq_aux_loss"：序列级辅助损失。
  • "global_aux_loss"：全局辅助损失。
  • "sinkhorn"：Sinkhorn 路由。
  • "none"：不使用负载均衡。
• moe_router_score_function: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "softmax" — 路由打分函数。
• moe_router_pre_softmax: bool = False — 在 softmax 之前打分。
• moe_router_num_groups: Optional[int] = None — Group-limited routing 的分组数（DeepSeek V2/V3 风格）。
• moe_router_group_topk: Optional[int] = None — 每组选中的 expert 数。
• moe_router_enable_expert_bias: bool = False — 动态专家偏置（aux-loss-free routing，DeepSeek V3）。
• moe_router_dtype: Optional[Literal["fp32", "fp64"]] = None — 路由计算的精度。
• moe_aux_loss_coeff: Union[float, List[float]] = 0.0 — 辅助损失的缩放系数。
• moe_z_loss_coeff: Optional[float] = None — z-loss 缩放系数（防止 router logits 过大）。

III.VI.III. Token Dispatch

• moe_token_dispatcher_type: Literal["allgather", "alltoall", "flex"] = "allgather" — token 分发方式：
  • "allgather"：使用 all-gather 收集所有 token 到 expert 所在的卡。
  • "alltoall"：使用 all-to-all 将 token 路由到对应 expert 的卡（对应 Para.txt I.V 的 a2a）。
  • "flex"：灵活分发，支持 DeepEP 或 HybridEP 后端。
• moe_flex_dispatcher_backend: Literal["deepep", "hybridep"] = "deepep" — flex 分发器的后端。

III.VI.IV. Expert 配置

• moe_ffn_hidden_size: Optional[int] = None — MoE expert 的 FFN 隐藏维度。默认取 ffn_hidden_size。
• moe_layer_freq: Union[int, List[int]] = 1 — MoE 层的出现频率。1 表示所有层都是 MoE；列表形式可指定哪些层是 MoE（如 [0, 1, 0] 表示第 0、2 层是 Dense，第 1 层是 MoE）。
• moe_grouped_gemm: bool = False — 使用 GroupedGEMM 加速多 expert 计算。
• moe_expert_capacity_factor: Optional[float] = None — expert 容量因子。超过上限的 token 被丢弃（对应 Para.txt I.V 中的 capacity + dropout）。

III.VI.V. Shared Expert

DeepSeek V2/V3 引入的共享 expert 机制，所有 token 都经过共享 expert，不经过 router。

• moe_shared_expert_intermediate_size: Optional[int] = None — 共享 expert 的 FFN 隐藏维度。
• moe_shared_expert_overlap: bool = False — 共享 expert 计算与 dispatcher 通讯重叠。

III.VI.VI. 其它 MoE 配置

• moe_latent_size: Optional[int] = None — MoE 的 latent 投影维度。
• moe_paged_stash: bool = False — paged stash，用于分页存储 routed expert 的激活值。

III.VII. 组合并行的约束与交互

III.VII.I. 乘积约束

\[N_\t{TP} \times N_\t{CP} \times N_\t{PP} \leq \t{world\_size} \]

且 world_size 必须能被 \(N_\t{TP} \times N_\t{CP} \times N_\t{PP}\) 整除。

III.VII.II. 常见约束汇总

条件	约束
SP = True	TP > 1
PP > 1	必须设置 pipeline_dtype
EP > 1	num_moe_experts ≠ None
VPP ≠ None	PP > 1
Hybrid CP	不兼容 Megatron FSDP
Torch FSDP2	不兼容 PP、EP、MCore DistributedOptimizer

III.VII.III. 并行度与通讯拓扑的适配

并行方案	通讯类型	适合的拓扑	Megatron 中的重叠配置
DP (ZeRO-1)	R-S + A-G	跨节点	overlap_grad_reduce, overlap_param_gather
PP	P2P	跨节点/跨集群	overlap_p2p_comm, batch_p2p_comm
TP	All-Reduce	节点内 (NVLink)	tp_comm_overlap, tp_comm_bulk_*
CP (ring)	P2P	跨节点	由 ring-attn 内部处理
CP (Ulysses)	A2A	节点内 (NVLink)	框架内部处理
EP	A2A	节点内 (NVLink)	overlap_moe_expert_parallel_comm

实践中，典型的分配策略是：TP 放在节点内（NVLink 高带宽），PP 放在跨节点（P2P 低延迟），CP/EP 根据具体需求选择拓扑。

IV. 真实部署分析

前三节分别阐述了配置体系的结构、参数的定义位置、以及多维并行的理论约束。本节将这些知识串联起来，完整地描述 一次实验从命令行到训练循环的全部执行过程——包括每一步调用了哪个函数、传入了哪些参数、以及为什么这样传。

IV.I. 单卡场景

如 II.I. 所述，有两种典型的启动方式，将分开描述。

IV.I.I. 最小独立循环

入口文件：examples/run_simple_mcore_train_loop.py

Step 1：环境初始化。入口是自定义的 initialize_distributed() 函数，其读取 torchrun 设置的环境变量 RANK、WORLD_SIZE、LOCAL_RANK，并调用：

torch.distributed.init_process_group(backend="nccl", rank=rank, world_size=world_size)
parallel_state.initialize_model_parallel(tensor_model_parallel_size=2, pipeline_model_parallel_size=1)

这一步建立了 NCCL 通讯组，并在 MCore 内部创建了 TP group（rank 0 和 1 被分到同一 TP 组）。注意此处 没有 调用 initialize_megatron()，因为该函数是 PretrainConfigContainer 体系的一部分。

Step 2：随机种子。model_parallel_cuda_manual_seed(123) 设置随机种子，在 DP 组内相同、TP 组内不同。

Step 3：模型构建。入口是自定义的 model_provider()，其直接实例化 TransformerConfig 和 GPTModel：

transformer_config = TransformerConfig(num_layers=2, hidden_size=12, num_attention_heads=4,use_cpu_initialization=True, pipeline_dtype=torch.float32,
)
gpt_model = GPTModel(config=transformer_config,transformer_layer_spec=get_gpt_layer_local_spec(),vocab_size=100, max_sequence_length=64,
)

此处 绕过了 core_transformer_config_from_args() 和 pretrain_cfg_container_from_args()，直接构造 config 对象。use_cpu_initialization=True 表示先在 CPU 上分配参数，再通过 .to(device) 搬到 GPU。

Step 4：DDP 包装。用 DistributedDataParallelConfig 和 DistributedDataParallel 包装模型：

ddp_config = DistributedDataParallelConfig(grad_reduce_in_fp32=False, overlap_grad_reduce=False,use_distributed_optimizer=False,
)
gpt_model = DistributedDataParallel(config=transformer_config, ddp_config=ddp_config, module=gpt_model)

单卡 TP 场景下，DDP 主要负责提供 finish_grad_sync() 方法供 finalize_model_grads() 调用。

Step 5：数据构建。入口是自定义的 get_train_data_iterator() 使用 MockGPTDataset 构造假数据：

config = GPTDatasetConfig(random_seed=0, sequence_length=64,tokenizer=MegatronTokenizer.from_pretrained(metadata_path={"library": "null-text"}, vocab_size=64,),
)
datasets = BlendedMegatronDatasetBuilder(MockGPTDataset, [1000, None, None], lambda: True, config).build()
train_dataloader = DataLoader(datasets[0], batch_size=8, shuffle=True)

Step 6：训练循环。每一步依次执行：

forward_backward_func = get_forward_backward_func()  # PP=1 时返回 forward_backward_no_pipeliningfor iteration in range(5):optim.zero_grad()losses_reduced = forward_backward_func(forward_step_func=forward_step_func,data_iterator=train_iterator, model=gpt_model,num_microbatches=1, seq_length=64, micro_batch_size=8,decoder_seq_length=64, forward_only=False,)finalize_model_grads([gpt_model])  # TP 组内 all-reduce 非 TP 参数的梯度optim.step()

get_forward_backward_func() from megatron.core.pipeline_parallel.schedules 根据当前的 PP 设置自动选择调度方案。PP = 1 时返回 forward_backward_no_pipelining，即简单的顺序 micro-batch 处理。

Step 7：Checkpoint 保存/加载。使用 MCore 的分布式 checkpoint API：

sharded_state_dict = model.sharded_state_dict(prefix="")
dist_checkpointing.save(sharded_state_dict=sharded_state_dict, checkpoint_dir=ckpt_path)
dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=ckpt_path)

[!TIP]

最小独立循环的意义在于：它展示了 GPTModel、TransformerConfig、get_forward_backward_func()、finalize_model_grads() 这四个 MCore 核心 API 的最小调用方式。当你需要将 Megatron 集成到自己的训练框架中时，这就是起点。

IV.I.II. 完整训练脚本（mock 数据）

目标：走通 PretrainConfigContainer 全流程，但不依赖真实数据集。

入口文件：pretrain_gpt.py

执行流程（对应 pretrain_gpt.py 的 __main__ 块）：

Step 1：参数解析。parse_and_validate_args() from megatron.training.arguments 被调用：

args = parse_and_validate_args(extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None,args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
)

Step 2：模型配置构造。通过 gpt_config_from_args(args) from megatron.training.argument_utils 将 args 映射为 GPTModelConfig（持有 TransformerConfig）。

Step 3：完整配置构造。通过 pretrain_cfg_container_from_args(args, model_cfg) from megatron.training.argument_utils 得到聚合的 PretrainConfigContainer 实例。

Step 4：进入训练。pretrain() from megatron.training.training 被调用，传入上述容器和回调函数：

pretrain(full_config,                          # PretrainConfigContainertrain_valid_test_datasets_provider,   # 数据集工厂函数partial(model_provider, gpt_builder), # 模型构建回调ModelType.encoder_or_decoder,         # 模型类型forward_step,                         # 前向计算回调store=store,                          # fault tolerance storeget_embedding_ranks=get_embedding_ranks,
)

Step 5：pretrain() 内部流程。

1. initialize_megatron() — 调用 torch.distributed.init_process_group() 和 parallel_state.initialize_model_parallel()，根据 tensor_model_parallel_size、pipeline_model_parallel_size 等参数创建所有进程组。
2. setup_model_and_optimizer() — 调用 model_provider() 构建 GPTModel，用 DistributedDataParallel 包装，创建 optimizer（DistributedOptimizer 当 use_distributed_optimizer=True）和 LR scheduler。
3. build_train_valid_test_data_iterators() — 调用 train_valid_test_datasets_provider() 构建数据集和 DataLoader。mock 数据模式下使用 MockGPTDataset。
4. train() — 执行训练循环，每步依次调用 forward_step()、backward、optimizer step。每隔 log_interval 步打印 loss 和 timing 信息。
5. evaluate_and_print_results() — 如果设置了 eval_interval，每隔指定步数运行验证。

Step 6：输出。训练日志输出到 stdout 和 --tensorboard-dir 指定的 TensorBoard 目录；checkpoint 保存到 --save 指定的目录。

IV.I.II. 完整训练脚本（真实数据）

使用真实数据时，需要额外准备 tokenizer 和预处理后的数据集。

数据准备（使用 tools/preprocess_data.py）：

python tools/preprocess_data.py \--input /path/to/raw.jsonl \--output-prefix /path/to/preprocessed \--tokenizer-type HuggingFaceTokenizer \--tokenizer-model /path/to/tokenizer.model \--append-eod

启动命令（与 mock 数据的区别仅在数据参数）：

torchrun --nproc_per_node=8 ... pretrain_gpt.py \... \--data-path /path/to/preprocessed_text_document \--tokenizer-type HuggingFaceTokenizer \--tokenizer-model /path/to/tokenizer.model \--vocab-size 128256 \--split '99,1,0' \--no-create-attention-mask-in-dataloader \--no-mmap-bin-files --num-workers 1 \...

[!TIP]

--split '99,1,0' 表示数据按 99%/1%/0% 分为 train/valid/test 三个子集。--no-create-attention-mask-in-dataloader 跳过 dataloader 中的 attention mask 构建（由模型内部处理），--no-mmap-bin-files 禁用 mmap 二进制缓存文件。

IV.II. 多卡场景

多卡训练的核心区别在于 分布式启动器（torchrun）的配置和 并行参数 的组合。以下分别描述两种典型场景。

IV.II.I. 单节点多卡（LLaMA 3 8B，8×H100）

目标：在单台 8 卡机器上训练 LLaMA 3 8B，启用 FP8 混合精度。

此时使用 torchrun 启动，parse_args() 会通过 os.getenv('RANK') 和 os.getenv('WORLD_SIZE') 读取相关配置。

环境变量 CUDA_DEVICE_MAX_CONNECTIONS=1：限制每个 CUDA stream 的连接数为 1，避免 kernel launch 过多导致的性能退化。这是 Megatron-LM 的标准设置。

执行流程与 IV.I.II 节完全一致。

IV.II.II. 多节点多卡（GPT-3 175B，128×A100）

目标：跨 16 个节点（每节点 8 卡，共 128 卡）训练 GPT-3 175B。

前置条件：所有节点之间需要网络互通（InfiniBand），且 MASTER_ADDR 指向的节点可被所有其它节点访问。

节点 0（主节点） 与 其它节点（node 1–15） 的启动命令的唯一区别，在于 --node_rank 不同。实践中通常使用 SLURM 的 srun 或脚本自动化多节点启动，而非手动在每台机器上执行 torchrun。