Update wave_tank_surgeon.py

wave_tank_surgeon.py

@@ -436,7 +436,7 @@ def collect_io_pairs(layer_info, model, dataloader, device, max_batches=50):
 
 
 def train_replacement(layer_info, inputs, outputs, device, ffn_grid=16, attn_grid=4,
-                      steps=500, lr=1e-3, batch_size=32):
+                      steps=5000, lr=1e-3, batch_size=32):
     """Train a wave tank replacement to mimic the original layer's behavior."""
     path = layer_info.path
     d_model = layer_info.d_model
@@ -568,11 +568,11 @@ def main():
     parser.add_argument('--model', required=True, help='HuggingFace model name or local path')
     parser.add_argument('--output', default=None, help='Output directory for replaced model')
     parser.add_argument('--ffn-grid', type=int, default=16, help='Grid size for FFN replacement (default: 16)')
-    parser.add_argument('--attn-grid', type=int, default=4, help='Grid size for attention replacement (default: 4)')
-    parser.add_argument('--steps', type=int, default=500, help='Training steps per layer (default: 500)')
+    parser.add_argument('--attn-grid', type=int, default=16, help='Grid size for attention replacement (default: 4)')
+    parser.add_argument('--steps', type=int, default=5000, help='Training steps per layer (default: 500)')
     parser.add_argument('--lr', type=float, default=1e-3, help='Learning rate (default: 1e-3)')
-    parser.add_argument('--batch-size', type=int, default=4, help='Batch size for IO collection (default: 4)')
-    parser.add_argument('--max-batches', type=int, default=50, help='Max batches for IO pair collection (default: 50)')
+    parser.add_argument('--batch-size', type=int, default=16, help='Batch size for IO collection (default: 4)')
+    parser.add_argument('--max-batches', type=int, default=500, help='Max batches for IO pair collection (default: 50)')
     parser.add_argument('--device', default=None, help='Device (default: auto-detect)')
     parser.add_argument('--skip-attn', action='store_true', help='Skip attention replacement (FFN only)')
     parser.add_argument('--skip-ffn', action='store_true', help='Skip FFN replacement (attention only)')