from __future__ import annotations import argparse import json import math import time from dataclasses import dataclass from pathlib import Path import torch import torch.nn.functional as F from huggingface_hub import hf_hub_download from safetensors.torch import save_file from transformers import PreTrainedModel from torch.optim import AdamW from torch.utils.data import DataLoader, Dataset from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer @dataclass class DistillConfig: teacher_path: str student_path: str output_dir: str dataset_path: str = "" dataset_split: str = "train" seq_len: int = 128 batch_size: int = 1 grad_accum: int = 4 lr: float = 3e-4 warmup_ratio: float = 0.05 max_steps: int = 1000 log_every: int = 25 save_every: int = 250 resume_from: str = "" ce_loss_weight: float = 1.0 mse_loss_weight: float = 1.0 max_grad_norm: float = 1.0 seed: int = 42 class TokenizedDataset(Dataset): def __init__(self, token_ids: list[int], seq_len: int): self.seq_len = seq_len self.examples = [] for i in range(0, len(token_ids) - seq_len - 1, seq_len): chunk = token_ids[i : i + seq_len + 1] if len(chunk) == seq_len + 1: self.examples.append(torch.tensor(chunk, dtype=torch.long)) def __len__(self) -> int: return len(self.examples) def __getitem__(self, idx: int) -> torch.Tensor: return self.examples[idx] def load_teacher(path: str, dtype: torch.dtype) -> tuple: model = AutoModel.from_pretrained(path, trust_remote_code=True, dtype=dtype) model.eval() for p in model.parameters(): p.requires_grad = False embed_weight = model.language_model.embed_tokens.weight return model, embed_weight def load_student(path: str, dtype: torch.dtype) -> tuple: config = AutoConfig.from_pretrained(path, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( path, trust_remote_code=True, dtype=dtype ) model.train() return model, config def tokenize_dataset(dataset_path: str, tokenizer, max_tokens: int = 500000) -> list[int]: """Read a raw text file and tokenize it.""" path = Path(dataset_path) if not path.exists(): raise FileNotFoundError(f"Dataset not found: {dataset_path}") text = path.read_text(encoding="utf-8") all_ids: list[int] = [] for paragraph in text.split("\n\n"): paragraph = paragraph.strip() if not paragraph: continue ids = tokenizer.encode(paragraph, add_special_tokens=False) all_ids.extend(ids) all_ids.append(tokenizer.eos_token_id) if len(all_ids) >= max_tokens: break return all_ids[:max_tokens] def compute_teacher_outputs(teacher, embed_weight, input_ids: torch.Tensor): with torch.no_grad(): out = teacher( input_ids=input_ids, output_hidden_states=True, use_cache=False, ) hidden = out.hidden_states[-1] logits = hidden @ embed_weight.T return logits, hidden def compute_student_outputs(student, input_ids: torch.Tensor): out = student( input_ids=input_ids, output_hidden_states=True, use_cache=False, ) return out.logits, out.hidden_states[-1] def distillation_loss( student_logits: torch.Tensor, teacher_logits: torch.Tensor, student_hidden: torch.Tensor, teacher_hidden: torch.Tensor, ce_weight: float, mse_weight: float, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: shift_student_logits = student_logits[:, :-1, :].contiguous() shift_teacher_logits = teacher_logits[:, :-1, :].contiguous() teacher_probs = F.softmax(shift_teacher_logits, dim=-1) student_log_probs = F.log_softmax(shift_student_logits, dim=-1) ce_loss = -(teacher_probs * student_log_probs).sum(dim=-1).mean() shift_student_hidden = student_hidden[:, :-1, :].contiguous() shift_teacher_hidden = teacher_hidden[:, :-1, :].contiguous() mse_loss = F.mse_loss(shift_student_hidden, shift_teacher_hidden) total = ce_weight * ce_loss + mse_weight * mse_loss return total, ce_loss.detach(), mse_loss.detach() def get_cosine_schedule_with_warmup(optimizer, warmup_steps: int, total_steps: int): def lr_lambda(current_step: int) -> float: if current_step < warmup_steps: return float(current_step) / float(max(1, warmup_steps)) progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps)) return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress))) return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) def save_checkpoint(student, student_config, output_dir: str, step: int): out = Path(output_dir) out.mkdir(parents=True, exist_ok=True) state_dict = {} for k, v in student.state_dict().items(): if k == "lm_head.weight" and "model.embed_tokens.weight" in state_dict: state_dict[k] = state_dict["model.embed_tokens.weight"].clone() else: state_dict[k] = v.cpu() save_file(state_dict, str(out / "model.safetensors"), metadata={"format": "pt"}) config_path = out / "config.json" config_path.write_text(json.dumps(student_config.to_dict(), indent=2, sort_keys=True) + "\n", encoding="utf-8") (out / "train_state.json").write_text(json.dumps({"step": step}) + "\n", encoding="utf-8") print(f" Checkpoint saved at step {step}: {out}", flush=True) def run_distill(args: DistillConfig) -> None: torch.manual_seed(args.seed) device = torch.device("cpu") dtype = torch.float32 print(f"Loading teacher from {args.teacher_path}...", flush=True) teacher, embed_weight = load_teacher(args.teacher_path, dtype) teacher_params = sum(p.numel() for p in teacher.parameters()) print(f" Teacher loaded: {teacher_params / 1e6:.1f}M params", flush=True) print(f"Loading student from {args.student_path}...", flush=True) student, student_config = load_student(args.student_path, dtype) student_params = sum(p.numel() for p in student.parameters()) print(f" Student loaded: {student_params / 1e6:.1f}M params", flush=True) print("Loading tokenizer...", flush=True) tokenizer = AutoTokenizer.from_pretrained(args.student_path, trust_remote_code=True) print(f"Tokenizing dataset ({args.dataset_path})...", flush=True) all_ids = tokenize_dataset(args.dataset_path, tokenizer) print(f" Total tokens: {len(all_ids)}", flush=True) dataset = TokenizedDataset(all_ids, args.seq_len) dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, drop_last=True) print(f" Dataset size: {len(dataset)} examples", flush=True) optimizer = AdamW( [p for p in student.parameters() if p.requires_grad], lr=args.lr, weight_decay=0.01, ) scheduler = get_cosine_schedule_with_warmup( optimizer, warmup_steps=int(args.max_steps * args.warmup_ratio), total_steps=args.max_steps, ) print(f"Starting distillation: {args.max_steps} steps, lr={args.lr}", flush=True) print(f" Loss weights: CE={args.ce_loss_weight}, MSE={args.mse_loss_weight}", flush=True) print(f" Batch size={args.batch_size}, grad_accum={args.grad_accum}, seq_len={args.seq_len}", flush=True) start_step = 0 if args.resume_from: ckpt_path = Path(args.resume_from) ckpt_file = ckpt_path / "model.safetensors" state_file = ckpt_path / "train_state.json" if ckpt_file.exists(): from safetensors.torch import load_file ckpt_state = load_file(str(ckpt_file)) missing, unexpected = student.load_state_dict(ckpt_state, strict=False) if state_file.exists(): start_step = json.loads(state_file.read_text())["step"] print(f" Resumed from {ckpt_path} at step {start_step}: missing={len(missing)}, unexpected={len(unexpected)}", flush=True) else: print(f" Warning: checkpoint not found at {ckpt_file}, starting from scratch", flush=True) step = start_step running_ce = 0.0 running_mse = 0.0 running_total = 0.0 start_time = time.time() steps_done = 0 student.train() while step < args.max_steps: for batch in dataloader: if step >= args.max_steps: break input_ids = batch.to(device) teacher_logits, teacher_hidden = compute_teacher_outputs(teacher, embed_weight, input_ids) student_logits, student_hidden = compute_student_outputs(student, input_ids) loss, ce_loss, mse_loss = distillation_loss( student_logits, teacher_logits, student_hidden, teacher_hidden, args.ce_loss_weight, args.mse_loss_weight, ) loss = loss / args.grad_accum loss.backward() running_ce += ce_loss.item() running_mse += mse_loss.item() running_total += loss.item() * args.grad_accum if (step + 1) % args.grad_accum == 0: torch.nn.utils.clip_grad_norm_(student.parameters(), args.max_grad_norm) optimizer.step() scheduler.step() optimizer.zero_grad() if (step + 1) % args.log_every == 0: elapsed = time.time() - start_time avg_ce = running_ce / args.log_every avg_mse = running_mse / args.log_every avg_total = running_total / args.log_every lr_now = scheduler.get_last_lr()[0] steps_per_sec = (step + 1) / elapsed eta = (args.max_steps - step - 1) / steps_per_sec if steps_per_sec > 0 else 0 print( f" Step {step + 1}/{args.max_steps} | " f"total={avg_total:.4f} ce={avg_ce:.4f} mse={avg_mse:.6f} | " f"lr={lr_now:.2e} | {elapsed:.0f}s elapsed, ~{eta:.0f}s remaining", flush=True, ) running_ce = 0.0 running_mse = 0.0 running_total = 0.0 if (step + 1) % args.save_every == 0: save_checkpoint(student, student_config, args.output_dir, step + 1) step += 1 final_dir = Path(args.output_dir) / "final" save_checkpoint(student, student_config, str(final_dir), step) print(f"\nDistillation complete. Final model saved to {final_dir}") print(f"Total time: {time.time() - start_time:.0f}s") def _main() -> None: parser = argparse.ArgumentParser(description="Knowledge distillation: Qwen3.5 teacher -> pruned student") parser.add_argument("--teacher-path", required=True) parser.add_argument("--student-path", required=True) parser.add_argument("--output-dir", required=True) parser.add_argument("--dataset-path", required=True, help="Path to raw text file") parser.add_argument("--seq-len", type=int, default=256) parser.add_argument("--batch-size", type=int, default=1) parser.add_argument("--grad-accum", type=int, default=4) parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--warmup-ratio", type=float, default=0.05) parser.add_argument("--max-steps", type=int, default=1000) parser.add_argument("--log-every", type=int, default=25) parser.add_argument("--save-every", type=int, default=250) parser.add_argument("--ce-loss-weight", type=float, default=1.0) parser.add_argument("--mse-loss-weight", type=float, default=1.0) parser.add_argument("--max-grad-norm", type=float, default=1.0) parser.add_argument("--resume-from", type=str, default="", help="Path to checkpoint dir to resume from") parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() run_distill(DistillConfig(**{k.replace("-", "_"): v for k, v in vars(args).items()})) if __name__ == "__main__": _main()