"""GRPO on GSM8K via forward_backward_custom; CENTER=1 adds the score-centering correction (martin-marek/score-centering@fdca22b README score_centering_loss), composed with IS (weight_fn = identity).""" import json, os, sys, time, random, hashlib, concurrent.futures as cf import numpy as np, torch, tinker from tinker import types from common import service_client import gsm8k_eval as ge MODEL = ge.MODEL GROUP, GROUPS, MAX_TOKENS, K, LR, RANK = 8, 32, 512, 128, 1e-4, 32 EPS = 1e-6 def weight_fn(name): return {"identity": lambda r: r, "ones": torch.ones_like}[name] def sc_terms(lp, d, wname): """lp [N, K+1] trainer logprobs (col 0 = sampled token, cols 1.. = sampler top-k ids). Paper formula.""" w = weight_fn(wname) token_logp, head_logp = lp[:, 0], lp[:, 1:] samp_tok = d["samp_lp"]; samp_head = d["samp_head"]; hm = d["head_mask"] p_head = torch.exp(head_logp) * hm; q_head = torch.exp(samp_head) * hm p_tail = torch.clamp(1 - p_head.sum(-1), min=EPS); q_tail = torch.clamp(1 - q_head.sum(-1), min=EPS) rho = q_tail / p_tail alpha = rho * w(1 / rho) head_weight = w(torch.exp(torch.where(hm > 0, head_logp - samp_head, torch.zeros_like(head_logp)))) # empty slots hold -1e9 residual = (q_head * head_weight - alpha[:, None] * p_head) * hm correction = (residual.detach() * head_logp).sum(-1) token_weight = w(torch.exp(token_logp - samp_tok)) return token_logp, token_weight, correction, residual def make_loss(center, store=None): def loss_fn(data, logprobs_list): total = 0.0; n = 0; stats = {"sc_abs_corr_sum": 0.0, "sc_vanilla_abs_corr_sum": 0.0, "resid_absmax": 0.0, "gap_sum": 0.0, "absgap_sum": 0.0, "ent_sum": 0.0} gaps = [] for datum, lp in zip(data, logprobs_list): d = {k: v.to_torch() for k, v in datum.loss_fn_inputs.items() if k not in ("target_tokens", "weights")} m = d["mask"]; A = d["adv"] lp32 = lp.float() if store is not None: store.append(lp32.detach().clone()) token_logp, tw, corr, resid = sc_terms(lp32, d, "identity") if center: tok = -A * (tw.detach() * token_logp - corr) # paper: -sg(A)*(sg(w)*logp - correction) else: tok = -A * torch.exp(token_logp - d["samp_lp"]) # stock importance_sampling surrogate total = total + (tok * m).sum() with torch.no_grad(): _, _, corr_v, _ = sc_terms(lp32, d, "ones") n += m.sum().item() stats["sc_abs_corr_sum"] += (corr.abs() * m).sum().item() stats["sc_vanilla_abs_corr_sum"] += (corr_v.abs() * m).sum().item() stats["resid_absmax"] = max(stats["resid_absmax"], (resid.abs() * m[:, None]).max().item()) g = (d["samp_lp"] - token_logp)[m > 0] gaps.append(g) stats["gap_sum"] += g.sum().item(); stats["absgap_sum"] += g.abs().sum().item() q = torch.exp(d["samp_head"]) * d["head_mask"]; qt = torch.clamp(1 - q.sum(-1), min=EPS) ent = -(q * torch.where(q > 0, torch.log(q), torch.zeros_like(q))).sum(-1) - qt * torch.log(qt) stats["ent_sum"] += (ent * m).sum().item() n = max(n, 1) ga = torch.cat(gaps).abs() if gaps else torch.zeros(1) metrics = {"n_tokens": float(n), "sc_abs_correction_mean": stats["sc_abs_corr_sum"] / n, "sc_vanilla_abs_correction_mean": stats["sc_vanilla_abs_corr_sum"] / n, "sc_identity_residual_absmax": stats["resid_absmax"], "logprob_gap_mean": stats["gap_sum"] / n, "logprob_absgap_mean": stats["absgap_sum"] / n, "logprob_absgap_p99": float(torch.quantile(ga.double(), 0.99)), "logprob_absgap_max": float(ga.max()), "sampler_topk_entropy_lb": stats["ent_sum"] / n} return total, metrics return loss_fn def topk_for(samp, full_tokens): params = types.SamplingParams(max_tokens=1, temperature=1.0) for attempt in range(6): try: r = samp.sample(types.ModelInput.from_ints(full_tokens), 1, params, include_prompt_logprobs=True, topk_prompt_logprobs=K).result() return r.topk_prompt_logprobs, r.prompt_logprobs except Exception as e: err = e; time.sleep(2 ** attempt) raise err def rollouts(samp, problems, pool): params = types.SamplingParams(max_tokens=MAX_TOKENS, temperature=1.0, stop=ge.STOP) def group(p): pt = ge.prompt_tokens(p["question"]) for attempt in range(6): try: r = samp.sample(types.ModelInput.from_ints(pt), GROUP, params).result(); break except Exception as e: err = e; time.sleep(2 ** attempt) else: raise err ref = ge.reference_answer(p["answer"]) seqs = [] for s in r.sequences: txt = ge.tokenizer().decode(s.tokens, skip_special_tokens=True) seqs.append({"prompt": pt, "tokens": list(s.tokens), "logprobs": list(s.logprobs), "reward": ge.is_correct(ge.parse_answer(txt), ref), "stop": str(s.stop_reason)}) return seqs groups = list(pool.map(group, problems)) return groups def build_datums(samp, groups, pool): flat = [] for g in groups: mean = sum(s["reward"] for s in g) / len(g) for s in g: s["adv"] = s["reward"] - mean if s["adv"] != 0 and len(s["tokens"]) > 0: flat.append(s) tops = list(pool.map(lambda s: topk_for(samp, s["prompt"] + s["tokens"]), flat)) datums, sampler_consistency = [], [] for s, (topk, plp) in zip(flat, tops): P, C = len(s["prompt"]), len(s["tokens"]) full = s["prompt"] + s["tokens"]; N = len(full) - 1 tgt = np.zeros((N, K + 1), dtype=np.int64); tgt[:, 0] = full[1:] tgt[:, 1:] = np.array(full[1:])[:, None] wts = np.zeros((N, K + 1), dtype=np.float32) head = np.full((N, K), -1e9, dtype=np.float32); hm = np.zeros((N, K), dtype=np.float32) slp = np.zeros(N, dtype=np.float32); mask = np.zeros(N, dtype=np.float32) for i in range(C): t = P + i - 1 entry = topk[P + i] or [] for j, (tid, l) in enumerate(entry[:K]): tgt[t, 1 + j] = tid; head[t, j] = l; hm[t, j] = 1.0 wts[t, :] = 1.0; slp[t] = s["logprobs"][i]; mask[t] = 1.0 if plp is not None and plp[P + i] is not None: sampler_consistency.append(abs(plp[P + i] - s["logprobs"][i])) datums.append(types.Datum(model_input=types.ModelInput.from_ints(full[:-1]), loss_fn_inputs={ "target_tokens": types.TensorData.from_numpy(tgt), "weights": types.TensorData.from_numpy(wts), "samp_lp": types.TensorData.from_numpy(slp), "samp_head": types.TensorData.from_numpy(head), "head_mask": types.TensorData.from_numpy(hm), "mask": types.TensorData.from_numpy(mask), "adv": types.TensorData.from_numpy(np.full(N, s["adv"], dtype=np.float32))})) return datums, flat, sampler_consistency def batches(seed, steps): train = ge.load("gsm8k_train.jsonl") idx = list(range(len(train))); random.Random(seed).shuffle(idx) return [[train[j] for j in idx[GROUPS * i: GROUPS * (i + 1)]] for i in range(steps)] def create_client(sc): delay = 60 t0 = time.time() while True: try: return sc.create_lora_training_client(base_model=MODEL, rank=RANK) except Exception as e: print(json.dumps({"create_client_error": repr(e)[:400], "elapsed": time.time() - t0}), flush=True) if time.time() - t0 > 900: raise time.sleep(delay); delay = min(delay * 2, 240) def train(center, seed, steps, log_dir, eval_every=0, sc=None, pool=None): sc = sc or service_client(); pool = pool or cf.ThreadPoolExecutor(64) os.makedirs(log_dir, exist_ok=True) tc = create_client(sc) loss_fn = make_loss(center) log = open(os.path.join(log_dir, "steps.jsonl"), "a") consumed = [] for i, probs in enumerate(batches(seed, steps)): t0 = time.time() samp = tc.save_weights_and_get_sampling_client(name=f"rx-b40a-c{center}-s{seed}-step{i}") t1 = time.time() groups = rollouts(samp, probs, pool); t2 = time.time() datums, flat, cons = build_datums(samp, groups, pool); t3 = time.time() consumed.append(",".join(p["id"] for p in probs)) rec = {"step": i, "center": center, "seed": seed, "batch_ids": consumed[-1], "mean_reward": float(np.mean([s["reward"] for g in groups for s in g])), "n_datums": len(datums), "mean_completion_tokens": float(np.mean([len(s["tokens"]) for g in groups for s in g])), "sampler_prompt_vs_sample_logprob_absdiff_max": float(max(cons)) if cons else None} if datums: fb = tc.forward_backward_custom(datums, loss_fn).result(); t4 = time.time() opt = tc.optim_step(types.AdamParams(learning_rate=LR)).result(); t5 = time.time() rec.update({"fb_metrics": fb.metrics, "optim_metrics": opt.metrics}) else: t4 = t5 = time.time() rec.update({"t_save": t1 - t0, "t_rollout": t2 - t1, "t_topk": t3 - t2, "t_fb": t4 - t3, "t_opt": t5 - t4, "t_step": t5 - t0}) log.write(json.dumps(rec) + "\n"); log.flush(); print(json.dumps(rec)[:1500], flush=True) return tc, consumed if __name__ == "__main__": center, seed, steps = int(os.environ["CENTER"]), int(os.environ["SEED"]), int(os.environ.get("STEPS", 100)) train(center, seed, steps, f"/tmp/rx-sc-{center}-{seed}")