From a229a11b0eb43e29922ec4b2576aa381bad0a0e5 Mon Sep 17 00:00:00 2001
From: coffee <hanluzhi0429@gmail.com>
Date: Thu, 22 May 2025 17:45:47 +0800
Subject: [PATCH] [25.05.22] add MoE support

---
 README.md                              | 3 +++
 src_llm_moe/modeling_dollyMoE.py       | 6 ++++--
 src_pretrain/pretrain_deepspeed_MoE.py | 6 +++---
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 8d6c2dd..634f163 100755
--- a/README.md
+++ b/README.md
@@ -4,6 +4,9 @@ LLM-Dolly is a custom LLM built from scratch
 
 ## 🗓️ 更新日志
 
+#### 2025-05-22
+- 📝 添加支持MoE模型，训练资源占用不稳定，Experts=8, experts_per_tok=4,GPU显存从60%-94%跳动,单个GPU利用率0%-100%跳动。
+
 #### 2025-05-20
 - 📝 添加支持jsonl文件训练.
 
diff --git a/src_llm_moe/modeling_dollyMoE.py b/src_llm_moe/modeling_dollyMoE.py
index 64941d1..bf900fa 100644
--- a/src_llm_moe/modeling_dollyMoE.py
+++ b/src_llm_moe/modeling_dollyMoE.py
@@ -803,7 +803,8 @@ class DollyMoEForCausalLM(DollyMoEPreTrainedModel, GenerationMixin):
 
         loss = None
         if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.vocab_size, **kwargs)
+        
         aux_loss = None
         if output_router_logits:
             aux_loss = load_balancing_loss_func(
@@ -812,7 +813,8 @@ class DollyMoEForCausalLM(DollyMoEPreTrainedModel, GenerationMixin):
                 self.num_experts_per_tok,
                 attention_mask,
             )
-        if labels is not None:
+
+            if labels is not None:
                 loss += self.router_aux_loss_coef * aux_loss.to(loss.device) 
         return MoeCausalLMOutputWithPast(
             loss=loss,
diff --git a/src_pretrain/pretrain_deepspeed_MoE.py b/src_pretrain/pretrain_deepspeed_MoE.py
index bfacf05..c9dca5e 100644
--- a/src_pretrain/pretrain_deepspeed_MoE.py
+++ b/src_pretrain/pretrain_deepspeed_MoE.py
@@ -29,7 +29,7 @@ if tokenizer.pad_token is None:
 
 if isinstance(resume_option, str) and os.path.exists(resume_option):
     logger.info(f"从指定检查点 {resume_option} 加载模型...")
-    model = DollyForCausalLM.from_pretrained(resume_option)
+    model = DollyMoEForCausalLM.from_pretrained(resume_option)
 else:
     logger.info("未找到检查点，从预训练模型加载...")
     logger.info(f"从预训练模型 {model_name} 加载...")
@@ -116,5 +116,5 @@ if accelerator.is_main_process:
     tokenizer.save_pretrained("/mnt/han.luzhi/dolly/chechpoints_MoE/final")
     logger.info("***** 训练完成，模型已保存 *****")
 
-##  deepspeed --num_gpus=4 --master_port=12345 pretrain_transformers.py
-#   deepspeed --include="localhost:1,2,3" --master_port=12345 pretrain_transformers.py
\ No newline at end of file
+## deepspeed --num_gpus=4 --master_port=12345 pretrain_transformers.py
+#  CUDA_VISIBLE_DEVICES=1,2,3 deepspeed --include="localhost:1,2,3" --master_port=12345 pretrain_transformers.py
\ No newline at end of file