drop jit for loss and remove bias sharding for tp

src/partitioning/partition.py

@@ -13,18 +13,9 @@
 from flax.core.frozen_dict import freeze
 from flax.traverse_util import flatten_dict, unflatten_dict
 from jax.sharding import Mesh, PartitionSpec
-from jax.sharding import PositionalSharding, NamedSharding
+from jax.sharding import NamedSharding
 from typing import Callable
 
-def setup_dp_mesh():
-    """
-    Creates jax device mesh for data-parallel training
-    """
-    devices = np.asarray(jax.devices())
-    mesh = Mesh(devices, ["dp"])
-
-    return mesh
-
 
 def _match(qs, ks):
     """Return True if regexes in qs match any window of strings in tuple ks."""
@@ -102,7 +93,7 @@ def _get_partition_rules_tp(mesh: Mesh):
         (("(query_proj|key_proj|value_proj)", "kernel"), NamedSharding(mesh,PartitionSpec(None, "mp"))),
         (("residual_out", "kernel"), NamedSharding(mesh,PartitionSpec("mp", None))),
         (("(query_proj|key_proj|value_proj)", "bias"), NamedSharding(mesh,PartitionSpec(None))),
-        (("residual_out", "bias"), NamedSharding(mesh,PartitionSpec("mp"))),
+        (("residual_out", "bias"), NamedSharding(mesh,PartitionSpec("None"))),
         # MLP
         (("fc_in", "kernel"), NamedSharding(mesh,PartitionSpec(None, "mp"))),
         (("fc_residual", "kernel"), NamedSharding(mesh,PartitionSpec("mp", None))),

src/utils/losses.py

@@ -1,12 +1,10 @@
 """
 Loss function
 """
-import flax.linen as nn
+import jax
 import jax.numpy as jnp
-from jax import jit
 
 
-@jit
 def cross_entropy_loss(labels: jnp.array, logits: jnp.array) -> jnp.array:
     """Standard Cross Entropy Loss function
 
@@ -19,5 +17,5 @@ def cross_entropy_loss(labels: jnp.array, logits: jnp.array) -> jnp.array:
     """
 
     return (
-        -jnp.sum(labels * nn.log_softmax(logits.astype(jnp.float32), axis=-1), axis=-1)
+        -jnp.sum(labels * jax.nn.log_softmax(logits.astype(jnp.float32), axis=-1), axis=-1)
     )