# Fine-tune with Unsloth SFT

!!! note "Prerequisites"

    Please make sure you have read [Write the First Algorithm](./write-first-algorithm.md). Although that recipe is based on a simple prompt tuning algorithm, it introduces the core concepts of Agent-lightning and you should be familiar with them before proceeding.

This recipe builds on [Write the First Algorithm](./write-first-algorithm.md). Instead of iterating on a prompt, we will fine-tune a large language model with [Unsloth](https://docs.unsloth.ai/)'s SFT Trainer and keep the whole loop inside Agent-lightning. The new pieces you will meet are the **LLM proxy**, the **trace-to-triplet adapter**, a [vLLM](https://github.com/vllm-project/vllm) inference endpoint, and an agent implemented with the [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/). The full sample code is available in the [`examples/unsloth`]({{ src("examples/unsloth") }}) folder.

!!! warning

    You need a GPU that can host the Unsloth base model and run vLLM. The sample defaults to `unsloth/Qwen3-4B-Instruct-2507`, which requires at least 16GB of GPU memory under 4-bit quantization.

## The Data and Serving Loop

To tune a large language model in Supervised Fine-Tuning (SFT), we commonly need a dataset with input/output samples. For example, the [TRL SFT Trainer](https://huggingface.co/docs/trl/sft_trainer) expects a dataset with samples like the following:

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

With supervised fine-tuning, the LLM learns to generate the "assistant" response as close as possible to the completion in the dataset.

Typically, the dataset used in SFT should be a curated set of samples. The samples can be either hand-written by humans, or generated by a more powerful model, which is known as [data distillation](https://docs.nvidia.com/nemo-framework/user-guide/24.12/modelalignment/knowledge-distillation.html). However, in this recipe, we use a different setup that relies on samples generated by the model itself. We use the reward emitted by the agent to select the top-performing samples.

Overall, the flow of the algorithm is an iteration of the following steps:

1. Serve the current checkpoint (with vLLM).
2. Publish the vLLM endpoint through the LLM proxy and let runners roll out some tasks with the current model.
3. Collect the traces from the rollouts and transform the highest-rewarded ones into a dataset that is acceptable for Unsloth SFT Trainer.
4. Launch Unsloth to fine-tune on the dataset and save a new checkpoint.

You will find the full source code of this iteration in `sft_one_iter` in [sft_algorithm.py]({{ src("examples/unsloth/sft_algorithm.py") }}). We will elaborate on each part below.

### Serving the Model with vLLM and Proxy

Most modern agents do not use the model directly; instead, they use an API like the OpenAI chat completions API to interact with the model. Therefore, we need a vLLM-based inference server launched before rollouts. The serving code looks like the following. See the `vllm_server` function in [sft_algorithm.py]({{ src("examples/unsloth/sft_algorithm.py") }}) if you want to see a more robust version.

```python
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"http://localhost:{port}/health"
start = time.time()
client = httpx.Client()

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

server_address = f"http://localhost:{port}/v1"

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

In this recipe, we do not expose the server address directly to the agent runners, because we want to install a "middleware" to collect the prompts and responses of all the requests. In general, it's up to you to decide whether to hide the vLLM server behind a proxy or not.

The "middleware" here is [`LLMProxy`][agentlightning.LLMProxy], which is an independent [LiteLLM](https://docs.litellm.ai/) server that forwards the requests to the vLLM server. It also exposes an OpenAI-compatible API that the runners can target without caring about where the model lives. The benefits of using the proxy are:

1. **Traces:** The proxy automatically logs the prompts and responses of all the requests into the store.
2. **Token IDs:** The proxy augments the requests so that the vLLM server can return the prompt and response token IDs (see more details in [Serving LLM](../deep-dive/serving-llm.md)).

The [`LLMProxy`][agentlightning.LLMProxy] accepts a list of model configurations, in the same syntax as LiteLLM's [`model_list`](https://docs.litellm.ai/docs/proxy/configs). Include a `hosted_vllm/` prefix to the models to activate LiteLLM's [vLLM integration](https://docs.litellm.ai/docs/providers/vllm).

```python
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()})
```

### Spawn Rollout and Collect Spans

Once the proxy is registered as a resource, the algorithm schedules work for the rollout runners. Each problem from a training dataset becomes a rollout with the proxy baked into its resources:

```python
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` ties every rollout to the `main_llm` proxy resource we just uploaded. The runners on the other side poll the store ([`LitAgentRunner.iter()`][agentlightning.LitAgentRunner.iter]) and execute the agent for each rollout. On the algorithm side we wait for completions with a non-blocking polling loop:

```python
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)
```

!!! note

    The `timeout=0.0` is needed here because this example uses a [`LightningStoreClient`][agentlightning.LightningStoreClient], and `wait_for_rollouts` establishes an HTTP connection to that store. Currently, only non-blocking wait requests are supported, which avoids holding the store connection open.

Once the rollouts complete, we terminate the vLLM server to free up GPU memory.

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

### Adapt the Spans to HuggingFace Dataset

[`LlmProxyTraceToTriplet`][agentlightning.LlmProxyTraceToTriplet] converts the proxy’s spans (which might be dozens to hundreds per rollout) into [`Triplet`][agentlightning.Triplet] objects that contain prompt/response token IDs plus an optional reward. The adapter may return multiple triplets per rollout (one per chat-completion call). To bias training toward successful reasoning chains the algorithm walks the triplets in reverse order, keeps the most recent reward, and turns each prompt/response pair into Hugging Face dataset rows:

```python
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,
            }
        )
