An API to train models
that are yours.

import numpy as np
import river_client as river
from collections import deque
from transformers import AutoTokenizer

MODEL         = "Qwen/Qwen3.6-35B-A3B-FP8"   # 35B MoE, FP8
LORA_RANK     = 8
LEARNING_RATE = 1e-4
WARMUP_STEPS  = 100          # linear LR warmup
PIPELINE_K    = 1            # PipelineRL depth (1 = on-policy)
GROUP         = 16           # generations per prompt
BATCH         = 48           # prompts per step
MAX_TOKENS    = 4096
EPS_MAX       = 6.0          # CISPO clip

client = river.Client(api_key="...")
tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)

SUFFIX = " Think step by step, then put your final answer inside \boxed{}."

def render(problem):
    messages = [{"role": "user", "content": problem + SUFFIX}]
    return tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

def grade(response, answer):
    ...   # your task reward, e.g. a symbolic math checker -> 0.0 / 1.0

def advantages(group_rewards):
    # ScaleRL: center each group by its own mean (the GRPO baseline), but
    # normalize by the BATCH std — not the per-group std that standard GRPO
    # uses. This drops GRPO's difficulty/length bias from per-group scaling.
    centered = [[r - np.mean(g) for r in g] for g in group_rewards]
    std = max(float(np.std([a for g in centered for a in g])), 1e-8)
    return [[a / std for a in g] for g in centered]

def build_cispo_data(prompts, groups, advs):
    ...   # pack prompt + sample tokens with old_logprobs and advantages

def lr_at(train_step):
    # linear warmup over the first WARMUP_STEPS optimizer steps, then hold
    if train_step < WARMUP_STEPS:
        return LEARNING_RATE * (train_step + 1) / WARMUP_STEPS
    return LEARNING_RATE

with client.session() as session:
    model = session.create_model(
        base_model=MODEL,
        lora=river.LoraConfig(
            rank=LORA_RANK,
            train_attn=True,     # attention projections
            train_mlp=True,      # MLP / MoE experts
            train_unembed=True,  # unembedding (lm_head)
        ),
    )

    pipeline = deque()                  # holds the last PIPELINE_K batches
    train_step = 0                      # optimizer steps; LR warmup keys on this

    for step in range(200):
        problems, answers = next_batch(BATCH)   # Polaris-53K, difficulty <= 3/8
        prompts = [render(p) for p in problems]

        # 1. Sample from an inference checkpoint of the current weights
        ckpt = model.save_weights(f"sample_{step}", mode="inference")
        groups = session.sample(prompts=prompts, base_model=MODEL,
                                checkpoint=ckpt, num_samples=GROUP,
                                max_tokens=MAX_TOKENS, stop=["<|im_end|>"])

        # 2. Grade, normalize advantages, drop zero-variance groups, build CISPO data
        rewards = [[grade(s.text, a) for s in g] for g, a in zip(groups, answers)]
        advs = advantages(rewards)

        # A group where every sample earned the same reward has all-zero
        # advantages -> no learning signal. Skip it (ScaleRL).
        keep = [any(a != 0.0 for a in g) for g in advs]
        data = build_cispo_data(
            [p for p, k in zip(prompts, keep) if k],
            [g for g, k in zip(groups,  keep) if k],
            [a for a, k in zip(advs,    keep) if k],
        )
        pipeline.append(data)

        # 3. Train once PIPELINE_K batches are queued, with a warmed-up LR
        if len(pipeline) >= PIPELINE_K:
            model.forward_backward(data=pipeline.popleft(), loss_fn="cispo", eps_max=EPS_MAX)
            model.optim_step(
                lr=lr_at(train_step), grad_clip_norm=1.0,
                beta1=0.9, beta2=0.95, eps=1e-15, weight_decay=0.01,  # ScaleRL Adam
            )
            train_step += 1

    model.save_weights("final")