> ## 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. This example trains GLM-5.3-Flash, not the quickstart's Qwen3.6-35B-A3B, 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.

## Prerequisites

* The Python project and `BASETEN_API_KEY` from the [quickstart](https://docs.baseten.co/loops/quickstart.md).
* [uv](https://docs.astral.sh/uv/) and Python 3.12 or later.

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 one training step over a batch: one `forward_backward()` on a list of `Datum`, then 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 the resumable state and the sampler weights: [`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
```

```text Output theme={"system"}
run_id=<RUN_ID>
prepared <N> examples
step <STEP> loss=<LOSS>
...
resumable state at weights/epoch-1-<SUFFIX>
sampler weights at sampler_weights/epoch-1-<SUFFIX>
completion='<TEXT>'
```

Keep the returned paths rather than constructing them, because checkpoint names can include a unique suffix.

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. For the REST API option, see [Deactivate a run](https://docs.baseten.co/loops/deactivate-a-run.md).

## Deploy the checkpoint

The sampler checkpoint is ready to serve. The deployment loads the adapter on the base weights, downloads the base weights with the `hf_access_token` in [workspace secrets](https://docs.baseten.co/organization/secrets.md), and bills for its GPU while it's live. With the [Baseten CLI](https://docs.baseten.co/reference/cli/baseten/overview.md#install) installed, list the run's checkpoints and deploy the one you want:

```bash theme={"system"}
baseten loops checkpoint list --run-id "<run_id>"
baseten loops checkpoint deploy --checkpoint-id "<checkpoint_id>"
```

The deploy command prompts for a model name, GPU type, GPU count, and the Hugging Face secret name. [Deploy a checkpoint](https://docs.baseten.co/loops/deploy-checkpoints.md) covers the REST API and dashboard paths.

## 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.