什么是配方?¶
本深入探讨将引导您了解 torchtune 中训练配方的设计。
什么是配方?
构成配方的核心组件是什么?
我应该如何构建新的配方?
什么是配方?¶
配方是 torchtune 用户的主要入口点。可以将它们视为用于训练和可选地评估 LLM 的“目标”端到端管道。每个配方都实现了一种训练方法(例如:完全微调),并应用了一组有意义的功能(例如:FSDP + 激活检查点 + 梯度累积 + 混合精度训练)到给定的模型系列(例如:Llama2)。
随着模型训练变得越来越复杂,在预测新的模型架构和训练方法的同时,还要考虑所有可能的权衡(例如:内存与模型质量),变得越来越困难。我们认为 a) 用户最适合根据其用例进行权衡,以及 b) 没有一种万能的解决方案。因此,配方旨在易于理解、扩展和调试,而不是所有可能设置的通用入口点。
根据您的用例和专业知识水平,您会经常发现自己修改现有配方(例如:添加新功能)或编写新的配方。torchtune 通过提供经过良好测试的模块化组件/构建块和通用实用程序(例如:WandB 记录 和 检查点)来简化配方编写。
配方设计
torchtune 中的配方旨在
简单。完全用原生 PyTorch 编写。
正确。对每个组件进行数值一致性验证,并与参考实现和基准进行广泛比较。
易于理解。每个配方提供有限数量的有意义功能,而不是隐藏在数百个标志后面的所有可能功能。代码重复优先于不必要的抽象。
易于扩展。不依赖于训练框架,也不存在实现继承。用户无需层层抽象就能弄清楚如何扩展核心功能。
- 面向各种用户。用户可以决定他们希望如何与 torchtune 配方交互
通过修改现有配置开始训练模型
修改现有配方以适应自定义情况
直接使用可用的构建块来编写全新的配方/训练范式
每个配方包含三个组件
可配置参数,通过 yaml 配置和命令行覆盖指定
配方脚本,入口点,将所有内容整合在一起,包括解析和验证配置、设置环境以及正确使用配方类
配方类,训练所需的核心逻辑,通过一组 API 向用户公开
在接下来的部分中,我们将更详细地了解这些组件中的每一个。有关完整的示例,请参阅 torchtune 中的 完整微调配方 以及相关的 配置。
配方不是什么?¶
单体训练器。配方不是一个单体训练器,旨在通过数百个标志来支持所有可能的功能。
通用入口点。配方不是为了支持所有可能的模型架构或微调方法。
外部框架的包装器。配方不是外部框架的包装器。这些完全用原生 PyTorch 编写,使用 torchtune 构建块。依赖项主要以附加实用程序或与周围生态系统(例如:EleutherAI 的评估工具)的互操作性的形式存在。
配方脚本¶
这是每个配方的主要入口点,并为用户提供对如何设置配方、如何训练模型以及如何使用后续检查点的控制。这包括
环境设置
解析和验证配置
训练模型
设置多阶段训练(例如:蒸馏),使用多个配方类
脚本通常应按以下顺序构建操作
初始化配方类,该类反过来初始化配方状态
加载和验证检查点以更新配方状态(如果恢复训练)
从检查点初始化配方组件(模型、分词器、优化器、损失和数据加载器)(如果适用)
训练模型
训练完成后清理配方状态
示例脚本如下所示
# Initialize the process group
init_process_group(backend="gloo" if cfg.device == "cpu" else "nccl")
# Setup the recipe and train the model
recipe = FullFinetuneRecipeDistributed(cfg=cfg)
recipe.setup(cfg=cfg)
recipe.train()
recipe.cleanup()
# Other stuff to do after training is complete
...
配方类¶
配方类承载训练模型的核心逻辑。每个类都实现相关的接口并公开一组 API。对于微调,此类的结构如下
初始化配方状态,包括种子、设备、数据类型、指标记录器、相关标志等
def __init__(...):
self._device = utils.get_device(device=params.device)
self._dtype = utils.get_dtype(dtype=params.dtype)
...
加载检查点,从检查点更新配方状态,初始化组件并从检查点加载状态字典
def setup(self, cfg: DictConfig):
ckpt_dict = self.load_checkpoint(cfg.checkpointer)
# Setup the model, including FSDP wrapping, setting up activation checkpointing and
# loading the state dict
self._model = self._setup_model(...)
self._tokenizer = self._setup_tokenizer(...)
# Setup Optimizer, including transforming for FSDP when resuming training
self._optimizer = self._setup_optimizer(...)
self._loss_fn = self._setup_loss(...)
self._sampler, self._dataloader = self._setup_data(...)
在所有时期运行前向和反向传播,并在每个时期结束时保存检查点
def train(...):
self._optimizer.zero_grad()
for curr_epoch in range(self.epochs_run, self.total_epochs):
for idx, batch in enumerate(self._dataloader):
...
with self._autocast:
logits = self._model(...)
...
loss = self._loss_fn(logits, labels)
if self.total_training_steps % self._log_every_n_steps == 0:
self._metric_logger.log_dict(...)
loss.backward()
self._optimizer.step()
self._optimizer.zero_grad()
# Update the number of steps when the weights are updated
self.total_training_steps += 1
self.save_checkpoint(epoch=curr_epoch)
清理配方状态
def cleanup(...)
self.metric_loggers.close()
...
使用配置运行配方¶
要使用一组用户定义的参数运行配方,您需要编写一个配置文件。您可以在我们的 配置深入解析 中了解有关配置的所有信息。
使用 parse
解析配置和 CLI¶
我们提供了一个方便的装饰器 parse()
,它包装您的配方以启用使用 tune
从命令行运行,并进行配置和 CLI 覆盖解析。
@config.parse
def recipe_main(cfg: DictConfig) -> None:
recipe = FullFinetuneRecipe(cfg=cfg)
recipe.setup(cfg=cfg)
recipe.train()
recipe.cleanup()
运行您的配方¶
您应该能够通过使用 tune
命令以及任何 CLI 覆盖,提供您自定义配方和自定义配置的直接路径来运行您的配方。
tune run <path/to/recipe> --config <path/to/config> k1=v1 k2=v2 ...