[feat] add MLA with decoupled RoPE (DeepSeek-V2 style)

- Add model_minimind_mla.py with Q/KV low-rank compression and decoupled RoPE
- Add kv_norm on kv_latent, q_compress + q_up + q_rope_proj for Q path
- Cache format: (kv_latent, k_rope) tuple, 176 floats/tok/layer (4.4x compression)
- Add --use_mla flag to all training scripts and utils
- Add benchmark_gqa_vs_mla.py script

eval_llm.py

@@ -5,6 +5,7 @@ import warnings
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
 from model.model_lora import *
 from trainer.trainer_utils import setup_seed, get_model_params
 warnings.filterwarnings('ignore')
@@ -12,14 +13,25 @@ warnings.filterwarnings('ignore')
 def init_model(args):
     tokenizer = AutoTokenizer.from_pretrained(args.load_from)
     if 'model' in args.load_from:
+        if args.use_mla:
+            config = MiniMindMLAConfig(
+                hidden_size=args.hidden_size,
+                num_hidden_layers=args.num_hidden_layers,
+                use_moe=bool(args.use_moe),
+                inference_rope_scaling=args.inference_rope_scaling,
+                kv_lora_rank=args.kv_lora_rank
+            )
+            model = MiniMindMLAForCausalLM(config)
+        else:
             model = MiniMindForCausalLM(MiniMindConfig(
                 hidden_size=args.hidden_size,
                 num_hidden_layers=args.num_hidden_layers,
                 use_moe=bool(args.use_moe),
                 inference_rope_scaling=args.inference_rope_scaling
             ))
-        moe_suffix = '_moe' if args.use_moe else ''
+        model_suffix = '_moe' if args.use_moe else ''
-        ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
+        if args.use_mla: model_suffix += '_mla'
+        ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{model_suffix}.pth'
         model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
         if args.lora_weight != 'None':
             apply_lora(model)
@@ -38,6 +50,8 @@ def main():
     parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     parser.add_argument('--inference_rope_scaling', default=False, action='store_true', help="启用RoPE位置编码外推（4倍，仅解决位置编码问题）")
     parser.add_argument('--max_new_tokens', default=8192, type=int, help="最大生成长度（注意：并非模型实际长文本能力）")
     parser.add_argument('--temperature', default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
@@ -62,8 +76,10 @@ def main():
     conversation = []
     model, tokenizer = init_model(args)
     input_mode = int(input('[0] 自动测试\n[1] 手动输入\n'))
+    #TextStreamer是transformers里的文本流式输出器，在model.generate生成过程中，边生成边把文本输出到终端
+    #tokenizer用来把tokne转成文本，skip_prompt=True不把输入提示词也打印出来，skip_special_tokens=True不打印特殊token
     streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    
+    #iter是python自带的迭代器，将一个函数包装成可迭代对象，第一个参数是一个函数，第二个参数是结束值，当input的输入值为''时结束迭代
     prompt_iter = prompts if input_mode == 0 else iter(lambda: input('💬: '), '')
     for prompt in prompt_iter:
         setup_seed(random.randint(0, 31415926))
@@ -71,10 +87,11 @@ def main():
         conversation = conversation[-args.historys:] if args.historys else []
         conversation.append({"role": "user", "content": prompt})
         if 'pretrain' in args.weight:
+            print(tokenizer.bos_token)
             inputs = tokenizer.bos_token + prompt
         else:
             inputs = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True, open_thinking=bool(args.open_thinking))
-        
+        #return_tensors=‘pt'是返回pytorch张量
         inputs = tokenizer(inputs, return_tensors="pt", truncation=True).to(args.device)
 
         print('🧠: ', end='')

model/model_minimind_mla.py

