🔬Modal Research Exchange

{
  "base_model": "Qwen/Qwen3.5-0.8B",
  "baseline": "The same RL loop, staleness, seeds and hyperparameters, but using Tinker's built-in importance_sampling policy-gradient loss instead of the centered-backward surrogate.",
  "batch_size": 128,
  "budget": {
    "gpu_hours": 2.0,
    "max_wall_minutes": 120
  },
  "command": "# custom tinker SDK script (custom parts: stale sampler refresh, auxiliary-token datums, centered loss)\nimport os, tinker, torch\nsc = tinker.ServiceClient(base_url=os.environ[\"TINKER_BASE_URL\"])\ntc = sc.create_lora_training_client(base_model=\"Qwen/Qwen3.5-0.8B\", rank=32)\nsampler = tc.save_weights_and_get_sampling_client(name=\"step0\")\n\ndef centered_loss(data, logprobs_list):\n    # data grouped as [target, aux1..aux8] per sequence; per-token advantages A in each target datum's loss_fn_inputs\n    loss = 0.0\n    for i in range(0, len(logprobs_list), 9):\n        tgt = logprobs_list[i]\n        aux = torch.stack(logprobs_list[i+1:i+9]).mean(0)\n        A = data[i].loss_fn_inputs[\"advantages\"].to_torch()\n        loss = loss + (-A * tgt + A * (aux - aux.detach())).sum()\n    return loss, {}\n\nfor step in range(100):\n    if step % 4 == 0:\n        sampler = tc.save_weights_and_get_sampling_client(name=f\"s{step}\")\n    batch = rollout_gsm8k(sampler, group_size=8, groups=16, max_tokens=512, topk_prompt_logprobs=20)  # cookbook gsm8k env + reward\n    datums = build_target_and_aux_datums(batch, n_aux=8)  # 8 sampled tokens per position from renormalized top-20\n    tc.forward_backward_custom(datums, centered_loss).result()\n    tc.optim_step(tinker.AdamParams(learning_rate=4e-5)).result()\n# baseline arm: same loop with tc.forward_backward(target_datums, loss_fn=\"importance_sampling\")",
  "dataset": {
    "hf_path": "openai/gsm8k",
    "n_examples": 12800,
    "name": "GSM8K",
    "split": "main/train"
  },
  "eval": {
    "benchmark": "GSM8K test",
    "metric": "exact-match accuracy (greedy)",
    "n_examples": 500
  },
  "hypothesis": "With rollouts generated by a policy 4 optimizer steps stale, the eight-token centered-backward loss reaches at least 4 points higher GSM8K test accuracy at step 100 than the importance-sampling policy-gradient loss, or avoids a collapse that the baseline shows.",
  "lora_rank": 32,
  "max_seq_len": 1024,
  "method": "Train a LoRA policy on GSM8K train with a 0/1 exact-match reward and group-mean advantages (group_size 8, 16 groups per batch). To make training off-policy, refresh the sampler weights only every 4 optimizer steps, so rollouts come from a policy up to 4 steps stale.\n\nAuxiliary tokens: for each completion position, request the sampler's top-20 prompt logprobs, renormalize them, and draw 8 auxiliary tokens. This approximates true sampler draws and is a documented deviation from the original algorithm.\n\nTreatment loss (custom, via forward_backward_custom): sum over tokens of -A*logp(target) + A*(mean_aux_logp - stopgrad(mean_aux_logp)). kl_coeff is set to 0 so only the centering differs from the baseline. Auxiliary logprobs come from 8 extra datums per sequence that share model_input but use the auxiliary tokens as target_tokens.\n\nBaseline loss: the built-in importance_sampling loss with sampler logprobs. All other settings are identical: model, data order, seeds, advantages, staleness, LR and renderer.\n\nEvaluation: sample greedily from saved checkpoints every 25 steps. Also log grad norm and mean policy/sampler logprob gap as stability signals.",
  "optimizer": {
    "lr": 4e-05,
    "name": "adam"
  },
  "steps": 100,
  "success_criterion": "Validated if the treatment's GSM8K accuracy at step 100 exceeds the baseline by at least 4 points (about 2x the standard error at n=500), or if the baseline's accuracy drops more than 10 points from its peak while the treatment's does not. Failed if the gap is under 2 points with no collapse difference.",
  "title": "Eight-token score centering vs importance-sampled policy gradient, 4-step-stale rollouts",
  "training_type": "rl"
}

