on-policy-self-distillation
gpt-5.6-sol
on-policy-self-distillation — graded by the task's verifier against a measured baseline.
- backend
- modal
- variant
- codex azure_ai/gpt-5.6-sol
- metric
- AIME24 avg@12 accuracy
- actions
- 596
- duration
- 7h 59m
- recorded
- 2026-08-03
run idopsd-aime24-4gpus__kuc95M7
Everything gpt-5.6-sol did on this task, in order.
download JSON ↗codexgpt-5.6-sol669 turns73 narration596 reasoning596 action596 observation1 final
#576
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolwait
args
{
"cell_id": "280",
"yield_time_ms": 60000,
"max_tokens": 300
}wait result
Script completed
Wall time 43.7 seconds
Output:
Processed prompts: 50%|█████ | 48/96 [14:20<10:54, 13.63s/it, est. speed input: 7.70 toks/s, output: 1030.96 toks/s]
Processed prompts: 71%|███████▏ | 60/84 [14:22<03:02, 7.59s/it, est. speed input: 7.10 toks/s, output: 1392.40 toks/s]
Processed prompts: 86%|████████▌ | 72/84 [14:26<01:01, 5.13s/it, est. speed input: 8.43 toks/s, output: 1781.34 toks/s]
Processed prompts: 62%|██████▎ | 60/96 [14:50<05:46, 9.63s/it, est. speed input: 8.91 toks/s, output: 1335.74 toks/s]#577
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.write_stdin({session_id: 58459, chars: \"\", yield_time_ms: 60000, max_output_tokens: 300}); text(r.output.slice(-600));\n"
}exec result
Script running with cell ID 281
Wall time 11.0 seconds
Output:
#578
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolwait
args
{
"cell_id": "281",
"yield_time_ms": 60000,
"max_tokens": 300
}wait result
Script completed
Wall time 45.0 seconds
Output:
| 72/96 [15:29<02:58, 7.44s/it, est. speed input: 10.14 toks/s, output: 1536.03 toks/s]
Processed prompts: 88%|████████▊ | 84/96 [16:10<01:13, 6.12s/it, est. speed input: 11.32 toks/s, output: 1770.06 toks/s]
Processed prompts: 100%|██████████| 96/96 [16:10<00:00, 6.12s/it, est. speed input: 13.48 toks/s, output: 2120.18 toks/s]
Processed prompts: 100%|██████████| 96/96 [16:10<00:00, 10.10s/it, est. speed input: 13.48 toks/s, output: 2120.18 toks/s]
FINAL avg@12 accuracy: 56.11% (202/360 solutions correct on 30 problems)
Results written to /workspace/eval_results/dev_eval_1785747609.json
#579
narration
The trajectory-balanced variant scored exactly 56.11% (202/360), identical to both prior mild front-loaded runs. Since it adds complexity without evidence of gain, I’m reverting that extra reduction and retaining the simpler method with three consistent full-eval outcomes.
#580
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"sed -n '35,125p' opsd_trainer.py && cat /workspace/.timer/remaining_secs && git status --short\",workdir:\"/workspace/repo\",yield_time_ms:10000,max_output_tokens:4000}); text(r.output);\n"
}exec result
Script completed
Wall time 0.1 seconds
Output:
from transformers.image_processing_utils import BaseImageProcessor
from transformers.integrations.integration_utils import is_wandb_available
from transformers.modeling_utils import PreTrainedModel
from transformers.processing_utils import ProcessorMixin
from transformers.tokenization_utils_base import PreTrainedTokenizerBase
from transformers.trainer_callback import TrainerCallback, TrainerControl, TrainerState
from transformers.trainer_utils import EvalPrediction
from transformers.utils import (
is_flash_attn_2_available,
is_liger_kernel_available,
is_peft_available,
is_rich_available,
)
from trl.data_utils import is_conversational, maybe_convert_to_chatml, pack_dataset, truncate_dataset
from trl.extras.profiling import profiling_decorator
from trl.extras.vllm_client import VLLMClient
from trl.import_utils import is_vllm_available
from trl.models import prepare_deepspeed
from trl.models.utils import unwrap_model_for_generation
from trl.trainer.sft_trainer import SFTTrainer
from trl.trainer.utils import (
DataCollatorForChatML,
disable_dropout_in_model,
empty_cache,
ensure_master_addr_port,
pad,
)
from trl.experimental.gold.gold_config import GOLDConfig
from data_collator import SelfDistillationDataCollator
if is_peft_available():
from peft import PeftConfig
if is_wandb_available():
import wandb
if is_vllm_available():
from vllm import LLM, SamplingParams
from vllm.sampling_params import GuidedDecodingParams
if is_rich_available():
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.text import Text
class EMAUpdateCallback(TrainerCallback):
"""Update EMA teacher weights after each optimizer step."""
def __init__(self, trainer):
self.trainer = trainer
def on_step_end(self, args, state: TrainerState, control: TrainerControl, **kwargs):
# Only update when the optimizer actually stepped (end of a gradient accumulation cycle)
if self.trainer.use_ema_teacher and self.trainer.accelerator.sync_gradients:
self.trainer._update_ema()
class GOLDVLLMSyncCallback(TrainerCallback):
"""Sync the model weights to vLLM after training steps when it's safe to do so."""
def __init__(self, trainer):
self.trainer = trainer
def on_step_end(self, args, state: TrainerState, control: TrainerControl, **kwargs):
"""Sync weights after training step when DeepSpeed is stable."""
if (
self.trainer.use_vllm
and state.global_step != self.trainer._last_vllm_sync_step
and state.global_step % self.trainer.vllm_sync_frequency == 0
):
# Check if this is a step where gradients are synchronized
# This happens at the end of gradient accumulation cycles
if (
hasattr(self.trainer.accelerator, "sync_gradients")
and self.trainer.accelerator.sync_gradients
):
self.trainer._move_model_to_vllm()
self.trainer._last_vllm_sync_step = state.global_step
class OPSDTrainer(SFTTrainer):
_tag_names = ["trl", "opsd"]
_name = "OPSD"
def __init__(
self,
model: PreTrainedModel | nn.Module | str | None = None,
4833
M README.md
M opsd_trainer.py
#581
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"rg -n \\\"token_counts|credit =|generalized_jsd_loss\\\" opsd_trainer.py\",workdir:\"/workspace/repo\",yield_time_ms:10000,max_output_tokens:1000}); text(r.output);\n"
}exec result
Script completed
Wall time 0.1 seconds
Output:
382: def generalized_jsd_loss(
479: credit = (1.5 - progress).clamp_min(0.5) * mask
480: credit = torch.where(lengths > 1, credit, mask.to(token_divergence.dtype))
489: token_counts = mask.sum(dim=-1).clamp_min(1)
490: per_trajectory = (token_divergence * credit).sum(dim=-1) / token_counts
677: # For JSD, keep logits (temperature will be applied in generalized_jsd_loss)
757: # Temperature is applied inside generalized_jsd_loss
758: loss = self.generalized_jsd_loss(
#582
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"sed -n '382,505p' opsd_trainer.py\",workdir:\"/workspace/repo\",yield_time_ms:10000,max_output_tokens:4000}); text(r.output);\n"
}exec result
Script completed
Wall time 0.1 seconds
Output:
def generalized_jsd_loss(
student_logits,
teacher_logits,
labels=None,
beta=0.5,
temperature=1.0,
reduction="batchmean",
logits_are_probs=False,
top_k=None,
token_clip=None,
):
"""
Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation using F.kl_div. See Eq. (1)
of https://huggingface.co/papers/2306.13649 for the definition.
Args:
student_logits:
Tensor of shape (batch_size, sequence_length, vocab_size)
teacher_logits:
Tensor of shape (batch_size, sequence_length, vocab_size)
labels:
Tensor of shape (batch_size, sequence_length) with -100 for padding tokens to ignore when computing
loss
beta:
Interpolation coefficient between 0 and 1 (default: 0.5)
temperature:
Softmax temperature (default: 1.0)
reduction:
Specifies the reduction to apply to the output (default: 'batchmean')
top_k:
If set, restricts the loss to only the top-k tokens of the teacher distribution. Both student and
teacher distributions are renormalized over these k tokens before computing JSD. This reduces memory
and focuses distillation on the teacher's most probable tokens. (default: None = full vocabulary)
token_clip:
if set, clips per-token divergence values to this maximum before reduction. Prevents style tokens from dominating the gradient signal over math tokens.
Returns:
loss: Scalar tensor with the generalized JSD loss
"""
if logits_are_probs:
student_log_probs = torch.log(student_logits.clamp_min(1e-8))
teacher_log_probs = torch.log(teacher_logits.clamp_min(1e-8))
else:
# Apply temperature scaling to logits before computing probabilities
student_logits = student_logits / temperature
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]
_, 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)
if beta == 0:
jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True)
elif beta == 1:
jsd = F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True)
else:
# Compute the log of the mixture distribution
# log(a + b) = log(exp(log(a)) + exp(log(b))) -> for mixture
beta = torch.tensor(beta, dtype=student_log_probs.dtype, device=student_log_probs.device)
mixture_log_probs = torch.logsumexp(
torch.stack([student_log_probs + torch.log1p(-beta), teacher_log_probs + torch.log(beta)]),
dim=0,
)
# Compute KL divergences using F.kl_div
# PyTorch differs from the standard mathematical definition, so the order of the probability distributions is swapped compared to that defined in the paper.
kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True)
kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True)
# Compute the Generalized Jensen-Shannon Divergence
jsd = beta * kl_teacher + (1 - beta) * kl_student
# Per-token clipping: cap each token's divergence value
if token_clip is not None:
jsd = jsd.clamp(max=token_clip)
# Masking and trajectory credit assignment. Early reasoning tokens are
# more likely to encode a consequential branch choice; late tokens are
# increasingly determined by the already-long prefix. Give early
# positions moderately more credit without changing the total loss
# scale. The linear 1.5 -> 0.5 schedule has mean one on every non-empty
# trajectory, so learning-rate and clipping behavior remain comparable
# to the original uniformly averaged OPSD objective.
if labels is not None:
mask = labels != -100
token_divergence = jsd.sum(dim=-1)
positions = mask.long().cumsum(dim=-1) - 1
lengths = mask.sum(dim=-1, keepdim=True).clamp_min(1)
progress = positions.to(token_divergence.dtype) / (lengths - 1).clamp_min(1)
credit = (1.5 - progress).clamp_min(0.5) * mask
credit = torch.where(lengths > 1, credit, mask.to(token_divergence.dtype))
# Apply reduction
if reduction == "batchmean":
if labels is not None:
# Match the paper's nested expectation: average tokens within
# each sampled trajectory first, then average trajectories. A
# flat token average silently gives long/capped generations
# more influence than concise ones.
token_counts = mask.sum(dim=-1).clamp_min(1)
per_trajectory = (token_divergence * credit).sum(dim=-1) / token_counts
valid_trajectories = mask.any(dim=-1)
return per_trajectory[valid_trajectories].mean()
return jsd.sum() / jsd.size(0)
elif reduction == "sum":
return (token_divergence * credit).sum() if labels is not None else jsd.sum()
elif reduction == "mean":
if labels is not None:
return (token_divergence * credit).sum() / credit.sum().clamp_min(1)
return jsd.mean()
else:
return token_divergence * credit if labels is not None else jsd
def _update_ema(self):
"""Update EMA parameters after an optimizer step.
#583
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const patch = \"*** Begin Patch\\n*** Update File: /workspace/repo/opsd_trainer.py\\n@@\\n # Apply reduction\\n if reduction == \\\"batchmean\\\":\\n if labels is not None:\\n- # Match the paper's nested expectation: average tokens within\\n- # each sampled trajectory first, then average trajectories. A\\n- # flat token average silently gives long/capped generations\\n- # more influence than concise ones.\\n- token_counts = mask.sum(dim=-1).clamp_min(1)\\n- per_trajectory = (token_divergence * credit).sum(dim=-1) / token_counts\\n- valid_trajectories = mask.any(dim=-1)\\n- return per_trajectory[valid_trajectories].mean()\\n+ return (token_divergence * credit).sum() / mask.sum().clamp_min(1)\\n return jsd.sum() / jsd.size(0)\\n*** End Patch\";\ntext(await tools.apply_patch(patch));\n"
}exec result
Script completed
Wall time 0.0 seconds
Output:
{}#584
narration
The simpler retained code is restored. With the remaining budget, I’m checking whether the repository already exposes a principled auxiliary objective that can complement early-credit OPSD; I’ll only run another full candidate if it stays within the same 100-step/data constraints and has a credible mechanism.
#585
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"rg -n \\\"lmbda|reason_first|tinker|teacher_thinking|solution|supervised|labels\\\" opsd_train.py opsd_trainer.py data_collator.py | head -240\",workdir:\"/workspace/repo\",yield_time_ms:10000,max_output_tokens:10000}); text(r.output);\n"
}exec result
Script completed
Wall time 0.1 seconds
Output:
data_collator.py:9: Teacher: sees problem + solution + transition prompt (with chat template)
data_collator.py:19: reason_first=True,
data_collator.py:21: teacher_thinking=True,
data_collator.py:25: self.reason_first = reason_first
data_collator.py:27: self.teacher_thinking = teacher_thinking
data_collator.py:29: # Prompt for reasoning about the solution before teaching
data_collator.py:30: self.reason_first_prompt = (
data_collator.py:32: "Please analyze this solution and explain the key reasoning steps and problem-solving strategies employed. "
data_collator.py:33: "Do NOT use <think> tags. Do NOT derive your own solution. "
data_collator.py:34: "Simply analyze and explain the reference solution provided above.\n"
data_collator.py:38: "\n\nAfter reading the reference solution above, make sure you truly understand "
data_collator.py:49: print(f"[DataCollator] Reason first mode: {self.reason_first}")
data_collator.py:58: teacher_reasoning_prompts = [] # NEW: for reason_first mode
data_collator.py:61: # Extract problem and solution from dataset
data_collator.py:64: solution = feature["solution"]
data_collator.py:76: if self.reason_first:
data_collator.py:77: # Reasoning prompt: ask teacher to analyze the solution
data_collator.py:82: f"{solution}\n"
data_collator.py:84: f"{self.reason_first_prompt}"
data_collator.py:99: f"Here is a reference solution to this problem:\n"
data_collator.py:100: f"=== Reference Solution Begin ===\n{solution}\n=== Reference Solution End ===\n"
data_collator.py:108: teacher_messages, tokenize=False, add_generation_prompt=True, enable_thinking=self.teacher_thinking
data_collator.py:141: if self.reason_first:
opsd_trainer.py:140: reason_first: bool = False,
opsd_trainer.py:146: teacher_thinking: bool = True,
opsd_trainer.py:159: reason_first=reason_first,
opsd_trainer.py:161: teacher_thinking=teacher_thinking,
opsd_trainer.py:181: self.lmbda = args.lmbda
opsd_trainer.py:188: self.reason_first = reason_first
opsd_trainer.py:223: if self.reason_first:
opsd_trainer.py:226: print("Teacher will first reason about the privileged solution, then evaluate student's response")
opsd_trainer.py:256: # Generation config for reasoning phase (when reason_first=True)
opsd_trainer.py:372: "solution",
opsd_trainer.py:385: labels=None,
opsd_trainer.py:402: labels:
opsd_trainer.py:473: if labels is not None:
opsd_trainer.py:474: mask = labels != -100
opsd_trainer.py:484: if labels is not None:
opsd_trainer.py:488: return (token_divergence * credit).sum() if labels is not None else jsd.sum()
opsd_trainer.py:490: if labels is not None:
opsd_trainer.py:494: return token_divergence * credit if labels is not None else jsd
opsd_trainer.py:651: shifted_labels = inputs["labels"][:, student_prompt_len:]
opsd_trainer.py:732: if shifted_labels is not None:
opsd_trainer.py:733: mask = shifted_labels != -100
opsd_trainer.py:754: labels=shifted_labels,
opsd_trainer.py:773: """Generate teacher's reasoning about the solution."""
opsd_trainer.py:861: new_labels = generated_tokens.clone()
opsd_trainer.py:864: new_labels[new_labels == pad_token_id] = -100
opsd_trainer.py:867: return generated_tokens, new_attention_mask, new_labels
opsd_trainer.py:1044: new_labels = new_input_ids.clone()
opsd_trainer.py:1047: new_labels[new_labels == pad_token_id] = -100
opsd_trainer.py:1056: return new_input_ids, new_attention_mask, new_labels, prompts_text_with_special, completion_texts
opsd_trainer.py:1309: If reason_first=True:
opsd_trainer.py:1310: 1. Generate teacher's reasoning about the solution
opsd_trainer.py:1323: if self.reason_first:
opsd_trainer.py:1325: print("REASONING PHASE: Teacher analyzing solution...")
opsd_trainer.py:1424: # Create labels for generation tokens
opsd_trainer.py:1426: labels = generated_ids.clone()
opsd_trainer.py:1427: for i in range(labels.shape[0]):
opsd_trainer.py:1429: labels[i, :actual_prompt_len] = -100 # Mask actual prompt
opsd_trainer.py:1432: labels[labels == self.processing_class.pad_token_id] = -100
opsd_trainer.py:1434: inputs["labels"] = labels
opsd_train.py:28: use_tinker_loss: bool = field(
opsd_train.py:57: reason_first: bool = field(
opsd_train.py:102: teacher_thinking: bool = field(
opsd_train.py:186: "lmbda": training_args.lmbda,
opsd_train.py:193: "use_tinker_loss": script_args.use_tinker_loss,
opsd_train.py:259: # Load the math dataset with ground truth solutions
opsd_train.py:276: use_thinking_machines_loss=script_args.use_tinker_loss,
opsd_train.py:278: reason_first=script_args.reason_first,
opsd_train.py:284: teacher_thinking=script_args.teacher_thinking,
#586
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"sed -n '1,125p' opsd_train.py; sed -n '620,770p' opsd_trainer.py; sed -n '1,180p' data_collator.py\",workdir:\"/workspace/repo\",yield_time_ms:10000,max_output_tokens:14000}); text(r.output);\n"
}exec result
Script completed
Wall time 0.1 seconds
Output:
import os
import wandb
from datasets import load_dataset
from transformers import AutoTokenizer, GenerationConfig
from trl import (
LogCompletionsCallback,
ModelConfig,
ScriptArguments,
TrlParser,
get_kbit_device_map,
get_peft_config,
get_quantization_config,
)
from trl.experimental.gold import GOLDConfig
from opsd_trainer import OPSDTrainer
from dataclasses import dataclass, field
# Enable logging in a Hugging Face Space
os.environ.setdefault("TRACKIO_SPACE_ID", "trl-trackio")
@dataclass
class CustomScriptArguments(ScriptArguments):
"""Extended script arguments with Thinking Machines loss option."""
use_tinker_loss: bool = field(
default=False,
metadata={
"help": "Use Thinking Machines style on-policy reverse KL loss instead of GKD's full-vocab JSD loss. "
"This is much more memory efficient (O(1) vs O(vocab_size) per token)."
},
)
fixed_teacher: bool = field(
default=False,
metadata={
"help": "Use the initial policy (step 0) as a fixed teacher. Only works with use_peft=True. "
"The teacher will use the base model without LoRA adapters, while the student updates."
},
)
run_config: str = field(
default=None,
metadata={
"help": "Run name for this experiment. Will be used for both the output directory "
"(appended to output_dir) and WandB run name. If not specified, will generate "
"automatic name based on hyperparameters."
},
)
presence_penalty: float = field(
default=0.0,
metadata={
"help": "Float that penalizes new tokens based on whether they appear in the generated text so far. "
"Values > 0 encourage the model to use new tokens, while values < 0 encourage the model to repeat tokens."
},
)
reason_first: bool = field(
default=False,
metadata={
"help": "Let the teacher model first rationalize (generate rationalization explictly) about the given reasoning first then act as teacher."
},
)
top_k_loss: int = field(
default=0,
metadata={
"help": "Restrict the JSD loss to only the top-k tokens of the teacher distribution. Both student and "
"teacher distributions are renormalized over these k tokens before computing JSD. "
"Set to 0 (default) to use the full vocabulary."
},
)
jsd_token_clip: float = field(
default=0.05,
metadata={
"help": "Clip the JSD loss for each token to a maximum value. This can improve stability by preventing "
"extremely high-loss stylistic tokens from dominating the training signal. Set to 0 for no clipping."
},
)
use_ema_teacher: bool = field(
default=False,
metadata={
"help": "Use an exponential moving average (EMA) of student weights as the teacher. "
"The EMA teacher is a smoothly-lagged version of the student, avoiding the teacher "
"collapsing to the current policy (dynamic) or staying frozen (fixed_teacher). "
"Mutually exclusive with fixed_teacher."
},
)
ema_decay: float = field(
default=0.999,
metadata={
"help": "EMA decay factor. Higher values make the teacher change more slowly. "
"Typical range: 0.99–0.9999. Only used when use_ema_teacher=True."
},
)
student_thinking: bool = field(
default=False,
metadata={
"help": "Whether to enable Qwen3 thinking mode for the student during rollout. "
"Default False (matches the main OPSD setup: student rolls out without <think>)."
},
)
teacher_thinking: bool = field(
default=True,
metadata={
"help": "Whether to enable Qwen3 thinking mode for the teacher when scoring student tokens. "
"Default True. Set to False for the matched non-thinking ablation (both nonthink)."
},
)
if __name__ == "__main__":
parser = TrlParser((CustomScriptArguments, GOLDConfig, ModelConfig))
script_args, training_args, model_args = parser.parse_args_and_config()
################
# WandB Run Name & Output Directory
################
# Format learning rate (e.g., 2e-4 -> "2e-4" or 0.0002 -> "2e-4")
lr_str = f"{training_args.learning_rate:.0e}".replace("e-0", "e-")
# Get number of processes from environment (set by accelerate launch)
num_processes = int(os.environ.get("WORLD_SIZE", 1))
# Calculate effective batch size
effective_batch_size = (
for name, param in name_to_param.items():
if name in saved:
param.data.copy_(saved[name])
else:
saved = {}
for name, param in unwrapped.named_parameters():
if not param.requires_grad or name not in self._ema_params:
continue
ema = self._ema_params[name]
if ema.device != param.data.device:
ema = ema.to(param.data.device)
self._ema_params[name] = ema
saved[name] = param.data
param.data = ema
try:
yield
finally:
for name, param in unwrapped.named_parameters():
if name in saved:
param.data = saved[name]
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
"""
Compute the self-distillation loss with memory-efficient log-prob extraction.
Memory optimization: Extract only needed log-probs immediately and free large tensors.
"""
# Get batch-level prompt lengths
student_prompt_len = inputs["student_prompt_length"]
teacher_prompt_len = inputs["teacher_prompt_length"]
sampled_token_ids = inputs["student_input_ids"][:, student_prompt_len:]
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"],
)
# Extract only what we need and convert to log-probs immediately
student_logits = outputs_student.logits[:, student_prompt_len - 1 : -1, :]
if self.use_thinking_machines_loss:
# For reverse KL, we only need log-probs of sampled tokens
student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
student_log_probs_sampled = torch.gather(
student_log_probs, dim=-1, index=sampled_token_ids.unsqueeze(-1)
).squeeze(-1)
del student_logits, student_log_probs # Free immediately!
else:
# For JSD, keep logits (temperature will be applied in generalized_jsd_loss)
student_logits_for_loss = student_logits
del student_logits
# Free the full outputs (but keep reference for return_outputs if needed)
if return_outputs:
# Create a minimal output object to return (just the loss, no logits)
class MinimalOutput:
def __init__(self):
self.loss = None
minimal_output = MinimalOutput()
del outputs_student
empty_cache()
# === TEACHER FORWARD - Extract log-probs immediately ===
# Choose teacher context based on mode:
# use_ema_teacher → swap in EMA weights temporarily
# fixed_teacher → disable LoRA adapters (base model = initial policy)
# default (dynamic)→ no-op, use current student weights
if self.use_ema_teacher:
adapter_context = self._ema_teacher_context(model)
elif self.fixed_teacher and is_peft_model(model):
adapter_context = self.accelerator.unwrap_model(model).disable_adapter()
else:
adapter_context = nullcontext()
with torch.no_grad(), adapter_context:
outputs_teacher = model(
input_ids=inputs["teacher_input_ids"],
attention_mask=inputs["teacher_attention_mask"],
)
teacher_logits = outputs_teacher.logits[:, teacher_prompt_len - 1 : -1, :]
if self.use_thinking_machines_loss:
teacher_log_probs = F.log_softmax(teacher_logits / self.temperature, dim=-1)
teacher_log_probs_sampled = torch.gather(
teacher_log_probs, dim=-1, index=sampled_token_ids.unsqueeze(-1)
).squeeze(-1)
del teacher_logits, teacher_log_probs # Free immediately!
else:
teacher_logits_for_loss = teacher_logits
del teacher_logits
del outputs_teacher
empty_cache()
# === COMPUTE LOSS with only small tensors ===
if self.use_thinking_machines_loss:
# Thinking Machines uses RL-style policy gradient:
# Advantage = log π_teacher(x) - log π_student(x)
# Loss = -E[Advantage * log π_student(x)]
#
# CRITICAL: advantage must be detached to prevent gradients flowing through it.
# We want: ∇θ L = -E[A(x) * ∇θ log π_student(x)]
# NOT: ∇θ L = -E[(T(x) - S(x)) * ∇θ S(x)] where both terms differentiate
advantage = (teacher_log_probs_sampled - student_log_probs_sampled).detach()
# Apply masking before computing loss
if shifted_labels is not None:
mask = shifted_labels != -100
advantage = advantage[mask]
student_log_probs_sampled_masked = student_log_probs_sampled[mask]
else:
student_log_probs_sampled_masked = student_log_probs_sampled
# Policy gradient loss: -advantage * log π_student
# Negative because we minimize loss (gradient descent), but want to maximize reward
loss = -(advantage * student_log_probs_sampled_masked).mean()
del (
student_log_probs_sampled,
teacher_log_probs_sampled,
advantage,
student_log_probs_sampled_masked,
)
else:
# Temperature is applied inside generalized_jsd_loss
loss = self.generalized_jsd_loss(
student_logits=student_logits_for_loss,
teacher_logits=teacher_logits_for_loss,
labels=shifted_labels,
beta=self.beta,
temperature=self.temperature, # Let the function handle temperature
top_k=self.top_k_loss,
token_clip=self.jsd_token_clip,
)
del student_logits_for_loss, teacher_logits_for_loss
empty_cache()
if return_outputs:
minimal_output.loss = loss
return (loss, minimal_output)
else:
return loss
def generate_teacher_reasoning(
import torch
class SelfDistillationDataCollator:
"""
Data collator for self-distillation that creates both student and teacher inputs.
Student: sees only the problem (with chat template)
Teacher: sees problem + solution + transition prompt (with chat template)
To enable batch-level operations (like original GKD), we pad prompts to the same length
within each batch, and track the actual (unpadded) prompt lengths for loss masking.
"""
def __init__(
self,
tokenizer,
max_length=2048,
reason_first=True,
student_thinking=False,
teacher_thinking=True,
):
self.tokenizer = tokenizer
self.max_length = max_length
self.reason_first = reason_first
self.student_thinking = student_thinking
self.teacher_thinking = teacher_thinking
# Prompt for reasoning about the solution before teaching
self.reason_first_prompt = (
"\n\nThe reference reasoning above arrives at the correct answer. "
"Please analyze this solution and explain the key reasoning steps and problem-solving strategies employed. "
"Do NOT use <think> tags. Do NOT derive your own solution. "
"Simply analyze and explain the reference solution provided above.\n"
)
# Prompt for transitioning to teaching mode after reasoning
self.transition_prompt = (
"\n\nAfter reading the reference solution above, make sure you truly understand "
"the reasoning behind each step — do not copy or paraphrase it. Now, using your "
"own words and independent reasoning, derive the same final answer to the problem above. "
"Think step by step, explore different approaches, and don't be afraid to backtrack "
"or reconsider if something doesn't work out:\n"
)
# Set padding side explicitly for consistency
print(f"[DataCollator] Original padding_side: {self.tokenizer.padding_side}")
self.tokenizer.padding_side = "right"
print(f"[DataCollator] Set padding_side to: {self.tokenizer.padding_side}")
print(f"[DataCollator] Reason first mode: {self.reason_first}")
def __call__(self, features):
batch_size = len(features)
# Prepare student and teacher prompts using chat template (matching evaluation)
student_prompts = []
teacher_prompts = []
teacher_reasoning_prompts = [] # NEW: for reason_first mode
for feature in features:
# Extract problem and solution from dataset
# Handle different possible column names
problem = feature["problem"]
solution = feature["solution"]
# 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{{}}."
student_messages = [{"role": "user", "content": student_user_message}]
# Apply chat template for student (matching evaluation)
student_prompt = self.tokenizer.apply_chat_template(
student_messages, tokenize=False, add_generation_prompt=True, enable_thinking=self.student_thinking
)
student_prompts.append(student_prompt)
if self.reason_first:
# Reasoning prompt: ask teacher to analyze the solution
reasoning_user_message = (
f"Problem: {problem}\n\n"
f"Here is a correct reasoning to this problem:"
f"=== Reference Reasoning Start ===\n"
f"{solution}\n"
f"=== Reference Reasoning End ===\n\n"
f"{self.reason_first_prompt}"
)
reasoning_messages = [{"role": "user", "content": reasoning_user_message}]
reasoning_prompt = self.tokenizer.apply_chat_template(
reasoning_messages, tokenize=False, add_generation_prompt=True
)
teacher_reasoning_prompts.append(reasoning_prompt)
# Teacher prompt will be constructed during training after reasoning
# For now, create placeholder (will be replaced in training_step)
teacher_prompts.append("") # Placeholder
else:
# Original teacher prompt (unchanged)
teacher_user_message = (
f"Problem: {problem}\n\n"
f"Here is a reference solution to this problem:\n"
f"=== Reference Solution Begin ===\n{solution}\n=== Reference Solution End ===\n"
f"{self.transition_prompt}\n"
f"Please reason step by step, and put your final answer within \\boxed{{}}."
)
teacher_messages = [{"role": "user", "content": teacher_user_message}]
# Apply chat template for teacher
teacher_prompt = self.tokenizer.apply_chat_template(
teacher_messages, tokenize=False, add_generation_prompt=True, enable_thinking=self.teacher_thinking
)
teacher_prompts.append(teacher_prompt)
# Tokenize WITHOUT padding first to get true lengths
student_encoded_no_pad = self.tokenizer(
student_prompts,
padding=False,
truncation=True,
max_length=self.max_length,
)
student_prompt_lengths = [len(ids) for ids in student_encoded_no_pad["input_ids"]]
# Find max lengths in this batch
max_student_prompt_len = max(student_prompt_lengths)
# Tokenize WITH padding to max length in batch
student_encoded = self.tokenizer(
student_prompts,
padding="max_length",
truncation=True,
max_length=max_student_prompt_len,
return_tensors="pt",
)
result = {
"student_prompts": student_encoded["input_ids"],
"student_prompt_attention_mask": student_encoded["attention_mask"],
"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),
}
if self.reason_first:
# Tokenize reasoning prompts
reasoning_encoded_no_pad = self.tokenizer(
teacher_reasoning_prompts,
padding=False,
truncation=True,
max_length=self.max_length,
)
reasoning_prompt_lengths = [len(ids) for ids in reasoning_encoded_no_pad["input_ids"]]
max_reasoning_prompt_len = max(reasoning_prompt_lengths)
reasoning_encoded = self.tokenizer(
teacher_reasoning_prompts,
padding="max_length",
truncation=True,
max_length=max_reasoning_prompt_len,
return_tensors="pt",
)
# Tokenize transition prompt (this will be appended after reasoning)
# Don't use chat template here - just the raw text
transition_text = f"\n{self.transition_prompt}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
transition_encoded = self.tokenizer(
[transition_text] * batch_size,
padding=False,
truncation=False,
return_tensors="pt",
)
result.update(
{
"teacher_reasoning_prompts": reasoning_encoded["input_ids"],
"teacher_reasoning_attention_mask": reasoning_encoded["attention_mask"],
"teacher_reasoning_prompt_length": max_reasoning_prompt_len,
"teacher_transition_tokens": transition_encoded["input_ids"],
}
)
else:
# Normal mode: tokenize teacher prompts
teacher_encoded_no_pad = self.tokenizer(
#587
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"find /workspace -maxdepth 2 -type f \\\\( -iname '*paper*txt' -o -iname '*extract*txt' \\\\) -print; pdftotext /workspace/paper.pdf - 2>/dev/null | rg -n -C 4 \\\"reason.first|rational|AIME 2024|Table 2|ablation|thinking\\\" | head -300\",workdir:\"/workspace\",yield_time_ms:10000,max_output_tokens:12000}); text(r.output);\n"
}exec result
Script completed
Wall time 0.1 seconds
Output:
#588
narration
I’m locking the validated front-loaded objective rather than spending the last reliable window on an unreplicated high-variance idea. Next I’m verifying its invariants (unit mean per trajectory, padding behavior, finite gradients), syntax, and the exact submission diff.
#589
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"python -m py_compile opsd_train.py opsd_trainer.py data_collator.py\\npython - <<'PY'\\nimport torch\\n\\ndef credit(labels, dtype=torch.float32):\\n mask = labels != -100\\n positions = mask.long().cumsum(dim=-1) - 1\\n lengths = mask.sum(dim=-1, keepdim=True).clamp_min(1)\\n progress = positions.to(dtype) / (lengths - 1).clamp_min(1)\\n weights = (1.5 - progress).clamp_min(0.5) * mask\\n return torch.where(lengths > 1, weights, mask.to(dtype)), mask\\n\\nlabels = torch.tensor([\\n [1, 2, 3, 4, -100, -100],\\n [1, 2, -100, -100, -100, -100],\\n [1, -100, -100, -100, -100, -100],\\n [-100, -100, -100, -100, -100, -100],\\n])\\nw, m = credit(labels)\\nassert torch.allclose(w.sum(-1), m.sum(-1).float())\\nassert torch.allclose(w[0, :4], torch.tensor([1.5, 7/6, 5/6, 0.5]))\\nx = torch.randn(4, 6, requires_grad=True)\\nloss = (x * w).sum() / m.sum().clamp_min(1)\\nloss.backward()\\nassert torch.isfinite(loss) and torch.isfinite(x.grad).all()\\nprint('credit invariant and gradient checks passed')\\nPY\\nawk 'NF && $1 !~ /^#/' /workspace/submission/recipe.env 2>/dev/null || true\\ngit diff --check -- opsd_trainer.py\\ngit diff --stat -- opsd_trainer.py\\ngit status --short\",workdir:\"/workspace/repo\",yield_time_ms:10000,max_output_tokens:5000}); text(r.output);\n"
}exec result
Script completed
Wall time 2.0 seconds
Output:
credit invariant and gradient checks passed
opsd_trainer.py | 25 ++++++++++++++++++++-----
1 file changed, 20 insertions(+), 5 deletions(-)
M README.md
M opsd_trainer.py
#590
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"git diff -- opsd_trainer.py; nl -ba opsd_trainer.py | sed -n '450,505p'; ls -la /workspace/submission /workspace/submission/recipe.env 2>/dev/null || true; cat /workspace/submission/recipe.env 2>/dev/null || true\",workdir:\"/workspace/repo\",yield_time_ms:10000,max_output_tokens:8000}); text(r.output);\n"
}exec result
Script completed
Wall time 0.1 seconds
Output:
diff --git a/opsd_trainer.py b/opsd_trainer.py
index c151bd7..7fd004e 100644
--- a/opsd_trainer.py
+++ b/opsd_trainer.py
@@ -463,20 +463,35 @@ class OPSDTrainer(SFTTrainer):
if token_clip is not None:
jsd = jsd.clamp(max=token_clip)
- # Masking
+ # Masking and trajectory credit assignment. Early reasoning tokens are
+ # more likely to encode a consequential branch choice; late tokens are
+ # increasingly determined by the already-long prefix. Give early
+ # positions moderately more credit without changing the total loss
+ # scale. The linear 1.5 -> 0.5 schedule has mean one on every non-empty
+ # trajectory, so learning-rate and clipping behavior remain comparable
+ # to the original uniformly averaged OPSD objective.
if labels is not None:
mask = labels != -100
- jsd = jsd[mask]
+ token_divergence = jsd.sum(dim=-1)
+ positions = mask.long().cumsum(dim=-1) - 1
+ lengths = mask.sum(dim=-1, keepdim=True).clamp_min(1)
+ progress = positions.to(token_divergence.dtype) / (lengths - 1).clamp_min(1)
+ credit = (1.5 - progress).clamp_min(0.5) * mask
+ credit = torch.where(lengths > 1, credit, mask.to(token_divergence.dtype))
# Apply reduction
if reduction == "batchmean":
- return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / jsd.size(0)
+ if labels is not None:
+ return (token_divergence * credit).sum() / mask.sum().clamp_min(1)
+ return jsd.sum() / jsd.size(0)
elif reduction == "sum":
- return jsd.sum()
+ return (token_divergence * credit).sum() if labels is not None else jsd.sum()
elif reduction == "mean":
+ if labels is not None:
+ return (token_divergence * credit).sum() / credit.sum().clamp_min(1)
return jsd.mean()
else:
- return jsd
+ return token_divergence * credit if labels is not None else jsd
def _update_ema(self):
"""Update EMA parameters after an optimizer step.
450 torch.stack([student_log_probs + torch.log1p(-beta), teacher_log_probs + torch.log(beta)]),
451 dim=0,
452 )
453
454 # Compute KL divergences using F.kl_div
455 # PyTorch differs from the standard mathematical definition, so the order of the probability distributions is swapped compared to that defined in the paper.
456 kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True)
457 kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True)
458
459 # Compute the Generalized Jensen-Shannon Divergence
460 jsd = beta * kl_teacher + (1 - beta) * kl_student
461
462 # Per-token clipping: cap each token's divergence value
463 if token_clip is not None:
464 jsd = jsd.clamp(max=token_clip)
465
466 # Masking and trajectory credit assignment. Early reasoning tokens are
467 # more likely to encode a consequential branch choice; late tokens are
468 # increasingly determined by the already-long prefix. Give early
469 # positions moderately more credit without changing the total loss
470 # scale. The linear 1.5 -> 0.5 schedule has mean one on every non-empty
471 # trajectory, so learning-rate and clipping behavior remain comparable
472 # to the original uniformly averaged OPSD objective.
473 if labels is not None:
474 mask = labels != -100
475 token_divergence = jsd.sum(dim=-1)
476 positions = mask.long().cumsum(dim=-1) - 1
477 lengths = mask.sum(dim=-1, keepdim=True).clamp_min(1)
478 progress = positions.to(token_divergence.dtype) / (lengths - 1).clamp_min(1)
479 credit = (1.5 - progress).clamp_min(0.5) * mask
480 credit = torch.where(lengths > 1, credit, mask.to(token_divergence.dtype))
481
482 # Apply reduction
483 if reduction == "batchmean":
484 if labels is not None:
485 return (token_divergence * credit).sum() / mask.sum().clamp_min(1)
486 return jsd.sum() / jsd.size(0)
487 elif reduction == "sum":
488 return (token_divergence * credit).sum() if labels is not None else jsd.sum()
489 elif reduction == "mean":
490 if labels is not None:
491 return (token_divergence * credit).sum() / credit.sum().clamp_min(1)
492 return jsd.mean()
493 else:
494 return token_divergence * credit if labels is not None else jsd
495
496 def _update_ema(self):
497 """Update EMA parameters after an optimizer step.
498
499 On the very first call this lazily initializes the EMA state as an exact copy of the
500 current (trainable) model parameters, then returns without applying a decay step.
501 Subsequent calls apply: ema = decay * ema + (1 - decay) * student.
502
503 Only trainable parameters are tracked (i.e. LoRA adapter weights for PEFT models,
504 or all parameters for full fine-tuning).
505
-rw-r--r-- 1 root root 653 Aug 3 08:45 /workspace/submission/recipe.env
/workspace/submission:
total 1
drwxr-xr-x 1 root root 60 Aug 3 04:41 .
drwxr-xr-x 1 root root 180 Aug 3 05:30 ..
-rw-r--r-- 1 root root 653 Aug 3 08:45 recipe.env
# OPSD method recipe — whitelisted hyper-parameter overrides (KEY=VALUE per line).
# Empty/all-commented == the OPSD baseline recipe on top of your /workspace/repo
# code changes. The scorer re-trains with THIS file + your repo for 100 steps.
# Budget knobs (max_steps, epochs, batch, gradient_accumulation, num_processes,
# model) are FIXED by the scorer and ignored here.
#
# learning_rate=5e-6
# lora_r=64
# lora_alpha=128
# beta=0
# jsd_token_clip=0.05
# temperature=1.1
# lmbda=1
# max_completion_length=1024 # <= 4096
# lr_scheduler_type=constant
# fixed_teacher=true
# use_ema_teacher=false
# reason_first=false
# teacher_thinking=false
#591
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"sed -n '1,240p' /workspace/train_capped.sh | rg -n -C 2 \\\"temperature|teacher_thinking|fixed_teacher|learning_rate|max_completion|lr_scheduler|jsd_token\\\"; rg -n \\\"temperature:|lmbda:\\\" /workspace/repo/opsd_train.py $(python - <<'PY'\\nimport inspect\\nfrom trl.experimental.gold import GOLDConfig\\nprint(inspect.getfile(GOLDConfig))\\nPY\\n) 2>/dev/null | head -50\",workdir:\"/workspace\",yield_time_ms:10000,max_output_tokens:5000}); text(r.output);\n"
}exec result
Script completed
Wall time 10.2 seconds
Output:
35-# ---- baseline method defaults (empty recipe == the OPSD baseline recipe) ----
36-declare -A CFG=(
37: [learning_rate]=5e-6 [max_grad_norm]=0.1 [weight_decay]=0
38: [lr_scheduler_type]=constant [warmup_ratio]=0
39- [lora_r]=64 [lora_alpha]=128 [lora_dropout]=0
40: [beta]=0 [jsd_token_clip]=0.05 [top_k_loss]=0
41: [temperature]=1.1 [top_p]=0.95 [top_k]=20
42: [lmbda]=1 [max_completion_length]=1024 [ema_decay]=0.999
43: [fixed_teacher]=true [use_ema_teacher]=false [use_tinker_loss]=false
44: [reason_first]=false [teacher_thinking]=false [student_thinking]=false
45-)
46:BOOLKEYS="fixed_teacher use_ema_teacher use_tinker_loss reason_first teacher_thinking student_thinking"
47-
48-# ---- overlay whitelisted knobs from recipe.env (budget/unknown keys ignored) ----
--
55-fi
56-
57:# ---- clamp max_completion_length so the fixed budget stays honest (<=4096) ----
58:mcl="${CFG[max_completion_length]}"; case "$mcl" in ''|*[!0-9]*) mcl=1024;; esac
59:if [ "$mcl" -gt 4096 ]; then echo "[train_capped] clamping max_completion_length $mcl -> 4096"; mcl=4096; fi
60:CFG[max_completion_length]="$mcl"
61-
62-# ---- assemble method args (value flags, then boolean store_true flags) ----
63-ARGS=()
64:for k in learning_rate max_grad_norm weight_decay lr_scheduler_type warmup_ratio \
65: lora_r lora_alpha lora_dropout beta jsd_token_clip top_k_loss \
66: temperature top_p top_k lmbda max_completion_length ema_decay; do
67- ARGS+=( "--$k" "${CFG[$k]}" )
68-done
<stdin>:2: TRLExperimentalWarning: You are importing from 'trl.experimental'. APIs here are unstable and may change or be removed without notice. Silence this warning by setting environment variable TRL_EXPERIMENTAL_SILENCE=1.
#592
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const patch = \"*** Begin Patch\\n*** Update File: /workspace/repo/opsd_trainer.py\\n@@\\n elif reduction == \\\"mean\\\":\\n if labels is not None:\\n- return (token_divergence * credit).sum() / credit.sum().clamp_min(1)\\n+ element_count = credit.sum() * jsd.size(-1)\\n+ return (token_divergence * credit).sum() / element_count.clamp_min(1)\\n return jsd.mean()\\n else:\\n- return token_divergence * credit if labels is not None else jsd\\n+ return jsd[mask] * credit[mask].unsqueeze(-1) if labels is not None else jsd\\n*** End Patch\";\ntext(await tools.apply_patch(patch));\n"
}exec result
Script completed
Wall time 0.0 seconds
Output:
{}#593
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"for f in /workspace/eval_results/dev_eval_*.json; do jq -r '[input_filename, (.accuracy // .avg_at_12 // .avg_accuracy // empty), (.correct // empty), (.total // empty)] | @tsv' \\\"$f\\\" 2>/dev/null; done | tail -20\\npython -m py_compile opsd_trainer.py\\npython - <<'PY'\\nimport ast\\nfrom pathlib import Path\\np = Path('/workspace/repo/opsd_trainer.py')\\nast.parse(p.read_text())\\nprint('AST and bytecode checks passed')\\nPY\\ngit diff --check -- opsd_trainer.py\\ncat /workspace/.timer/remaining_secs\",workdir:\"/workspace/repo\",yield_time_ms:10000,max_output_tokens:6000}); text(r.output);\n"
}exec result
Script completed
Wall time 0.4 seconds
Output:
/workspace/eval_results/dev_eval_1785734021.json
/workspace/eval_results/dev_eval_1785736095.json
/workspace/eval_results/dev_eval_1785738016.json
/workspace/eval_results/dev_eval_1785739883.json
/workspace/eval_results/dev_eval_1785741829.json
/workspace/eval_results/dev_eval_1785743764.json
/workspace/eval_results/dev_eval_1785745652.json
/workspace/eval_results/dev_eval_1785747609.json
AST and bytecode checks passed
4713
#594
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"jq 'keys' /workspace/eval_results/dev_eval_1785747609.json; jq '.' /workspace/eval_results/dev_eval_1785747609.json | head -80\",workdir:\"/workspace\",yield_time_ms:10000,max_output_tokens:5000}); text(r.output);\n"
}exec result
Script completed
Wall time 0.1 seconds
Output:
[
"average_at_n_pct",
"benchmark",
"checkpoint_kind",
"checkpoint_path",
"data_parallel_size",
"format_rate_pct",
"frozen_settings",
"metric",
"num_problems",
"pass_at_n_pct",
"results",
"seed",
"total_correct",
"total_solutions"
]
{
"benchmark": "aime24",
"metric": "avg@12",
"average_at_n_pct": 56.111111111111114,
"pass_at_n_pct": 83.33333333333333,
"format_rate_pct": 97.5,
"num_problems": 30,
"total_solutions": 360,
"total_correct": 202,
"checkpoint_path": "/workspace/runs/earlycredit_trajbalanced/capped_seed42/checkpoint-100",
"checkpoint_kind": "lora",
"seed": 20260610,
"data_parallel_size": 4,
"frozen_settings": {
"temperature": 1,
"top_p": 0.95,
"top_k": -1,
"min_p": 0,
"presence_penalty": 0,
"max_new_tokens": 38912,
"enable_thinking": true,
"val_n": 12,
"max_model_len": 40960
},
"results": [
{
"problem_id": 60,
"ground_truth": "204",
"num_correct": 12,
"n": 12,
"majority_vote_correct": true,
"generations": [
{
"predicted_answer": "204",
"correct": true,
"formatted": true
},
{
"predicted_answer": "204",
"correct": true,
"formatted": true
},
{
"predicted_answer": "204",
"correct": true,
"formatted": true
},
{
"predicted_answer": "204",
"correct": true,
"formatted": true
},
{
"predicted_answer": "204",
"correct": true,
"formatted": true
},
{
"predicted_answer": "204",
"correct": true,
"formatted": true
},
{
"predicted_answer": "204",
"correct": true,
"formatted": true
},
{
"predicted_answer": "204",
"correct": true,
"formatted": true
},
{
"predicted_answer": "204",
"correct": true,
"formatted": true
},
{
"predicted_answer": "204",
"correct": true,
#595
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"for f in /workspace/eval_results/dev_eval_*.json; do jq -r '[input_filename, .checkpoint_path, .average_at_n_pct, .total_correct] | @tsv' \\\"$f\\\"; done | tail -20\",workdir:\"/workspace\",yield_time_ms:10000,max_output_tokens:5000}); text(r.output);\n"
}exec result
Script completed
Wall time 0.3 seconds
Output:
/workspace/eval_results/dev_eval_1785734021.json /workspace/runs/contest2048/capped_seed42/checkpoint-100 51.94444444444444 187
/workspace/eval_results/dev_eval_1785736095.json /workspace/runs/earlycredit/capped_seed42/checkpoint-100 56.111111111111114 202
/workspace/eval_results/dev_eval_1785738016.json /workspace/runs/earlycredit_linear/capped_seed42/checkpoint-100 54.72222222222222 197
/workspace/eval_results/dev_eval_1785739883.json /workspace/runs/earlycredit_repeat2/capped_seed42/checkpoint-100 56.111111111111114 202
/workspace/eval_results/dev_eval_1785741829.json /workspace/runs/earlycredit_strong/capped_seed42/checkpoint-100 56.94444444444444 205
/workspace/eval_results/dev_eval_1785743764.json /workspace/runs/earlycredit_strong_repeat2/capped_seed42/checkpoint-100 55.27777777777778 199
/workspace/eval_results/dev_eval_1785745652.json /workspace/runs/earlycredit_clip025/capped_seed42/checkpoint-100 54.44444444444444 196
/workspace/eval_results/dev_eval_1785747609.json /workspace/runs/earlycredit_trajbalanced/capped_seed42/checkpoint-100 56.111111111111114 202
#596
narration
One evidence-backed optimization remains worthwhile: the linear-decay run reduced the average update size and lost 1.39 points, while constant 5e-6 was stable. I’m testing a modest 7.5e-6 constant rate with the validated objective; the fallback remains fully reproducible if it does not clear 56.11 convincingly.
#597
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const patch = \"*** Begin Patch\\n*** Update File: /workspace/submission/recipe.env\\n@@\\n-# learning_rate=5e-6\\n+learning_rate=7.5e-6\\n*** End Patch\";\ntext(await tools.apply_patch(patch));\n"
}exec result
Script completed
Wall time 0.0 seconds
Output:
{}#598
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.exec_command({cmd:\"SEED=42 OUT_DIR=/workspace/runs/earlycredit_lr75 PORT=13841 bash /workspace/train_capped.sh\",workdir:\"/workspace\",yield_time_ms:1000,max_output_tokens:500}); text(JSON.stringify(r));\n"
}exec result
Script completed
Wall time 1.2 seconds
Output:
{"chunk_id":"b29ac3","wall_time_seconds":1.001478324,"session_id":28348,"original_token_count":0,"output":""}#599
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolexec
args
{
"input": "const r = await tools.write_stdin({session_id:28348,chars:\"\",yield_time_ms:60000,max_output_tokens:250}); text(r.output.slice(-450));\n"
}exec result
Script running with cell ID 298
Wall time 11.0 seconds
Output:
#600
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolwait
args
{
"cell_id": "298",
"yield_time_ms": 60000,
"max_tokens": 250
}wait result
Script completed
Wall time 45.3 seconds
Output:
eration done - elapsed time: 2.95s, prompts: 4, total tokens: 2219, avg length: 554.8, speed: 751.3 tok/s
vLLM generation done - elapsed time: 3.19s, prompts: 4, total tokens: 2305, avg length: 576.2, speed: 723.0 tok/s
vLLM generation done - elapsed time: 3.31s, prompts: 4, total tokens: 3634, avg length: 908.5, speed: 1098.8 tok/s
vLLM generation done - elapsed time: 3.34s, prompts: 4, total tokens: 3871, avg length: 967.8, speed: 1158.1 tok/s