目录

什么是配方?

本深入讲解将引导您了解 torchtune 中训练配方(training-recipes)的设计。

本深入探讨将涵盖以下内容
  • 什么是示例?

  • 构成一个示例的核心组件有哪些?

  • 我应该如何构建新示例?

示例是 torchtune 用户的主要入口。这些可以被视为针对训练(并可选择性地评估)大语言模型(LLM)的“定向”端到端流程。每个示例实现一种训练方法(例如:全量微调),并结合一组有意义的功能(例如: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 类,训练所需的核心逻辑,通过一组 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
...

Recipe Class

Recipe 类承载了模型训练的核心逻辑。每个类都实现了相应的接口,并暴露了一组 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()
    ...

运行带有配置的配方

要使用一组用户定义的参数运行配方,你需要编写一个配置文件。 你可以在我们的 配置深入解析 中了解有关配置的所有内容。

配置和CLI解析使用 parse

我们提供了一个方便的装饰器 parse(),它将您的配方包装起来,以支持通过命令行使用 tune 运行,并解析配置和 CLI 覆盖。

@config.parse
def recipe_main(cfg: DictConfig) -> None:
    recipe = FullFinetuneRecipe(cfg=cfg)
    recipe.setup(cfg=cfg)
    recipe.train()
    recipe.cleanup()

运行您的配方

您可以使用 tune 命令并提供自定义配方和配置的直接路径来运行您的配方,同时可以添加任何命令行覆盖选项:

tune run <path/to/recipe> --config <path/to/config> k1=v1 k2=v2 ...

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源