```

!!! note

    You might notice that the dataset format used here differs from the format described in the **SFT Trainer** documentation. According to the documentation, dataset samples should be provided as plain text strings or message objects.

    As a matter of fact, this example leverages some [undocumented behavior](https://github.com/huggingface/trl/blob/e0eec055b412c48ad754149c475a87a8fca34fb4/trl/trainer/sft_trainer.py#L887) in the SFT Trainer implementation. When the dataset already includes a `"input_ids"` column, the Trainer automatically marks it as `is_processed` and skips the internal tokenization step.

    Since we already have spans with token IDs generated by the [`LLMProxy`][agentlightning.LLMProxy], providing them directly avoids unnecessary [**re-tokenization** and related complications](../deep-dive/serving-llm.md). This approach will both save processing time and increase consistency between training and inference.

After aggregating every rollout we shuffle, sort by reward, and keep the top fraction (e.g., 50%) before shuffling again. The resulting list feeds directly into `datasets.Dataset.from_list`, which is the format Unsloth’s SFT trainer expects.

```python
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)
```

### Launch Unsloth Training

The heavy lifting happens in [`trl.SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer) (see [unsloth_helper.py]({{ src("examples/unsloth/unsloth_helper.py") }}) on how it's used). We launch it in a fresh process created with `multiprocessing.get_context("spawn")` so CUDA memory is reliably reclaimed when training ends. Launching it in the same process will also work for the first iteration, but we found that the memory won't be freed properly for subsequent vLLM serving.

```python
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)
```

Inside the `unsloth_training` subprocess, Unsloth loads the previous checkpoint in 4-bit, applies LoRA adapters, and forwards the Hugging Face dataset to [`trl.SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer) with the configuration defined in [`SFTConfig`](https://huggingface.co/docs/trl/sft_trainer#trl.SFTConfig) (batch size, accumulation steps, learning rate, etc.). The merged 16-bit weights are saved under `models/version_<iteration + 1>` so the next iteration can immediately serve them with vLLM.

```python
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")
```

## Math Agent: OpenAI Agents SDK with MCP

We build an agent with the [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/) to wire a calculator MCP tool and an OpenAI-compatible chat completion model together. The agent aims to solve a math problem and returns a reward indicating whether the answer is correct or not. The runner injects the `LLM` resource supplied by the algorithm side:

```python
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"])
```

!!! tip

    You can test the agent with a dry run:

    ```python
    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))
    ```

## Run this Recipe

The full runnable script for this recipe resides in [`examples/unsloth`]({{ src("examples/unsloth") }}) folder.

Before running this example, install `unsloth`, `vllm`, and the other libraries used in the examples (the project uses CUDA tooling, TRL, rich, datasets, etc.). We tested with `unsloth==2025.10.1`. `unsloth==2025.10.2` and `2025.10.3` are not working because of an [issue](https://github.com/unslothai/unsloth/issues/3451) we have been investigating with the unsloth team.

It's recommended to download the base model before running the example, such that the first iteration and subsequent iterations can both load from local checkpoints.

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

The repository already contains `examples/unsloth/data_gsmhard.jsonl` (which is a very small subset of the [GSM-hard math dataset](https://huggingface.co/datasets/reasoning-machines/gsm-hard) for demonstration purposes).

### Run Manually

Similar to the [Write the First Algorithm](./write-first-algorithm.md) recipe, you can open three terminals and start each component in parallel.

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

In this case, [`sft_rollout_runners.py`]({{ src("examples/unsloth/sft_rollout_runners.py") }}) is a simple spawner implemented in Python that spawns 4 runners in parallel. The runners all connect to the same store server executing in another terminal.

```python
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("http://localhost:4747")
spawn_runners(store=store, n_runners=4)
```

!!! tip

    Try to swap [`OtelTracer`][agentlightning.OtelTracer] in the runners with other tracers like [`AgentOpsTracer`][agentlightning.AgentOpsTracer]. Try to use a different adapter at the algorithm side such as [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] to see what happens.

### Run Everything with Trainer

We also show how to wrap everything into a single script using [`Trainer`][agentlightning.Trainer]. [`sft_allinone.py`]({{ src("examples/unsloth/sft_allinone.py") }}) wires the same components together, replacing the manual management of runners above.

```python
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())
```

You might wonder where the initialization of [`Adapter`][agentlightning.Adapter] happens in this code. It turns out that [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] is the default adapter in [`Trainer`][agentlightning.Trainer], so we don't need to create one manually.

Now you can run the example with:

```bash
python examples/unsloth/sft_allinone.py
```

It starts an [`InMemoryLighningStore`][agentlightning.InMemoryLightningStore] for you, launches four worker processes, iterates the SFT loop, and prints the final checkpoint path when done. Adjust `max_iterations`, `train_triplet_fraction`, `n_runners`, or the proxy port to match your hardware or training goals. If you already run an external store or proxy you can also pass those objects into [`Trainer`][agentlightning.Trainer] instead of relying on the [Trainer-managed defaults][debug-with-external-store].

!!! info

    As a future plan, we might graduate this example into a more powerful SFT algorithm bundled into [Algorithm Zoo](../algorithm-zoo/index.md). Currently, this `UnslothSupervisedFinetuning` is still for demo purposes.
