changed per device batch size to 0.125

Author: prishajain1

src/maxdiffusion/tests/wan_magcache_test.py

@@ -161,7 +161,7 @@ def setUpClass(cls):
             "skip_jax_distributed_system=True",
             "weights_dtype=bfloat16",
             "activations_dtype=bfloat16",
-            "per_device_batch_size=0.25",
+            "per_device_batch_size=0.125",
             "ici_data_parallelism=2",
             "ici_fsdp_parallelism=1",
             "ici_context_parallelism=4",
@@ -321,7 +321,7 @@ def setUpClass(cls):
             "skip_jax_distributed_system=True",
             "weights_dtype=bfloat16",
             "activations_dtype=bfloat16",
-            "per_device_batch_size=0.25",
+            "per_device_batch_size=0.125",
             "ici_data_parallelism=2",
             "ici_fsdp_parallelism=1",
             "ici_context_parallelism=4",