update training script

scripts/train.py

@@ -5,6 +5,7 @@ from datasets import load_dataset
 from myolmoe.modeling_myolmoe import MyOlmoeForCausalLM, OlmoeConfig
 from torch.utils.data import Dataset
 import os
+from tqdm import tqdm
 
 class CustomDataset(Dataset):
     def __init__(self, tokenizer, dataset_name="allenai/tulu-v2-sft-mixture", max_length=512):
@@ -40,67 +41,67 @@ def expand_model_with_small_experts(base_model):
     # Create new model with expanded architecture
     expanded_model = MyOlmoeForCausalLM(config)
     
-    # 1. Copy all non-expert weights exactly
-    base_state_dict = base_model.state_dict()
-    expanded_state_dict = expanded_model.state_dict()
+    # 1. First load all compatible weights (excluding expert layers)
+    base_model_state_dict = base_model.state_dict()
+    expanded_model_state_dict = expanded_model.state_dict()
     
     # Copy all non-expert parameters
-    for name, param in base_state_dict.items():
-        if "experts" not in name:  # Skip expert-specific parameters
-            expanded_state_dict[name].copy_(param)
-    
-    # 2. Copy the original experts' weights
-    for i in range(config.num_experts):
-        # Copy gate_proj weights
-        expanded_state_dict[f'model.layers.{i}.mlp.experts.{i}.gate_proj.weight'].copy_(
-            base_state_dict[f'model.layers.{i}.mlp.experts.{i}.gate_proj.weight'][:config.small_expert_intermediate_size]
-        )
-        # Copy up_proj weights
-        expanded_state_dict[f'model.layers.{i}.mlp.experts.{i}.up_proj.weight'].copy_(
-            base_state_dict[f'model.layers.{i}.mlp.experts.{i}.up_proj.weight'][:config.small_expert_intermediate_size]
-        )
-        # Copy down_proj weights (need to handle output dimension differently)
-        expanded_state_dict[f'model.layers.{i}.mlp.experts.{i}.down_proj.weight'].copy_(
-            base_state_dict[f'model.layers.{i}.mlp.experts.{i}.down_proj.weight'][:,:config.small_expert_intermediate_size]
-        )
-    
-    # 3. Initialize the gate layer for all experts (original + small)
-    # The original gate had shape (hidden_size, num_experts)
-    # New gate needs shape (hidden_size, num_experts + num_small_experts)
-    for i in range(config.num_hidden_layers):
-        original_gate = base_state_dict[f'model.layers.{i}.mlp.gate.weight']
-        new_gate = expanded_state_dict[f'model.layers.{i}.mlp.gate.weight']
+    for name, param in tqdm(base_model_state_dict.items(), desc="Copying base weights"):
+        if "mlp.experts" not in name and "mlp.gate" not in name:
+            if name in expanded_model_state_dict:
+                expanded_model_state_dict[name].copy_(param)
+    
+    # 2. Handle expert layers
+    for layer_idx in tqdm(range(config.num_hidden_layers), desc="Processing expert layers"):
+        # Copy original expert weights
+        for expert_idx in range(config.num_experts):
+            # Get the original expert weights
+            gate_proj_weight = base_model_state_dict[f'model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight']
+            up_proj_weight = base_model_state_dict[f'model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight']
+            down_proj_weight = base_model_state_dict[f'model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight']
+            
+            # Copy to new model (slicing for small experts not needed here since we're copying original experts)
+            expanded_model_state_dict[f'model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight'].copy_(gate_proj_weight)
+            expanded_model_state_dict[f'model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight'].copy_(up_proj_weight)
+            expanded_model_state_dict[f'model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight'].copy_(down_proj_weight)
+        
+        # Handle gate layer expansion
+        original_gate_weight = base_model_state_dict[f'model.layers.{layer_idx}.mlp.gate.weight']
+        new_gate_weight = expanded_model_state_dict[f'model.layers.{layer_idx}.mlp.gate.weight']
         
         # Copy original gate weights
-        new_gate[:, :config.num_experts].copy_(original_gate)
+        new_gate_weight[:, :config.num_experts].copy_(original_gate_weight)
         
-        # Initialize small experts gate weights (could use different initialization)
+        # Initialize small experts gate weights
         torch.nn.init.normal_(
-            new_gate[:, config.num_experts:], 
+            new_gate_weight[:, config.num_experts:], 
             mean=0.0, 
             std=config.initializer_range
         )
     
-    # Load the combined state dict into the new model
-    expanded_model.load_state_dict(expanded_state_dict)
+    # 3. Load the combined state dict into the new model
+    expanded_model.load_state_dict(expanded_model_state_dict, strict=False)
     
     return expanded_model
 
 def main():
     # Load base model (with only 64 experts)
     model_path = "myolmoe"
+    print("Loading base model...")
     base_model = MyOlmoeForCausalLM.from_pretrained(model_path)
     
     # Verify base model has only 64 experts
     print(f"Base model has {base_model.config.num_experts} experts")
     
     # Expand model to include small experts
+    print("Expanding model with small experts...")
     model = expand_model_with_small_experts(base_model)
     
     # Verify expanded model
     print(f"Expanded model has {model.config.num_experts} regular experts and {model.config.num_small_experts} small experts")
     
     # Prepare dataset
+    print("Preparing dataset...")
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     dataset = CustomDataset(tokenizer)
     
@@ -118,11 +119,7 @@ def main():
         eval_steps=500,
         fp16=True,
         gradient_checkpointing=True,
-        report_to="tensorboard",
-        # Important: Only train the new parameters initially
-        # Freeze original experts first, then unfreeze later
-        # You may want to modify this based on your training strategy
-        freeze_existing_experts=True  
+        report_to="tensorboard"
     )
     
     # Custom trainer to handle expert freezing
@@ -134,7 +131,7 @@ def main():
             if self.freeze_existing:
                 # Freeze all original expert parameters
                 for name, param in self.model.named_parameters():
-                    if "experts" in name and "small_experts" not in name:
+                    if "mlp.experts" in name and "small_experts" not in name:
                         param.requires_grad = False
                 print("Frozen original experts, only training small experts")
     
@@ -143,17 +140,20 @@ def main():
         args=training_args,
         train_dataset=dataset,
         eval_dataset=dataset,
-        freeze_existing_experts=training_args.freeze_existing_experts
+        freeze_existing_experts=True
     )
     
     # Train
+    print("Starting training...")
     trainer.train()
     
     # Save final model
     output_dir = "./final_model"
     os.makedirs(output_dir, exist_ok=True)
+    print(f"Saving final model to {output_dir}")
     model.save_pretrained(output_dir)
     tokenizer.save_pretrained(output_dir)
+    print("Training complete!")
 
 if __name__ == "__main__":
     main()