跳到内容

使用 Unsloth SFT 进行微调

先决条件

请确保您已阅读 编写第一个算法。虽然该 Recipe 基于一个简单的提示调整算法,但它介绍了 Agent-lightning 的核心概念,您应该在继续之前熟悉它们。

此 Recipe 基于 编写第一个算法。我们不会迭代提示,而是使用 Unsloth 的 SFT Trainer 微调大型语言模型,并将整个循环保留在 Agent-lightning 内部。您将遇到的新内容是 LLM 代理trace-to-triplet 适配器vLLM 推理端点,以及使用 OpenAI Agents SDK 实现的 Agent。完整的示例代码可在 examples/unsloth 文件夹中找到。

警告

您需要一个可以托管 Unsloth 基础模型并运行 vLLM 的 GPU。示例默认使用 unsloth/Qwen3-4B-Instruct-2507,在 4 位量化下至少需要 16GB 的 GPU 内存。

数据与服务循环

为了在监督微调 (SFT) 中微调大型语言模型,我们通常需要一个包含输入/输出样本的数据集。例如,TRL SFT Trainer 期望一个包含如下样本的数据集

{"messages": [{"role": "user", "content": "What color is the sky?"},
              {"role": "assistant", "content": "It is blue."}]}

通过监督微调,LLM 学习生成与数据集中 completion 尽可能接近的“助手”响应。

通常,用于 SFT 的数据集应该是一个策划的样本集合。这些样本可以由人类手动编写,也可以由更强大的模型生成,这被称为 数据蒸馏。但是,在本 Recipe 中,我们使用不同的设置,该设置依赖于模型自身生成的样本。我们使用 Agent 发出的奖励来选择表现最佳的样本。

总而言之,该算法的流程是以下步骤的迭代

  1. 服务当前检查点(使用 vLLM)。
  2. 通过 LLM 代理发布 vLLM 端点,并让 runners 使用当前模型执行一些任务。
  3. 收集 rollout 中的 traces,并将最高奖励的 traces 转换为 Unsloth SFT Trainer 可接受的数据集。
  4. 启动 Unsloth,使用数据集进行微调,并保存新的检查点。

您可以在 sft_algorithm.py 中的 sft_one_iter 中找到此迭代的完整源代码。我们将在下面详细介绍每个部分。

使用 vLLM 和代理服务模型

大多数现代 Agent 不直接使用模型;而是使用 OpenAI chat completions API 等 API 与模型交互。因此,在 rollout 之前,我们需要启动一个基于 vLLM 的推理服务器。服务代码如下所示。如果您想查看更健壮的版本,请参阅 sft_algorithm.py 中的 vllm_server 函数。

from openai import OpenAI

vllm_process = subprocess.Popen([
    "vllm", "serve", model_path, "--port", str(port),
    "--enable-auto-tool-choice", "--tool-call-parser", "hermes"
])

# Wait for the server to be ready
url = f"https://:{port}/health"
start = time.time()
client = httpx.Client()

while True:
    if client.get(url).status_code == 200:
        break

server_address = f"https://:{port}/v1"

# Try using the vLLM server
openai = OpenAI(base_url=server_address)
...

在本 Recipe 中,我们不会直接将服务器地址暴露给 Agent runners,因为我们希望安装一个“中间件”来收集所有请求的 prompts 和 responses。通常,是否将 vLLM 服务器隐藏在代理后面由您决定。

这里的“中间件”是 LLMProxy,它是一个独立的 LiteLLM 服务器,它将请求转发到 vLLM 服务器。它还暴露了一个 OpenAI 兼容的 API,runners 可以针对该 API,而无需关心模型的位置。使用代理的好处是

  1. Traces: 代理会自动将所有请求的 prompts 和 responses 记录到 store 中。
  2. Token IDs: 代理增强请求,以便 vLLM 服务器可以返回 prompt 和 response token IDs(有关更多详细信息,请参阅 Serving LLM)。

LLMProxy 接受模型配置列表,语法与 LiteLLM 的 model_list 相同。在模型前包含 hosted_vllm/ 前缀以激活 LiteLLM 的 vLLM 集成

import agentlightning as agl

