gpt-5.6-sol

Script running with cell ID 117
Wall time 61.0 seconds
Output:

Script running with cell ID 117
Wall time 61.0 seconds
Output:

Script running with cell ID 117
Wall time 61.0 seconds
Output:

Script running with cell ID 117
Wall time 61.0 seconds
Output:

Script running with cell ID 117
Wall time 61.0 seconds
Output:

Script running with cell ID 117
Wall time 61.0 seconds
Output:

Script running with cell ID 117
Wall time 61.0 seconds
Output:

Script running with cell ID 117
Wall time 61.0 seconds
Output:

Script completed
Wall time 28.8 seconds
Output:
Warning: truncated output (original token count: 226)
Total output lines: 21

     | 0/96 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]


 [03:22<08:29,  7.08s/it, est. speed input: 19.11 toks/s, output: 862.73 toks…146 tokens truncated…<05:56,  9.91s/it, est. speed input: 8.49 toks/s, output: 1318.83 toks/s]
0 problems)
Results written to /workspace/eval_results/dev_eval_1785711242.json

EXIT0

Script completed
Wall time 0.1 seconds
Output:
57.22222222222222 206
[(60, 12), (61, 12), (62, 0), (63, 0), (64, 5), (65, 5), (66, 11), (67, 12), (68, 11), (69, 11), (70, 7), (71, 12), (72, 12), (73, 0), (74, 8), (75, 8), (76, 9), (77, 2), (78, 5), (79, 12), (80, 3), (81, 0), (82, 8), (83, 12), (84, 12), (85, 0), (86, 11), (87, 6), (88, 0), (89, 0)]
4597

Script completed
Wall time 0.1 seconds
Output:
fatal: /workspace/submission/recipe.env: '/workspace/submission/recipe.env' is outside repository at '/workspace/repo'

Script completed
Wall time 0.1 seconds
Output:
diff --git a/data_collator.py b/data_collator.py
index 0829108..221b555 100644
--- a/data_collator.py
+++ b/data_collator.py
@@ -55,6 +55,7 @@ class SelfDistillationDataCollator:
         # Prepare student and teacher prompts using chat template (matching evaluation)
         student_prompts = []
         teacher_prompts = []
+        anchor_sequences = []
         teacher_reasoning_prompts = []  # NEW: for reason_first mode
 
         for feature in features:
@@ -73,6 +74,13 @@ class SelfDistillationDataCollator:
             )
             student_prompts.append(student_prompt)
 
+            # Expert trajectory used as a weak anchor for OPSD.  It starts at
+            # exactly the same non-privileged student prompt; unlike the
+            # teacher context below it never exposes the solution in the user
+            # message.  The reference solution is available in every row of
+            # the released OPSD dataset.
+            anchor_sequences.append(student_prompt + solution.strip() + self.tokenizer.eos_token)
+
             if self.reason_first:
                 # Reasoning prompt: ask teacher to analyze the solution
                 reasoning_user_message = (
@@ -138,6 +146,25 @@ class SelfDistillationDataCollator:
             "student_prompt_lengths_per_example": torch.tensor(student_prompt_lengths),
         }
 
