Upload folder using huggingface_hub

tools/backup/train_poe_distill_lora_linear/cos_beta.py

@@ -0,0 +1,783 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""LoRA training with an online product-of-experts distillation target.
+
+This script trains on fixed pi_ref rollouts, but computes full-vocabulary
+teacher/ref distributions online:
+
+    pi_star(. | s) proportional to pi_T(. | s)^beta * pi_ref(. | s)^(1-beta)
+    beta = alpha / (alpha + 1)
+
+The trainable model is pi_ref plus LoRA adapters.  The frozen pi_ref
+distribution is obtained by disabling the adapter on the same model, avoiding a
+second copy of the 4B reference model.
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import os
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+from peft import LoraConfig, TaskType, get_peft_model
+from torch.nn.utils.rnn import pad_sequence
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    Trainer,
+    TrainingArguments,
+    set_seed,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Product-of-experts LoRA distillation on fixed rollouts.")
+
+    parser.add_argument("--student-model", default=os.environ.get("SFT_CHECKPOINT"), required=False)
+    parser.add_argument("--teacher-model", default=os.environ.get("TEACHER_MODEL", "Qwen/Qwen3-8B"))
+    parser.add_argument("--train-data", default="data/rollouts/dapo-math-17k-qwen3-4b-sft-rollouts.parquet")
+    parser.add_argument("--output-dir", default="checkpoints/qwen3-4b-poe-distill-lora")
+
+    parser.add_argument("--alpha", type=float, default=1.0)
+    parser.add_argument(
+        "--beta-start",
+        type=float,
+        default=None,
+        help="Initial beta. If unset, uses alpha / (alpha + 1) as a fixed beta.",
+    )
+    parser.add_argument(
+        "--beta-end",
+        type=float,
+        default=None,
+        help="Final beta. If unset, uses alpha / (alpha + 1) as a fixed beta.",
+    )
+    parser.add_argument(
+        "--beta-schedule-steps",
+        type=int,
+        default=None,
+        help="Number of optimizer steps used to ramp beta from beta-start to beta-end.",
+    )
+    parser.add_argument("--beta-schedule", choices=["linear", "cosine"], default="linear")
+    parser.add_argument(
+        "--beta-hold-steps",
+        type=int,
+        default=0,
+        help=(
+            "Keep beta fixed at beta-start for this many optimizer steps before "
+            "transitioning to beta-end. This enables schedules such as: hold "
+            "beta=1.0 for 100 steps, then decay to 0.5 over 10 steps."
+        ),
+    )
+    parser.add_argument(
+        "--beta-transition-steps",
+        type=int,
+        default=None,
+        help=(
+            "Number of optimizer steps used to transition beta from beta-start to "
+            "beta-end after beta-hold-steps. If unset, falls back to "
+            "beta-schedule-steps / max_steps for backward compatibility."
+        ),
+    )
+    parser.add_argument(
+        "--hold-transition-schedule",
+        choices=["linear", "cosine"],
+        default="linear",
+        help="Schedule shape used by the hold-then-transition beta/LR schedules.",
+    )
+    parser.add_argument(
+        "--lr-start",
+        type=float,
+        default=None,
+        help="Initial LR for custom hold-then-transition scheduling. If unset, uses --learning-rate.",
+    )
+    parser.add_argument(
+        "--lr-end",
+        type=float,
+        default=None,
+        help="Final LR after the custom transition. If unset, custom LR scheduling is disabled.",
+    )
+    parser.add_argument(
+        "--lr-hold-steps",
+        type=int,
+        default=None,
+        help="Keep LR fixed at lr-start for this many optimizer steps. If unset, uses beta-hold-steps.",
+    )
+    parser.add_argument(
+        "--lr-transition-steps",
+        type=int,
+        default=None,
+        help=(
+            "Number of optimizer steps used to transition LR from lr-start to lr-end. "
+            "If unset, uses beta-transition-steps."
+        ),
+    )
+    parser.add_argument(
+        "--loss-type",
+        choices=["full_vocab", "sampled_token"],
+        default="full_vocab",
+        help=(
+            "full_vocab matches the normalized PoE distribution over the whole vocab. "
+            "sampled_token uses an OPD-style sampled-token surrogate with a PoE advantage."
+        ),
+    )
+    parser.add_argument(
+        "--advantage-normalization",
+        choices=["none", "batch", "sequence"],
+        default="batch",
+        help="Only used by --loss-type sampled_token.",
+    )
+    parser.add_argument(
+        "--advantage-clip",
+        type=float,
+        default=None,
+        help="Symmetric clamp for sampled-token advantages. Example: 5.0.",
+    )
+    parser.add_argument(
+        "--use-ppo-clip",
+        action="store_true",
+        default=False,
+        help=(
+            "Only used by --loss-type sampled_token. Use PPO-style ratio clipping "
+            "with the frozen reference log-prob as the old rollout log-prob."
+        ),
+    )
+    parser.add_argument(
+        "--ppo-clip-low",
+        type=float,
+        default=0.2,
+        help="Only used when --use-ppo-clip is set. Lower PPO clip epsilon.",
+    )
+    parser.add_argument(
+        "--ppo-clip-high",
+        type=float,
+        default=0.2,
+        help="Only used when --use-ppo-clip is set. Upper PPO clip epsilon.",
+    )
+    parser.add_argument(
+        "--sampled-loss-reduction",
+        choices=["per_sample", "per_token"],
+        default="per_sample",
+        help=(
+            "Only used by --loss-type sampled_token. per_sample averages each response "
+            "first, then averages across batch; per_token averages over all response tokens."
+        ),
+    )
+    parser.add_argument(
+        "--positive-advantages-only",
+        action="store_true",
+        default=False,
+        help="Only reinforce sampled tokens with positive PoE advantages.",
+    )
+    parser.add_argument("--max-length", type=int, default=4096)
+    parser.add_argument("--distill-chunk-size", type=int, default=128)
+    parser.add_argument("--max-train-samples", type=int, default=None)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--num-train-epochs", type=float, default=1.0)
+    parser.add_argument("--max-steps", type=int, default=-1)
+    parser.add_argument("--per-device-train-batch-size", type=int, default=1)
+    parser.add_argument("--gradient-accumulation-steps", type=int, default=16)
+    parser.add_argument("--learning-rate", type=float, default=2e-5)
+    parser.add_argument("--weight-decay", type=float, default=0.0)
+    parser.add_argument("--adam-beta1", type=float, default=0.9)
+    parser.add_argument("--adam-beta2", type=float, default=0.999)
+    parser.add_argument("--adam-epsilon", type=float, default=1e-8)
+    parser.add_argument("--warmup-ratio", type=float, default=0.03)
+    parser.add_argument("--lr-scheduler-type", default="cosine")
+    parser.add_argument("--logging-steps", type=int, default=1)
+    parser.add_argument("--save-steps", type=int, default=100)
+    parser.add_argument("--save-total-limit", type=int, default=0)
+    parser.add_argument("--bf16", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--fp16", action="store_true", default=False)
+    parser.add_argument("--gradient-checkpointing", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--report-to", default="none")
+
+    parser.add_argument("--lora-r", type=int, default=64)
+    parser.add_argument("--lora-alpha", type=int, default=128)
+    parser.add_argument("--lora-dropout", type=float, default=0.05)
+    parser.add_argument(
+        "--freeze-lora-b-after-step",
+        type=int,
+        default=None,
+        help="Freeze all LoRA B matrices once global_step reaches this value. Example: 20.",
+    )
+    parser.add_argument(
+        "--lora-target-modules",
+        default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
+        help="Comma-separated LoRA target modules.",
+    )
+
+    parser.add_argument("--trust-remote-code", action="store_true", default=True)
+    parser.add_argument(
+        "--attn-implementation",
+        default=None,
+        choices=[None, "eager", "sdpa", "flash_attention_2"],
+        help="Forwarded to from_pretrained when set.",
+    )
+
+    args = parser.parse_args()
+    if not args.student_model:
+        raise ValueError("Pass --student-model or set SFT_CHECKPOINT to the Qwen3-4B SFT checkpoint.")
+    if args.alpha <= 0:
+        raise ValueError("--alpha must be positive.")
+    fixed_beta = args.alpha / (args.alpha + 1.0)
+    if args.beta_start is None:
+        args.beta_start = fixed_beta
+    if args.beta_end is None:
+        args.beta_end = fixed_beta
+    if not 0.0 <= args.beta_start <= 1.0:
+        raise ValueError("--beta-start must be in [0, 1].")
+    if not 0.0 <= args.beta_end <= 1.0:
+        raise ValueError("--beta-end must be in [0, 1].")
+    if args.beta_schedule_steps is not None and args.beta_schedule_steps <= 0:
+        raise ValueError("--beta-schedule-steps must be positive when set.")
+    if args.beta_hold_steps < 0:
+        raise ValueError("--beta-hold-steps must be non-negative.")
+    if args.beta_transition_steps is not None and args.beta_transition_steps <= 0:
+        raise ValueError("--beta-transition-steps must be positive when set.")
+    if args.lr_start is None:
+        args.lr_start = args.learning_rate
+    if args.lr_hold_steps is None:
+        args.lr_hold_steps = args.beta_hold_steps
+    if args.lr_transition_steps is None:
+        args.lr_transition_steps = args.beta_transition_steps
+    if args.lr_hold_steps is not None and args.lr_hold_steps < 0:
+        raise ValueError("--lr-hold-steps must be non-negative when set.")
+    if args.lr_end is not None:
+        if args.lr_start <= 0.0 or args.lr_end <= 0.0:
+            raise ValueError("--lr-start and --lr-end must be positive when using custom LR scheduling.")
+        if args.lr_transition_steps is None or args.lr_transition_steps <= 0:
+            raise ValueError("--lr-transition-steps must be positive when using --lr-end.")
+    if args.advantage_clip is not None and args.advantage_clip <= 0:
+        raise ValueError("--advantage-clip must be positive when set.")
+    if args.ppo_clip_low < 0 or args.ppo_clip_high < 0:
+        raise ValueError("--ppo-clip-low and --ppo-clip-high must be non-negative.")
+    if args.freeze_lora_b_after_step is not None and args.freeze_lora_b_after_step < 0:
+        raise ValueError("--freeze-lora-b-after-step must be non-negative when set.")
+    if args.fp16 and args.bf16:
+        args.bf16 = False
+    return args
+
+
+def first_assistant_index(messages: list[dict[str, str]]) -> int:
+    for idx, message in enumerate(messages):
+        if message.get("role") == "assistant":
+            return idx
+    raise ValueError("Rollout row has no assistant message.")
+
+
+def tokenize_rollout(example: dict[str, Any], tokenizer: AutoTokenizer, max_length: int) -> dict[str, Any]:
+    messages = example["messages"]
+    assistant_idx = first_assistant_index(messages)
+    prompt_messages = messages[:assistant_idx]
+    full_messages = messages[: assistant_idx + 1]
+
+    prompt_text = tokenizer.apply_chat_template(
+        prompt_messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=True,
+    )
+    full_text = tokenizer.apply_chat_template(
+        full_messages,
+        tokenize=False,
+        add_generation_prompt=False,
+        enable_thinking=True,
+    )
+
+    prompt_ids = tokenizer.encode(prompt_text, add_special_tokens=False)
+    input_ids = tokenizer.encode(full_text, add_special_tokens=False)
+
+    if len(input_ids) > max_length:
+        input_ids = input_ids[:max_length]
+
+    # Mask is aligned to labels=input_ids[1:].  A label predicts token position
+    # j=i+1, so it belongs to the response when j >= len(prompt_ids).
+    label_len = max(len(input_ids) - 1, 0)
+    loss_mask = [1 if i + 1 >= len(prompt_ids) else 0 for i in range(label_len)]
+
+    if sum(loss_mask) == 0:
+        # Drop examples where truncation removed the assistant response.
+        return {"input_ids": [], "loss_mask": []}
+
+    return {"input_ids": input_ids, "loss_mask": loss_mask}
+
+
+@dataclass
+class DistillCollator:
+    pad_token_id: int
+
+    def __call__(self, features: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
+        input_ids = [torch.tensor(f["input_ids"], dtype=torch.long) for f in features]
+        loss_masks = [torch.tensor(f["loss_mask"], dtype=torch.float32) for f in features]
+        lengths = torch.tensor([x.size(0) for x in input_ids], dtype=torch.long)
+
+        padded_input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.pad_token_id)
+        # loss_mask is one shorter than input_ids because it aligns to shifted labels.
+        padded_loss_masks = pad_sequence(loss_masks, batch_first=True, padding_value=0.0)
+        positions = torch.arange(padded_input_ids.size(1)).unsqueeze(0)
+        attention_mask = (positions < lengths.unsqueeze(1)).long()
+        return {
+            "input_ids": padded_input_ids,
+            "attention_mask": attention_mask,
+            "loss_mask": padded_loss_masks,
+        }
+
+
+def hold_then_transition_value(
+    *,
+    step: int,
+    start: float,
+    end: float,
+    hold_steps: int,
+    transition_steps: int | None,
+    schedule: str,
+) -> float:
+    """Return start during hold, then interpolate start -> end.
+
+    Step is an optimizer global_step, not a micro-batch step.  With gradient
+    accumulation, global_step advances only after one optimizer update.
+    """
+    if step < hold_steps:
+        return start
+    if transition_steps is None or transition_steps <= 0:
+        return end
+
+    local_step = step - hold_steps
+    progress = min(max(local_step / transition_steps, 0.0), 1.0)
+    if schedule == "cosine":
+        progress = 0.5 - 0.5 * torch.cos(torch.tensor(progress * torch.pi)).item()
+    elif schedule != "linear":
+        raise ValueError(f"Unknown schedule: {schedule}")
+    return start + (end - start) * progress
+
+
+class PoEDistillTrainer(Trainer):
+    def __init__(
+        self,
+        *args: Any,
+        teacher_model: torch.nn.Module,
+        beta_start: float,
+        beta_end: float,
+        beta_schedule_steps: int | None,
+        beta_schedule: str,
+        beta_hold_steps: int,
+        beta_transition_steps: int | None,
+        hold_transition_schedule: str,
+        loss_type: str,
+        advantage_normalization: str,
+        advantage_clip: float | None,
+        positive_advantages_only: bool,
+        use_ppo_clip: bool,
+        ppo_clip_low: float,
+        ppo_clip_high: float,
+        sampled_loss_reduction: str,
+        distill_chunk_size: int,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.teacher_model = teacher_model
+        self.teacher_model.to(self.args.device)
+        self.teacher_model.eval()
+        self.beta_start = beta_start
+        self.beta_end = beta_end
+        self.beta_schedule_steps = beta_schedule_steps
+        self.beta_schedule = beta_schedule
+        self.beta_hold_steps = beta_hold_steps
+        self.beta_transition_steps = beta_transition_steps
+        self.hold_transition_schedule = hold_transition_schedule
+        self.loss_type = loss_type
+        self.advantage_normalization = advantage_normalization
+        self.advantage_clip = advantage_clip
+        self.positive_advantages_only = positive_advantages_only
+        self.use_ppo_clip = use_ppo_clip
+        self.ppo_clip_low = ppo_clip_low
+        self.ppo_clip_high = ppo_clip_high
+        self.sampled_loss_reduction = sampled_loss_reduction
+        self.distill_chunk_size = distill_chunk_size
+
+    def current_beta(self) -> float:
+        # New mode: hold beta_start for beta_hold_steps, then transition to beta_end.
+        if self.beta_hold_steps > 0 or self.beta_transition_steps is not None:
+            return hold_then_transition_value(
+                step=self.state.global_step,
+                start=self.beta_start,
+                end=self.beta_end,
+                hold_steps=self.beta_hold_steps,
+                transition_steps=self.beta_transition_steps,
+                schedule=self.hold_transition_schedule,
+            )
+
+        # Backward-compatible old behavior: directly schedule beta_start -> beta_end.
+        schedule_steps = self.beta_schedule_steps
+        if schedule_steps is None:
+            schedule_steps = self.state.max_steps if self.state.max_steps > 0 else None
+        if schedule_steps is None or schedule_steps == 0:
+            return self.beta_end
+
+        progress = min(max(self.state.global_step / schedule_steps, 0.0), 1.0)
+        if self.beta_schedule == "cosine":
+            progress = 0.5 - 0.5 * torch.cos(torch.tensor(progress * torch.pi)).item()
+        return self.beta_start + (self.beta_end - self.beta_start) * progress
+
+    @staticmethod
+    def gather_token_logprobs(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        logits = logits.float()
+        token_logits = logits.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)
+        return token_logits - logits.logsumexp(dim=-1)
+
+    def normalize_advantages(self, advantages: torch.Tensor, loss_mask: torch.Tensor) -> torch.Tensor:
+        if self.advantage_normalization == "none":
+            return advantages
+
+        if self.advantage_normalization == "batch":
+            denom = loss_mask.sum().clamp_min(1.0)
+            mean = (advantages * loss_mask).sum() / denom
+            var = (((advantages - mean) * loss_mask) ** 2).sum() / denom
+            return (advantages - mean) / torch.sqrt(var + 1e-6)
+
+        denom = loss_mask.sum(dim=1, keepdim=True).clamp_min(1.0)
+        mean = (advantages * loss_mask).sum(dim=1, keepdim=True) / denom
+        var = (((advantages - mean) * loss_mask) ** 2).sum(dim=1, keepdim=True) / denom
+        return (advantages - mean) / torch.sqrt(var + 1e-6)
+
+    def compute_loss(
+        self,
+        model: torch.nn.Module,
+        inputs: dict[str, torch.Tensor],
+        return_outputs: bool = False,
+        **_: Any,
+    ):
+        loss_mask = inputs.pop("loss_mask")
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs["attention_mask"]
+        labels = input_ids[:, 1:]
+
+        with torch.no_grad():
+            teacher_logits = self.teacher_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                use_cache=False,
+            ).logits[:, :-1, :].detach()
+
+            adapter_owner = model.module if hasattr(model, "module") else model
+            disable_adapter = getattr(adapter_owner, "disable_adapter", None)
+            ref_context = disable_adapter() if disable_adapter is not None else contextlib.nullcontext()
+            with ref_context:
+                ref_logits = adapter_owner(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    use_cache=False,
+                ).logits[:, :-1, :].detach()
+
+        student_outputs = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
+        student_logits = student_outputs.logits[:, :-1, :]
+        if teacher_logits.size(-1) != student_logits.size(-1) or ref_logits.size(-1) != student_logits.size(-1):
+            raise ValueError(
+                "Teacher, reference, and student vocab sizes must match for full-vocab PoE distillation. "
+                f"Got teacher={teacher_logits.size(-1)}, ref={ref_logits.size(-1)}, "
+                f"student={student_logits.size(-1)}."
+            )
+
+        total_loss = student_logits.new_zeros(())
+        total_tokens = loss_mask.sum().clamp_min(1.0)
+        beta = self.current_beta()
+
+        if self.loss_type == "sampled_token":
+            teacher_logp = self.gather_token_logprobs(teacher_logits, labels)
+            ref_logp = self.gather_token_logprobs(ref_logits, labels)
+            student_logp = self.gather_token_logprobs(student_logits, labels)
+
+            poe_score = beta * teacher_logp + (1.0 - beta) * ref_logp
+            advantages = poe_score - student_logp.detach()
+            advantages = self.normalize_advantages(advantages, loss_mask)
+            if self.advantage_clip is not None:
+                advantages = advantages.clamp(min=-self.advantage_clip, max=self.advantage_clip)
+            if self.positive_advantages_only:
+                advantages = advantages.clamp_min(0.0)
+
+            advantages = advantages.detach()
+
+            if self.use_ppo_clip:
+                # The rollouts are generated by the frozen reference/SFT policy, so ref_logp
+                # is used as the old rollout log-prob. This mirrors the PPO-style clipped
+                # policy loss used in RL frameworks such as slime.
+                ratio = torch.exp(student_logp - ref_logp.detach())
+                ratio_clipped = ratio.clamp(1.0 - self.ppo_clip_low, 1.0 + self.ppo_clip_high)
+                pg_loss_unclipped = -ratio * advantages
+                pg_loss_clipped = -ratio_clipped * advantages
+                token_loss = torch.maximum(pg_loss_unclipped, pg_loss_clipped)
+            else:
+                # Direct OPD-style sampled-token surrogate.
+                token_loss = -advantages * student_logp
+
+            if self.sampled_loss_reduction == "per_token":
+                loss = (token_loss * loss_mask).sum() / total_tokens
+            else:
+                # Per-sample mean: each response contributes equally regardless of length.
+                seq_loss = (token_loss * loss_mask).sum(dim=1) / loss_mask.sum(dim=1).clamp_min(1.0)
+                loss = seq_loss.mean()
+
+            return (loss, student_outputs) if return_outputs else loss
+
+        seq_len = student_logits.size(1)
+        for start in range(0, seq_len, self.distill_chunk_size):
+            end = min(start + self.distill_chunk_size, seq_len)
+            mask = loss_mask[:, start:end]
+            if mask.sum() == 0:
+                continue
+
+            teacher_logp = F.log_softmax(teacher_logits[:, start:end, :].float(), dim=-1)
+            ref_logp = F.log_softmax(ref_logits[:, start:end, :].float(), dim=-1)
+            student_logp = F.log_softmax(student_logits[:, start:end, :].float(), dim=-1)
+
+            poe_logits = beta * teacher_logp + (1.0 - beta) * ref_logp
+            target_probs = F.softmax(poe_logits, dim=-1)
+            token_ce = -(target_probs * student_logp).sum(dim=-1)
+            total_loss = total_loss + (token_ce * mask).sum()
+
+        loss = total_loss / total_tokens
+        return (loss, student_outputs) if return_outputs else loss
+
+
+class FreezeLoRABCallback(TrainerCallback):
+    def __init__(self, freeze_after_step: int | None) -> None:
+        self.freeze_after_step = freeze_after_step
+        self.frozen = False
+
+    def on_step_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        model: torch.nn.Module | None = None,
+        **kwargs: Any,
+    ) -> TrainerControl:
+        if self.freeze_after_step is None or self.frozen or model is None:
+            return control
+        if state.global_step < self.freeze_after_step:
+            return control
+
+        frozen_params = 0
+        module = model.module if hasattr(model, "module") else model
+        for name, param in module.named_parameters():
+            if ".lora_B." in name or "lora_B." in name:
+                param.requires_grad_(False)
+                frozen_params += param.numel()
+
+        self.frozen = True
+        if args.process_index == 0:
+            print(f"[PoE Distill] Froze LoRA B at global_step={state.global_step} ({frozen_params} params).")
+        return control
+
+
+class HoldThenTransitionLRCallback(TrainerCallback):
+    """Custom LR schedule: hold lr_start, transition to lr_end, then keep lr_end.
+
+    Disable the Hugging Face scheduler interaction by setting
+    --lr-scheduler-type constant and --warmup-ratio 0.0 in the launcher when
+    using --lr-end.  This callback sets optimizer param-group LRs directly at
+    each optimizer step.
+    """
+
+    def __init__(
+        self,
+        lr_start: float,
+        lr_end: float | None,
+        lr_hold_steps: int,
+        lr_transition_steps: int | None,
+        schedule: str,
+    ) -> None:
+        self.lr_start = lr_start
+        self.lr_end = lr_end
+        self.lr_hold_steps = lr_hold_steps
+        self.lr_transition_steps = lr_transition_steps
+        self.schedule = schedule
+
+    def current_lr(self, step: int) -> float:
+        if self.lr_end is None:
+            return self.lr_start
+        return hold_then_transition_value(
+            step=step,
+            start=self.lr_start,
+            end=self.lr_end,
+            hold_steps=self.lr_hold_steps,
+            transition_steps=self.lr_transition_steps,
+            schedule=self.schedule,
+        )
+
+    def on_train_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        optimizer: torch.optim.Optimizer | None = None,
+        **kwargs: Any,
+    ) -> TrainerControl:
+        return self._set_lr(args, state, control, optimizer)
+
+    def on_step_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        optimizer: torch.optim.Optimizer | None = None,
+        **kwargs: Any,
+    ) -> TrainerControl:
+        return self._set_lr(args, state, control, optimizer)
+
+    def _set_lr(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        optimizer: torch.optim.Optimizer | None,
+    ) -> TrainerControl:
+        if optimizer is None or self.lr_end is None:
+            return control
+
+        lr = self.current_lr(state.global_step)
+        for group in optimizer.param_groups:
+            group["lr"] = lr
+
+        if args.process_index == 0 and state.global_step % max(args.logging_steps, 1) == 0:
+            print(f"[PoE Distill] global_step={state.global_step}, custom_lr={lr:.3e}")
+        return control
+
+
+class BetaLoggingCallback(TrainerCallback):
+    """Log beta occasionally without changing training behavior."""
+
+    def __init__(self, trainer_ref_getter) -> None:
+        self.trainer_ref_getter = trainer_ref_getter
+
+    def on_step_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs: Any,
+    ) -> TrainerControl:
+        trainer = self.trainer_ref_getter()
+        if trainer is not None and args.process_index == 0 and state.global_step % max(args.logging_steps, 1) == 0:
+            print(f"[PoE Distill] global_step={state.global_step}, beta={trainer.current_beta():.6f}")
+        return control
+
+
+def main() -> None:
+    args = parse_args()
+    set_seed(args.seed)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.student_model, trust_remote_code=args.trust_remote_code)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    raw_dataset = load_dataset("parquet", data_files=args.train_data, split="train")
+    if args.max_train_samples is not None:
+        raw_dataset = raw_dataset.select(range(min(args.max_train_samples, len(raw_dataset))))
+
+    train_dataset = raw_dataset.map(
+        lambda ex: tokenize_rollout(ex, tokenizer, args.max_length),
+        remove_columns=raw_dataset.column_names,
+        desc="Tokenizing pi_ref rollouts",
+    ).filter(lambda ex: len(ex["input_ids"]) > 0, desc="Dropping empty responses")
+
+    model_kwargs = {
+        "torch_dtype": torch.bfloat16 if args.bf16 else (torch.float16 if args.fp16 else torch.float32),
+        "trust_remote_code": args.trust_remote_code,
+    }
+    if args.attn_implementation is not None:
+        model_kwargs["attn_implementation"] = args.attn_implementation
+
+    student = AutoModelForCausalLM.from_pretrained(args.student_model, **model_kwargs)
+    teacher = AutoModelForCausalLM.from_pretrained(args.teacher_model, **model_kwargs)
+    teacher.eval()
+    teacher.requires_grad_(False)
+
+    if args.gradient_checkpointing:
+        student.gradient_checkpointing_enable()
+        student.config.use_cache = False
+        teacher.config.use_cache = False
+
+    lora_config = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        target_modules=[m.strip() for m in args.lora_target_modules.split(",") if m.strip()],
+    )
+    student = get_peft_model(student, lora_config)
+    student.print_trainable_parameters()
+
+    training_args = TrainingArguments(
+        output_dir=args.output_dir,
+        num_train_epochs=args.num_train_epochs,
+        max_steps=args.max_steps,
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        learning_rate=args.learning_rate,
+        weight_decay=args.weight_decay,
+        adam_beta1=args.adam_beta1,
+        adam_beta2=args.adam_beta2,
+        adam_epsilon=args.adam_epsilon,
+        warmup_ratio=args.warmup_ratio,
+        lr_scheduler_type=args.lr_scheduler_type,
+        logging_steps=args.logging_steps,
+        save_steps=args.save_steps,
+        save_total_limit=args.save_total_limit,
+        bf16=args.bf16,
+        fp16=args.fp16,
+        gradient_checkpointing=args.gradient_checkpointing,
+        remove_unused_columns=False,
+        report_to=[] if args.report_to == "none" else args.report_to.split(","),
+    )
+
+    trainer = PoEDistillTrainer(
+        model=student,
+        args=training_args,
+        train_dataset=train_dataset,
+        data_collator=DistillCollator(pad_token_id=tokenizer.pad_token_id),
+        tokenizer=tokenizer,
+        teacher_model=teacher,
+        beta_start=args.beta_start,
+        beta_end=args.beta_end,
+        beta_schedule_steps=args.beta_schedule_steps,
+        beta_schedule=args.beta_schedule,
+        beta_hold_steps=args.beta_hold_steps,
+        beta_transition_steps=args.beta_transition_steps,
+        hold_transition_schedule=args.hold_transition_schedule,
+        loss_type=args.loss_type,
+        advantage_normalization=args.advantage_normalization,
+        advantage_clip=args.advantage_clip,
+        positive_advantages_only=args.positive_advantages_only,
+        use_ppo_clip=args.use_ppo_clip,
+        ppo_clip_low=args.ppo_clip_low,
+        ppo_clip_high=args.ppo_clip_high,
+        sampled_loss_reduction=args.sampled_loss_reduction,
+        distill_chunk_size=args.distill_chunk_size,
+        callbacks=[
+            FreezeLoRABCallback(args.freeze_lora_b_after_step),
+            HoldThenTransitionLRCallback(
+                lr_start=args.lr_start,
+                lr_end=args.lr_end,
+                lr_hold_steps=args.lr_hold_steps,
+                lr_transition_steps=args.lr_transition_steps,
+                schedule=args.hold_transition_schedule,
+            ),
+        ],
+    )
+    trainer.train()
+    trainer.save_model(args.output_dir)
+    tokenizer.save_pretrained(args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

tools/convert_fsdp_to_hf.py

@@ -0,0 +1,181 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+import pickle
+import shutil
+import time
+
+import torch
+import torch.distributed.checkpoint as dist_cp
+from transformers import AutoConfig, AutoModelForCausalLM
+from typing_extensions import override
+
+
+class UnpicklerWrapper(pickle.Unpickler):
+    @override
+    def find_class(self, mod_name, name):
+        class DummyClass:
+            def __init__(self, *args, **kwargs):
+                pass
+
+        if mod_name.startswith("megatron") or mod_name.startswith("glm"):
+            return DummyClass
+        return super().find_class(mod_name, name)
+
+
+class WrappedStorageReader(dist_cp.FileSystemReader):
+    @override
+    def read_metadata(self):
+        path = self.fs.concat_path(self.path, ".metadata")
+        with self.fs.create_stream(path, "rb") as metadata_file:
+            metadata = UnpicklerWrapper(metadata_file).load()
+        if getattr(metadata, "storage_meta", None) is None:
+            metadata.storage_meta = dist_cp.StorageMeta()
+        metadata.storage_meta.load_id = self.load_id
+        if metadata.planner_data is None:
+            metadata.planner_data = {}
+        return metadata
+
+
+class EmptyStateDictLoadPlanner(dist_cp.default_planner.DefaultLoadPlanner):
+    @override
+    def set_up_planner(
+        self,
+        state_dict: dist_cp.metadata.STATE_DICT_TYPE,
+        metadata: dist_cp.metadata.Metadata | None = None,
+        is_coordinator: bool = False,
+    ) -> None:
+        for k, v in metadata.state_dict_metadata.items():
+            if "optimizer" in k:
+                continue
+            print(f"find {k} in torch_dist ckpt")
+            if isinstance(v, dist_cp.metadata.TensorStorageMetadata):
+                v = torch.empty(v.size, dtype=v.properties.dtype)  # type: ignore[assignment]
+            state_dict[k] = v
+        super().set_up_planner(state_dict, metadata, is_coordinator)
+
+
+def _detect_model_dir(input_dir: str) -> str:
+    model_dir = os.path.join(input_dir, "model")
+    return model_dir if os.path.isdir(model_dir) else input_dir
+
+
+def _load_fsdp_state_dict(input_dir: str) -> dict[str, torch.Tensor]:
+    state_dict: dict[str, torch.Tensor] = {}
+    dist_cp.state_dict_loader._load_state_dict(
+        state_dict,
+        storage_reader=WrappedStorageReader(input_dir),
+        planner=EmptyStateDictLoadPlanner(),
+        no_dist=True,
+    )
+    return state_dict
+
+
+def _get_candidate_prefixes(keys: list[str]) -> list[str]:
+    predefined = [
+        "model_state.model.",
+        "model_state.",
+        "model.",
+        "module.",
+        "",
+    ]
+
+    detected: set[str] = set()
+    for key in keys:
+        for prefix in predefined:
+            if prefix and key.startswith(prefix):
+                detected.add(prefix)
+
+    # Always keep empty string as a fall back option for exact match.
+    detected.add("")
+    # Preserve predefined order while keeping only detected prefixes.
+    return [p for p in predefined if p in detected]
+
+
+def _strip_best_prefix(keys: list[str], target_keys: set[str]) -> tuple[str, int]:
+    best_prefix = ""
+    best_match = -1
+
+    for prefix in _get_candidate_prefixes(keys):
+        mapped_keys = {k.removeprefix(prefix) for k in keys}
+        match_count = len(mapped_keys & target_keys)
+        if match_count > best_match:
+            best_match = match_count
+            best_prefix = prefix
+
+    return best_prefix, best_match
+
+
+def _convert_fsdp_to_hf(
+    origin_hf_dir: str,
+    input_dir: str,
+    output_dir: str,
+) -> None:
+    print(f"loading FSDP model from {input_dir}")
+    t = time.time()
+    state_dict = _load_fsdp_state_dict(input_dir)
+    print(f"FSDP model loaded in {time.time()-t:.2f} sec.")
+
+    tensor_items = {k: v for k, v in state_dict.items() if isinstance(v, torch.Tensor)}
+
+    config = AutoConfig.from_pretrained(origin_hf_dir, trust_remote_code=True)
+    hf_model = AutoModelForCausalLM.from_config(config)
+    target_keys = set(hf_model.state_dict().keys())
+
+    best_prefix, best_match = _strip_best_prefix(list(tensor_items.keys()), target_keys)
+    total_keys = len(tensor_items)
+
+    print(f"Using prefix '{best_prefix}' for key mapping. " f"Matched {best_match}/{total_keys} parameter keys.")
+
+    model_state = {k.removeprefix(best_prefix): v for k, v in tensor_items.items()}
+
+    if not model_state:
+        raise ValueError(
+            "No model weights found in checkpoint. "
+            "Please pass the checkpoint directory (e.g. iter_xxx or iter_xxx/model)."
+        )
+
+    missing, unexpected = hf_model.load_state_dict(model_state, strict=False)
+    print(f"Missing keys: {missing}\nUnexpected keys: {unexpected}")
+
+    os.makedirs(output_dir, exist_ok=True)
+    hf_model.save_pretrained(output_dir, safe_serialization=True)
+    print(f"Model weights saved to {output_dir}")
+
+
+def copy_assets(origin_hf_dir: str, output_dir: str) -> None:
+    for filename in os.listdir(origin_hf_dir):
+        if filename == "model.safetensors.index.json" or filename.endswith(".safetensors"):
+            continue
+        origin_filename = os.path.join(origin_hf_dir, filename)
+        if not os.path.isfile(origin_filename):
+            print(f"Skip {filename}, not a file.")
+            continue
+        src, dst = origin_filename, os.path.join(output_dir, filename)
+        print(f"copy from {src} to {dst}")
+        shutil.copy(src, dst)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input-dir", type=str, required=True)
+    parser.add_argument("--output-dir", type=str, required=True)
+    parser.add_argument(
+        "--origin-hf-dir",
+        type=str,
+        required=True,
+        help="The original Hugging Face model directory to load config/tokenizer assets.",
+    )
+    parser.add_argument(
+        "-f", "--force", action="store_true", help="Force overwrite the output directory if it exists."
+    )
+    args = parser.parse_args()
+
+    if os.path.exists(args.output_dir) and not args.force:
+        raise ValueError(f"Output directory {args.output_dir} already exists. Use --force to overwrite it.")
+
+    model_dir = _detect_model_dir(args.input_dir)
+    _convert_fsdp_to_hf(args.origin_hf_dir, model_dir, args.output_dir)
+    copy_assets(args.origin_hf_dir, args.output_dir)

tools/convert_hf_to_fp8.py

@@ -0,0 +1,256 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+python tools/convert_hf_to_fp8.py [-h] [--model-dir MODEL_DIR] [--save-dir SAVE_DIR] [--strategy {block,channel,tensor}] [--block-size [BLOCK_SIZE ...]]
+                           [--max-workers MAX_WORKERS]
+
+options:
+  -h, --help            show this help message and exit
+  --model-dir MODEL_DIR
+                        Path to the directory of the HF safetensors model.
+  --save-dir SAVE_DIR   Path to the directory to save the converted model.
+  --strategy {block,channel,tensor}
+  --block-size [BLOCK_SIZE ...]
+                        eg. --block-size 32 32
+  --max-workers MAX_WORKERS
+                        Number of worker threads for parallel processing
+"""
+
+import argparse
+import gc
+import json
+import os
+import shutil
+import threading
+from concurrent.futures import ThreadPoolExecutor
+
+import safetensors
+import safetensors.torch
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+FP8_INFO = torch.finfo(torch.float8_e4m3fn)
+FP8_MAX, FP8_MIN = FP8_INFO.max, FP8_INFO.min
+
+
+def ceildiv(a, b):
+    return -(-a // b)
+
+
+def block_fp8(weight, block_size):
+
+    # per block quant
+    block_n, block_k = block_size[0], block_size[1]
+
+    shape_0, shape_1 = weight.shape
+
+    n_tiles = ceildiv(shape_0, block_n)
+    k_tiles = ceildiv(shape_1, block_k)
+
+    q_weight = F.pad(
+        weight,
+        (0, k_tiles * block_k - shape_1, 0, n_tiles * block_n - shape_0),
+        mode="constant",
+        value=0.0,
+    )
+
+    qweight = q_weight.reshape(n_tiles, block_n, k_tiles, block_k)
+    block_max = torch.max(torch.abs(qweight), dim=1, keepdim=True)[0]
+    block_max = torch.max(block_max, dim=3, keepdim=True)[0]
+
+    scale = block_max.to(torch.float32) / FP8_MAX
+    qweight = (
+        (qweight / scale)
+        .clamp(min=FP8_MIN, max=FP8_MAX)
+        .reshape((n_tiles * block_n, k_tiles * block_k))
+        .to(torch.float8_e4m3fn)
+    )
+    qweight = qweight[:shape_0, :shape_1].clone().detach()
+    scale = scale.squeeze()
+
+    return qweight, scale
+
+
+def channel_fp8(weight):
+    channel_max = torch.max(weight.abs(), dim=-1, keepdim=True)[0]
+    scale = channel_max.clamp(min=1e-12).to(torch.float32) / FP8_MAX
+    qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX)
+    qweight = qweight.to(torch.float8_e4m3fn)
+    return qweight, scale
+
+
+def tensor_fp8(weight):
+    scale = weight.abs().max().clamp(min=1e-12).to(torch.float32) / FP8_MAX
+    qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX)
+    qweight = qweight.to(torch.float8_e4m3fn)
+    scale = scale.view(1)
+    return qweight, scale
+
+
+def quant_fp8(weight, strategy, block_size=None):
+    if strategy == "tensor":
+        return tensor_fp8(weight)
+    elif strategy == "channel":
+        return channel_fp8(weight)
+    else:
+        return block_fp8(weight, block_size)
+
+
+class ConversionResult:
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.weight_map = {}
+        self.param_count = 0
+        self.modules_to_not_convert = []
+
+    def add_result(self, filename, q_weights, module_names):
+        with self.lock:
+            for k, v in q_weights.items():
+                self.weight_map[k] = filename
+                self.param_count += len(v)
+            self.modules_to_not_convert.extend(module_names)
+
+
+def process_file(input_path, output_path, filename, strategy, block_size, result_collector):
+    if not filename.endswith(".safetensors"):
+        return
+
+    print(f"Processing {filename}, memory usage: {torch.cuda.memory_allocated()}")
+    weights = {}
+    q_weights = {}
+
+    with safetensors.safe_open(os.path.join(input_path, filename), framework="pt", device="cuda") as f:
+        for k in f.keys():
+            weights[k] = f.get_tensor(k)
+
+    modules_to_not_convert = []
+    for key in weights.keys():
+        if (
+            "weight" in key
+            and "layernorm" not in key
+            and "embed" not in key
+            and "router" not in key
+            and "mlp.gate." not in key
+            and "norm" not in key
+            and "lm_head" not in key
+            and "eh_proj" not in key
+        ):
+            qw, s = quant_fp8(weights[key], strategy, block_size)
+            q_weights[key] = qw
+            if block_size:
+                scale_name = key.replace(".weight", ".weight_scale_inv")
+            else:
+                scale_name = key.replace(".weight", ".weight_scale")
+            q_weights[scale_name] = s
+        else:
+            modules_to_not_convert.append(key.replace(".weight", ""))
+            q_weights[key] = weights[key]
+
+    safetensors.torch.save_file(q_weights, os.path.join(output_path, filename), metadata={"format": "pt"})
+
+    result_collector.add_result(filename, q_weights, modules_to_not_convert)
+
+
+def convert_fp8(input_path, output_path, strategy, block_size=None, max_workers=4):
+    input_path = os.path.abspath(input_path)
+    os.makedirs(output_path, exist_ok=True)
+
+    for filename in os.listdir(input_path):
+        if not filename.endswith(".safetensors") and not os.path.isdir(os.path.join(input_path, filename)):
+            shutil.copyfile(os.path.join(input_path, filename), os.path.join(output_path, filename))
+
+    safetensors_files = [f for f in os.listdir(input_path) if f.endswith(".safetensors")]
+
+    result_collector = ConversionResult()
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = []
+        for filename in safetensors_files:
+            future = executor.submit(
+                process_file, input_path, output_path, filename, strategy, block_size, result_collector
+            )
+            futures.append(future)
+
+        for future in tqdm(futures, desc="Processing files"):
+            future.result()
+
+    if strategy == "block" or strategy == "tensor":
+        quantization_config = {
+            "activation_scheme": "dynamic",
+            "fmt": "e4m3",
+            "quant_method": "fp8",
+        }
+        if block_size:
+            quantization_config["weight_block_size"] = block_size
+        if len(result_collector.modules_to_not_convert) > 0:
+            quantization_config["modules_to_not_convert"] = list(set(result_collector.modules_to_not_convert))
+    else:
+        quant_group = {
+            "group_0": {
+                "input_activations": {
+                    "actorder": None,
+                    "block_structure": None,
+                    "dynamic": True,
+                    "group_size": None,
+                    "num_bits": 8,
+                    "observer": None,
+                    "observer_kwargs": {},
+                    "strategy": "token",
+                    "symmetric": True,
+                    "type": "float",
+                },
+                "output_activations": None,
+                "targets": ["Linear"],
+                "weights": {
+                    "actorder": None,
+                    "block_structure": None,
+                    "dynamic": False,
+                    "group_size": None,
+                    "num_bits": 8,
+                    "observer": "minmax",
+                    "observer_kwargs": {},
+                    "strategy": strategy,
+                    "symmetric": True,
+                    "type": "float",
+                },
+            },
+        }
+        quantization_config = {
+            "config_groups": quant_group,
+            "format": "float-quantized",
+            "ignore": list(set(result_collector.modules_to_not_convert)),
+            "quant_method": "compressed-tensors",
+            "quantization_status": "compressed",
+        }
+
+    config_path = os.path.join(input_path, "config.json")
+    if os.path.exists(config_path):
+        cfg = json.load(open(config_path))
+        cfg["quantization_config"] = quantization_config
+        json.dump(cfg, open(os.path.join(output_path, "config.json"), "w"), indent=2)
+
+    index_dict = {"weight_map": result_collector.weight_map, "metadata": {"total_size": result_collector.param_count}}
+    json.dump(index_dict, open(os.path.join(output_path, "model.safetensors.index.json"), "w"), indent=2)
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir", type=str, help="Path to the directory of the HF safetensors model.")
+    parser.add_argument("--save-dir", type=str, help="Path to the directory to save the converted model.")
+    parser.add_argument("--strategy", type=str, default="block", choices=["block", "channel", "tensor"])
+    parser.add_argument("--block-size", type=int, nargs="*", default=None, help="eg. --block-size 32 32")
+    parser.add_argument("--max-workers", type=int, default=1, help="Number of worker threads for parallel processing")
+    args = parser.parse_args()
+
+    if not os.path.exists(args.save_dir):
+        print(f"Creating directory {args.save_dir}")
+        os.makedirs(args.save_dir)
+    elif not os.path.isdir(args.save_dir):
+        raise ValueError("The save_dir should be a directory.")
+
+    convert_fp8(args.model_dir, args.save_dir, args.strategy, args.block_size, args.max_workers)

tools/convert_hf_to_torch_dist.py

@@ -0,0 +1,140 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import gc
+import os
+import shutil
+
+import torch
+import torch.distributed as dist
+from megatron.core.enums import ModelType
+from megatron.training.arguments import parse_args, validate_args
+from megatron.training.checkpointing import get_checkpoint_name, get_checkpoint_tracker_filename, save_checkpoint
+from megatron.training.training import get_model
+
+import slime_plugins.mbridge  # noqa: F401
+from mbridge import AutoBridge
+from slime.backends.megatron_utils.arguments import set_default_megatron_args
+from slime.backends.megatron_utils.initialize import init
+from slime.backends.megatron_utils.model_provider import get_model_provider_func
+from slime.utils.logging_utils import configure_logger
+from slime.utils.memory_utils import print_memory
+
+
+def add_convertion_args(parser):
+    """Add conversion arguments to the parser"""
+    parser.add_argument("--hf-checkpoint", type=str, required=True, help="HuggingFace model path")
+    parser.add_argument(
+        "--megatron-to-hf-mode",
+        choices=["raw", "bridge"],
+        default="raw",
+        help="The method to convert megatron weights to hugging face weights for SGLang.",
+    )
+    try:
+        parser.add_argument("--padded-vocab-size", type=int, default=None)
+    except Exception:
+        pass
+    return parser
+
+
+def get_args():
+    args = parse_args(add_convertion_args)
+    args = set_default_megatron_args(args)
+
+    # set to pass megatron validate_args
+    args.save_interval = 1
+    args.micro_batch_size = 1
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    args.global_batch_size = int(os.environ.get("WORLD_SIZE", "1"))
+
+    assert world_size <= args.num_layers, (
+        f"World size {world_size} must be less than or equal to number of layers {args.num_layers}. "
+        "You are using too many GPUs for this conversion."
+    )
+
+    def ceildiv(a, b):
+        return -(a // -b)
+
+    if args.pipeline_model_parallel_size == 1 and world_size > 1:
+        pp_size = world_size
+        while True:
+            args.pipeline_model_parallel_size = pp_size
+            args.decoder_last_pipeline_num_layers = args.num_layers - ceildiv(
+                args.num_layers, args.pipeline_model_parallel_size
+            ) * (args.pipeline_model_parallel_size - 1)
+
+            if args.decoder_last_pipeline_num_layers > 0:
+                break
+
+            if pp_size % 2 == 0:
+                pp_size //= 2
+            else:
+                raise ValueError(
+                    f"Cannot find a valid pipeline model parallel size for {args.num_layers} layers and {world_size} GPUs."
+                )
+    print(
+        f"Using pipeline model parallel size: {args.pipeline_model_parallel_size}, decoder last pipeline num layers: {args.decoder_last_pipeline_num_layers}"
+    )
+
+    validate_args(args)
+    return args
+
+
+def main():
+    if torch.version.hip:
+        import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module
+        from slime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync
+
+        filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync
+        print("[ROCm] Applied FileSystemWriterAsync patch for HIP compatibility")
+
+    configure_logger()
+
+    # Initialize distributed environment
+    world_size = int(os.getenv("WORLD_SIZE") or os.getenv("SLURM_NTASKS") or 1)
+    local_rank = int(os.getenv("LOCAL_RANK") or os.getenv("SLURM_LOCALID") or 0)
+    global_rank = int(os.getenv("RANK") or os.getenv("SLURM_PROCID") or 0)
+
+    torch.cuda.set_device(local_rank)
+    os.environ.setdefault("WORLD_SIZE", str(world_size))
+    os.environ.setdefault("RANK", str(global_rank))
+    os.environ.setdefault("LOCAL_RANK", str(local_rank))
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "12355")
+    dist.init_process_group(
+        backend="nccl",
+        world_size=world_size,
+        rank=global_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
+    args = get_args()
+    init(args)
+    model = get_model(get_model_provider_func(args), ModelType.encoder_or_decoder, wrap_with_ddp=False)
+
+    # Load model
+    hf_model_path = args.hf_checkpoint
+    bridge = AutoBridge.from_pretrained(hf_model_path, trust_remote_code=True)
+    bridge.load_weights(model, hf_model_path, memory_efficient=True)
+    print(f"Model loaded: {hf_model_path}")
+
+    print_memory("after loading model")
+    torch.cuda.synchronize()
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    save_checkpoint(1, model, None, None, 0)
+
+    if dist.get_rank() == 0:
+        # change to release ckpt
+        tracker_filename = get_checkpoint_tracker_filename(args.save)
+        with open(tracker_filename, "w") as f:
+            f.write("release")
+        source_dir = get_checkpoint_name(args.save, 1, False, return_base_dir=True)
+        target_dir = get_checkpoint_name(args.save, -1, True, return_base_dir=True)
+        shutil.move(source_dir, target_dir)
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

tools/convert_k2_thinking_int4_to_bf16.py

@@ -0,0 +1,245 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Usage:
+------
+python convert_k2_thinking_int4_to_bf16.py [-h] --model-dir MODEL_DIR [--output-dir OUTPUT_DIR]
+                           [--files FILE [FILE ...]] [--config-path CONFIG_PATH]
+                           [--overwrite]
+options:
+  -h, --help            Show this help message and exit.
+  --model-dir MODEL_DIR Path to the directory of the HF safetensors quantized model.
+  --output-dir OUTPUT_DIR
+                        Path to the directory to save the converted BF16 model.
+                        Default: <model-dir>_bf16
+  --files FILE [FILE ...]
+                        Specific safetensors filenames to convert (relative to model-dir).
+                        Convert all if omitted.
+  --config-path CONFIG_PATH
+                        Path to config.json to extract group_size (default: model-dir/config.json).
+  --overwrite           Rewrite output files even if they already exist.
+
+
+Example:
+--------
+python convert_k2_thinking_int4_to_bf16.py --model-dir /Kimi-K2-Thinking --output-dir /Kimi-K2-Thinking-bf16
+"""
+
+import argparse
+import json
+import os
+import shutil
+from collections import defaultdict
+
+import torch
+from compressed_tensors.compressors import unpack_from_int32
+from safetensors.torch import safe_open, save_file
+from tqdm import tqdm
+
+
+def _load_config(model_dir: str, config_path: str | None) -> tuple[int, int, int]:
+    """Read config.json and return hidden_size, inter_size, and group_size."""
+    cfg_path = config_path or os.path.join(model_dir, "config.json")
+    with open(cfg_path) as f:
+        cfg = json.load(f)
+    hidden_size = int(cfg.get("hidden_size"))
+    inter_size = int(cfg.get("moe_intermediate_size"))
+    group_size = int(
+        cfg.get("quantization_config", {})
+        .get("config_groups", {})
+        .get("group_0", {})
+        .get("weights", {})
+        .get("group_size", 128)
+    )
+    return hidden_size, inter_size, group_size
+
+
+def _dequantize_tensor(
+    weight_packed: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_shape: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    """Unpack int32 quantized tensor and multiply with scales to create BF16 tensor."""
+    if isinstance(weight_shape, torch.Tensor):
+        shape = tuple(int(v) for v in weight_shape.view(-1).tolist())
+    else:
+        shape = tuple(weight_shape)
+
+    weight = unpack_from_int32(weight_packed, 4, shape)
+
+    if group_size > 0:
+        scale = weight_scale.to(torch.float32)
+        if scale.dim() == 1:
+            scale = scale.unsqueeze(1)
+        scales = torch.repeat_interleave(scale, repeats=group_size, dim=1)
+    else:
+        scales = weight_scale.to(torch.float32)
+
+    if scales.shape != weight.shape:
+        if scales.numel() == weight.numel():
+            scales = scales.reshape_as(weight)
+        else:
+            raise ValueError(f"Scale shape {scales.shape} incompatible with weight shape {weight.shape}")
+
+    bf16 = (weight.to(torch.float32) * scales).to(torch.bfloat16)
+    return bf16.contiguous()
+
+
+def _is_quantized_weight_key(key: str) -> bool:
+    """Check if the key is a quantized MoE expert weight key."""
+    if ".mlp.experts." not in key or ".shared_experts." in key:
+        return False
+    suffixes = ("weight_packed", "weight_scale", "weight_shape")
+    for proj in ("gate_proj", "up_proj", "down_proj"):
+        for suffix in suffixes:
+            if key.endswith(f".{proj}.{suffix}"):
+                return True
+    return False
+
+
+def convert_file(
+    input_path: str,
+    output_path: str,
+    group_size: int,
+    skip_existing: bool = True,
+):
+    """Convert a single safetensors file from quantized format to BF16 (GPU accelerated)."""
+    if skip_existing and os.path.exists(output_path):
+        return
+
+    tensors = {}
+    expert_buffers = defaultdict(lambda: defaultdict(dict))
+
+    # Load weights directly on GPU
+    with safe_open(input_path, framework="pt", device="cuda") as reader:
+        keys = list(reader.keys())
+        for key in keys:
+            tensor = reader.get_tensor(key)
+            if not _is_quantized_weight_key(key):
+                tensors[key] = tensor
+                continue
+            parts = key.split(".")
+            try:
+                expert_idx = parts.index("experts")
+            except ValueError:
+                tensors[key] = tensor
+                continue
+            prefix = ".".join(parts[: expert_idx + 2])
+            project = parts[-2]
+            suffix = parts[-1]
+            expert_buffers[prefix][project][suffix] = tensor
+
+    # Convert quantized weights
+    for prefix, components in expert_buffers.items():
+        for proj_name in ["gate_proj", "up_proj", "down_proj"]:
+            proj_data = components.get(proj_name, {})
+            required = {"weight_packed", "weight_scale", "weight_shape"}
+
+            if not required.issubset(proj_data.keys()):
+                # Keep quantized tensors if incomplete
+                for suffix, value in proj_data.items():
+                    tensors[f"{prefix}.{proj_name}.{suffix}"] = value
+                continue
+
+            # Dequantize to BF16
+            bf16_weight = _dequantize_tensor(
+                proj_data["weight_packed"].to(torch.int32),
+                proj_data["weight_scale"].to(torch.float32),
+                proj_data["weight_shape"],
+                group_size,
+            )
+            tensors[f"{prefix}.{proj_name}.weight"] = bf16_weight.to(torch.bfloat16)
+
+    # Save converted file (moved to CPU for compatibility)
+    cpu_tensors = {k: v.cpu() for k, v in tensors.items()}
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    save_file(cpu_tensors, output_path)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Convert GPTQ MoE experts to BF16 weights.")
+    parser.add_argument("--model-dir", required=True, help="Directory containing safetensors checkpoints.")
+    parser.add_argument(
+        "--output-dir",
+        default=None,
+        help="Destination BF16 model directory (default: <model-dir>_bf16).",
+    )
+    parser.add_argument(
+        "--files",
+        nargs="+",
+        default=None,
+        help="Optional specific safetensor files to convert.",
+    )
+    parser.add_argument(
+        "--config-path",
+        default=None,
+        help="Path to config.json if not in model-dir.",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Overwrite existing BF16 files.",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    model_dir = os.path.abspath(args.model_dir)
+    output_dir = os.path.abspath(args.output_dir or f"{model_dir}_bf16")
+
+    if not os.path.isdir(model_dir):
+        raise FileNotFoundError(f"Model directory not found: {model_dir}")
+
+    _, _, group_size = _load_config(model_dir, args.config_path)
+
+    # Collect target files
+    if args.files:
+        targets = [os.path.join(model_dir, fname) for fname in args.files]
+    else:
+        targets = [
+            os.path.join(model_dir, name) for name in sorted(os.listdir(model_dir)) if name.endswith(".safetensors")
+        ]
+
+    if not targets:
+        print("No safetensors checkpoints found.")
+        return
+
+    # Convert with progress bar
+    for path in tqdm(targets, desc="Converting weights", unit="file"):
+        if not os.path.isfile(path):
+            continue
+        rel = os.path.relpath(path, model_dir)
+        output_path = os.path.join(output_dir, rel)
+        convert_file(path, output_path, group_size, skip_existing=not args.overwrite)
+
+    # Copy config/json/py/tokenizer
+    for fname in os.listdir(model_dir):
+        src_path = os.path.join(model_dir, fname)
+        dst_path = os.path.join(output_dir, fname)
+        if fname == "model.safetensors.index.json":
+            continue
+        if fname.endswith(".json") or fname.endswith(".py") or fname.startswith("tokenizer"):
+            shutil.copy2(src_path, dst_path)
+
+    # Generate new index
+    new_index_path = os.path.join(output_dir, "model.safetensors.index.json")
+    weight_map = {}
+    for fname in sorted(os.listdir(output_dir)):
+        if not fname.endswith(".safetensors"):
+            continue
+        safetensor_path = os.path.join(output_dir, fname)
+        with safe_open(safetensor_path, framework="pt") as reader:
+            for key in reader.keys():
+                weight_map[key] = fname
+
+    with open(new_index_path, "w") as f:
+        json.dump({"metadata": {}, "weight_map": weight_map}, f, indent=2)
+
+    print(f"\nSuccessful! Output saved to: {output_dir}")
+
+
+if __name__ == "__main__":
+    main()

tools/convert_to_hf.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+from megatron.core import mpu
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+import slime.backends.megatron_utils as megatron_utils
+from slime.backends.megatron_utils import update_weight_utils
+from slime.utils.arguments import parse_args
+
+
+def add_checkpoint_args(parser):
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default=None,
+        help="Directory to save the converted HF model.",
+    )
+    parser.add_argument(
+        "--check-same",
+        action="store_true",
+        default=False,
+        help="Check if the converted model is the same as the original model.",
+    )
+    return parser
+
+
+def main(args):
+    megatron_utils.init(args)
+
+    pp_size = mpu.get_pipeline_model_parallel_world_size()
+    ep_size = mpu.get_expert_model_parallel_world_size()
+
+    is_save_rank = (
+        mpu.get_data_parallel_rank(with_context_parallel=True) == 0 and mpu.get_tensor_model_parallel_rank() == 0
+    )
+
+    # Setup the model and optimizer
+    args.no_load_optim = True
+    args.no_load_rng = True
+    model, _, _, _ = megatron_utils.initialize_model_and_optimizer(args)
+
+    hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
+    model_name = type(hf_config).__name__.lower()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
+
+    vocab_size = tokenizer.vocab_size if args.vocab_size is None else args.vocab_size
+
+    param_infos = update_weight_utils.get_param_infos(args, model)
+
+    state_dict = {}
+    rank = dist.get_rank()
+    for info in param_infos:
+        if dist.get_rank() == info.src_rank:
+            for name_, param_ in update_weight_utils.named_parameters(args, model):
+                if name_ == info.name:
+                    param = param_
+                    break
+        else:
+            param = torch.empty(info.shape, dtype=info.dtype, device=torch.cuda.current_device())
+
+        if pp_size > 1:
+            if info.src_rank in dist.get_process_group_ranks(mpu.get_pipeline_model_parallel_group()):
+                torch.distributed.broadcast(param, src=info.src_rank, group=mpu.get_pipeline_model_parallel_group())
+
+        # broadcast params across ep ranks
+        if ep_size > 1:
+            if ".experts." in info.name:
+                src_rank = (
+                    info.src_rank
+                    if info.src_rank in dist.get_process_group_ranks(mpu.get_expert_model_parallel_group())
+                    else rank
+                )
+                torch.distributed.broadcast(param, src=src_rank, group=mpu.get_expert_model_parallel_group())
+
+        for key, value in info.attrs.items():
+            setattr(param, key, value)
+
+        param = update_weight_utils.all_gather_param(info.name, param)
+        param = update_weight_utils.remove_padding(info.name, param, vocab_size)
+        # use torch.distributed
+        if is_save_rank:
+            converted_named_tensors = update_weight_utils.convert_to_hf(args, model_name, info.name, param)
+            for name, param in converted_named_tensors:
+                state_dict[name] = param.cpu()
+        del param
+
+    if is_save_rank:
+        hf_model = AutoModelForCausalLM.from_pretrained(
+            args.hf_checkpoint, torch_dtype="auto", device_map="cpu", trust_remote_code=True
+        )
+
+        if args.check_same:
+            for name, param in hf_model.named_parameters():
+                if name in state_dict:
+                    assert (
+                        param.shape == state_dict[name].shape
+                    ), f"Shape mismatch for {name}: {param.shape} vs {state_dict[name].shape}"
+                    assert torch.all(param == state_dict[name]), f"Value mismatch for {name}"
+                else:
+                    print(f"Warning: {name} not found in state_dict")
+
+        if args.output_dir:
+            tokenizer.save_pretrained(args.output_dir)
+            print(hf_model.load_state_dict(state_dict, strict=False))
+            hf_model.save_pretrained(args.output_dir)
+
+    dist.barrier()
+
+
+if __name__ == "__main__":
+    args = parse_args(add_custom_arguments=add_checkpoint_args)
+    main(args)

tools/convert_torch_dist_to_hf.py

@@ -0,0 +1,219 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import json
+import os
+import pickle
+import re
+import shutil
+import time
+
+
+import safetensors.torch
+import torch
+import torch.distributed.checkpoint as dist_cp
+from transformers import AutoConfig
+from typing_extensions import override
+
+from slime.backends.megatron_utils.megatron_to_hf import convert_to_hf, remove_padding
+
+
+class UnpicklerWrapper(pickle.Unpickler):
+    @override
+    def find_class(self, mod_name, name):
+        class DummyClass:
+            def __init__(self, *args, **kwargs):
+                pass
+
+        if mod_name.startswith("megatron") or mod_name.startswith("glm"):
+            return DummyClass
+        return super().find_class(mod_name, name)
+
+
+pickle.Unpickler = UnpicklerWrapper
+
+
+class WrappedStorageReader(dist_cp.FileSystemReader):
+    @override
+    def read_metadata(self):
+        path = self.fs.concat_path(self.path, ".metadata")
+        with self.fs.create_stream(path, "rb") as metadata_file:
+            metadata = UnpicklerWrapper(metadata_file).load()
+        if getattr(metadata, "storage_meta", None) is None:
+            metadata.storage_meta = dist_cp.StorageMeta()
+        metadata.storage_meta.load_id = self.load_id
+        if metadata.planner_data is None:
+            metadata.planner_data = {}
+        return metadata
+
+
+class EmptyStateDictLoadPlanner(dist_cp.default_planner.DefaultLoadPlanner):
+    @override
+    def set_up_planner(
+        self,
+        state_dict: dist_cp.metadata.STATE_DICT_TYPE,
+        metadata: dist_cp.metadata.Metadata | None = None,
+        is_coordinator: bool = False,
+    ) -> None:
+        for k, v in metadata.state_dict_metadata.items():
+            if "optimizer" in k or "_state" in k:
+                continue
+            print(f"find {k} in torch_dist ckpt")
+            if isinstance(v, dist_cp.metadata.TensorStorageMetadata):
+                v = torch.empty(v.size, dtype=v.properties.dtype)  # type: ignore[assignment]
+            state_dict[k] = v
+        super().set_up_planner(state_dict, metadata, is_coordinator)
+
+
+def get_expert_param(args, name, param):
+    if ".experts." not in name:
+        yield name, param
+        return
+
+    num_experts = args.num_experts
+    match = re.search(r"mlp.experts\.(.+)\.weight(\d+)", name)
+    if not match:
+        assert param.shape[0] == num_experts
+        for expert_id in range(num_experts):
+            expert_name = name.replace(".experts.experts.", ".experts.") + str(expert_id)
+            expert_param = param[expert_id]
+            yield expert_name, expert_param
+    else:
+        yield name, param
+
+
+def get_layer_param(args, name, param):
+    if ".layers." not in name:
+        yield name, param
+        return
+
+    num_layers = args.num_layers
+    match = re.search(r"\.layers\.(\d+)\.", name)
+    if not match:
+        assert param.shape[0] == num_layers
+        for layer_id in range(num_layers):
+            layer_name = name.replace(".layers.", f".layers.{layer_id}.")
+            layer_param = param[layer_id]
+            yield from get_expert_param(args, layer_name, layer_param)
+    else:
+        yield from get_expert_param(args, name, param)
+
+
+def get_named_params(args, state_dict):
+    for name, param in state_dict.items():
+        name = f"module.module.{name}"
+        yield from get_layer_param(args, name, param)
+
+
+def save_tensors(args, model_name, state_dict, output_dir, chunk_size, vocab_size=None):
+    # for slime update_weight compatible
+    args.sglang_enable_ep_moe = False
+
+    print(f"start saving to {output_dir}")
+    os.makedirs(output_dir, exist_ok=True)
+    # 2GB
+    current_size = 0
+    total_size = 0
+    modeltensors = [{}]
+    for name, param in get_named_params(args, state_dict):
+        if vocab_size:
+            param = remove_padding(name, param, vocab_size)
+        converted_named_tensors = convert_to_hf(args, model_name, name, param)
+        for converted_name, converted_param in converted_named_tensors:
+            tensor_size = converted_param.numel() * converted_param.element_size()
+            if tensor_size + current_size > chunk_size:
+                modeltensors.append({})
+                current_size = 0
+            modeltensors[-1][converted_name] = converted_param
+            current_size += tensor_size
+            total_size += tensor_size
+
+    metadata = {"metadata": {"total_size": total_size}, "weight_map": {}}
+
+    num_files = len(modeltensors)
+    for i, tensors in enumerate(modeltensors):
+        filename = f"model-{i:05d}-of-{num_files:05d}.safetensors"
+        for key in tensors.keys():
+            metadata["weight_map"][key] = filename
+    index_filepath = os.path.join(output_dir, "model.safetensors.index.json")
+    json.dump(metadata, open(index_filepath, "w"), indent=2)
+    print(f"{index_filepath} saved.")
+
+    for i, tensors in enumerate(modeltensors):
+        filename = f"model-{i:05d}-of-{num_files:05d}.safetensors"
+        t = time.time()
+        filepath = os.path.join(output_dir, filename)
+        safetensors.torch.save_file(tensors, filepath)
+        print(f"{filename} saved in {time.time() - t:.2f} sec.")
+
+
+def copy_assets(origin_hf_dir, output_dir):
+    for filename in os.listdir(origin_hf_dir):
+        if filename == "model.safetensors.index.json" or filename.endswith(".safetensors"):
+            continue
+        origin_filename = os.path.join(origin_hf_dir, filename)
+        if not os.path.isfile(origin_filename):
+            print(f"Skip {filename}, not a file.")
+            continue
+        src, dst = origin_filename, os.path.join(output_dir, filename)
+        print(f"copy from {src} to {dst}")
+        shutil.copy(src, dst)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-name", type=str, default=None)
+    parser.add_argument("--input-dir", type=str, required=True)
+    parser.add_argument("--output-dir", type=str, required=True)
+    parser.add_argument(
+        "--origin-hf-dir",
+        type=str,
+        default=None,
+        help="use the origin hf dir to copy files like tokenizer, config.json, etc.",
+    )
+    parser.add_argument(
+        "-f", "--force", action="store_true", help="Force overwrite the output directory if it exists."
+    )
+    parser.add_argument(
+        "--chunk-size",
+        type=int,
+        default=5 * 1024**3,
+        help="Chunk size for saving tensors, default is 2GB.",
+    )
+    parser.add_argument(
+        "--vocab-size",
+        type=int,
+        default=None,
+        help="Vocab size for removing padding, if applicable. If not provided, no padding will be removed.",
+    )
+    args = parser.parse_args()
+
+    if os.path.exists(args.output_dir) and not args.force:
+        raise ValueError(f"Output directory {args.output_dir} already exists. Use --force to overwrite it.")
+
+    if args.model_name is None and args.origin_hf_dir is None:
+        raise ValueError(
+            "Either --model-name or --origin-hf-dir must be provided, so that we can know the name of the params."
+        )
+
+    if args.model_name is None:
+        hf_config = AutoConfig.from_pretrained(args.origin_hf_dir, trust_remote_code=True)
+        args.model_name = type(hf_config).__name__.lower()
+
+    state_dict = {}
+    print(f"loading model from {args.input_dir}")
+    t = time.time()
+    megatron_args = torch.load(os.path.join(args.input_dir, "common.pt"), weights_only=False)["args"]
+    dist_cp.state_dict_loader._load_state_dict(
+        state_dict,
+        storage_reader=WrappedStorageReader(args.input_dir),
+        planner=EmptyStateDictLoadPlanner(),
+        no_dist=True,
+    )
+    print(f"model loaded in {time.time()-t:.2f} sec.")
+
+    save_tensors(megatron_args, args.model_name, state_dict, args.output_dir, args.chunk_size, args.vocab_size)
+
+    if args.origin_hf_dir:
+        copy_assets(args.origin_hf_dir, args.output_dir)

tools/eval_aime2024_vllm.py

@@ -0,0 +1,494 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Evaluate a HF-format model on AIME 2024 using vLLM.
+
+Example:
+
+CUDA_VISIBLE_DEVICES=0,1,2,3 python tools/eval_aime_vllm.py \
+    --model /workspace/Lightning-OPD/checkpoints/qwen3-4b-lightning-opd-hf \
+    --num-gpus 4 \
+    --output outputs/eval_aime2024_qwen3_4b_lightning_opd.jsonl
+
+Paper-style AIME setting:
+    temperature = 0.6
+    top_p = 0.95
+    max_tokens = 32768
+    n_samples = 32
+    metric = average pass@1
+"""
+
+import argparse
+import json
+import os
+import re
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    # Model / hardware
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Path to HF-format model checkpoint.",
+    )
+    parser.add_argument(
+        "--num-gpus",
+        type=int,
+        default=1,
+        help="Tensor parallel size for vLLM.",
+    )
+    parser.add_argument(
+        "--gpu-ids",
+        type=str,
+        default=None,
+        help='Optional CUDA_VISIBLE_DEVICES, e.g. "0,1,2,3". Must match --num-gpus.',
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bfloat16",
+        choices=["auto", "float16", "bfloat16", "float32"],
+    )
+    parser.add_argument(
+        "--gpu-memory-utilization",
+        type=float,
+        default=0.90,
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        default=True,
+    )
+
+    # Dataset
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="AI-MO/aimo-validation-aime",
+        help="HF dataset name for AIME 2024.",
+    )
+    parser.add_argument(
+        "--split",
+        type=str,
+        default=None,
+        help="Dataset split. If not set, use the first available split.",
+    )
+    parser.add_argument(
+        "--hf-cache",
+        type=str,
+        default=None,
+        help="Optional HuggingFace cache dir.",
+    )
+
+    # Generation settings from paper
+    parser.add_argument("--n-samples", type=int, default=32)
+    parser.add_argument("--temperature", type=float, default=0.6)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument("--max-tokens", type=int, default=32768)
+    parser.add_argument("--seed", type=int, default=42)
+
+    # Prompt / tokenizer behavior
+    parser.add_argument(
+        "--no-chat-template",
+        action="store_true",
+        help="Do not apply tokenizer chat template; use raw prompt.",
+    )
+    parser.add_argument(
+        "--enable-thinking",
+        action="store_true",
+        default=False,
+        help="Pass enable_thinking=True to Qwen3 chat template if supported.",
+    )
+    parser.add_argument(
+        "--disable-thinking",
+        action="store_true",
+        default=False,
+        help="Pass enable_thinking=False to Qwen3 chat template if supported.",
+    )
+    parser.add_argument(
+        "--assistant-prefix",
+        type=str,
+        default=None,
+        help=(
+            "Optional text appended after the assistant generation prompt. "
+            "Example for no-thinking style: '<think>\\n\\n</think>\\n\\n'"
+        ),
+    )
+    parser.add_argument(
+        "--prompt-template",
+        type=str,
+        default="train",
+        choices=["train", "paper", "chat"],
+        help=(
+            "Prompt template to use. "
+            "'train' uses the exact training ChatML prompt; "
+            "'paper' uses the paper-style ChatML prompt; "
+            "'chat' uses build_problem_prompt plus tokenizer.apply_chat_template."
+        ),
+    )
+
+    # Output / debug
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="outputs/eval_aime2024_vllm.jsonl",
+        help="Path to save per-sample generations and scores.",
+    )
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=None,
+        help="Limit number of problems for quick debugging.",
+    )
+    parser.add_argument(
+        "--print-examples",
+        type=int,
+        default=3,
+        help="Print first N examples with predictions.",
+    )
+
+    return parser.parse_args()
+
+
+def normalize_answer(x: Any) -> Optional[str]:
+    if x is None:
+        return None
+
+    s = str(x).strip()
+    s = s.replace("$", "")
+    s = s.replace(",", "")
+    s = s.replace("\\,", "")
+    s = s.strip()
+
+    # Remove simple latex wrappers.
+    s = s.replace("\\text", "")
+    s = s.replace("\\mathrm", "")
+    s = s.strip("{}").strip()
+
+    # AIME answers are integers from 0 to 999.
+    m = re.search(r"-?\d+", s)
+    if m is None:
+        return None
+
+    try:
+        return str(int(m.group(0)))
+    except ValueError:
+        return None
+
+
+def extract_last_boxed(text: str) -> Optional[str]:
+    """
+    Extract the last \\boxed{...}. Handles simple nested braces better than regex.
+    """
+    marker = r"\boxed{"
+    positions = [m.start() for m in re.finditer(re.escape(marker), text)]
+    if not positions:
+        return None
+
+    for start in reversed(positions):
+        i = start + len(marker)
+        depth = 1
+        chars = []
+
+        while i < len(text):
+            ch = text[i]
+            if ch == "{":
+                depth += 1
+                chars.append(ch)
+            elif ch == "}":
+                depth -= 1
+                if depth == 0:
+                    return "".join(chars)
+                chars.append(ch)
+            else:
+                chars.append(ch)
+            i += 1
+
+    return None
+
+
+def extract_answer(text: str) -> Optional[str]:
+    # Prefer boxed answer, because the actual training prompt asks for \boxed{$Answer}.
+    boxed = extract_last_boxed(text)
+    if boxed is not None:
+        ans = normalize_answer(boxed)
+        if ans is not None:
+            return ans
+
+    # Fallback: final Answer: line.
+    matches = re.findall(r"Answer:\s*([^\n]+)", text, flags=re.IGNORECASE)
+    if matches:
+        ans = normalize_answer(matches[-1])
+        if ans is not None:
+            return ans
+
+    # No valid final answer.
+    return None
+
+def build_exact_training_prompt(problem: str) -> str:
+    return (
+        "<|im_start|>user\n"
+        "Solve the following math problem step by step. "
+        "The last line of your response should be of the form "
+        "Answer: \\boxed{$Answer} where $Answer is the answer to the problem.\n\n"
+        f"{problem}\n\n"
+        "Remember to put your answer on its own line after \"Answer:\"."
+        "<|im_end|>\n"
+        "<|im_start|>assistant\n\n\n\n"
+    )
+
+def build_problem_prompt(problem: str) -> str:
+    return (
+        f"{problem}\n\n"
+        "Please reason step by step, and put your final answer within \\boxed{}."
+    )
+
+def build_paper_math_eval_prompt(problem: str) -> str:
+    return (
+        "<|im_start|>user\n"
+        f"Question: {problem}\n"
+        "Please reason step by step, and put your final answer within \\boxed{}.\n"
+        "<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+def get_field(ex: Dict[str, Any], candidates: List[str]) -> Any:
+    for key in candidates:
+        if key in ex and ex[key] is not None:
+            return ex[key]
+    raise KeyError(f"Cannot find any of {candidates}. Example keys: {list(ex.keys())}")
+
+# def build_problem_prompt(problem: str) -> str:
+#     return (
+#         "Solve the following math problem step by step. "
+#         "The last line of your response should be of the form "
+#         "Answer: $Answer (without quotes) where $Answer is the answer to the problem.\n\n"
+#         f"{problem}\n\n"
+#         'Remember to put your answer on its own line after "Answer:".'
+#     )
+
+def apply_chat_template(tokenizer, prompt: str, args) -> str:
+    if args.no_chat_template:
+        return prompt
+
+    messages = [{"role": "user", "content": prompt}]
+
+    kwargs = {
+        "tokenize": False,
+        "add_generation_prompt": True,
+    }
+
+    # Qwen3 tokenizer may support enable_thinking.
+    if args.enable_thinking and args.disable_thinking:
+        raise ValueError("Do not set both --enable-thinking and --disable-thinking.")
+
+    if args.enable_thinking:
+        kwargs["enable_thinking"] = True
+    elif args.disable_thinking:
+        kwargs["enable_thinking"] = False
+
+    try:
+        text = tokenizer.apply_chat_template(messages, **kwargs)
+    except TypeError:
+        # Some tokenizers do not accept enable_thinking.
+        kwargs.pop("enable_thinking", None)
+        text = tokenizer.apply_chat_template(messages, **kwargs)
+
+    if args.assistant_prefix is not None:
+        text += args.assistant_prefix.encode("utf-8").decode("unicode_escape")
+
+    return text
+
+
+def main():
+    args = parse_args()
+
+    if args.gpu_ids is not None:
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
+
+    if args.hf_cache is not None:
+        os.environ["HF_HOME"] = args.hf_cache
+        os.environ["HF_DATASETS_CACHE"] = str(Path(args.hf_cache) / "datasets")
+        os.environ["HF_HUB_CACHE"] = str(Path(args.hf_cache) / "hub")
+
+    # Import after CUDA_VISIBLE_DEVICES is set.
+    from datasets import load_dataset
+    from transformers import AutoTokenizer
+    from vllm import LLM, SamplingParams
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    print("=" * 80)
+    print("AIME 2024 vLLM Evaluation")
+    print("=" * 80)
+    print(f"model:                  {args.model}")
+    print(f"dataset:                {args.dataset}")
+    print(f"num_gpus / TP size:     {args.num_gpus}")
+    print(f"CUDA_VISIBLE_DEVICES:   {os.environ.get('CUDA_VISIBLE_DEVICES')}")
+    print(f"n_samples/problem:      {args.n_samples}")
+    print(f"temperature:            {args.temperature}")
+    print(f"top_p:                  {args.top_p}")
+    print(f"max_tokens:             {args.max_tokens}")
+    print(f"prompt_template:        {args.prompt_template}")
+    print(f"assistant_prefix:       {repr(args.assistant_prefix)}")
+    print(f"output:                 {args.output}")
+    print("=" * 80)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model,
+        trust_remote_code=args.trust_remote_code,
+    )
+
+    dataset_dict = load_dataset(args.dataset)
+    split = args.split or list(dataset_dict.keys())[0]
+    data = dataset_dict[split]
+
+    # AI-MO/aimo-validation-aime contains multiple AIME years.
+    # Keep only AIME 2024 examples.
+    if "url" in data.column_names:
+        data = data.filter(lambda ex: "2024_AIME" in ex["url"])
+    else:
+        raise ValueError(
+            "Expected a `url` column for filtering AIME 2024, "
+            f"but got columns: {data.column_names}"
+        )
+
+    print(f"Filtered AIME 2024 examples: {len(data)}")
+    for i in range(min(3, len(data))):
+        print(i, data[i].get("url", "NO_URL"), data[i].get("answer", "NO_ANSWER"))
+
+    assert len(data) == 30, f"Expected 30 AIME 2024 problems, got {len(data)}"
+
+    if args.limit is not None:
+        data = data.select(range(min(args.limit, len(data))))
+
+    print(f"Loaded split: {split}")
+    print(f"Number of problems: {len(data)}")
+    print(f"First example keys: {list(data[0].keys())}")
+
+    prompts = []
+    examples = []
+
+    for idx, ex in enumerate(data):
+        problem = get_field(ex, ["problem", "question", "prompt"])
+        gold_raw = get_field(ex, ["answer", "final_answer", "target", "solution"])
+        gold = normalize_answer(gold_raw)
+
+        if args.prompt_template == "train":
+            # Raw ChatML prompt matching the OPD training parquet.
+            # Do NOT call apply_chat_template again, otherwise ChatML will be nested.
+            full_prompt = build_exact_training_prompt(problem)
+        elif args.prompt_template == "paper":
+            # Raw ChatML prompt matching the paper-style math evaluation prompt.
+            # Do NOT call apply_chat_template again.
+            full_prompt = build_paper_math_eval_prompt(problem)
+        elif args.prompt_template == "chat":
+            # Normal user-content prompt; tokenizer will add ChatML.
+            raw_prompt = build_problem_prompt(problem)
+            full_prompt = apply_chat_template(tokenizer, raw_prompt, args)
+        else:
+            raise ValueError(f"Unknown prompt_template: {args.prompt_template}")
+
+        # Important: when using raw ChatML prompts, apply_chat_template() is bypassed.
+        # Therefore --assistant-prefix must be appended here, not only inside apply_chat_template().
+        # This is useful for Qwen3 no-thinking style, e.g.:
+        #   --assistant-prefix '<think>\\n\\n</think>\\n\\n'
+        if args.prompt_template in {"train", "paper"} and args.assistant_prefix is not None:
+            full_prompt += args.assistant_prefix.encode("utf-8").decode("unicode_escape")
+
+        prompts.append(full_prompt)
+        examples.append(
+            {
+                "idx": idx,
+                "problem": problem,
+                "gold_raw": gold_raw,
+                "gold": gold,
+                "prompt": full_prompt,
+            }
+        )
+
+    llm = LLM(
+        model=args.model,
+        tensor_parallel_size=args.num_gpus,
+        dtype=args.dtype,
+        trust_remote_code=args.trust_remote_code,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        seed=args.seed,
+    )
+
+    sampling_params = SamplingParams(
+        n=args.n_samples,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        #repetition_penalty=1.05,
+        max_tokens=args.max_tokens,
+        stop=["<|im_end|>"],
+    )
+
+    outputs = llm.generate(prompts, sampling_params)
+
+    total = 0
+    correct = 0
+    per_problem_records = []
+
+    with output_path.open("w", encoding="utf-8") as f:
+        for ex, out in zip(examples, outputs):
+            sample_records = []
+            problem_correct = 0
+
+            for sample_id, completion in enumerate(out.outputs):
+                text = completion.text
+                pred = extract_answer(text)
+                is_correct = pred == ex["gold"]
+
+                total += 1
+                correct += int(is_correct)
+                problem_correct += int(is_correct)
+
+                record = {
+                    "idx": ex["idx"],
+                    "sample_id": sample_id,
+                    "gold": ex["gold"],
+                    "gold_raw": str(ex["gold_raw"]),
+                    "pred": pred,
+                    "correct": is_correct,
+                    "completion": text,
+                    "finish_reason": completion.finish_reason,
+                    "stop_reason": getattr(completion, "stop_reason", None),
+                }
+                sample_records.append(record)
+                f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+            problem_acc = problem_correct / max(1, len(out.outputs))
+            per_problem_records.append(problem_acc)
+
+            if ex["idx"] < args.print_examples:
+                print("-" * 80)
+                print(f"Problem {ex['idx']}")
+                print(f"Gold: {ex['gold']} | Correct samples: {problem_correct}/{len(out.outputs)}")
+                print(f"First pred: {sample_records[0]['pred']}")
+                print(f"First completion preview:\n{sample_records[0]['completion'][:1000]}")
+
+    avg_pass1_micro = correct / total if total > 0 else 0.0
+    avg_pass1_macro = sum(per_problem_records) / len(per_problem_records) if per_problem_records else 0.0
+
+    print("=" * 80)
+    print("Final results")
+    print("=" * 80)
+    print(f"Problems:                 {len(examples)}")
+    print(f"Samples per problem:       {args.n_samples}")
+    print(f"Total samples:             {total}")
+    print(f"Correct samples:           {correct}")
+    print(f"Average pass@1 micro:      {avg_pass1_micro:.6f} ({100 * avg_pass1_micro:.2f}%)")
+    print(f"Average pass@1 macro:      {avg_pass1_macro:.6f} ({100 * avg_pass1_macro:.2f}%)")
+    print(f"Saved generations to:      {output_path}")
+    print("=" * 80)
+
+
+if __name__ == "__main__":
+    main()

tools/eval_aime2025_vllm.py

@@ -0,0 +1,488 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Evaluate a HF-format model on AIME 2025 using vLLM.
+
+Example:
+
+CUDA_VISIBLE_DEVICES=0,1,2,3 python tools/eval_aime_vllm.py \
+    --model /workspace/Lightning-OPD/checkpoints/qwen3-4b-lightning-opd-hf \
+    --num-gpus 4 \
+    --output outputs/eval_aime_2025_qwen3_4b_lightning_opd.jsonl
+
+Paper-style AIME setting:
+    temperature = 0.6
+    top_p = 0.95
+    max_tokens = 32768
+    n_samples = 32
+    metric = average pass@1
+"""
+
+import argparse
+import json
+import os
+import re
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    # Model / hardware
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Path to HF-format model checkpoint.",
+    )
+    parser.add_argument(
+        "--num-gpus",
+        type=int,
+        default=1,
+        help="Tensor parallel size for vLLM.",
+    )
+    parser.add_argument(
+        "--gpu-ids",
+        type=str,
+        default=None,
+        help='Optional CUDA_VISIBLE_DEVICES, e.g. "0,1,2,3". Must match --num-gpus.',
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bfloat16",
+        choices=["auto", "float16", "bfloat16", "float32"],
+    )
+    parser.add_argument(
+        "--gpu-memory-utilization",
+        type=float,
+        default=0.90,
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        default=True,
+    )
+
+    # Dataset
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="MathArena/aime_2025",
+        help="HF dataset name for AIME 2025.",
+    )
+    parser.add_argument(
+        "--split",
+        type=str,
+        default=None,
+        help="Dataset split. If not set, use the first available split.",
+    )
+    parser.add_argument(
+        "--hf-cache",
+        type=str,
+        default=None,
+        help="Optional HuggingFace cache dir.",
+    )
+
+    # Generation settings from paper
+    parser.add_argument("--n-samples", type=int, default=32)
+    parser.add_argument("--temperature", type=float, default=0.6)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument("--max-tokens", type=int, default=32768)
+    parser.add_argument("--seed", type=int, default=42)
+
+    # Prompt / tokenizer behavior
+    parser.add_argument(
+        "--no-chat-template",
+        action="store_true",
+        help="Do not apply tokenizer chat template; use raw prompt.",
+    )
+    parser.add_argument(
+        "--enable-thinking",
+        action="store_true",
+        default=False,
+        help="Pass enable_thinking=True to Qwen3 chat template if supported.",
+    )
+    parser.add_argument(
+        "--disable-thinking",
+        action="store_true",
+        default=False,
+        help="Pass enable_thinking=False to Qwen3 chat template if supported.",
+    )
+    parser.add_argument(
+        "--assistant-prefix",
+        type=str,
+        default=None,
+        help=(
+            "Optional text appended after the assistant generation prompt. "
+            "Example for no-thinking style: '<think>\\n\\n</think>\\n\\n'"
+        ),
+    )
+    parser.add_argument(
+        "--prompt-template",
+        type=str,
+        default="train",
+        choices=["train", "paper", "chat"],
+        help=(
+            "Prompt template to use. "
+            "'train' uses the exact training ChatML prompt; "
+            "'paper' uses the paper-style ChatML prompt; "
+            "'chat' uses build_problem_prompt plus tokenizer.apply_chat_template."
+        ),
+    )
+
+    # Output / debug
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="outputs/eval_aime_2025_vllm.jsonl",
+        help="Path to save per-sample generations and scores.",
+    )
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=None,
+        help="Limit number of problems for quick debugging.",
+    )
+    parser.add_argument(
+        "--print-examples",
+        type=int,
+        default=3,
+        help="Print first N examples with predictions.",
+    )
+
+    return parser.parse_args()
+
+
+def normalize_answer(x: Any) -> Optional[str]:
+    if x is None:
+        return None
+
+    s = str(x).strip()
+    s = s.replace("$", "")
+    s = s.replace(",", "")
+    s = s.replace("\\,", "")
+    s = s.strip()
+
+    # Remove simple latex wrappers.
+    s = s.replace("\\text", "")
+    s = s.replace("\\mathrm", "")
+    s = s.strip("{}").strip()
+
+    # AIME answers are integers from 0 to 999.
+    m = re.search(r"-?\d+", s)
+    if m is None:
+        return None
+
+    try:
+        return str(int(m.group(0)))
+    except ValueError:
+        return None
+
+
+def extract_last_boxed(text: str) -> Optional[str]:
+    """
+    Extract the last \\boxed{...}. Handles simple nested braces better than regex.
+    """
+    marker = r"\boxed{"
+    positions = [m.start() for m in re.finditer(re.escape(marker), text)]
+    if not positions:
+        return None
+
+    for start in reversed(positions):
+        i = start + len(marker)
+        depth = 1
+        chars = []
+
+        while i < len(text):
+            ch = text[i]
+            if ch == "{":
+                depth += 1
+                chars.append(ch)
+            elif ch == "}":
+                depth -= 1
+                if depth == 0:
+                    return "".join(chars)
+                chars.append(ch)
+            else:
+                chars.append(ch)
+            i += 1
+
+    return None
+
+
+def extract_answer(text: str) -> Optional[str]:
+    # Prefer boxed answer, because the actual training prompt asks for \boxed{$Answer}.
+    boxed = extract_last_boxed(text)
+    if boxed is not None:
+        ans = normalize_answer(boxed)
+        if ans is not None:
+            return ans
+
+    # Fallback: final Answer: line.
+    matches = re.findall(r"Answer:\s*([^\n]+)", text, flags=re.IGNORECASE)
+    if matches:
+        ans = normalize_answer(matches[-1])
+        if ans is not None:
+            return ans
+
+    # No valid final answer.
+    return None
+
+def build_exact_training_prompt(problem: str) -> str:
+    return (
+        "<|im_start|>user\n"
+        "Solve the following math problem step by step. "
+        "The last line of your response should be of the form "
+        "Answer: \\boxed{$Answer} where $Answer is the answer to the problem.\n\n"
+        f"{problem}\n\n"
+        "Remember to put your answer on its own line after \"Answer:\"."
+        "<|im_end|>\n"
+        "<|im_start|>assistant\n\n\n\n"
+    )
+
+def build_problem_prompt(problem: str) -> str:
+    return (
+        f"{problem}\n\n"
+        "Please reason step by step, and put your final answer within \\boxed{}."
+    )
+
+def build_paper_math_eval_prompt(problem: str) -> str:
+    return (
+        "<|im_start|>user\n"
+        f"Question: {problem}\n"
+        "Please reason step by step, and put your final answer within \\boxed{}.\n"
+        "<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+def get_field(ex: Dict[str, Any], candidates: List[str]) -> Any:
+    for key in candidates:
+        if key in ex and ex[key] is not None:
+            return ex[key]
+    raise KeyError(f"Cannot find any of {candidates}. Example keys: {list(ex.keys())}")
+
+# def build_problem_prompt(problem: str) -> str:
+#     return (
+#         "Solve the following math problem step by step. "
+#         "The last line of your response should be of the form "
+#         "Answer: $Answer (without quotes) where $Answer is the answer to the problem.\n\n"
+#         f"{problem}\n\n"
+#         'Remember to put your answer on its own line after "Answer:".'
+#     )
+
+def apply_chat_template(tokenizer, prompt: str, args) -> str:
+    if args.no_chat_template:
+        return prompt
+
+    messages = [{"role": "user", "content": prompt}]
+
+    kwargs = {
+        "tokenize": False,
+        "add_generation_prompt": True,
+    }
+
+    # Qwen3 tokenizer may support enable_thinking.
+    if args.enable_thinking and args.disable_thinking:
+        raise ValueError("Do not set both --enable-thinking and --disable-thinking.")
+
+    if args.enable_thinking:
+        kwargs["enable_thinking"] = True
+    elif args.disable_thinking:
+        kwargs["enable_thinking"] = False
+
+    try:
+        text = tokenizer.apply_chat_template(messages, **kwargs)
+    except TypeError:
+        # Some tokenizers do not accept enable_thinking.
+        kwargs.pop("enable_thinking", None)
+        text = tokenizer.apply_chat_template(messages, **kwargs)
+
+    if args.assistant_prefix is not None:
+        text += args.assistant_prefix.encode("utf-8").decode("unicode_escape")
+
+    return text
+
+
+def main():
+    args = parse_args()
+
+    if args.gpu_ids is not None:
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
+
+    if args.hf_cache is not None:
+        os.environ["HF_HOME"] = args.hf_cache
+        os.environ["HF_DATASETS_CACHE"] = str(Path(args.hf_cache) / "datasets")
+        os.environ["HF_HUB_CACHE"] = str(Path(args.hf_cache) / "hub")
+
+    # Import after CUDA_VISIBLE_DEVICES is set.
+    from datasets import load_dataset
+    from transformers import AutoTokenizer
+    from vllm import LLM, SamplingParams
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    print("=" * 80)
+    print("AIME 2025 vLLM Evaluation")
+    print("=" * 80)
+    print(f"model:                  {args.model}")
+    print(f"dataset:                {args.dataset}")
+    print(f"num_gpus / TP size:     {args.num_gpus}")
+    print(f"CUDA_VISIBLE_DEVICES:   {os.environ.get('CUDA_VISIBLE_DEVICES')}")
+    print(f"n_samples/problem:      {args.n_samples}")
+    print(f"temperature:            {args.temperature}")
+    print(f"top_p:                  {args.top_p}")
+    print(f"max_tokens:             {args.max_tokens}")
+    print(f"prompt_template:        {args.prompt_template}")
+    print(f"assistant_prefix:       {repr(args.assistant_prefix)}")
+    print(f"output:                 {args.output}")
+    print("=" * 80)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model,
+        trust_remote_code=args.trust_remote_code,
+    )
+
+    # Load HF dataset or local JSON/JSONL.
+    if args.dataset.endswith(".json") or args.dataset.endswith(".jsonl"):
+        dataset_dict = load_dataset("json", data_files=args.dataset)
+    else:
+        dataset_dict = load_dataset(args.dataset)
+
+    split = args.split or list(dataset_dict.keys())[0]
+    data = dataset_dict[split]
+
+    assert len(data) == 30, f"Expected 30 AIME 2025 problems, got {len(data)}"
+
+    print(f"Loaded AIME 2025 examples: {len(data)}")
+    for i in range(min(3, len(data))):
+        print(i, {k: data[i].get(k, None) for k in data.column_names[:6]})
+
+    if args.limit is not None:
+        data = data.select(range(min(args.limit, len(data))))
+
+    print(f"Loaded split: {split}")
+    print(f"Number of problems: {len(data)}")
+    print(f"First example keys: {list(data[0].keys())}")
+
+    prompts = []
+    examples = []
+
+    for idx, ex in enumerate(data):
+        problem = get_field(ex, ["problem", "question", "prompt"])
+        gold_raw = get_field(ex, ["answer", "final_answer", "target", "solution"])
+        gold = normalize_answer(gold_raw)
+
+        if args.prompt_template == "train":
+            # Raw ChatML prompt matching the OPD training parquet.
+            # Do NOT call apply_chat_template again, otherwise ChatML will be nested.
+            full_prompt = build_exact_training_prompt(problem)
+        elif args.prompt_template == "paper":
+            # Raw ChatML prompt matching the paper-style math evaluation prompt.
+            # Do NOT call apply_chat_template again.
+            full_prompt = build_paper_math_eval_prompt(problem)
+        elif args.prompt_template == "chat":
+            # Normal user-content prompt; tokenizer will add ChatML.
+            raw_prompt = build_problem_prompt(problem)
+            full_prompt = apply_chat_template(tokenizer, raw_prompt, args)
+        else:
+            raise ValueError(f"Unknown prompt_template: {args.prompt_template}")
+
+        # Important: when using raw ChatML prompts, apply_chat_template() is bypassed.
+        # Therefore --assistant-prefix must be appended here, not only inside apply_chat_template().
+        # This is useful for Qwen3 no-thinking style, e.g.:
+        #   --assistant-prefix '<think>\\n\\n</think>\\n\\n'
+        if args.prompt_template in {"train", "paper"} and args.assistant_prefix is not None:
+            full_prompt += args.assistant_prefix.encode("utf-8").decode("unicode_escape")
+
+        prompts.append(full_prompt)
+        examples.append(
+            {
+                "idx": idx,
+                "problem": problem,
+                "gold_raw": gold_raw,
+                "gold": gold,
+                "prompt": full_prompt,
+            }
+        )
+
+    llm = LLM(
+        model=args.model,
+        tensor_parallel_size=args.num_gpus,
+        dtype=args.dtype,
+        trust_remote_code=args.trust_remote_code,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        seed=args.seed,
+    )
+
+    sampling_params = SamplingParams(
+        n=args.n_samples,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        max_tokens=args.max_tokens,
+        stop=["<|im_end|>"],
+    )
+
+    outputs = llm.generate(prompts, sampling_params)
+
+    total = 0
+    correct = 0
+    per_problem_records = []
+
+    with output_path.open("w", encoding="utf-8") as f:
+        for ex, out in zip(examples, outputs):
+            sample_records = []
+            problem_correct = 0
+
+            for sample_id, completion in enumerate(out.outputs):
+                text = completion.text
+                pred = extract_answer(text)
+                is_correct = pred == ex["gold"]
+
+                total += 1
+                correct += int(is_correct)
+                problem_correct += int(is_correct)
+
+                record = {
+                    "idx": ex["idx"],
+                    "sample_id": sample_id,
+                    "gold": ex["gold"],
+                    "gold_raw": str(ex["gold_raw"]),
+                    "pred": pred,
+                    "correct": is_correct,
+                    "completion": text,
+                    "finish_reason": completion.finish_reason,
+                    "stop_reason": getattr(completion, "stop_reason", None),
+                }
+                sample_records.append(record)
+                f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+            problem_acc = problem_correct / max(1, len(out.outputs))
+            per_problem_records.append(problem_acc)
+
+            if ex["idx"] < args.print_examples:
+                print("-" * 80)
+                print(f"Problem {ex['idx']}")
+                print(f"Gold: {ex['gold']} | Correct samples: {problem_correct}/{len(out.outputs)}")
+                print(f"First pred: {sample_records[0]['pred']}")
+                print(f"First completion preview:\n{sample_records[0]['completion'][:1000]}")
+
+    avg_pass1_micro = correct / total if total > 0 else 0.0
+    avg_pass1_macro = sum(per_problem_records) / len(per_problem_records) if per_problem_records else 0.0
+
+    print("=" * 80)
+    print("Final results")
+    print("=" * 80)
+    print(f"Problems:                 {len(examples)}")
+    print(f"Samples per problem:       {args.n_samples}")
+    print(f"Total samples:             {total}")
+    print(f"Correct samples:           {correct}")
+    print(f"Average pass@1 micro:      {avg_pass1_micro:.6f} ({100 * avg_pass1_micro:.2f}%)")
+    print(f"Average pass@1 macro:      {avg_pass1_macro:.6f} ({100 * avg_pass1_macro:.2f}%)")
+    print(f"Saved generations to:      {output_path}")
+    print("=" * 80)
+
+
+if __name__ == "__main__":
+    main()

