快速上手：训一个 tiny GPT

nanotron 的入门用 examples/config_tiny_llama.yaml，一个 ~50M 参数的 mini Llama。本章把它跑通。

3.1看 tiny config

# examples/config_tiny_llama.yaml（简化）
general:
  project: debug
  run: tiny_llama
  seed: 42

checkpoints:
  checkpoint_interval: 1000
  checkpoints_path: ./checkpoints

model:
  model_config:
    hidden_size: 16
    intermediate_size: 64
    num_attention_heads: 4
    num_hidden_layers: 2
    num_key_value_heads: 4
    vocab_size: 256
    max_position_embeddings: 2048
    is_llama_config: true

parallelism:
  dp: 2
  pp: 1
  tp: 2
  pp_engine: 1f1b
  tp_mode: ALL_REDUCE

tokens:
  batch_accumulation_per_replica: 1
  micro_batch_size: 2
  sequence_length: 128
  train_steps: 50

optimizer:
  optimizer_factory:
    adam_beta1: 0.9
    adam_beta2: 0.95
    adam_eps: 1.0e-08
    learning_rate_scheduler:
      learning_rate: 1.0e-3
      lr_warmup_steps: 10
      lr_decay_starting_step: 50

data_stages:
  - name: stable_phase
    start_training_step: 1
    data:
      dataset:
        dataset_overwrite_cache: false
        dataset_processing_num_proc_per_process: 1
        hf_dataset_or_datasets: stas/openwebtext-10k
        hf_dataset_splits: train
        text_column_name: text
      num_loading_workers: 1
      seed: 42

3.2启动

# 4 卡（2 DP × 2 TP × 1 PP）
torchrun --nproc_per_node=4 run_train.py \
    --config-file examples/config_tiny_llama.yaml

跑 50 step，loss 应该从 5.5 降到 ~3.5。整个流程几分钟。

3.3启动脚本拆解

# run_train.py（约 50 行）
import argparse
from nanotron.config import Config, get_config_from_file
from nanotron.trainer import DistributedTrainer

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config-file", required=True)
    args = parser.parse_args()

    config = get_config_from_file(args.config_file)
    trainer = DistributedTrainer(config)
    trainer.train(dataloader_builder)

if __name__ == "__main__":
    main()

所有复杂度被封到 DistributedTrainer 里。

3.4多卡 / 多机

# 单机 8 卡
torchrun --nproc_per_node=8 run_train.py --config-file ...

# 多机
NNODES=2 NODE_RANK=$RANK MASTER_ADDR=10.0.0.5 MASTER_PORT=12345 \
torchrun --nproc_per_node=8 --nnodes=$NNODES \
    --node_rank=$NODE_RANK \
    --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \
    run_train.py --config-file ...

启动方式跟普通 PyTorch DDP 一致。

3.5训练日志

[rank0] step: 10  loss: 5.21  tokens/s/gpu: 12345
[rank0] step: 20  loss: 4.82  ...

# wandb（如开）
wandb_logger:
  project: my-nanotron
  name: tiny-run

3.6常见踩坑

3.7这章你需要带走的

• 启动两步：选 yaml + torchrun run_train.py；
• tiny config 一两分钟看到效果，验证环境；
• 分布式三维：dp × tp × pp = world_size；
• 多机沿用 PyTorch torchrun 风格。