跳到主要内容

数据截至 (上游 commit bd2a0fc7c314)

torchtune — 架构与原理

30 秒导读: torchtune 是 PyTorch 官方的 LLM 微调库。它的核心主张是「反框架」:每种训练方式(SFT、LoRA、DPO、知识蒸馏、QAT、GRPO)都是 recipes/ 下一份从上到下能读完的普通 Python 脚本,官方明说你应该复制它、改它,而不是继承它。YAML 配置只负责一件事——用 _component_: torchtune.models.llama3_1.llama3_1_8b 这样的点路径声明「装配哪个函数/类」,由 config.instantiate 递归变成真对象。分布式与显存优化全部以 torchtune.training 里的显式工具函数写进脚本,没有藏在背后的 Trainer 抽象。

维护状态提醒: 该项目已于 2025 年停止活跃开发(见 README.md:1 的停维公告)。本 teardown 针对最后的主线状态,其设计仍极有参考价值。

1. 这是什么(零基础也能懂)​

一句话定义​

torchtune 是一个用原生 PyTorch 写的大语言模型后训练(post-training)库:你给它一个开源模型(Llama / Qwen / Gemma / Mistral / Phi……)和一份数据,它提供一整套「训练配方(recipe)」把监督微调、LoRA、DPO 这类训练跑起来,从单卡到多机多卡。

它要解决谁的什么问题​

假设你想在自己的 8 张卡上微调一个 Llama-3.1-8B:

  • 用 HuggingFace 全家桶(Trainer + PEFT + Accelerate + TRL)当然能跑,但框架层很厚——出了 bug 要在好几层回调和封装里翻。
  • 自己从零写训练脚本,则要把 FSDP 分片、activation checkpointing、checkpoint 格式转换、变长序列的 loss 归一化这些坑全部自己踩一遍。

torchtune 的答案是第三条路:把「正确的训练脚本长什么样」直接给你看。 每个 recipe 就是答案本身——能跑、能读、能复制走改成你自己的。所以它的源码组织原则是「recipes are meant to be read and copied, not subclassed」。

它能做什么​

能力具体支持
训练配方全量/LoRA 微调(单卡、多卡、多机)、DPO、PPO、GRPO(dev)、知识蒸馏、QAT、生成、EleutherAI 评测(recipes/ 目录)
模型Llama2/3/3.1/3.2/3.3/4、Gemma/Gemma2、Mistral、Phi3/Phi4、Qwen2/2.5/3、Smol、T5(torchtune/models/)
并行FSDP2(fully_shard)、张量并行(TP)、上下文并行(CP),由 ParallelDims 组合(torchtune/training/_distributed.py:62)
显存技术activation checkpointing(全量/选择性)、activation offloading、optimizer-in-backward、fp8 训练、QLoRA(NF4)
配置YAML + 命令行点路径覆盖(model.lora_rank=16),_component_ 声明式装配
互操作直接读写 HF / Meta 原生 checkpoint 格式,LoRA 可导出为 PEFT 格式

用起来什么样​

安装后得到一个 tune 命令(pyproject.toml:48-49 注册的 entry point):

# 单卡 LoRA 微调,内置配方 + 内置配置名
tune run lora_finetune_single_device --config llama3_2/3B_lora_single_device

# 4 卡全量微调:torchrun 参数原样透传
tune run --nproc_per_node 4 full_finetune_distributed --config llama3_1/8B_full

# 命令行点路径覆盖 YAML 里的任意字段
tune run lora_finetune_single_device --config llama3_1/8B_lora_single_device model.lora_rank=16

配置长这样(摘自 recipes/configs/llama3_1/8B_lora.yaml):

model:
  _component_: torchtune.models.llama3_1.lora_llama3_1_8b
  lora_attn_modules: ['q_proj', 'v_proj', 'output_proj']
  apply_lora_to_mlp: True
  lora_rank: 8
  lora_alpha: 16
loss:
  _component_: torchtune.modules.loss.LinearCrossEntropyLoss
optimizer:
  _component_: torch.optim.AdamW
  lr: 3e-4