+        anchor_encoded = self.tokenizer(
+            anchor_sequences,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="pt",
+        )
+        anchor_labels = anchor_encoded["input_ids"].clone()
+        for i, prompt_len in enumerate(student_prompt_lengths):
+            anchor_labels[i, :prompt_len] = -100
+        anchor_labels[anchor_encoded["attention_mask"] == 0] = -100
+        result.update(
+            {
+                "anchor_input_ids": anchor_encoded["input_ids"],
+                "anchor_attention_mask": anchor_encoded["attention_mask"],
+                "anchor_labels": anchor_labels,
+            }
+        )
+
         if self.reason_first:
             # Tokenize reasoning prompts
             reasoning_encoded_no_pad = self.tokenizer(
diff --git a/opsd_train.py b/opsd_train.py
index 8526a68..b4be7a4 100644
--- a/opsd_train.py
+++ b/opsd_train.py
@@ -266,6 +266,19 @@ if __name__ == "__main__":
     dataset = load_dataset("siyanzhao/Openthoughts_math_30k_opsd")
     train_dataset = dataset["train"]
 
+    # Reasoning-effort curriculum.  `generated_token_count` is metadata shipped
+    # with the released dataset and measures how much reasoning was needed to
+    # obtain its verified solution.  Under a 100-step budget, uniformly drawing
+    # many short arithmetic items wastes privileged-teacher calls.  Retain the
+    # broad source mix but train on examples requiring at least 3k reasoning
+    # tokens, a model-native difficulty signal rather than a benchmark label.
+    train_dataset = train_dataset.filter(
+        lambda example: example["generated_token_count"] >= 3000,
+        desc="Selecting high-reasoning-effort OPSD examples",
+    )
+    if os.environ.get("LOCAL_RANK", "0") == "0":
+        print(f"Reasoning-effort curriculum: {len(train_dataset)} training rows")
+
     trainer = OPSDTrainer(
         model=model_args.model_name_or_path,
         args=training_args,
diff --git a/opsd_trainer.py b/opsd_trainer.py
index c151bd7..8011b79 100644
--- a/opsd_trainer.py
+++ b/opsd_trainer.py
@@ -428,15 +428,31 @@ class OPSDTrainer(SFTTrainer):
             teacher_logits = teacher_logits / temperature
 
             if top_k is not None and top_k > 0:
-                # Restrict to top-k tokens of the teacher distribution and renormalize.
-                # Shape: [batch, seq_len, top_k]
+                # Teacher-support KL: retain the teacher's most actionable
+                # alternatives and aggregate every remaining vocabulary item
+                # into one exact "other" bucket.  Renormalizing only the
+                # selected logits (the old implementation) discarded the
+                # student's probability leakage outside teacher support and
+                # was not a projection of the full distributions.
                 _, top_k_indices = torch.topk(teacher_logits, k=top_k, dim=-1)
-                student_logits = torch.gather(student_logits, dim=-1, index=top_k_indices)
-                teacher_logits = torch.gather(teacher_logits, dim=-1, index=top_k_indices)
-
-            # Compute log probabilities for student and probabilities for teacher
-            student_log_probs = F.log_softmax(student_logits, dim=-1)
-            teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)
+                student_full_log_probs = F.log_softmax(student_logits, dim=-1)
+                teacher_full_log_probs = F.log_softmax(teacher_logits, dim=-1)
+                student_head = torch.gather(student_full_log_probs, dim=-1, index=top_k_indices)
+                teacher_head = torch.gather(teacher_full_log_probs, dim=-1, index=top_k_indices)
+                # Accumulate retained mass in fp32: in bf16, 1-1e-6 rounds to
+                # one and log1p(-mass) can become -inf.
+                student_mass = student_head.float().exp().sum(dim=-1, keepdim=True).clamp(max=1 - 1e-5)
+                teacher_mass = teacher_head.float().exp().sum(dim=-1, keepdim=True).clamp(max=1 - 1e-5)
+                student_tail = torch.log1p(-student_mass)
+                teacher_tail = torch.log1p(-teacher_mass)
+                student_head = student_head.float()
+                teacher_head = teacher_head.float()
+                student_log_probs = torch.cat([student_head, student_tail], dim=-1)
+                teacher_log_probs = torch.cat([teacher_head, teacher_tail], dim=-1)
+            else:
+                # Compute log probabilities for student and teacher
+                student_log_probs = F.log_softmax(student_logits, dim=-1)
+                teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)
 
         if beta == 0:
             jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True)
@@ -459,7 +475,9 @@ class OPSDTrainer(SFTTrainer):
             # Compute the Generalized Jensen-Shannon Divergence
             jsd = beta * kl_teacher + (1 - beta) * kl_student
 
-        # Per-token clipping: cap each token's divergence value
+        # Per-token clipping: cap each vocabulary contribution.  Although the
+        # individual KL summands are signed, retaining this released behavior
+        # is important for the paper's point-wise robustness heuristic.
         if token_clip is not None:
             jsd = jsd.clamp(max=token_clip)
 
@@ -636,13 +654,39 @@ class OPSDTrainer(SFTTrainer):
         shifted_labels = inputs["labels"][:, student_prompt_len:]
 
         # === STUDENT FORWARD - Extract log-probs immediately ===
-        outputs_student = model(
-            input_ids=inputs["student_input_ids"],
-            attention_mask=inputs["student_attention_mask"],
-        )
+        # Put the on-policy and correct anchor trajectories in one forward.
+        # Besides being cheaper than two calls, this is required by ZeRO-2,
+        # whose overlapping reducer cannot visit the same LoRA parameter from
+        # two checkpointed graphs in one backward pass.
+        use_anchor = self.lmbda > 0 and "anchor_input_ids" in inputs and not self.use_thinking_machines_loss
+        student_batch_size = inputs["student_input_ids"].shape[0]
+        student_seq_len = inputs["student_input_ids"].shape[1]
+        if use_anchor:
+            anchor_seq_len = inputs["anchor_input_ids"].shape[1]
+            combined_seq_len = max(student_seq_len, anchor_seq_len)
+            pad_id = self.processing_class.pad_token_id
+            student_ids = F.pad(inputs["student_input_ids"], (0, combined_seq_len - student_seq_len), value=pad_id)
+            student_mask = F.pad(inputs["student_attention_mask"], (0, combined_seq_len - student_seq_len), value=0)
+            anchor_ids = F.pad(inputs["anchor_input_ids"], (0, combined_seq_len - anchor_seq_len), value=pad_id)
+            anchor_mask = F.pad(inputs["anchor_attention_mask"], (0, combined_seq_len - anchor_seq_len), value=0)
+            forward_ids = torch.cat([student_ids, anchor_ids], dim=0)
+            forward_mask = torch.cat([student_mask, anchor_mask], dim=0)
+        else:
+            forward_ids = inputs["student_input_ids"]
+            forward_mask = inputs["student_attention_mask"]
+
+        outputs_student = model(input_ids=forward_ids, attention_mask=forward_mask)
 
         # Extract only what we need and convert to log-probs immediately
