import argparse import dataclasses import asyncio import hashlib import json import math import os from pathlib import Path import shutil import time import numpy as np import tinker import torch from tinker_cookbook.completers import TinkerTokenCompleter from tinker_cookbook.recipes.math_rl import math_env, math_grading from tinker_cookbook.rl.data_processing import assemble_training_data, compute_advantages from tinker_cookbook.rl.rollouts import do_group_rollout ROOT = Path(__file__).resolve().parent MODEL = "Qwen/Qwen3.5-2B" RENDERER = "qwen3_5_disable_thinking" def kernel_baselines(features, rewards, current_count, bandwidth=0.5, threshold=1e-3): x = np.asarray(features, dtype=np.float64) r = np.asarray(rewards, dtype=np.float64) if x.shape != (len(r), 2) or len(r) < 2 or not 0 < current_count <= len(r): raise ValueError("invalid window") if not np.isfinite(x).all() or not np.isfinite(r).all() or not np.isfinite(bandwidth) or bandwidth <= 0: raise ValueError("invalid numbers") scale = x.std(axis=0) z = (x - x.mean(axis=0)) / np.where(scale > 0, scale, 1) start = len(r) - current_count distances = ((z[start:, None, :] - z[None, :, :]) ** 2).sum(axis=2) weights = np.exp(-distances / (2 * bandwidth**2)) weights[np.arange(current_count), np.arange(start, len(r))] = 0 mass = weights.sum(axis=1) fallback = (r.sum() - r[start:]) / (len(r) - 1) baseline = np.divide(weights @ r, mass, out=fallback.copy(), where=mass >= threshold) square = (weights**2).sum(axis=1) neighbors = np.divide(mass**2, square, out=np.zeros_like(mass), where=square > 0) return baseline, neighbors def advantages(groups, arm, features, history): rewards = np.array([r for g in groups for r in g.get_total_rewards()], dtype=np.float64) if arm == "grpo_gs8": result = compute_advantages(groups) flat = np.concatenate([a.numpy().astype(np.float64) for a in result]) return result, rewards - flat, np.zeros(len(rewards)) window = (history + list(zip(features, rewards)))[-256:] baseline, neighbors = kernel_baselines([x for x, _ in window], [r for _, r in window], len(rewards)) flat = rewards - baseline result, offset = [], 0 for group in groups: size = len(group.trajectories_G) result.append(torch.tensor(flat[offset:offset + size], dtype=torch.float64)) offset += size return result, baseline, neighbors def guard_grader(): original = math_grading._is_int math_grading._is_int = lambda x: math.isfinite(x) and original(x) def service(): return tinker.ServiceClient(base_url=os.environ["TINKER_BASE_URL"], default_headers_overrides={ "Modal-Key": os.environ["MODAL_PROXY_TOKEN_ID"], "Modal-Secret": os.environ["MODAL_PROXY_TOKEN_SECRET"], }) async def dataset(batch_size, group_size): train, test = await math_env.Gsm8kDatasetBuilder(batch_size=batch_size, group_size=group_size, model_name_for_tokenizer=MODEL, renderer_name=RENDERER, seed=0)() train.ds = train.ds.select(range(2560)) return train, test async def sample_batch(sampler, builders, temperature=1.0): async def prompt_feature(builder): env = builder.env_thunk() prompt, _ = await env.initial_observation() logprobs = await sampler.compute_logprobs_async(prompt) finite = [x for x in logprobs if x is not None] if not finite or not np.isfinite(finite).all(): raise ValueError("missing or nonfinite prompt logprobs") return [math.log(prompt.length), float(np.mean(finite))] features = await asyncio.gather(*(prompt_feature(b) for b in builders)) policy = TinkerTokenCompleter(sampler, max_tokens=512, temperature=temperature) groups = await asyncio.gather(*(do_group_rollout(b, policy) for b in builders)) expanded = [f for f, g in zip(features, groups) for _ in g.trajectories_G] return groups, expanded def datums(groups, adv): data, metadata = assemble_training_data(groups, adv, include_provenance=False) return [tinker.Datum(model_input=d.model_input, loss_fn_inputs={k: v for k, v in d.loss_fn_inputs.items() if k != "mask"}) for d in data], metadata def serialized(value): if isinstance(value, tinker.ModelInput): return dict(chunks=[dict(tokens=value.to_ints())]) if isinstance(value, tinker.TensorData): return dict(data=value.data, dtype=value.dtype, shape=value.shape, sparse_crow_indices=value.sparse_crow_indices, sparse_col_indices=value.sparse_col_indices) if dataclasses.is_dataclass(value): return {field.name: serialized(getattr(value, field.name)) for field in dataclasses.fields(value)} if isinstance(value, dict): return {key: serialized(item) for key, item in value.items()} if isinstance(value, (list, tuple)): return [serialized(item) for item in value] if isinstance(value, (np.ndarray, np.generic)): return value.tolist() if hasattr(value, "model_dump"): return value.model_dump(mode="json", exclude_unset=True) return value def write_json(path, obj): path.write_text(json.dumps(serialized(obj), indent=2, allow_nan=False) + "\n") async def train(arm): group_size = 8 if arm == "grpo_gs8" else 1 ds, _ = await dataset(64 // group_size, group_size) history = [] persistent = ROOT / "logs" / arm transient = Path("/tmp/rx-gs1") / arm for p in (persistent, transient): p.mkdir(parents=True, exist_ok=True) sc = service() try: checkpoint = json.loads((ROOT / "init-save-response.json").read_text())["path"] tc = await sc.create_training_client_from_state_async(checkpoint, base_model=MODEL, optimizer=tinker.AdamOptimizerConfig()) for step in range(40): begin = time.monotonic() sampler = await tc.save_weights_and_get_sampling_client_async() sampled = time.monotonic() groups, features = await sample_batch(sampler, ds.get_batch(step)) rollout_end = time.monotonic() adv, baseline, neighbors = advantages(groups, arm, features, history) rewards = [r for g in groups for r in g.get_total_rewards()] history = (history + list(zip(features, rewards)))[-256:] data, _ = datums(groups, adv) fb_future = await tc.forward_backward_async(data, loss_fn="importance_sampling") fb = await fb_future.result_async() optim_future = await tc.optim_step_async(tinker.AdamParams(learning_rate=4e-5, beta1=0.9, beta2=0.95, eps=1e-8)) await optim_future.result_async() flat_adv = np.concatenate([a.numpy() for a in adv]) record = dict(step=step, rewards=rewards, executed_group_size=[len(g.trajectories_G) for g in groups], features=features, kernel_effective_neighbors=neighbors.tolist(), b_i=baseline.tolist(), advantages=flat_adv.tolist(), frac_nonzero_advantage=float(np.mean(np.abs(flat_adv) > 1e-6)), timings=dict(save=sampled-begin, rollout=rollout_end-sampled, update=time.monotonic()-rollout_end), metrics=fb.metrics) path = transient / f"step-{step:03d}.json" write_json(path, record) shutil.copy2(path, persistent / path.name) saved_future = await tc.save_weights_for_sampler_async(name=f"{arm}-final") saved = await saved_future.result_async() write_json(persistent / "final.json", dict(path=saved.path)) finally: await sc.close("success") def main(): parser = argparse.ArgumentParser() parser.add_argument("--arm", choices=["grpo_gs8", "gs1_kernel"], default="gs1_kernel") parser.add_argument("command", choices=["preflight", "train"]) args = parser.parse_args() guard_grader() if args.command == "preflight": from run_preflight import preflight asyncio.run(preflight()) else: asyncio.run(train(args.arm)) if __name__ == "__main__": main()