读法:每个 _component_ 就是一个 Python 可调用对象的点路径,其余键是它的 kwargs。训练时这些配置块被递归地变成真模型、真损失、真优化器(机制见第 1 章)。

一句话直觉​

把 torchtune 想成一本「可以执行的食谱」。 框架是给你端上一道菜(API 背后的黑盒),torchtune 是把菜谱全文给你:每一步用什么料(组件)、什么火候(超参)、什么刀工(分布式技巧),全部写在明面上,你可以照抄开小灶。

2. 顶层全景(它大概怎么转)​

2.1 一张顶层图​

怎么读这张图: 从左到右是一次 tune run 的生命周期;中间扇出的四个部件全部由配置装配,再汇进训练循环。

┌───────────┐   ┌────────────────┐   ┌──────────────────────┐
│  tune CLI │──►│  recipe 脚本    │──►│  config.instantiate  │
│  tune run │   │  @config.parse │   │  _component_ 递归装配 │
└───────────┘   └────────────────┘   └──────────┬───────────┘
                                                 │
              ┌──────────────┬───────────────────┼─────────────┐
              ▼              ▼                   ▼             ▼
         模型 builder    tokenizer         dataset/       optimizer
         (models/)       (tokenizers/)     dataloader     + loss
              │                            (datasets/)
              ▼
        recipe.train() 训练循环
        FSDP/TP/CP · AC/offload · token 归一化 loss
              │
              ▼
        CheckpointClient ── 存回 HF/Meta 格式

2.2 部件职责​

部件干什么在哪个文件
tune CLI子命令入口:run / ls / cp / download / cat / validate;run 把 torchrun 参数嫁接进自己的 argparsertorchtune/_cli/tune.py、torchtune/_cli/run.py:26(Run)
recipe 注册表内置 recipe 名 → 脚本路径 + 配套 config 名单torchtune/_recipe_registry.py:17(Recipe)、:701(get_all_recipes)
recipe 脚本一个训练配方的全部流程:__init__ → setup → train → cleanuprecipes/full_finetune_distributed.py:53(FullFinetuneRecipeDistributed)等
配置系统YAML + CLI 覆盖合成 DictConfig;instantiate 递归装配 _component_torchtune/config/_parse.py:68(parse)、torchtune/config/_instantiate.py:69(instantiate)
模型组件两级 builder(型号 builder → 组件 builder)拼装 TransformerDecodertorchtune/models/llama3_1/_model_builders.py:19、torchtune/modules/transformer.py:331
数据管线消息 → token+mask 的 transform、sample packing、block causal mask collatetorchtune/datasets/_sft.py:19、torchtune/datasets/_packed.py:17、torchtune/data/_collate.py:562
分布式/显存工具device mesh、逐层 FSDP、AC、offloading、optimizer-in-backwardtorchtune/training/_distributed.py、torchtune/training/memory.py
checkpoint加载 HF/Meta 权重并改名转置,保存时分片写回;LoRA 导出 PEFTtorchtune/training/checkpointing/_checkpointer.py:369、torchtune/models/convert_weights.py

2.3 主线走一遍(一次分布式全量微调)​

以 tune run --nproc_per_node 4 full_finetune_distributed --config llama3_1/8B_full 为例,对应 recipes/full_finetune_distributed.py:

  1. CLI 解析。 Run._run_cmd(torchtune/_cli/run.py:168)在注册表里查到 recipe 与 config 的真实路径;检测到带 torchrun 参数就交给 _run_distributed(torchtune/_cli/run.py:83),等价于 torchrun 起 4 个进程跑同一份脚本。
  2. 配置合成。 脚本入口 recipe_main(recipes/full_finetune_distributed.py:1152)被 @config.parse 包住:YAML 加载后与 key=value 形式的 CLI 覆盖合并成一棵 DictConfig 树。
  3. 装配。 setup()(recipes/full_finetune_distributed.py:325)依次 config.instantiate 出 checkpointer、模型、tokenizer、loss、optimizer、lr_scheduler、dataloader——每个都来自 YAML 的一个 _component_ 块。
  4. 分布式模型就位。 模型先在 meta device 上实例化(只占结构不占显存,recipes/full_finetune_distributed.py:611-612),然后按 ParallelDims 建的 device mesh 做 TP 切分与逐层 FSDP fully_shard,最后 load_from_full_model_state_dict 把完整权重分片灌进去(recipes/full_finetune_distributed.py:702)。
  5. 训练循环。 train()(recipes/full_finetune_distributed.py:949)逐 batch:前向算 loss 并乘上本 batch 的有效 token 数、backward;累积够步数后跨 rank all_reduce token 总数,用 scale_grads_ 按 world_size / num_tokens 缩放梯度,再 optimizer.step()。变长序列 + 梯度累积下这是数学上正确的平均方式(第 4 章展开)。
  6. 落盘。 CheckpointClient.save_checkpoint 把模型权重从 torchtune 命名转换回 HF 命名与排布,按加载时记住的 _weight_map 分片写回 safetensors,中途 checkpoint 附带 optimizer/dataloader 状态用于续训(torchtune/training/checkpointing/_checkpoint_client.py:403)。

