API docs
Guides

The training loop

Drive gradients yourself with rendered token IDs, forward-backward passes and optimiser steps.

Managed fine-tuning is the right tool when a supervised job on a JSONL file is all you need. This guide is for the other case: you are doing RL, DPO, distillation or a custom objective, and you want to drive every gradient yourself.

The low-level surface gives you a small set of primitives. A training run holds LoRA state for a base model on allocated GPUs. Forward-backward computes gradients on batches you send. Optim-step applies them. Checkpoints make the result durable, and samplers generate from it. Everything in this guide is those calls in a loop.

Tokens are your job

Every training and sampling route takes token IDs, not text. Rendering, chat templates and loss masking sit with you, the caller. The compatibility routes render server-side, which makes them the easy on-ramp and the wrong tool for RL: in a training loop you need to know exactly which tokens the model saw and which tokens contribute to the loss.

Render a conversation does the heavy lifting. Send a messages array and it applies the model's chat template server-side, returning token_ids plus a weights array, the per-token loss mask that marks assistant tokens as trainable and everything else as context. See Tokens and rendering for how the mask works.

The shape of the loop

POST /v1/render                                  token_ids + weights
POST /v1/training-runs                           run_3Bz

repeat:
  POST /v1/training-runs/run_3Bz/forward-backward
  POST /v1/training-runs/run_3Bz/optim-step

POST /v1/training-runs/run_3Bz/checkpoints       nucleus://...
POST /v1/samplers                                sampler_id
DELETE /v1/training-runs/run_3Bz

The rest of this guide builds that sketch into a complete supervised fine-tuning loop, then shows the RL and custom-loss variants.

Render the dataset

Render each conversation once, up front. The response carries everything a training step needs: the token IDs to feed the model and the weights to mask the loss.

import requests

BASE = 'https://api.nucleus-ai.io/v1'
HEADERS = {'Authorization': f'Bearer {API_KEY}'}
MODEL = 'Qwen/Qwen3-8B'

def render(messages):
    resp = requests.post(
        f'{BASE}/render',
        headers=HEADERS,
        json={'model': MODEL, 'messages': messages},
    )
    return resp.json()  # {'token_ids': [...], 'weights': [...]}

dataset = [render(row['messages']) for row in conversations]

Create a training run

A training run allocates GPUs and returns the run ID that every training route is keyed on. You choose the LoRA rank, the seed, and which weight matrices train.

resp = requests.post(
    f'{BASE}/training-runs',
    headers=HEADERS,
    json={
        'base_model': MODEL,
        'rank': 32,
        'seed': 42,
        'train_mlp': True,
        'train_attn': True,
        'train_unembed': False,
    },
)
run = resp.json()
run_id = run['training_run_id']  # "run_3Bz"

rank defaults to 32 and train_mlp, train_attn and train_unembed all default to true, so the minimal body is just base_model. Set seed for reproducible runs and user_metadata for your own tags. The same route also accepts resume fields (from_path, restore_optimizer, weights_access_token) instead of the new-run fields; send one set or the other, not both. Resuming is covered below.

Creating a run allocates hardware, so retrying a timed-out request can allocate twice. Send an Idempotency-Key header if your client retries automatically. See Errors and retries.

Forward-backward, then optim-step

One training step is two calls. Forward-backward accumulates gradients for a batch. Optimizer step applies them with AdamW and advances the step counter.

data is an array of examples, each pairing model_input (the token IDs to run through the model) with loss_fn_inputs (whatever the chosen loss needs). For cross-entropy that is target_tokens, the rendered tokens shifted by one position, plus the weights mask aligned to them. Built-in losses cover supervised fine-tuning and the common RL objectives: cross-entropy, importance sampling, PPO, CISPO and DRO. Name the loss in loss_fn and pass any loss-specific settings in loss_fn_config.

BATCH = 8

