通过 NVIDIA DGX Spark 进行全参数微调，解锁 Llama 3.1-8B 的推理能力

本地运行大语言模型（LLM）的隐秘乐趣是什么？

下载模型权重，通宵进行各种实验，房间里可能还会变得暖烘烘的，然后“砰”的一声，你就亲手打造出了一个能在台式机上运行、性能卓越的小型模型，这种魔力简直令人着迷。

通常这需要庞大的 GPU 机器和成堆的线缆；但在我们这里，是一个非常可爱的小盒子，刚好能塞进显示器支架下的空间，还能让我们的手保持温暖。说真的，DGX Spark 看起来非常有趣！

在本篇博客中，我们将分享一套秘籍，利用 DGX Spark 小盒子在合成数据上对 Llama 3.1-8B-Instruct 进行全参数微调，从而解锁大语言模型的“推理”能力。得益于统一内存架构，我们能够完全在本地生成合成思维链，并以此对模型进行微调。

红色文字显示了经过微调后的模型在 NVIDIA DGX Spark 上新增的行为表现。

整套方案在 DGX Spark 上离线运行，耗时不到一天。我们成功实现了 Llama-3.1-8B-Instruct 的全参数微调，且没有任何问题，上下文长度为 16k token，批大小（batch size）为 16。

我们计划在未来的博客中分享关于 Llama 70B 的更多实验，以及更令人兴奋的 FP4 实验主题。敬请期待！

为 Llama-3.1-8B 增加推理行为。得益于推理时间的扩展，大语言模型的推理和思考能力在实践中展现出了巨大的提升。

我们提出了一个问题：我们能否通过针对特定主题的合成思维链微调，来创造出这种推理行为？

我们提示在本地运行的 Llama 3.3-70B-Instruct，将思维链（Chain of Thought）添加到现有的对话中。

数据生成

注意：我们对整个 ToolACE 数据集生成了合成思维链（CoT），该数据集包含 1.1 万个对话对。

此功能可通过 Synthetic-Data-Kit 开箱即用，它提供了直观的命令行界面（CLI），用于准备和丰富您的数据集以进行 LLM 微调。

我们可以在 DGX Spark 上使用 vLLM 在本地运行此功能，并使用以下方法生成思维链响应。

我们使用一条简单的命令。

synthetic-data-kit create --type cot-enhance /path/to/dataset

然后我们可以在配置文件中创建一个自定义提示，并按如下方式使用它。

# cot_tools_config.yaml

vllm:
  api_base: "https://:8000/v1"
  model: "unsloth/Meta-Llama-3.3-70B-Instruct"
  max_retries: 3
  retry_delay: 1.0

generation:
  temperature: 0.2   # Lower temperature for more consistent reasoning
  top_p: 0.95
  max_tokens: 16384   # Allow for longer outputs to accommodate CoT reasoning

# The most important part - our custom Chain of Thought prompt
prompts:
  cot_enhancement: |
    You are a highly intelligent AI with an IQ of 170, and your job is to
enhance existing conversation examples. Remember to return the entire
conversation as is, BUT

BUT, we will add Chain of Thought and planning to "Assistant" messages
whenever they return a tool call.

    Remember, ONLY when an assistant message returns a tool call will we add
thinking and reasoning traces before it to add logic. Otherwise, we don't touch the conversation history.

    Remember to return the entire message, but only enhance the assistant
messages whenever a tool is called in the conversation by adding thoughts.

    Please keep in mind that we are not modifying anything in the example, nor
are we changing what it does. We are only adding CoT every time a tool gets
called in the conversation.

    Think out loud and maximize your tokens when adding CoT.

    For example, if you see:

    "from": "assistant",
    "value": "<tool>[Some API(param=\"value\")]</tool>"

    Change it to:

    "from": "assistant",
    "value": "Let me think about this request. I need to gather X information
