Fine-tuning

Start from a recipe: a base model, a dataset and a few settings. train starts a job on Relay’s GPUs and returns at once.

from relay.sft.config import SFTConfig
from relay.sft.train import train

config = SFTConfig(
    model="pi05",
    dataset_name="lerobot/pusht_image",
    batch_size=32,
    learning_rate=2.5e-5,
)

job = train(config)
job.wait()
print(job.status, job.step, job.loss)   # completed 6000 0.012

To write the training loop yourself instead, see Training loops.

The config

str

required

The base model: "pi05".

str

required

A LeRobot dataset on the Hugging Face Hub, such as "lerobot/pusht_image". See Datasets.

int

default:"128"

Samples per step, 1 to 512.

float

default:"2.5e-5"

The peak learning rate, after warmup.

int

default:"6000"

Optimizer steps, 1 to 100,000.

str

default:"lora"

"lora" trains low-rank adapters; "full" trains every weight of the targets.

int

default:"32"

The adapters’ rank, 1 to 256.

float

LoRA’s alpha: the update is scaled by alpha / rank. By default, the rank.

list[str]

Where training goes. By default the action expert for LoRA, and the whole model in full. See Where training goes.

dict[str, str]

Model camera → dataset feature, such as {"base_0_rgb": "observation.images.top"}. By default, the dataset’s image features, in order.

str

A name for the job.

SFTConfig(model="pi05", dataset_name=..., finetune="full")                                 # every weight
SFTConfig(model="pi05", dataset_name=..., targets=["all"])                                 # LoRA everywhere
SFTConfig(model="pi05", dataset_name=..., targets=["language.mlp", "action_expert.attn"])  # LoRA at select places

The rest is fixed: AdamW (β 0.9, 0.95), a linear warmup over the first tenth of the steps (at most 200), then cosine decay to zero; gradients clipped at norm 1; chunks of 10 actions; checkpoints every 1,000 steps.

Datasets

A dataset is a LeRobot dataset on the Hugging Face Hub, in format v2.1 (Relay reads the repo’s v2.1 revision), with:

its cameras stored as images ("dtype": "image" in meta/info.json). Datasets that store video aren’t supported yet;
the state in observation.state and the actions in action, each at most 32 wide;
the tasks in meta/tasks.jsonl: each frame’s task is its prompt.

Relay downloads a dataset once. If it has no norm_stats.json, Relay computes one: state and actions are normalized so their 1st and 99th percentiles map to −1 and 1, then padded with zeros to 32. Send states, and read actions, in those units when you serve the result.

The job

train returns a FineTuneJob, as the API last reported it:

Attribute	Is
`id`	`"ft-..."`
`status`	`queued` (its GPU is starting), `running`, `cancel_requested`, `completed`, `cancelled` or `error`
`step`, `loss`	progress: the last optimizer step and its loss
`error`	why it failed
`output_dir`	where the trained weights are, once completed: `model.safetensors`, adapters merged in
and the config’s fields	`model`, `dataset_name`, `batch_size`, `learning_rate`, `steps`, `finetune`, `lora_rank`, `lora_alpha`, `targets`, `cameras`, `name`

job.refresh()          # bring it up to date
job.wait(poll=30)      # until completed or cancelled; RuntimeError if it failed
job.cancel()
job.events()           # [{"step": 10, "loss": 0.21, "grad_norm": 0.9, "lr": 2.4e-6, "skipped": 0.0}, ..., {"step": 6000, "event": "finished"}]

Find jobs again:

from relay.sft import list_jobs, retrieve

retrieve("ft-09dbfed4f795")
list_jobs()            # yours, newest first

A failed request raises RuntimeError with the API’s message, such as POST /v1/fine-tunes: 400 bad batch_size: 0.

Get started

Inference

Training

Reference

The config

Datasets

The job