unfreeze only gate and experts

scripts/train.py

@@ -88,7 +88,22 @@ def main():
         warmup_ratio=0.1,
         max_grad_norm=1.0,
     )
+    # Freeze all parameters
+    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_frequency)):
+            param.requires_grad = True
+            print(f"Unfreezing small expert parameter: {name}")
+        
+        # Unfreeze gating network parameters
+        if "mlp.gate" in name:
+            param.requires_grad = True
+            print(f"Unfreezing gating network parameter: {name}")
+            
     # Trainer
     trainer = Trainer(
         model=model,