Train on a chat dataset with the Loops SDK

Continue from the quickstart 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, chat-format conversations that teach the model pirate dialect. The Jobs quickstart 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.

Turn a dataset into training data

Each training example becomes a Datum: 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:

uv add datasets

Save the following code as train_dataset.py:

train_dataset.py

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:

train_dataset.py

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() preserves training and optimizer state, while save_weights_for_sampler() publishes the adapter for sampling and deployment. Append:

train_dataset.py

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() 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:

train_dataset.py

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:

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.

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 installed, end the run:

Deactivate the run, and confirm when prompted:

baseten loops run deactivate --run-id "<run_id>"

Check that the run’s status.name is INACTIVE:

baseten loops run describe --run-id "<run_id>" --output json

Saved checkpoints remain available. See Deactivate a run for the REST API option.

Next steps

Deploy a checkpoint: Serve the returned sampler checkpoint as an inference endpoint.
Hyperparameters: Apply a learning-rate schedule to this training loop.
Loss functions: Swap cross_entropy for importance_sampling, ppo, and other objectives when you move past SFT.

Examples

Post-training

Models

Engines

Custom Docker servers

Custom Python models

Chains

Baseten for Model Labs

Train on a chat dataset with the Loops SDK

Turn a dataset into training data

Run the training loop

Save a resumable checkpoint

Sample from the checkpoint

Shut down the run

Next steps

​Turn a dataset into training data

​Run the training loop

​Save a resumable checkpoint

​Sample from the checkpoint

​Shut down the run

​Next steps

Turn a dataset into training data

Run the training loop

Save a resumable checkpoint

Sample from the checkpoint

Shut down the run

Next steps