Chapter 05

配置系统:arguments.py 全字段拆解

📌 本章字段名与默认值基于 megatron/training/arguments.py · commit 432d76b2b8bf (v0.15.0)。 这是 Megatron-LM 1000+ 行的"argparse 总入口",所有 CLI flag 都在这里注册。

Megatron 不像 LLaMA-Factory 用 YAML,也不像 DeepSpeed 用 JSON,它用纯命令行 argparse。一份 175B 训练脚本里轻松出现 60+ 个 flag。本章给你一份"凭印象查得到"的索引。

5.1 字段总入口在哪

文件角色
megatron/training/arguments.py 训练 / 微调主入口,1000+ 行
megatron/post_training/arguments.py 后训练(RL / SFT)专用
megatron/elastification/arguments.py 弹性扩缩容专用
megatron/core/transformer/transformer_config.pyMegatron-Core 的 TransformerConfig dataclass
megatron/core/model_parallel_config.py ModelParallelConfig dataclass

训练脚本里所有 --xxx-yyy flag 都在第一个文件里通过 add_argument() 注册。找不到字段就 grep 它

grep -n '"--<flag>"' megatron/training/arguments.py

5.2 字段按"功能段"分类

Megatron 的 1000+ flag 被组织成 ~15 个 add_xxx_args() 函数,每个函数管一类:

注册函数管什么
_add_network_size_args 模型结构尺寸
_add_regularization_args weight decay / dropout
_add_training_args batch / iter / dataloader
_add_initialization_args 权重初始化
_add_learning_rate_args 学习率与调度
_add_checkpointing_args 保存 / 恢复
_add_mixed_precision_args fp16 / bf16 / fp8 / fp4
_add_distributed_args TP / PP / CP / EP 与通信
_add_data_args 数据集与 tokenizer
_add_validation_args 评估
_add_logging_args tensorboard / wandb
_add_inference_args 推理
_add_transformer_engine_args TransformerEngine 集成
_add_retro_args Retro 检索增强(特殊)
_add_experimental_args 实验性功能

5.3 网络尺寸(最常改的一组)

flag默认说明
--num-layers 必填 Transformer 层数 $L$
--hidden-size 必填 hidden dim $H$
--num-attention-heads 必填 attention head 数
--num-query-groups NoneGQA / MQA 时设;MQA 设 1
--kv-channels H/num_headshead dim 单独设
--ffn-hidden-size 4*HFFN 中间维
--seq-length 必填 训练样本长度
--max-position-embeddings 必填 位置编码上限
--position-embedding-type learned_absoluterope / alibi / relative / learned
--rotary-base 10000 RoPE base,长上下文常调到 1e6
--rotary-percent 1.0 RoPE 维度占比
--normalization LayerNormLayerNorm / RMSNorm
--swiglu off 用 SwiGLU 激活
--untie-embeddings-and-output-weightsoff 解耦 input/output embedding
--vocab-size 由 tokenizer 自动推

典型 Llama-2 7B 风格组合:

--num-layers 32 --hidden-size 4096 --num-attention-heads 32 \
--num-query-groups 8 --kv-channels 128 --ffn-hidden-size 11008 \
--seq-length 4096 --max-position-embeddings 4096 \
--position-embedding-type rope --normalization RMSNorm --swiglu \
--untie-embeddings-and-output-weights

5.4 并行配置(核心)

flag默认说明
--tensor-model-parallel-size 1 TP 度
--pipeline-model-parallel-size 1 PP 度
--virtual-pipeline-model-parallel-size Noneinterleaved PP(见第 8 章)
--context-parallel-size 1 CP 度(见第 9 章)
--expert-model-parallel-size 1 EP 度(MoE)
--sequence-parallel offSP(不是 CP!)见第 9 章
--num-layers-per-virtual-pipeline-stage interleaved 时每个 vstage 多少层
--standalone-embedding-stage offembedding 单独占一个 stage
--distributed-backend ncclnccl / gloo
--use-distributed-optimizer offZeRO-1 等价物
--overlap-grad-reduce off反向通信与计算重叠
--overlap-param-gather off前向参数 gather 重叠

