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

​

Prerequisites

• The Python project and BASETEN_API_KEY from the quickstart.
• uv and Python 3.12 or later.

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

Output

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.

​

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:

1. Deactivate the run, and confirm when prompted:

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

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

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

Saved checkpoints remain available. For the REST API option, see Deactivate a run.

​

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, and bills for its GPU while it’s live. With the Baseten CLI installed, list the run’s checkpoints and deploy the one you want:

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 covers the REST API and dashboard paths.

​

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.