From cc29d9a3518d7956ae477e3f227ac726bcbad3e5 Mon Sep 17 00:00:00 2001
From: dieu <515648571@qq.com>
Date: Fri, 5 Dec 2025 13:55:20 +0800
Subject: [PATCH 1/2] feat: add merge_lora.py to support merging LoRA weights
 into base model

---
 scripts/call_merged_model.py |  45 +++++++
 scripts/merge_lora_model.py  | 239 +++++++++++++++++++++++++++++++++++
 2 files changed, 284 insertions(+)
 create mode 100644 scripts/call_merged_model.py
 create mode 100644 scripts/merge_lora_model.py

diff --git a/scripts/call_merged_model.py b/scripts/call_merged_model.py
new file mode 100644
index 0000000..8f7e5a6
--- /dev/null
+++ b/scripts/call_merged_model.py
@@ -0,0 +1,45 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+model_path = "minimind/out/merged_model"
+
+# 加载 tokenizer（通常不需要 trust_remote_code）
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+print("EOS token:", tokenizer.eos_token)
+print("EOS token ID:", tokenizer.eos_token_id)
+print("PAD token ID:", tokenizer.pad_token_id)
+
+# 加载模型：必须加 trust_remote_code=True
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    dtype=torch.float16,
+    trust_remote_code=True
+).to("cuda")  # 或 "cpu"
+
+# 推理
+messages = [
+    {"role": "user", "content": "你好"}
+]
+prompt = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True  # ← 关键！添加生成触发符
+)
+print("Prompt:", repr(prompt))
+inputs = tokenizer(
+    prompt,
+    return_tensors="pt",
+    return_token_type_ids=False  # ← 关键！
+).to(model.device)
+# 关键：用 autocast 确保所有计算在 float16 下进行
+with torch.autocast(device_type="cuda", dtype=torch.float16):
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=100,
+        do_sample=True,
+        temperature=0.3,
+        top_p=0.9
+    )
+print("Generated token IDs:", outputs[0].tolist())
+print("Decoded with special tokens:", tokenizer.decode(outputs[0], skip_special_tokens=False))
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
\ No newline at end of file
diff --git a/scripts/merge_lora_model.py b/scripts/merge_lora_model.py
new file mode 100644
index 0000000..3e2cfce
--- /dev/null
+++ b/scripts/merge_lora_model.py
@@ -0,0 +1,239 @@
+import os
+import sys
+__package__ = "scripts"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import torch
+import warnings
+from transformers import AutoTokenizer
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+import argparse
+warnings.filterwarnings('ignore', category=UserWarning)
+
+# MiniMind 配置（需与训练时一致）
+lm_config = MiniMindConfig(
+    hidden_size=512,
+    num_hidden_layers=8,
+    max_position_embeddings=32768,
+    vocab_size=6400,  # 确保与 tokenizer 一致
+    use_moe=False,
+    num_attention_heads=8,
+    num_key_value_heads=2,
+    rms_norm_eps=1e-5,
+    rope_theta=10000.0
+)
+
+
+def merge_lora_weights(model, lora_state_dict):
+    """
+    将 LoRA 权重合并到基础模型中
+    LoRA 权重格式: {module_name}.lora.A.weight 和 {module_name}.lora.B.weight
+    合并公式: W_new = W_old + B @ A
+    
+    注意: 根据 model_lora.py，LoRA 只应用于方阵线性层 (shape[0] == shape[1])
+    """
+    # 清理 state_dict 中的 'module.' 前缀（如果有）
+    cleaned_lora_dict = {k.replace('module.', ''): v for k, v in lora_state_dict.items()}
+    
+    # 收集所有 LoRA 权重
+    lora_weights = {}
+    for key in cleaned_lora_dict.keys():
+        if '.lora.' in key:
+            # 提取模块名称和权重类型
+            parts = key.split('.lora.')
+            if len(parts) == 2:
+                module_name = parts[0]
+                weight_type = parts[1]  # 'A.weight' 或 'B.weight'
+                
+                if module_name not in lora_weights:
+                    lora_weights[module_name] = {}
+                lora_weights[module_name][weight_type] = cleaned_lora_dict[key]
+    
+    print(f"找到 {len(lora_weights)} 个应用了 LoRA 的模块")
+    
+    # 合并 LoRA 权重到对应的线性层
+    merged_count = 0
+    skipped_count = 0
+    
+    for name, module in model.named_modules():
+        # 只处理线性层
+        if isinstance(module, torch.nn.Linear):
+            # 检查是否有对应的 LoRA 权重
+            if name in lora_weights:
+                lora_data = lora_weights[name]
+                if 'A.weight' in lora_data and 'B.weight' in lora_data:
+                    A = lora_data['A.weight']  # shape: [rank, in_features]
+                    B = lora_data['B.weight']  # shape: [out_features, rank]
+                    
+                    # 验证形状匹配
+                    original_weight = module.weight.data  # shape: [out_features, in_features]
+                    expected_out, expected_in = original_weight.shape
+                    
+                    # A: [rank, in_features], B: [out_features, rank]
+                    if A.shape[1] != expected_in or B.shape[0] != expected_out:
+                        print(f"  ⚠ 跳过 {name}: 形状不匹配 (A: {A.shape}, B: {B.shape}, 期望: [{expected_out}, {expected_in}])")
+                        skipped_count += 1
+                        continue
+                    
+                    # 合并: W_new = W_old + B @ A
+                    # B @ A: [out_features, rank] @ [rank, in_features] = [out_features, in_features]
+                    lora_delta = torch.matmul(B, A)  # shape: [out_features, in_features]
+                    module.weight.data = original_weight + lora_delta.to(original_weight.device)
+                    merged_count += 1
+                    print(f"  ✓ 合并 LoRA 权重到: {name} (形状: {original_weight.shape}, rank: {A.shape[0]})")
+    
+    print(f"成功合并 {merged_count} 个 LoRA 权重")
+    if skipped_count > 0:
+        print(f"跳过 {skipped_count} 个不匹配的权重")
+    
+    return model
+
+
+def merge_lora_to_minimind(lora_model_path, base_model_path, output_path, config=None, dtype=torch.float16):
+    """
+    将 MiniMind + LoRA 权重合并并保存为 safetensors 格式
+    
+    Args:
+        lora_model_path: LoRA 权重文件路径 (.pth)
+        base_model_path: 基础模型文件路径 (.pth)
+        output_path: 输出目录路径
+        config: MiniMindConfig 配置对象
+        dtype: 输出模型的数据类型
+    """
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    # 使用提供的配置或默认配置
+    if config is None:
+        config = lm_config
+    
+    print("=" * 60)
+    print("开始合并 LoRA 权重到 MiniMind 模型")
+    print("=" * 60)
+    
+    # 1. 加载基础模型
+    print(f"\n1. 加载基础模型: {base_model_path}")
+    model = MiniMindForCausalLM(config)
+    
+    if os.path.isfile(base_model_path):
+        state_dict = torch.load(base_model_path, map_location='cpu')
+        
+        # 如果 checkpoint 包含 'model' 键，则提取模型权重
+        if isinstance(state_dict, dict) and 'model' in state_dict:
+            state_dict = state_dict['model']
+            print("  ℹ 检测到完整 checkpoint，提取模型权重")
+        
+        # 清理 state_dict 中的 'module.' 前缀（如果有，用于 DDP 模型）
+        cleaned_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
+        
+        # 过滤出模型相关的键（排除 optimizer、scheduler 等）
+        model_keys = set(model.state_dict().keys())
+        filtered_state_dict = {k: v for k, v in cleaned_state_dict.items() if k in model_keys}
+        
+        missing_keys, unexpected_keys = model.load_state_dict(filtered_state_dict, strict=False)
+        print("  ✓ 基础模型加载完成")
+        if missing_keys:
+            print(f"  ⚠ 警告: {len(missing_keys)} 个键未找到（已忽略）")
+        if unexpected_keys:
+            print(f"  ℹ 信息: {len(unexpected_keys)} 个额外键（已忽略）")
+    else:
+        raise FileNotFoundError(f"基础模型文件不存在: {base_model_path}")
+    
+    # 2. 加载 LoRA 权重
+    print(f"\n2. 加载 LoRA 权重: {lora_model_path}")
+    if not os.path.isfile(lora_model_path):
+        raise FileNotFoundError(f"LoRA 权重文件不存在: {lora_model_path}")
+    
+    lora_state_dict = torch.load(lora_model_path, map_location='cpu')
+    print(f"  ✓ LoRA 权重加载完成 (包含 {len(lora_state_dict)} 个键)")
+    
+    # 3. 合并 LoRA 权重
+    print(f"\n3. 合并 LoRA 权重到基础模型...")
+    model = merge_lora_weights(model, lora_state_dict)
+    
+    # 4. 转换精度
+    print(f"\n4. 转换模型精度为 {dtype}...")
+    model = model.to(dtype)
+    model.eval()
+    
+    # 5. 保存模型为 safetensors 格式
+    print(f"\n5. 保存模型到: {output_path}")
+    os.makedirs(output_path, exist_ok=True)
+    
+    # 注册模型类以便 transformers 能够识别
+    MiniMindConfig.register_for_auto_class()
+    MiniMindForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+    
+    # 保存模型（使用 safe_serialization=True 保存为 safetensors）
+    model.save_pretrained(output_path, safe_serialization=False)
+    print("  ✓ 模型已保存为 safetensors 格式")
+    
+    # 6. 保存 tokenizer
+    print(f"\n6. 保存 tokenizer...")
+    tokenizer_path = os.path.join(os.path.dirname(__file__), '..', 'model')
+    if os.path.exists(tokenizer_path):
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+        tokenizer.save_pretrained(output_path)
+        print("  ✓ Tokenizer 保存完成")
+    else:
+        print(f"  ⚠ 警告: Tokenizer 路径不存在: {tokenizer_path}")
+    
+    # 7. 打印参数量
+    print(f"\n7. 模型信息:")
+    model_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"  总参数量: {model_params / 1e6:.2f} M = {model_params / 1e9:.2f} B")
+    print(f"  可训练参数: {trainable_params / 1e6:.2f} M")
+    
+    print("\n" + "=" * 60)
+    print(f"✅ LoRA 已合并并保存到: {output_path}")
+    print("=" * 60)
+
+
+if __name__ == '__main__':
+    # 配置路径
+    parser = argparse.ArgumentParser(description="Merge LoRA weights into base MiniMind model")
+    parser.add_argument(
+        "--lora_model_path",
+        type=str,
+        required=True,
+        # default="/home/dieu/minimind/out/lora/lora_identity_512.pth",
+        help="Path to the LoRA weights (.pth file)"
+    )
+    parser.add_argument(
+        "--base_model_path",
+        type=str,
+        required=True,
+        # default="/home/dieu/minimind/out/full_sft_512.pth",
+        help="Path to the base model (.pth file)"
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        required=True,
+        # default="/home/dieu/minimind/out/merged_model_test",
+        help="Directory to save the merged model"
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="float16",
+        choices=["float16", "float32", "bfloat16"],
+        help="Data type for the merged model (default: float16)"
+    )
+
+    args = parser.parse_args()
+
+    # 映射 dtype 字符串到 torch.dtype
+    dtype_map = {
+        "float16": torch.float16,
+        "float32": torch.float32,
+        "bfloat16": torch.bfloat16,
+    }
+    dtype = dtype_map[args.dtype]
+    # 调用合并函数
+    merge_lora_to_minimind(
+        lora_model_path=args.lora_model_path,
+        base_model_path=args.base_model_path,
+        output_path=args.output_path,
+        config=lm_config,
+        dtype=dtype
+    )