@@ -0,0 +1,241 @@
+"""
+MiniMind with Multi-head Latent Attention (MLA)
+基于 DeepSeek-V2 的低秩 KV 压缩注意力机制（解耦 RoPE 架构）
+将 GQA 中独立的 K/V 投影替换为"压缩→缓存→还原"的低秩路径
+位置信息通过独立的低维投影承载（decoupled rope），K 的 content 部分不施 RoPE
+KV-Cache 从 GQA 的 768 floats/token/layer 压缩到 kv_lora_rank + rope_dim floats（默认 128+48=176，压缩约 4.4 倍）
+"""
+import math, torch, torch.nn.functional as F
+from torch import nn
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+from model.model_minimind import (
+    MiniMindConfig, RMSNorm, precompute_freqs_cis, repeat_kv,
+    FeedForward, MOEFeedForward
+)
+
+
+class MiniMindMLAConfig(MiniMindConfig):
+    model_type = "minimind_mla"
+
+    def __init__(self, kv_lora_rank=128, q_lora_rank=256, rope_dim=None, **kwargs):
+        super().__init__(**kwargs)
+        self.kv_lora_rank = kv_lora_rank
+        self.q_lora_rank = q_lora_rank
+        # hidden_size=768 → head_dim=96 → rope_dim=48
+        self.rope_dim = rope_dim if rope_dim is not None else self.hidden_size // self.num_attention_heads // 2
+
+
+def _apply_rope(x, cos, sin, unsqueeze_dim=1):
+    """单张量 RoPE，MLA 中 Q 和 K 序列长度不同，需要分别应用"""
+    def rotate_half(x): return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
+    return ((x * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(x) * sin.unsqueeze(unsqueeze_dim))).to(x.dtype)
+
+
+class MLAAttention(nn.Module):
+    def __init__(self, config: MiniMindMLAConfig):
+        super().__init__()
+        self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
+        self.n_local_heads = config.num_attention_heads
+        self.n_local_kv_heads = self.num_key_value_heads
+        self.n_rep = self.n_local_heads // self.n_local_kv_heads
+        self.head_dim = config.head_dim
+        self.kv_lora_rank = config.kv_lora_rank
+        self.q_lora_rank = config.q_lora_rank
+        self.rope_dim = config.rope_dim
+        self.is_causal = True
+
+        # Q 低秩压缩路径 (DeepSeek-V2 Eq.12-14)
+        self.q_compress = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False)
+        self.q_up = nn.Linear(config.q_lora_rank, config.num_attention_heads * self.head_dim, bias=False)
+        self.q_rope_proj = nn.Linear(config.q_lora_rank, config.num_attention_heads * self.rope_dim, bias=False)
+
+        # KV 低秩压缩路径 (DeepSeek-V2 Eq.9-11)
+        self.kv_compress = nn.Linear(config.hidden_size, config.kv_lora_rank, bias=False)
+        self.kv_up_k = nn.Linear(config.kv_lora_rank, self.num_key_value_heads * self.head_dim, bias=False)
+        self.kv_up_v = nn.Linear(config.kv_lora_rank, self.num_key_value_heads * self.head_dim, bias=False)
+
+        # K 的解耦 RoPE 投影 (DeepSeek-V2 Eq.15)
+        self.k_rope_proj = nn.Linear(config.hidden_size, config.rope_dim, bias=False)
+
+        # 输出投影
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+
+        # 归一化
+        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.kv_norm = RMSNorm(config.kv_lora_rank, eps=config.rms_norm_eps)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.dropout = config.dropout
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn
+
+    def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
+        bsz, seq_len, _ = x.shape
+        cos, sin = position_embeddings  # [seq_len, rope_dim]
+
+        # ── Q Path ──
+        q_latent = self.q_compress(x)  # [bsz, seq, q_lora_rank]
+        # Q content: q_compress → q_up → q_norm (NO RoPE)
+        q_content = self.q_up(q_latent).view(bsz, seq_len, self.n_local_heads, self.head_dim)
+        q_content = self.q_norm(q_content)  # [bsz, seq, 8, 96]
+        # Q rope: q_compress → q_rope_proj → RoPE
+        q_rope = self.q_rope_proj(q_latent).view(bsz, seq_len, self.n_local_heads, self.rope_dim)
+        q_rope = _apply_rope(q_rope, cos, sin)  # [bsz, seq, 8, 48]
+        # 拼接 content + rope
+        xq = torch.cat([q_content, q_rope], dim=-1)  # [bsz, seq, 8, 144]
+
+        # ── K Path ──
+        # Content: kv_compress → norm → cache → kv_up_k (NO RoPE!)
+        kv_latent = self.kv_norm(self.kv_compress(x))  # [bsz, seq, kv_lora_rank]
+        # Rope: k_rope_proj → RoPE (共享跨所有 head)
+        k_rope = self.k_rope_proj(x).unsqueeze(2)  # [bsz, seq, 1, rope_dim]
+        k_rope = _apply_rope(k_rope, cos, sin).squeeze(2)  # [bsz, seq, rope_dim]
+
+        # 拼接缓存
+        if past_key_value is not None:
+            past_kv_latent, past_k_rope = past_key_value
+            kv_latent = torch.cat([past_kv_latent, kv_latent], dim=1)
+            k_rope = torch.cat([past_k_rope, k_rope], dim=1)
+        past_kv = (kv_latent, k_rope) if use_cache else None
+
+        total_len = kv_latent.shape[1]
+        # 从 latent 还原 K content (NO RoPE!)
+        k_content = self.kv_up_k(kv_latent).view(bsz, total_len, self.n_local_kv_heads, self.head_dim)  # [bsz, total, 4, 96]
+        # K rope 扩展 head 维度后拼接
+        k_rope_expanded = k_rope.unsqueeze(2).expand(-1, -1, self.n_local_kv_heads, -1)  # [bsz, total, 4, rope_dim]
+        xk = torch.cat([k_content, k_rope_expanded], dim=-1)  # [bsz, total, 4, 144]
+
+        # ── V Path ──
+        xv = self.kv_up_v(kv_latent).view(bsz, total_len, self.n_local_kv_heads, self.head_dim)  # [bsz, total, 4, 96]
+
+        # ── Attention ──
+        scale = math.sqrt(self.head_dim + self.rope_dim)
+        xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
+        if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
+            output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)
+        else:
+            scores = (xq @ xk.transpose(-2, -1)) / scale
+            if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
+            if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
+            output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
+        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
+        output = self.resid_dropout(self.o_proj(output))
+        return output, past_kv
+
+
+class MiniMindMLABlock(nn.Module):
+    def __init__(self, layer_id: int, config: MiniMindMLAConfig):
+        super().__init__()
+        self.self_attn = MLAAttention(config)
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
+
+    def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
+        residual = hidden_states
+        hidden_states, present_key_value = self.self_attn(
+            self.input_layernorm(hidden_states), position_embeddings,
+            past_key_value, use_cache, attention_mask
+        )
+        hidden_states += residual
+        hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
+        return hidden_states, present_key_value
+
+
+class MiniMindMLAModel(nn.Module):
+    def __init__(self, config: MiniMindMLAConfig):
+        super().__init__()
+        self.config = config
+        self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.dropout)
+        self.layers = nn.ModuleList([MiniMindMLABlock(l, config) for l in range(self.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.rope_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling)
+        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
+        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+
+    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
+        batch_size, seq_length = input_ids.shape
+        if hasattr(past_key_values, 'layers'): past_key_values = None
+        past_key_values = past_key_values or [None] * len(self.layers)
+        # MLA 的 cache 是 tuple (kv_latent, k_rope)，取 kv_latent 的长度
+        start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
+        hidden_states = self.dropout(self.embed_tokens(input_ids))
+        # Recompute RoPE buffers lost during meta-device init (transformers>=5.x)
+        if self.freqs_cos[0, 0] == 0:
+            freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.rope_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
+            self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
+        # MLA: RoPE 只作用于当前 token 的 rope 投影（Q rope 和 K rope 使用相同位置范围）
+        position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length])
+        presents = []
+        for layer, past_key_value in zip(self.layers, past_key_values):
+            hidden_states, present = layer(
+                hidden_states,
+                position_embeddings,
+                past_key_value=past_key_value,
+                use_cache=use_cache,
+                attention_mask=attention_mask
+            )
+            presents.append(present)
+        hidden_states = self.norm(hidden_states)
+        aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
+        return hidden_states, presents, aux_loss
+
+
+class MiniMindMLAForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MiniMindMLAConfig
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: MiniMindMLAConfig = None):
+        self.config = config or MiniMindMLAConfig()
+        super().__init__(self.config)
+        self.model = MiniMindMLAModel(self.config)
+        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        if self.config.tie_word_embeddings: self.model.embed_tokens.weight = self.lm_head.weight
+        self.post_init()
+
+    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs):
+        hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        loss = None
+        if labels is not None:
+            x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
+            loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
+        return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states)
+
+    @torch.inference_mode()
+    def generate(self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=2, streamer=None, use_cache=True, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs):
+        input_ids = kwargs.pop("input_ids", inputs).repeat(num_return_sequences, 1)
+        attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
+        past_key_values = kwargs.pop("past_key_values", None)
+        finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device)
+        if streamer: streamer.put(input_ids.cpu())
+        for _ in range(max_new_tokens):
+            # MLA: past_key_values[0] 是 tuple (kv_latent, k_rope)，取 kv_latent 的长度
+            past_len = past_key_values[0][0].shape[1] if past_key_values else 0
+            outputs = self.forward(input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, **kwargs)
+            attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
+            logits = outputs.logits[:, -1, :] / temperature
+            if repetition_penalty != 1.0:
+                for i in range(input_ids.shape[0]):
+                    seen = torch.unique(input_ids[i]); score = logits[i, seen]; logits[i, seen] = torch.where(score > 0, score / repetition_penalty, score * repetition_penalty)
+            if top_k > 0:
+                logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
+                mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
+                logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')
+            next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
+            if eos_token_id is not None: next_token = torch.where(finished.unsqueeze(-1), next_token.new_full((next_token.shape[0], 1), eos_token_id), next_token)
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+            past_key_values = outputs.past_key_values if use_cache else None
+            if streamer: streamer.put(next_token.cpu())
+            if eos_token_id is not None:
+                finished |= next_token.squeeze(-1).eq(eos_token_id)
+                if finished.all(): break
+        if streamer: streamer.end()
+        if kwargs.get("return_kv"): return {'generated_ids': input_ids, 'past_kv': past_key_values}
+        return input_ids