3. 阅读地图(建议顺序)​

五章由浅入深。时间有限的话,读 01 → 04 就能抓住 torchtune 区别于 HF 全家桶的全部要害。

顺序章节讲什么适合谁
101-recipes-and-config.mdrecipe 哲学、tune CLI、_component_ + instantiate 配置装配所有人必读,这是 torchtune 的立身之本
202-model-components.mdbuilder 模式、TransformerDecoder、RoPE、消息掩码、sample packing想知道「模型从哪来、数据怎么变成 tensor」的人
303-checkpointing.mdHF/Meta 权重格式互转、_permute 的由来、Checkpointer 家族、PEFT 导出要跟 HF 生态互操作、或做权重转换的人
404-training-loop-distributed.md训练循环精读、token 归一化、ParallelDims、FSDP/TP/CP、AC/offloading、optimizer-in-backward关心吞吐、显存、分布式正确性的人
505-peft-lora.md(未写)LoRALinear、LoRA builder、参数冻结校验、DoRA、QLoRA做参数高效微调的人

4. 巧妙之处(可借鉴的技术)​

每条先白话点出妙处,细节与完整引用在对应章节。

  1. 把「不许继承」写进接口契约。 FTRecipeInterface 的 docstring 明说「torchtune strictly prohibits implementation inheritance in the codebase」(torchtune/recipe_interfaces.py:10-22)——接口只约定方法名,不承载复用,鼓励 copy-paste-modify。这是对「框架腐化」的制度性防御。
  2. 配置即装配图。 _component_ 点路径 + 递归 instantiate(torchtune/config/_instantiate.py:27),让 YAML 能直接引用任意可调用对象——包括第三方库(bitsandbytes.optim.PagedAdamW8bit 就这么出现在官方 config 里),库自己不用写任何插件机制。
  3. RoPE 交错差异一次转置抹平。 HF 与 Meta 的 q/k 权重在 RoPE 维度上排列不同,torchtune 在加载时用 _permute(torchtune/models/convert_weights.py:151)做一次 reshape-transpose 对齐,训练全程不用管两套约定。
  4. RoPE 不进 state dict,meta device 友好。 RoPE 余弦缓存注册为 persistent=False 的 buffer(torchtune/modules/position_embeddings.py:46-66),模型可以在 meta device 上实例化,FSDP 分片完、加载权重前再统一 rope_init()(recipes/full_finetune_distributed.py:694-698)。
  5. loss 按 token 数归一化。 变长序列下逐样本平均是错的;recipe 里 loss 乘 current_num_tokens,梯度累积结束时 all_reduce 总 token 数再 scale_grads_ 缩放(recipes/full_finetune_distributed.py:1016、:1035,torchtune/training/_grad_scaler.py:43)。
  6. 把 lm_head 搬进 loss 函数。 LinearCrossEntropyLoss.set_model_output 让模型跳过输出投影只产 hidden states,loss 内部先丢掉 pad/ignored token 再分块投影算交叉熵(torchtune/modules/loss/cross_entropy_loss.py:67、:117、:147)——vocab 12 万时 logits 是显存大户,这一手省掉整份 [b, s, vocab] 张量。
  7. sample packing + flex attention。 多条短样本拼成一条定长序列,用 block causal document mask 防止跨样本注意力(torchtune/data/_collate.py:562、torchtune/modules/attention_utils.py:133),padding 浪费降到接近零。
  8. optimizer-in-backward。 给每个参数挂 register_post_accumulate_grad_hook,梯度一就绪就地对它 step()(torchtune/training/memory.py:222),省掉同时持有全部梯度的显存峰值——单机低显存场景的杀招。
  9. activation offloading 的「输出头例外」。 offloading 默认把大激活搬到 CPU,但输出投影(lm_head)的激活搬过去马上要搬回来,纯亏;于是给 model.output 挂 hook 就地进入 NoOp 上下文豁免它(torchtune/training/_activation_offloading.py:412-431 附近)。

