river-client-training

Train with river_client (the current API)

river_client is the Python client for River's training API: create a LoRA model on remote GPU workers inside a session, push training data through train_step, sample from the live weights, checkpoint. This skill is the short path to writing a correct script with the current API (train_step era, package ≥ 0.5).

Setup and connection

pip install river-client        # PyPI, Python 3.12+

import os

import river_client as river

client = river.Client(api_key=os.environ["RIVER_API_KEY"], endpoint="api.river.ai")

Everything happens inside a session (the GPU allocation) and a model:

with client.session() as session:
    model = session.create_model(
        base_model="Qwen/Qwen3.6-35B-A3B-FP8",
        lora=river.LoraConfig(rank=16, train_unembed=True),  # see note below
    )
    ...
# GPUs freed automatically on exit.

Any keyword arguments to client.session(...) become session tags — arbitrary string key→value metadata stamped on the session:

with client.session(experiment="grpo-math", run="lr4e-5-r16") as session:
    ...

Tags don't change behavior; they exist so runs can be found later: the River Console can filter training runs by tag, and tags travel with the run for correlating against an external experiment tracker (a common convention is client.session(wandb_project=..., wandb_name=...)). Tag long-running experiments — an untagged session is hard to tell apart from every other one once you have dozens. Key/value counts and lengths are bounded server-side, so keep them short labels, not payloads.

LoraConfig knobs: rank (max 32), train_attn / train_mlp (both default on), train_unembed (default off), seed (reproducible adapter init).

Why train_unembed=True for RL. The unembedding (lm_head) is the final hidden-state → vocab-logits matrix; this flag adds a LoRA adapter on it. With the head frozen, a policy update can only change token probabilities indirectly, by bending hidden states through the trunk adapters. The RL losses (importance_sampling / ppo / cispo) are exactly "push probability toward or away from these sampled tokens, weighted by advantage" — a gradient that lands first on the logits — and RL's signal is one scalar reward per sequence, so you want it expressible as directly as possible. A head adapter gives that direct lever (suppress premature EOS, boost format tokens, sharpen the distribution) without spending trunk capacity — and without it, updates burn more KL (kl, mean_ratio drift) for less reward gain. For SFT the dense per-token cross-entropy signal usually makes trunk adapters sufficient, which is why the flag defaults off; turn it on for RL loops.