scripts/benchmark_gqa_vs_mla.py

@@ -0,0 +1,557 @@
+"""
+GQA vs MLA Benchmark 脚本
+对比 MiniMind GQA 与 MLA (Multi-head Latent Attention) 的性能差异：
+  1. 模型参数量 & KV-Cache 理论大小
+  2. 推理速度 (Prefill / Decode)
+  3. KV-Cache 显存占用
+  4. 训练吞吐量
+  5. 输出质量 (PPL)
+
+Usage:
+    # 性能测试（随机权重）
+    python scripts/benchmark_gqa_vs_mla.py [--device cuda:0] [--hidden_size 768]
+
+    # 加载 checkpoint 测 PPL（必须提供 checkpoint 路径）
+    python scripts/benchmark_gqa_vs_mla.py --device cuda:0 \
+        --gqa_ckpt out/pretrain_768.pth \
+        --mla_ckpt out/pretrain_768_mla.pth \
+        --data_path dataset/pretrain_data.jsonl \
+        --tokenizer_path model
+"""
+
+import argparse
+import gc
+import math
+import os
+import sys
+import time
+
+import torch
+
+sys.path.append(".")
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
+
+
+# ──────────────────────────── helpers ────────────────────────────
+
+def count_params(model):
+    return sum(p.numel() for p in model.parameters())
+
+
+def attention_params_per_layer(model):
+    """统计第一个 transformer block 中 self_attn 的参数量"""
+    block = model.model.layers[0]
+    return sum(p.numel() for p in block.self_attn.parameters())
+
+
+def fmt_num(n):
+    if n >= 1e6:
+        return f"{n / 1e6:.2f}M"
+    if n >= 1e3:
+        return f"{n / 1e3:.2f}K"
+    return str(n)
+
+
+def fmt_mem(bytes_val):
+    if bytes_val >= 1 << 30:
+        return f"{bytes_val / (1 << 30):.2f} GB"
+    if bytes_val >= 1 << 20:
+        return f"{bytes_val / (1 << 20):.2f} MB"
+    return f"{bytes_val:.0f} B"
+
+
+def warmup_forward(model, input_ids, device, n=3):
+    """warmup: forward n 次，不计时"""
+    model.train(False)
+    for _ in range(n):
+        with torch.no_grad():
+            model(input_ids.to(device))
+    torch.cuda.synchronize(device)
+
+
+def cleanup_gpu(device):
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats(device)
+
+
+# ────────────────────── model factories ──────────────────────────
+
+def build_gqa(hidden_size):
+    config = MiniMindConfig(hidden_size=hidden_size)
+    return MiniMindForCausalLM(config)
+
+
+def build_mla(hidden_size, kv_lora_rank):
+    config = MiniMindMLAConfig(hidden_size=hidden_size, kv_lora_rank=kv_lora_rank)
+    return MiniMindMLAForCausalLM(config)
+
+
+# ────────────────────── Benchmark 1: 模型信息 ────────────────────
+
+def benchmark_model_info(hidden_size):
+    print("\n" + "=" * 70)
+    print("  Benchmark 1: 模型信息对比")
+    print("=" * 70)
+
+    models = {
+        "GQA": build_gqa(hidden_size),
+        "MLA(128)": build_mla(hidden_size, 128),
+        "MLA(64)": build_mla(hidden_size, 64),
+        "MLA(256)": build_mla(hidden_size, 256),
+    }
+
+    # GQA: KV-Cache floats/token/layer = 2 * num_kv_heads * head_dim
+    gqa_config = MiniMindConfig(hidden_size=hidden_size)
+    kv_cache_gqa = 2 * gqa_config.num_key_value_heads * gqa_config.head_dim
+
+    header = f"{'指标':<28}{'GQA':>12}{'MLA(128)':>12}{'MLA(64)':>12}{'MLA(256)':>12}"
+    print(header)
+    print("-" * len(header))
+
+    total_params = {}
+    attn_params = {}
+    kv_cache = {}
+
+    for name, m in models.items():
+        total_params[name] = count_params(m)
+        attn_params[name] = attention_params_per_layer(m)
+        if "MLA" in name:
+            rank = int(name.split("(")[1].rstrip(")"))
+            kv_cache[name] = rank + m.config.rope_dim  # kv_latent + k_rope
+        else:
+            kv_cache[name] = kv_cache_gqa
+
+    row = f"{'总参数量':<28}" + "".join(f"{fmt_num(total_params[n]):>12}" for n in models)
+    print(row)
+    row = f"{'注意力参数量/层':<28}" + "".join(f"{fmt_num(attn_params[n]):>12}" for n in models)
+    print(row)
+    row = f"{'KV-Cache (floats/tok/layer)':<28}" + "".join(f"{kv_cache[n]:>12}" for n in models)
+    print(row)
+    row = f"{'KV-Cache 压缩比':<28}" + "".join(
+        f"{kv_cache_gqa / kv_cache[n]:>11.1f}x" for n in models
+    )
+    print(row)
+
+    return models, kv_cache, kv_cache_gqa
+
+
+# ────────────────────── Benchmark 2: 推理速度 ────────────────────
+
+def benchmark_inference_speed(models, device):
+    print("\n" + "=" * 70)
+    print("  Benchmark 2: 推理速度 (Prefill / Decode)")
+    print("=" * 70)
+
+    prompt_lens = [128, 512, 1024]
+    batch_sizes = [1, 4]
+    max_new_tokens = 128
+
+    for bs in batch_sizes:
+        print(f"\n--- batch_size = {bs} ---")
+        header = (f"{'模型':<16}{'prompt_len':>12}")
+        for pl in prompt_lens:
+            header += f"{'Prefill(ms)':>14}"
+        for pl in prompt_lens:
+            header += f"{'Decode(tok/s)':>16}"
+        print(header)
+        print("-" * len(header))
+
+        for name, model in models.items():
+            model = model.half().to(device).eval()
+            row = f"{name:<16}"
+            prefill_times = []
+
+            for pl in prompt_lens:
+                input_ids = torch.randint(0, 6400, (bs, pl), device=device)
+                warmup_forward(model, input_ids, device, n=2)
+
+                torch.cuda.synchronize(device)
+                t0 = time.perf_counter()
+                with torch.no_grad():
+                    model(input_ids)
+                torch.cuda.synchronize(device)
+                t1 = time.perf_counter()
+                ms = (t1 - t0) * 1000
+                prefill_times.append(ms)
+                row += f"{ms:>14.2f}"
+
+            for i, pl in enumerate(prompt_lens):
+                input_ids = torch.randint(0, 6400, (bs, pl), device=device)
+                attn_mask = torch.ones(bs, pl, device=device)
+                warmup_forward(model, input_ids, device, n=2)
+
+                torch.cuda.synchronize(device)
+                t0 = time.perf_counter()
+                with torch.no_grad():
+                    model.generate(
+                        input_ids=input_ids,
+                        attention_mask=attn_mask,
+                        max_new_tokens=max_new_tokens,
+                        use_cache=True,
+                    )
+                torch.cuda.synchronize(device)
+                t1 = time.perf_counter()
+                tokens_per_sec = max_new_tokens * bs / (t1 - t0)
+                row += f"{tokens_per_sec:>16.1f}"
+
+            print(row)
+            del model
+            cleanup_gpu(device)
+
+
+# ────────────────────── Benchmark 3: KV-Cache 显存 ───────────────
+
+def benchmark_kv_cache_memory(models, device):
+    print("\n" + "=" * 70)
+    print("  Benchmark 3: KV-Cache 显存占用")
+    print("=" * 70)
+
+    context_lens = [256, 512, 1024, 2048, 4096]
+    max_new_tokens = 64
+    bs = 1
+
+    header = f"{'模型':<16}" + "".join(f"{'ctx='+str(c):>12}" for c in context_lens)
+    print(header)
+    print("-" * len(header))
+
+    for name, model in models.items():
+        model = model.half().to(device).eval()
+        row = f"{name:<16}"
+        mem_results = []
+
+        for ctx_len in context_lens:
+            cleanup_gpu(device)
+            input_ids = torch.randint(0, 6400, (bs, ctx_len), device=device)
+            attn_mask = torch.ones(bs, ctx_len, device=device)
+
+            # 记录 generate 前的显存
+            torch.cuda.synchronize(device)
+            mem_before = torch.cuda.max_memory_allocated(device)
+
+            with torch.no_grad():
+                model.generate(
+                    input_ids=input_ids,
+                    attention_mask=attn_mask,
+                    max_new_tokens=max_new_tokens,
+                    use_cache=True,
+                )
+
+            torch.cuda.synchronize(device)
+            mem_after = torch.cuda.max_memory_allocated(device)
+            delta = mem_after - mem_before
+            mem_results.append(delta)
+            row += f"{fmt_mem(delta):>12}"
+
+        print(row)
+
+        # 打印显存增长趋势
+        print(f"  {'增长趋势':<14}", end="")
+        for i in range(1, len(mem_results)):
+            ratio = mem_results[i] / mem_results[i - 1] if mem_results[i - 1] > 0 else 0
+            print(f"{'×' + f'{ratio:.2f}':>12}", end="")
+        print()
+
+        del model
+        cleanup_gpu(device)
+
+
+# ────────────────────── Benchmark 4: 训练吞吐量 ─────────────────
+
+def benchmark_training_throughput(models, device):
+    print("\n" + "=" * 70)
+    print("  Benchmark 4: 训练吞吐量 (forward + backward)")
+    print("=" * 70)
+
+    batch_size = 4
+    seq_len = 340
+    n_steps = 10
+
+    header = f"{'模型':<16}{'step时间(ms)':>16}{'steps/s':>12}{'tokens/s':>14}"
+    print(header)
+    print("-" * len(header))
+
+    for name, model in models.items():
+        model = model.float().to(device).train()
+        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
+
+        # warmup
+        for _ in range(3):
+            input_ids = torch.randint(0, 6400, (batch_size, seq_len), device=device)
+            labels = torch.randint(0, 6400, (batch_size, seq_len), device=device)
+            outputs = model(input_ids, labels=labels)
+            outputs.loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize(device)
+        t0 = time.perf_counter()
+
+        for _ in range(n_steps):
+            input_ids = torch.randint(0, 6400, (batch_size, seq_len), device=device)
+            labels = torch.randint(0, 6400, (batch_size, seq_len), device=device)
+            outputs = model(input_ids, labels=labels)
+            outputs.loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize(device)
+        t1 = time.perf_counter()
+
+        total_time = t1 - t0
+        step_ms = total_time / n_steps * 1000
+        steps_per_sec = n_steps / total_time
+        tokens_per_sec = batch_size * seq_len * n_steps / total_time
+
+        print(f"{name:<16}{step_ms:>16.2f}{steps_per_sec:>12.2f}{tokens_per_sec:>14.1f}")
+
+        del model, optimizer
+        cleanup_gpu(device)
+
+
+# ────────────────────── Benchmark 5: 输出质量 (PPL) ──────────────
+
+def benchmark_perplexity(args, device, hidden_size):
+    print("\n" + "=" * 70)
+    print("  Benchmark 5: 输出质量 (Perplexity)")
+    print("=" * 70)
+
+    from datasets import load_dataset
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
+    samples = load_dataset("json", data_files=args.data_path, split="train")
+    # 取前 eval_samples 条，截取文本
+    max_seq = 512
+    all_tokens = []
+    for i in range(min(args.eval_samples, len(samples))):
+        text = str(samples[i].get("text", ""))
+        ids = tokenizer(text, add_special_tokens=False, max_length=max_seq, truncation=True).input_ids
+        ids = [tokenizer.bos_token_id] + ids + [tokenizer.eos_token_id]
+        all_tokens.append(ids)
+
+    results = {}
+
+    for ckpt_name, ckpt_path, config_cls, config_kwargs in [
+        ("GQA", args.gqa_ckpt, MiniMindConfig, {}),
+        ("MLA(128)", args.mla_ckpt, MiniMindMLAConfig, {"kv_lora_rank": 128}),
+        ("MLA(64)", args.mla_ckpt_r64, MiniMindMLAConfig, {"kv_lora_rank": 64}),
+        ("MLA(256)", args.mla_ckpt_r256, MiniMindMLAConfig, {"kv_lora_rank": 256}),
+    ]:
+        if ckpt_path is None:
+            continue
+        if not os.path.isfile(ckpt_path):
+            print(f"  [SKIP] {ckpt_name}: checkpoint not found at {ckpt_path}")
+            continue
+
+        config_kw = {"hidden_size": hidden_size}
+        config_kw.update(config_kwargs)
+        config = config_cls(**config_kw)
+        if isinstance(config, MiniMindMLAConfig):
+            model = MiniMindMLAForCausalLM(config)
+        else:
+            model = MiniMindForCausalLM(config)
+
+        state_dict = torch.load(ckpt_path, map_location="cpu")
+        model.load_state_dict(state_dict, strict=True)
+        model = model.half().to(device).eval()
+
+        total_loss = 0.0
+        total_tokens = 0
+        batch_input_ids = []
+        batch_labels = []
+
+        for ids in all_tokens:
+            pad_len = max_seq - len(ids)
+            padded = ids + [tokenizer.pad_token_id] * pad_len
+            labels = ids[1:] + [tokenizer.pad_token_id] * (pad_len + 1)
+            # 标记 padding 为 -100
+            labels = [t if t != tokenizer.pad_token_id else -100 for t in labels]
+            batch_input_ids.append(padded)
+            batch_labels.append(labels)
+
+            if len(batch_input_ids) == args.eval_batch_size:
+                input_t = torch.tensor(batch_input_ids, dtype=torch.long, device=device)
+                label_t = torch.tensor(batch_labels, dtype=torch.long, device=device)
+                with torch.no_grad():
+                    out = model(input_t, labels=label_t)
+                # loss 是 cross_entropy 的 mean，乘以有效 token 数得到总 loss
+                n_valid = (label_t != -100).sum().item()
+                total_loss += out.loss.item() * n_valid
+                total_tokens += n_valid
+                batch_input_ids, batch_labels = [], []
+
+        # 处理剩余样本
+        if batch_input_ids:
+            input_t = torch.tensor(batch_input_ids, dtype=torch.long, device=device)
+            label_t = torch.tensor(batch_labels, dtype=torch.long, device=device)
+            with torch.no_grad():
+                out = model(input_t, labels=label_t)
+            n_valid = (label_t != -100).sum().item()
+            total_loss += out.loss.item() * n_valid
+            total_tokens += n_valid
+
+        avg_nll = total_loss / total_tokens
+        ppl = math.exp(avg_nll)
+        results[ckpt_name] = ppl
+        print(f"  {ckpt_name:<16} Loss: {avg_nll:.4f}  PPL: {ppl:.2f}  (eval on {total_tokens} tokens)")
+
+        del model, state_dict
+        cleanup_gpu(device)
+
+    if not results:
+        print("  没有可用的 checkpoint，跳过 PPL 测试")
+        print("  提示：使用 --gqa_ckpt / --mla_ckpt 参数指定权重路径")
+        return {}
+
+    # Markdown PPL 表格
+    print()
+    cols = " | ".join(results.keys())
+    print(f"| Metric | {cols} |")
+    sep = "|".join(["------"] * len(results))
+    print(f"|--------|{sep}|")
+    print(f"| PPL | " + " | ".join(f"{results[k]:.2f}" for k in results) + " |")
+    print()
+    return results
+
+
+# ────────────────────── Markdown 汇总 ────────────────────────────
+
+def print_markdown_summary(models_info, hidden_size):
+    print("\n" + "=" * 70)
+    print("  Markdown 汇总（可复制到简历 / README）")
+    print("=" * 70)
+    print()
+
+    gqa_config = MiniMindConfig(hidden_size=hidden_size)
+    kv_cache_gqa = 2 * gqa_config.num_key_value_heads * gqa_config.head_dim
+
+    print("### Model Info")
+    print()
+    print(f"| Metric | GQA | MLA(128) | MLA(64) | MLA(256) |")
+    print(f"|--------|-----|----------|---------|----------|")
+
+    total_params = {}
+    attn_params = {}
+    kv_cache = {}
+    for name, m in models_info.items():
+        total_params[name] = count_params(m)
+        attn_params[name] = attention_params_per_layer(m)
+        if "MLA" in name:
+            rank = int(name.split("(")[1].rstrip(")"))
+            kv_cache[name] = rank + m.config.rope_dim  # kv_latent + k_rope
+        else:
+            kv_cache[name] = kv_cache_gqa
+
+    row = "| Total Params |"
+    for n in models_info:
+        row += f" {fmt_num(total_params[n])} |"
+    print(row)
+
+    row = "| Attn Params/Layer |"
+    for n in models_info:
+        row += f" {fmt_num(attn_params[n])} |"
+    print(row)
+
+    row = "| KV-Cache (floats/tok/layer) |"
+    for n in models_info:
+        row += f" {kv_cache[n]} |"
+    print(row)
+
+    row = "| KV-Cache Compression Ratio |"
+    for n in models_info:
+        row += f" {kv_cache_gqa / kv_cache[n]:.1f}x |"
+    print(row)
+    print()
+
+    print("> hidden_size={}, num_layers={}, num_attention_heads={}, num_kv_heads={}, head_dim={}".format(
+        gqa_config.hidden_size, gqa_config.num_hidden_layers,
+        gqa_config.num_attention_heads, gqa_config.num_key_value_heads, gqa_config.head_dim
+    ))
+    print()
+
+
+# ────────────────────── main ─────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description="GQA vs MLA Benchmark")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--hidden_size", type=int, default=768)
+    # Benchmark 5 (PPL) 参数
+    parser.add_argument("--gqa_ckpt", type=str, default=None, help="GQA checkpoint 路径 (e.g. out/pretrain_768.pth)")
+    parser.add_argument("--mla_ckpt", type=str, default=None, help="MLA checkpoint 路径 (e.g. out/pretrain_768_mla.pth)")
+    parser.add_argument("--mla_ckpt_r64", type=str, default=None, help="MLA(64) checkpoint 路径")
+    parser.add_argument("--mla_ckpt_r256", type=str, default=None, help="MLA(256) checkpoint 路径")
+    parser.add_argument("--data_path", type=str, default="dataset/pretrain_hq.jsonl", help="评估数据集路径")
+    parser.add_argument("--tokenizer_path", type=str, default="model", help="tokenizer 路径")
+    parser.add_argument("--eval_samples", type=int, default=500, help="评估样本数")
+    parser.add_argument("--eval_batch_size", type=int, default=4, help="评估 batch size")
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+    hidden_size = args.hidden_size
+
+    has_ckpt = any([args.gqa_ckpt, args.mla_ckpt, args.mla_ckpt_r64, args.mla_ckpt_r256])
+
+    print("=" * 70)
+    print(f"  MiniMind GQA vs MLA Benchmark")
+    print(f"  hidden_size={hidden_size}  device={device}")
+    if has_ckpt:
+        print(f"  PPL evaluation: ENABLED (data={args.data_path})")
+    print("=" * 70)
+
+    # Benchmark 1: 模型信息（CPU 即可）
+    models, kv_cache, kv_cache_gqa = benchmark_model_info(hidden_size)
+
+    if device.type == "cuda":
+        # Benchmark 2: 推理速度
+        models_gpu = {
+            "GQA": build_gqa(hidden_size),
+            "MLA(128)": build_mla(hidden_size, 128),
+            "MLA(64)": build_mla(hidden_size, 64),
+            "MLA(256)": build_mla(hidden_size, 256),
+        }
+        benchmark_inference_speed(models_gpu, device)
+
+        # Benchmark 3: KV-Cache 显存
+        models_gpu = {
+            "GQA": build_gqa(hidden_size),
+            "MLA(128)": build_mla(hidden_size, 128),
+            "MLA(64)": build_mla(hidden_size, 64),
+            "MLA(256)": build_mla(hidden_size, 256),
+        }
+        benchmark_kv_cache_memory(models_gpu, device)
+
+        # Benchmark 4: 训练吞吐量
+        models_gpu = {
+            "GQA": build_gqa(hidden_size),
+            "MLA(128)": build_mla(hidden_size, 128),
+            "MLA(64)": build_mla(hidden_size, 64),
+            "MLA(256)": build_mla(hidden_size, 256),
+        }
+        benchmark_training_throughput(models_gpu, device)
+
+        # Benchmark 5: PPL（需要提供 checkpoint 路径）
+        if has_ckpt:
+            benchmark_perplexity(args, device, hidden_size)
+        else:
+            print("\n" + "=" * 70)
+            print("  Benchmark 5: 输出质量 (PPL) — SKIPPED")
+            print("=" * 70)
+            print("  未指定 checkpoint，跳过 PPL 测试。")
+            print("  使用方法: python scripts/benchmark_gqa_vs_mla.py \\")
+            print("      --gqa_ckpt out/pretrain_768.pth \\")
+            print("      --mla_ckpt out/pretrain_768_mla.pth \\")
+            print("      --data_path dataset/pretrain_hq.jsonl")
+    else:
+        print("\n[SKIP] Benchmark 2-5 需要 CUDA 设备，当前 device=cpu")
+
+    # Markdown 汇总（使用 Benchmark 1 的 CPU 模型）
+    print_markdown_summary(models, hidden_size)
+
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()

