fixed unit tests for paritioning

Author: ajkv-google

recml/core/training/partitioning_test.py

@@ -40,7 +40,8 @@ def test_data_parallelism(
       self, partitioner_cls: type[partitioning.Partitioner]
   ):
     if partitioner_cls is partitioning.ModelParallelPartitioner:
-      kwargs = {"axes": [("data", -1), ("model", 1)], "dp_axes": 1}
+      devs = np.array(jax.devices()).reshape(-1, 1)
+      kwargs = {"axes": [("data", -1), ("model", 1)], "dp_axes": 1, "devices": devs}
     else:
       kwargs = {}
     partitioner = partitioner_cls(**kwargs)
@@ -112,8 +113,12 @@ def _eval_step(
     )
 
   def test_model_parallelism(self):
+    devs = np.array(jax.devices()).reshape(1, -1)
+    
     partitioner = partitioning.ModelParallelPartitioner(
-        axes=[("data", 1), ("model", jax.device_count())], dp_axes=1
+        axes=[("data", 1), ("model", jax.device_count())], 
+        dp_axes=1,
+        devices=devs
     )
 
     inputs = np.zeros((128, 16), dtype=np.float32)

recml/examples/train_hstu_jax.py

@@ -198,7 +198,7 @@ def experiment() -> fdl.Config[recml.Experiment]:
     """Defines the experiment structure using Fiddle configs"""
 
     max_seq_len = 50
-    batch_size = 128
+    batch_size = 64
 
     model_cfg = fdl.Config(
         HSTUModelConfig,