🔬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_rank

32

custom_probe_datums

9

custom_probe_backward_seconds

50.68

custom_probe_grad_norm

84.78

packed_probe_grad_norm

82.7

comparison_steps_completed

0

evaluation_examples_completed

0