scripts/convert_model.py

@@ -9,11 +9,17 @@ import transformers
 import warnings
 from transformers import AutoTokenizer, AutoModelForCausalLM, Qwen3Config, Qwen3ForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
 from model.model_lora import apply_lora, merge_lora
 
 warnings.filterwarnings('ignore', category=UserWarning)
 
 def convert_torch2transformers_minimind(torch_path, transformers_path, dtype=torch.float16):
+    if isinstance(lm_config, MiniMindMLAConfig):
+        MiniMindMLAConfig.register_for_auto_class()
+        MiniMindMLAForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+        lm_model = MiniMindMLAForCausalLM(lm_config)
+    else:
         MiniMindConfig.register_for_auto_class()
         MiniMindForCausalLM.register_for_auto_class("AutoModelForCausalLM")
         lm_model = MiniMindForCausalLM(lm_config)
@@ -104,6 +110,9 @@ def convert_transformers2torch(transformers_path, torch_path):
 
 def convert_merge_base_lora(base_torch_path, lora_path, merged_torch_path):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if isinstance(lm_config, MiniMindMLAConfig):
+        lm_model = MiniMindMLAForCausalLM(lm_config).to(device)
+    else:
         lm_model = MiniMindForCausalLM(lm_config).to(device)
     state_dict = torch.load(base_torch_path, map_location=device)
     lm_model.load_state_dict(state_dict, strict=False)
