什么是食谱?¶
本次深入探讨将引导您了解torchtune中训练食谱的设计。
什么是食谱?
构成食谱的核心组成部分是什么?
我应该如何构建一个新的配方?
Recipes 是 torchtune 用户的主要入口点。这些可以被视为“有针对性的”端到端管道,用于训练和可选地评估 LLMs。每个 recipe 实现了一种训练方法(例如:全微调),并应用了一组有意义的特性(例如:FSDP + 激活检查点 + 梯度累积 + 混合精度训练)到给定的模型系列(例如:Llama2)。
随着模型训练变得越来越复杂,预测新的模型架构和训练方法变得越来越困难,同时还要考虑所有可能的权衡(例如:内存与模型质量)。我们相信 a) 用户最适合根据他们的使用场景做出特定的权衡,b) 没有一种适用于所有情况的解决方案。因此,配方旨在易于理解、扩展和调试,而不是适用于所有可能设置的通用入口点。
根据您的使用案例和专业水平,您会经常发现自己需要修改现有的配方(例如:添加新功能)或编写新的配方。torchtune 通过提供经过良好测试的模块化组件/构建块和通用工具(例如:WandB 日志记录 和 检查点),使编写配方变得容易。
食谱设计
torchtune 中的配方设计为:
简单。完全用原生PyTorch编写。
正确。对每个组件进行数值奇偶校验,并与参考实现和基准进行广泛比较。
易于理解。每个配方提供一组有限的有意义的特性,而不是隐藏在数百个标志背后的所有可能特性。代码重复优于不必要的抽象。
易于扩展。不依赖训练框架,也不实现继承。用户不需要通过层层抽象来弄清楚如何扩展核心功能。
- Accessible to a spectrum of Users. Users can decide how they want to interact with torchtune recipes:
通过修改现有配置开始训练模型
修改现有配方以适应自定义情况
直接使用可用的构建块来编写全新的配方/训练范式
每个配方由三个部分组成:
可配置参数,通过yaml配置和命令行覆盖指定
Recipe Script,入口点,将所有内容整合在一起,包括解析和验证配置、设置环境以及正确使用配方类
Recipe Class,训练所需的核心逻辑,通过一组API向用户公开
在接下来的部分中,我们将更详细地了解这些组件中的每一个。 有关完整的工作示例,请参考 torchtune 中的 完整微调配方 以及相关的 配置。
什么是食谱不是的?¶
单体训练器。 一个配方不是一个旨在通过数百个标志支持每个可能功能的单体训练器。
通用入口点。 一个配方不旨在支持所有可能的模型架构或微调方法。
围绕外部框架的封装。 一个配方不应该是围绕外部框架的封装。这些完全使用torchtune构建块以原生PyTorch编写。依赖项主要以附加实用程序或与周围生态系统(例如: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 = training.get_dtype(dtype=params.dtype, device=self._device)
...
加载检查点,从检查点更新配方状态,初始化组件并从检查点加载状态字典
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(...)
在所有epoch中向前和向后运行,并在每个epoch结束时保存检查点
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.global_step % 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.global_step += 1
self.save_checkpoint(epoch=curr_epoch)
清理配方状态
def cleanup(...)
self.metric_loggers.close()
...
使用配置运行配方¶
要使用一组用户定义的参数运行配方,您需要编写一个配置文件。 您可以在我们的配置深入探讨中了解所有关于配置的信息。
使用parse进行配置和CLI解析¶
我们提供了一个方便的装饰器 parse(),它包装了你的配方,以便能够通过命令行使用 调优 运行,并支持配置和命令行覆盖解析。
@config.parse
def recipe_main(cfg: DictConfig) -> None:
recipe = FullFinetuneRecipe(cfg=cfg)
recipe.setup(cfg=cfg)
recipe.train()
recipe.cleanup()
运行你的配方¶
你应该能够通过使用调优命令并提供自定义配方和自定义配置的直接路径来运行你的配方,同时可以使用任何CLI覆盖:
tune run <path/to/recipe> --config <path/to/config> k1=v1 k2=v2 ...