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.py | Megatron-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 | None | GQA / MQA 时设;MQA 设 1 |
--kv-channels | H/num_heads | head dim 单独设 |
--ffn-hidden-size | 4*H | FFN 中间维 |
--seq-length | 必填 | 训练样本长度 |
--max-position-embeddings | 必填 | 位置编码上限 |
--position-embedding-type | learned_absolute | rope / alibi / relative / learned |
--rotary-base | 10000 | RoPE base,长上下文常调到 1e6 |
--rotary-percent | 1.0 | RoPE 维度占比 |
--normalization | LayerNorm | LayerNorm / RMSNorm |
--swiglu | off | 用 SwiGLU 激活 |
--untie-embeddings-and-output-weights | off | 解耦 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 | None | interleaved PP(见第 8 章) |
--context-parallel-size | 1 | CP 度(见第 9 章) |
--expert-model-parallel-size | 1 | EP 度(MoE) |
--sequence-parallel | off | SP(不是 CP!)见第 9 章 |
--num-layers-per-virtual-pipeline-stage | — | interleaved 时每个 vstage 多少层 |
--standalone-embedding-stage | off | embedding 单独占一个 stage |
--distributed-backend | nccl | nccl / gloo |
--use-distributed-optimizer | off | ZeRO-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-fp32 | fp32 累积梯度 |
5.8 attention 与 transformer 后端
| flag | 默认 | 说明 |
--attention-backend | auto | auto / flash / fused / unfused / local |
--attention-softmax-in-fp32 | off | 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_engine | local / 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-type | aux_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-size | shared expert(DeepSeek 风格) |
5.13 字段查不到怎么办
grep -n '"--<flag>"' megatron/training/arguments.py;
- 找不到 → 试
megatron/core/transformer/transformer_config.py 里的 dataclass 字段;
- 还找不到 → 多半被改名了,
git log --all -S "<flag>" megatron/training/arguments.py 看历史;
- 都没有 → 拼写错;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.py 的 parse_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_args | attention 研究开关 |
| 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 就快。