Merge 864bd67437 into 6fc918beb6

Expose num_experts / num_experts_per_tok / seed on the trainer CLIs

README.md

@@ -1599,6 +1599,15 @@ MiniMind 的数据规模远小于表中其他模型，且训练比例偏向中
 | [SmolLM2-135M](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct) | HuggingFace | 135M | 24.44 / 24.71 | 58.50 / 68.17 / 32.80 / 43.15 / 39.46 |
 | [Aquila-135M](https://huggingface.co/BAAI/Aquila-135M-Instruct) | BAAI | 135M | 25.19 / 25.10 | 54.59 / 67.52 / 34.40 / 41.67 / 39.66 |
 
+> **这张表的分辨率**：`lm_eval` 对每个指标都会同时输出一个 `stderr`，它主要由测试集大小决定，因此下面的数值对表中所有模型都近似适用（随准确率略有浮动）：
+>
+> | 数据集 | obqa | ceval | piqa | siqa | arc | hellaswag | cmmlu |
+> |---|---|---|---|---|---|---|---|
+> | 样本量 | 500 | ~1.3k | 1838 | 1954 | 2376 | 10042 | 11582 |
+> | ± 标准误 | **1.9** | 1.2 | 1.2 | 1.1 | 0.9 | 0.5 | **0.4** |
+>
+> 比较两个模型时，差值的标准误约为单项的 `√2` 倍。因此 obqa 上小于约 5 个百分点、ceval / piqa / siqa 上小于约 3 个百分点的差异，都不足以区分两个模型——这也解释了为什么本表中同一量级的模型排序不宜过度解读。
+
 <details>
 <summary><strong>补充说明（来源/无污染/复现）</strong></summary>
 

README_en.md

@@ -1597,6 +1597,15 @@ MiniMind is trained on far less data than the other models listed here, and its
 | [SmolLM2-135M](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct) | HuggingFace | 135M | 24.44 / 24.71 | 58.50 / 68.17 / 32.80 / 43.15 / 39.46 |
 | [Aquila-135M](https://huggingface.co/BAAI/Aquila-135M-Instruct) | BAAI | 135M | 25.19 / 25.10 | 54.59 / 67.52 / 34.40 / 41.67 / 39.66 |
 
+> **The resolution of this table**: `lm_eval` reports a `stderr` alongside every metric. It is set mainly by the size of the test set, so the values below apply approximately to every model in the table (they shift slightly with accuracy):
+>
+> | dataset | obqa | ceval | piqa | siqa | arc | hellaswag | cmmlu |
+> |---|---|---|---|---|---|---|---|
+> | samples | 500 | ~1.3k | 1838 | 1954 | 2376 | 10042 | 11582 |
+> | ± stderr | **1.9** | 1.2 | 1.2 | 1.1 | 0.9 | 0.5 | **0.4** |
+>
+> When comparing two models, the standard error of the difference is about `√2` times the single-model value. So a gap smaller than roughly 5 points on obqa, or roughly 3 points on ceval / piqa / siqa, does not separate two models — which is why the ordering of similarly-sized models in this table should not be over-read.
+
 <details>
 <summary><strong>Additional note (source / no contamination / reproduction)</strong></summary>
 

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

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()

trainer/train_full_sft.py

@@ -99,6 +99,7 @@ if __name__ == "__main__":
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=768, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--seed', default=42, type=int, help="随机种子（DDP下每个rank为seed+rank，每轮为seed+epoch）")
     parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl", help="训练数据路径")
     parser.add_argument('--from_weight', default='pretrain', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
@@ -110,7 +111,7 @@ if __name__ == "__main__":
     # ========== 1. 初始化环境和随机种子 ==========
     local_rank = init_distributed_mode()
     if dist.is_initialized(): args.device = f"cuda:{local_rank}"
-    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+    setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0))
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
@@ -157,7 +158,7 @@ if __name__ == "__main__":
     # ========== 8. 开始训练 ==========
     for epoch in range(start_epoch, args.epochs):
         train_sampler and train_sampler.set_epoch(epoch)
-        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        setup_seed(args.seed + epoch); 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=True)

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()

trainer/train_pretrain.py

@@ -98,6 +98,7 @@ if __name__ == "__main__":
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--seed', default=42, type=int, help="随机种子（DDP下每个rank为seed+rank，每轮为seed+epoch）")
     parser.add_argument("--data_path", type=str, default="../dataset/pretrain_t2t_mini.jsonl", help="预训练数据路径")
     parser.add_argument('--from_weight', default='none', type=str, help="基于哪个权重训练，为none则从头开始")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
@@ -109,7 +110,7 @@ if __name__ == "__main__":
     # ========== 1. 初始化环境和随机种子 ==========
     local_rank = init_distributed_mode()
     if dist.is_initialized(): args.device = f"cuda:{local_rank}"
-    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+    setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0))
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
@@ -156,7 +157,7 @@ if __name__ == "__main__":
     # ========== 8. 开始训练 ==========
     for epoch in range(start_epoch, args.epochs):
         train_sampler and train_sampler.set_epoch(epoch)
-        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        setup_seed(args.seed + epoch); 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=True)