From c37d924b475cf17bc37ec20921dce4b3ad948e7b Mon Sep 17 00:00:00 2001
From: dieu <515648571@qq.com>
Date: Mon, 8 Dec 2025 15:25:34 +0800
Subject: [PATCH 2/2] [fix]:Fixed the gradient explosion problem that occurred
 during pre-training of the MOE module.

---
 model/model_minimind.py   |  7 ++--
 trainer/train_pretrain.py | 84 ++++++++++++++++++++++++++++++---------
 trainer/trainer_utils.py  | 28 ++++++++++++-
 3 files changed, 96 insertions(+), 23 deletions(-)

diff --git a/model/model_minimind.py b/model/model_minimind.py
index 259f0af..a522b03 100755
--- a/model/model_minimind.py
+++ b/model/model_minimind.py
@@ -34,7 +34,7 @@ class MiniMindConfig(PretrainedConfig):
             n_routed_experts: int = 4,
             n_shared_experts: int = 1,
             scoring_func: str = 'softmax',
-            aux_loss_alpha: float = 0.1,
+            aux_loss_alpha: float = 0.01,
             seq_aux: bool = True,
             norm_topk_prob: bool = True,
             **kwargs
@@ -249,7 +249,8 @@ class MoEGate(nn.Module):
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
-        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        # init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        init.normal_(self.weight, mean=0.0, std=0.01)
 
     def forward(self, hidden_states):
         bsz, seq_len, h = hidden_states.shape
