训练循环：从 train.py 拆解每一步

TorchTitan 的精髓全在 torchtitan/train.py（约 500 行）。本章逐段拆解。

8.1train.py 主流程

def main():
    # 1) 读 config
    config = JobConfig.from_args()

    # 2) 初始化分布式
    init_distributed(config)

    # 3) 构造 mesh
    parallel_dims = ParallelDims(...)
    mesh = parallel_dims.build_mesh(device_type="cuda")

    # 4) 模型 + tokenizer
    model = build_model(config.model, parallel_dims)
    tokenizer = build_tokenizer(config.model.tokenizer_path)

    # 5) 应用并行
    parallelize_llama(model, mesh, parallel_dims, config)

    # 6) Optimizer / Scheduler
    optimizer = build_optimizer(model, config.optimizer)
    scheduler = build_lr_scheduler(optimizer, config.lr_scheduler)

    # 7) Data loader
    data_loader = build_data_loader(config.training, tokenizer, parallel_dims)

    # 8) Checkpoint
    ckpt_mgr = CheckpointManager(model, optimizer, scheduler, ...)
    if config.checkpoint.enable:
        ckpt_mgr.load()

    # 9) 训练循环
    for step in range(start_step, config.training.steps):
        train_step(model, optimizer, scheduler, data_loader)
        if step % config.checkpoint.interval == 0:
            ckpt_mgr.save(step)

8.2训练一步：train_step()

def train_step(model, optimizer, scheduler, data_loader):
    batch = next(data_loader)
    input_ids, labels = batch["input_ids"], batch["labels"]

    # CP context（如果开了）
    cp_ctx = context_parallel(mesh["cp"], buffers=[input_ids, labels]) \
             if cp_degree > 1 else nullcontext()

    # PP path
    if pp_degree > 1:
        # 走 pipeline schedule
        loss = pp_schedule.step(input_ids, target=labels)
    else:
        with cp_ctx:
            logits = model(input_ids)
            loss = cross_entropy_loss(logits, labels)
            loss.backward()

    # 梯度 clip + step
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()

    return loss

8.3关键模块对应文件

8.3parallelize_llama 内部

def parallelize_llama(model, mesh, parallel_dims, config):
    # 1) TP / SP（如开启）
    if parallel_dims.tp_enabled:
        apply_tp(model, mesh["tp"],
                 loss_parallel=config.training.enable_loss_parallel,
                 enable_async_tp=config.experimental.enable_async_tensor_parallel)

    # 2) Float8（如开启）
    if config.float8.enable_float8_linear:
        convert_to_float8_linear(model, ...)

    # 3) Activation checkpoint
    if config.activation_checkpoint.mode != "none":
        apply_ac(model, config.activation_checkpoint)

    # 4) Compile（如开启）
    if config.compile.enable:
        apply_compile(model)

    # 5) FSDP2
    if parallel_dims.dp_shard_enabled:
        apply_fsdp(model, mesh["dp_shard"], ...)

    # 6) PP（如开启，会改变 model 结构）
    if parallel_dims.pp_enabled:
        model = apply_pp(model, mesh["pp"], config)

    return model

顺序很重要：先 TP（改 weight 切法），再 float8，再 ac，再 compile，最后 FSDP2 包外层。PP 因为切层，放最后。

8.4compute / comm overlap

TorchTitan 不显式管 overlap，全靠 PyTorch 自身：

• FSDP2 自动 prefetch 下一层 weight；
• TP allreduce 跟下一个 op overlap；
• async TP（experimental.enable_async_tensor_parallel）把通信切成微块跟计算 overlap。

8.5train.py 与 Trainer 类源码定位

TorchTitan v2026 的 train 入口已经从单一 train.py 拆为 "torchtitan/train.py + torchtitan/trainer.py + torchtitan/config/manager.py" 三件套：

8.6启动脚本与 dry-run 模式

run_train.sh（line 28-47）使用 torchrun ... -m torchtitan.train --module ... --config ...。两个调试模式：

• COMM_MODE=fake_backend：完全跳过通信，单进程跑通流程（验证逻辑）；
• COMM_MODE=local_tensor：单 GPU 模拟多卡（验证 sharding 语义）。

这两个模式跟 v2026 之前不一样——以前要改源码自己 mock，现在原生支持。

8.7这章你需要带走的

• train.py ≈ 500 行，9 步主流程；
• 核心是 build_* 系列 + parallelize_llama + train_step；
• parallelize_llama 顺序：TP → Float8 → AC → Compile → FSDP → PP；
• compute/comm overlap 全由 PyTorch 自身做。