llm_proxy = agl.LLMProxy(port=port, store=store)
model_list = [
    {
        "model_name": "Qwen3-4B-Instruct",
        "litellm_params": {"model": f"hosted_vllm/{model_path}", "api_base": server_address},
    }
]
llm_proxy.update_model_list(model_list)
# If the proxy is not running, it will start automatically.
await llm_proxy.restart()
# Add the proxy as a resource to the store so that the runners can access it via URL.
resource_update = await store.add_resources({"main_llm": llm_proxy.as_resource()})

生成 Rollout 并收集 Spans

注册代理作为资源后,算法为 rollout runners 安排工作。训练数据集中的每个问题都成为一个带有代理烘焙到其资源中的 rollout

rollouts: list[Rollout] = []
for sample in train_dataset:
    rollouts.append(
        await store.enqueue_rollout(
            input=sample,
            mode="train",
            resources_id=resources_update.resources_id,
        )
    )

resources_id 将每个 rollout 绑定到我们刚刚上传的 main_llm 代理资源。另一侧的 runners 轮询 store (LitAgentRunner.iter()) 并为每个 rollout 执行 Agent。在算法侧,我们使用非阻塞轮询循环等待完成

completed_rollouts: list[Rollout] = []
while True:
    completed_rollouts = await store.wait_for_rollouts(
        rollout_ids=[r.rollout_id for r in rollouts],
        timeout=0.0,
    )
    if len(completed_rollouts) == len(rollouts):
        break
    await asyncio.sleep(5.0)

注意

这里需要 timeout=0.0,因为此示例使用 LightningStoreClient,并且 wait_for_rollouts 建立到该 store 的 HTTP 连接。目前,仅支持非阻塞等待请求,这避免了保持 store 连接打开。

Rollout 完成后,我们终止 vLLM 服务器以释放 GPU 内存。

vllm_process.terminate()
vllm_process.join(timeout=10.0)

将 Spans 适配为 HuggingFace 数据集

LlmProxyTraceToTriplet 将代理的 spans(每个 rollout 可能有数十到数百个)转换为 Triplet 对象,其中包含 prompt/response token IDs 以及可选的奖励。适配器可以为每个 rollout 返回多个 triplets(每个 chat-completion 调用一个)。为了偏向于成功的推理链,算法会反向遍历 triplets,保留最新的奖励,并将每对 prompt/response 转换为 Hugging Face 数据集行

all_triplets = []
data_adapter = agl.LlmProxyTraceToTriplet()

for rollout in completed_rollouts:
    spans = await store.query_spans(rollout.rollout_id, "latest")
    triplets = data_adapter.adapt(spans)

    recent_reward = None
    for triplet in reversed(triplets):
        if triplet.reward is not None:
            recent_reward = triplet.reward
        if recent_reward is None:
            continue

        input_ids = triplet.prompt["token_ids"] + triplet.response["token_ids"]
        # We don't train on prompt tokens, so they are masked out by setting to -100.
        labels = [-100] * len(triplet.prompt["token_ids"]) + triplet.response["token_ids"]
        # This matches the dataset format required by the Unsloth SFT trainer.
        all_triplets.append(
            {
                "input_ids": input_ids,
                "attention_mask": [1] * len(input_ids),
                "labels": labels,
                "reward": recent_reward,
            }
        )

注意

您可能会注意到,这里使用的数据集格式与 SFT Trainer 文档中描述的格式不同。根据文档,数据集样本应作为纯文本字符串或消息对象提供。

事实上,此示例利用了 SFT Trainer 实现中的一些 未记录的行为。当数据集已经包含 "input_ids" 列时,Trainer 会自动将其标记为 is_processed 并跳过内部 tokenization 步骤。

由于我们已经使用 LLMProxy 生成了带有 token IDs 的 spans,因此直接提供它们可以避免不必要的 重新 tokenization 和相关复杂性。这种方法既可以节省处理时间,又可以提高训练和推理之间的一致性。

在聚合每个 rollout 后,我们会 shuffle、按奖励排序,并保留前一部分(例如,50%),然后再 shuffle。生成的列表直接输入 datasets.Dataset.from_list,这是 Unsloth 的 SFT trainer 期望的格式。

