debug to train script

scripts/train.py

@@ -9,6 +9,7 @@ from tqdm import tqdm
 
 class CustomDataset(Dataset):
     def __init__(self, tokenizer, dataset_name="allenai/tulu-v2-sft-mixture", max_length=512):
+        print(f"# DEBUG: Loading dataset '{dataset_name}' with max length {max_length}")
         self.dataset = load_dataset(dataset_name, split="train")  # Use train split
         self.tokenizer = tokenizer
         self.max_length = max_length
@@ -26,6 +27,10 @@ class CustomDataset(Dataset):
             truncation=True,
             return_tensors="pt"
         )
+        # DEBUG: Print the first few token IDs for inspection
+        if idx == 0:
+            print(f"# DEBUG: Sample input text: {text[:100]}")
+            print(f"# DEBUG: Tokenized input_ids[:10]: {encoding['input_ids'][0][:10]}")
         return {
             "input_ids": encoding["input_ids"].squeeze(),
             "attention_mask": encoding["attention_mask"].squeeze(),
@@ -33,84 +38,69 @@ class CustomDataset(Dataset):
         }
 
 def expand_model_with_small_experts(base_model):
-    # Create new config with small experts
+    print("# DEBUG: Expanding model with small experts...")
     config = base_model.config
-    config.num_small_experts = 64  # Add 64 small experts
-    config.small_expert_intermediate_size = config.intermediate_size // 2  # Half size
-    
-    # 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()
     
-    # Copy all non-expert parameters
+    print("# DEBUG: Copying non-expert parameters...")
     for name, param in base_state_dict.items():
-        if "experts" not in name and "gate" not in name:  # Skip expert-specific and gate parameters
-            expanded_state_dict[name].copy_(param)
+        if "experts" not in name and "gate" not in name:
+            if name in expanded_state_dict:
+                expanded_state_dict[name].copy_(param)
+            else:
+                print(f"# DEBUG: Skipped non-expert param {name} (not found in expanded model)")
     
-    # 2. Copy the original experts' weights
+    print("# DEBUG: Copying expert weights...")
     for i in range(config.num_experts):
-        # Copy expert weights
         for proj in ['gate_proj', 'up_proj', 'down_proj']:
-            if f'model.layers.{i}.mlp.experts.{i}.{proj}.weight' in base_state_dict:
-                orig_weight = base_state_dict[f'model.layers.{i}.mlp.experts.{i}.{proj}.weight']
+            key = f'model.layers.{i}.mlp.experts.{i}.{proj}.weight'
+            if key in base_state_dict:
+                orig_weight = base_state_dict[key]
                 if proj == 'down_proj':
-                    # For down_proj, we need to handle the output dimension
-                    expanded_state_dict[f'model.layers.{i}.mlp.experts.{i}.{proj}.weight'].copy_(
-                        orig_weight[:, :config.small_expert_intermediate_size]
-                    )
+                    expanded_state_dict[key].copy_(orig_weight[:, :config.small_expert_intermediate_size])
                 else:
-                    # For gate_proj and up_proj, we take the first half
-                    expanded_state_dict[f'model.layers.{i}.mlp.experts.{i}.{proj}.weight'].copy_(
-                        orig_weight[:config.small_expert_intermediate_size]
-                    )
-    
-    # 3. Handle the gate layer expansion
+                    expanded_state_dict[key].copy_(orig_weight[:config.small_expert_intermediate_size])
+                print(f"# DEBUG: Copied {proj} weights for expert {i}")
+            else:
+                print(f"# DEBUG: Missing {key} in base model")
+
+    print("# DEBUG: Expanding and initializing gate weights...")
     for i in range(config.num_hidden_layers):
         gate_key = f'model.layers.{i}.mlp.gate.weight'
         if gate_key in base_state_dict:
             original_gate = base_state_dict[gate_key]
             new_gate = expanded_state_dict[gate_key]
-            
-            # Copy original gate weights for the first 64 experts
             new_gate[:, :config.num_experts].copy_(original_gate)
-            
-            # Initialize small experts gate weights with lower magnitude
-            torch.nn.init.normal_(
-                new_gate[:, config.num_experts:], 
-                mean=0.0, 
-                std=config.initializer_range * 0.1  # Smaller initialization
-            )
-    
-    # Load the combined state dict into the new model
-    expanded_model.load_state_dict(expanded_state_dict, strict=False)  # Note strict=False
+            torch.nn.init.normal_(new_gate[:, config.num_experts:], mean=0.0, std=config.initializer_range * 0.1)
+            print(f"# DEBUG: Initialized gate for layer {i}")
+        else:
+            print(f"# DEBUG: Missing gate weight {gate_key}")
+
+    print("# DEBUG: Loading expanded state dict into model...")
+    expanded_model.load_state_dict(expanded_state_dict, strict=False)
     
     return expanded_model
 
 def main():
-    # Load base model (with only 64 experts)
     model_path = "myolmoe"
-    print("Loading base model...")
+    print("# DEBUG: 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...")
+    print(f"# DEBUG: Base model has {base_model.config.num_experts} experts")
+
+    print("# DEBUG: Calling expand_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...")
+    print(f"# DEBUG: Expanded model has {model.config.num_experts} regular experts and {model.config.num_small_experts} small experts")
+
+    print("# DEBUG: Loading tokenizer and preparing dataset...")
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     dataset = CustomDataset(tokenizer)
-    
-    # Training arguments
+
+    print("# DEBUG: Setting up training arguments...")
     training_args = TrainingArguments(
         output_dir="./output",
         per_device_train_batch_size=4,
@@ -126,20 +116,22 @@ def main():
         gradient_checkpointing=True,
         report_to="tensorboard"
     )
-    
-    # Custom trainer to handle expert freezing
+
     class MoETrainer(Trainer):
         def __init__(self, *args, **kwargs):
             self.freeze_existing = kwargs.pop('freeze_existing_experts', False)
             super().__init__(*args, **kwargs)
-            
+
             if self.freeze_existing:
-                # Freeze all original expert parameters
+                print("# DEBUG: Freezing original expert parameters...")
+                frozen_count = 0
                 for name, param in self.model.named_parameters():
                     if "mlp.experts" in name and "small_experts" not in name:
                         param.requires_grad = False
-                print("Frozen original experts, only training small experts")
-    
+                        frozen_count += 1
+                print(f"# DEBUG: Total frozen expert parameters: {frozen_count}")
+
+    print("# DEBUG: Initializing trainer...")
     trainer = MoETrainer(
         model=model,
         args=training_args,
@@ -147,18 +139,16 @@ def main():
         eval_dataset=dataset,
         freeze_existing_experts=True
     )
-    
-    # Train
-    print("Starting training...")
+
+    print("# DEBUG: 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}")
+    print(f"# DEBUG: Saving final model to {output_dir}...")
     model.save_pretrained(output_dir)
     tokenizer.save_pretrained(output_dir)
-    print("Training complete!")
+    print("# DEBUG: Training complete!")
 
 if __name__ == "__main__":
-    main()
+    main()