smolgym

Training API · Small models

You write
the loop.
Our GPUs.

SmolGym is post-training for small language models. Your Python script keeps the data, the reward, and the loss. We run sampling, the backward pass, and the optimizer. The API is compatible with the Tinker SDK. The models stay small, and custom weights are welcome.

SFTRLLoRACustom weights

Start a run See one step

A Python loop

You write

The GPUs

We run

Small, or yours

The model

Downloadable

The weights

01

The split

Your recipe, our machines

You own the training loop. SmolGym owns the part that is annoying to run: loading the model, sampling it, computing gradients, and keeping the checkpoint. Your script can sit on a laptop.

You control

  • Data and environmentsPrompts, tasks, tool traces, and the sandbox those tasks run in.
  • Rewards and verifiersAn exact match, a unit test, a rubric, or a judge you call.
  • The lossSupervised, reinforcement, preference, or a loss you write.
  • The modelA small open model on the rack, or weights you upload.

SmolGym runs

  • The machinesGPUs sized for a small model, not a frontier cluster.
  • The samplerCompletions from the weights you are training.
  • The trainerforward_backward() and optim_step(), for the loss you sent.
  • The checkpointState you can resume, and weights you can download.
02

The step

Tinker SDK compatible

One reinforcement step is sample(), your_reward(),forward_backward(), then optim_step(). Supervised training uses the same loop without sample(). The calls match the Tinker SDK. SmolGym runs them for small models, and for models you bring.

One steptrain.pyThe call shape
training = service.create_lora_training_client(
    base_model="rack/small",  # or weights you uploaded
)

for step in range(n_steps):
    sampling = training.save_weights_and_get_sampling_client()
    responses = sampling.sample(prompts, params).result()
    rewards = your_reward(responses)

    training.forward_backward(batch, loss_fn="importance_sampling")
    training.optim_step(AdamParams(learning_rate=1e-4))

training.save_state("run-1")
  1. sample()

    Draw completions from the weights you are training.

  2. your_reward()

    Score them. The verifier, the sandbox, and the judge stay in your code.

  3. forward_backward()

    Send the batch and the loss. We compute the gradients.

  4. optim_step()

    Apply the update on the GPUs that hold the model.

  5. save_state()

    Keep optimizer state, and download the weights when you want them.

Full loopOne step, then againBlue moves weights
The training loopYour script sends prompts to the sampler, scores the completions, and sends the batch to the trainer. The trainer writes the weights. Sync copies those weights back to the sampler for the next step. save_state downloads them.Your machineOur GPUs01train.pyThe loop lives here.A laptop is enough.03your_reward()You score the completions.Weights do not change.A verifier, a test, a judge.02sample()Reads a copy of W.W · read only04forward_backward()05optim_step()Writes W from your loss.W · blue cells are written06syncCopies written W ontothe sampler.save_state()Download W when you want.promptscompletionsscores + batchwritten Wto sampler
03

Weights

Who writes W

Two copies of the weights sit on the GPUs. The sampler reads one, to generate. The trainer writes the other, when you call optim_step(). Sync copies the written weights onto the sampler, so the next sample() sees the update. Your script sends the loss. It does not hold W.

  1. ReadSampler

    Uses the last synced W. Does not train.

  2. WriteTrainer

    optim_step() is the only place W changes.

  3. NoneYour script

    Scores and sends the batch. W stays on the GPU.

LoRAThe adapter moves. The base stays.

Several runs can share the loaded base. The step writes the small matrices only.

Base · held

Adapter · written

Full updateEvery parameter moves.

Same step, larger write. Only when the whole model still fits the floor.

All of W · written

One adapter rowW ← W − ∇Llr = 1, so the arithmetic is visible
W0.120.04-0.210.080.15-0.02
∇L0.03-0.010.050.00-0.020.01
W′0.090.05-0.260.080.17-0.03

The loss is yours. The subtraction runs on the trainer.

04

Methods

What the loop can be
  1. 01

    Supervised

    A batch and a loss. No sampling. The script sends examples, we run forward_backward() and optim_step().

  2. 02

    Reinforcement

    Sample, score, update. The reward can be an exact match, a test, a rubric, or a judge.

  3. 03

    Preferences

    A pairwise loss, if that is the update you want. The loop is still yours.

  4. 04

    LoRA

    An adapter on a small base. Several runs can share the loaded model.

  5. 05

    Full update

    Every parameter, when an adapter is not enough, and the model still fits the floor.

  6. 06

    Your environment

    Multi-turn tools, code execution, a search sandbox. You run the episode. We call sample() each turn and train on the signal you send back.

05

Models

Rack and custom
Rack models and custom weights on the SmolGym floor
RackCustom
ModelA small open model we hostWeights you upload
FitOne training machineThe same machine
UpdateLoRA, or the full modelLoRA, or the full model
AfterDownload the weightsDownload the weights
06

Start a run.

A session

Send the model, whether the loop is supervised or reinforcement, and what the reward checks. If it fits the floor, you get a session: your script, our GPUs, weights you can download.

Start a run
swapnil@rycerz.es