@@ -126,17 +135,33 @@ def convert_json_to_jinja(json_file_path, output_path):
 
 
 if __name__ == '__main__':
-    lm_config = MiniMindConfig(hidden_size=768, num_hidden_layers=8, max_seq_len=8192, use_moe=False)
+    import argparse
+    parser = argparse.ArgumentParser(description="MiniMind Model Conversion")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
+    parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀")
+    parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
+    args = parser.parse_args()
+
+    if args.use_mla:
+        lm_config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, max_seq_len=8192, use_moe=bool(args.use_moe), kv_lora_rank=args.kv_lora_rank)
+    else:
+        lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, max_seq_len=8192, use_moe=bool(args.use_moe))
+    model_suffix = '_moe' if lm_config.use_moe else ''
+    if isinstance(lm_config, MiniMindMLAConfig): model_suffix += '_mla'
 
     # convert torch to transformers
-    torch_path = f"../out/full_sft_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+    torch_path = f"../{args.save_dir}/{args.weight}_{lm_config.hidden_size}{model_suffix}.pth"
     transformers_path = '../minimind-3'
     convert_torch2transformers(torch_path, transformers_path)
 
     # # merge lora
-    # base_torch_path = f"../out/full_sft_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+    # base_torch_path = f"../{args.save_dir}/{args.weight}_{lm_config.hidden_size}{model_suffix}.pth"
-    # lora_path = f"../out/lora_identity_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+    # lora_path = f"../{args.save_dir}/lora_identity_{lm_config.hidden_size}{model_suffix}.pth"
-    # merged_torch_path = f"../out/merge_identity_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+    # merged_torch_path = f"../{args.save_dir}/merge_identity_{lm_config.hidden_size}{model_suffix}.pth"
     # convert_merge_base_lora(base_torch_path, lora_path, merged_torch_path)
 
     # convert_transformers2torch(transformers_path, torch_path)

