gpt-5.6-sol

Script completed
Wall time 33.0 seconds
Output:

Processed prompts:  50%|▌| 48/96 [18:11<17:16, 21.60s/it, est. speed input: 6.2
Processed prompts:  86%|▊| 72/84 [18:22<01:43,  8.59s/it, est. speed input: 7.9
Processed prompts: 100%|█| 84/84 [18:39<00:00,  6.24s/it, est. speed input: 8.8
Processed prompts: 100%|█| 84/84 [18:39<00:00,  6.24s/it, est. speed input: 8.8
Processed prompts: 100%|█| 84/84 [18:39<00:00, 13.33s/it, est. speed input: 8.8

Processed prompts:  62%|▋| 60/96 [18:41<08:49, 14.70s/it, est. speed input: 7.2

Script running with cell ID 244
Wall time 11.0 seconds
Output:

Script completed
Wall time 33.6 seconds
Output:

Script running with cell ID 245
Wall time 11.0 seconds
Output:

Script completed
Wall time 27.9 seconds
Output:

Processed prompts:  75%|▊| 72/96 [20:16<04:57, 12.40s/it, est. speed input: 8.6
Processed prompts:  88%|▉| 84/96 [20:44<01:49,  9.10s/it, est. speed input: 9.3

Script running with cell ID 246
Wall time 11.0 seconds
Output:

Script completed
Wall time 5.8 seconds
Output:

Processed prompts: 100%|█| 96/96 [21:55<00:00,  8.08s/it, est. speed input: 9.9
Processed prompts: 100%|█| 96/96 [21:55<00:00,  8.08s/it, est. speed input: 9.9
Processed prompts: 100%|█| 96/96 [21:55<00:00, 13.70s/it, est. speed input: 9.9

FINAL avg@12 accuracy: 55.00% (198/360 solutions correct on 30 problems)
Results written to /workspace/eval_results/frontier_v3_aops_full.json

Script completed
Wall time 0.0 seconds
Output:
{}

Script completed
Wall time 0.2 seconds
Output:
diff --git a/data_collator.py b/data_collator.py
index 0829108..ffb6d9d 100644
--- a/data_collator.py
+++ b/data_collator.py
@@ -56,12 +56,14 @@ class SelfDistillationDataCollator:
         student_prompts = []
         teacher_prompts = []
         teacher_reasoning_prompts = []  # NEW: for reason_first mode
+        answer_texts = []
 
         for feature in features:
             # Extract problem and solution from dataset
             # Handle different possible column names
             problem = feature["problem"]
             solution = feature["solution"]
+            answer_texts.append(str(feature.get("Answer", "")))
 
             # Student prompt: just the problem with instruction (matching evaluation format)
             student_user_message = f"Problem: {problem}\n\nPlease reason step by step, and put your final answer within \\boxed{{}}."
@@ -136,6 +138,8 @@ class SelfDistillationDataCollator:
             "student_prompt_length": max_student_prompt_len,  # Single value for batch!
             # Keep individual lengths for proper masking
             "student_prompt_lengths_per_example": torch.tensor(student_prompt_lengths),
+            # Kept as Python strings until training_step grades the sampled rollout.
+            "answer_texts": answer_texts,
         }
 
         if self.reason_first:
diff --git a/opsd_train.py b/opsd_train.py
index 8526a68..374d8f5 100644
--- a/opsd_train.py
+++ b/opsd_train.py
@@ -1,4 +1,5 @@
 import os
+import re
 import wandb
 
 from datasets import load_dataset
@@ -266,6 +267,26 @@ if __name__ == "__main__":
     dataset = load_dataset("siyanzhao/Openthoughts_math_30k_opsd")
     train_dataset = dataset["train"]
 
+    # Frontier curriculum for AIME-style reasoning.  OPSD only has 3,200 examples of
+    # budget, so spending them on elementary MATH/AMC items gives the privileged
+    # teacher little useful disagreement.  Select hard olympiad/AoPS questions with
+    # direct integer answers (AIME's answer regime).  Dataset.select is an in-memory
+    # view and does not create or alter dataset cache files.
+    def is_aime_style_answer(answer):
+        answer = re.sub(r"\\(?:text|mathrm)\{([^{}]*)\}", r"\1", str(answer))
+        answer = answer.replace("$", "").strip()
+        return bool(re.fullmatch(r"[+]?\d{1,3}", answer)) and 0 <= int(answer) <= 999
+
+    sources = train_dataset["source"]
+    answers = train_dataset["Answer"]
+    frontier_indices = [
+        i
+        for i, (source, answer) in enumerate(zip(sources, answers))
+        if source in {"olympiads", "aops_forum"} and is_aime_style_answer(answer)
+    ]
+    train_dataset = train_dataset.select(frontier_indices)
+    print(f"Verified frontier curriculum: selected {len(train_dataset)} / {len(dataset['train'])} examples")
+
     trainer = OPSDTrainer(
         model=model_args.model_name_or_path,
         args=training_args,
diff --git a/opsd_trainer.py b/opsd_trainer.py
index c151bd7..706bc0f 100644
--- a/opsd_trainer.py
+++ b/opsd_trainer.py
@@ -14,6 +14,7 @@
 
 import os
 import random
+import re
 import textwrap
 import warnings
 from collections import defaultdict, deque
@@ -370,6 +371,7 @@ class OPSDTrainer(SFTTrainer):
         required_columns = [
             "problem",
             "solution",
+            "Answer",
         ]
         if self._signature_columns is None:
             self._signature_columns = required_columns
@@ -389,6 +391,10 @@ class OPSDTrainer(SFTTrainer):
         logits_are_probs=False,
         top_k=None,
         token_clip=None,
+        sampled_token_ids=None,
+        outcome_correct=None,
+        outcome_weight=0.0,
+        position_scale=256.0,
     ):
         """
         Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation using F.kl_div. See Eq. (1)
@@ -463,20 +469,76 @@ class OPSDTrainer(SFTTrainer):
         if token_clip is not None:
             jsd = jsd.clamp(max=token_clip)
 
-        # Masking
+        # Sum vocabulary contributions into a token loss.  We retain the paper's
+        # pointwise vocabulary clipping, but discount late continuation tokens:
+        # after a long shared prefix the teacher nearly copies the student and the
+        # remaining tokens are much less likely to be a meaningful decision point.
+        token_jsd = jsd.sum(dim=-1)
+
         if labels is not None:
             mask = labels != -100
-            jsd = jsd[mask]
+        else:
+            mask = torch.ones_like(token_jsd, dtype=torch.bool)
+
+        positions = torch.arange(token_jsd.shape[1], device=token_jsd.device, dtype=token_jsd.dtype)
+        position_weights = 1.0 / (1.0 + positions / position_scale)
+        weights = mask.to(token_jsd.dtype) * position_weights.unsqueeze(0)
 
         # Apply reduction
         if reduction == "batchmean":
-            return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / jsd.size(0)
+            loss = (token_jsd * weights).sum() / weights.sum().clamp_min(1.0)
         elif reduction == "sum":
-            return jsd.sum()
+            loss = (token_jsd * weights).sum()
         elif reduction == "mean":
-            return jsd.mean()
+            loss = token_jsd[mask].mean()
         else:
-            return jsd
+            return token_jsd
+
+        # Verified on-policy self-training complements the dense KL signal.  It
+        # reinforces only trajectories that actually reach the reference answer;
+        # incorrect trajectories still receive dense OPSD feedback instead of being
+        # discarded as in pure rejection sampling/STaR.
+        if outcome_weight > 0 and sampled_token_ids is not None and outcome_correct is not None:
+            correct_mask = mask & outcome_correct.to(torch.bool).unsqueeze(1)
+            if correct_mask.any():
+                sampled_log_probs = torch.gather(
+                    student_log_probs, dim=-1, index=sampled_token_ids.unsqueeze(-1)
+                ).squeeze(-1)
+                correct_weights = correct_mask.to(token_jsd.dtype) * position_weights.unsqueeze(0)
+                rejection_sft = -(sampled_log_probs * correct_weights).sum() / correct_weights.sum().clamp_min(1.0)
+                loss = loss + outcome_weight * rejection_sft
+
+        return loss
+
+    @staticmethod
+    def _extract_boxed_answer(text):
+        """Extract the final balanced \\boxed{...} expression."""
+        idx = text.rfind("\\boxed")
+        if idx < 0:
+            return None
+        left = text.find("{", idx)
+        if left < 0:
+            return None
+        depth = 0
+        for pos in range(left, len(text)):
+            if text[pos] == "{":
+                depth += 1
+            elif text[pos] == "}":
+                depth -= 1
+                if depth == 0:
+                    return text[left + 1 : pos]
+        return None
+
+    @staticmethod
+    def _normalize_direct_answer(answer):
+        if answer is None:
+            return None
+        answer = re.sub(r"\\(?:text|mathrm)\{([^{}]*)\}", r"\1", str(answer))
+        answer = answer.replace("$", "").replace("\\,", "").replace(",", "")
+        answer = re.sub(r"\s+", "", answer).lower()
+        if re.fullmatch(r"[+]?\d+", answer):
+            return str(int(answer))
+        return answer
 
     def _update_ema(self):
         """Update EMA parameters after an optimizer step.
@@ -741,6 +803,9 @@ class OPSDTrainer(SFTTrainer):
                 temperature=self.temperature,  # Let the function handle temperature
                 top_k=self.top_k_loss,
                 token_clip=self.jsd_token_clip,
+                sampled_token_ids=sampled_token_ids,
+                outcome_correct=inputs.get("outcome_correct"),
+                outcome_weight=self.lmbda,
             )
             del student_logits_for_loss, teacher_logits_for_loss
 
@@ -1418,6 +1483,17 @@ class OPSDTrainer(SFTTrainer):
 
         inputs["labels"] = labels
 
+        # Grade the sampled on-policy response before the forward pass.  The
+        # curriculum uses direct numeric answers, making exact normalization both
+        # fast and reliable (no symbolic grader or external verifier is involved).
+        outcome_correct = [
+            self._normalize_direct_answer(self._extract_boxed_answer(completion))
+            == self._normalize_direct_answer(answer)
+            for completion, answer in zip(completion_texts, inputs["answer_texts"])
+        ]
+        inputs["outcome_correct"] = torch.tensor(outcome_correct, dtype=torch.bool, device=generated_ids.device)
+        self._metrics["train"]["rollout_accuracy"].append(sum(outcome_correct) / max(1, len(outcome_correct)))
+
         # Log prompt and completion texts
         self._textual_logs["prompt"].extend(gather_object(prompt_texts))
         self._textual_logs["completion"].extend(gather_object(completion_texts))
lmbda=0.01
lr_scheduler_type=linear
cache-093442ee98f5e5e5.arrow 472555832
cache-55f01c4829935f27.arrow 297287784
cache-a0c8e8754bd7bb38.arrow 288198312
cache-a47d60ea3fed15a4.arrow 182299336
cache-ab5fafcaaba4eebd.arrow 472555832
cache-d162f0b4cc47513f.arrow 297287784
9276

Script completed
Wall time 0.2 seconds
Output:
2026-08-02 17:14:28.6370304490 1f33e9dc2e8a1c639ca74f8024ad4a9f1f5eae62_builder.lock 0
2026-08-02 17:14:36.6819176270 openthoughts_math_30k_opsd-train-00000-of-00002.arrow 537225960
2026-08-02 17:14:36.7019174960 openthoughts_math_30k_opsd-train-00001-of-00002.arrow 116664552
2026-08-02 17:14:36.7039174830 dataset_info.json 1562
2026-08-02 17:14:36.7049412940 1f33e9dc2e8a1c639ca74f8024ad4a9f1f5eae62.incomplete_info.lock 0
_opt_hf_cache_datasets_siyanzhao___openthoughts_math_30k_opsd_default_0.0.0_1f33e9dc2e8a1c639ca74f8024ad4a9f1f5eae62.lock 0
v1 [56.388888888888886, 53.888888888888886] mean 55.138888888888886
v2 [55.55555555555556, 54.44444444444444] mean 55.0
v3 [55.0] mean 55.0

Script completed
Wall time 0.1 seconds
Output:
2026-08-03 04:38:34.0678932140 /opt/hf_cache/modules/__init__.py 0
2026-08-03 04:38:34.0680008830 /opt/hf_cache/modules/datasets_modules/__init__.py 0
2026-08-03 07:21:59.2641770310 /opt/hf_cache/datasets/_opt_hf_cache_datasets_siyanzhao___openthoughts_math_30k_opsd_default_0.0.0_1f33e9dc2e8a1c639ca74f8024ad4a9f1f5eae62.lock 0

Script completed
Wall time 3.4 seconds
Output:
da57b57c57b35fe5d98a61273603ce1a7513a42368b6ed209c22409338369ead  /opt/hf_cache/datasets/siyanzhao___openthoughts_math_30k_opsd/default/0.0.0/1f33e9dc2e8a1c639ca74f8024ad4a9f1f5eae62/openthoughts_math_30k_opsd-train-00000-of-00002.arrow
727ba2b8c54d190651709e5a5f323e844fc90dcbf03309197950ef070595f2b5  /opt/hf_cache/datasets/siyanzhao___openthoughts_math_30k_opsd/default/0.0.0/1f33e9dc2e8a1c639ca74f8024ad4a9f1f5eae62/openthoughts_math_30k_opsd-train-00001-of-00002.arrow
1535163387e4c5da3a21aa51766d31261806e335d1cb141104cb285aa7b3d76c  /opt/hf_cache/datasets/siyanzhao___openthoughts_math_30k_opsd/default/0.0.0/1f33e9dc2e8a1c639ca74f8024ad4a9f1f5eae62/dataset_info.json

Script completed
Wall time 1.2 seconds
Output:
{}{"chunk_id":"fe5474","wall_time_seconds":1.001657935,"session_id":89779,"original_token_count":0,"output":""}

Script running with cell ID 253
Wall time 11.0 seconds
Output:

Script completed
Wall time 32.4 seconds
Output:
Warning: truncated output (original token count: 207)
Total output lines: 10

Warning: truncated output (original token count: 9091)
Total output lines: 94

<string>:24: FutureWarning: `torch_dtype` is deprecated and will be removed in version 0.27.0, please use `dtype` instead.

================================================================================
RUN CONFIGURATION
======================================================…27 tokens truncated…███▋   | 5856/8822 [00:38<00:16, 175.37 examples/s]
Tokenizing train dataset:  67%|██████▋   | 5877/8822 [00:38<00:16, 177.88 examples/s]
Tokenizing train dataset:  67%|██████▋   | 5905/8822 [00:38<00:16, 178.30 examples/s]
Tokenizing train dataset:  67%|██████▋   | 5925/8822 [00:39<00:15, 181.21 examples/s]

Script running with cell ID 254
Wall time 11.0 seconds
Output:

Script completed
Wall time 30.2 seconds
Output:
Warning: truncated output (original token count: 207)
Total output lines: 5

Warning: truncated output (original token count: 6431)
Total output lines: 105


Tokenizing train dataset:  67%|██████▋   | 5945/8822 [00:39<00:16, 178.52 examples/s]
Tokenizing train dataset:  68%|██████▊   | 5965/8822 [00:39<00:15, 182.05 examples/s]
Tokenizing train dataset:  68%|██████▊   | 5984/8822 [00:39<00:15…27 tokens truncated…                                              
{'loss': -0.0032, 'grad_norm': 0.1673268973827362, 'learning_rate': 4.95e-06, 'on_policy_loss': -0.0032, 'rollout_accuracy': 0.375, 'epoch': 0.01}

  2%|▏         | 2/100 [00:18<14:46,  9.04s/it]vLLM generation done - elapsed time: 1.78s, prompts: 4, total tokens: 1279, avg length: 319.8, speed: 720.0 tok/s

Script running with cell ID 255
Wall time 11.0 seconds
Output:

Script completed
Wall time 32.1 seconds
Output:
Warning: truncated output (original token count: 187)
Total output lines: 8

Warning: truncated output (original token count: 2750)
Total output lines: 154

vLLM generation done - elapsed time: 3.68s, prompts: 4, total tokens: 3051, avg length: 762.8, speed: 829.1 tok/s
vLLM generation done - elapsed time: 3.69s, prompts: 4, total tokens: 3269, avg length: 817.2, speed: 886.7 tok/s
vLLM gener…27 tokens truncated…apsed time: 3.66s, prompts: 4, total tokens: 3093, avg length: 773.2, speed: 845.2 tok/s
vLLM generation done - elapsed time: 3.67s, prompts: 4, total tokens: 3473, avg length: 868.2, speed: 947.2 tok/s
vLLM generation done - elapsed time: 3.70s, prompts: 4, total tokens: 3088, avg length: 772.0, speed: 834.4 tok/s

Script running with cell ID 256
Wall time 11.0 seconds
Output:

Script completed
Wall time 20.2 seconds
Output:
Warning: truncated output (original token count: 187)
Total output lines: 8

Warning: truncated output (original token count: 3265)
Total output lines: 195

vLLM generation done - elapsed time: 3.57s, prompts: 4, total tokens: 2704, avg length: 676.0, speed: 756.6 tok/s
vLLM generation done - elapsed time: 3.70s, prompts: 4, total tokens: 3280, avg length: 820.0, speed: 886.4 tok/s
vLLM gener…27 tokens truncated…psed time: 3.72s, prompts: 4, total tokens: 2768, avg length: 692.0, speed: 744.5 tok/s
vLLM generation done - elapsed time: 3.78s, prompts: 4, total tokens: 3241, avg length: 810.2, speed: 857.0 tok/s
vLLM generation done - elapsed time: 3.84s, prompts: 4, total tokens: 3916, avg length: 979.0, speed: 1019.7 tok/s

Script running with cell ID 257
Wall time 11.0 seconds
Output: