Axolotl 为微调大语言模型的开发者提供了单台服务器的训练配置。扩散语言模型可在不使用单独插件的情况下训练,也可使用包含图像的数据学习响应偏好和奖励。

1. 在配备 8 张卡的服务器上微调超大模型

Reflection 的 Beam 是一个混合专家模型,总参数量为 5010 亿,每个 token 激活 230 亿参数。其结构根据输入仅让部分专家参与计算,面向软件开发及长期执行多个步骤的智能体任务。

此次发布支持在配备 8 张 H200 的单节点上进行 QLoRA 微调。这种方式以低精度存储原始权重,并训练小规模的适配权重,并不意味着重新训练整个模型或在低配置设备上运行。

配置提供了分散训练状态的 FSDP2 方式,以及将模型层分配到不同设备的 device_map 方式。分布式方式与微调对象应分开看待。

注意力是决定输入的哪些部分相互参考的计算。两种分布式方式均提供微调注意力和专家权重的配置。device_map 还提供仅微调注意力的配置。

device_map 配置以单进程运行。当适配权重跨多个设备时,负责更新权重的 adamw_torch_4bit 优化器会失败,因此必须使用文档指定的 adamw_torch。

负责选择路径的路由器必须排除在微调训练对象之外。文档警告,训练路由器会导致专家选择不稳定。

2. 扩散模型纳入基础训练路径

扩散语言模型通过学习还原被遮蔽或混入噪声的 token 来训练。使用新增的顶层配置 diffusion_lm 和现有对话数据集后,各模型的配置档案会确定噪声过程、注意力布局和损失。

根据扩散模型指南,支持范围包括 Nemotron-Labs-Diffusion 的 30 亿和 80 亿参数模型,以及 Dream v0。

DiffusionGemma 也受支持,但这些模型的基础训练仅限于训练小规模适配权重的 LoRA 方式。

Nemotron 支持的变长注意力,会仅汇集不同长度文档中的实际 token 进行计算。该路径需要 Torch 2.14 或更高版本。文档的安装部分给出了如下环境安装示例。

python -m venv .venv-diffusion
source .venv-diffusion/bin/activate
python -m pip install --upgrade pip
python -m pip install torch==2.14.0  # 또는 가속기에 맞는 휠
python -m pip install -e .

30 亿参数模型的入门示例使用 32 条合成对话执行两次优化器更新。其用途是确认数据准备、训练和保存流程,而非评估质量。

将多篇文档合并到一行的打包操作主要用于节省内存。文档指出,吞吐量取决于文档长度和注意力实现,因此建议直接测量开启与关闭打包时的表现。

3. 图像响应也可应用偏好与奖励

根据发布说明,现在可在图像数据集上运行 DPO 和 KTO。这为包含图像的输入提供了按偏好调整响应的训练路径。

也支持 GRPO 这类利用生成响应奖励的强化学习。不过,尚未给出图像训练的质量提升数值。

若要配合使用图像强化学习的生成服务器,还需使用 vLLM 0.30.0 或更高版本。

扩散模型的图像输入采用独立路径。Nemotron-Labs-Diffusion-VLM-8B 仅通过专用插件接收图像,该插件训练模型在指定答案栏中填写标签或分数等内容。常规监督微调仅接收文本,并拒绝多模态数据集。

4. 收益在于整合,成本在于检查运行环境

基础扩散训练以顶层配置和日常使用的对话数据集取代原先的插件指定方式,减少了单独配置插件的负担。旧版扩散配置会在发出警告后转换,因此迁移时仍需检查警告及各模型的支持范围。

B200 的性能说明属于另一项比较。在设备间交换数据时,使用无需额外安装的默认通信而非 DeepEP 的路径,在内部小规模测试中慢了约 25%。

计算输入间引用关系的实现也存在差异。同一说明将 FlashAttention 2 描述为在该模型规模下比 FlashAttention 4 慢约 2 倍的替代方案。这两个数值比较的是该设备上的实现选择,并非对前述服务器配置的速度或成本估算。

本版本的价值在于将适用范围扩展至大模型的适配训练及扩散、图像训练,并减轻基础扩散训练中单独配置插件的负担。