scripts/eval_toolcall.py

@@ -12,6 +12,7 @@ from datetime import datetime
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from openai import OpenAI
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
 from trainer.trainer_utils import setup_seed, get_model_params
 warnings.filterwarnings('ignore')
 
@@ -57,9 +58,14 @@ TEST_CASES = [
 def init_model(args):
     tokenizer = AutoTokenizer.from_pretrained(args.load_from)
     if 'model' in args.load_from:
+        if args.use_mla:
+            config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe), kv_lora_rank=args.kv_lora_rank)
+            model = MiniMindMLAForCausalLM(config)
+        else:
             model = MiniMindForCausalLM(MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe)))
-        moe_suffix = '_moe' if args.use_moe else ''
+        model_suffix = '_moe' if args.use_moe else ''
-        ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
+        if args.use_mla: model_suffix += '_mla'
+        ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{model_suffix}.pth'
         model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
     else:
         model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
@@ -208,6 +214,8 @@ def main():
     parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     parser.add_argument('--max_new_tokens', default=512, type=int, help="最大生成长度")
     parser.add_argument('--temperature', default=0.9, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
     parser.add_argument('--top_p', default=0.9, type=float, help="nucleus采样阈值（0-1）")

scripts/serve_openai_api.py

@@ -18,6 +18,7 @@ from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
 from model.model_lora import apply_lora, load_lora
 
 warnings.filterwarnings('ignore')
@@ -28,8 +29,20 @@ app = FastAPI()
 def init_model(args):
     tokenizer = AutoTokenizer.from_pretrained(args.load_from)
     if 'model' in args.load_from:
-        moe_suffix = '_moe' if args.use_moe else ''
+        model_suffix = '_moe' if args.use_moe else ''
-        ckp = f'../{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
+        if args.use_mla: model_suffix += '_mla'
+        ckp = f'../{args.save_dir}/{args.weight}_{args.hidden_size}{model_suffix}.pth'
+        if args.use_mla:
+            config = MiniMindMLAConfig(
+                hidden_size=args.hidden_size,
+                num_hidden_layers=args.num_hidden_layers,
+                max_seq_len=args.max_seq_len,
+                use_moe=bool(args.use_moe),
+                inference_rope_scaling=args.inference_rope_scaling,
+                kv_lora_rank=args.kv_lora_rank
+            )
+            model = MiniMindMLAForCausalLM(config)
+        else:
             model = MiniMindForCausalLM(MiniMindConfig(
                 hidden_size=args.hidden_size,
                 num_hidden_layers=args.num_hidden_layers,
@@ -237,6 +250,8 @@ if __name__ == "__main__":
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=8192, type=int, help="最大序列长度")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     parser.add_argument('--inference_rope_scaling', default=False, action='store_true', help="启用RoPE位置编码外推（4倍，仅解决位置编码问题）")
     parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
     args = parser.parse_args()

trainer/train_agent.py

@@ -22,8 +22,9 @@ from torch.utils.data import DataLoader, DistributedSampler
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from transformers import AutoTokenizer
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
 from dataset.lm_dataset import AgentRLDataset
-from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel, get_model_suffix
 from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps
 
 warnings.filterwarnings('ignore')
@@ -349,8 +350,8 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
 
         if (step % args.save_interval == 0 or step == iters) and is_main_process():
             model.eval()
-            moe_suffix = '_moe' if lm_config.use_moe else ''
+            model_suffix = get_model_suffix(lm_config)
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{model_suffix}.pth'
             raw_model = model.module if isinstance(model, DistributedDataParallel) else model
             raw_model = getattr(raw_model, '_orig_mod', raw_model)
             state_dict = raw_model.state_dict()
@@ -387,6 +388,8 @@ if __name__ == "__main__":
     parser.add_argument('--hidden_size', default=768, type=int, help="模型隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="模型层数")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE")
+    parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     parser.add_argument('--max_seq_len', default=1024, type=int, help="最大序列长度")
     parser.add_argument("--max_gen_len", type=int, default=768, help="单次最大生成长度")
     parser.add_argument("--max_total_len", type=int, default=2500, help="训练侧最终总长度上界")
@@ -416,7 +419,9 @@ if __name__ == "__main__":
     setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
 
     os.makedirs(args.save_dir, exist_ok=True)
-    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+    lm_config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+                               max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe),
+                               kv_lora_rank=args.kv_lora_rank) if args.use_mla else MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
                                max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None
 

