初始化项目，由ModelHub XC社区提供模型
Model: ayh015/myLightningOPD Source: Original Platform
This commit is contained in:
ModelHub XC
783
tools/backup/train_poe_distill_lora_linear/cos_beta.py
Normal file
783
tools/backup/train_poe_distill_lora_linear/cos_beta.py
Normal file
@@ -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()
|
||||
181
tools/convert_fsdp_to_hf.py
Normal file
181
tools/convert_fsdp_to_hf.py
Normal file
@@ -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)
|
||||
256
tools/convert_hf_to_fp8.py
Normal file
256
tools/convert_hf_to_fp8.py
Normal file
@@ -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)
|
||||
140
tools/convert_hf_to_torch_dist.py
Normal file
140
tools/convert_hf_to_torch_dist.py
Normal file
@@ -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()
|
||||
245
tools/convert_k2_thinking_int4_to_bf16.py
Normal file
245
tools/convert_k2_thinking_int4_to_bf16.py
Normal file
@@ -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()
|
||||
116
tools/convert_to_hf.py
Normal file
116
tools/convert_to_hf.py
Normal file
@@ -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)
|
||||
219
tools/convert_torch_dist_to_hf.py
Normal file
219
tools/convert_torch_dist_to_hf.py
Normal file
@@ -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)
|
||||
494
tools/eval_aime2024_vllm.py
Normal file
494
tools/eval_aime2024_vllm.py
Normal file
@@ -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()
|
||||
488
tools/eval_aime2025_vllm.py
Normal file
488
tools/eval_aime2025_vllm.py
Normal file
@@ -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()
|
||||
486
tools/eval_hmmt2025_vllm.py
Normal file
486
tools/eval_hmmt2025_vllm.py
Normal file
@@ -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()
|
||||
118
tools/fp8_cast_bf16.py
Normal file
118
tools/fp8_cast_bf16.py
Normal file
@@ -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)
|
||||
53
tools/merge_poe_lora.py
Normal file
53
tools/merge_poe_lora.py
Normal file
@@ -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()
|
||||
604
tools/train_poe_distill_lora.py
Normal file
604
tools/train_poe_distill_lora.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user