跳到主要内容

数据截至 (上游 commit bd2a0fc7c314)

01 · recipe 与配置系统

这一章讲什么: torchtune 最与众不同的两件事:① recipe(训练配方)的组织哲学——脚本即文档、禁止继承;② 配置系统——一份 YAML 如何通过 _component_ 点路径被递归装配成模型、优化器、数据集。读完你会知道 tune run ... --config llama3_1/8B_full 这行命令背后发生了什么。

1. 它要解决的小问题​

微调框架通常长成一个「Trainer 类 + 几十个子类 + 回调钩子」的继承树。想改一个行为,你得先搞清楚该 override 哪一层。

torchtune 的作者们反过来问:能不能让每个训练配方就是一份从上到下读完的脚本? 难点在于:不抽象,十几个 recipe 之间会有大量重复代码;抽象了,又走回继承树的老路。

2. 思路:复制是特性,不是坏味道​

torchtune 的选择是容忍重复、拒绝继承。这条纪律直接写进接口契约里——FTRecipeInterface 的 docstring(torchtune/recipe_interfaces.py:10-22)有三句关键的话:

  • 「torchtune strictly prohibits implementation inheritance in the codebase」——代码库里严禁实现继承。
  • 「Minimizing code duplication is not the goal. Recipe-writers are encouraged to copy-paste-modify.」——最小化重复不是目标,鼓励复制-粘贴-修改。
  • 接口只约束方法名(setup / train / save_checkpoint / cleanup),不承载任何共享实现。

配套的工作流也围绕「复制」设计:tune cp 子命令(torchtune/_cli/cp.py:18 的 Copy)存在的全部意义就是把内置 recipe 或 config 拷到你的本地目录让你改。复制粘贴从「坏实践」升格为一等公民。

代价与收益的权衡很直白:重复代码确实存在(比如各 recipe 的 _setup_profiler 几乎一样),但每份脚本都是自包含的答案——读 lora_finetune_single_device.py 不需要先读另外三个基类。

3. tune CLI:从命令行到 recipe 进程​

3.1 入口与子命令​

安装包在 pyproject.toml:48-49 注册了 tune = "torchtune._cli.tune:main"。TuneCLIParser(torchtune/_cli/tune.py:18)挂六个子命令:

子命令干什么
tune run跑一个 recipe(单进程直接跑,多进程走 torchrun)
tune ls列出全部内置 recipe 及其配套 config
tune cp把内置 recipe/config 拷到本地(改造的起点)
tune download从 HF Hub / Meta 下载模型权重
tune cat打印某个内置 config 的内容
tune validate校验一份 config 能被正确解析

3.2 注册表:recipe 名怎么找到文件​

内置配方登记在 _ALL_RECIPES(torchtune/_recipe_registry.py:24)这个纯数据列表里,每条是一个 Recipe dataclass(torchtune/_recipe_registry.py:17):

Recipe(name="full_finetune_single_device",
       file_path="full_finetune_single_device.py",
       configs=[Config("llama3_2/1B_full_single_device", ...), ...],
       supports_distributed=True)

get_all_recipes()(torchtune/_recipe_registry.py:701)返回这份列表,tune ls 和 tune run 都靠它解析名字。注册表是数据而非插件机制——加一个内置 recipe 就是往列表里加一条。

3.3 tune run 的两条路径​

Run._run_cmd(torchtune/_cli/run.py:168)的分发逻辑:

tune run [--nproc_per_node N ...] <recipe> --config <config> [k=v ...]
              │
              ▼
      recipe 是内置名? ──否──► 当作本地文件/模块(自定义 recipe)
              │是
              ▼
   查注册表拿到绝对路径 + supports_distributed
              │
     带了 torchrun 参数?(_is_distributed_args)
        │              │
       是              否
        ▼              ▼
 _run_distributed   _run_single_device
 (torchrun 语义)    (runpy 直接跑)