tools/eval_hmmt2025_vllm.py

@@ -0,0 +1,486 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Evaluate a HF-format model on HMMT Feb 2025 using vLLM.
+
+Example:
+
+CUDA_VISIBLE_DEVICES=0,1,2,3 python tools/eval_aime_vllm.py \
+    --model /workspace/Lightning-OPD/checkpoints/qwen3-4b-lightning-opd-hf \
+    --num-gpus 4 \
+    --output outputs/eval_hmmt_feb_2025_qwen3_4b_lightning_opd.jsonl
+
+Paper-style AIME setting:
+    temperature = 0.6
+    top_p = 0.95
+    max_tokens = 32768
+    n_samples = 32
+    metric = average pass@1
+"""
+
+import argparse
+import json
+import os
+import re
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    # Model / hardware
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Path to HF-format model checkpoint.",
+    )
+    parser.add_argument(
+        "--num-gpus",
+        type=int,
+        default=1,
+        help="Tensor parallel size for vLLM.",
+    )
+    parser.add_argument(
+        "--gpu-ids",
+        type=str,
+        default=None,
+        help='Optional CUDA_VISIBLE_DEVICES, e.g. "0,1,2,3". Must match --num-gpus.',
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bfloat16",
+        choices=["auto", "float16", "bfloat16", "float32"],
+    )
+    parser.add_argument(
+        "--gpu-memory-utilization",
+        type=float,
+        default=0.90,
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        default=True,
+    )
+
+    # Dataset
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="MathArena/hmmt_feb_2025",
+        help="HF dataset name for HMMT February 2025.",
+    )
+    parser.add_argument(
+        "--split",
+        type=str,
+        default=None,
+        help="Dataset split. If not set, use the first available split.",
+    )
+    parser.add_argument(
+        "--hf-cache",
+        type=str,
+        default=None,
+        help="Optional HuggingFace cache dir.",
+    )
+
+    # Generation settings from paper
+    parser.add_argument("--n-samples", type=int, default=32)
+    parser.add_argument("--temperature", type=float, default=0.6)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument("--max-tokens", type=int, default=32768)
+    parser.add_argument("--seed", type=int, default=42)
+
+    # Prompt / tokenizer behavior
+    parser.add_argument(
+        "--no-chat-template",
+        action="store_true",
+        help="Do not apply tokenizer chat template; use raw prompt.",
+    )
+    parser.add_argument(
+        "--enable-thinking",
+        action="store_true",
+        default=False,
+        help="Pass enable_thinking=True to Qwen3 chat template if supported.",
+    )
+    parser.add_argument(
+        "--disable-thinking",
+        action="store_true",
+        default=False,
+        help="Pass enable_thinking=False to Qwen3 chat template if supported.",
+    )
+    parser.add_argument(
+        "--assistant-prefix",
+        type=str,
+        default=None,
+        help=(
+            "Optional text appended after the assistant generation prompt. "
+            "Example for no-thinking style: '<think>\\n\\n</think>\\n\\n'"
+        ),
+    )
+    parser.add_argument(
+        "--prompt-template",
+        type=str,
+        default="train",
+        choices=["train", "paper", "chat"],
+        help=(
+            "Prompt template to use. "
+            "'train' uses the exact training ChatML prompt; "
+            "'paper' uses the paper-style ChatML prompt; "
+            "'chat' uses build_problem_prompt plus tokenizer.apply_chat_template."
+        ),
+    )
+
+    # Output / debug
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="outputs/eval_hmmt_feb_2025_vllm.jsonl",
+        help="Path to save per-sample generations and scores.",
+    )
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=None,
+        help="Limit number of problems for quick debugging.",
+    )
+    parser.add_argument(
+        "--print-examples",
+        type=int,
+        default=3,
+        help="Print first N examples with predictions.",
+    )
+
+    return parser.parse_args()
+
+
+def normalize_answer(x: Any) -> Optional[str]:
+    if x is None:
+        return None
+
+    s = str(x).strip()
+    s = s.replace("$", "")
+    s = s.replace(",", "")
+    s = s.replace("\\,", "")
+    s = s.strip()
+
+    # Remove simple latex wrappers.
+    s = s.replace("\\text", "")
+    s = s.replace("\\mathrm", "")
+    s = s.strip("{}").strip()
+
+    # AIME answers are integers from 0 to 999.
+    m = re.search(r"-?\d+", s)
+    if m is None:
+        return None
+
+    try:
+        return str(int(m.group(0)))
+    except ValueError:
+        return None
+
+
+def extract_last_boxed(text: str) -> Optional[str]:
+    """
+    Extract the last \\boxed{...}. Handles simple nested braces better than regex.
+    """
+    marker = r"\boxed{"
+    positions = [m.start() for m in re.finditer(re.escape(marker), text)]
+    if not positions:
+        return None
+
+    for start in reversed(positions):
+        i = start + len(marker)
+        depth = 1
+        chars = []
+
+        while i < len(text):
+            ch = text[i]
+            if ch == "{":
+                depth += 1
+                chars.append(ch)
+            elif ch == "}":
+                depth -= 1
+                if depth == 0:
+                    return "".join(chars)
+                chars.append(ch)
+            else:
+                chars.append(ch)
+            i += 1
+
+    return None
+
+
+def extract_answer(text: str) -> Optional[str]:
+    # Prefer boxed answer, because the actual training prompt asks for \boxed{$Answer}.
+    boxed = extract_last_boxed(text)
+    if boxed is not None:
+        ans = normalize_answer(boxed)
+        if ans is not None:
+            return ans
+
+    # Fallback: final Answer: line.
+    matches = re.findall(r"Answer:\s*([^\n]+)", text, flags=re.IGNORECASE)
+    if matches:
+        ans = normalize_answer(matches[-1])
+        if ans is not None:
+            return ans
+
+    # No valid final answer.
+    return None
+
+def build_exact_training_prompt(problem: str) -> str:
+    return (
+        "<|im_start|>user\n"
+        "Solve the following math problem step by step. "
+        "The last line of your response should be of the form "
+        "Answer: \\boxed{$Answer} where $Answer is the answer to the problem.\n\n"
+        f"{problem}\n\n"
+        "Remember to put your answer on its own line after \"Answer:\"."
+        "<|im_end|>\n"
+        "<|im_start|>assistant\n\n\n\n"
+    )
+
+def build_problem_prompt(problem: str) -> str:
+    return (
+        f"{problem}\n\n"
+        "Please reason step by step, and put your final answer within \\boxed{}."
+    )
+
+def build_paper_math_eval_prompt(problem: str) -> str:
+    return (
+        "<|im_start|>user\n"
+        f"Question: {problem}\n"
+        "Please reason step by step, and put your final answer within \\boxed{}.\n"
+        "<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+def get_field(ex: Dict[str, Any], candidates: List[str]) -> Any:
+    for key in candidates:
+        if key in ex and ex[key] is not None:
+            return ex[key]
+    raise KeyError(f"Cannot find any of {candidates}. Example keys: {list(ex.keys())}")
+
+# def build_problem_prompt(problem: str) -> str:
+#     return (
+#         "Solve the following math problem step by step. "
+#         "The last line of your response should be of the form "
+#         "Answer: $Answer (without quotes) where $Answer is the answer to the problem.\n\n"
+#         f"{problem}\n\n"
+#         'Remember to put your answer on its own line after "Answer:".'
+#     )
+
+def apply_chat_template(tokenizer, prompt: str, args) -> str:
+    if args.no_chat_template:
+        return prompt
+
+    messages = [{"role": "user", "content": prompt}]
+
+    kwargs = {
+        "tokenize": False,
+        "add_generation_prompt": True,
+    }
+
+    # Qwen3 tokenizer may support enable_thinking.
+    if args.enable_thinking and args.disable_thinking:
+        raise ValueError("Do not set both --enable-thinking and --disable-thinking.")
+
+    if args.enable_thinking:
+        kwargs["enable_thinking"] = True
+    elif args.disable_thinking:
+        kwargs["enable_thinking"] = False
+
+    try:
+        text = tokenizer.apply_chat_template(messages, **kwargs)
+    except TypeError:
+        # Some tokenizers do not accept enable_thinking.
+        kwargs.pop("enable_thinking", None)
+        text = tokenizer.apply_chat_template(messages, **kwargs)
+
+    if args.assistant_prefix is not None:
+        text += args.assistant_prefix.encode("utf-8").decode("unicode_escape")
+
+    return text
+
+
+def main():
+    args = parse_args()
+
+    if args.gpu_ids is not None:
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
+
+    if args.hf_cache is not None:
+        os.environ["HF_HOME"] = args.hf_cache
+        os.environ["HF_DATASETS_CACHE"] = str(Path(args.hf_cache) / "datasets")
+        os.environ["HF_HUB_CACHE"] = str(Path(args.hf_cache) / "hub")
+
+    # Import after CUDA_VISIBLE_DEVICES is set.
+    from datasets import load_dataset
+    from transformers import AutoTokenizer
+    from vllm import LLM, SamplingParams
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    print("=" * 80)
+    print("HMMT Feb 2025 vLLM Evaluation")
+    print("=" * 80)
+    print(f"model:                  {args.model}")
+    print(f"dataset:                {args.dataset}")
+    print(f"num_gpus / TP size:     {args.num_gpus}")
+    print(f"CUDA_VISIBLE_DEVICES:   {os.environ.get('CUDA_VISIBLE_DEVICES')}")
+    print(f"n_samples/problem:      {args.n_samples}")
+    print(f"temperature:            {args.temperature}")
+    print(f"top_p:                  {args.top_p}")
+    print(f"max_tokens:             {args.max_tokens}")
+    print(f"prompt_template:        {args.prompt_template}")
+    print(f"assistant_prefix:       {repr(args.assistant_prefix)}")
+    print(f"output:                 {args.output}")
+    print("=" * 80)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model,
+        trust_remote_code=args.trust_remote_code,
+    )
+
+    # Load HF dataset or local JSON/JSONL.
+    if args.dataset.endswith(".json") or args.dataset.endswith(".jsonl"):
+        dataset_dict = load_dataset("json", data_files=args.dataset)
+    else:
+        dataset_dict = load_dataset(args.dataset)
+
+    split = args.split or list(dataset_dict.keys())[0]
+    data = dataset_dict[split]
+
+    print(f"Loaded HMMT Feb 2025 examples: {len(data)}")
+    for i in range(min(3, len(data))):
+        print(i, {k: data[i].get(k, None) for k in data.column_names[:6]})
+
+    if args.limit is not None:
+        data = data.select(range(min(args.limit, len(data))))
+
+    print(f"Loaded split: {split}")
+    print(f"Number of problems: {len(data)}")
+    print(f"First example keys: {list(data[0].keys())}")
+
+    prompts = []
+    examples = []
+
+    for idx, ex in enumerate(data):
+        problem = get_field(ex, ["problem", "question", "prompt"])
+        gold_raw = get_field(ex, ["answer", "final_answer", "target", "solution"])
+        gold = normalize_answer(gold_raw)
+
+        if args.prompt_template == "train":
+            # Raw ChatML prompt matching the OPD training parquet.
+            # Do NOT call apply_chat_template again, otherwise ChatML will be nested.
+            full_prompt = build_exact_training_prompt(problem)
+        elif args.prompt_template == "paper":
+            # Raw ChatML prompt matching the paper-style math evaluation prompt.
+            # Do NOT call apply_chat_template again.
+            full_prompt = build_paper_math_eval_prompt(problem)
+        elif args.prompt_template == "chat":
+            # Normal user-content prompt; tokenizer will add ChatML.
+            raw_prompt = build_problem_prompt(problem)
+            full_prompt = apply_chat_template(tokenizer, raw_prompt, args)
+        else:
+            raise ValueError(f"Unknown prompt_template: {args.prompt_template}")
+
+        # Important: when using raw ChatML prompts, apply_chat_template() is bypassed.
+        # Therefore --assistant-prefix must be appended here, not only inside apply_chat_template().
+        # This is useful for Qwen3 no-thinking style, e.g.:
+        #   --assistant-prefix '<think>\\n\\n</think>\\n\\n'
+        if args.prompt_template in {"train", "paper"} and args.assistant_prefix is not None:
+            full_prompt += args.assistant_prefix.encode("utf-8").decode("unicode_escape")
+
+        prompts.append(full_prompt)
+        examples.append(
+            {
+                "idx": idx,
+                "problem": problem,
+                "gold_raw": gold_raw,
+                "gold": gold,
+                "prompt": full_prompt,
+            }
+        )
+
+    llm = LLM(
+        model=args.model,
+        tensor_parallel_size=args.num_gpus,
+        dtype=args.dtype,
+        trust_remote_code=args.trust_remote_code,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        seed=args.seed,
+    )
+
+    sampling_params = SamplingParams(
+        n=args.n_samples,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        max_tokens=args.max_tokens,
+        stop=["<|im_end|>"],
+    )
+
+    outputs = llm.generate(prompts, sampling_params)
+
+    total = 0
+    correct = 0
+    per_problem_records = []
+
+    with output_path.open("w", encoding="utf-8") as f:
+        for ex, out in zip(examples, outputs):
+            sample_records = []
+            problem_correct = 0
+
+            for sample_id, completion in enumerate(out.outputs):
+                text = completion.text
+                pred = extract_answer(text)
+                is_correct = pred == ex["gold"]
+
+                total += 1
+                correct += int(is_correct)
+                problem_correct += int(is_correct)
+
+                record = {
+                    "idx": ex["idx"],
+                    "sample_id": sample_id,
+                    "gold": ex["gold"],
+                    "gold_raw": str(ex["gold_raw"]),
+                    "pred": pred,
+                    "correct": is_correct,
+                    "completion": text,
+                    "finish_reason": completion.finish_reason,
+                    "stop_reason": getattr(completion, "stop_reason", None),
+                }
+                sample_records.append(record)
+                f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+            problem_acc = problem_correct / max(1, len(out.outputs))
+            per_problem_records.append(problem_acc)
+
+            if ex["idx"] < args.print_examples:
+                print("-" * 80)
+                print(f"Problem {ex['idx']}")
+                print(f"Gold: {ex['gold']} | Correct samples: {problem_correct}/{len(out.outputs)}")
+                print(f"First pred: {sample_records[0]['pred']}")
+                print(f"First completion preview:\n{sample_records[0]['completion'][:1000]}")
+
+    avg_pass1_micro = correct / total if total > 0 else 0.0
+    avg_pass1_macro = sum(per_problem_records) / len(per_problem_records) if per_problem_records else 0.0
+
+    print("=" * 80)
+    print("Final results")
+    print("=" * 80)
+    print(f"Problems:                 {len(examples)}")
+    print(f"Samples per problem:       {args.n_samples}")
+    print(f"Total samples:             {total}")
+    print(f"Correct samples:           {correct}")
+    print(f"Average pass@1 micro:      {avg_pass1_micro:.6f} ({100 * avg_pass1_micro:.2f}%)")
+    print(f"Average pass@1 macro:      {avg_pass1_macro:.6f} ({100 * avg_pass1_macro:.2f}%)")
+    print(f"Saved generations to:      {output_path}")
+    print("=" * 80)
+
+
+if __name__ == "__main__":
+    main()

tools/fp8_cast_bf16.py

@@ -0,0 +1,118 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Adapt from https://github.com/alibaba/Pai-Megatron-Patch/blob/2b201af08336dea0403df7c6b497c964cf5a2e75/toolkits/model_checkpoints_convertor/deepseek/fp8_cast_bf16.py
+import json
+import os
+from argparse import ArgumentParser
+from glob import glob
+
+import torch
+import triton
+import triton.language as tl
+from safetensors.torch import load_file, save_file
+from tqdm import tqdm
+
+
+@triton.jit
+def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr):
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    n = tl.cdiv(N, BLOCK_SIZE)
+    offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    offs = offs_m[:, None] * N + offs_n[None, :]
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    x = tl.load(x_ptr + offs, mask=mask).to(tl.float32)
+    s = tl.load(s_ptr + pid_m * n + pid_n)
+    y = x * s
+    tl.store(y_ptr + offs, y, mask=mask)
+
+
+def weight_dequant(x: torch.Tensor, s: torch.Tensor, block_size: int = 128) -> torch.Tensor:
+    assert x.is_contiguous() and s.is_contiguous()
+    assert x.dim() == 2 and s.dim() == 2
+    M, N = x.size()
+    y = torch.empty_like(x, dtype=torch.get_default_dtype())
+
+    def grid(meta):
+        return (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"]))
+
+    weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE=block_size)
+    return y
+
+
+def main(fp8_path, bf16_path):
+    torch.set_default_dtype(torch.bfloat16)
+    os.makedirs(bf16_path, exist_ok=True)
+    os.system("cp -rf " + fp8_path + "/config.json " + bf16_path)
+    os.system("cp -rf " + fp8_path + "/*.py " + bf16_path)
+    os.system("cp -rf " + fp8_path + "/tokenizer* " + bf16_path)
+    os.system("cp -rf " + fp8_path + "/chat_template* " + bf16_path)
+    model_index_file = os.path.join(fp8_path, "model.safetensors.index.json")
+    with open(model_index_file) as f:
+        model_index = json.load(f)
+    weight_map = model_index["weight_map"]
+
+    # Cache for loaded safetensor files
+    loaded_files = {}
+    fp8_weight_names = []
+
+    # Helper function to get tensor from the correct file
+    def get_tensor(tensor_name):
+        file_name = weight_map[tensor_name]
+        if file_name not in loaded_files:
+            file_path = os.path.join(fp8_path, file_name)
+            loaded_files[file_name] = load_file(file_path, device="cuda")
+        return loaded_files[file_name][tensor_name]
+
+    safetensor_files = list(glob(os.path.join(fp8_path, "*.safetensors")))
+    safetensor_files.sort()
+    for safetensor_file in tqdm(safetensor_files):
+        print(f"Handling file: {safetensor_file}")
+        file_name = os.path.basename(safetensor_file)
+        current_state_dict = load_file(safetensor_file, device="cuda")
+        loaded_files[file_name] = current_state_dict
+
+        new_state_dict = {}
+        for weight_name, weight in current_state_dict.items():
+            if weight_name.endswith("_scale_inv"):
+                continue
+            elif weight.element_size() == 1:  # FP8 weight
+                scale_inv_name = f"{weight_name}_scale_inv"
+                try:
+                    # Get scale_inv from the correct file
+                    scale_inv = get_tensor(scale_inv_name)
+                    fp8_weight_names.append(weight_name)
+                    new_state_dict[weight_name] = weight_dequant(weight, scale_inv)
+                except KeyError:
+                    print(f"Warning: Missing scale_inv tensor for {weight_name}, skipping conversion")
+                    new_state_dict[weight_name] = weight
+            else:
+                new_state_dict[weight_name] = weight
+
+        new_safetensor_file = os.path.join(bf16_path, file_name)
+        save_file(new_state_dict, new_safetensor_file)
+
+        # Memory management: keep only the 2 most recently used files
+        if len(loaded_files) > 2:
+            oldest_file = next(iter(loaded_files))
+            del loaded_files[oldest_file]
+            torch.cuda.empty_cache()
+
+    # Update model index
+    new_model_index_file = os.path.join(bf16_path, "model.safetensors.index.json")
+    for weight_name in fp8_weight_names:
+        scale_inv_name = f"{weight_name}_scale_inv"
+        if scale_inv_name in weight_map:
+            weight_map.pop(scale_inv_name)
+    with open(new_model_index_file, "w") as f:
+        json.dump({"metadata": {}, "weight_map": weight_map}, f, indent=2)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--input-fp8-hf-path", type=str, required=True)
+    parser.add_argument("--output-bf16-hf-path", type=str, required=True)
+    args = parser.parse_args()
+    main(args.input_fp8_hf_path, args.output_bf16_hf_path)

tools/merge_poe_lora.py

@@ -0,0 +1,53 @@
+import argparse
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base-model", required=True)
+    parser.add_argument("--adapter", required=True)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
+    args = parser.parse_args()
+
+    dtype_map = {
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+        "float32": torch.float32,
+    }
+    dtype = dtype_map[args.dtype]
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.base_model,
+        trust_remote_code=True,
+    )
+
+    base_model = AutoModelForCausalLM.from_pretrained(
+        args.base_model,
+        torch_dtype=dtype,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+
+    model = PeftModel.from_pretrained(
+        base_model,
+        args.adapter,
+        torch_dtype=dtype,
+    )
+
+    model = model.merge_and_unload()
+
+    model.save_pretrained(
+        args.output_dir,
+        safe_serialization=True,
+        max_shard_size="4GB",
+    )
+    tokenizer.save_pretrained(args.output_dir)
+
+    print(f"Saved merged model to {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()

tools/train_poe_distill_lora.py

@@ -0,0 +1,604 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""LoRA training with an online product-of-experts distillation target.
+
+This script trains on fixed pi_ref rollouts, but computes full-vocabulary
+teacher/ref distributions online:
+
+    pi_star(. | s) proportional to pi_T(. | s)^beta * pi_ref(. | s)^(1-beta)
+    beta = alpha / (alpha + 1)
+
+The trainable model is pi_ref plus LoRA adapters.  The frozen pi_ref
+distribution is obtained by disabling the adapter on the same model, avoiding a
+second copy of the 4B reference model.
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import os
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+from peft import LoraConfig, TaskType, get_peft_model
+from torch.nn.utils.rnn import pad_sequence
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    Trainer,
+    TrainingArguments,
+    set_seed,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Product-of-experts LoRA distillation on fixed rollouts.")
+
+    parser.add_argument("--student-model", default=os.environ.get("SFT_CHECKPOINT"), required=False)
+    parser.add_argument("--teacher-model", default=os.environ.get("TEACHER_MODEL", "Qwen/Qwen3-8B"))
+    parser.add_argument("--train-data", default="data/rollouts/dapo-math-17k-qwen3-4b-sft-rollouts.parquet")
+    parser.add_argument("--output-dir", default="checkpoints/qwen3-4b-poe-distill-lora")
+
+    parser.add_argument("--alpha", type=float, default=1.0)
+    parser.add_argument(
+        "--beta-start",
+        type=float,
+        default=None,
+        help="Initial beta. If unset, uses alpha / (alpha + 1) as a fixed beta.",
+    )
+    parser.add_argument(
+        "--beta-end",
+        type=float,
+        default=None,
+        help="Final beta. If unset, uses alpha / (alpha + 1) as a fixed beta.",
+    )
+    parser.add_argument(
+        "--beta-schedule-steps",
+        type=int,
+        default=None,
+        help="Number of optimizer steps used to ramp beta from beta-start to beta-end.",
+    )
+    parser.add_argument("--beta-schedule", choices=["linear", "cosine"], default="linear")
+    parser.add_argument(
+        "--beta-hold-steps",
+        type=int,
+        default=0,
+        help="Keep beta fixed at beta-start for this many optimizer steps before scheduling.",
+    )
+    parser.add_argument(
+        "--beta-transition-steps",
+        type=int,
+        default=None,
+        help="Number of optimizer steps used to move beta from beta-start to beta-end after beta-hold-steps.",
+    )
+    parser.add_argument(
+        "--lr-start",
+        type=float,
+        default=None,
+        help="Initial LR for custom hold-then-transition schedule. If unset, uses --learning-rate.",
+    )
+    parser.add_argument(
+        "--lr-end",
+        type=float,
+        default=None,
+        help="Final LR after custom transition. If unset, custom LR scheduling is disabled.",
+    )
+    parser.add_argument(
+        "--lr-hold-steps",
+        type=int,
+        default=None,
+        help="Keep LR fixed at lr-start for this many optimizer steps. If unset, uses beta-hold-steps.",
+    )
+    parser.add_argument(
+        "--lr-transition-steps",
+        type=int,
+        default=None,
+        help="Number of optimizer steps used to move LR from lr-start to lr-end. If unset, uses beta-transition-steps.",
+    )
+    parser.add_argument(
+        "--hold-transition-schedule",
+        choices=["linear", "cosine"],
+        default="linear",
+        help="Schedule type for hold-then-transition beta/LR.",
+    )
+    parser.add_argument(
+        "--loss-type",
+        choices=["full_vocab", "sampled_token"],
+        default="full_vocab",
+        help=(
+            "full_vocab matches the normalized PoE distribution over the whole vocab. "
+            "sampled_token uses an OPD-style sampled-token surrogate with a PoE advantage."
+        ),
+    )
+    parser.add_argument(
+        "--advantage-normalization",
+        choices=["none", "batch", "sequence"],
+        default="batch",
+        help="Only used by --loss-type sampled_token.",
+    )
+    parser.add_argument(
+        "--advantage-clip",
+        type=float,
+        default=None,
+        help="Symmetric clamp for sampled-token advantages. Example: 5.0.",
+    )
+    parser.add_argument(
+        "--use-ppo-clip",
+        action="store_true",
+        default=False,
+        help=(
+            "Only used by --loss-type sampled_token. Use PPO-style ratio clipping "
+            "with the frozen reference log-prob as the old rollout log-prob."
+        ),
+    )
+    parser.add_argument(
+        "--ppo-clip-low",
+        type=float,
+        default=0.2,
+        help="Only used when --use-ppo-clip is set. Lower PPO clip epsilon.",
+    )
+    parser.add_argument(
+        "--ppo-clip-high",
+        type=float,
+        default=0.2,
+        help="Only used when --use-ppo-clip is set. Upper PPO clip epsilon.",
+    )
+    parser.add_argument(
+        "--sampled-loss-reduction",
+        choices=["per_sample", "per_token"],
+        default="per_sample",
+        help=(
+            "Only used by --loss-type sampled_token. per_sample averages each response "
+            "first, then averages across batch; per_token averages over all response tokens."
+        ),
+    )
+    parser.add_argument(
+        "--positive-advantages-only",
+        action="store_true",
+        default=False,
+        help="Only reinforce sampled tokens with positive PoE advantages.",
+    )
+    parser.add_argument("--max-length", type=int, default=4096)
+    parser.add_argument("--distill-chunk-size", type=int, default=128)
+    parser.add_argument("--max-train-samples", type=int, default=None)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--num-train-epochs", type=float, default=1.0)
+    parser.add_argument("--max-steps", type=int, default=-1)
+    parser.add_argument("--per-device-train-batch-size", type=int, default=1)
+    parser.add_argument("--gradient-accumulation-steps", type=int, default=16)
+    parser.add_argument("--learning-rate", type=float, default=2e-5)
+    parser.add_argument("--weight-decay", type=float, default=0.0)
+    parser.add_argument("--adam-beta1", type=float, default=0.9)
+    parser.add_argument("--adam-beta2", type=float, default=0.999)
+    parser.add_argument("--adam-epsilon", type=float, default=1e-8)
+    parser.add_argument("--warmup-ratio", type=float, default=0.03)
+    parser.add_argument("--lr-scheduler-type", default="cosine")
+    parser.add_argument("--logging-steps", type=int, default=1)
+    parser.add_argument("--save-steps", type=int, default=100)
+    parser.add_argument("--save-total-limit", type=int, default=0)
+    parser.add_argument("--bf16", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--fp16", action="store_true", default=False)
+    parser.add_argument("--gradient-checkpointing", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--report-to", default="none")
+
+    parser.add_argument("--lora-r", type=int, default=64)
+    parser.add_argument("--lora-alpha", type=int, default=128)
+    parser.add_argument("--lora-dropout", type=float, default=0.05)
+    parser.add_argument(
+        "--freeze-lora-b-after-step",
+        type=int,
+        default=None,
+        help="Freeze all LoRA B matrices once global_step reaches this value. Example: 20.",
+    )
+    parser.add_argument(
+        "--lora-target-modules",
+        default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
+        help="Comma-separated LoRA target modules.",
+    )
+
+    parser.add_argument("--trust-remote-code", action="store_true", default=True)
+    parser.add_argument(
+        "--attn-implementation",
+        default=None,
+        choices=[None, "eager", "sdpa", "flash_attention_2"],
+        help="Forwarded to from_pretrained when set.",
+    )
+
+    args = parser.parse_args()
+    if not args.student_model:
+        raise ValueError("Pass --student-model or set SFT_CHECKPOINT to the Qwen3-4B SFT checkpoint.")
+    if args.alpha <= 0:
+        raise ValueError("--alpha must be positive.")
+    fixed_beta = args.alpha / (args.alpha + 1.0)
+    if args.beta_start is None:
+        args.beta_start = fixed_beta
+    if args.beta_end is None:
+        args.beta_end = fixed_beta
+    if not 0.0 <= args.beta_start <= 1.0:
+        raise ValueError("--beta-start must be in [0, 1].")
+    if not 0.0 <= args.beta_end <= 1.0:
+        raise ValueError("--beta-end must be in [0, 1].")
+    if args.beta_schedule_steps is not None and args.beta_schedule_steps <= 0:
+        raise ValueError("--beta-schedule-steps must be positive when set.")
+    if args.advantage_clip is not None and args.advantage_clip <= 0:
+        raise ValueError("--advantage-clip must be positive when set.")
+    if args.ppo_clip_low < 0 or args.ppo_clip_high < 0:
+        raise ValueError("--ppo-clip-low and --ppo-clip-high must be non-negative.")
+    if args.freeze_lora_b_after_step is not None and args.freeze_lora_b_after_step < 0:
+        raise ValueError("--freeze-lora-b-after-step must be non-negative when set.")
+    if args.fp16 and args.bf16:
+        args.bf16 = False
+    return args
+
+
+def first_assistant_index(messages: list[dict[str, str]]) -> int:
+    for idx, message in enumerate(messages):
+        if message.get("role") == "assistant":
+            return idx
+    raise ValueError("Rollout row has no assistant message.")
+
+
+def tokenize_rollout(example: dict[str, Any], tokenizer: AutoTokenizer, max_length: int) -> dict[str, Any]:
+    messages = example["messages"]
+    assistant_idx = first_assistant_index(messages)
+    prompt_messages = messages[:assistant_idx]
+    full_messages = messages[: assistant_idx + 1]
+
+    prompt_text = tokenizer.apply_chat_template(
+        prompt_messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=True,
+    )
+    full_text = tokenizer.apply_chat_template(
+        full_messages,
+        tokenize=False,
+        add_generation_prompt=False,
+        enable_thinking=True,
+    )
+
+    prompt_ids = tokenizer.encode(prompt_text, add_special_tokens=False)
+    input_ids = tokenizer.encode(full_text, add_special_tokens=False)
+
+    if len(input_ids) > max_length:
+        input_ids = input_ids[:max_length]
+
+    # Mask is aligned to labels=input_ids[1:].  A label predicts token position
+    # j=i+1, so it belongs to the response when j >= len(prompt_ids).
+    label_len = max(len(input_ids) - 1, 0)
+    loss_mask = [1 if i + 1 >= len(prompt_ids) else 0 for i in range(label_len)]
+
+    if sum(loss_mask) == 0:
+        # Drop examples where truncation removed the assistant response.
+        return {"input_ids": [], "loss_mask": []}
+
+    return {"input_ids": input_ids, "loss_mask": loss_mask}
+
+
+@dataclass
+class DistillCollator:
+    pad_token_id: int
+
+    def __call__(self, features: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
+        input_ids = [torch.tensor(f["input_ids"], dtype=torch.long) for f in features]
+        loss_masks = [torch.tensor(f["loss_mask"], dtype=torch.float32) for f in features]
+        lengths = torch.tensor([x.size(0) for x in input_ids], dtype=torch.long)
+
+        padded_input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.pad_token_id)
+        # loss_mask is one shorter than input_ids because it aligns to shifted labels.
+        padded_loss_masks = pad_sequence(loss_masks, batch_first=True, padding_value=0.0)
+        positions = torch.arange(padded_input_ids.size(1)).unsqueeze(0)
+        attention_mask = (positions < lengths.unsqueeze(1)).long()
+        return {
+            "input_ids": padded_input_ids,
+            "attention_mask": attention_mask,
+            "loss_mask": padded_loss_masks,
+        }
+
+
+class PoEDistillTrainer(Trainer):
+    def __init__(
+        self,
+        *args: Any,
+        teacher_model: torch.nn.Module,
+        beta_start: float,
+        beta_end: float,
+        beta_schedule_steps: int | None,
+        beta_schedule: str,
+        loss_type: str,
+        advantage_normalization: str,
+        advantage_clip: float | None,
+        positive_advantages_only: bool,
+        use_ppo_clip: bool,
+        ppo_clip_low: float,
+        ppo_clip_high: float,
+        sampled_loss_reduction: str,
+        distill_chunk_size: int,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.teacher_model = teacher_model
+        self.teacher_model.to(self.args.device)
+        self.teacher_model.eval()
+        self.beta_start = beta_start
+        self.beta_end = beta_end
+        self.beta_schedule_steps = beta_schedule_steps
+        self.beta_schedule = beta_schedule
+        self.loss_type = loss_type
+        self.advantage_normalization = advantage_normalization
+        self.advantage_clip = advantage_clip
+        self.positive_advantages_only = positive_advantages_only
+        self.use_ppo_clip = use_ppo_clip
+        self.ppo_clip_low = ppo_clip_low
+        self.ppo_clip_high = ppo_clip_high
+        self.sampled_loss_reduction = sampled_loss_reduction
+        self.distill_chunk_size = distill_chunk_size
+
+    def current_beta(self) -> float:
+        schedule_steps = self.beta_schedule_steps
+        if schedule_steps is None:
+            schedule_steps = self.state.max_steps if self.state.max_steps > 0 else None
+        if schedule_steps is None or schedule_steps == 0:
+            return self.beta_end
+
+        progress = min(max(self.state.global_step / schedule_steps, 0.0), 1.0)
+        if self.beta_schedule == "cosine":
+            progress = 0.5 - 0.5 * torch.cos(torch.tensor(progress * torch.pi)).item()
+        return self.beta_start + (self.beta_end - self.beta_start) * progress
+
+    @staticmethod
+    def gather_token_logprobs(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        logits = logits.float()
+        token_logits = logits.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)
+        return token_logits - logits.logsumexp(dim=-1)
+
+    def normalize_advantages(self, advantages: torch.Tensor, loss_mask: torch.Tensor) -> torch.Tensor:
+        if self.advantage_normalization == "none":
+            return advantages
+
+        if self.advantage_normalization == "batch":
+            denom = loss_mask.sum().clamp_min(1.0)
+            mean = (advantages * loss_mask).sum() / denom
+            var = (((advantages - mean) * loss_mask) ** 2).sum() / denom
+            return (advantages - mean) / torch.sqrt(var + 1e-6)
+
+        denom = loss_mask.sum(dim=1, keepdim=True).clamp_min(1.0)
+        mean = (advantages * loss_mask).sum(dim=1, keepdim=True) / denom
+        var = (((advantages - mean) * loss_mask) ** 2).sum(dim=1, keepdim=True) / denom
+        return (advantages - mean) / torch.sqrt(var + 1e-6)
+
+    def compute_loss(
+        self,
+        model: torch.nn.Module,
+        inputs: dict[str, torch.Tensor],
+        return_outputs: bool = False,
+        **_: Any,
+    ):
+        loss_mask = inputs.pop("loss_mask")
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs["attention_mask"]
+        labels = input_ids[:, 1:]
+
+        with torch.no_grad():
+            teacher_logits = self.teacher_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                use_cache=False,
+            ).logits[:, :-1, :].detach()
+
+            adapter_owner = model.module if hasattr(model, "module") else model
+            disable_adapter = getattr(adapter_owner, "disable_adapter", None)
+            ref_context = disable_adapter() if disable_adapter is not None else contextlib.nullcontext()
+            with ref_context:
+                ref_logits = adapter_owner(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    use_cache=False,
+                ).logits[:, :-1, :].detach()
+
+        student_outputs = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
+        student_logits = student_outputs.logits[:, :-1, :]
+        if teacher_logits.size(-1) != student_logits.size(-1) or ref_logits.size(-1) != student_logits.size(-1):
+            raise ValueError(
+                "Teacher, reference, and student vocab sizes must match for full-vocab PoE distillation. "
+                f"Got teacher={teacher_logits.size(-1)}, ref={ref_logits.size(-1)}, "
+                f"student={student_logits.size(-1)}."
+            )
+
+        total_loss = student_logits.new_zeros(())
+        total_tokens = loss_mask.sum().clamp_min(1.0)
+        beta = self.current_beta()
+
+        if self.loss_type == "sampled_token":
+            teacher_logp = self.gather_token_logprobs(teacher_logits, labels)
+            ref_logp = self.gather_token_logprobs(ref_logits, labels)
+            student_logp = self.gather_token_logprobs(student_logits, labels)
+
+            poe_score = beta * teacher_logp + (1.0 - beta) * ref_logp
+            advantages = poe_score - student_logp.detach()
+            advantages = self.normalize_advantages(advantages, loss_mask)
+            if self.advantage_clip is not None:
+                advantages = advantages.clamp(min=-self.advantage_clip, max=self.advantage_clip)
+            if self.positive_advantages_only:
+                advantages = advantages.clamp_min(0.0)
+
+            advantages = advantages.detach()
+
+            if self.use_ppo_clip:
+                # The rollouts are generated by the frozen reference/SFT policy, so ref_logp
+                # is used as the old rollout log-prob. This mirrors the PPO-style clipped
+                # policy loss used in RL frameworks such as slime.
+                ratio = torch.exp(student_logp - ref_logp.detach())
+                ratio_clipped = ratio.clamp(1.0 - self.ppo_clip_low, 1.0 + self.ppo_clip_high)
+                pg_loss_unclipped = -ratio * advantages
+                pg_loss_clipped = -ratio_clipped * advantages
+                token_loss = torch.maximum(pg_loss_unclipped, pg_loss_clipped)
+            else:
+                # Direct OPD-style sampled-token surrogate.
+                token_loss = -advantages * student_logp
+
+            if self.sampled_loss_reduction == "per_token":
+                loss = (token_loss * loss_mask).sum() / total_tokens
+            else:
+                # Per-sample mean: each response contributes equally regardless of length.
+                seq_loss = (token_loss * loss_mask).sum(dim=1) / loss_mask.sum(dim=1).clamp_min(1.0)
+                loss = seq_loss.mean()
+
+            return (loss, student_outputs) if return_outputs else loss
+
+        seq_len = student_logits.size(1)
+        for start in range(0, seq_len, self.distill_chunk_size):
+            end = min(start + self.distill_chunk_size, seq_len)
+            mask = loss_mask[:, start:end]
+            if mask.sum() == 0:
+                continue
+
+            teacher_logp = F.log_softmax(teacher_logits[:, start:end, :].float(), dim=-1)
+            ref_logp = F.log_softmax(ref_logits[:, start:end, :].float(), dim=-1)
+            student_logp = F.log_softmax(student_logits[:, start:end, :].float(), dim=-1)
+
+            poe_logits = beta * teacher_logp + (1.0 - beta) * ref_logp
+            target_probs = F.softmax(poe_logits, dim=-1)
+            token_ce = -(target_probs * student_logp).sum(dim=-1)
+            total_loss = total_loss + (token_ce * mask).sum()
+
+        loss = total_loss / total_tokens
+        return (loss, student_outputs) if return_outputs else loss
+
+
+class FreezeLoRABCallback(TrainerCallback):
+    def __init__(self, freeze_after_step: int | None) -> None:
+        self.freeze_after_step = freeze_after_step
+        self.frozen = False
+
+    def on_step_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        model: torch.nn.Module | None = None,
+        **kwargs: Any,
+    ) -> TrainerControl:
+        if self.freeze_after_step is None or self.frozen or model is None:
+            return control
+        if state.global_step < self.freeze_after_step:
+            return control
+
+        frozen_params = 0
+        module = model.module if hasattr(model, "module") else model
+        for name, param in module.named_parameters():
+            if ".lora_B." in name or "lora_B." in name:
+                param.requires_grad_(False)
+                frozen_params += param.numel()
+
+        self.frozen = True
+        if args.process_index == 0:
+            print(f"[PoE Distill] Froze LoRA B at global_step={state.global_step} ({frozen_params} params).")
+        return control
+
+
+def main() -> None:
+    args = parse_args()
+    set_seed(args.seed)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.student_model, trust_remote_code=args.trust_remote_code)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    raw_dataset = load_dataset("parquet", data_files=args.train_data, split="train")
+    if args.max_train_samples is not None:
+        raw_dataset = raw_dataset.select(range(min(args.max_train_samples, len(raw_dataset))))
+
+    train_dataset = raw_dataset.map(
+        lambda ex: tokenize_rollout(ex, tokenizer, args.max_length),
+        remove_columns=raw_dataset.column_names,
+        desc="Tokenizing pi_ref rollouts",
+    ).filter(lambda ex: len(ex["input_ids"]) > 0, desc="Dropping empty responses")
+
+    model_kwargs = {
+        "torch_dtype": torch.bfloat16 if args.bf16 else (torch.float16 if args.fp16 else torch.float32),
+        "trust_remote_code": args.trust_remote_code,
+    }
+    if args.attn_implementation is not None:
+        model_kwargs["attn_implementation"] = args.attn_implementation
+
+    student = AutoModelForCausalLM.from_pretrained(args.student_model, **model_kwargs)
+    teacher = AutoModelForCausalLM.from_pretrained(args.teacher_model, **model_kwargs)
+    teacher.eval()
+    teacher.requires_grad_(False)
+
+    if args.gradient_checkpointing:
+        student.gradient_checkpointing_enable()
+        student.config.use_cache = False
+        teacher.config.use_cache = False
+
+    lora_config = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        target_modules=[m.strip() for m in args.lora_target_modules.split(",") if m.strip()],
+    )
+    student = get_peft_model(student, lora_config)
+    student.print_trainable_parameters()
+
+    training_args = TrainingArguments(
+        output_dir=args.output_dir,
+        num_train_epochs=args.num_train_epochs,
+        max_steps=args.max_steps,
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        learning_rate=args.learning_rate,
+        weight_decay=args.weight_decay,
+        adam_beta1=args.adam_beta1,
+        adam_beta2=args.adam_beta2,
+        adam_epsilon=args.adam_epsilon,
+        warmup_ratio=args.warmup_ratio,
+        lr_scheduler_type=args.lr_scheduler_type,
+        logging_steps=args.logging_steps,
+        save_steps=args.save_steps,
+        save_total_limit=args.save_total_limit,
+        bf16=args.bf16,
+        fp16=args.fp16,
+        gradient_checkpointing=args.gradient_checkpointing,
+        remove_unused_columns=False,
+        report_to=[] if args.report_to == "none" else args.report_to.split(","),
+    )
+
+    trainer = PoEDistillTrainer(
+        model=student,
+        args=training_args,
+        train_dataset=train_dataset,
+        data_collator=DistillCollator(pad_token_id=tokenizer.pad_token_id),
+        tokenizer=tokenizer,
+        teacher_model=teacher,
+        beta_start=args.beta_start,
+        beta_end=args.beta_end,
+        beta_schedule_steps=args.beta_schedule_steps,
+        beta_schedule=args.beta_schedule,
+        loss_type=args.loss_type,
+        advantage_normalization=args.advantage_normalization,
+        advantage_clip=args.advantage_clip,
+        positive_advantages_only=args.positive_advantages_only,
+        use_ppo_clip=args.use_ppo_clip,
+        ppo_clip_low=args.ppo_clip_low,
+        ppo_clip_high=args.ppo_clip_high,
+        sampled_loss_reduction=args.sampled_loss_reduction,
+        distill_chunk_size=args.distill_chunk_size,
+        callbacks=[FreezeLoRABCallback(args.freeze_lora_b_after_step)],
+    )
+    trainer.train()
+    trainer.save_model(args.output_dir)
+    tokenizer.save_pretrained(args.output_dir)
+
+
+if __name__ == "__main__":
+    main()