"""Restoration preflight, correctness fixtures/parity, and 3-step pilots of both arms.""" import asyncio, json, pickle, sys, time import numpy as np, torch, tinker from tinker import types from datasets import load_dataset from tinker_cookbook.completers import TinkerTokenCompleter from tinker_cookbook.rl.rollouts import do_group_rollout from tinker_cookbook.rl.data_processing import compute_advantages, assemble_training_data from tinker_cookbook.rl.train import _remove_mask from tinker_cookbook.recipes.math_rl.math_env import Gsm8kDataset, MathEnv from common import * import gs1, gs1_reference PILOT_STEPS = int(sys.argv[1]) if len(sys.argv) > 1 else 3 log_f = open(f"{OUT}/preflight.log", "a") def log(*a): s = " ".join(str(x) for x in a) print(s, flush=True); log_f.write(s + "\n"); log_f.flush() def td(x): return np.asarray(x.to_numpy() if hasattr(x, "to_numpy") else x, dtype=np.float64) def info_dict(info): oc = getattr(info, "optimizer_config", None) return {"model_name": getattr(info, "model_name", None), "lora_rank": getattr(info, "lora_rank", None), "is_lora": getattr(info, "is_lora", None), "optimizer_config": oc.model_dump() if hasattr(oc, "model_dump") else (str(oc) if oc is not None else None)} def fixed_batch(tok, ds): data, toks = [], [] for row in ds.select(range(8)): t = tok.encode(row["question"] + "\n" + row["answer"]) toks.append(t) data.append(types.Datum(model_input=types.ModelInput.from_ints(t[:-1]), loss_fn_inputs={"target_tokens": t[1:], "weights": [1.0] * (len(t) - 1)})) return data, sha256_json(toks) async def initial_output(tc, data): r = await (await tc.forward_async(data, "cross_entropy")).result_async() lp = np.concatenate([np.asarray(o["logprobs"].to_numpy(), np.float32) for o in r.loss_fn_outputs]) return lp, sha256_bytes(lp.tobytes()) def is_terms(lp, old, adv): ratio = np.exp(lp - old) return -(ratio * adv), -(ratio * adv) # per-token loss, d loss / d logprob def parity(ref, act): ref, act = np.asarray(ref, np.float64), np.asarray(act, np.float64) d = act - ref cos = float(ref @ act / (np.linalg.norm(ref) * np.linalg.norm(act))) if np.linalg.norm(ref) and np.linalg.norm(act) else (1.0 if not d.any() else 0.0) return {"max_abs_error": float(np.abs(d).max()), "reference_max_abs": float(np.abs(ref).max()), "relative_l2_error": float(np.linalg.norm(d) / max(np.linalg.norm(ref), 1e-300)), "cosine": cos, "n": int(ref.size)} async def correctness(P, arm, groups, data, prompts, corpus): """Forward fixed step-0 rollouts once through the untrained init client; compare loss/grad vectors.""" fixture = [[(tr.transitions[0].ob.to_ints(), tr.transitions[0].ac.tokens, list(tr.transitions[0].ac.logprobs)) for tr in g.trajectories_G] for g in groups] rewards = [g.get_total_rewards() for g in groups] fixed_hash = sha256_json({"fixture": fixture, "rewards": rewards, "prompts": prompts}) if arm == "grpo_gs8_seed0": ref_data, _ = assemble_training_data(groups, compute_advantages(groups)) ref_tok_adv = [td(d.loss_fn_inputs["advantages"]) for d in ref_data] label = "baseline_vs_stock_cookbook_compute_advantages" else: a_ref = gs1_reference.gs1_advantages([r[0] for r in rewards], prompts, corpus, TAU) ref_data = data ref_tok_adv = [td(d.loss_fn_inputs["mask"]) * a_ref[i] for i, d in enumerate(data)] label = "gs1_vs_float64_reference" for a, b in zip(ref_data, data): assert a.model_input.to_ints() == b.model_input.to_ints() assert np.array_equal(td(a.loss_fn_inputs["logprobs"]), td(b.loss_fn_inputs["logprobs"])) fr = await (await P.forward_async([_remove_mask(d) for d in ref_data], "importance_sampling")).result_async() fa = await (await P.forward_async([_remove_mask(d) for d in data], "importance_sampling")).result_async() lp = [td(o["logprobs"]) for o in fr.loss_fn_outputs] old = [td(d.loss_fn_inputs["logprobs"]) for d in data] act_adv = [td(d.loss_fn_inputs["advantages"]) for d in data] Lr, Gr, La, Ga = [], [], [], [] for i in range(len(data)): lr_, gr_ = is_terms(lp[i], old[i], ref_tok_adv[i]); la_, ga_ = is_terms(lp[i], old[i], act_adv[i]) Lr.append(lr_.sum()); La.append(la_.sum()); Gr.append(gr_); Ga.append(ga_) res = {"path": label, "arm": arm, "fixed_input_sha256": fixed_hash, "fixed_input": f"step-0 pilot rollouts of {arm} ({len(data)} datums) sampled from shared init weights; one server forward (importance_sampling) on the never-updated init client for logprobs", "precision": "float64 local loss/grad over server float32 logprobs; advantages sent as float32", "normalization": "importance_sampling sum over tokens per datum (Tinker loss:sum), no per-batch normalization", "loss_total": {"reference": float(np.sum(Lr)), "actual": float(np.sum(La))}, "loss_per_datum": parity(Lr, La), "gradient_wrt_logprobs": parity(np.concatenate(Gr), np.concatenate(Ga)), "server_loss_sum": {"reference": fr.metrics.get("loss:sum"), "actual": fa.metrics.get("loss:sum")}} log("CORRECTNESS", json.dumps(res)) with open(f"{OUT}/fixture_{arm}.pkl", "wb") as f: pickle.dump({"fixture": fixture, "rewards": rewards, "prompts": prompts}, f) return res async def pilot_arm(arm, tc, P, ds, kernel, corpus, tok, results): cfg = ARMS[arm] rec = results[arm] = {"steps": []} for step in range(PILOT_STEPS): t0 = time.time() sc = await tc.save_weights_and_get_sampling_client_async(name=f"{RUN}-{arm}-pilot-s{step}") policy = TinkerTokenCompleter(sc, max_tokens=MAX_TOKENS, temperature=TEMPERATURE) builders = ds.get_batch(step) prompts = [b.env_thunk.args[0] for b in builders] groups = await asyncio.gather(*[do_group_rollout(b, policy) for b in builders]) t_roll = time.time() - t0 rewards = [g.get_total_rewards() for g in groups] s = {"step": step, "executed_group_sizes": sorted({len(g.trajectories_G) for g in groups}), "rollouts": sum(len(g.trajectories_G) for g in groups), "n_prompts": len(groups), "mean_reward": float(np.mean([x for r in rewards for x in r])), "frac_correct": float(np.mean([x > 0.5 for r in rewards for x in r])), "rollout_errors": sum(len(g.rollout_errors or []) for g in groups)} if arm == "gs1_seed0": adv, diag = gs1.gs1_advantages([r[0] for r in rewards], prompts, kernel) adv_P = [torch.tensor([a], dtype=torch.float32) for a in adv] s.update(n_eff=diag["n_eff"].tolist(), abs_baseline_minus_uniform_loo=np.abs(diag["baseline"] - diag["uniform_loo"]).tolist()) else: adv_P = [torch.tensor(a, dtype=torch.float32) for a in gs1.grpo_advantages(rewards)] data, _ = assemble_training_data(groups, adv_P) s["n_datums"] = len(data) s["rollout_records"] = [{"prompt": prompts[i], "rewards": rewards[i], "responses": [tok.decode(tr.transitions[0].ac.tokens) for tr in g.trajectories_G]} for i, g in enumerate(groups)] if step == 0: rec["correctness"] = await correctness(P, arm, groups, data, prompts, corpus) t1 = time.time() fb = await tc.forward_backward_async([_remove_mask(d) for d in data], "importance_sampling") op = await tc.optim_step_async(ADAM) fbr, opr = await fb.result_async(), await op.result_async() s.update(train_seconds=time.time() - t1, rollout_seconds=t_roll, step_seconds=time.time() - t0, fb_metrics=dict(fbr.metrics), optim_metrics=dict(opr.metrics or {})) log(arm, "step", step, {k: v for k, v in s.items() if k not in ("rollout_records", "n_eff", "abs_baseline_minus_uniform_loo")}) rec["steps"].append(s) dump("pilot_results.json", results) async def main(): t_start = time.time() renderer, tok = renderer_and_tokenizer() train = load_dataset("openai/gsm8k", "main", split="train") corpus = list(train["question"]) kernel = gs1.TfidfKernel(corpus, TAU) svc = service_client() P = await with_capacity_retry(lambda: svc.create_lora_training_client_async(base_model=MODEL, rank=LORA_RANK), log=log) init_path = (await (await P.save_state_async(name=f"{RUN}-init")).result_async()).path log("init state", init_path) arms = {} for arm in ARMS: arms[arm] = await with_capacity_retry(lambda: svc.create_training_client_from_state_async(init_path), log=log) fb_data, fb_hash = fixed_batch(tok, train) lpP, hP = await initial_output(P, fb_data) prov = {"version": 1, "arms": []} for arm, tc in arms.items(): lp, h = await initial_output(tc, fb_data) info = info_dict(await tc.get_info_async()) prov["arms"].append({ "arm": arm, "base_checkpoint": MODEL, "init_checkpoint": init_path, "init_method": "create_lora_training_client(rank=32; server rejects per-model seeds, so seed 0 not applied) -> save_state -> create_training_client_from_state (fresh optimizer state)", "adapter_sha256": None, "evidence_ref": "out/restoration.json", "before_first_update": True, "restore_count": 1, "optimizer_checks": [{"event": "restore_from_init_state", "expected": {"family": "adam", "lr": LR, "beta1": 0.9, "beta2": 0.95, "eps": 1e-8, "state": "fresh"}, "actual": info["optimizer_config"], "source": "TrainingClient.get_info().optimizer_config (server)" if info["optimizer_config"] is not None else "server did not report optimizer config", "verified": None}], "initial_outputs": {"fixed_batch": "first 8 GSM8K train rows (question+answer), cross_entropy forward", "fixed_batch_sha256": fb_hash, "logprobs_sha256": h, "max_abs_diff_vs_init_client": float(np.abs(lp.astype(np.float64) - lpP).max()), "init_client_logprobs_sha256": hP}, "server_info": info}) dump("restoration.json", prov) log("restoration", json.dumps([{k: a[k] for k in ("arm", "initial_outputs", "server_info")} for a in prov["arms"]])) results = {"init_path": init_path} dss = {arm: Gsm8kDataset(batch_size=c["groups"], group_size=c["group_size"], renderer=renderer, convo_prefix=MathEnv.standard_fewshot_prefix(), split="train", seed=0) for arm, c in ARMS.items()} await asyncio.gather(*[pilot_arm(arm, arms[arm], P, dss[arm], kernel, corpus, tok, results) for arm in ARMS]) results["wall_seconds"] = time.time() - t_start dump("pilot_results.json", results) for arm, tc in arms.items(): p = (await (await tc.save_state_async(name=f"{RUN}-{arm}-pilot-final")).result_async()).path results.setdefault("pilot_final_states", {})[arm] = p dump("pilot_results.json", results) log("DONE wall", results["wall_seconds"]) asyncio.run(main())