两个实现细节值得记住:

  • torchrun 参数是「偷」来的。 _add_arguments(torchtune/_cli/run.py:61)直接遍历 torch.distributed.run 自己的 argparser actions,把 --nproc_per_node 等参数原样嫁接到 tune run 上,只把 training_script 改名为 recipe。所以 torchrun 新增参数时 tune run 自动跟进。
  • 单机路径用 runpy。 _run_single_device(torchtune/_cli/run.py:103)把 recipe 文件用 runpy.run_path(..., run_name="__main__") 在当前进程里跑起来;分布式路径 _run_distributed(torchtune/_cli/run.py:83)则把参数回填后调 torchrun 的 run(args),且没指定 rendezvous 时强制 standalone=True,让多个训练任务可以并行起跑。

自定义 recipe 不在注册表里也能跑:传文件路径即可,_convert_to_dotpath 把它转成模块点路径用 python -m 语义执行(torchtune/_cli/run.py:146-157)。

4. @config.parse:YAML 与命令行怎么合成一棵树​

每个 recipe 的入口都长一个样(recipes/full_finetune_distributed.py:1152-1165):

@config.parse
def recipe_main(cfg: DictConfig) -> None:
    config.log_config(recipe_name="FullFinetuneRecipeDistributed", cfg=cfg)
    recipe = FullFinetuneRecipeDistributed(cfg=cfg)
    recipe.setup(cfg=cfg)
    recipe.train()
    recipe.cleanup()

@config.parse(torchtune/config/_parse.py:68)这个装饰器做三件事:

  1. 用 TuneRecipeArgumentParser(torchtune/config/_parse.py:20)解析。它是 argparse 的子类,内置一个必填的 --config 参数;parse_known_args(torchtune/config/_parse.py:43)先正常解析拿到 yaml 路径,然后把 yaml 内容设为 argparse 的默认值再解析一遍(torchtune/config/_parse.py:58-64)——这样命令行参数天然覆盖 yaml 值。
  2. 用 _merge_yaml_and_cli_args(torchtune/config/_utils.py:121)把结果与 key=value 点路径覆盖(model.lora_rank=16)合并成一棵 DictConfig。OmegaConf 的 dotlist 语法让 a.b.c=v 直接写到嵌套层级。
  3. 调 recipe_main(conf) 并 sys.exit(torchtune/config/_parse.py:99)。

这个合并器有三个贴心的小设计:

语法行为位置
model=torchtune.models.lora_llama2_7b若覆盖目标带 _component_,自动改写成 model._component_=...torchtune/config/_utils.py:182-184
~metric_logger~ 前缀表示从 yaml 里删除这个键(组件本身不可删)torchtune/config/_utils.py:155-160
x=None字符串 "None" 转成 OmegaConf 的 nulltorchtune/config/_utils.py:186-188

坑: 命令行不接受任何额外的 --flag(只允许 --config 和 k=v),否则直接 ValueError(torchtune/config/_parse.py:52-56)。这是故意的——recipe 的可配置面完全由 yaml schema 定义,不允许脚本作者私加 flag。

5. _component_ + instantiate:配置即装配图​

这是整个库的心脏机制。规则只有一条:任何含 _component_ 键的字典,都会被替换成「该键指向的可调用对象、以其余键为 kwargs 调用后」的返回值。

5.1 直觉示例​

# 示意,非源码
cfg = {
    "model": {
        "_component_": "torchtune.models.llama3_1.llama3_1_8b",  # 装配哪个函数
        # 其余键 = kwargs
    },
    "optimizer": {
        "_component_": "torch.optim.AdamW",
        "lr": 3e-4,
    },
}
# instantiate(cfg["optimizer"], model.parameters())
#   等价于 torch.optim.AdamW(model.parameters(), lr=3e-4)

重点看:_component_ 的值可以是任何可 import 的点路径——torch 的类、torchtune 的 builder、甚至第三方库。官方 config 里就出现 bitsandbytes.optim.PagedAdamW8bit(recipes/configs/llama3_1/8B_full_single_device.yaml)。torchtune 因此不需要任何插件注册机制:能被 import,就能被配置。

5.2 真实实现:递归的两层​