using Tool Y.
    To do this, I need to set the parameter to 'value' because of reason Z.
    <tool>[Some API(param=\"value\")]</tool>"

    BEGIN WORK NOW. Enhance the assistant's messages with detailed Chain of
Thought reasoning before each tool call:
    {conversations}

注意：我们对整个 ToolACE 数据集生成了合成思维链，该数据集包含 1.1 万个对话对。

synthetic-data-kit -c cot_tools_config.yaml create 
test_files/conversation_example.json \
  --type cot-enhance \
  –output enhanced_output/

利用 NVIDIA DGX Spark 解锁本地全参数微调

大语言模型的微调已经非常成熟，在进行监督微调（SFT）时，我们可以调整很多参数。

在我们的实验中，我们采用全参数微调（Full-Fine Tuning）来展示 NVIDIA DGX Spark 128GB 统一内存的强大威力。

128GB 统一内存空间充裕！

DGX Spark 最棒的一点是：训练所需的所有可用内存都以 128GB 统一接口的形式呈现。因此，在进行监督微调时，我们可以假设拥有一个 128GB 的内存设备，而不必把时间浪费在处理卸载（offloading）设置上。

这使我们能够直接运行 Llama-3.1-8B 的全参数微调，而无需为了将一切塞入有限的设备内存而不断进行配置实验。

上下文长度

更大的内存使我们能够在更长的上下文中进行实验，从而更容易教授诸如工具调用（tool-calling）之类的任务。

对于我们的用例，我们在合成推理轨迹上对 Llama 进行微调。这些轨迹可能会非常长！在我们的案例中，实验运行在 16k token 的长度下。

随着 LLM 序列长度的增加，内存需求会呈二次方增长。

这是 DGX Spark 大放异彩的另一个领域：即使在 16k 上下文长度下进行全参数微调，我们的内存峰值使用率也大约只有 80%（见下文结果部分中的设备使用图）。

批大小（Batch Size）

另一个黄金法则是在 2 的幂次方范围内最大化批大小，以实现更快的 LLM 收敛。在我们的实验中，我们有足够的空间将 batch_size 设置为 16——这真是太棒了！

为什么要进行全参数微调？

得益于 128GB 的内存，我们能够完全在 DGX Spark 上运行 8B 的全参数微调，因此我们决定采用这种路径以获得最大性能，并在下文报告了我们的实验结果。

全参数微调

我们使用 TorchTune 进行全参数微调实验。Torchtune 提供了一个直观的命令行界面，可以与配置文件交互，实现单行命令运行实验。

完整的配置和详细信息可在此处获取。

我们在下面解释了关键配置。

tune run full_finetune_single_device --config fft-8b.yaml

我们在配置中设置的关键项：

序列长度（Seq_len）: 16384
批大小（Batch_size）: 16
迭代次数（Epochs）: 3
数据类型（Dtype）: bf16

整个微调流水线（从数据生成到全参数微调）仅需一天时间，每个 Epoch 平均耗时仅 8 小时左右。这真的非常令人印象深刻！

注意：考虑到整个运行过程中内存峰值使用率约为 80%，你甚至可以进一步压榨出更多性能。

结果

下面我们展示了核心训练运行图，以及在函数调用基准测试中的性能表现。

对于我们的方案，我们使用 ToolACE 数据集，并以 BFCLv3 以及 v4（近期发布）作为基准结果。

BFCL 衡量以下指标：

多轮工具调用（Multi-Turn Tool Calling）
支持并行工具调用
执行代理调用的能力（v4）
基于开放式查询进行网页搜索的能力（v4）

结论

我们展示了一套在 DGX Spark 上运行的端到端方案，利用其统一内存对 Llama-3.1-8B-Instruct 进行了全参数微调。

我们生成了合成思维和思维链，随后对模型进行了微调，以提升其在上述基准测试中的性能。

通过在 NVIDIA DGX Spark 上进行完全微调来解锁 Llama 3.1-8B 的推理能力