"""LoRA RL with a verifiable reward on a Spindle server through the stock Tinker SDK.

Reads TINKER_BASE_URL, TINKER_API_KEY, MODAL_PROXY_TOKEN_ID and
MODAL_PROXY_TOKEN_SECRET. Each arm in CONFIG["arms"] trains a fresh LoRA on
GSM8K train prompts: sample answers from the current weights, grade them by
exact match on the final number, turn rewards into advantages with the arm's
baseline, and train with the "importance_sampling" loss. Exact-match accuracy
on the test split (greedy decoding) is measured before and after each arm.
Prints one final line `RESULT {...}`: the last arm's metrics at the top level
and every arm under "arms".

Baselines:
- group_mean: GRPO style, reward minus the mean reward of the prompt's group.
- loo_similarity: GS1 style, reward minus a leave-one-out mean over the batch
  weighted by softmax(cosine(TF-IDF(prompt_i), TF-IDF(prompt_j)) / temperature).
- batch_mean: reward minus the plain leave-one-out batch mean.
"""

import json
import math
import os
import re
import time
from collections import Counter

import tinker
from datasets import load_dataset
from tinker import types

# CONFIG-START
CONFIG = {
    "base_model": "Qwen/Qwen3.5-2B",
    "lora_rank": 32,
    "dataset": "openai/gsm8k",
    "dataset_config": "main",
    "split": "train",
    "eval_split": "test",
    "n_eval": 256,
    "steps": 100,
    "lr": 4e-05,
    "max_tokens": 512,
    "temperature": 1.0,
    "similarity_temperature": 0.1,
    "max_wall_seconds": 2700,
    "arms": [
        {
            "name": "grpo",
            "baseline": "group_mean",
            "n_prompts_per_step": 16,
            "group_size": 8
        }
    ]
}
# CONFIG-END

STARTED = time.monotonic()
INSTRUCTION = "\n\nSolve the problem step by step, then write the final answer on its own line as '#### <number>'."


def elapsed():
    return time.monotonic() - STARTED


def load_problems(split, n):
    """Return n (question, gold answer) pairs; GSM8K keeps the gold number after '####'."""
    rows = load_dataset(CONFIG["dataset"], CONFIG["dataset_config"], split=split, streaming=True)
    problems = []
    for row in rows:
        question = row.get("question") or row.get("problem") or row.get("prompt")
        answer = row.get("answer") or row.get("solution") or ""
        gold = answer.split("####")[-1].strip() if "####" in answer else last_number(answer)
        if question and gold:
            problems.append((question, normalize_number(gold)))
        if len(problems) == n:
            break
    return problems


def last_number(text):
    numbers = re.findall(r"-?\d[\d,]*\.?\d*", text)
    return numbers[-1] if numbers else ""


def normalize_number(text):
    text = text.replace(",", "").replace("$", "").strip().rstrip(".")
    try:
        value = float(text)
    except ValueError:
        return text
    return str(int(value)) if value == int(value) else str(value)


def grade(completion, gold):
    """1.0 if the completion's final answer matches the gold number exactly, else 0.0."""
    answer = completion.split("####")[-1] if "####" in completion else completion
    predicted = last_number(answer)
    return 1.0 if predicted and normalize_number(predicted) == gold else 0.0


