训练循环与 checkpoint

trainer.py（约 800 行）是 nanotron 的"中枢神经"。本章把训练循环和 ckpt 流程过一遍。

10.1DistributedTrainer 核心

# src/nanotron/trainer.py
class DistributedTrainer:
    def __init__(self, config):
        self.config = config
        self.parallel_context = ParallelContext(
            data_parallel_size=config.parallelism.dp,
            pipeline_parallel_size=config.parallelism.pp,
            tensor_parallel_size=config.parallelism.tp,
        )

        self.model = self._init_model()
        self.optimizer, self.lr_scheduler = self._init_optimizer_and_scheduler()
        self.loaded_checkpoint = self._load_checkpoint_if_exists()

    def train(self, dataloader_builder):
        for stage in self.config.data_stages:
            if self.current_step < stage.start_training_step:
                continue
            dataloader = dataloader_builder(stage, self.parallel_context)
            self._train_stage(dataloader)

    def _train_stage(self, dataloader):
        for batch in dataloader:
            if self.current_step >= self.config.tokens.train_steps:
                return
            self._training_step(batch)
            self.current_step += 1
            self._log_metrics()
            if self.current_step % ckpt_interval == 0:
                self._save_checkpoint()

10.2_training_step

def _training_step(self, batch):
    # 1) Optionally split batch into micro_batches
    micro_batches = split_batch(batch, self.config.tokens.batch_accumulation_per_replica)

    # 2) Forward + backward
    if self.pp_size > 1:
        loss = self.pp_engine.step(self.model, micro_batches, self.optimizer)
    else:
        loss = 0
        for mb in micro_batches:
            with self.tp_context():
                output = self.model(mb)
                local_loss = self.loss_fn(output, mb.target)
                local_loss.backward()
                loss += local_loss.item()

    # 3) DP sync gradients
    if self.dp_size > 1:
        self.sync_gradients()

    # 4) Clip + step
    torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=clip_grad)
    self.optimizer.step()
    self.lr_scheduler.step()
    self.optimizer.zero_grad()

10.3checkpoint 设计

nanotron 的 ckpt 格式按 (pp, tp, dp_shard) 三维存：

./checkpoints/iteration_1000/
├── config.yaml
├── pp-0-of-2/
│   ├── tp-0-of-4/
│   │   ├── optimizer/
│   │   │   └── dp-0-of-2/...     # ZeRO-1 切到 dp_rank
│   │   ├── model_weights.safetensors
│   │   └── ...
│   └── tp-1-of-4/...
├── pp-1-of-2/...
└── metadata.json

save 是 distributed save（每 rank 写自己那份），快、不阻塞。

10.4save / load

# save
trainer._save_checkpoint()
# 内部：
#   1) 每 rank 把自己那份 weight 写到对应子目录
#   2) DP 范围内只 rank 0 写 weight（其他 rank 是副本）
#   3) optimizer state 每 dp_rank 各自写

# load
trainer._load_checkpoint_if_exists()
# 自动按当前 ParallelContext 决定从哪几个文件读

10.5resume 包含

10.6跨 mesh resume

不像 TorchTitan dcp 支持任意 reshard，nanotron 的 ckpt 跟 (pp, tp) 结构强绑定。要换 tp_size 必须先用 tools/converter 脚本重新切。

10.7训练循环中的 hooks

nanotron 没有 NeMo / MMEngine 那种 Hook 系统，所有"切点"都直接写在 trainer.py 里（更简单但灵活度低）。常见自定义点：

• _log_metrics：每 N step 打日志（含 tensorboard / wandb）；
• _save_checkpoint：每 N step 存 ckpt；
• before_optim_step / after_optim_step：梯度统计、scaling 等。

10.8这章你需要带走的

• DistributedTrainer ~800 行，是 nanotron 的中枢；
• 训练流程：split micro_batch → forward/backward → DP sync → clip + step；
• checkpoint 按 (pp, tp, dp) 三维存，distributed save；
• resume 默认完整（含 data sampler 位置）；
• 跨 mesh resume 需要 tools 重新切 ckpt。