Chapter 07

配置:toml + ArgumentParser

📌 commit af33f76读 train_configs/llama3_8b.toml 各段

TorchTitan 配置 = toml 默认 + argparse 覆盖。本章把所有 toml section 列清楚。

7.1toml 全部 section

section典型字段
[job]dump_folder / description
[profiling]enable_profiling / profile_freq
[metrics]enable_tensorboard / log_freq
[model]name / flavor / norm_type / tokenizer_path
[optimizer]name / lr / fused
[lr_scheduler]warmup_steps / decay_type
[training]batch_size / seq_len / steps / dataset / dp_*/tp_*
[parallelism]新版集中放并行字段(替代部分 training)
[experimental]pp / cp / async_tp 等实验特性
[checkpoint]enable / folder / interval / async_mode
[float8]enable / recipe_name / filter_fqns
[activation_checkpoint]mode / selective_ac_option
[compile]enable / model_backend

7.2关键字段速查

model

[model]
name = "llama3"                  # 模型族
flavor = "8B"                    # 8B / 70B / 405B 任一 flavor
norm_type = "rmsnorm"
tokenizer_path = "./tokenizer/tokenizer.model"

training

[training]
batch_size = 1                   # 单 GPU 微批
seq_len = 8192
steps = 1000
dataset = "c4"
data_parallel_replicate_degree = 1   # HSDP 外层
data_parallel_shard_degree = -1       # -1 = 自动算
tensor_parallel_degree = 1
context_parallel_degree = 1
enable_loss_parallel = false
warmup_steps = 200
max_norm = 1.0
compile = false

experimental(含 PP / async TP)

[experimental]
pipeline_parallel_degree = 1
pipeline_parallel_schedule = "1F1B"
pipeline_parallel_microbatches = 8
enable_async_tensor_parallel = false

activation_checkpoint

[activation_checkpoint]
mode = "selective"               # none / selective / full
selective_ac_option = "op"       # op / 2  数字 = 每 N 层选一层

checkpoint

[checkpoint]
enable_checkpoint = true
folder = "checkpoint"
interval = 500
async_mode = "async"             # disabled / async / async_with_pinned_mem
keep_latest_k = 3
model_weights_only = false       # true = 只存 model(推理用)

float8

[float8]
enable_float8_linear = false
recipe_name = ""                  # rowwise / tensorwise / mxfp8

7.3命令行 override

每个字段都能命令行覆盖:

torchrun ... torchtitan/train.py \
    --job.config_file ./train_configs/llama3_8b.toml \
    --training.steps 100 \
    --training.batch_size 2 \
    --training.tensor_parallel_degree 4 \
    --experimental.pipeline_parallel_degree 2 \
    --float8.enable_float8_linear true \
    --compile.enable true

7.4添加新字段

torchtitan/config_manager.py,所有字段都用 argparse 注册。要加字段就在对应 group 内 add_argument。

# 假设要加 --training.my_new_field
parser.add_argument(
    "--training.my_new_field",
    type=int,
    default=42,
    help="...",
)

之后 toml 里加 my_new_field = 100 也会被读到。

7.5跟 YAML 对比

toml(TorchTitan)YAML(Axolotl / NeMo 1.x)
嵌套section 风格层级缩进
类型显式(123 vs "123")自动推断(易错)
注释支持(#)支持(#)
命令行 overrideargparse 完美兼容Hydra 也行但更重
Python 标准库tomllib (3.11+)PyYAML 第三方

7.6常见踩坑

现象处理
命令行 override 不生效段名 / 字段名拼写错;用 --help
toml 类型错(int 写成 str)tomllib 区分 1"1",要看清
dp_shard_degree=-1 算错显式指定
checkpoint folder 写满keep_latest_k 限制

7.7这章你需要带走的