def prompt_tokens(tokenizer, question):
    ids = tokenizer.apply_chat_template(
        [{"role": "user", "content": question + INSTRUCTION}],
        tokenize=True,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    return list(ids["input_ids"] if hasattr(ids, "keys") else ids)


def stop_tokens(tokenizer):
    ids = {tokenizer.eos_token_id}
    im_end = tokenizer.convert_tokens_to_ids("<|im_end|>")
    if isinstance(im_end, int) and im_end != tokenizer.unk_token_id:
        ids.add(im_end)
    return sorted(i for i in ids if i is not None)


def sample_all(sampler, tokenizer, prompts, num_samples, temperature):
    """Sample num_samples completions per prompt concurrently; returns one list of sequences per prompt."""
    params = types.SamplingParams(
        max_tokens=CONFIG["max_tokens"], temperature=temperature, stop=stop_tokens(tokenizer)
    )
    futures = [
        sampler.sample(prompt=types.ModelInput.from_ints(ids), num_samples=num_samples, sampling_params=params)
        for ids in prompts
    ]
    return [future.result().sequences for future in futures]


def eval_accuracy(training, tokenizer, problems):
    sampler = training.save_weights_and_get_sampling_client()
    prompts = [prompt_tokens(tokenizer, q) for q, _ in problems]
    groups = sample_all(sampler, tokenizer, prompts, 1, 0.0)
    correct = sum(
        grade(tokenizer.decode(group[0].tokens, skip_special_tokens=True), gold)
        for group, (_, gold) in zip(groups, problems)
    )
    return correct / len(problems)


def tfidf_vectors(texts):
    docs = [Counter(re.findall(r"[a-z0-9]+", text.lower())) for text in texts]
    df = Counter(term for doc in docs for term in doc)
    n = len(docs)
    vectors = []
    for doc in docs:
        total = sum(doc.values()) or 1
        vectors.append({t: (c / total) * (math.log((n + 1) / (df[t] + 1)) + 1) for t, c in doc.items()})
    return vectors


def cosine(a, b):
    dot = sum(a[t] * b[t] for t in a if t in b)
    norm = math.sqrt(sum(v * v for v in a.values())) * math.sqrt(sum(v * v for v in b.values()))
    return dot / norm if norm else 0.0


def advantages_for(arm, questions, rewards):
    """One advantage per rollout. questions[i] is the prompt text of rollout i."""
    baseline = arm["baseline"]
    n = len(rewards)
    if n < 2:
        return [0.0] * n
    if baseline == "group_mean":
        advantages = []
        group = arm["group_size"]
        for start in range(0, n, group):
            chunk = rewards[start : start + group]
            mean = sum(chunk) / len(chunk)
            advantages.extend(r - mean for r in chunk)
        return advantages
    if baseline == "batch_mean":
        total = sum(rewards)
        return [r - (total - r) / (n - 1) for r in rewards]
    if baseline == "loo_similarity":
        vectors = tfidf_vectors(questions)
        tau = CONFIG["similarity_temperature"]
        advantages = []
        for i in range(n):
            weights = [0.0 if j == i else math.exp((cosine(vectors[i], vectors[j]) - 1.0) / tau) for j in range(n)]
            total = sum(weights)
            if total == 0:
                weights, total = [0.0 if j == i else 1.0 for j in range(n)], n - 1
            advantages.append(rewards[i] - sum(w * r for w, r in zip(weights, rewards)) / total)
        return advantages
    raise ValueError(f"unknown baseline {baseline!r}")


def make_datum(prompt_ids, sequence, advantage):
    full = prompt_ids + list(sequence.tokens)
    n_prompt = len(prompt_ids) - 1
    return types.Datum(
        model_input=types.ModelInput.from_ints(full[:-1]),
        loss_fn_inputs={
            "target_tokens": full[1:],
            "logprobs": [0.0] * n_prompt + list(sequence.logprobs),
            "advantages": [0.0] * n_prompt + [advantage] * len(sequence.tokens),
        },
    )


def train_arm(service, arm, train_problems, eval_problems, budget_seconds, eval_before):
    """Train one arm; its wall budget starts once the training client is ready (trainer cold start is not charged)."""
    name = arm["name"]
    training = service.create_lora_training_client(base_model=CONFIG["base_model"], rank=CONFIG["lora_rank"])
    tokenizer = training.get_tokenizer()
    deadline = time.monotonic() + budget_seconds
    print(f"[{name}] training client ready after {elapsed():.0f}s; {budget_seconds:.0f}s budget for this arm", flush=True)
    if eval_before is None:
        eval_before = eval_accuracy(training, tokenizer, eval_problems)
        print(f"[{name}] eval accuracy before training: {eval_before:.3f} ({elapsed():.0f}s)", flush=True)

    n_prompts, group = arm["n_prompts_per_step"], arm["group_size"]
    rewards_per_step, completion_lengths = [], []
    cursor = 0
    for step in range(CONFIG["steps"]):
        if time.monotonic() > deadline:
            print(f"[{name}] stopping at step {step}: wall-time budget reached", flush=True)
            break
        step_started = time.monotonic()
        batch = [train_problems[(cursor + i) % len(train_problems)] for i in range(n_prompts)]
        cursor += n_prompts
        prompts = [prompt_tokens(tokenizer, q) for q, _ in batch]
        sampler = training.save_weights_and_get_sampling_client()
        groups = sample_all(sampler, tokenizer, prompts, group, CONFIG["temperature"])
        sampled = time.monotonic()

        rollouts, questions, rewards = [], [], []
        for (question, gold), ids, sequences in zip(batch, prompts, groups):
            for sequence in sequences:
                rollouts.append((ids, sequence))
                questions.append(question)
                rewards.append(grade(tokenizer.decode(sequence.tokens, skip_special_tokens=True), gold))
                completion_lengths.append(len(sequence.tokens))
        advantages = advantages_for(arm, questions, rewards)
        datums = [make_datum(ids, seq, adv) for (ids, seq), adv in zip(rollouts, advantages) if adv != 0.0]
        mean_reward = sum(rewards) / len(rewards)
        rewards_per_step.append(mean_reward)
        if datums:
            output = training.forward_backward(datums, "importance_sampling").result()
            training.optim_step(types.AdamParams(learning_rate=CONFIG["lr"])).result()
            loss = output.metrics.get("loss:sum")
        else:
            loss = None
        print(
            f"[{name}] step {step}: mean reward {mean_reward:.3f} over {len(rewards)} rollouts, "
            f"{len(datums)} with nonzero advantage, loss:sum {loss}, "
            f"sampling {sampled - step_started:.1f}s, total {time.monotonic() - step_started:.1f}s",
            flush=True,
        )

    eval_after = eval_accuracy(training, tokenizer, eval_problems)
    print(f"[{name}] eval accuracy after training: {eval_after:.3f} ({elapsed():.0f}s)", flush=True)
    return {
        "train_reward_first": rewards_per_step[0] if rewards_per_step else None,
        "train_reward_last": rewards_per_step[-1] if rewards_per_step else None,
        "train_reward_mean": sum(rewards_per_step) / len(rewards_per_step) if rewards_per_step else None,
        "eval_accuracy_before": eval_before,
        "eval_accuracy_after": eval_after,
        "steps": len(rewards_per_step),
        "n_prompts_per_step": n_prompts,
        "group_size": group,
        "rollouts_per_step": n_prompts * group,
        "baseline": arm["baseline"],
        "mean_completion_tokens": sum(completion_lengths) / len(completion_lengths) if completion_lengths else None,
    }


def main():
    service = tinker.ServiceClient(
        base_url=os.environ["TINKER_BASE_URL"],
        api_key=os.environ["TINKER_API_KEY"],
        default_headers_overrides={
            "Modal-Key": os.environ["MODAL_PROXY_TOKEN_ID"],
            "Modal-Secret": os.environ["MODAL_PROXY_TOKEN_SECRET"],
        },
    )
    rollouts = max(arm["n_prompts_per_step"] * arm["group_size"] for arm in CONFIG["arms"])
    n_train = max(arm["n_prompts_per_step"] for arm in CONFIG["arms"]) * CONFIG["steps"]
    train_problems = load_problems(CONFIG["split"], n_train)
    eval_problems = load_problems(CONFIG["eval_split"], CONFIG["n_eval"])
    print(f"loaded {len(train_problems)} train / {len(eval_problems)} eval problems; up to {rollouts} rollouts per step", flush=True)

    arms = {}
    eval_before = None
    per_arm_seconds = CONFIG["max_wall_seconds"] / len(CONFIG["arms"])
    for arm in CONFIG["arms"]:
        arms[arm["name"]] = train_arm(service, arm, train_problems, eval_problems, per_arm_seconds, eval_before)
        eval_before = arms[arm["name"]]["eval_accuracy_before"]

    last = arms[CONFIG["arms"][-1]["name"]]
    result = {
        **last,
        "arms": arms,
        "base_model": CONFIG["base_model"],
        "lora_rank": CONFIG["lora_rank"],
        "lr": CONFIG["lr"],
        "max_tokens": CONFIG["max_tokens"],
        "dataset": CONFIG["dataset"],
        "n_eval": len(eval_problems),
        "wall_seconds": round(elapsed(), 1),
    }
    print("RESULT " + json.dumps(result), flush=True)


if __name__ == "__main__":
    main()

lora_rank

32

lr

4e-05

max_tokens

512

n_eval

256

wall_seconds

1,637.4

baseline

group_mean

base_model

Qwen/Qwen3.5-2B

dataset

openai/gsm8k

command

python experiment_run_bffef575d5.py

config

{'base_model': 'Qwen/Qwen3.5-2B', 'lora_rank': 32, 'dataset': 'openai/gsm8k', 'dataset_config': 'main', 'split': 'train', 'eval_split': 'test', 'n_eval': 256, 'steps': 100, 'lr': 4e-05, 'max_tokens': 512, 'temperature': 1.0, 'similarity_temperature': 0.1, 'max_wall_seconds': 2700, 'arms': [{'name': 'grpo', 'baseline': 'group_mean', 'n_prompts_per_step': 16, 'group_size': 8}]}

modal_call_id

fc-01M6FQEKREWRWXA5YMTSW9G3G3

script_source

default script (Claude script rejected: contains 'import torch')