-        student_logits = outputs_student.logits[:, student_prompt_len - 1 : -1, :]
+        student_logits = outputs_student.logits[:student_batch_size, student_prompt_len - 1 : student_seq_len - 1, :]
+
+        if use_anchor:
+            anchor_logits = outputs_student.logits[student_batch_size:, :anchor_seq_len, :]
+            anchor_loss = F.cross_entropy(
+                anchor_logits[:, :-1, :].reshape(-1, anchor_logits.shape[-1]),
+                inputs["anchor_labels"][:, 1:].reshape(-1),
+                ignore_index=-100,
+            )
 
         if self.use_thinking_machines_loss:
             # For reverse KL, we only need log-probs of sampled tokens
@@ -733,7 +777,7 @@ class OPSDTrainer(SFTTrainer):
             )
         else:
             # Temperature is applied inside generalized_jsd_loss
-            loss = self.generalized_jsd_loss(
+            opsd_loss = self.generalized_jsd_loss(
                 student_logits=student_logits_for_loss,
                 teacher_logits=teacher_logits_for_loss,
                 labels=shifted_labels,
@@ -744,6 +788,18 @@ class OPSDTrainer(SFTTrainer):
             )
             del student_logits_for_loss, teacher_logits_for_loss
 
+            # OPSD only visits student trajectories.  A weak maximum-
+            # likelihood anchor on the known-correct trajectory prevents
+            # self-distillation from polishing an incorrect branch forever.
+            # `lmbda` is the anchor coefficient; setting it to zero recovers
+            # pure OPSD exactly.
+            if use_anchor:
+                loss = opsd_loss + self.lmbda * anchor_loss
+                self._metrics["train"]["opsd_kl"].append(float(opsd_loss.detach()))
+                self._metrics["train"]["anchor_ce"].append(float(anchor_loss.detach()))
+            else:
+                loss = opsd_loss
+
         empty_cache()
 
         if return_outputs:
RECIPE
# Difficulty-curriculum OPSD with the released optimizer/model settings.
lmbda=0
top_k_loss=0
lora_r=64
lora_alpha=128
max_completion_length=1024
lr_scheduler_type=linear

Script completed
Wall time 0.0 seconds
Output:
{}

Script completed
Wall time 0.2 seconds
Output:
 README.md     |  4 +---
 opsd_train.py | 13 +++++++++++++
 2 files changed, 14 insertions(+), 3 deletions(-)
diff --git a/opsd_train.py b/opsd_train.py
index 8526a68..b4be7a4 100644
--- a/opsd_train.py
+++ b/opsd_train.py
@@ -266,6 +266,19 @@ if __name__ == "__main__":
     dataset = load_dataset("siyanzhao/Openthoughts_math_30k_opsd")
     train_dataset = dataset["train"]
 
+    # Reasoning-effort curriculum.  `generated_token_count` is metadata shipped
+    # with the released dataset and measures how much reasoning was needed to
+    # obtain its verified solution.  Under a 100-step budget, uniformly drawing
+    # many short arithmetic items wastes privileged-teacher calls.  Retain the
+    # broad source mix but train on examples requiring at least 3k reasoning
+    # tokens, a model-native difficulty signal rather than a benchmark label.
+    train_dataset = train_dataset.filter(
+        lambda example: example["generated_token_count"] >= 3000,
+        desc="Selecting high-reasoning-effort OPSD examples",
+    )
+    if os.environ.get("LOCAL_RANK", "0") == "0":
+        print(f"Reasoning-effort curriculum: {len(train_dataset)} training rows")
+
     trainer = OPSDTrainer(
         model=model_args.model_name_or_path,
         args=training_args,
 M README.md
 M opsd_train.py
4537

Script completed
Wall time 0.0 seconds
Output:
{}

Script completed
Wall time 1.2 seconds
Output:
SESSION_ID=70713

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

Script running with cell ID 125
Wall time 61.0 seconds
Output:

Script running with cell ID 125
Wall time 61.0 seconds
Output:

Script running with cell ID 125
Wall time 61.0 seconds
Output:

Script running with cell ID 125
Wall time 61.0 seconds
Output:

Script running with cell ID 125
Wall time 61.0 seconds
Output:

Script running with cell ID 125
Wall time 61.0 seconds
Output: