diff --git a/trainer/opd_utils.py b/trainer/opd_utils.py new file mode 100644 index 0000000..65add25 --- /dev/null +++ b/trainer/opd_utils.py @@ -0,0 +1,148 @@ +"""Core tensor utilities for Generalized Knowledge Distillation (GKD).""" + +from typing import Dict, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import Tensor + + +def build_completion_mask( + completion_ids: Tensor, + base_mask: Optional[Tensor], + eos_token_id: Optional[int], + pad_token_id: Optional[int] = None, +) -> Tensor: + """Keep valid completion tokens up to and including the first EOS. + + MiniMind's native ``generate`` pads finished rows with repeated EOS tokens + until every row is complete. Those repeated tokens are valid model inputs, + but they must not contribute to the distillation objective. + """ + if completion_ids.ndim != 2: + raise ValueError(f"completion_ids must be 2D, got {tuple(completion_ids.shape)}") + if base_mask is None: + mask = torch.ones_like(completion_ids, dtype=torch.bool) + else: + if base_mask.shape != completion_ids.shape: + raise ValueError( + f"base_mask shape {tuple(base_mask.shape)} does not match " + f"completion_ids shape {tuple(completion_ids.shape)}" + ) + mask = base_mask.bool().clone() + + # Some tokenizers intentionally use EOS as PAD. In that case the first EOS + # still belongs to the completion and is handled by the cumulative mask. + if pad_token_id is not None and pad_token_id != eos_token_id: + mask &= completion_ids.ne(pad_token_id) + + if eos_token_id is not None: + eos_hits = completion_ids.eq(eos_token_id) & mask + eos_seen_before = eos_hits.long().cumsum(dim=1) - eos_hits.long() + mask &= eos_seen_before.eq(0) + return mask + + +def _masked_mean(values: Tensor, mask: Tensor) -> Tensor: + mask = mask.to(device=values.device, dtype=values.dtype) + return (values * mask).sum() / mask.sum().clamp_min(1.0) + + +def generalized_jsd_loss( + student_logits: Tensor, + teacher_logits: Tensor, + completion_mask: Tensor, + beta: float = 0.5, + temperature: float = 1.0, +) -> Tuple[Tensor, Dict[str, Tensor]]: + """Compute the exact full-vocabulary generalized JSD from GKD. + + With teacher distribution ``p`` and student distribution ``q``, the + mixture is ``m = beta * p + (1 - beta) * q`` and the objective is + ``beta * KL(p || m) + (1 - beta) * KL(q || m)``. Its limiting cases are + forward KL at ``beta=0`` and reverse KL at ``beta=1``. + + The teacher is always treated as a fixed target. Only completion tokens + selected by ``completion_mask`` contribute to the token-mean loss. + """ + if student_logits.ndim != 3 or teacher_logits.ndim != 3: + raise ValueError("student_logits and teacher_logits must be [batch, time, vocab]") + if student_logits.shape != teacher_logits.shape: + raise ValueError( + "student and teacher logits must have identical shapes, got " + f"{tuple(student_logits.shape)} and {tuple(teacher_logits.shape)}" + ) + if completion_mask.shape != student_logits.shape[:2]: + raise ValueError( + f"completion_mask shape {tuple(completion_mask.shape)} does not match " + f"logits shape {tuple(student_logits.shape[:2])}" + ) + if not 0.0 <= beta <= 1.0: + raise ValueError(f"beta must be in [0, 1], got {beta}") + if temperature <= 0: + raise ValueError(f"temperature must be positive, got {temperature}") + + # Float32 log-softmax keeps the divergence stable under fp16/bf16 model + # autocast, while preserving gradients back to the student parameters. + student_logps = F.log_softmax(student_logits.float() / temperature, dim=-1) + teacher_logps = F.log_softmax( + teacher_logits.detach().float() / temperature, dim=-1 + ) + student_probs = student_logps.exp() + teacher_probs = teacher_logps.exp() + + forward_kl = (teacher_probs * (teacher_logps - student_logps)).sum(dim=-1) + reverse_kl = (student_probs * (student_logps - teacher_logps)).sum(dim=-1) + + if beta == 0.0: + per_token_loss = forward_kl + elif beta == 1.0: + per_token_loss = reverse_kl + else: + beta_tensor = student_logps.new_tensor(beta) + mixture_logps = torch.logsumexp( + torch.stack( + ( + student_logps + torch.log1p(-beta_tensor), + teacher_logps + torch.log(beta_tensor), + ) + ), + dim=0, + ) + teacher_to_mixture = ( + teacher_probs * (teacher_logps - mixture_logps) + ).sum(dim=-1) + student_to_mixture = ( + student_probs * (student_logps - mixture_logps) + ).sum(dim=-1) + per_token_loss = ( + beta_tensor * teacher_to_mixture + + (1.0 - beta_tensor) * student_to_mixture + ) + + mask = completion_mask.bool() + loss = _masked_mean(per_token_loss, mask) + top1_agreement = _masked_mean( + student_logps.detach() + .argmax(dim=-1) + .eq(teacher_logps.argmax(dim=-1)) + .float(), + mask, + ) + metrics = { + "divergence": loss.detach(), + "forward_kl": _masked_mean(forward_kl.detach(), mask), + "reverse_kl": _masked_mean(reverse_kl.detach(), mask), + "top1_agreement": top1_agreement.detach(), + "valid_tokens": mask.sum().detach(), + } + return loss, metrics + + +def use_on_policy_batch(lmbda: float, random_value: float) -> bool: + """Return the Algorithm 1 branch for a pre-sampled value in ``[0, 1)``.""" + if not 0.0 <= lmbda <= 1.0: + raise ValueError(f"lmbda must be in [0, 1], got {lmbda}") + if not 0.0 <= random_value < 1.0: + raise ValueError(f"random_value must be in [0, 1), got {random_value}") + return random_value < lmbda diff --git a/trainer/rollout_engine.py b/trainer/rollout_engine.py index 3a1aab0..66d731b 100644 --- a/trainer/rollout_engine.py +++ b/trainer/rollout_engine.py @@ -52,7 +52,10 @@ class RolloutEngine(ABC): tokenizer = None @abstractmethod - def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult: + def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, + max_new_tokens: int, temperature: float = 0.8, + top_p: Optional[float] = None, top_k: Optional[int] = None, + calculate_logps: bool = True) -> RolloutResult: pass @abstractmethod @@ -68,9 +71,15 @@ class TorchRolloutEngine(RolloutEngine): self.device = device self.autocast_ctx = autocast_ctx - def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult: + def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, + max_new_tokens: int, temperature: float = 0.8, + top_p: Optional[float] = None, top_k: Optional[int] = None, + calculate_logps: bool = True) -> RolloutResult: model = self.policy_model.module if isinstance(self.policy_model, DistributedDataParallel) else self.policy_model ctx = self.autocast_ctx if self.autocast_ctx else nullcontext() + sampling_kwargs = {} + if top_p is not None: sampling_kwargs['top_p'] = top_p + if top_k is not None: sampling_kwargs['top_k'] = top_k with torch.no_grad(), ctx: output_ids = model.generate( input_ids=prompt_ids.repeat_interleave(num_generations, dim=0), @@ -81,11 +90,20 @@ class TorchRolloutEngine(RolloutEngine): num_return_sequences=1, pad_token_id=self.tokenizer.pad_token_id, eos_token_id=self.tokenizer.eos_token_id, + **sampling_kwargs, ).clone() # [B*num_gen, P+R] prompt_len = prompt_ids.size(1) completion_ids = output_ids[:, prompt_len:] # [B*num_gen, R] full_mask = (output_ids != self.tokenizer.pad_token_id).long() - per_token_logps = compute_per_token_logps(self.policy_model, output_ids, completion_ids.size(1), attention_mask=full_mask) + if calculate_logps: + per_token_logps = compute_per_token_logps( + self.policy_model, output_ids, completion_ids.size(1), attention_mask=full_mask + ) + else: + per_token_logps = torch.zeros( + output_ids.size(0), completion_ids.size(1), + device=output_ids.device, dtype=torch.float32 + ) completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True) return RolloutResult(output_ids, completion_ids, per_token_logps, completions, prompt_ids.new_full((output_ids.size(0),), prompt_len), @@ -104,7 +122,10 @@ class SGLangRolloutEngine(RolloutEngine): self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) self.http = requests - def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult: + def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, + max_new_tokens: int, temperature: float = 0.8, + top_p: Optional[float] = None, top_k: Optional[int] = None, + calculate_logps: bool = True) -> RolloutResult: # 去除左侧 padding tokens，只保留有效 token input_ids_list = [] for ids, mask in zip(prompt_ids, attention_mask): @@ -119,8 +140,10 @@ class SGLangRolloutEngine(RolloutEngine): "max_new_tokens": max_new_tokens, "stop_token_ids": [self.tokenizer.eos_token_id] if self.tokenizer.eos_token_id else [], }, - "return_logprob": True, + "return_logprob": calculate_logps, } + if top_p is not None: payload['sampling_params']['top_p'] = top_p + if top_k is not None: payload['sampling_params']['top_k'] = top_k if top_k > 0 else -1 resp = self.http.post(f"{self.base_url}/generate", json=payload, timeout=self.timeout) resp.raise_for_status() diff --git a/trainer/train_opd.py b/trainer/train_opd.py new file mode 100644 index 0000000..5629af7 --- /dev/null +++ b/trainer/train_opd.py @@ -0,0 +1,755 @@ +import os +import sys + +__package__ = "trainer" +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771) +import argparse +import gc +import json +import math +import random +import warnings +from contextlib import nullcontext +from pathlib import Path + +import torch +import torch.distributed as dist +from datasets import load_dataset +from torch import optim +from torch.nn.parallel import DistributedDataParallel +from torch.optim.lr_scheduler import CosineAnnealingLR +from torch.utils.data import DataLoader, Dataset, DistributedSampler +from transformers import AutoTokenizer + +from model.model_minimind import MiniMindConfig, MiniMindForCausalLM +from trainer.opd_utils import ( + build_completion_mask, + generalized_jsd_loss, + use_on_policy_batch, +) +from trainer.rollout_engine import create_rollout_engine +from trainer.trainer_utils import ( + Logger, + SkipBatchSampler, + init_distributed_mode, + is_main_process, + lm_checkpoint, + setup_seed, +) + +warnings.filterwarnings("ignore") +REPO_ROOT = Path(__file__).resolve().parents[1] + + +class OPDDataset(Dataset): + """Raw conversational samples used by both GKD data branches.""" + + def __init__(self, data_path): + self.samples = load_dataset("json", data_files=str(data_path), split="train") + + def __len__(self): + return len(self.samples) + + def __getitem__(self, index): + sample = self.samples[index] + conversations = sample.get("conversations", sample.get("messages")) + if not isinstance(conversations, list) or not conversations: + raise ValueError( + "OPD data requires a non-empty 'conversations' or 'messages' list" + ) + return {"conversations": [dict(message) for message in conversations]} + + +def collate_conversations(batch): + return {"conversations": [item["conversations"] for item in batch]} + + +def resolve_device(device): + if device != "auto": + return device + if torch.cuda.is_available(): + return "cuda:0" + if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): + return "mps" + return "cpu" + + +def resolve_path(value): + path = Path(value).expanduser() + if path.is_absolute(): + return path.resolve() + return (REPO_ROOT / path).resolve() + + +def validate_args(parsed_args): + if not 0.0 <= parsed_args.lmbda <= 1.0: + raise ValueError("--lmbda must be in [0, 1]") + if not 0.0 <= parsed_args.beta <= 1.0: + raise ValueError("--beta must be in [0, 1]") + if parsed_args.temperature <= 0 or parsed_args.distill_temperature <= 0: + raise ValueError("sampling and distillation temperatures must be positive") + if not 0.0 < parsed_args.rollout_top_p <= 1.0: + raise ValueError("--rollout_top_p must be in (0, 1]") + if parsed_args.rollout_top_k < 0: + raise ValueError("--rollout_top_k must be zero (disabled) or positive") + if parsed_args.num_generations < 1 or parsed_args.max_gen_len < 1: + raise ValueError("--num_generations and --max_gen_len must be positive") + if parsed_args.max_seq_len < 1 or parsed_args.accumulation_steps < 1: + raise ValueError("--max_seq_len and --accumulation_steps must be positive") + if parsed_args.max_train_steps < 0: + raise ValueError("--max_train_steps must be zero or positive") + if parsed_args.log_interval < 1 or parsed_args.save_interval < 1: + raise ValueError("log/save intervals must be positive") + if parsed_args.rollout_top_p < 1.0 or parsed_args.rollout_top_k > 0: + Logger( + "[OPD WARNING] top-p/top-k truncation changes the behavior policy; " + "top_p=1 and top_k=0 match the GKD baseline" + ) + + +def _torch_load(path): + try: + return torch.load(path, map_location="cpu", weights_only=True) + except TypeError: + return torch.load(path, map_location="cpu") + + +def _load_transformers_state_dict(model_dir): + single_safetensors = model_dir / "model.safetensors" + safetensors_index = model_dir / "model.safetensors.index.json" + pytorch_bin = model_dir / "pytorch_model.bin" + if single_safetensors.exists() or safetensors_index.exists(): + try: + from safetensors.torch import load_file + except ImportError as exc: + raise ImportError( + "Loading a Transformers safetensors model requires safetensors" + ) from exc + if single_safetensors.exists(): + return load_file(str(single_safetensors), device="cpu") + if safetensors_index.exists(): + index = json.loads(safetensors_index.read_text(encoding="utf-8")) + state_dict = {} + for shard_name in sorted(set(index["weight_map"].values())): + state_dict.update(load_file(str(model_dir / shard_name), device="cpu")) + return state_dict + if pytorch_bin.exists(): + return _torch_load(pytorch_bin) + raise FileNotFoundError( + f"No model.safetensors, model.safetensors.index.json, or pytorch_model.bin in {model_dir}" + ) + + +def resolve_model_source(source, config): + direct = Path(source).expanduser() + if not direct.is_absolute(): + direct = REPO_ROOT / direct + if direct.exists(): + return direct.resolve() + moe_suffix = "_moe" if config.use_moe else "" + weight_path = REPO_ROOT / "out" / f"{source}_{config.hidden_size}{moe_suffix}.pth" + if weight_path.exists(): + return weight_path.resolve() + raise FileNotFoundError( + f"Model source {source!r} was not found directly or as {weight_path}" + ) + + +def load_state_dict(source, config): + source = resolve_model_source(source, config) + state_dict = ( + _load_transformers_state_dict(source) if source.is_dir() else _torch_load(source) + ) + if isinstance(state_dict, dict) and isinstance(state_dict.get("model"), dict): + state_dict = state_dict["model"] + if not isinstance(state_dict, dict) or not state_dict: + raise ValueError(f"Invalid or empty state dict from {source}") + for prefix in ("module.", "_orig_mod."): + if all(key.startswith(prefix) for key in state_dict): + state_dict = {key[len(prefix):]: value for key, value in state_dict.items()} + if "lm_head.weight" not in state_dict and "model.embed_tokens.weight" in state_dict: + state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"] + if "model.embed_tokens.weight" not in state_dict and "lm_head.weight" in state_dict: + state_dict["model.embed_tokens.weight"] = state_dict["lm_head.weight"] + return state_dict, source + + +def validate_directory_config(source, config, label): + source_path = Path(source).expanduser() + if not source_path.is_absolute(): + source_path = REPO_ROOT / source_path + config_path = source_path / "config.json" if source_path.is_dir() else None + if not config_path or not config_path.exists(): + return + saved = json.loads(config_path.read_text(encoding="utf-8")) + expected = { + "hidden_size": config.hidden_size, + "num_hidden_layers": config.num_hidden_layers, + "vocab_size": config.vocab_size, + "num_attention_heads": config.num_attention_heads, + "num_key_value_heads": config.num_key_value_heads, + } + mismatches = { + key: (saved.get(key), value) + for key, value in expected.items() + if saved.get(key) is not None and saved.get(key) != value + } + if mismatches: + raise ValueError(f"{label} config does not match CLI architecture: {mismatches}") + + +def load_native_model(config, source, device, label, model_dtype=None): + validate_directory_config(source, config, label) + state_dict, resolved_source = load_state_dict(source, config) + loaded_model = MiniMindForCausalLM(config) + loaded_model.load_state_dict(state_dict, strict=True) + del state_dict + gc.collect() + if model_dtype is None: + loaded_model = loaded_model.to(device) + else: + loaded_model = loaded_model.to(device=device, dtype=model_dtype) + Logger(f"Loaded {label} strictly from {resolved_source}") + return loaded_model + + +def load_tokenizer(tokenizer_path, vocab_size): + tokenizer = AutoTokenizer.from_pretrained( + str(resolve_path(tokenizer_path)), trust_remote_code=True + ) + if len(tokenizer) != vocab_size: + raise ValueError( + f"Tokenizer size {len(tokenizer)} does not match model vocab {vocab_size}" + ) + if tokenizer.pad_token_id is None or tokenizer.eos_token_id is None: + raise ValueError("OPD requires tokenizer.pad_token_id and tokenizer.eos_token_id") + tokenizer.padding_side = "left" + tokenizer.truncation_side = "left" + return tokenizer + + +def _parse_json_field(value): + if not isinstance(value, str): + return value + try: + return json.loads(value) + except json.JSONDecodeError: + return value + + +def render_sample(tokenizer, conversations): + """Render one prompt and its optional fixed completion.""" + messages = [] + tools = None + for raw_message in conversations: + message = {key: value for key, value in dict(raw_message).items() if value is not None} + if message.get("tools"): + parsed_tools = _parse_json_field(message["tools"]) + if isinstance(parsed_tools, list): + tools = parsed_tools + if message.get("tool_calls"): + message["tool_calls"] = _parse_json_field(message["tool_calls"]) + messages.append(message) + + has_completion = messages[-1].get("role") == "assistant" and any( + messages[-1].get(field) for field in ("content", "reasoning_content", "tool_calls") + ) + prompt_messages = messages[:-1] if messages[-1].get("role") == "assistant" else messages + template_kwargs = {"tools": tools} if tools else {} + prompt_text = tokenizer.apply_chat_template( + prompt_messages, + tokenize=False, + add_generation_prompt=True, + **template_kwargs, + ) + full_text = None + if has_completion: + full_text = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=False, + **template_kwargs, + ) + return prompt_text, full_text + + +def tokenize_prompts(tokenizer, prompt_texts, max_seq_len, device): + encoded = tokenizer( + prompt_texts, + add_special_tokens=False, + return_tensors="pt", + padding=True, + truncation=True, + max_length=max_seq_len, + return_token_type_ids=False, + ) + return {key: value.to(device) for key, value in encoded.items()} + + +def build_off_policy_batch( + tokenizer, prompt_texts, full_texts, max_seq_len, max_gen_len, device +): + """Tokenize fixed prompt-completion pairs and mask all prompt tokens.""" + sequences = [] + label_masks = [] + for prompt_text, full_text in zip(prompt_texts, full_texts): + if full_text is None: + raise ValueError( + "Off-policy GKD selected a sample without a fixed assistant completion; " + "use --lmbda 1 for prompt-only data" + ) + prompt_ids = tokenizer(prompt_text, add_special_tokens=False).input_ids + full_ids = tokenizer(full_text, add_special_tokens=False).input_ids + if full_ids[:len(prompt_ids)] != prompt_ids: + raise ValueError( + "The rendered prompt is not a token prefix of the fixed conversation; " + "the tokenizer/chat template is incompatible with prompt-completion GKD" + ) + completion_ids = full_ids[len(prompt_ids):] + prompt_ids = prompt_ids[-max_seq_len:] + completion_ids = completion_ids[:max_gen_len] + if not completion_ids: + raise ValueError("A fixed completion became empty after tokenization/truncation") + sequences.append(prompt_ids + completion_ids) + label_masks.append([False] * len(prompt_ids) + [True] * len(completion_ids)) + + max_length = max(len(sequence) for sequence in sequences) + input_ids = [] + attention_mask = [] + labels = [] + for sequence, label_mask in zip(sequences, label_masks): + padding = max_length - len(sequence) + input_ids.append(sequence + [tokenizer.pad_token_id] * padding) + attention_mask.append([1] * len(sequence) + [0] * padding) + labels.append( + [token if keep else -100 for token, keep in zip(sequence, label_mask)] + + [-100] * padding + ) + return ( + torch.tensor(input_ids, dtype=torch.long, device=device), + torch.tensor(attention_mask, dtype=torch.long, device=device), + torch.tensor(labels, dtype=torch.long, device=device), + ) + + +def distributed_mean(value): + result = value.detach().float().clone() + if dist.is_initialized(): + dist.all_reduce(result, op=dist.ReduceOp.SUM) + result /= dist.get_world_size() + return result.item() + + +def choose_on_policy(lmbda, device): + if dist.is_initialized(): + draw = torch.zeros((), device=device) + if dist.get_rank() == 0: + draw.fill_(random.random()) + dist.broadcast(draw, src=0) + random_value = draw.item() + else: + random_value = random.random() + return use_on_policy_batch(lmbda, random_value) + + +def save_student(lm_config, epoch, step, wandb): + if not is_main_process(): + return + model.eval() + raw_model = model.module if isinstance(model, DistributedDataParallel) else model + raw_model = getattr(raw_model, "_orig_mod", raw_model) + moe_suffix = "_moe" if lm_config.use_moe else "" + output_path = resolve_path(args.save_dir) / ( + f"{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth" + ) + state_dict = raw_model.state_dict() + torch.save({key: value.half().cpu() for key, value in state_dict.items()}, output_path) + lm_checkpoint( + lm_config, + weight=args.save_weight, + model=model, + optimizer=optimizer, + epoch=epoch, + step=step, + wandb=wandb, + save_dir=str(resolve_path(args.checkpoint_dir)), + scheduler=scheduler, + scaler=scaler, + ) + model.train() + del state_dict + + +def gkd_train_epoch( + epoch, loader, iters, rollout_engine, teacher_model, lm_config, + start_step=0, wandb=None +): + teacher_model.eval().requires_grad_(False) + model.train() + on_policy_batches = 0 + processed_batches = 0 + + for step, batch in enumerate(loader, start=start_step + 1): + if args.max_train_steps > 0 and processed_batches >= args.max_train_steps: + break + processed_batches += 1 + on_policy = choose_on_policy(args.lmbda, args.device) + on_policy_batches += int(on_policy) + rendered = [render_sample(tokenizer, item) for item in batch["conversations"]] + prompt_texts = [item[0] for item in rendered] + full_texts = [item[1] for item in rendered] + + if on_policy: + prompt_inputs = tokenize_prompts( + tokenizer, prompt_texts, args.max_seq_len, args.device + ) + model.eval() + rollout_result = rollout_engine.rollout( + prompt_ids=prompt_inputs["input_ids"], + attention_mask=prompt_inputs["attention_mask"], + num_generations=args.num_generations, + max_new_tokens=args.max_gen_len, + temperature=args.temperature, + top_p=args.rollout_top_p, + top_k=args.rollout_top_k, + calculate_logps=False, + ) + model.train() + input_ids = rollout_result.output_ids.to(args.device) + completion_ids = rollout_result.completion_ids.to(args.device) + completion_mask = build_completion_mask( + completion_ids, + rollout_result.completion_mask.to(args.device), + eos_token_id=tokenizer.eos_token_id, + pad_token_id=tokenizer.pad_token_id, + ) + prompt_mask = prompt_inputs["attention_mask"].repeat_interleave( + args.num_generations, dim=0 + ) + attention_mask = torch.cat( + (prompt_mask, completion_mask.to(prompt_mask.dtype)), dim=1 + ) + with autocast_ctx: + student_output = model( + input_ids, + attention_mask=attention_mask, + logits_to_keep=completion_ids.size(1) + 1, + ) + student_logits = student_output.logits[:, :-1, :] + with torch.no_grad(), autocast_ctx: + teacher_logits = teacher_model( + input_ids, + attention_mask=attention_mask, + logits_to_keep=completion_ids.size(1) + 1, + ).logits[:, :-1, :] + avg_completion_length = completion_mask.sum(dim=1).float().mean() + debug_completion = rollout_result.completions[0] + else: + input_ids, attention_mask, labels = build_off_policy_batch( + tokenizer, + prompt_texts, + full_texts, + args.max_seq_len, + args.max_gen_len, + args.device, + ) + completion_mask = labels[:, 1:].ne(-100) + with autocast_ctx: + student_output = model(input_ids, attention_mask=attention_mask) + student_logits = student_output.logits[:, :-1, :] + with torch.no_grad(), autocast_ctx: + teacher_logits = teacher_model( + input_ids, attention_mask=attention_mask + ).logits[:, :-1, :] + avg_completion_length = completion_mask.sum(dim=1).float().mean() + debug_completion = full_texts[0] + + gkd_loss, gkd_metrics = generalized_jsd_loss( + student_logits, + teacher_logits, + completion_mask, + beta=args.beta, + temperature=args.distill_temperature, + ) + aux_loss = student_output.aux_loss if lm_config.use_moe else gkd_loss.new_zeros(()) + total_loss = gkd_loss + aux_loss + scaler.scale(total_loss / args.accumulation_steps).backward() + + should_step = step % args.accumulation_steps == 0 or processed_batches == iters + if should_step: + scaler.unscale_(optimizer) + if args.grad_clip > 0: + torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + scaler.step(optimizer) + scaler.update() + scheduler.step() + optimizer.zero_grad(set_to_none=True) + + if args.debug_mode and is_main_process() and step % args.debug_interval == 0: + Logger(f"[OPD DEBUG] on_policy={on_policy}, prompt={prompt_texts[0]!r}") + Logger(f"[OPD DEBUG] completion={debug_completion!r}") + + if step % args.log_interval == 0 or processed_batches == iters: + values = { + "loss/gkd": distributed_mean(gkd_loss), + "loss/aux": distributed_mean(aux_loss), + "loss/total": distributed_mean(total_loss), + "distill/forward_kl": distributed_mean(gkd_metrics["forward_kl"]), + "distill/reverse_kl": distributed_mean(gkd_metrics["reverse_kl"]), + "distill/top1_agreement": distributed_mean(gkd_metrics["top1_agreement"]), + "distill/valid_tokens": distributed_mean(gkd_metrics["valid_tokens"]), + "distill/on_policy": float(on_policy), + "distill/on_policy_fraction": on_policy_batches / processed_batches, + "rollout/avg_completion_length": distributed_mean(avg_completion_length), + "train/learning_rate": optimizer.param_groups[0]["lr"], + } + Logger( + f"Epoch:[{epoch + 1}/{args.epochs}]({step}/{start_step + iters}), " + f"GKD Loss: {values['loss/gkd']:.4f}, " + f"Forward KL: {values['distill/forward_kl']:.4f}, " + f"Reverse KL: {values['distill/reverse_kl']:.4f}, " + f"Top1 Agree: {values['distill/top1_agreement']:.4f}, " + f"On Policy: {int(on_policy)}, LR: {values['train/learning_rate']:.8f}" + ) + if wandb and is_main_process(): + wandb.log(values, step=epoch * iters + step) + + if should_step and (step % args.save_interval == 0 or processed_batches == iters): + save_student(lm_config, epoch, step, wandb) + + del input_ids, attention_mask, completion_mask + del student_output, student_logits, teacher_logits + del gkd_loss, gkd_metrics, aux_loss, total_loss + if on_policy: + del prompt_inputs, completion_ids, prompt_mask, rollout_result + else: + del labels + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="MiniMind Generalized/On-Policy Knowledge Distillation" + ) + parser.add_argument("--data_path", type=str, default="dataset/rlaif.jsonl") + parser.add_argument("--student_model", type=str, default="full_sft") + parser.add_argument("--teacher_model", type=str, default="full_sft") + parser.add_argument("--tokenizer_path", type=str, default="model") + parser.add_argument("--save_dir", type=str, default="out") + parser.add_argument("--checkpoint_dir", type=str, default="checkpoints") + parser.add_argument("--save_weight", type=str, default="opd") + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch_size", type=int, default=2) + parser.add_argument("--learning_rate", type=float, default=1e-6) + parser.add_argument("--device", type=str, default="auto") + parser.add_argument( + "--dtype", type=str, default="bfloat16", + choices=["float32", "float16", "bfloat16"] + ) + parser.add_argument("--num_workers", type=int, default=4) + parser.add_argument("--accumulation_steps", type=int, default=1) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--log_interval", type=int, default=1) + parser.add_argument("--save_interval", type=int, default=100) + parser.add_argument("--max_seq_len", type=int, default=768) + parser.add_argument("--max_gen_len", type=int, default=256) + parser.add_argument("--max_train_steps", type=int, default=0) + parser.add_argument( + "--lmbda", type=float, default=1.0, + help="fraction of student-generated batches; 1.0 is pure OPD" + ) + parser.add_argument( + "--beta", type=float, default=0.5, + help="generalized JSD coefficient; 0=forward KL, 1=reverse KL" + ) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--distill_temperature", type=float, default=1.0) + parser.add_argument("--rollout_top_p", type=float, default=1.0) + parser.add_argument("--rollout_top_k", type=int, default=0) + parser.add_argument("--num_generations", type=int, default=1) + parser.add_argument("--student_hidden_size", type=int, default=768) + parser.add_argument("--student_num_layers", type=int, default=8) + parser.add_argument("--student_use_moe", type=int, default=0, choices=[0, 1]) + parser.add_argument("--teacher_hidden_size", type=int, default=768) + parser.add_argument("--teacher_num_layers", type=int, default=8) + parser.add_argument("--teacher_use_moe", type=int, default=1, choices=[0, 1]) + parser.add_argument("--from_resume", type=int, default=0, choices=[0, 1]) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--use_wandb", action="store_true") + parser.add_argument("--wandb_project", type=str, default="MiniMind-OPD") + parser.add_argument("--wandb_run_name", type=str, default="") + parser.add_argument( + "--wandb_mode", type=str, default="cloud", + choices=["cloud", "local", "disabled"] + ) + parser.add_argument("--wandb_logdir", type=str, default="swanlog") + parser.add_argument("--use_compile", type=int, default=0, choices=[0, 1]) + parser.add_argument("--debug_mode", action="store_true") + parser.add_argument("--debug_interval", type=int, default=20) + args = parser.parse_args() + validate_args(args) + + local_rank = init_distributed_mode() + if dist.is_initialized(): + args.device = f"cuda:{local_rank}" + else: + args.device = resolve_device(args.device) + rank = dist.get_rank() if dist.is_initialized() else 0 + setup_seed(args.seed + rank) + resolve_path(args.save_dir).mkdir(parents=True, exist_ok=True) + resolve_path(args.checkpoint_dir).mkdir(parents=True, exist_ok=True) + + max_position_embeddings = max(32768, args.max_seq_len + args.max_gen_len) + student_config = MiniMindConfig( + hidden_size=args.student_hidden_size, + num_hidden_layers=args.student_num_layers, + use_moe=bool(args.student_use_moe), + max_position_embeddings=max_position_embeddings, + ) + teacher_config = MiniMindConfig( + hidden_size=args.teacher_hidden_size, + num_hidden_layers=args.teacher_num_layers, + use_moe=bool(args.teacher_use_moe), + max_position_embeddings=max_position_embeddings, + ) + if student_config.vocab_size != teacher_config.vocab_size: + raise ValueError("GKD requires student and teacher to share the same vocabulary") + tokenizer = load_tokenizer(args.tokenizer_path, student_config.vocab_size) + + device_type = args.device.split(":", 1)[0] + mixed_dtype = { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + }[args.dtype] + if device_type == "mps" and mixed_dtype == torch.bfloat16: + Logger("[OPD WARNING] MPS uses float16 for the frozen teacher") + mixed_dtype = torch.float16 + teacher_dtype = mixed_dtype if device_type in {"cuda", "mps"} else torch.float32 + autocast_ctx = ( + torch.cuda.amp.autocast(dtype=mixed_dtype) + if device_type == "cuda" and mixed_dtype != torch.float32 + else nullcontext() + ) + + model = load_native_model( + student_config, args.student_model, args.device, "student" + ) + teacher_model = load_native_model( + teacher_config, + args.teacher_model, + args.device, + "teacher", + model_dtype=teacher_dtype, + ).eval().requires_grad_(False) + Logger(f"Student params: {sum(p.numel() for p in model.parameters()) / 1e6:.3f}M") + Logger(f"Teacher params: {sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f}M") + + checkpoint = ( + lm_checkpoint( + student_config, + weight=args.save_weight, + save_dir=str(resolve_path(args.checkpoint_dir)), + ) + if args.from_resume == 1 else None + ) + train_ds = OPDDataset(resolve_path(args.data_path)) + train_sampler = ( + DistributedSampler(train_ds, shuffle=True, seed=args.seed) + if dist.is_initialized() else None + ) + count_loader = DataLoader( + train_ds, + batch_size=args.batch_size, + sampler=train_sampler, + collate_fn=collate_conversations, + ) + full_iters = len(count_loader) + iters = min(full_iters, args.max_train_steps) if args.max_train_steps > 0 else full_iters + optimizer_steps_per_epoch = math.ceil(iters / args.accumulation_steps) + total_optimizer_steps = max(1, optimizer_steps_per_epoch * args.epochs) + optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) + scheduler = CosineAnnealingLR( + optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10 + ) + scaler = torch.cuda.amp.GradScaler( + enabled=(device_type == "cuda" and mixed_dtype == torch.float16) + ) + + start_epoch, start_step = 0, 0 + if checkpoint: + model.load_state_dict(checkpoint["model"], strict=True) + optimizer.load_state_dict(checkpoint["optimizer"]) + if "scheduler" in checkpoint: + scheduler.load_state_dict(checkpoint["scheduler"]) + if "scaler" in checkpoint: + scaler.load_state_dict(checkpoint["scaler"]) + start_epoch = checkpoint["epoch"] + start_step = checkpoint.get("step", 0) + Logger(f"Resuming GKD from epoch={start_epoch}, step={start_step}") + + if args.use_compile == 1: + model = torch.compile(model) + Logger("torch.compile enabled") + if dist.is_initialized(): + model = DistributedDataParallel(model, device_ids=[local_rank]) + rollout_engine = create_rollout_engine( + engine_type="torch", + policy_model=model, + tokenizer=tokenizer, + device=args.device, + autocast_ctx=autocast_ctx, + ) + + wandb = None + if args.use_wandb and is_main_process(): + import swanlab as wandb + wandb_id = checkpoint.get("wandb_id") if checkpoint else None + wandb.init( + project=args.wandb_project, + name=args.wandb_run_name or ( + f"GKD-L{args.lmbda}-B{args.beta}-BS{args.batch_size}x" + f"{args.accumulation_steps}-LR{args.learning_rate}" + ), + config=vars(args), + mode=args.wandb_mode, + logdir=str(resolve_path(args.wandb_logdir)), + id=wandb_id if args.wandb_mode == "cloud" else None, + resume="must" if wandb_id and args.wandb_mode == "cloud" else None, + ) + + optimizer.zero_grad(set_to_none=True) + for epoch in range(start_epoch, args.epochs): + if train_sampler: + train_sampler.set_epoch(epoch) + setup_seed(args.seed + epoch + rank) + indices = torch.randperm(len(train_ds)).tolist() + skip = start_step if epoch == start_epoch and start_step > 0 else 0 + batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) + loader = DataLoader( + train_ds, + batch_sampler=batch_sampler, + num_workers=args.num_workers, + pin_memory=(device_type == "cuda"), + collate_fn=collate_conversations, + ) + if skip > 0: + Logger(f"Epoch [{epoch + 1}/{args.epochs}]: skipping {skip} completed batches") + gkd_train_epoch( + epoch, + loader, + iters, + rollout_engine, + teacher_model, + student_config, + start_step=skip, + wandb=wandb, + ) + start_step = 0 + + if wandb and is_main_process() and hasattr(wandb, "finish"): + wandb.finish() + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group()