Chapter 05

Hydra 配置系统:YAML / override / interpolation

📌 commit 2ea3e0f1.x 时代的灵魂,2.x 仍可用

NeMo 1.x 全部用 Hydra + YAML 驱动训练,2.x 改成 Python recipe,但 YAML 路径仍然完整保留。本章讲三件事:YAML 结构、命令行 override、配置组合。

5.1Hydra 是什么

Hydra 是 Meta 开源的配置框架,三个核心特性:

5.2NeMo 的 YAML 结构

典型 GPT 预训练 YAML(NeMo 1.x 风格,在 examples/nlp/language_modeling/conf/):

# megatron_gpt_config.yaml
name: megatron_gpt
restore_from_path: null

trainer:
  devices: 8
  num_nodes: 1
  accelerator: gpu
  precision: bf16
  max_steps: 100000
  log_every_n_steps: 10
  val_check_interval: 2000

model:
  micro_batch_size: 4
  global_batch_size: 128
  tensor_model_parallel_size: 2
  pipeline_model_parallel_size: 1
  encoder_seq_length: 4096
  num_layers: 24
  hidden_size: 2048
  ffn_hidden_size: 8192
  num_attention_heads: 16
  init_method_std: 0.014
  use_flash_attention: True
  normalization: rmsnorm

  data:
    data_prefix:
      - 0.7
      - /data/pile/pile_text_document
      - 0.3
      - /data/code/code_text_document
    splits_string: 998,1,1
    seq_length: 4096

  optim:
    name: distributed_fused_adam
    lr: 1e-4
    weight_decay: 0.1
    betas: [0.9, 0.95]

5.3命令行 override

# 简单 override
python megatron_gpt_pretraining.py \
    --config-path conf --config-name megatron_gpt_config \
    model.hidden_size=4096 \
    trainer.max_steps=200000

# 添加新字段
python ... +model.new_field=hello

# 删除字段
python ... ~model.optim.weight_decay

# list override(注意双 [[ ]])
python ... model.data.data_prefix='[0.5, /a, 0.5, /b]'

5.4config group:组合不同配方

NeMo 把"模型架构 / 优化器 / 数据"拆成不同子目录,用 group 引用:

conf/
├── megatron_gpt_config.yaml         # 主配置
├── model/
│   ├── llama3_8b.yaml
│   ├── llama3_70b.yaml
│   └── mixtral_8x7b.yaml
└── optim/
    ├── adam.yaml
    └── adamw.yaml
# 命令行切换 model group
python train.py +model=llama3_8b +optim=adamw

# 主 YAML 里写
defaults:
  - model: llama3_8b
  - optim: adamw
  - _self_

5.5interpolation:引用其他字段

model:
  hidden_size: 4096
  ffn_hidden_size: ${eval:${model.hidden_size} * 4}    # = 16384
  num_attention_heads: 32
  kv_channels: ${eval:${model.hidden_size} / ${model.num_attention_heads}}  # = 128

trainer:
  precision: bf16
  log_dir: /logs/${name}/${oc.env:USER}
  resume_from: ${trainer.log_dir}/last.ckpt

${...} 是字段引用,${eval:...} 是数学求值,${oc.env:VAR} 是环境变量。

5.6NeMo 2.x 怎么对接 Hydra

2.x 默认用 Python recipe。如果项目坚持用 YAML,仍可调用:

from nemo.collections.llm.recipes import gpt3_5b
from hydra import compose, initialize

with initialize(config_path="conf"):
    cfg = compose(config_name="megatron_gpt_config",
                  overrides=["model.hidden_size=2048"])

# 把 OmegaConf 转成 recipe 等价
recipe = gpt3_5b.pretrain_recipe_from_cfg(cfg)

大多数 NeMo 2.x 项目不用 Hydra,直接 Python 函数 + dataclass。

5.7这章你需要带走的