关系仍然是:

$W = \mathrm{TP} \times \mathrm{PP} \times \mathrm{DP} \times \mathrm{CP}$。

5.5 batch 与梯度累积

flag默认说明
--micro-batch-size 必填 单卡 micro batch
--global-batch-size 必填 全局 batch
--rampup-batch-size None渐进式 ramp up batch(如 "16 16 5000000")
--step-batch-size-schedule 按训过的 token 数切换 batch(175B 脚本用过)
--train-iters 训练步数;和 --train-samples 二选一
--train-samples 训练样本数

必须满足:

$B_\text{global} = b \cdot m \cdot d$,其中 $b$ = micro batch,$m$ = micro batch 数(自动算),$d$ = 数据并行度。

5.6 优化器与学习率

flag默认说明
--optimizer adam adam / sgd
--adam-beta1 0.9
--adam-beta2 0.999 大模型推荐 0.95
--adam-eps 1e-8
--lr 必填 峰值学习率
--min-lr 0 decay 后的下限
--lr-decay-style linear linear / cosine / WSD / constant / inverse-square-root
--lr-warmup-iters 0 warmup 步数
--lr-warmup-fraction warmup 占比(与 iters 二选一)
--lr-decay-iters decay 持续步数
--weight-decay 0.01 L2 系数
--clip-grad 1.0 梯度裁剪
--init-method-std 0.02 权重初始化标准差

5.7 混合精度

flag说明
--fp16 fp16 训练(默认)
--bf16 bf16;A100/H100 推荐
--fp8 H100 fp8 训练,要 TransformerEngine
--fp8-format e4m3 / hybrid / e5m2
--fp4-recipe fp4 训练(Blackwell+)
--loss-scale 静态 loss scale,0 = 动态
--initial-loss-scale 动态 scale 初值
--accumulate-allreduce-grads-in-fp32fp32 累积梯度

5.8 attention 与 transformer 后端

flag默认说明
--attention-backend autoauto / flash / fused / unfused / local
--attention-softmax-in-fp32off softmax 用 fp32 计算
--attention-dropout 0.1
--hidden-dropout 0.1
--no-position-embedding off 仅 RoPE 时关掉绝对位置
--use-flash-attn off 显式启用 flash-attn(也可走 backend)
--transformer-impl transformer_enginelocal / transformer_engine

5.9 重算与显存

flag说明
--recompute-activations 把 attention 内激活重算(默认 selective)
--recompute-granularity selective / full
--recompute-method uniform / block
--recompute-num-layers full 时每多少层 ckpt 一次
--distribute-saved-activations 把激活按 TP 切(节省显存)
--check-for-nan-in-loss-and-grad每步检查 NaN(debug 时开)

5.10 数据相关

flag说明
--data-path indexed dataset 路径前缀(不含 .bin/.idx)
--split "949,50,1" 训/验/测百分比
--vocab-file GPT-2 词表 json
--merge-file BPE merges
--tokenizer-type GPT2BPETokenizer / Llama3Tokenizer / HuggingFaceTokenizer
--tokenizer-model SentencePiece 模型路径
--data-cache-path 数据 cache 目录
--num-workers DataLoader worker 数
--seed
--reset-position-ids 跨 sample 重置 position id
--reset-attention-mask 跨 sample 重置 attention mask
--eod-mask-loss 不在 EOD token 计 loss

5.11 checkpoint 与日志

