fixes to smoke test batch sizing

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -637,6 +637,20 @@ def main(argv):
     devices_array = create_device_mesh(config)
     mesh = Mesh(devices_array, config.mesh_axes)
 
+    # Auto-adjust batch sharding axes if batch_size is not divisible by (data * fsdp)
+    data_size = mesh.shape.get('data', 1)
+    fsdp_size = mesh.shape.get('fsdp', 1)
+    if config.batch_size % (data_size * fsdp_size) != 0:
+        print(f"⚠️ Warning: batch_size ({config.batch_size}) is not divisible by FSDP*Data mesh size ({fsdp_size * data_size}).")
+        print("  Automatically falling back to sharding batch dimension across 'data' axis only to prevent JAX SPMD errors.")
+        new_rules = []
+        for rule in config.logical_axis_rules:
+            if rule[0] in ('activation_batch', 'conv_batch'):
+                new_rules.append([rule[0], 'data'])
+            else:
+                new_rules.append(rule)
+        pyconfig._config.keys['logical_axis_rules'] = tuple(new_rules)
+    
     # 3. Load or generate initial latents
     # Unpacked shape: (batch_size, 32, height//8, width//8)
     latents_unpacked = load_or_generate_latents(config)

src/maxdiffusion/generate_flux2klein_9B.py

@@ -654,6 +654,20 @@ def main(argv):
     devices_array = create_device_mesh(config)
     mesh = Mesh(devices_array, config.mesh_axes)
 
+    # Auto-adjust batch sharding axes if batch_size is not divisible by (data * fsdp)
+    data_size = mesh.shape.get('data', 1)
+    fsdp_size = mesh.shape.get('fsdp', 1)
+    if config.batch_size % (data_size * fsdp_size) != 0:
+        print(f"⚠️ Warning: batch_size ({config.batch_size}) is not divisible by FSDP*Data mesh size ({fsdp_size * data_size}).")
+        print("  Automatically falling back to sharding batch dimension across 'data' axis only to prevent JAX SPMD errors.")
+        new_rules = []
+        for rule in config.logical_axis_rules:
+            if rule[0] in ('activation_batch', 'conv_batch'):
+                new_rules.append([rule[0], 'data'])
+            else:
+                new_rules.append(rule)
+        pyconfig._config.keys['logical_axis_rules'] = tuple(new_rules)
+    
     # 3. Load or generate initial latents
     # Unpacked shape: (batch_size, 32, height//8, width//8)
     latents_unpacked = load_or_generate_latents(config)