trainer/train_distillation.py

@@ -15,8 +15,9 @@ from torch import optim
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from model.model_minimind import MiniMindConfig
+from model.model_minimind_mla import MiniMindMLAConfig
 from dataset.lm_dataset import SFTDataset
-from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
+from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler, get_model_suffix
 
 warnings.filterwarnings('ignore')
 
@@ -122,8 +123,8 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st
 
         if (step % args.save_interval == 0 or step == iters) and is_main_process():
             model.eval()
-            moe_suffix = '_moe' if lm_config_student.use_moe else ''
+            model_suffix = get_model_suffix(lm_config_student)
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config_student.hidden_size}{moe_suffix}.pth'
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config_student.hidden_size}{model_suffix}.pth'
             raw_model = model.module if isinstance(model, DistributedDataParallel) else model
             raw_model = getattr(raw_model, '_orig_mod', raw_model)
             state_dict = raw_model.state_dict()
@@ -165,6 +166,8 @@ if __name__ == "__main__":
     parser.add_argument('--teacher_num_layers', default=8, type=int, help="教师模型隐藏层数量")
     parser.add_argument('--student_use_moe', default=0, type=int, choices=[0, 1], help="学生模型是否使用MoE（0=否，1=是）")
     parser.add_argument('--teacher_use_moe', default=1, type=int, choices=[0, 1], help="教师模型是否使用MoE（0=否，1=是）")
+    parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     parser.add_argument('--from_student_weight', default='full_sft', type=str, help="学生模型基于哪个权重")
     parser.add_argument('--from_teacher_weight', default='full_sft', type=str, help="教师模型基于哪个权重")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
@@ -182,7 +185,7 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
-    lm_config_student = MiniMindConfig(hidden_size=args.student_hidden_size, num_hidden_layers=args.student_num_layers, use_moe=bool(args.student_use_moe))
+    lm_config_student = MiniMindMLAConfig(hidden_size=args.student_hidden_size, num_hidden_layers=args.student_num_layers, use_moe=bool(args.student_use_moe), kv_lora_rank=args.kv_lora_rank) if args.use_mla else MiniMindConfig(hidden_size=args.student_hidden_size, num_hidden_layers=args.student_num_layers, use_moe=bool(args.student_use_moe))
     lm_config_teacher = MiniMindConfig(hidden_size=args.teacher_hidden_size, num_hidden_layers=args.teacher_num_layers, use_moe=bool(args.teacher_use_moe))
     ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     

trainer/train_dpo.py

@@ -15,8 +15,9 @@ from torch import optim
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from model.model_minimind import MiniMindConfig
+from model.model_minimind_mla import MiniMindMLAConfig
 from dataset.lm_dataset import DPODataset
-from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
+from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler, get_model_suffix
 
 warnings.filterwarnings('ignore')
 
@@ -106,8 +107,8 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=
 
         if (step % args.save_interval == 0 or step == iters) and is_main_process():
             model.eval()
-            moe_suffix = '_moe' if lm_config.use_moe else ''
+            model_suffix = get_model_suffix(lm_config)
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{model_suffix}.pth'
             raw_model = model.module if isinstance(model, DistributedDataParallel) else model
             raw_model = getattr(raw_model, '_orig_mod', raw_model)
             state_dict = raw_model.state_dict()
@@ -145,6 +146,8 @@ if __name__ == "__main__":
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=1024, 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('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     parser.add_argument("--data_path", type=str, default="../dataset/dpo.jsonl", help="DPO训练数据路径")
     parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
@@ -161,6 +164,9 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    if args.use_mla:
+        lm_config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe), kv_lora_rank=args.kv_lora_rank)
+    else:
         lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     

trainer/train_full_sft.py

@@ -14,8 +14,9 @@ from torch import optim, nn
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from model.model_minimind import MiniMindConfig
+from model.model_minimind_mla import MiniMindMLAConfig
 from dataset.lm_dataset import SFTDataset
-from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
+from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler, get_model_suffix
 
 warnings.filterwarnings('ignore')
 
@@ -59,8 +60,8 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
 
         if (step % args.save_interval == 0 or step == iters) and is_main_process():
             model.eval()
-            moe_suffix = '_moe' if lm_config.use_moe else ''
+            model_suffix = get_model_suffix(lm_config)
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{model_suffix}.pth'
             raw_model = model.module if isinstance(model, DistributedDataParallel) else model
             raw_model = getattr(raw_model, '_orig_mod', raw_model)
             state_dict = raw_model.state_dict()
@@ -98,6 +99,8 @@ 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('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     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=是）")
@@ -113,6 +116,9 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    if args.use_mla:
+        lm_config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe), kv_lora_rank=args.kv_lora_rank)
+    else:
         lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     

trainer/train_grpo.py

@@ -20,8 +20,9 @@ from torch.utils.data import DataLoader, DistributedSampler
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from transformers import AutoModel
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
 from dataset.lm_dataset import RLAIFDataset
-from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel, get_model_suffix
 from trainer.rollout_engine import create_rollout_engine
 
 warnings.filterwarnings('ignore')
@@ -178,8 +179,8 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
 
         if (step % args.save_interval == 0 or step == iters) and is_main_process():
             model.eval()
-            moe_suffix = '_moe' if lm_config.use_moe else ''
+            model_suffix = get_model_suffix(lm_config)
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{model_suffix}.pth'
             raw_model = model.module if isinstance(model, DistributedDataParallel) else model
             raw_model = getattr(raw_model, '_orig_mod', raw_model)
             state_dict = raw_model.state_dict()
@@ -219,6 +220,8 @@ if __name__ == "__main__":
     parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     parser.add_argument('--max_seq_len', default=768, type=int, help="Prompt最大长度")
     parser.add_argument("--max_gen_len", type=int, default=1024, help="生成的最大长度")
     parser.add_argument("--data_path", type=str, default="../dataset/rlaif.jsonl", help="RLAIF数据路径")
@@ -249,7 +252,9 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
-    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+    lm_config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+                               max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe),
+                               kv_lora_rank=args.kv_lora_rank) if args.use_mla else MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
                                max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     