机制分两半,都在 torchtune/config/_instantiate.py:

  • instantiate(torchtune/config/_instantiate.py:69)做准备工作:把 dict 转 DictConfig、深拷贝并解开 struct/readonly 限制(:141-146)、把调用方额外给的 kwargs merge 进去(:150)、解析 ${...} 插值(OmegaConf.resolve,:153),然后交给递归函数。docstring 注明它「Based on Hydra's instantiate utility」(:83)。
  • _instantiate_node(torchtune/config/_instantiate.py:27)是递归本体:字典里有 _component_ 就先递归实例化所有子值,再调用组件(:53-62);没有 _component_ 的普通字典/列表就逐元素递归(:47-52、:63-64);标量原样返回(:65-66)。
instantiate(cfg.model)
      │
      ▼
_instantiate_node({'_component_': ..., 'lora_attn_modules': [...], ...})
      │
      ├─ 先递归处理每个子值(子值里可能还嵌着 _component_)
      │
      ▼
_get_component_from_path("torchtune.models.llama3_1.lora_llama3_1_8b")
      │   import 模块、取属性(torchtune/config/_utils.py:37)
      ▼
_create_component(fn, args, kwargs)   # 就是 fn(*args, **kwargs)(:18-24)

嵌套即组合:dataset 的 kwargs 里可以再嵌一个 _component_(比如自定义 message transform),装配会自底向上完成。

5.3 三个让「自定义组件」能工作的细节​

  • 当前目录可 import。 instantiate 开头把 os.getcwd() 加进 sys.path(torchtune/config/_instantiate.py:138-139),所以你可以在自己的项目目录里写 my_module.MyDataset 并在 yaml 里引用。
  • 调用方 globals 兜底。 点路径解析失败时,_get_component_from_path 会去调用栈上一层的 globals 里找这个名字——instantiate 通过 inspect.currentframe().f_back.f_globals 把它传下去(torchtune/config/_instantiate.py:157-161)。于是 recipe 文件里定义的局部类也能被 _component_: MyLocalClass 引用。
  • 插值先解析再装配。 ${output_dir} 这类 OmegaConf 引用在 _instantiate_node 之前就被 OmegaConf.resolve 解析掉(torchtune/config/_instantiate.py:153),组件拿到的都是具体值。

5.4 坑​

坑表现位置
忘了写 _component_InstantiationError,报错会打印整份 configtorchtune/config/_instantiate.py:130-135
点路径拼错 / 没安装对应包同样到 instantiate 时才炸,不在 tune validate 阶段保证类型正确(它只做解析层面检查)torchtune/config/_utils.py:37
yaml 里用了保留键 config断言失败:「Cannot use 'config' within a config file」torchtune/config/_parse.py:59
把 recipe 的隐式约定当配置recipe 代码里 cfg.get("xxx", default) 的键不出现在 yaml 里也能生效——读 recipe 源码才是完整 schema,yaml 只是常用子集各 recipe 的 __init__(如 recipes/full_finetune_distributed.py:144-243)

6. 代码地图​

主题文件路径符号名
CLI 入口与子命令挂载torchtune/_cli/tune.pyTuneCLIParser、main
tune run 分发torchtune/_cli/run.pyRun._run_cmd、Run._run_distributed、Run._run_single_device、Run._add_arguments
复制到本地(改造起点)torchtune/_cli/cp.pyCopy._cp_cmd
recipe/config 注册表torchtune/_recipe_registry.pyRecipe、Config、_ALL_RECIPES、get_all_recipes
recipe 接口契约torchtune/recipe_interfaces.pyFTRecipeInterface
yaml+CLI 合成torchtune/config/_parse.pyparse、TuneRecipeArgumentParser.parse_known_args
点路径覆盖合并torchtune/config/_utils.py_merge_yaml_and_cli_args、_remove_key_by_dotpath
递归装配torchtune/config/_instantiate.pyinstantiate、_instantiate_node、_create_component
点路径解析torchtune/config/_utils.py_get_component_from_path
recipe 入口样例recipes/full_finetune_distributed.pyrecipe_main