[WIP] partial nnx impl

MaxText/layers/models.py

@@ -144,7 +144,7 @@ def __call__(
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )
 
-    return layer_output, None if cfg.scan_layers else layer_output
+    return (layer_output, None) if cfg.scan_layers else layer_output
 
 class SequentialBlockDecoderLayers(nn.Module):
   """Sequential unscanned series of decoder layers."""
@@ -264,6 +264,8 @@ def __call__(
     y = self.shared_embedding(decoder_input_tokens.astype("int32"))
     y = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(y, deterministic=deterministic)
     y = y.astype(cfg.dtype)
+
+    self.sow("intermediates", "rdyro_shared_embedding", y)
 
     if cfg.use_untrainable_positional_embedding:
       y = PositionalEmbedding(cfg.base_emb_dim)(y, decoder_positions)
@@ -337,6 +339,7 @@ def __call__(
         policy=policy,
         static_argnums=(-1, -2, -3, -4, -5),
     )
+    self.sow("intermediates", "rdyro_before_scan", y)
     if cfg.using_pipeline_parallelism:
         if cfg.num_layers_per_pipeline_stage == 1:
           stage_module = BlockLayer(config=cfg, mesh=mesh, quant=self.quant)
@@ -370,6 +373,8 @@ def __call__(
               deterministic,
               model_mode,
           )
+          self.sow("intermediates", f"rdyro_after_layer_{lyr}", y)
+    self.sow("intermediates", "rdyro_after_scan", y)
 
     y = self.get_norm_layer()(
         dtype=cfg.dtype,
@@ -379,6 +384,7 @@ def __call__(
         kernel_axes=("norm",),
     )(y)
     y = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(y, deterministic=deterministic)
+    self.sow("intermediates", "rdyro_after_norm", y)
 
     # [batch, length, emb_dim] -> [batch, length, vocab_size]
     if cfg.logits_via_embedding:
@@ -401,6 +407,7 @@ def __call__(
       )(
           y
       )  # We do not quantize the logits matmul.
+    self.sow("intermediates", "rdyro_after_logits", logits)
     logits = nn.with_logical_constraint(logits, ("activation_embed_and_logits_batch", "activation_length", "activation_vocab"))
     logits = logits.astype(jnp.float32)
     return logits

MaxText/nnx_layers/__init__.py

@@ -0,0 +1,47 @@
+import jax
+from flax import nnx
+from flax import linen as nn
+from flax.core import meta
+from flax.linen.spmd import LogicallyPartitioned
+
+def _maybe_unbox_value(x):
+  return x
+  if isinstance(x, meta.Partitioned):
+    return x.unbox()
+  else:
+    return x
+  #return x.unbox() if isinstance(x, meta.Partitioned) else x
+
+class LinenToNNX(nnx.Module):
+  def __init__(self, module: nn.Module, rngs=None):
+    assert rngs is not None, "You must provide `rngs=`"
+    self.linen_module = module
+    self.initialized, self.rngs, self.linen_state = False, rngs, None
+    self.rngs = None
+    # flax's flags to keep track of train / eval
+    self.use_running_average, self.deterministic = False, False
+
+  #@nnx.jit
+  def __call__(self, *args, **kw):
+    if not self.initialized:
+      rngs = nnx.Rngs(0)
+      #self.linen_state = self.linen_module.init(self.rngs(), *args, **kw)
+      self.linen_state = self.linen_module.init(rngs(), *args, **kw)
+      self.linen_state["params"] = jax.tree.map(lambda x: nnx.Param(
+        _maybe_unbox_value(x)), self.linen_state["params"], 
+        is_leaf=lambda x: isinstance(x, meta.Partitioned))
+      for key in (set(self.linen_state.keys()) - set(["params"])):
+        self.linen_state[key] = jax.tree.map(
+          lambda x: nnx.Variable(_maybe_unbox_value(x)), self.linen_state[key], 
+          is_leaf=lambda x: isinstance(x, meta.Partitioned))
+      self.initialized = True
+      #del self.rngs
+    mutable_keys = [k for k in self.linen_state.keys() if k != "params"]
+    linen_state = jax.tree.map(lambda x: x.value, self.linen_state)
+    ret, updates = self.linen_module.apply(linen_state, *args, **kw, 
+                                           mutable=mutable_keys)
+    #print(f"Update keys: {mutable_keys}")
+    if not self.use_running_average or not self.deterministic:
+      updates = jax.tree.map(lambda x: nnx.Variable(x), updates)
+      nnx.update(self, nnx.state({"linen_state": updates}))
+    return ret