Supervised fine-tuning

Your first SFT run

Train a small, inspectable task from input–output examples, then generate with the trained model and save its weights. Read Learning from examples for the meaning of the loss and token masks.

Before running, complete authentication, confirm access to Qwen/Qwen3.6-35B-A3B-FP8, and install the dependencies:

pip install river-client transformers

Supervised fine-tuning (SFT)

This example verifies a complete training loop with a small, inspectable task: prefix test- to every word of the input (e.g. "hello world" → "test-hello test-world"). The model can learn this pattern from a small dataset in about 15 steps. In your application, replace these pairs with examples of the behavior you want the model to learn.

1. Build a tiny dataset

Each example is a (prompt, completion) pair. We tokenize both and concatenate them into one input_ids sequence. Per-token weights mask the prompt with 0.0 so the loss is taken only on the completion, and target_tokens is the next-token target at each position (position i predicts ids[i+1], so the mask is offset by one).

import os
import river_client as river
from transformers import AutoTokenizer

client = river.Client(api_key=os.environ["RIVER_API_KEY"])
BASE = "Qwen/Qwen3.6-35B-A3B-FP8"
tok = AutoTokenizer.from_pretrained(BASE)
EOS = tok.eos_token_id

def target_for(x):
    return " ".join("test-" + w for w in x.split())

train_inputs = [
    "hello world",
    "good morning sunshine",
    "the quick brown fox",
    "i love programming",
    "open the pod bay doors",
    "river flows to the sea",
]

def render(x):
    return f"Input: {x}\nOutput:"

def make_datum(x):
    prompt_ids = tok(render(x), add_special_tokens=False)["input_ids"]
    completion_ids = tok(" " + target_for(x), add_special_tokens=False)["input_ids"] + [EOS]
    ids = prompt_ids + completion_ids
    target_tokens = ids[1:] + [EOS]
    weights = [0.0] * (len(prompt_ids) - 1) + [1.0] * (len(completion_ids) + 1)
    return {"input_ids": ids, "target_tokens": target_tokens, "weights": weights}

batch = [make_datum(x) for x in train_inputs]

2. Create a session and train

client.session() opens a training session (a context manager that frees the model on exit). session.create_model(...) creates a LoRA model to train, sized by river.LoraConfig(rank=...) (the default supported rank range is 1–32). Each step: forward_backward(data, loss_fn="cross_entropy") runs the forward pass and accumulates gradients (returns .metrics["loss"]), and optim_step(lr=...) applies an AdamW update and advances model.step.

with client.session(project="sft-prefix") as session:
    model = session.create_model(base_model=BASE, lora=river.LoraConfig(rank=32))
    print("model_id:", model.model_id)

    for step in range(15):
        fb = model.forward_backward(batch, loss_fn="cross_entropy")
        model.optim_step(lr=2e-4, grad_clip_norm=1.0)
        print(f"step {model.step:2d}  loss={fb.metrics['loss']:.4f}")

The loss collapses within a handful of steps:

model_id: ba1c7208-3c2a-442d-98f2-8aea7bb6980e:model:1
step  1  loss=34.4821
step  2  loss=19.3106
step  3  loss=9.1079
step  4  loss=3.2211
step  5  loss=0.1230
step  6  loss=0.0441
step  7  loss=0.0037
step  8  loss=0.0001
...
step 15  loss=0.0000

3. Sample from the trained model

model.sample(prompt, ...) generates from the model's current in-memory weights (no checkpoint needed) and returns list[list[Sample]] (per-prompt × per-sample). We sample greedily (temperature=0.0), still inside the session:

    for x in ["hello world", "the lazy dog sleeps"]:
        out = model.sample(render(x), max_tokens=16, temperature=0.0, stop=["\n"])
        print(f"{x!r} -> {out[0][0].text!r}")

It learned the rule, and even generalizes to the unseen "the lazy dog sleeps":

'hello world' -> ' test-hello test-world'
'the lazy dog sleeps' -> ' test-the test-lazy test-dog test-sleeps'

4. Save a checkpoint and sample from it

model.save_weights(name, mode="inference") saves the LoRA as a checkpoint, and session.sample(..., checkpoint=ckpt) samples from a saved checkpoint — it loads the LoRA, generates, then unloads (no live model needed):

    ckpt = model.save_weights("prefix", mode="inference")
    print("saved:", ckpt.path)

    for x in ["hello world", "the lazy dog sleeps"]:
        out = session.sample(render(x), base_model=BASE, checkpoint=ckpt,
                             max_tokens=16, temperature=0.0, stop=["\n"])
        print(f"{x!r} -> {out[0][0].text!r}")

The checkpoint reproduces the same behavior:

saved: river://6501216f-c72d-4186-a37c-b65bee62bf58/sampler_weights/prefix
'hello world' -> ' test-hello test-world'
'the lazy dog sleeps' -> ' test-the test-lazy test-dog test-sleeps'

Install the dependencies and run the whole thing:

pip install river-client transformers
export RIVER_API_KEY="rv_..."
python sft.py

Full runnable script (sft.py)

import os
import river_client as river
from transformers import AutoTokenizer

client = river.Client(api_key=os.environ["RIVER_API_KEY"])
BASE = "Qwen/Qwen3.6-35B-A3B-FP8"
tok = AutoTokenizer.from_pretrained(BASE)
EOS = tok.eos_token_id

def target_for(x):
    return " ".join("test-" + w for w in x.split())

train_inputs = [
    "hello world",
    "good morning sunshine",
    "the quick brown fox",
    "i love programming",
    "open the pod bay doors",
    "river flows to the sea",
]

def render(x):
    return f"Input: {x}\nOutput:"

def make_datum(x):
    prompt_ids = tok(render(x), add_special_tokens=False)["input_ids"]
    completion_ids = tok(" " + target_for(x), add_special_tokens=False)["input_ids"] + [EOS]
    ids = prompt_ids + completion_ids
    target_tokens = ids[1:] + [EOS]
    weights = [0.0] * (len(prompt_ids) - 1) + [1.0] * (len(completion_ids) + 1)
    return {"input_ids": ids, "target_tokens": target_tokens, "weights": weights}

batch = [make_datum(x) for x in train_inputs]

with client.session(project="sft-prefix") as session:
    model = session.create_model(base_model=BASE, lora=river.LoraConfig(rank=32))
    print("model_id:", model.model_id)

    # Train
    for step in range(15):
        fb = model.forward_backward(batch, loss_fn="cross_entropy")
        model.optim_step(lr=2e-4, grad_clip_norm=1.0)
        print(f"step {model.step:2d}  loss={fb.metrics['loss']:.4f}")

    # Sample from the live trained weights
    for x in ["hello world", "the lazy dog sleeps"]:
        out = model.sample(render(x), max_tokens=16, temperature=0.0, stop=["\n"])
        print(f"{x!r} -> {out[0][0].text!r}")

    # Save an inference checkpoint and sample from it
    ckpt = model.save_weights("prefix", mode="inference")
    print("saved:", ckpt.path)
    for x in ["hello world", "the lazy dog sleeps"]:
        out = session.sample(render(x), base_model=BASE, checkpoint=ckpt,
                             max_tokens=16, temperature=0.0, stop=["\n"])
        print(f"{x!r} -> {out[0][0].text!r}")

Evaluate beyond the demonstration

This tiny dataset is useful for checking that the training loop works. Its training loss and a few sample outputs are not a broad evaluation of the model. Keep a separate set of inputs, compare the starting and trained models under the same generation settings, and inspect failures. The SFT chapter explains how to interpret those results.