[25.05.22] add MoE support

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文件训练.
 

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,

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 --num_gpus=4 --master_port=12345 pretrain_transformers.py
-#   deepspeed --include="localhost:1,2,3" --master_port=12345 pretrain_transformers.py
+#  CUDA_VISIBLE_DEVICES=1,2,3 deepspeed --include="localhost:1,2,3" --master_port=12345 pretrain_transformers.py