
TorchTitan 配置编程指南用 Python 函数描述完整训练运行【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan导读TorchTitan 的配置体系围绕一个核心思想构建一次训练运行由完整配置full configuration描述——一个返回完整Trainer.Config的 Python 函数其中包含模型、各维并行度以及全部优化选项。本指南将带你理解这套配置即编程configuration programming的设计哲学掌握用--module/--config选择配置、编写自定义配置函数、管理配置存放位置以及理解为何命令行 flags 被冻结并最终淘汰。读完后你将能够为任意集群拓扑与模型组合写出自己的可复现训练配置并懂得如何在配置树中引入新特性。什么是 full configuration配置即函数在 TorchTitan 中一次运行不是靠一长串命令行参数拼出来的而是由一个返回完整Trainer.Config的 Python 函数完整描述。Trainer.Config是 torchtitan/trainer.py 中定义的总配置容器它聚合了模型规格model_spec、并行度parallelism、训练超参training、优化器、数据加载器、检查点、编译、通信、调试等全部子配置。因为配置用 Python 编写所以搭建一个新的运行本质上就是用 TorchTitan 的组件做配置编程组合、派生、变换配置对象而不是记忆一堆 CLI 开关。配置函数以torchtitan_recipes.tests.features中的llama3_debugmodel_fsdp2_cp2这类形式存在选择它只需两个参数--module定义配置函数的模块--config模块内的配置函数名。NGPU4 MODULEtorchtitan_recipes.tests.features CONFIGllama3_debugmodel_fsdp2_cp2 ./run_train.sh这条命令的解析链路在 torchtitan/config/manager.py 的ConfigManager._load_config中先扫描参数提取--module与--config按支持的模型/实验缩写 → 完整模块路径的顺序导入模块再从模块中取出同名可调用对象并执行得到初始化的Trainer.Config对象随后剩余的 CLI 参数交给 tyro 做覆盖与合并。并行度在配置里世界大小在环境里一个关键设计是并行度parallelism degrees写在配置里但世界大小world size不在配置里。因此启动时的NGPU必须与配置中所有并行度的乘积严格一致。# 假设配置中 data_parallel_shard_degree × context_parallel_degree 4 NGPU4 MODULEtorchtitan_recipes.tests.features CONFIGllama3_debugmodel_fsdp2_cp2 ./run_train.shrun_train.sh脚本见仓库根目录 run_train.sh负责把环境变量转换成实际的分布式启动命令正常训练时以torchrun --nproc_per_node${NGPU}拉起torchtitan.train模块当COMM_MODEfake_backend时则走python3 -m torchtitan.train --comm.modefake_backend --training.steps 1的干跑路径用假进程组在单 GPU 上做配置与模型搭建的验证无需 NCCL 与多卡。改变行为的方式写自己的配置函数TorchTitan 明确建议想改变任何行为直接修改配置——写一个返回新Trainer.Config的函数而不是去堆 CLI flags。派生出新配置往往只需要几行# torchtitan_recipes/my_runs.py def llama3_debugmodel_fsdp2_cp4() - Trainer.Config: config llama3_debugmodel_fsdp2_cp2() config.parallelism.context_parallel_degree 4 config.training.steps 100 return config这里config.parallelism是 torchtitan/config/configs.py 中定义的ParallelismConfigconfig.training是TrainingConfig。llama3_debugmodel_fsdp2_cp2这类基础函数位于torchtitan_recipes.tests集成测试套件所用配置每个套件一个模块或各模型的config_registry.py如 torchtitan/models/llama3/config_registry.py 中的llama3_debugmodel。派生时直接沿用基础配置的字段并覆写既减少重复也让新旧配置的差异一目了然。旧的--section.optionflags兼容而非推荐--section.option形式的 CLI flags仍然可用且优先级高于配置ConfigManager的注释明确写出优先级CLI args config_registry function defaults但保留它们只是为了不破坏已有脚本——它们不再是配置运行的方式并且终将被删除。例如NGPU8 ./run_train.sh --parallelism.context_parallel_degree 2这种覆盖通过 tyro 的tyro.cli(config_cls, argsargs, defaultloaded_config, registrycustom_registry)完成见 torchtitan/config/manager.py 的parse_args。新项目请直接以配置函数为准把 CLI 覆盖仅当作临时调试手段。配置放在哪里torchtitan_recipes 包完整配置集中存放在独立的顶层包 torchtitan_recipes 中一个配置函数就是一个recipe。它之所以与torchtitan平级而不放在其内部是因为两者承载的东西不同torchtitan提供模型定义与实现各项优化的类torchtitan_recipes只负责挑选这些组件的组合。配置还天然绑定到某一台集群 某一次运行其变更节奏与库不同因此发布一个配置并不等于承诺发布一个类。需要留意一个安装细节torchtitan_recipes是第二个顶层包如果你做过可编辑安装pip install -e .之后才出现该包那么在仓库根目录之外运行会出现import torchtitan_recipes失败。此时只需重新执行一次pip install -e .从仓库根目录运行run_train.sh与 CI 的做法则两种方式均可。编写自己的配置为一个集群写一个函数不同的集群通常意味着不同的分片布局sharding layout也就意味着不同的配置。这不需要改任何代码在torchtitan_recipes中、以模型命名的模块里添加一个函数然后在命令行指名即可# torchtitan_recipes/llama3.py def llama3_8b_fsdp8_tp2_h200() - Trainer.Config: model_spec model_registry(8B, attn_backendflex) return Trainer.Config( model_specmodel_spec, parallelismParallelismConfig( data_parallel_shard_degree8, tensor_parallel_degree2, ), ... )ParallelismConfig中的字段含义与约束来自 torchtitan/config/configs.py 的 docstring包括字段默认值说明data_parallel_replicate_degree1权重复制的数据并行度1 且data_parallel_shard_degree1 时为 HSDP否则为 DDP1表示禁用data_parallel_shard_degree-1权重分片的数据并行度1 时为 FSDP配合 replicate 度则为 HSDP-1表示使用DP_REPLICATE/SP/PP 之后剩余的所有 rank只有该字段可以为负1表示禁用tensor_parallel_degree1张量并行度1表示禁用enable_sequence_parallelTrue是否在张量并行中启用序列并行pipeline_parallel_degree1流水线并行度物理 rank 数1表示禁用pipeline_parallel_schedule1F1B流水线调度循环调度如Interleaved1F1B需要pipeline_parallel_degreerank 数且 split_pointsstage 数-1num_pp_microbatches1每个 DP rank 每个梯度累积迭代的流水线微批数PP 关闭时忽略context_parallel_degree1上下文并行度1表示禁用context_parallel_load_balancerheadtailCP 负载均衡器headtailSDPA 用、ptrrFlexInnerAttention 用或Noneexpert_parallel_degree1专家并行度对非 MoE 模型无影响要求稠密区与稀疏区覆盖相同 rankdp_shard * cp * tp efsdp * ep--module接受任意可导入的模块因此把配置放在本仓库之外也完全可行工作方式不变MODULEmy_company_configs.experiments CONFIGllama3_ablation_7 ./run_train.sh仓库内现成的 recipe 示例可参考 torchtitan_recipes/muse_glimmer.py其中muse_glimmer_30b_allgather_cp8通过apply_transforms(config, [ContextParallelTransform(inner_attentionKVAllGatherCPFlexInnerAttention)])构造 CP 配置muse_glimmer_30b_ulysses_cp2则演示了模型只有两个 KV head、Ulysses CP 只能到 2这类硬件/架构约束如何直接写进配置注释。配置树的完整形态Trainer.Config 的字段一次运行的完整配置由Trainer.Configtorchtitan/trainer.py 中的嵌套 dataclass继承自Configurable.Config承载。从源码可见其标准字段model_spec、hf_assets_path、dump_folder、profiler、metrics、tokenizer、dataloader、optimizer、lr_scheduler、training、parallelism、checkpoint、activation_checkpoint、compile、comm、validator、debug、sdc_replayer、override、loss。其中若干高频字段值得说明均为 torchtitan/config/configs.py 中 docstring 的原意training.num_tokens_per_microbatch_per_dp_rank默认16384单个 DP rank 一次前向处理的输入 token 槽位数在 CP/TP 分片之前计training.num_tokens_per_train_step默认-1跨 DP rank、PP 微批与梯度累积的全局 token 槽位数默认由前者 × 微批数 × DP 度推导training.steps默认10000训练步数training.dtype默认float32可选bfloat16全 bf16 训练时参数/梯度/优化器状态都放 bf16不再额外保留 fp32 权重副本RoPE 与 logits 仍在 fp32training.mixed_precision_param/mixed_precision_reduce混合精度下参数与规约的 dtypecompile.enable/compile.components默认[model, loss]torch.compile开关与编译组件enable_async_tensor_parallel还要求model在组件中debug.print_config与debug.save_config_file打印配置或导出配置文件便于核对。这些 dataclass 都强制dataclass(kw_onlyTrue, slotsTrue)并由 torchtitan/config/configurable.py 的Configurable.__init_subclass__在类定义时强制检查——这是整个配置体系对每个组件都自带一个规范配置的约束保证。Transforming model configs变换模型配置树Model config transforms 是修改已完成的模型配置树的官方支持方式之一。它们是可选的用户完全可以使用自己的工具。核心用法先设置好全部训练选项然后一次性调用apply_transformsconfig muse_glimmer_30b() config.parallelism.context_parallel_degree 8 config apply_transforms( config, [ContextParallelTransform(inner_attentionKVAllGatherCPFlexInnerAttention)], )apply_transforms会深拷贝 trainer 配置、依次应用变换并校验结果失败时原配置保持不变transform_model_config_则用于没有 trainer 配置、只有裸ModelSpec的场景原地改写并返回根节点。从 torchtitan/config/transform/README.md 的分工约定看model_registry负责选择基础架构、注意力算法与注意力元数据格式transform 负责替换或包装树中的配置例如上下文并行、TP GEMM 后端、MoE 通信后端、量化与 LoRA 都属于 transform 的职责范围。Validation单字段校验与跨区校验的分工配置校验遵循清晰的分层约定只涉及单个配置字段的简单不变量写在对应配置类的__post_init__里。例如TrainingConfig.__post_init__会拒绝num_tokens_per_microbatch_per_dp_rank 0、num_tokens_per_train_step非法取值与负的max_normParallelismConfig.__post_init__会校验context_parallel_load_balancer的取值、fsdp_symm_mem_scope的 GPU 算力要求以及 pipeline schedule 名称是否有效。跨配置区section的检查放在 torchtitan/config/validation.py 中。其中validate_context_parallel遍历模型中每个BaseAttention.Config校验 CP 开启时注意力必须是CPInnerAttention系列、CP1 时不能残留 CP 注意力、同一模型内 CP 后端必须一致以及 Ulysses 后端下n_heads/n_kv_heads必须能被tensor_parallel_degree * context_parallel_degree整除。Trainer.Config.__post_init__会在自身局部校验之后调用这些跨区检查见 torchtitan/trainer.py 中__post_init__末尾对validate_context_parallel的调用形成局部校验在前、全局校验兜底的两层防线。命令行选项被冻结新特性走配置树--section.option这组 CLI flags 的数量将不再增长。新特性把它们的旋钮knobs放进配置树因此引入新特性的方式是新增一个配置而不是新增一个 CLI flag。已经存在于命令行上的东西会继续工作——那是为了向后兼容而不是因为它是推荐路径。最终目标是彻底移除这些 flags只保留--module与--config甚至完全移除 tyro。被冻结的是CLI而不是配置 dataclass。新字段仍然可以加入配置树加在所属组件上、加在模型上或者——对于那些没有更好归属的少数选项如training.num_tokens_per_microbatch_per_dp_rank——加在 torchtitan/config/configs.py 中该文件并未封闭但添加前需与维护者讨论。tyro.conf.Suppress让配置树字段不变成 CLI组件配置上的字段或configs.py中的字段默认都会成为 CLI 选项除非你用tyro.conf.Suppress注解它——正是这个注解保证了 CLI 不增长同时配置仍然可以设置该字段new_job_level_knob: Annotated[int, tyro.conf.Suppress] 3模型配置中的字段则什么都不需要因为Trainer.Config.model_spec本身已经用tyro.conf.Suppress注解整个模型配置树都被排除在命令行之外。在 torchtitan/trainer.py 中可以看到model_spec的实际注解写法Annotated[ModelSpec, tyro.conf.Suppress, tyro.conf.arg(constructorstr)]Trainer.Config.sdc_replayer同样用tyro.conf.Suppress隐藏只能在 recipe 里以编程方式启用。什么该放进 torchtitan_recipes仓库随包发布的内容是刻意的小集合见 torchtitan_recipestests/——集成测试套件运行的配置每个套件一个模块在特定硬件上验证过的golden 配置以硬件命名保证仅凭名称即可复现基准运行演示新特性的配置。仓库不会发布模型 × 并行度 × 优化的全组合——那是指数级的。你的运行是你自己的配置可以加到本仓库但不提交也可以放在自己的包里并把--module指向它从已发布配置派生只需要几行如本文第二节所示。config_registry 的迁移方向各模型目录下的config_registry.py模块用--module model --config function选择例如 torchtitan/models/llama3/config_registry.py 的llama3_debugmodel是同一件事的早期位置。它们会继续工作但终将被删除其中保存的模型规格基线llama3_8b之类会迁入torchtitan_recipes届时命令行变为--module torchtitan_recipes.llama3。目前没有为迁移提供 shim 的计划——在每个模型目录里做个 re-export 无非是给每个配置再起一个名字。小结从 flags 走向配置编程TorchTitan 的配置体系完成了一次明确的方法论迁移从靠 CLI 参数描述运行转向靠 Python 函数描述运行。配置即函数意味着派生、复用与版本管理都复用 Python 语言本身的能力torchtitan_recipes作为配置的独立顶层包把组件的实现与组件的组合解耦校验分层、CLI 冻结与tyro.conf.Suppress机制则共同保证了 CLI 表面积不膨胀让新特性以新配置而非新 flag的形式进入系统。对使用者来说掌握这一套约定就是掌握了在任意集群上快速产出可复现、可对比训练运行的核心技能。【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考