@@ -313,7 +314,7 @@ class MOEFeedForward(nn.Module):
         flat_topk_idx = topk_idx.view(-1)
         if self.training:
             x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
-            y = torch.empty_like(x, dtype=torch.float16)
+            y = torch.empty_like(x, dtype=x.dtype)
             for i, expert in enumerate(self.experts):
                 y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype)  # 确保类型一致
             y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py
index a926288..5de170b 100644
--- a/trainer/train_pretrain.py
+++ b/trainer/train_pretrain.py
@@ -27,19 +27,32 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
         X = X.to(args.device)
         Y = Y.to(args.device)
         loss_mask = loss_mask.to(args.device)
-        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        # lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate, args.warmup_steps, args.warmup_ratio)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
         with autocast_ctx:
             res = model(X)
-            loss = loss_fct(
+            # loss = loss_fct(
+            #     res.logits.view(-1, res.logits.size(-1)),
+            #     Y.view(-1)
+            # ).view(Y.size())
+            #
+            # loss = (loss * loss_mask).sum() / loss_mask.sum()
+            # loss += res.aux_loss
+            # loss = loss / args.accumulation_steps
+            # --- 【新增监控 1】拆分 Loss 方便观察 ---
+            raw_loss = loss_fct(
                 res.logits.view(-1, res.logits.size(-1)),
                 Y.view(-1)
             ).view(Y.size())
 
-            loss = (loss * loss_mask).sum() / loss_mask.sum()
-            loss += res.aux_loss
+            main_loss = (raw_loss * loss_mask).sum() / loss_mask.sum()
+            aux_loss = res.aux_loss  # 获取辅助损失
+
+            # 这里的比例很关键，如果 aux_loss 比 main_loss 还大，模型就练歪了
+            loss = main_loss + aux_loss
             loss = loss / args.accumulation_steps
 
         scaler.scale(loss).backward()
@@ -59,10 +72,42 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
             current_loss = loss.item() * args.accumulation_steps
             current_lr = optimizer.param_groups[-1]['lr']
             eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60
-            
-            Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:')
-            
-            if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_min})
+
+            # --- 【新增监控 2】打印 Gate 梯度和 Aux Loss ---
+            # 1. 解包 DDP 模型，获取原始 model
+            raw_model = model.module if hasattr(model, 'module') else model
+
+            # 2. 统计 MoE Gate 的梯度范数 (检查路由是否还有梯度)
+            gate_grad_sum = 0.0
+            gate_count = 0
+            for name, param in raw_model.named_parameters():
+                # 筛选出所有门控层的权重
+                if "gate.weight" in name and param.grad is not None:
+                    gate_grad_sum += param.grad.norm().item()
+                    gate_count += 1
+            avg_gate_grad = gate_grad_sum / gate_count if gate_count > 0 else 0.0
+
+            # 3. 打印详细日志
+            # Main: 主任务Loss (越低越好)
+            # Aux: 负载均衡Loss (太低说明可能只顾着平均分配了，太高说明完全没分配好)
+            # G_Grad: Gate梯度 (如果是 0.0000，说明路由彻底死了，没有任何学习)
+            debug_info = f"Main:{main_loss.item():.4f} Aux:{aux_loss.item():.4f} G_Grad:{avg_gate_grad:.6f}"
+
+            Logger(
+                f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}) Loss:{current_loss:.4f} {debug_info} lr:{current_lr:.8f} Time:{eta_min}m')
+
+            if wandb:
+                wandb.log({
+                    "total_loss": current_loss,
+                    "main_loss": main_loss.item(),
+                    "aux_loss": aux_loss.item(),
+                    "gate_grad_norm": avg_gate_grad,
+                    "lr": current_lr,
+                    "epoch_Time": eta_min
+                })
+            # Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:')
+            #
+            # if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_min})
 
         if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
             model.eval()
