> ## Documentation Index
> Fetch the complete documentation index at: https://docs.baseten.co/llms.txt
> Use this file to discover all available pages before exploring further.

# Train on a chat dataset with the Loops SDK

> Train a LoRA adapter on a chat dataset with serverless Loops, save resumable state, and sample from the checkpoint.

Continue from the [quickstart](https://docs.baseten.co/loops/quickstart.md) with a chat dataset: prepare batches, train a LoRA adapter, save resumable state, and sample from the trained checkpoint. This example uses the same GLM-5.3-Flash model and explicitly selects serverless Loops, so it bills for tokens and fails if serverless can't accept the run.

The example dataset is [pirate-ultrachat-10k](https://huggingface.co/datasets/winglian/pirate-ultrachat-10k), chat-format conversations that teach the model pirate dialect. The [Jobs quickstart](https://docs.baseten.co/training/getting-started.md) uses the same dataset.

Use the Python project and `BASETEN_API_KEY` from the quickstart. For a model or configuration that needs dedicated capacity, see [Dedicated Loops](https://docs.baseten.co/loops/availability-modes/dedicated.md).

## Turn a dataset into training data

Each training example becomes a [`Datum`](https://docs.baseten.co/reference/sdk/loops/types): input tokens plus loss targets that mask the prompt and supervise the answer, with the same label shift the quickstart uses. For chat data, render the conversation with the tokenizer's chat template and treat the final assistant message as the answer.

Install the dataset library in your project:

```bash theme={"system"}
uv add datasets
```

Save the following code as `train_dataset.py`:

```python train_dataset.py theme={"system"}
from datasets import load_dataset

from baseten.loops import (
    AdamParams,
    Datum,
    ModelInput,
    SamplingParams,
    ServiceClient,
    TensorData,
)

BASE_MODEL = "zai-org/GLM-5.3-Flash"
MAX_SEQ_LEN = 16384

service_client = ServiceClient()
training_client = service_client.create_lora_training_client(
    base_model=BASE_MODEL,
    rank=32,
    max_seq_len=MAX_SEQ_LEN,
    scope="serverless",
)
print(f"run_id={training_client.run_id}")
tokenizer = training_client.get_tokenizer()


def to_datum(example):
    messages = example["messages"]
    if not messages or messages[-1]["role"] != "assistant":
        return None
    prompt = tokenizer.apply_chat_template(
        messages[:-1], tokenize=False, add_generation_prompt=True
    )
    p = tokenizer.encode(prompt, add_special_tokens=False)
    a = tokenizer.encode(messages[-1]["content"], add_special_tokens=False)
    full = p + a
    if not p or not a or len(full) - 1 > MAX_SEQ_LEN:
        return None
    tokens = full[:-1]
    targets = [-100] * (len(p) - 1) + list(a)
    return Datum(
        model_input=ModelInput.from_ints(tokens),
        loss_fn_inputs={
            "target_tokens": TensorData(
                data=targets, dtype="int64", shape=[len(targets)]
            )
        },
    )


dataset = load_dataset("winglian/pirate-ultrachat-10k", split="train[:64]")
data = [datum for ex in dataset if (datum := to_datum(ex)) is not None]
if not data:
    raise ValueError("No examples fit the sequence limit with a final assistant answer")
print(f"prepared {len(data)} examples")
```

The script reads 64 conversations and skips ones without a final assistant answer or whose tokenized input exceeds the sequence limit. Keep this small slice for a first run; choose training data and a held-out evaluation before scaling up.

## Run the training loop

Each iteration is the quickstart's round trip over a batch: one `forward_backward()` on a list of `Datum`, one `optim_step()`. Append:

```python train_dataset.py theme={"system"}
BATCH_SIZE = 8

for step, start in enumerate(range(0, len(data), BATCH_SIZE), 1):
    batch = data[start : start + BATCH_SIZE]
    fb = training_client.forward_backward(data=batch).result(timeout=600.0)
    training_client.optim_step(AdamParams(learning_rate=4e-5)).result(timeout=600.0)
    print(f"step {step} loss={fb.loss:.4f}")
```

## Save a resumable checkpoint

As in the quickstart, save both types of checkpoint: [`save_state()`](https://docs.baseten.co/reference/sdk/loops/training-client) preserves training and optimizer state, while `save_weights_for_sampler()` publishes the adapter for sampling and deployment. Append:

```python train_dataset.py theme={"system"}
state = training_client.save_state(name="epoch-1").result(timeout=600.0)
save_resp = training_client.save_weights_for_sampler(name="epoch-1").result(
    timeout=600.0
)
print(f"resumable state at {state.path}")
print(f"sampler weights at {save_resp.path}")
```

To resume later, provision a training client and call [`load_state_with_optimizer()`](https://docs.baseten.co/reference/sdk/loops/training-client) with the saved `state.path`.

## Sample from the checkpoint

Create a sampling client from the returned sampler checkpoint, then generate a completion without creating an inference deployment. The serverless client stays pinned to that checkpoint; publish again and create a new client to sample later weights. Append:

```python train_dataset.py theme={"system"}
sampling_client = training_client.create_sampling_client(model_path=save_resp.path)
prompt = tokenizer.apply_chat_template(
    [{"role": "user", "content": "How do I learn Python?"}],
    tokenize=False,
    add_generation_prompt=True,
)
sample = sampling_client.sample(
    prompt=ModelInput.from_ints(tokenizer.encode(prompt, add_special_tokens=False)),
    num_samples=1,
    sampling_params=SamplingParams(max_tokens=48),
)
print(f"completion={tokenizer.decode(sample.sequences[0].tokens)!r}")
sampling_client.close()
training_client.close()
```

Run the complete script:

```bash theme={"system"}
uv run python train_dataset.py
```

Check that it finishes without an error and prints the run ID, number of prepared examples, loss for each batch, both returned checkpoint paths, and a completion. Trainer state paths use `weights/`; sampler paths use `sampler_weights/`. Keep the returned paths rather than constructing them, because checkpoint names can include a unique suffix.

A completion checks that sampling works. A small training slice doesn't establish a change in dialect, and batch losses need not fall at every step. To measure whether training helped, compare the checkpoint with the base model on a [held-out evaluation](https://docs.baseten.co/post-training/evaluate.md).

## Shut down the run

Closing the SDK clients releases local connections and stops the training client's keepalive. An idle serverless run incurs no token charges. To end the run and its paired sampler explicitly, use the printed run ID. With the [Baseten CLI](https://docs.baseten.co/reference/cli/baseten/overview.md#install) installed, end the run:

1. Deactivate the run, and confirm when prompted:

   ```bash theme={"system"}
   baseten loops run deactivate --run-id "<run_id>"
   ```

2. Check that the run's `status.name` is `INACTIVE`:

   ```bash theme={"system"}
   baseten loops run describe --run-id "<run_id>" --output json
   ```

Saved checkpoints remain available. See [Deactivate a run](https://docs.baseten.co/loops/deactivate-a-run.md) for the REST API option.

## Next steps

* **[Deploy a checkpoint](https://docs.baseten.co/loops/deploy-checkpoints.md)**: Serve the returned sampler checkpoint as an inference endpoint.
* **[Hyperparameters](https://docs.baseten.co/loops/hyperparameters.md)**: Apply a learning-rate schedule to this training loop.
* **[Loss functions](https://docs.baseten.co/loops/loss-functions.md)**: Swap `cross_entropy` for `importance_sampling`, `ppo`, and other objectives when you move past SFT.