5. 边界与局限​

诚实清单,大多写在 recipe 的 docstring 里:

  • 已停止维护。 2025 年官方宣布 wind down(README.md:1);新模型、新 PyTorch API 不会再跟进。
  • 模型覆盖面窄于 HF transformers。 只有十几族模型,每族都要手写 builder 与权重转换——这是「简单实现」哲学的代价。
  • 只支持 FSDP,不支持 DDP;不支持 fp16 全量训练与混合精度。 见 FullFinetuneRecipeDistributed 的 docstring(recipes/full_finetune_distributed.py:63、:89)。
  • optimizer_in_bwd 与梯度累积、梯度裁剪互斥,配置冲突时 __init__ 直接抛错(recipes/full_finetune_distributed.py:222-232)。
  • activation offloading 只支持 CUDA/XPU,且必须先开 activation checkpointing;多模态模型直接 NotImplementedError。 见 recipes/full_finetune_distributed.py:255-263 与 torchtune/training/_activation_offloading.py:435。
  • 只支持 map-style 数据集,流式数据集不在支持范围(recipes/full_finetune_distributed.py:804 的 docstring)。
  • checkpoint 以 epoch 为单位;步级保存是后加的 save_every_n_steps,旧文档的「mid-epoch checkpointing is currently not supported」说明仍留在 docstring 里(recipes/full_finetune_distributed.py:108)。
  • 功能刻意少于 HF 栈。 没有 callback 系统、没有 Trainer 生命周期钩子——这是设计选择,不是缺口(源卡片 gotcha)。

6. 横向对比​

同书架上与 torchtune 相邻的几个项目,取舍各不相同:

维度torchtune(本篇)trlacceleratepeft
定位完整训练配方库(可读可复制)RLHF/DPO 算法训练器库分布式启动/混合精度抽象层参数高效微调方法库
抽象哲学反框架:脚本即文档Trainer 子类 + 回调包住 launch 与 device 的最薄层给任意 nn.Module 注入 adapter
配置YAML + _component_ 装配dataclass *Config命令行 + AcceleratorPeftConfig dataclass
LoRA 实现自实现 LoRALinear(本库内)依赖 peft不管主场
分布式FSDP2/TP/CP 显式写进脚本委托 accelerate主场(封装 launch)不管
底层直接站在 pytorch 的 FSDP2/DTensor 上transformers + acceleratetransformers 的下游transformers 的旁挂

一句话:torchtune 是把 trl/accelerate/peft 三者叠起来才能覆盖的领域,用「零抽象 + 原生 PyTorch」重做了一遍——牺牲覆盖面与扩展点,换来「一个文件读懂一次训练」。

7. 代码地图(入口级)​

每章末尾有自己的细粒度地图,这里只列从零开始读源码的五个入口:

主题文件路径符号名
一切的开始:CLI 与 recipe 分发torchtune/_cli/run.pyRun._run_cmd、Run._run_distributed
最典型的 recipe(分布式全量)recipes/full_finetune_distributed.pyFullFinetuneRecipeDistributed.setup、FullFinetuneRecipeDistributed.train
配置装配torchtune/config/_instantiate.pyinstantiate、_instantiate_node
模型拼装样例torchtune/models/llama3_1/_component_builders.pyllama3_1、lora_llama3_1
训练循环本体torchtune/modules/transformer.py + recipes/full_finetune_distributed.pyTransformerDecoder.forward、FullFinetuneRecipeDistributed.train