@@ -84,25 +129,28 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Pretraining")
     parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
     parser.add_argument('--save_weight', default='pretrain', type=str, help="保存权重的前缀名")
-    parser.add_argument("--epochs", type=int, default=1, help="训练轮数（建议1轮zero或2-6轮充分训练）")
-    parser.add_argument("--batch_size", type=int, default=32, help="batch size")
-    parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
+    parser.add_argument("--epochs", type=int, default=3, help="训练轮数（建议1轮zero或2-6轮充分训练）")
+    parser.add_argument("--batch_size", type=int, default=64, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率")
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
     parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
-    parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数")
-    parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
-    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--num_workers", type=int, default=64, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=16, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1, help="梯度裁剪阈值")
     parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
-    parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
+    parser.add_argument("--save_interval", type=int, default=500, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=640, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=512, type=int, help="训练的最大截断长度")
-    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--use_moe', default=1, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
     parser.add_argument("--data_path", type=str, default="../dataset/pretrain_hq.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=是）")
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
-    parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain", help="wandb项目名")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain-640-moe-wramup", help="wandb项目名")
+    parser.add_argument("--warmup_steps", type=int, default=3000, help="Warmup步数，0表示不使用warmup（建议为总步数的1-5%）")
+    parser.add_argument("--warmup_ratio", type=float, default=0.01,
+                        help="Warmup起始学习率比例，起始lr = lr * warmup_ratio")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py
index 3697ed3..a7e5058 100644
--- a/trainer/trainer_utils.py
+++ b/trainer/trainer_utils.py
@@ -21,8 +21,32 @@ def Logger(content):
         print(content)
 
 
-def get_lr(current_step, total_steps, lr):
-    return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
+def get_lr(current_step, total_steps, lr, warmup_steps=0, warmup_ratio=0.01):
+    """
+    学习率调度：Warmup + Cosine Decay
+    
+    Args:
+        current_step: 当前步数
+        total_steps: 总步数
+        lr: 目标学习率
+        warmup_steps: warmup步数，0表示不使用warmup
+        warmup_ratio: warmup起始学习率比例，起始lr = lr * warmup_ratio
+    
+    Returns:
+        当前步数对应的学习率
+    """
+    if warmup_steps > 0 and current_step < warmup_steps:
+        # Warmup阶段：线性从 warmup_ratio * lr 增长到 lr
+        return warmup_ratio * lr + (1 - warmup_ratio) * lr * (current_step / warmup_steps)
+    else:
+        # Cosine Decay阶段：从 lr 按余弦函数衰减到 lr / 10
+        if warmup_steps > 0:
+            cosine_steps = current_step - warmup_steps
+            cosine_total = total_steps - warmup_steps
+        else:
+            cosine_steps = current_step
+            cosine_total = total_steps
+        return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * cosine_steps / cosine_total))
 
 
 def init_distributed_mode():