From bc3d1372486787f7caaa2051e0d65383fc152881 Mon Sep 17 00:00:00 2001 From: baihaochen <1291308544@qq.com> Date: Tue, 12 May 2026 18:38:47 +0800 Subject: [PATCH] [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 | 37 ++- model/model_minimind_mla.py | 241 ++++++++++++++ scripts/benchmark_gqa_vs_mla.py | 557 ++++++++++++++++++++++++++++++++ scripts/convert_model.py | 43 ++- scripts/eval_toolcall.py | 14 +- scripts/serve_openai_api.py | 33 +- trainer/train_agent.py | 13 +- trainer/train_distillation.py | 11 +- trainer/train_dpo.py | 14 +- trainer/train_full_sft.py | 14 +- trainer/train_grpo.py | 13 +- trainer/train_lora.py | 14 +- trainer/train_ppo.py | 43 +-- trainer/train_pretrain.py | 12 +- trainer/trainer_utils.py | 20 +- 15 files changed, 998 insertions(+), 81 deletions(-) create mode 100644 model/model_minimind_mla.py create mode 100644 scripts/benchmark_gqa_vs_mla.py diff --git a/eval_llm.py b/eval_llm.py index 9359ca8..f3a5c42 100755 --- a/eval_llm.py +++ b/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: - 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 '' - ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth' + 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 + )) + model_suffix = '_moe' if args.use_moe else '' + 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='') diff --git a/model/model_minimind_mla.py b/model/model_minimind_mla.py new file mode 100644 index 0000000..089072d --- /dev/null +++ b/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 diff --git a/scripts/benchmark_gqa_vs_mla.py b/scripts/benchmark_gqa_vs_mla.py new file mode 100644 index 0000000..ea29bab --- /dev/null +++ b/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() diff --git a/scripts/convert_model.py b/scripts/convert_model.py index a4162a3..c2415ff 100644 --- a/scripts/convert_model.py +++ b/scripts/convert_model.py @@ -9,14 +9,20 @@ 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): - MiniMindConfig.register_for_auto_class() - MiniMindForCausalLM.register_for_auto_class("AutoModelForCausalLM") - lm_model = MiniMindForCausalLM(lm_config) + 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) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') state_dict = torch.load(torch_path, map_location=device) lm_model.load_state_dict(state_dict, strict=False) @@ -104,7 +110,10 @@ 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') - lm_model = MiniMindForCausalLM(lm_config).to(device) + 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) apply_lora(lm_model) @@ -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" - # lora_path = f"../out/lora_identity_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth" - # merged_torch_path = f"../out/merge_identity_{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"../{args.save_dir}/lora_identity_{lm_config.hidden_size}{model_suffix}.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) diff --git a/scripts/eval_toolcall.py b/scripts/eval_toolcall.py index da078b4..71cec61 100644 --- a/scripts/eval_toolcall.py +++ b/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: - 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 '' - ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth' + 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))) + model_suffix = '_moe' if args.use_moe else '' + 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）") diff --git a/scripts/serve_openai_api.py b/scripts/serve_openai_api.py index 0550d4a..3327438 100644 --- a/scripts/serve_openai_api.py +++ b/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,15 +29,27 @@ 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 '' - ckp = f'../{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth' - model = MiniMindForCausalLM(MiniMindConfig( - 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 - )) + model_suffix = '_moe' if args.use_moe else '' + 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, + max_seq_len=args.max_seq_len, + use_moe=bool(args.use_moe), + inference_rope_scaling=args.inference_rope_scaling + )) model.load_state_dict(torch.load(ckp, map_location=device), strict=True) if args.lora_weight != 'None': apply_lora(model) @@ -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() diff --git a/trainer/train_agent.py b/trainer/train_agent.py index e652c3e..e8fea0b 100644 --- a/trainer/train_agent.py +++ b/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 '' - ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + model_suffix = get_model_suffix(lm_config) + 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 diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index 63b6025..13c5887 100644 --- a/trainer/train_distillation.py +++ b/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 '' - ckp = f'{args.save_dir}/{args.save_weight}_{lm_config_student.hidden_size}{moe_suffix}.pth' + model_suffix = get_model_suffix(lm_config_student) + 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 diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index 558d108..14818d4 100644 --- a/trainer/train_dpo.py +++ b/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 '' - ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + model_suffix = get_model_suffix(lm_config) + 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,7 +164,10 @@ 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)) + 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 # ========== 3. 设置混合精度 ========== diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index 51c755b..637f278 100644 --- a/trainer/train_full_sft.py +++ b/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 '' - ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + model_suffix = get_model_suffix(lm_config) + 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,7 +116,10 @@ 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)) + 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 # ========== 3. 设置混合精度 ========== diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 52f7c7d..c368b0a 100755 --- a/trainer/train_grpo.py +++ b/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 '' - ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + model_suffix = get_model_suffix(lm_config) + 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 diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 974bf37..b036781 100644 --- a/trainer/train_lora.py +++ b/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 '' - lora_save_path = f'{args.save_dir}/{args.lora_name}_{lm_config.hidden_size}{moe_suffix}.pth' + model_suffix = get_model_suffix(lm_config) + 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,7 +110,10 @@ 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)) + 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 # ========== 3. 设置混合精度 ========== diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 228fb68..e5e5c60 100644 --- a/trainer/train_ppo.py +++ b/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,20 +33,21 @@ 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 -class CriticModel(MiniMindForCausalLM): - def __init__(self, params): - super().__init__(params) - # 替换lm_head为输出单一价值的线性层 - self.value_head = nn.Linear(params.hidden_size, 1) +def _make_critic_model(base_class): + class CriticModel(base_class): + def __init__(self, params): + super().__init__(params) + # 替换lm_head为输出单一价值的线性层 + self.value_head = nn.Linear(params.hidden_size, 1) - def forward(self, input_ids=None, attention_mask=None, **kwargs): - # 使用基础模型获取隐藏状态 - outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs) - hidden_states = self.model.norm(outputs[0]) - # 使用value_head获取价值估计 - values = self.value_head(hidden_states).squeeze(-1) - return values + def forward(self, input_ids=None, attention_mask=None, **kwargs): + # 使用基础模型获取隐藏状态 + outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs) + hidden_states = self.model.norm(outputs[0]) + # 使用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 '' - ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + model_suffix = get_model_suffix(lm_config) + 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 '' - ckp = f'{args.save_dir}/{base_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + model_suffix = get_model_suffix(lm_config) + 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) diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 1698294..ec191b6 100644 --- a/trainer/train_pretrain.py +++ b/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,7 +117,10 @@ 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)) + 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 # ========== 3. 设置混合精度 ========== diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index 4f8c758..0dc7b4b 100644 --- a/trainer/trainer_utils.py +++ b/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) - model = MiniMindForCausalLM(lm_config) + 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 '' - weight_path = f'{save_dir}/{from_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + model_suffix = get_model_suffix(lm_config) + 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)