[2025.05.22] add MoE support

.gitignore

@@ -1,6 +1,9 @@
 corpus/wikipedia-cn-20230720-filtered.json
 corpus/wikipedia.txt
-corpus/*.jsonl
+corpus/*/*.jsonl
 dolly_env_local
 llm_dolly.egg-info
 src_llm/__pycache__
+weights_MoE
+weights_llm
+weights_tokenizer/*

__init__.py

@@ -1 +1,5 @@
 __version__ = "0.0.1"
+from src_llm_moe.configuration_dollyMoE import DollyMoEConfig
+from src_llm_moe.modeling_dollyMoE import DollyMoEModel, DollyMoEPreTrainedModel
+from src_llm.configuration_dolly import DollyConfig
+from src_llm.modeling_dolly import DollyModel, DollyPreTrainedModel

src_llm/configuration_dolly.py

@@ -28,11 +28,11 @@ class DollyConfig(PretrainedConfig):
     def __init__(
         self,
         vocab_size=151936,
-        hidden_size=4096,
-        intermediate_size=22016,
-        num_hidden_layers=32,
-        num_attention_heads=32,
-        num_key_value_heads=32,
+        hidden_size=1024,
+        intermediate_size=3072,
+        num_hidden_layers=28,
+        num_attention_heads=16,
+        num_key_value_heads=8,
         head_dim=128,
         hidden_act="silu",
         max_position_embeddings=32768,
@@ -40,7 +40,7 @@ class DollyConfig(PretrainedConfig):
         rms_norm_eps=1e-6,
         use_cache=True,
         tie_word_embeddings=False,
-        rope_theta=10000.0,
+        rope_theta=1000000.0,
         rope_scaling=None,
         attention_bias=False,
         use_sliding_window=False,

src_llm_moe/__init__.py

@@ -0,0 +1,2 @@
+from .configuration_dollyMoE import DollyMoEConfig
+from .modeling_dollyMoE import DollyMoEModel, DollyMoEPreTrainedModel

src_llm_moe/configuration_dollyMoE.py

@@ -0,0 +1,95 @@
+from transformers import PretrainedConfig
+from transformers import logging
+from transformers.modeling_rope_utils import rope_config_validation
+logger = logging.get_logger(__name__)
+
+
+class DollyMoEConfig(PretrainedConfig):
+    model_type = "dolly_moe"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.experts.*.gate_proj": "colwise",
+        "layers.*.mlp.experts.*.up_proj": "colwise",
+        "layers.*.mlp.experts.*.down_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=1024,
+        intermediate_size=3072,
+        num_hidden_layers=28,
+        num_attention_heads=16,
+        num_key_value_heads=8,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=1000000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        attention_dropout=0.0,
+        decoder_sparse_step=1,
+        moe_intermediate_size=768,
+        num_experts_per_tok=4,
+        num_experts=8,
+        norm_topk_prob=False,
+        output_router_logits=False,
+        router_aux_loss_coef=0.001,
+        mlp_only_layers=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if use_sliding_window else None
+        self.max_window_layers = max_window_layers
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+        self.decoder_sparse_step = decoder_sparse_step
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.num_experts = num_experts
+        self.norm_topk_prob = norm_topk_prob
+        self.output_router_logits = output_router_logits
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+__all__ = ["DollyMoEConfig"]

src_pretrain/pretrain_deepspeed_MoE.py

@@ -0,0 +1,120 @@
+import torch
+from src_llm_moe.modeling_dollyMoE import DollyMoEForCausalLM
+from transformers import (
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+    AutoTokenizer,
+    set_seed
+)
+from datasets import load_dataset
+from accelerate import Accelerator
+import os
+import logging
+import glob
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+accelerator = Accelerator()
+set_seed(42)
+
+model_name = "/mnt/han.luzhi/dolly/weights_MoE"  
+resume_option = None                     
+# resume_option = True                    
+# resume_option = "/mnt/han.luzhi/dolly_llm/checkpoints_ds/checkpoint-500"  
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+if isinstance(resume_option, str) and os.path.exists(resume_option):
+    logger.info(f"从指定检查点 {resume_option} 加载模型...")
+    model = DollyForCausalLM.from_pretrained(resume_option)
+else:
+    logger.info("未找到检查点，从预训练模型加载...")
+    logger.info(f"从预训练模型 {model_name} 加载...")
+    model = DollyMoEForCausalLM.from_pretrained(model_name)
+model.gradient_checkpointing_enable()
+
+data_files = glob.glob("/mnt/han.luzhi/dolly/corpus/*.jsonl")
+logger.info(f"找到训练语料: {data_files}")
+
+dataset = load_dataset('json', data_files=data_files, split='train')
+# dataset = load_dataset('text', data_files={'train': '/mnt/han.luzhi/dolly_llm/corpus/wikipedia.txt'})
+
+def preprocess_function(examples):
+    return tokenizer(
+        examples['text'],
+        truncation=True,
+        max_length=512,
+        padding='max_length',
+        return_tensors='pt'
+    )
+
+tokenized_datasets = dataset.map(
+    preprocess_function,
+    batched=True,
+    # remove_columns=['text'], 
+    remove_columns=dataset.column_names,
+    num_proc=16  
+)
+
+data_collator = DataCollatorForLanguageModeling(
+    tokenizer=tokenizer,
+    mlm=False 
+)
+
+training_args = TrainingArguments(
+    output_dir="/mnt/han.luzhi/dolly/checkpoints_MoE",
+    deepspeed="/mnt/han.luzhi/dolly/deepspeed_config/ds_stage2.json",  
+    do_eval=False,
+    save_strategy="steps",
+    save_steps=500,
+    learning_rate=1e-4,
+    per_device_train_batch_size=16,
+    num_train_epochs=2,
+    logging_dir="/mnt/han.luzhi/dolly/logs",
+    logging_steps=20,
+    save_total_limit=2,
+    fp16=True,  
+    dataloader_num_workers=3,
+    gradient_accumulation_steps=2,
+    report_to="tensorboard",
+    optim="adamw_torch",
+    lr_scheduler_type="cosine",
+    warmup_steps=500,
+    weight_decay=0.01,
+    local_rank=int(os.environ.get("LOCAL_RANK", 0)),
+    disable_tqdm=False,
+    logging_first_step=True,
+    log_level="info",
+    ddp_find_unused_parameters=False,
+    overwrite_output_dir=True,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets,
+    processing_class=tokenizer, 
+    data_collator=data_collator,
+)
+
+logger.info("***** 开始训练 *****")
+logger.info(f" 进程数 = {accelerator.num_processes}")
+logger.info(f" 设备 = {accelerator.device}")
+logger.info(f" 总批大小 = {training_args.per_device_train_batch_size * accelerator.num_processes * training_args.gradient_accumulation_steps}")
+
+if resume_option !=None:
+    trainer.train(resume_from_checkpoint=resume_option)
+else:
+    trainer.train()
+
+
+if accelerator.is_main_process: 
+    trainer.save_model("/mnt/han.luzhi/dolly/chechpoints_MoE/final")
+    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

test/compute_dollyMoE.py

@@ -0,0 +1,18 @@
+from transformers.modeling_utils import PreTrainedModel
+import torch
+from torch import nn
+from src_llm_moe.configuration_dollyMoE import DollyMoEConfig
+from src_llm_moe.modeling_dollyMoE import DollyMoEModel,DollyMoEForCausalLM
+model = DollyMoEForCausalLM(DollyMoEConfig())
+model.save_pretrained("./weights_MoE")
+model.config.save_pretrained("./weights_MoE")
+
+config = DollyMoEConfig.from_pretrained("./weights_MoE")
+model = DollyMoEForCausalLM.from_pretrained("./weights_MoE", config=config)
+
+input_ids = torch.tensor([[101, 1024, 102]]).long() 
+outputs = model(input_ids)
+print(outputs)
+
+model_size = sum(t.numel() for t in model.parameters())
+print(f"Model Size: {model_size/1000**2:.1f}M parameters")