trainer/train_lora.py

@@ -14,9 +14,10 @@ from torch import optim, nn
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from model.model_minimind import MiniMindConfig
+from model.model_minimind_mla import MiniMindMLAConfig
 from dataset.lm_dataset import SFTDataset
 from model.model_lora import save_lora, apply_lora
-from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
+from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler, get_model_suffix
 
 warnings.filterwarnings('ignore')
 
@@ -58,8 +59,8 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
 
         if (step % args.save_interval == 0 or step == iters) and is_main_process():
             model.eval()
-            moe_suffix = '_moe' if lm_config.use_moe else ''
+            model_suffix = get_model_suffix(lm_config)
-            lora_save_path = f'{args.save_dir}/{args.lora_name}_{lm_config.hidden_size}{moe_suffix}.pth'
+            lora_save_path = f'{args.save_dir}/{args.lora_name}_{lm_config.hidden_size}{model_suffix}.pth'
             # LoRA只保存LoRA权重
             save_lora(model, lora_save_path)
             lm_checkpoint(lm_config, weight=args.lora_name, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
@@ -92,6 +93,8 @@ 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('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     parser.add_argument("--data_path", type=str, default="../dataset/lora_medical.jsonl", help="LoRA训练数据路径")
     parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练，默认full_sft")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
@@ -107,6 +110,9 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    if args.use_mla:
+        lm_config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe), kv_lora_rank=args.kv_lora_rank)
+    else:
         lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
     ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None
     

trainer/train_ppo.py

@@ -19,8 +19,9 @@ from torch.utils.data import DataLoader, DistributedSampler
 from torch.nn.utils import clip_grad_norm_
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
 from dataset.lm_dataset import RLAIFDataset
-from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel, get_model_suffix
 from trainer.rollout_engine import create_rollout_engine
 
 warnings.filterwarnings('ignore')
@@ -32,8 +33,8 @@ def rep_penalty(text, n=3, cap=0.5):
     return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
 
 
-# 自定义的Critic模型，继承自MiniMindLM
+def _make_critic_model(base_class):
-class CriticModel(MiniMindForCausalLM):
+    class CriticModel(base_class):
         def __init__(self, params):
             super().__init__(params)
             # 替换lm_head为输出单一价值的线性层
@@ -46,6 +47,7 @@ class CriticModel(MiniMindForCausalLM):
             # 使用value_head获取价值估计
             values = self.value_head(hidden_states).squeeze(-1)
             return values
+    return CriticModel
 
 
 def calculate_rewards(prompts, responses, reward_model):
@@ -272,8 +274,8 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
 
         if (step % args.save_interval == 0 or step == iters) and is_main_process():
             actor_model.eval()
-            moe_suffix = '_moe' if lm_config.use_moe else ''
+            model_suffix = get_model_suffix(lm_config)
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{model_suffix}.pth'
             raw_actor = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
             raw_actor = getattr(raw_actor, '_orig_mod', raw_actor)
             actor_state = raw_actor.state_dict()
@@ -310,6 +312,8 @@ if __name__ == "__main__":
     parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     parser.add_argument('--max_seq_len', default=768, type=int, help="Prompt最大长度")
     parser.add_argument("--max_gen_len", type=int, default=1024, help="生成的最大长度")
     parser.add_argument("--data_path", type=str, default="../dataset/rlaif.jsonl", help="RLAIF数据路径")
@@ -344,7 +348,7 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
-    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
+    lm_config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe), kv_lora_rank=args.kv_lora_rank) if args.use_mla else MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========
@@ -367,9 +371,10 @@ if __name__ == "__main__":
     actor_model, tokenizer = init_model(lm_config, base_weight, device=args.device)
     ref_model, _ = init_model(lm_config, base_weight, device=args.device)
     ref_model = ref_model.eval().requires_grad_(False)
-    moe_suffix = '_moe' if lm_config.use_moe else ''
+    model_suffix = get_model_suffix(lm_config)
-    ckp = f'{args.save_dir}/{base_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+    ckp = f'{args.save_dir}/{base_weight}_{lm_config.hidden_size}{model_suffix}.pth'
     state_dict = torch.load(ckp, map_location=args.device)
+    CriticModel = _make_critic_model(MiniMindMLAForCausalLM if isinstance(lm_config, MiniMindMLAConfig) else MiniMindForCausalLM)
     critic_model = CriticModel(lm_config)
     critic_model.load_state_dict(state_dict, strict=False)
     critic_model = critic_model.to(args.device)

trainer/train_pretrain.py

@@ -14,8 +14,9 @@ from torch import optim, nn
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from model.model_minimind import MiniMindConfig
+from model.model_minimind_mla import MiniMindMLAConfig
 from dataset.lm_dataset import PretrainDataset
-from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
+from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler, get_model_suffix
 
 warnings.filterwarnings('ignore')
 
@@ -60,7 +61,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
 
         if (step % args.save_interval == 0 or step == iters) and is_main_process():
             model.eval()
-            moe_suffix = '_moe' if lm_config.use_moe else ''
+            moe_suffix = get_model_suffix(lm_config)
             ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
             raw_model = model.module if isinstance(model, DistributedDataParallel) else model
             raw_model = getattr(raw_model, '_orig_mod', raw_model)
@@ -99,6 +100,8 @@ 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('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）")
+    parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）")
     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=是）")
@@ -114,6 +117,9 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    if args.use_mla:
+        lm_config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe), kv_lora_rank=args.kv_lora_rank)
+    else:
         lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     

trainer/trainer_utils.py

@@ -14,6 +14,15 @@ from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import Sampler
 from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
 from model.model_minimind import MiniMindForCausalLM
+from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
+
+def get_model_suffix(lm_config):
+    """根据模型配置返回 checkpoint 文件名后缀"""
+    suffix = ''
+    if getattr(lm_config, 'use_moe', False): suffix += '_moe'
+    if isinstance(lm_config, MiniMindMLAConfig): suffix += '_mla'
+    return suffix
+
 
 def get_model_params(model, config):
     total = sum(p.numel() for p in model.parameters()) / 1e6
@@ -63,7 +72,7 @@ def setup_seed(seed: int):
 
 def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs):
     os.makedirs(save_dir, exist_ok=True)
-    moe_path = '_moe' if lm_config.use_moe else ''
+    moe_path = get_model_suffix(lm_config)
     ckp_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}.pth'
     resume_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}_resume.pth'
 
@@ -125,11 +134,14 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc
 
 def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device='cuda'):
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+    if isinstance(lm_config, MiniMindMLAConfig):
+        model = MiniMindMLAForCausalLM(lm_config)
+    else:
         model = MiniMindForCausalLM(lm_config)
 
     if from_weight!= 'none':
-        moe_suffix = '_moe' if lm_config.use_moe else ''
+        model_suffix = get_model_suffix(lm_config)
-        weight_path = f'{save_dir}/{from_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+        weight_path = f'{save_dir}/{from_weight}_{lm_config.hidden_size}{model_suffix}.pth'
         weights = torch.load(weight_path, map_location=device)
         model.load_state_dict(weights, strict=False)