for step in range(len(dataset) // BATCH):
    batch = dataset[step * BATCH:(step + 1) * BATCH]

    fb = requests.post(
        f'{BASE}/training-runs/{run_id}/forward-backward',
        headers=HEADERS,
        json={
            'data': [
                {
                    'model_input': ex['token_ids'][:-1],
                    'loss_fn_inputs': {
                        'target_tokens': ex['token_ids'][1:],
                        'weights': ex['weights'][1:],
                    },
                }
                for ex in batch
            ],
            'loss_fn': 'cross_entropy',
        },
    ).json()

    opt = requests.post(
        f'{BASE}/training-runs/{run_id}/optim-step',
        headers=HEADERS,
        json={'learning_rate': 8e-5},
    ).json()

    print(opt['step'], fb['loss'], fb['metrics']['grad_norm'])

learning_rate is required on every optim-step: the schedule is yours, so warm-up, decay and restarts are a matter of what you pass each step. betas, eps and weight_decay are optional. Note that weight_decay defaults to 0.0, unlike PyTorch's AdamW, which catches people out. If you expect regularisation, pass it explicitly.

To measure validation loss without touching gradients, send the same body to Forward pass instead.

Read the loss curve

The platform records loss, grad_norm and lr per step on the server, so the curve survives a dropped connection or a restarted client. You do not need to log every forward-backward response yourself; fetch the history from Get run metrics whenever you need it.

resp = requests.get(
    f'{BASE}/training-runs/{run_id}/metrics',
    headers=HEADERS,
    params={'from_step': 0, 'keys': 'loss,grad_norm'},
)
series = resp.json()['series']
# [{'step': 238, 'loss': 0.74, 'grad_norm': 1.12, 'lr': 8e-5}, ...]

Checkpoint as you go

Training state lives on GPUs until you persist it. Write a training checkpoint every few hundred steps so a failure costs you minutes, not the run.

if opt['step'] % 200 == 0:
    ckpt = requests.post(
        f'{BASE}/training-runs/{run_id}/checkpoints',
        headers=HEADERS,
        json={'type': 'training', 'name': f"{opt['step']:06d}"},
    ).json()
    # ckpt['path'] == "nucleus://run_3Bz:train:0/weights/000200"

After a failure you have two ways back, both keyed on the checkpoint's nucleus:// path:

  • New run. Create a training run with from_path instead of the new-run fields. Pass restore_optimizer: true to restore the AdamW state alongside the weights, and weights_access_token if the checkpoint requires one.
  • Live run. If the run is still up and you only want different weights in it, call Load state with path to swap weights into the running session without opening a new run.
resp = requests.post(
    f'{BASE}/training-runs',
    headers=HEADERS,
    json={
        'from_path': 'nucleus://run_3Bz:train:0/weights/000200',
        'restore_optimizer': True,
    },
)
run_id = resp.json()['training_run_id']

Sample from the run

Sampling needs a sampler checkpoint, a lighter artefact than a training checkpoint. Write one, then open a sampler on its path. Like every low-level route, Generate samples speaks token IDs: render the prompt first, and decode the output with Detokenize token IDs.

ckpt = requests.post(
    f'{BASE}/training-runs/{run_id}/checkpoints',
    headers=HEADERS,
    json={'type': 'sampler', 'name': '000240'},
).json()
# ckpt['path'] == "nucleus://run_3Bz:train:0/sampler_weights/000240"

sampler = requests.post(
    f'{BASE}/samplers',
    headers=HEADERS,
    json={'model_path': ckpt['path']},
).json()
sampler_id = sampler['sampler_id']

prompt = render([{'role': 'user', 'content': 'Summarise the incident report.'}])
out = requests.post(
    f'{BASE}/samplers/{sampler_id}/sample',
    headers=HEADERS,
    json={
        'prompt': {'token_ids': prompt['token_ids']},
        'num_samples': 1,
        'sampling_params': {'max_tokens': 256},
    },
).json()

text = requests.post(
    f'{BASE}/detokenize',
    headers=HEADERS,
    json={'model': MODEL, 'token_ids': out['samples'][0]['token_ids']},
).json()['text']

For SFT this is how you spot-check the model as it trains. For RL it is the inner loop itself: checkpoint, sample rollouts from the current policy, score them, then feed the scores back through forward-backward with an RL loss such as importance sampling or PPO, and repeat. Because samplers are opened from checkpoints, the policy you sample from is always an explicit, versioned artefact rather than whatever happens to be in GPU memory.

Close the run

Runs and samplers hold GPU allocations, and allocations are billed while they exist. When training is done, close both to release the hardware and stop the meter.

requests.delete(f'{BASE}/samplers/{sampler_id}', headers=HEADERS)
requests.delete(f'{BASE}/training-runs/{run_id}', headers=HEADERS)

Close a training run keeps your checkpoints by default (keep_checkpoints defaults to true), so everything you persisted outlives the run. A run left open keeps its allocation whether or not you are stepping it, so close runs you are not actively using and resume from a checkpoint later.

Bring your own loss

When the built-in losses do not fit, compute the loss yourself and let the platform do the heavy tensor work. Custom-loss forward-backward is a two-leg exchange:

  1. Leg one. Send your batch with loss_type_input: "logprobs". The platform runs the forward pass and returns an exchange_id, the per-token log probabilities, and expires_in, the number of seconds the exchange stays open.
  2. Leg two. Differentiate your loss with respect to those log probabilities in any framework you like, then submit the gradients against the exchange_id before it expires. The platform completes the backward pass with your gradients, accumulating them exactly as a built-in loss would. Call optim-step as usual.
leg1 = requests.post(
    f'{BASE}/training-runs/{run_id}/forward-backward-custom',
    headers=HEADERS,
    json={
        'data': [{'model_input': ex['token_ids']} for ex in batch],
        'loss_type_input': 'logprobs',
    },
).json()
# {'exchange_id': 'xch_5Tn', 'logprobs': [[-0.21, -1.84, ...]], 'expires_in': 60}

grads = my_loss_backward(leg1['logprobs'])  # your loss, your framework

Submit grads against leg1['exchange_id'] before expires_in lapses, then call optim-step as usual. See Custom-loss forward-backward for the exchange in full.

The exchange holds intermediate state on the GPUs, so it expires quickly; compute your gradients promptly and treat an expired exchange as a step to redo, not an error to retry.

Where next

The training loop | Nucleus AI API