TorchTune:PyTorch的微调库

在各种应用中,大型语言模型LLMs)的使用越来越广泛。然而,针对特定任务对这些强大的模型进行微调可能是一个复杂且资源密集的任务。TorchTune,一个全新的PyTorch库,直面这一挑战,提供了一个直观且可扩展的解决方案。PyTorch发布了alpha版本的TorchTune,这是一个PyTorch原生库,用于轻松微调大型语言模型。根据PyTorch的设计原则,它提供了可组合和模块化的构建块,以及易于扩展的训练配方,用于微调各种消费级和专业GPU上的大型语言技术,如LORA和QLORA。

目录

  • 为什么使用TorchTune?
  • TorchTune工作流程
  • 微调配方
  • TorchTune的设计是什么?
  • 与LLM的集成
  • 快速开始

为什么使用TorchTune?

在过去的一年里,对开放的大型语言模型(LLMs)的兴趣激增。对这些尖端模型进行微调以适应特定应用已成为一项关键技术。然而,这种适应过程可能很复杂,需要在各个阶段进行广泛的定制,包括数据和模型选择、量化、评估和推理。此外,这些模型的庞大规模在资源受限的消费级GPU上进行微调时提出了重大挑战。

当前的解决方案往往通过抽象层次隐藏关键组件,阻碍了定制和优化。这种缺乏透明度使得理解不同元素如何相互作用以及需要修改哪些元素以实现所需功能变得困难。TorchTune通过赋予开发者对整个微调过程的细粒度控制和可见性来解决这一挑战,使他们能够根据特定要求和约束定制LLMs

TorchTune工作流程

TorchTune支持以下微调工作流程:下载和准备数据集和模型检查点;使用支持不同模型架构、参数高效微调(PEFT)技术和更多的可组合构建块自定义训练;记录进度和指标以洞察训练过程;微调后量化模型;在流行的基准测试上评估微调模型;运行本地推理以测试微调模型;与流行的生产推理系统兼容的检查点。

微调配方

TorchTune提供以下微调配方。训练微调方法包括分布式训练(1至8个GPU)、单设备/低内存(1个GPU)和单设备(1个GPU)。所有配方都在各种设置上进行了测试,包括具有24GB VRAM的商品GPU以及数据中心中找到的更强大的选项。

单GPU配方公开了一些在分布式版本中不可用的内存优化,包括支持来自bitsandbytes的低精度优化器,以及融合优化器步骤与后向步骤以减少梯度的内存占用(参见示例配置)。对于内存受限的设置,建议使用单设备配置作为起点。例如,默认QLORA配置的峰值内存使用量约为9.3GB。同样,在单设备上使用batch_size=2的LoRA,峰值内存使用量约为17.1GB。这两个都使用dtype=bf16和AdamW作为优化器。

TorchTune的设计是什么?

TorchTune的设计是可扩展的:承认微调技术的快速发展和多样化的用户需求,TorchTune优先考虑易于扩展性。其配方利用模块化组件和易于修改的训练循环。最小化的抽象确保了用户对微调过程的控制。每个配方都是自包含的(少于600行代码!)并且不需要外部训练器或框架,进一步促进了透明度和定制。

民主化微调:TorchTune通过迎合不同专业水平的用户,促进包容性。其直观的配置文件易于修改,允许用户在没有广泛的编码知识的情况下自定义设置。此外,内存高效的配方使得在现成的消费级GPU(例如,24GB)上进行微调,消除了对昂贵数据中心硬件的需求。

开源生态系统集成:认识到充满活力的开源LLM生态系统,PyTorch的TorchTune优先考虑与各种工具和资源的互操作性。这种灵活性为用户提供了对微调过程和模型部署的更大控制权。

与LLM的集成

  • Hugging Face Hub:利用Hugging Face Hub上可用的大量开源模型和数据集进行微调。通过tunedownload CLI命令流线化集成,便于立即启动微调任务。
  • PyTorchFSDP:通过利用PyTorch FSDP的能力,实现分布式训练。这迎合了利用多GPU设置的增长趋势,通常包括像NVIDIA的3090/4090系列这样的消费级卡。TorchTune提供由FSDP支持的分布式训练配方,以利用这些硬件配置。
  • Weights & Biases:与Weights & Biases AI平台集成,全面记录训练指标和模型检查点。这集中了配置细节、性能指标和模型版本,方便监控和分析微调运行。
  • EleutherAI的LM评估套件:认识到模型评估的关键作用,TorchTune包括一个由EleutherAI的LM评估套件支持的流线化评估配方。这为用户提供了直接访问一系列既定LLM基准的途径。为了进一步增强评估体验,计划在未来几个月与EleutherAI密切合作,建立更深入、更原生的集成。
  • ExecuTorch:通过促进向ExecuTorch的无缝导出,实现在各种移动和边缘设备上高效推理微调模型。
  • torchao:提供由torchao的量化API支持的简单后训练配方,使微调模型能够高效地转换为低精度格式(例如,4位或8位),以减少内存占用并加快推理速度。
  1. 下载模型:按照官方仓库的说明确保有权访问Llama2模型权重。一旦确认了访问权限,可以运行以下命令将权重下载到本地机器。这也会下载分词器模型和负责任的使用指南。
  2. 运行微调配方:Llama2 7B + LoRA在单GPU上使用以下命令运行微调配方。对于分布式训练,tune CLI与集成。Llama2 7B + LoRA在两个GPU上使用以下命令运行微调配方。确保将任何torchrun命令放在配方规格之前。在此之后的任何CLI参数将覆盖配置,并不会影响分布式训练。
  3. 修改配置:可以通过两种方式修改配置:配置覆盖和更新本地副本。
沪ICP备2024098111号-1
上海秋旦网络科技中心:上海市奉贤区金大公路8218号1幢 联系电话:17898875485