from datasets import Dataset as HuggingFaceDataset

random.shuffle(all_triplets)
all_triplets.sort(key=lambda x: x["reward"], reverse=True)
sliced_triplets = all_triplets[: max(1, int(len(all_triplets) * triplet_fraction))]
# Shuffle the sliced triplets again
random.shuffle(sliced_triplets)

sft_dataset = HuggingFaceDataset.from_list(sliced_triplets)

启动 Unsloth 训练

繁重的计算发生在 trl.SFTTrainer (参见 unsloth_helper.py 了解其使用方法)。我们使用 multiprocessing.get_context("spawn") 创建的新进程启动它,以便在训练结束后可靠地回收 CUDA 内存。在同一进程中启动它对于第一次迭代也是可行的,但我们发现后续 vLLM 服务无法正确释放内存。

context = multiprocessing.get_context("spawn")
unsloth_process = context.Process(
    target=unsloth_training,
    args=(model_path, sft_dataset, next_model_path),
    daemon=True,
)
unsloth_process.start()
unsloth_process.join(timeout=600.0)

unsloth_training 子进程中,Unsloth 以 4 位加载先前的检查点,应用 LoRA 适配器,并将 Hugging Face 数据集转发到 trl.SFTTrainer,配置定义在 SFTConfig (批大小、累积步数、学习率等) 中。合并后的 16 位权重保存在 models/version_<iteration + 1> 下,以便下一次迭代可以立即使用 vLLM 服务它们。

from unsloth import FastLanguageModel
# TRL is patched by unsloth.
from trl import SFTConfig, SFTTrainer

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_path,
    load_in_4bit=True,  # 4 bit quantization to reduce memory
)

# Config the model to use LoRA
model = FastLanguageModel.get_peft_model(
    model,
    r=32,
    ...
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=sft_dataset,
    ...
)

# This is the heaviest step.
trainer_stats = trainer.train()

# Save in 16-bit for vLLM inference later
model.save_pretrained_merged(next_model_path, tokenizer, save_method="merged_16bit")

数学代理:OpenAI Agents SDK 与 MCP

我们使用 OpenAI Agents SDK 构建一个代理,将计算器 MCP 工具和兼容 OpenAI 的聊天完成模型连接起来。该代理旨在解决一个数学问题,并返回一个奖励,指示答案是否正确。运行器注入算法侧提供的 LLM 资源

import os
from typing import TypedDict

import agentlightning as agl
from agents import Agent, ModelSettings, OpenAIChatCompletionsModel, Runner as OpenAIRunner
from agents.mcp import MCPServerStdio
from openai import AsyncOpenAI

class GsmProblem(TypedDict):
    input: str
    target: float

def compute_reward(result: str, target: float) -> float:
    ...

@agl.rollout
async def math_agent(task: GsmProblem, llm: agl.LLM) -> float:
    async with MCPServerStdio(
        name="Calculator via uvx",
        params={"command": "uvx", "args": ["mcp-server-calculator"]},
    ) as server:
        agent = Agent(
            name="Assistant",
            instructions=(
                "Use the calculator tool for every question. "
                "Return only the numeric answer wrapped like ### <answer> ###."
            ),
            mcp_servers=[server],
            model=OpenAIChatCompletionsModel(
                model=llm.model,
                openai_client=AsyncOpenAI(
                    base_url=llm.endpoint,
                    api_key=llm.api_key or "dummy",
                ),
            ),
            model_settings=ModelSettings(
                temperature=llm.sampling_parameters.get("temperature", 0.0),
            ),
        )
        result = await OpenAIRunner.run(agent, task["input"])
    return compute_reward(result.final_output, task["target"])

提示

你可以使用模拟运行来测试代理

import asyncio

llm = agl.LLM(
    endpoint=os.environ["OPENAI_BASE_URL"],
    api_key=os.environ["OPENAI_API_KEY"],
    model="gpt-4.1-mini",
)
asyncio.run(math_agent({"input": "What is 1 + 1?", "target": 2.0}, llm))

运行此示例

此示例的可运行脚本位于 examples/unsloth 文件夹中。

