"""SFT Product-based Martingale Training""" import os import sys import json import datetime import wandb # Multiprocessing spawn method set at module level, before any CUDA ops import torch as t print(f"After imports torch: CUDA = {t.cuda.is_initialized()}") # Add project root to Python path sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # Math # Model & Dataset import multiprocessing multiprocessing.set_start_method('spawn') # Set before any CUDA ops #from trl import SFTConfig, SFTTrainer #print(f"After trl: CUDA = {t.cuda.is_initialized()}") from peft import get_peft_model, LoraConfig print(f"After peft: CUDA = {t.cuda.is_initialized()}") from torch.utils.data import Dataset, DataLoader print(f"After dataloader: CUDA = {t.cuda.is_initialized()}") from transformers import (AutoModelForCausalLM, AutoTokenizer, ) print(f"After imports transformer: CUDA = {t.cuda.is_initialized()}") # from unsloth import FastLanguageModel # Commented out - using standard transformers to avoid in-place ops # Martingale Modules from scripts.run_reasoning import run_reasoning from core.algo.martingale import MartingaleStrategy from core.algo.accuracy import GroundTruthAccuracy from core.algo.sycoreason import SycophanticReasoning from core.domain.forecasting import Forecasting from core.domain.changemyview import CMV from core.domain.openreview import OpenReview from core.domain.research import Research from core.policy.localmodel import LocalModel #from core.policy.apimodel import APIModel # We do not use API models in training. from scripts.setup import setup from core.reasoning.debate import SelfDebate from core.reasoning.cot import ChainOfThought # Utils from utils.io_utils import dump_file, load_file, complete_path, load_file_for_run, dump_file_for_run, complete_create_dir_path from utils.cuda_utils import check_cuda_status from training.convert_sft_data import run_convert_sft_data # NEP to understand what precision is desirable. (someting like pf16) # Env variables # A custom Dataset for our purpose here class SFTDataset(Dataset): def __init__(self, processed_data, tokenized_inputs, question_lengths): self.processed_data = processed_data # List of (actual_delta, prior) tuples self.tokenized_inputs = tokenized_inputs # Dictionary of tensors self.question_lengths = question_lengths # List of ints: #tokens in question alone def __len__(self): return len(self.processed_data) def __getitem__(self, idx): # Get the tuple (actual_delta, prior) actual_delta, prior, y = self.processed_data[idx] # Get the corresponding tokenized tensors for this index item = { "actual_delta": actual_delta, "prior": prior, "y": y, "input_ids": self.tokenized_inputs["input_ids"][idx], "attention_mask": self.tokenized_inputs["attention_mask"][idx], "question_length": self.question_lengths[idx], # int: where answer tokens start } return item def data_preprocessing(sft_data, tokenizer): """Prep dataloader for finetuning and removing text & dict from it for now.""" priors = [item["prior"] for item in sft_data] actual_deltas = [item["actual_delta"] for item in sft_data] ys = [float(item.get("y", item.get("y_resolve", 0.0))) for item in sft_data] processed_data = list(zip(actual_deltas, priors, ys)) reasoning_texts = [f'{item["question+context"]}\n\n---\n\n{item["answer"]}' for item in sft_data] # Use the tokenizer to convert the text to tensors tokenizer.pad_token = tokenizer.eos_token tokenizer.truncation_side = "left" # truncate questions first while preserving answers. tokenized_inputs = tokenizer( reasoning_texts, padding=True, # Pad shorter texts to the length of the longest one in the batch truncation=True, # Truncate texts that are longer than the model's max length max_length=int(os.getenv("TRAIN_CONTEXT_LEN", 4092)), return_tensors="pt", return_offsets_mapping=True, ) offset_mapping = tokenized_inputs.pop("offset_mapping") # [N, seq_len, 2]; remove before model use # Find where answer tokens start using character offsets (robust against BPE context merges). separator = "\n\n---\n\n" question_lengths = [] for text, offsets in zip(reasoning_texts, offset_mapping): sep_char_pos = text.find(separator) if sep_char_pos == -1: question_lengths.append(0) continue answer_start_char = sep_char_pos + len(separator) found_idx = 0 for tok_idx, (char_start, char_end) in enumerate(offsets.tolist()): if char_start >= answer_start_char and char_end > char_start: found_idx = tok_idx break question_lengths.append(found_idx) dataset = SFTDataset(processed_data, tokenized_inputs, question_lengths) # NEP later you want need to look into data manually to check out belief eval results (e.g., whether the first belief should be taken as prior.) train_dataloader = DataLoader(dataset, batch_size=int(os.getenv("TRAIN_BATCH_SIZE", "4")), # reduced from 8 (B200 ~180GB) to fit 32B model shuffle=True) print(f'Training dataloader is ready. Example: {next(iter(train_dataloader))}') # Note: we don't need usual train&eval split since we have our own martingale eval. return train_dataloader # Apply Lora def toggle_lora(model): # Avoid double-apply lora if hasattr(model, 'peft_config'): print("Model already has PEFT adapters, skipping LoRA application") return model else: print(f"GPU memory before LoRA: {t.cuda.memory_allocated() / 1024**3:.1f} GB") # Add LoRA adapters using standard PEFT (no Unsloth) peft_config = LoraConfig( r=32, # LoRA rank optimized for B200 (80GB VRAM) lora_alpha=32, # For lora: effective_lr = lr × (α/r), so we maintain the same lr with the original. lora_dropout=0.1, # to prevent overfitting bias="none", task_type="CAUSAL_LM", target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] ) model = get_peft_model(model, peft_config) """ # UNSLOTH VERSION (commented out - has in-place ops incompatible with DIY loss) model = FastLanguageModel.get_peft_model( model, r = 128, target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], lora_alpha = 16, lora_dropout = 0.1, bias = "none", use_gradient_checkpointing = False, random_state = 3407, ) """ print(f"GPU memory after LoRA: {t.cuda.memory_allocated() / 1024**3:.1f} GB") print("LORA turned on with standard PEFT (no Unsloth in-place ops)!") return model # The differentiable OLS function def differentiable_ols(priors, deltas): """Train a Linear Regressor""" # Convert to float32 for OLS calculation (inverse doesn't support float16) original_dtype = priors.dtype priors_f32 = priors.float() deltas_f32 = deltas.float() ones = t.ones_like(priors_f32).unsqueeze(1) priors_reshaped = priors_f32.unsqueeze(1) X = t.cat([ones, priors_reshaped], dim=1) try: xtx_inv = t.inverse(X.T @ X + 1e-6 * t.eye(2, device=X.device, dtype=X.dtype)) xt_y = X.T @ deltas_f32.unsqueeze(1) coefs = xtx_inv @ xt_y intercept, slope = coefs[0], coefs[1] # Convert back to original dtype return intercept.squeeze().to(original_dtype), slope.squeeze().to(original_dtype) except t.linalg.LinAlgError: return t.tensor(0.0, device=priors.device, dtype=original_dtype), t.tensor(0.0, device=priors.device, dtype=original_dtype) # main training loop def train(model, train_dataloader, trained_model_dir, sft_data, lr=1e-6, epoch=0, tokenizer=None): # Enable anomaly detection to find in-place operation issues (disable in production - retains all activations) # t.autograd.set_detect_anomaly(True) model.train() # required for gradient checkpointing (GC skips when model.training=False) print(f"Starting training with lr={lr}") print(f"Total batches: {len(train_dataloader)}") print(f"Model device: {model.device}") # Set up loss logging import os import json from datetime import datetime log_dir = "training/temp" os.makedirs(log_dir, exist_ok=True) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") loss_log_file = f"{log_dir}/training_log_{timestamp}.json" # L1 mode: per-batch OLS in reward + batch-level centering (env var L1_BATCH_OLS=1). # Rationale: per-epoch OLS target lags behind the policy; per-batch target moves with policy. # All other gradient flow unchanged (REINFORCE with detached reward). L1_BATCH_OLS = int(os.getenv("L1_BATCH_OLS", "0")) # [MenoClaw 2026-05-21 | KL-safe reward redesign - env-gated, default OFF] # KL_SAFE_REWARD=1 enables: (a) content-token attribution (credit the full # reasoning trace + answer, excluding chat-template/formatting tokens), and # (b) a sequence-level KL-to-base penalty (k3 estimator) so a strong step # cannot collapse the policy into a non-forecaster. Default path unchanged. KL_SAFE = int(os.getenv("KL_SAFE_REWARD", "0")) KL_BETA = float(os.getenv("KL_BETA", "0.1")) _structural_ids = None if KL_SAFE and tokenizer is not None: _sid = set(getattr(tokenizer, "all_special_ids", []) or []) for _ws in ["\n", "\n\n", "\n\n\n", " ", " "]: try: _sid.update(tokenizer.encode(_ws, add_special_tokens=False)) except Exception: pass _structural_ids = t.tensor(sorted(_sid), dtype=t.long) if _sid else None print(f"[KL-safe] ON: KL-to-base beta={KL_BETA}; content-token attribution " f"excludes {len(_sid)} structural/formatting token ids") # [MenoClaw 2026-05-21 | informativeness term - env-gated, default OFF] # The base reward -(delta*predicted_bias) gives delta~0 a reward of ~0, so an # inert "stop updating" policy is never penalised (KL-safe eval: delta=0 frac # 17%->31%). INFO_TERM=1 adds info_penalty = INFO_COEF * relu(INFO_EPS - |delta|): # a substantive update is neutral, near-zero inertia is penalised. INFO_TERM = int(os.getenv("INFO_TERM", "0")) INFO_COEF = float(os.getenv("INFO_COEF", "1.0")) INFO_EPS = float(os.getenv("INFO_EPS", "0.02")) def _info_penalty(_abs_d): return INFO_COEF * t.clamp(INFO_EPS - _abs_d, min=0.0) if INFO_TERM: print(f"[info-term] ON: info_penalty = INFO_COEF={INFO_COEF} * relu(INFO_EPS={INFO_EPS} - |delta|)") # [MenoClaw 2026-06-02 | A/B/C reward modes - env-gated, default A = existing behavior] REWARD_MODE = os.getenv("REWARD_MODE", "A").upper() # A=martingale, B=vanilla Brier -(y-post)^2, C=lambda*A+(1-lambda)*B (z-scored) LAMBDA_C = float(os.getenv("LAMBDA_C", "0.5")) print(f"[reward-mode] REWARD_MODE={REWARD_MODE} LAMBDA_C={LAMBDA_C}") training_log = { "start_time": timestamp, "lr": lr, "ppo_clip_eps": float(os.getenv("PPO_CLIP_EPS", "0.0")), "l1_batch_ols": L1_BATCH_OLS, "total_batches": len(train_dataloader), "reward_mean_epoch": None, # filled after OLS below "steps": [] } import torch.nn.functional as F # Each epoch , we train Linear Regressor with newly acquired inference data. with t.no_grad(): # No gradients needed for this part all_priors = t.tensor([item["prior"] for item in sft_data], dtype=t.float16) all_actual_deltas = t.tensor([item["actual_delta"] for item in sft_data], dtype=t.float16) intercept, slope = differentiable_ols(all_priors, all_actual_deltas) # Epoch-level reward mean for centering: makes E[reward] = 0 so net gradient # direction is zero and updates distinguish above/below-average samples only. all_predicted_bias = (intercept + slope * all_priors).float() all_ys = t.tensor([float(item.get("y", item.get("y_resolve", 0.0))) for item in sft_data], dtype=t.float16) _A = -(all_actual_deltas.float() * all_predicted_bias) if INFO_TERM: _A = _A - _info_penalty(all_actual_deltas.float().abs()) _post = (all_priors.float() + all_actual_deltas.float()).clamp(0.0, 1.0) # [MenoClaw 2026-06-03 | judge-in-loop] Brier from the model's OWN stated P, not the # judge-derived (prior+delta): an LLM judge moderates beliefs and launders calibration, # so the Brier/B/C term must use stated_p. Martingale term _A stays on the judge prior/delta. # Env-gated BRIER_FROM_STATED; falls back to prior+delta when the seed has no stated_p. if int(os.getenv("BRIER_FROM_STATED", "0")) and all(("stated_p" in _it) for _it in sft_data): _post_brier = t.tensor([float(_it["stated_p"]) for _it in sft_data], dtype=t.float16).float().clamp(0.0, 1.0) print(f"[brier-src] stated-P meanP={_post_brier.mean():.3f} (judge-in-loop: martingale=judge, Brier=stated)") else: _post_brier = _post _B = -((all_ys.float() - _post_brier) ** 2) _muA, _sdA = _A.mean(), _A.std().clamp(min=1e-6) _muB, _sdB = _B.mean(), _B.std().clamp(min=1e-6) if REWARD_MODE == "B": all_rewards = _B elif REWARD_MODE == "C": all_rewards = LAMBDA_C*((_A-_muA)/_sdA) + (1.0-LAMBDA_C)*((_B-_muB)/_sdB) else: all_rewards = _A print(f"[reward-mode {REWARD_MODE}] epoch: A_mean={_A.mean():.4f} B_mean={_B.mean():.4f} post_mean={_post.mean():.3f} y_mean={all_ys.float().mean():.3f}") reward_mean_epoch = all_rewards.mean() print(f"Epoch reward mean (for centering): {reward_mean_epoch:.6f}") print(f"L1_BATCH_OLS mode: {L1_BATCH_OLS} ({'per-batch OLS + batch-level centering' if L1_BATCH_OLS else 'per-epoch OLS + epoch-level centering (default)'})") training_log["reward_mean_epoch"] = float(reward_mean_epoch) trainable_params = [p for p in model.parameters() if p.requires_grad] optimizer = t.optim.AdamW(trainable_params, lr=lr) def _mem(): return t.cuda.memory_allocated() / 1024**3 # PPO: pre-compute reference log_probs under the initial (pre-update) policy ppo_eps = float(os.getenv("PPO_CLIP_EPS", "0.0")) ref_log_probs = {} if ppo_eps > 0 or KL_SAFE: print(f"Reference log_prob pre-pass (ppo_eps={ppo_eps}, KL_SAFE={KL_SAFE})...") model.eval() with t.no_grad(): for ref_i, ref_batch in enumerate(train_dataloader): ref_input_ids = ref_batch['input_ids'].to(model.device) ref_attn = ref_batch['attention_mask'].to(model.device) ref_qlens = ref_batch['question_length'] ref_logits = model(input_ids=ref_input_ids, attention_mask=ref_attn).logits ref_shift_labels = ref_input_ids[:, 1:].clone() ref_shift_logits = ref_logits[:, :-1, :] ref_tgt = ref_shift_logits.gather(2, ref_shift_labels.unsqueeze(-1)).squeeze(-1).float() ref_log_Z = ref_shift_logits[:, :, :8192].float().logsumexp(dim=-1) for _v in range(8192, ref_shift_logits.shape[-1], 8192): ref_log_Z = t.logaddexp(ref_log_Z, ref_shift_logits[:, :, _v:_v+8192].float().logsumexp(dim=-1)) ref_tlp = ref_tgt - ref_log_Z ref_amask = t.zeros_like(ref_tlp) for b, qlen in enumerate(ref_qlens): if qlen > 0: ref_amask[b, qlen - 1:] = 1.0 ref_amask = ref_amask * ref_attn[:, 1:] if int(os.getenv("ANSWER_TOKEN_CREDIT", "0")): _Kr = int(os.getenv("ANSWER_CREDIT_TOKENS", "12")) _r2 = t.zeros_like(ref_amask) for _rb in range(ref_amask.shape[0]): _rnz = ref_amask[_rb].nonzero(as_tuple=True)[0] if _rnz.numel() > 0: _rl = int(_rnz[-1].item()); _rf = int(_rnz[0].item()) _r2[_rb, max(_rf, _rl - _Kr + 1):_rl + 1] = 1.0 ref_amask = _r2 if KL_SAFE and _structural_ids is not None: _rstruct = t.isin(ref_shift_labels, _structural_ids.to(ref_shift_labels.device)) ref_amask = ref_amask * (~_rstruct).float() ref_valid = ref_amask.sum(dim=-1).clamp(min=1.0) ref_log_probs[ref_i] = ((ref_tlp * ref_amask).sum(dim=-1) / ref_valid).cpu() if ref_i % 100 == 0: print(f" PPO pre-pass: {ref_i}/{len(train_dataloader)} batches") model.train() print("PPO reference log_probs ready.") total_loss = 0.0 for i, batch in enumerate(train_dataloader): if i == 0: print(f"[MEM] before batch-0 input: {_mem():.2f} GiB") # Extract all data from batch input_ids = batch['input_ids'].to(model.device) # [batch_size, seq_len] attention_mask = batch['attention_mask'].to(model.device) # [batch_size, seq_len] actual_deltas = batch['actual_delta'].to(model.device) # [batch_size] priors = batch['prior'].to(model.device) # [batch_size] question_lengths = batch['question_length'] # list of ints (CPU) if i == 0: print(f"[MEM] after inputs to GPU: {_mem():.2f} GiB") print(f"[{i}] 1a. Input shapes: input_ids={input_ids.shape}, attention_mask={attention_mask.shape}") # 1. FORWARD PASS lm_outputs = model(input_ids=input_ids, attention_mask=attention_mask) logits = lm_outputs.logits # [batch_size, seq_len, vocab_size] if i == 0: print(f"[MEM] after forward: {_mem():.2f} GiB") print(f"[{i}] 1d. Logits shape: {logits.shape}, dtype: {logits.dtype}") print(f"[{i}] 1e. Has NaN in logits: {t.isnan(logits).any()}") # 2. LOG PROB OVER ANSWER TOKENS ONLY (REINFORCE) print(f"[{i}] 2. Computing log_prob over answer tokens...") # Shift for next-token prediction alignment. # log_prob = target_logit - log Σ_v exp(logit_v) (= log_softmax at target position) # Float32 required to avoid bf16 saturation on confident CoT tokens. # Materializing [B, L, V] float32 at B=16 costs ~40 GB and OOMs, so: # 1. gather target logit first (tiny: [B, L]) # 2. accumulate logsumexp over vocab in 8192-token chunks (~2 GB each) shift_labels = input_ids[:, 1:].clone() # [batch_size, seq_len-1] shift_logits = logits[:, :-1, :] # [batch_size, seq_len-1, vocab] bf16 target_logits = shift_logits.gather(2, shift_labels.unsqueeze(-1)).squeeze(-1).float() # [batch_size, seq_len-1] float32 _VOCAB_CHUNK = 8192 log_Z = shift_logits[:, :, :_VOCAB_CHUNK].float().logsumexp(dim=-1) for _v in range(_VOCAB_CHUNK, shift_logits.shape[-1], _VOCAB_CHUNK): log_Z = t.logaddexp(log_Z, shift_logits[:, :, _v:_v + _VOCAB_CHUNK].float().logsumexp(dim=-1)) token_log_probs = target_logits - log_Z # [batch_size, seq_len-1] float32 del logits, shift_logits, target_logits, log_Z # Mask to answer tokens only (positions after question boundary) answer_mask = t.zeros_like(token_log_probs) for b, qlen in enumerate(question_lengths): if qlen > 0: answer_mask[b, qlen - 1:] = 1.0 # -1: shift alignment # Mask the padding token shift_attention_mask = attention_mask[:,1:] # shifted attention_mask [b, seq-1] answer_mask = answer_mask * shift_attention_mask # [MenoClaw 2026-05-18 | env-gated, default OFF | ANSWER_TOKEN_CREDIT] # The mask above spans the whole CoT+answer, so the REINFORCE log_prob is a # sequence-mean over hundreds-thousands of tokens -> credit assignment is # smeared off the decision (the trailing probability). When enabled, restrict # the mask to the last ANSWER_CREDIT_TOKENS real tokens (the answer float), # clamped to the question boundary. Single isolated variable; default path # is byte-identical to before. if int(os.getenv("ANSWER_TOKEN_CREDIT", "0")): _K = int(os.getenv("ANSWER_CREDIT_TOKENS", "12")) _atc = t.zeros_like(answer_mask) for _b in range(answer_mask.shape[0]): _nz = answer_mask[_b].nonzero(as_tuple=True)[0] if _nz.numel() > 0: _last = int(_nz[-1].item()) _first = int(_nz[0].item()) _start = max(_first, _last - _K + 1) _atc[_b, _start:_last + 1] = 1.0 answer_mask = _atc if KL_SAFE and _structural_ids is not None: _struct = t.isin(shift_labels, _structural_ids.to(shift_labels.device)) answer_mask = answer_mask * (~_struct).float() valid_tokens_per_seq = answer_mask.sum(dim=-1).clamp(min=1.0) log_prob = (token_log_probs * answer_mask).sum(dim=-1) / valid_tokens_per_seq # [batch_size] if i <= 5: ql = [int(q) for q in question_lengths] print(f"[{i}] 2a. question_lengths: min={min(ql)}, max={max(ql)}, mean={sum(ql)/len(ql):.1f} (sanity: should be >> 0)") print(f"[{i}] 2b. log_prob: mean={log_prob.mean():.4f}, has_nan={t.isnan(log_prob).any()}") print(f"[{i}] 2c. answer_mask tokens: {answer_mask[0].sum().item():.0f} / {answer_mask.shape[1]}") # 3. REINFORCE LOSS print(f"[{i}] 3. REINFORCE loss calculation...") if L1_BATCH_OLS: # L1: refit OLS on this batch's (prior, delta) pairs (no epoch lag). with t.no_grad(): batch_intercept, batch_slope = differentiable_ols(priors.float(), actual_deltas.float()) predicted_bias_q_p = (batch_intercept + batch_slope * priors).float() else: predicted_bias_q_p = (intercept + slope * priors).float() # reward by REWARD_MODE (A=martingale, B=vanilla Brier, C=combined z-scored) _A_b = -(actual_deltas.float() * predicted_bias_q_p.detach()) if INFO_TERM: _info_pen = _info_penalty(actual_deltas.float().abs()) _A_b = _A_b - _info_pen if i <= 15: _zf = (actual_deltas.float().abs() < INFO_EPS).float().mean() print(f'[{i}] info-term: penalty mean={_info_pen.mean():.6f} frac|d| 0: ref_lp = ref_log_probs[i].to(model.device) ratio = (log_prob - ref_lp).exp() clipped_ratio = ratio.clamp(1.0 - ppo_eps, 1.0 + ppo_eps) loss = -t.min(ratio * reward, clipped_ratio * reward).mean() else: ratio = None loss = -(reward * log_prob).mean() if KL_SAFE: _ref_lp = ref_log_probs[i].to(log_prob.device) _lr = log_prob - _ref_lp _kl = (t.exp(-_lr) - 1.0 + _lr).mean() loss = loss + KL_BETA * _kl if i <= 15: print(f'[{i}] KL(policy||base)={_kl.item():.6f} beta={KL_BETA}') total_loss += loss.item() if i <= 15: print(f"[{i}] 3a. Values summary:") print(f"[{i}] reward: mean={reward.mean():.6f}, std={reward.std() if reward.numel() > 1 else 0:.6f}") print(f"[{i}] predicted_bias_q_p: mean={predicted_bias_q_p.mean():.6f}") print(f"[{i}] log_prob: mean={log_prob.mean():.6f}") print(f"[{i}] loss: {loss.item():.6f}") # 4. BACKWARD PASS print(f"[{i}] 4. Backward pass...") optimizer.zero_grad() loss.backward() total_norm = sum(p.grad.data.norm(2).item() ** 2 for p in model.parameters() if p.grad is not None) ** 0.5 print(f"[{i}] 4a. Gradients: Loss={loss.item():.6f}, Grad_norm={total_norm:.4f}, OLS: intercept={intercept:.4f}, slope={slope:.4f}") # Log step data — core metrics every step, detailed diagnostics for first 15 step_data = { "step": i, "loss": loss.item(), "grad_norm": total_norm, "ols_intercept": float(batch_intercept) if L1_BATCH_OLS else float(intercept), "ols_slope": float(batch_slope) if L1_BATCH_OLS else float(slope), "log_prob_mean": float(log_prob.detach().mean()), "reward_mean": float(reward.mean()), } if ratio is not None: step_data["ratio_mean"] = float(ratio.detach().mean()) step_data["ratio_clip_frac"] = float( ((ratio < 1.0 - ppo_eps) | (ratio > 1.0 + ppo_eps)).float().mean() ) if i <= 15: # Detailed diagnostics for first 15 steps only step_data.update({ "reward_std": float(reward.std()) if reward.numel() > 1 else 0.0, "predicted_bias_mean": float(predicted_bias_q_p.mean()), "predicted_bias_std": float(predicted_bias_q_p.std()) if predicted_bias_q_p.numel() > 1 else 0.0, }) training_log["steps"].append(step_data) if wandb.run is not None: wlog = { "train/loss": loss.item(), "train/grad_norm": total_norm, "train/log_prob_mean": float(log_prob.mean()), "train/reward_mean": float(reward.mean()), "train/ols_intercept": float(intercept), "train/ols_slope": float(slope), } if ratio is not None: wlog["train/ratio_mean"] = float(ratio.detach().mean()) wlog["train/ratio_clip_frac"] = step_data["ratio_clip_frac"] wandb.log(wlog, step=epoch * len(train_dataloader) + i) # 5. GRADIENT CLIPPING & WEIGHT UPDATE print(f"[{i}] 5. Gradient clipping and weight update...") trainable_params = [p for p in model.parameters() if p.requires_grad] t.nn.utils.clip_grad_norm_(trainable_params, max_norm=1.0) optimizer.step() if i % 20 == 0 and t.cuda.is_available(): print(f"[{i}] GPU memory: {t.cuda.memory_allocated() / 1024**3:.2f} GB") print(f"=== TRAINING STEP {i} END ===\n") avg_loss = total_loss / len(train_dataloader) print(f"Training completed. Average loss: {avg_loss:.6f}") # Save final training log training_log["end_time"] = datetime.now().strftime("%Y%m%d_%H%M%S") training_log["avg_loss"] = avg_loss training_log["final_step"] = len(train_dataloader) - 1 dump_file_for_run(loss_log_file, training_log, run_id=os.environ.get("RUN_ID")) print(f"Saving model to {trained_model_dir}") model.save_pretrained(trained_model_dir) print("Model saved successfully") return trained_model_dir, loss_log_file def evaluate(run_id): # Martingale Setup with SFT-ed policy """ Args: run_id: we pass on a new run_id each time """ # Ensure inference env vars are set (safe to call multiple times via setdefault) os.environ.setdefault("SELF_REPORT", "1") os.environ.setdefault("DOUBLE_SAMPLE", "1") os.environ.setdefault("DECOUPLE_TRAJECTORY_BELIEFS", "1") os.environ.setdefault("RECOMPUTE_TRAJECTORIES", "never") os.environ.setdefault("RECOMPUTE_BELIEFS", "missing") os.environ.setdefault("USE_BATCHED_INFERENCE", "1") os.environ.setdefault("INFERENCE_BATCH_SIZE", "32") os.environ.setdefault("CHECKPOINT_SIZE", "100") # To generate training files into data/runs/batch-martingale-training # Generate run_id if not provided via environment #run_id = os.environ.get("RUN_ID", "BATCH-ChainOfThought-Forecasting-llama-3") num_trajectories = int(os.getenv("NUM_TRAJ", "2282")) expected_priors = [None] if run_id is None: import datetime timestamp = datetime.datetime.now().strftime("%Y-%m-%d-%H%M%S") suffix = os.environ.get("RUN_ID_SUFFIX", "") run_id = timestamp + suffix print(f'===Prep Martingale Evaluation for Reasoning Traj {run_id}') if int(os.environ.get("DEBUG", "0")): # Debug mode: break at the first error import pdb import sys import traceback print(f'Entering debug mode of Martingale Eval...') try: run_reasoning( setup.algo, setup.policy, setup.domain, setup.reasoning_mode, run_id, num_trajectories, expected_priors, False ) except Exception as e: print(e) extype, value, tb = sys.exc_info() traceback.print_exc() pdb.post_mortem(tb) else: # Normal mode: run the script directly print(f'Entering Normal mode of Martingale Eval, no debug performed') run_reasoning( setup.algo, setup.policy, setup.domain, setup.reasoning_mode, run_id, num_trajectories, expected_priors, False ) def comparison(old_run_id: str, new_run_id: str): """ Comparing two policies. Note that this approach takes coef in training but here we evaluate r2 (still report coef tho) Args: old_run_id (str): path to martingale eval results on old policy new_run_id (str): path to martingale eval results on new policy. We could set up env RUNS_SUBDIR=batch-martingale-training """ file_path = 'bias-eval-results-MartingaleStrategy.json' old_data = load_file_for_run(file_path, old_run_id) new_data = load_file_for_run(file_path, new_run_id) old_martingale_score_coef = old_data["loss_details"]["linear_regression"]["coef"][0] new_martingale_score_coef = new_data["loss_details"]["linear_regression"]["coef"][0] old_martingale_score_r2 = old_data["loss_details"]["linear_regression"]["r2_score"] new_martingale_score_r2 = new_data["loss_details"]["linear_regression"]["r2_score"] if old_martingale_score_r2 >= new_martingale_score_r2: print(f"Success! Martingale Score of newly trained policy is {new_martingale_score_r2}, smaller than the previous one: {old_martingale_score_r2}") else: print(f"Martingale Score increased from {old_martingale_score_r2} to {new_martingale_score_r2}. Rejected! ") # NEP At some point you may need async pipeline too (but why?) def main(): check_cuda_status("Script start") max_epoch = int(os.getenv("MAX_EPOCH", 5)) # NEP For now we go through each round of training set only once. epoch, on the other hand, means how many times you go through the entire training set. We steal the term, but not the concept. trained_model_dir = "./training/trained_models" base_model_dir = os.environ.get("MODEL_BASE_DIR", "/tmp/models_tmp/base_model_qwen") # RUN_ID is the inference run whose data seeds epoch 0 training (fresh start only). # On resume, run_id is always overridden from training_state.json (post_run_id of last # complete epoch), so this env var has no effect after the first epoch. run_id = os.environ.get("RUN_ID", "BaseModel-Complete-DirectInference-Forecasting-Qwen3-32B-2026-04-01-194158-deduped") import json as _json state_path = os.path.join(trained_model_dir, "training_state.json") def _load_state(): if os.path.exists(state_path): return _json.load(open(state_path)) return {"epochs": []} def _save_state(state): with open(state_path, "w") as f: _json.dump(state, f, indent=2) _plot_ts = datetime.datetime.now().strftime("%Y-%m-%d") _plot_batch = os.getenv("TRAIN_BATCH_SIZE", "1") _plot_lr = os.getenv("TRAIN_LR", "1e-6") _plot_scale = os.getenv("REWARD_SCALE", "1.0") _plot_tags = os.getenv("PLOT_TAGS", "") # e.g. "centered,float32" — appended verbatim _plot_label = f"{_plot_ts}_batch{_plot_batch}_lr{_plot_lr}_scale{_plot_scale}" if _plot_tags: _plot_label += f"_{_plot_tags}" plots_dir = os.path.join("visuals", "training", _plot_label) os.makedirs(plots_dir, exist_ok=True) print(f"Plots will be saved to: {plots_dir}") print(f"GPU memory total: {t.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB") print(f"GPU memory allocated: {t.cuda.memory_allocated() / 1024**3:.1f} GB") print(f"GPU memory cached: {t.cuda.memory_reserved() / 1024**3:.1f} GB") print(f"GPU memory free: {(t.cuda.get_device_properties(0).total_memory - t.cuda.memory_reserved()) / 1024**3:.1f} GB") # --- One-time base model download guard --- if not os.path.exists(base_model_dir) or not os.listdir(base_model_dir): print("Base model not found, downloading via download_basemodel.py...") import subprocess subprocess.run([sys.executable, "scripts/train_scripts/download_basemodel.py"], check=True) print(f"Base model saved at {base_model_dir}. Run the script again to start training.") return # --- Load model once: fresh start or crash recovery --- os.makedirs(trained_model_dir, exist_ok=True) epoch_checkpoints = sorted( [d for d in os.listdir(trained_model_dir) if d.startswith("epoch_")], key=lambda d: int(d.split("_")[1]) ) if epoch_checkpoints: latest = epoch_checkpoints[-1] start_epoch = int(latest.split("_")[1]) + 1 latest_checkpoint_dir = os.path.join(trained_model_dir, latest) print(f"Crash recovery: resuming from epoch {start_epoch} (checkpoint: {latest_checkpoint_dir})") state = _load_state() _complete = [e for e in state["epochs"] if e["status"] == "complete"] _pending = [e for e in state["epochs"] if e["status"] == "train_done"] if _complete: # Always use the most recent completed epoch's inference run as training data source. # This ensures each training epoch trains on data from the previous epoch's inference, # regardless of what RUN_ID env var is set to. run_id = _complete[-1]["post_run_id"] print(f"Crash recovery: restored run_id to post-epoch-{_complete[-1]['epoch']} eval: {run_id}") elif _pending: run_id = _pending[-1]["pre_run_id"] print(f"Crash recovery: pending eval for epoch {_pending[-1]['epoch']}, pre_run_id: {run_id}") tokenizer = AutoTokenizer.from_pretrained(base_model_dir) if not _pending: # No pending eval — load training model now for next epoch model = AutoModelForCausalLM.from_pretrained( base_model_dir, device_map="auto", torch_dtype=t.bfloat16, attn_implementation="sdpa" ) from peft import PeftModel model = PeftModel.from_pretrained(model, latest_checkpoint_dir, is_trainable=True) # If _pending: defer model load until after eval — SGLang needs the full GPU else: start_epoch = 0 state = _load_state() _pending = [] print("Fresh start: loading base model...") model = AutoModelForCausalLM.from_pretrained( base_model_dir, device_map="auto", torch_dtype=t.bfloat16, attn_implementation="sdpa" ) _first_param_dtype = next(model.parameters()).dtype print(f"[DIAG] First param dtype after load: {_first_param_dtype}") check_cuda_status("After loading base model") tokenizer = AutoTokenizer.from_pretrained(base_model_dir) if int(os.environ.get("LORA", "1")): # NEP I didn't see there was a place where we could change this parameter model = toggle_lora(model) tokenizer.save_pretrained(os.path.join(trained_model_dir, "tokenizer")) if not int(os.environ.get("EXCLUDE_BASE_EVAL", "1")): print("Running Martingale Eval on base model...") check_cuda_status("Before evaluate()") evaluate(run_id) check_cuda_status("After evaluate()") else: print("Skipping Martingale Eval on base model since it's done already.") if not _pending: # Model is loaded — configure for training print(f"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}") if hasattr(model, "enable_input_require_grads"): model.enable_input_require_grads() model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) _base = model.base_model.model if hasattr(model, 'base_model') else model _sample_layer = list(_base.modules())[1] print(f"[DIAG] GC active on base model: {getattr(_base, 'is_gradient_checkpointing', 'N/A')}") print(f"[DIAG] Sample layer gradient_checkpointing: {getattr(_sample_layer, 'gradient_checkpointing', 'N/A')}") # Complete any epoch that trained but didn't get its eval (crash between train and eval) if _pending: _p = _pending[-1] _ts = datetime.datetime.now().strftime("%Y-%m-%d-%H%M%S") run_id = f'Finetuned-{_ts}-Slow-Fast-Thinking-Forecasting-Qwen3-32B' os.environ["RUN_ID"] = run_id complete_create_dir_path(run_id) print(f"Completing missed eval for epoch {_p['epoch']}...") _pending_ckpt = os.path.join(trained_model_dir, f"epoch_{_p['epoch']}") os.environ["LORA_PATH"] = _pending_ckpt setup.policy.lora_path = _pending_ckpt evaluate(run_id) import asyncio; asyncio.run(setup.policy.stop_backend()) model = AutoModelForCausalLM.from_pretrained(base_model_dir, device_map="auto", torch_dtype=t.bfloat16, attn_implementation="sdpa") from peft import PeftModel as _PeftModel model = _PeftModel.from_pretrained(model, _pending_ckpt, is_trainable=True) model.enable_input_require_grads() model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) _p["post_run_id"] = run_id _p["status"] = "complete" _save_state(state) print(f"Recovered epoch {_p['epoch']} eval complete. run_id={run_id}") wandb.init( project=os.getenv("WANDB_PROJECT", "sft-product-based"), name=os.getenv("WANDB_RUN_NAME", None), resume="allow", config={ "lr": float(os.getenv("TRAIN_LR", "1e-6")), "batch_size": int(os.getenv("TRAIN_BATCH_SIZE", "4")), "reward_scale": float(os.getenv("REWARD_SCALE", "1.0")), "ppo_clip_eps": float(os.getenv("PPO_CLIP_EPS", "0.0")), "max_epoch": max_epoch, "lora_r": 32, "lora_alpha": 32, "context_len": int(os.getenv("TRAIN_CONTEXT_LEN", "4092")), "base_model": base_model_dir, "seed_run_id": run_id, }, ) for i in range(start_epoch, max_epoch): # Step 1: Prep SFT data of current epoch # Generate SFT data (with priors and deltas in same file) if not os.path.exists(complete_path("sft-data.json", run_id)): # NEP 需要check 是否已经有sft & 以及注意每一轮都需要新的SFT在同一个folder # Check which format of data files exist # NEW FORMAT: Single file with self-reported beliefs (reasoning-record.json) # OLD FORMAT: Two files (reasoning-trajectories-raw.json + reasoning-beliefs-MartingaleStrategy-gpt-4o.json) if os.path.exists(complete_path("reasoning-record.json", run_id)): # NEW FORMAT: Self-reported beliefs print("Using NEW FORMAT: reasoning-record.json with self-reported beliefs") run_convert_sft_data(run_id, "reasoning-record.json", None, "sft-data.json") elif os.path.exists(complete_path("reasoning-trajectories-raw.json", run_id)): # OLD FORMAT: External judge print("Using OLD FORMAT: reasoning-trajectories-raw.json + reasoning-beliefs") run_convert_sft_data(run_id, "reasoning-trajectories-raw.json", "reasoning-beliefs-MartingaleStrategy-gpt-4o.json", "sft-data.json") else: raise FileNotFoundError(f"No reasoning data files found for run_id: {run_id}. Expected either 'reasoning-record.json' or 'reasoning-trajectories-raw.json'") sft_data = load_file_for_run("sft-data.json", run_id) # Step 2: Train a LinearRegressor, and then train SFT model against it. train_dataloader = data_preprocessing(sft_data, tokenizer) print(f'After data preprocessing, first item in train_dataloader is {next(iter(train_dataloader))}') epoch_checkpoint_dir = os.path.join(trained_model_dir, f"epoch_{i}") _pre_run_id = run_id _, loss_log_file = train( model, train_dataloader, epoch_checkpoint_dir, sft_data, lr=float(os.getenv("TRAIN_LR", "1e-6")), epoch=i, tokenizer=tokenizer, ) state["epochs"].append({"epoch": i, "pre_run_id": _pre_run_id, "post_run_id": None, "status": "train_done"}) _save_state(state) # Monitor plot: step-level training metrics (loss, reward, log_prob, grad_norm) try: from scripts.visualization.training.plot_reinforce_training import plot as plot_reinforce plot_reinforce( log_path=complete_path(loss_log_file, _pre_run_id), out_path=os.path.join(plots_dir, f"epoch_{i}_monitor_reinforce_training"), ) except Exception as e: print(f"Monitor plot failed (non-fatal): {e}") # Step 3: Martingale Eval on fine-tuned model (skip if TRAIN_ONLY=1) if os.getenv("TRAIN_ONLY", "0") == "1": print("TRAIN_ONLY=1: skipping eval. Epoch marked train_done.") continue old_run_id = run_id # Updating run_id timestamp = datetime.datetime.now().strftime("%Y-%m-%d-%H%M%S") run_id = f'Finetuned-{timestamp}-Slow-Fast-Thinking-Forecasting-Qwen3-32B' os.environ["RUN_ID"] = run_id new_dir = complete_create_dir_path(run_id) print(f'New data dir created for saving martingale eval data for fine-tuned model at {new_dir}') # New martingale eval on fine-tuned model # Free training model from GPU so SGLang has full memory for model weights + KV cache print("Freeing training model from GPU for SGLang eval...") import gc; model.cpu(); del model; gc.collect(); gc.collect(); t.cuda.synchronize(); t.cuda.empty_cache() check_cuda_status("After freeing training model (before eval)") # SGLang loads base model + LoRA adapter directly (no merge needed) os.environ["LORA_PATH"] = epoch_checkpoint_dir setup.policy.lora_path = epoch_checkpoint_dir print("=== Starting Martingale Evaluation on Newly Trained Model === ") evaluate(run_id) import asyncio; asyncio.run(setup.policy.stop_backend()) t.cuda.synchronize(); t.cuda.empty_cache() check_cuda_status("After stop_backend (before model reload)") state["epochs"][-1]["post_run_id"] = run_id state["epochs"][-1]["status"] = "complete" _save_state(state) # Reload training model (clean base + same LoRA weights) for next epoch print("Reloading model for next epoch...") model = AutoModelForCausalLM.from_pretrained( base_model_dir, device_map="auto", torch_dtype=t.bfloat16, attn_implementation="sdpa", ) from peft import PeftModel as _PeftModel model = _PeftModel.from_pretrained(model, epoch_checkpoint_dir, is_trainable=True) model.enable_input_require_grads() model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) # Generate reasoning-record-intersect.json for both runs (required by comparative plots and stats) try: from scripts.visualization.training.intersect_runs import intersect intersect(f"run-{old_run_id}", f"run-{run_id}") except Exception as e: print(f"Intersect step failed (non-fatal): {e}") try: comparison(old_run_id, run_id) except Exception as e: print(f"comparison() failed (non-fatal): {e}") # Comparative stats (martingale score, Brier, t-test) across training print("=== Running Comparative Stats ===") os.environ["RUN_ID"] = old_run_id os.environ["NEW_RUN_ID"] = run_id try: from training.comparative_stats import main as comparative_stats_main comparative_stats_main() except Exception as e: print(f"comparative_stats failed (non-fatal): {e}") # Comparative plots: prior-delta scatter and Brier distribution _runs_base = os.path.join("data/runs/batch-martingale-training") _pre_dir = os.path.join(_runs_base, f"run-{old_run_id}") _post_dir = os.path.join(_runs_base, f"run-{run_id}") try: from scripts.visualization.training.scatter_prior_delta_with_training import ( load_and_match_data_self_report, create_scatter_plot ) _pre_data, _post_data, _pids = load_and_match_data_self_report(_pre_dir, _post_dir) for focus in ('pre', 'post'): _out = os.path.join(plots_dir, f"epoch_{i}_comparative_prior_delta_scatter_focus_{focus}.png") if not os.path.exists(_out): create_scatter_plot(_pre_data, _post_data, _pids, _out, focus=focus) except Exception as e: print(f"Scatter plot failed (non-fatal): {e}") try: from scripts.visualization.training.brier_delta_distribution import ( compute_deltas, create_brier_delta_chart ) for belief_idx, label in ((0, 'prior'), (1, 'posterior')): _out = os.path.join(plots_dir, f"epoch_{i}_comparative_brier_distribution_{label}.png") if not os.path.exists(_out): _f_before, _f_after, _delta = compute_deltas(_pre_dir, _post_dir, belief_index=belief_idx) create_brier_delta_chart(_delta, n_bins=20, output_path=_out, belief_label=label) except Exception as e: print(f"Brier plot failed (non-fatal): {e}") t.cuda.empty_cache() if __name__ == "__main__": main() wandb.finish()