update training script

scripts/train.py

@@ -61,7 +61,6 @@ def main():
         tokenized["labels"] = tokenized["input_ids"].copy()
         return tokenized
 
-    
     tokenized_dataset = dataset.map(
         tokenize_function,
         batched=True,
@@ -74,7 +73,7 @@ def main():
         output_dir="./output",
         per_device_train_batch_size=1,
         gradient_accumulation_steps=8,
-        learning_rate=1e-5,
+        learning_rate=1e-4,  # Higher LR for expert training
         num_train_epochs=1,
         logging_dir="./logs",
         logging_steps=10,
@@ -88,21 +87,28 @@ def main():
         warmup_ratio=0.1,
         max_grad_norm=1.0,
     )
-    # Freeze all parameters
+    
+    # Freeze all parameters first
     for param in model.parameters():
         param.requires_grad = False
     
     # Unfreeze only the small experts and their gating networks
     for name, param in model.named_parameters():
         # Unfreeze small expert layers
-        if "mlp.experts" in name and any(f"mlp.experts.{i}." in name for i in range(0, config.num_experts, config.small_expert_count)):
+        if "mlp.small_experts" in name:
             param.requires_grad = True
             print(f"Unfreezing small expert parameter: {name}")
         
-        # Unfreeze gating network parameters
-        if "mlp.gate" in name:
+        # Unfreeze small gating network parameters
+        if "mlp.small_gate" in name:
             param.requires_grad = True
-            print(f"Unfreezing gating network parameter: {name}")
+            print(f"Unfreezing small gate parameter: {name}")
+
+    # Create custom data collator to handle router logits
+    def data_collator(features):
+        batch = default_data_collator(features)
+        batch["output_router_logits"] = True  # Ensure we get router logits for aux loss
+        return batch
 
     # Trainer
     trainer = Trainer(
@@ -110,14 +116,24 @@ def main():
         args=training_args,
         train_dataset=tokenized_dataset,
         tokenizer=tokenizer,
-        data_collator=default_data_collator,
+        data_collator=data_collator,
     )
     
     # Train
     trainer.train()
     
-    # Save
-    trainer.save_model("./final_model")
+    # Save only the small experts and gates
+    print("Saving only small experts and gates...")
+    small_expert_state_dict = {
+        name: param for name, param in model.named_parameters()
+        if "mlp.small_experts" in name or "mlp.small_gate" in name
+    }
+    
+    os.makedirs("./final_model", exist_ok=True)
+    torch.save(small_expert_state_dict, "./final_model/small_experts_and_gates.bin")
+    
+    # Also save config
+    config.save_pretrained("./final_model")
 
 if __name__ == "__main__":
     main()