在运行此示例之前,请安装 unslothvllm 和示例中使用的其他库(该项目使用 CUDA 工具、TRL、rich、datasets 等)。我们使用 unsloth==2025.10.1 进行测试。unsloth==2025.10.22025.10.3 无法正常工作,因为我们正在与 unsloth 团队调查一个 问题

建议在运行示例之前下载基础模型,以便第一次迭代和后续迭代都可以从本地检查点加载。

hf download unsloth/Qwen3-4B-Instruct-2507 --local-dir models/version_0

该存储库已经包含 examples/unsloth/data_gsmhard.jsonl(这是 GSM-hard 数学数据集的一个非常小的子集,仅用于演示目的)。

手动运行

编写第一个算法 示例类似,你可以打开三个终端并并行启动每个组件。

agl store --port 4747
python examples/unsloth/sft_rollout_runners.py
python examples/unsloth/sft_algorithm.py

在这种情况下,sft_rollout_runners.py 是一个用 Python 实现的简单生成器,并行生成 4 个运行器。所有运行器都连接到另一个终端中执行的相同存储服务器。

import agentlightning as agl

def run_rollout(store: agl.LightningStore, worker_id: int) -> None:
    # Since the server side has already used LiteLLM proxy to collect traces,
    # a simple OtelTracer to collect the rewards is enough.
    tracer = agl.OtelTracer()

    runner = agl.LitAgentRunner(tracer=tracer)

    with runner.run_context(agent=math_agent, store=store, worker_id=worker_id):
        asyncio.run(runner.iter())


def spawn_runners(store: agl.LightningStore, n_runners: int) -> None:
    runners = [
        multiprocessing.Process(target=run_rollout, args=(store, worker_id))
        for worker_id in range(n_runners)
    ]
    for runner in runners:
        runner.start()

    for runner in runners:
        runner.join()


store = agl.LightningStoreClient("https://:4747")
spawn_runners(store=store, n_runners=4)

提示

尝试将运行器中的 OtelTracer 与其他跟踪器(如 AgentOpsTracer)互换。尝试在算法侧使用不同的适配器,例如 TracerTraceToTriplet,看看会发生什么。

使用 Trainer 运行所有内容

我们还展示了如何使用 Trainer 将所有内容包装到一个脚本中。sft_allinone.py 将相同的组件连接在一起,取代了上面对运行器的手动管理。

class UnslothSupervisedFinetuning(agl.Algorithm):

    async def run(
        self,
        train_dataset: Optional[Dataset[GsmProblem]] = None,
        val_dataset: Optional[Dataset[GsmProblem]] = None,
    ):
        # Use the store, llm_proxy, and adapter from the trainer
        store = self.get_store()
        llm_proxy = self.get_llm_proxy()
        data_adapter = self.get_adapter()

        for iteration in range(self.max_iterations):
            ...  # Same logic as sft_algorithm.py

algo = UnslothSupervisedFinetuning(
    max_iterations=2,
    vllm_port=12316,
    train_triplet_fraction=0.5,
    initial_model_path="models/version_0",
)

# The LLM proxy can be created before Trainer
trainer = Trainer(
    n_runners=4,
    algorithm=algo,
    llm_proxy=LLMProxy(port=12358),
)

trainer.fit(math_agent, load_math_dataset())

你可能想知道 Adapter 在此代码中是如何初始化的。事实证明,TracerTraceToTripletTrainer 中的默认适配器,因此我们不需要手动创建一个。

现在你可以使用以下命令运行示例

python examples/unsloth/sft_allinone.py

它会为你启动一个 InMemoryLighningStore,启动四个工作进程,迭代 SFT 循环,并在完成后打印最终检查点路径。调整 max_iterationstrain_triplet_fractionn_runners 或代理端口以匹配你的硬件或训练目标。如果你已经运行了外部存储或代理,你也可以将这些对象传递给 Trainer,而不是依赖 Trainer 管理的默认值

信息

作为未来的计划,我们可能会将此示例升级到打包到 Algorithm Zoo 中的更强大的 SFT 算法。目前,此 UnslothSupervisedFinetuning 仍然仅用于演示目的。