配置：dataclass 驱动 + YAML

nanotron 的配置系统是"dataclass + YAML"：每个 section 对应一个 dataclass，YAML 自动校验。本章拆它的 8 个核心 dataclass。

7.1八大 dataclass

每个都在 src/nanotron/config/config.py。

7.2核心字段速读

parallelism

parallelism:
  dp: 2
  pp: 2
  tp: 4
  pp_engine: 1f1b           # 1f1b / afab / interleaved
  tp_mode: REDUCE_SCATTER    # ALL_REDUCE / REDUCE_SCATTER
  tp_linear_async_communication: true
  expert_parallel_size: 1   # MoE

model

model:
  model_config:
    hidden_size: 4096
    intermediate_size: 14336
    num_attention_heads: 32
    num_hidden_layers: 32
    num_key_value_heads: 8
    vocab_size: 128256
    max_position_embeddings: 8192
    is_llama_config: true
    rope_theta: 500000.0
  init_method:
    std: 0.025
  dtype: bfloat16
  make_vocab_size_divisible_by: 1

tokens

tokens:
  batch_accumulation_per_replica: 4
  micro_batch_size: 1
  sequence_length: 8192
  train_steps: 50000
  val_check_interval: 1000
  limit_val_batches: 5
  limit_test_batches: 0

optimizer

optimizer:
  zero_stage: 1
  weight_decay: 0.1
  clip_grad: 1.0
  accumulate_grad_in_fp32: true
  learning_rate_scheduler:
    learning_rate: 3.0e-4
    lr_warmup_steps: 2000
    lr_warmup_style: linear
    lr_decay_steps: 48000
    lr_decay_style: cosine
    min_decay_lr: 3.0e-5
  optimizer_factory:
    name: adam_w
    adam_beta1: 0.9
    adam_beta2: 0.95
    adam_eps: 1.0e-8
    torch_adam_is_fused: true

7.3data_stages：分阶段训练

nanotron 支持多阶段数据混合切换，常用于"先 web 再 code 再 instruction"风格的预训练：

data_stages:
  - name: stage_web
    start_training_step: 1
    data:
      dataset:
        hf_dataset_or_datasets: c4
        text_column_name: text
      num_loading_workers: 4

  - name: stage_code_mix
    start_training_step: 30000      # ★ 第 30k 步切换
    data:
      dataset:
        hf_dataset_or_datasets:
          - c4
          - bigcode/the-stack
        text_column_name: text
        weights: [0.7, 0.3]

  - name: stage_instruction
    start_training_step: 45000
    data:
      dataset: ...

这种"data schedule" 是 nanotron 比同类框架更便捷的一处。

7.4读 config 的入口

from nanotron.config import get_config_from_file

config = get_config_from_file("./config.yaml")
print(config.parallelism.dp)        # 类型安全：dataclass 字段
print(config.tokens.train_steps)

内部用 dacite 把 YAML 转 dataclass，类型不对会立即报错。

7.5跟其他配置系统对比

7.6这章你需要带走的

• nanotron 配置 = YAML + dataclass 校验；
• 8 个核心 section：general / checkpoints / parallelism / model / tokens / optimizer / data_stages / logging；
• multi-stage 数据是 nanotron 独家亮点；
• 看不懂某个字段就翻 src/nanotron/config/config.py 的 dataclass。