flag说明
--save 保存目录
--save-interval 多少步保存一次
--load 恢复 / 继续训练目录
--no-load-optim 恢复时不加载 optimizer 状态
--no-load-rng 恢复时不加载 RNG
--ckpt-format torch / torch_dist / zarr(推荐 torch_dist
--async-save 异步保存(H100 节省时间)
--log-interval 多少步打 log
--eval-interval 多少步评测
--eval-iters 评测步数
--tensorboard-dir TB 日志目录
--wandb-project WandB 项目名

5.12 MoE 字段

flag说明
--num-experts 专家总数
--moe-router-topk 每 token 路由 top-k
--moe-router-load-balancing-typeaux_loss / sinkhorn / none
--moe-aux-loss-coeff 负载均衡 loss 系数
--moe-grouped-gemm 分组 GEMM 优化
--moe-token-dispatcher-type alltoall / allgather
--moe-expert-capacity-factor capacity factor
--moe-shared-expert-intermediate-sizeshared expert(DeepSeek 风格)

5.13 字段查不到怎么办

  1. grep -n '"--<flag>"' megatron/training/arguments.py
  2. 找不到 → 试 megatron/core/transformer/transformer_config.py 里的 dataclass 字段;
  3. 还找不到 → 多半被改名了,git log --all -S "<flag>" megatron/training/arguments.py 看历史;
  4. 都没有 → 拼写错;Megatron 对未知 flag 会抛 error: unrecognized arguments不会 silently 忽略(比 LLaMA-Factory 友好)。

5.14 一份"高频开关"速查

你想做什么关键 flag
训 Llama 风格 --swiglu --normalization RMSNorm --position-embedding-type rope --untie-embeddings-and-output-weights
训 GQA --num-query-groups 8
开 SP --sequence-parallel(需 TP>1)
开长上下文 --seq-length 32768 --max-position-embeddings 32768 --context-parallel-size 4 --rotary-base 1000000
开 ZeRO-1 风格 --use-distributed-optimizer
开通信 overlap --overlap-grad-reduce --overlap-param-gather
开 fp8 --fp8 --fp8-format hybrid
开 selective recompute--recompute-activations --recompute-granularity selective
开 full recompute --recompute-activations --recompute-granularity full --recompute-method uniform --recompute-num-layers 1
分布式 checkpoint --ckpt-format torch_dist --async-save

5.15 arguments.py 全部 33 个 _add_*_args 函数

megatron/training/arguments.pyparse_args() 顺序调用 33 个分组注册器。完整地图(行号取自 commit ab1a0d8):

行号函数注册什么 CLI
1788_add_transformer_engine_args TE 内核编译开关
1805_add_inference_args 推理期参数(不影响训练)
2004_add_network_size_args hidden / num_layers / kv_channels / seq_len
2171_add_straggler_detector_args 慢节点检测
2179_add_workload_inspector_server_args性能监控 server
2185_add_inprocess_restart_args 同进程重启(不退出)
2234_add_one_logger_args NVIDIA OneLogger
2263_add_ft_package_args 容错框架
2279_add_logging_args tensorboard / wandb / log_interval
2288_add_regularization_args dropout / weight_decay
2355_add_rl_args ★ RL(reward model / PPO)
2503_add_training_args ★ batch / lr / micro_batch / grad_accum
2604_add_rerun_machine_args resume 状态机
2613_add_initialization_args init_method / seed
2625_add_learning_rate_args lr_schedule / min_lr / decay
2650_add_checkpointing_args ★ ckpt_format / async_save / save_interval
2680_add_mixed_precision_args fp16 / bf16 / fp8
2711_add_distributed_args ★ TP / PP / CP / EP / VP
2842_add_validation_args eval_iters / eval_interval
2851_add_tokenizer_args tokenizer_type / vocab_file
2860_add_data_args data_path / blend / split
2983_add_autoresume_args 自动从最新 ckpt 续训
2995_add_biencoder_args 双塔检索
3052_add_vision_args 多模态 vision encoder
3119_add_moe_args ★ MoE 全部 12 个 flag
3150_add_mla_args ★ Multi-head Latent Attention
3180_add_experimental_attention_variant_argsattention 研究开关
3194_add_heterogeneous_args ★ 异构硬件训练
3244_add_experimental_args Mamba / CP / packed_seq
3324_add_msc_args 多 stream 计算
3335_add_kitchen_quantization_arguments★ 训练期量化
3365_add_sft_args ★ SFT instruction-tuning
3373_add_fault_injector_args 故障注入测试

带 ★ 的是训练时最常调的 7 组,记住它们大致行号,grep 就快。