[flax:examples:lm1b_nnx] Update example to work internally. #jax-fixit.

PiperOrigin-RevId: 838909002

Author: danielsuo

examples/lm1b_nnx/configs/default.py

@@ -107,10 +107,10 @@ class Config:
   # Parallelism
   mesh_axes: tuple[str, ...] = ('data', 'fsdp', 'tensor')
   axis_rules: MeshRules = MeshRules(
-    embed='fsdp',
-    mlp='tensor',
-    kv='tensor',
-    vocab='tensor',
+      embed=None,
+      mlp=None,
+      kv=None,
+      vocab=None,
   )
   data_sharding: tuple[str, ...] = ('data',)
   # One axis for each parallelism type may hold a placeholder (-1)

examples/lm1b_nnx/models.py

@@ -66,6 +66,16 @@ def replace(self, **kwargs):
     return dataclasses.replace(self, **kwargs)
 
 
+def maybe_with_partitioning(
+    init_fn: nnx.Initializer,
+    names: tuple[str | None, ...],
+    config: TransformerConfig,
+) -> nnx.Initializer:
+  if all(name is None for name in names):
+    return init_fn
+  return nnx.with_partitioning(init_fn, names)
+
+
 def shift_right(x: jax.Array, axis: int = 1):
   """Shift the input to the right by padding and slicing on axis."""
   pad_widths: list[tuple[int, int]] = [(0, 0)] * len(x.shape)
@@ -202,32 +212,36 @@ def __init__(self, config: TransformerConfig, *, rngs: nnx.Rngs):
     self.config = config
 
     self.linear1 = nnx.Linear(
-      config.emb_dim,
-      config.mlp_dim,
-      dtype=config.dtype,
-      kernel_init=nnx.with_partitioning(
-        config.kernel_init,
-        config.axis_rules('embed', 'mlp'),
-      ),
-      bias_init=nnx.with_partitioning(
-        config.bias_init,
-        config.axis_rules('mlp'),
-      ),
-      rngs=rngs,
+        config.emb_dim,
+        config.mlp_dim,
+        dtype=config.dtype,
+        kernel_init=maybe_with_partitioning(
+            config.kernel_init,
+            config.axis_rules('embed', 'mlp'),
+            config,
+        ),
+        bias_init=maybe_with_partitioning(
+            config.bias_init,
+            config.axis_rules('mlp'),
+            config,
+        ),
+        rngs=rngs,
     )
     self.linear2 = nnx.Linear(
-      config.mlp_dim,
-      config.emb_dim,
-      dtype=config.dtype,
-      kernel_init=nnx.with_partitioning(
-        config.kernel_init,
-        config.axis_rules('mlp', 'embed'),
-      ),
-      bias_init=nnx.with_partitioning(
-        config.bias_init,
-        config.axis_rules('embed'),
-      ),
-      rngs=rngs,
+        config.mlp_dim,
+        config.emb_dim,
+        dtype=config.dtype,
+        kernel_init=maybe_with_partitioning(
+            config.kernel_init,
+            config.axis_rules('mlp', 'embed'),
+            config,
+        ),
+        bias_init=maybe_with_partitioning(
+            config.bias_init,
+            config.axis_rules('embed'),
+            config,
+        ),
+        rngs=rngs,
     )
     self.dropout = nnx.Dropout(rate=config.dropout_rate)
 
@@ -252,47 +266,51 @@ def __init__(self, config: TransformerConfig, *, rngs: nnx.Rngs):
     self.config = config
 
     self.ln1 = nnx.LayerNorm(
-      num_features=config.emb_dim,
-      dtype=config.dtype,
-      bias_init=nnx.with_partitioning(
-        nnx.initializers.zeros_init(),
-        config.axis_rules('embed'),
-      ),
-      scale_init=nnx.with_partitioning(
-        nnx.initializers.ones_init(),
-        config.axis_rules('embed'),
-      ),
-      rngs=rngs,
+        num_features=config.emb_dim,
+        dtype=config.dtype,
+        bias_init=maybe_with_partitioning(
+            nnx.initializers.zeros_init(),
+            config.axis_rules('embed'),
+            config,
+        ),
+        scale_init=maybe_with_partitioning(
+            nnx.initializers.ones_init(),
+            config.axis_rules('embed'),
+            config,
+        ),
+        rngs=rngs,
     )
     self.ln2 = nnx.LayerNorm(
-      num_features=config.emb_dim,
-      dtype=config.dtype,
-      bias_init=nnx.with_partitioning(
-        nnx.initializers.zeros_init(),
-        config.axis_rules('embed'),
-      ),
-      scale_init=nnx.with_partitioning(
-        nnx.initializers.ones_init(),
-        config.axis_rules('embed'),
-      ),
-      rngs=rngs,
+        num_features=config.emb_dim,
+        dtype=config.dtype,
+        bias_init=maybe_with_partitioning(
+            nnx.initializers.zeros_init(),
+            config.axis_rules('embed'),
+            config,
+        ),
+        scale_init=maybe_with_partitioning(
+            nnx.initializers.ones_init(),
+            config.axis_rules('embed'),
+            config,
+        ),
+        rngs=rngs,
     )
     self.attention = nnx.MultiHeadAttention(
-      num_heads=config.num_heads,
-      in_features=config.emb_dim,
-      qkv_features=config.qkv_dim,
-      dtype=config.dtype,
-      kernel_init=nnx.with_partitioning(
-        config.kernel_init, config.axis_rules('embed', 'kv')
-      ),
-      bias_init=nnx.with_partitioning(
-        config.bias_init, config.axis_rules('embed')
-      ),
-      use_bias=False,
-      broadcast_dropout=False,
-      dropout_rate=config.attention_dropout_rate,
-      rngs=rngs,
-      keep_rngs=False,
+        num_heads=config.num_heads,
+        in_features=config.emb_dim,
+        qkv_features=config.qkv_dim,
+        dtype=config.dtype,
+        kernel_init=maybe_with_partitioning(
+            config.kernel_init, config.axis_rules('embed', 'kv'), config
+        ),
+        bias_init=maybe_with_partitioning(
+            config.bias_init, config.axis_rules('embed'), config
+        ),
+        use_bias=False,
+        broadcast_dropout=False,
+        dropout_rate=config.attention_dropout_rate,
+        rngs=rngs,
+        keep_rngs=False,
     )
     self.mlp = MlpBlock(config=config, rngs=rngs)
     self.dropout = nnx.Dropout(rate=config.dropout_rate)
@@ -348,13 +366,14 @@ def __init__(
     # Target Embedding
     if self.shared_embedding is None:
       self.output_embed = nnx.Embed(
-        num_embeddings=config.output_vocab_size,
-        features=config.emb_dim,
-        embedding_init=nnx.with_partitioning(
-          nnx.initializers.normal(stddev=1.0),
-          config.axis_rules('vocab', 'embed'),
-        ),
-        rngs=rngs,
+          num_embeddings=config.output_vocab_size,
+          features=config.emb_dim,
+          embedding_init=maybe_with_partitioning(
+              nnx.initializers.normal(stddev=1.0),
+              config.axis_rules('vocab', 'embed'),
+              config,
+          ),
+          rngs=rngs,
       )
     else:
       self.output_embed = self.shared_embedding
@@ -366,28 +385,28 @@ def __init__(
       setattr(self, f'encoderdecoderblock_{idx}', layer)
 
     self.encoderdecoder_norm = nnx.LayerNorm(
-      num_features=config.emb_dim,
-      dtype=config.dtype,
-      bias_init=nnx.with_partitioning(
-        nnx.initializers.zeros_init(), config.axis_rules('embed')
-      ),
-      scale_init=nnx.with_partitioning(
-        nnx.initializers.ones_init(), config.axis_rules('embed')
-      ),
-      rngs=rngs,
-    )
-    if not config.logits_via_embedding:
-      self.logitdense = nnx.Linear(
-        in_features=config.emb_dim,
-        out_features=config.output_vocab_size,
+        num_features=config.emb_dim,
         dtype=config.dtype,
-        kernel_init=nnx.with_partitioning(
-          config.kernel_init, config.axis_rules('embed', 'vocab')
+        bias_init=maybe_with_partitioning(
+            nnx.initializers.zeros_init(), config.axis_rules('embed'), config
         ),
-        bias_init=nnx.with_partitioning(
-          config.bias_init, config.axis_rules('vocab')
+        scale_init=maybe_with_partitioning(
+            nnx.initializers.ones_init(), config.axis_rules('embed'), config
         ),
         rngs=rngs,
+    )
+    if not config.logits_via_embedding:
+      self.logitdense = nnx.Linear(
+          in_features=config.emb_dim,
+          out_features=config.output_vocab_size,
+          dtype=config.dtype,
+          kernel_init=maybe_with_partitioning(
+              config.kernel_init, config.axis_rules('embed', 'vocab'), config
+          ),
+          bias_init=maybe_with_partitioning(
+              config.bias_init, config.axis_rules('vocab'), config
+          ),
+          rngs=rngs,
       )
     else:
       self.logitdense = None

examples/lm1b_nnx/models_test.py

@@ -33,13 +33,8 @@
 
 jax.config.update('jax_disable_most_optimizations', True)
 
-# add project_root to import lm1b Linen model
-# "/path/to/flax/examples/lm1b_nnx/models_test.py" -> "/path/to/flax"
-project_root = str(Path(__file__).absolute().parents[2])
-sys.path.append(project_root)
-from examples.lm1b.models import TransformerLM as TransformerLinen  # type: ignore[import-error]
-
-sys.path.pop()
+# Import lm1b Linen model for compatibility testing
+from flax.examples.lm1b.models import TransformerLM as TransformerLinen
 
 
 @dataclasses.dataclass(unsafe_hash=True)
@@ -51,23 +46,23 @@ class CompatTransformerConfig(TransformerConfig):
 def get_transformer_config(**kwargs):
   base_config = default.get_config()
   config = CompatTransformerConfig(
-    vocab_size=base_config.vocab_size,
-    output_vocab_size=base_config.vocab_size,
-    logits_via_embedding=base_config.logits_via_embedding,
-    dtype=jnp.bfloat16 if base_config.use_bfloat16 else jnp.float32,
-    emb_dim=base_config.emb_dim,
-    num_heads=base_config.num_heads,
-    num_layers=base_config.num_layers,
-    qkv_dim=base_config.qkv_dim,
-    mlp_dim=base_config.mlp_dim,
-    max_len=max(
-      base_config.max_target_length, base_config.max_eval_target_length
-    ),
-    dropout_rate=base_config.dropout_rate,
-    attention_dropout_rate=base_config.attention_dropout_rate,
-    kernel_init=nnx.initializers.xavier_uniform(),
-    bias_init=nnx.initializers.normal(stddev=1e-6),
-    **kwargs,
+      vocab_size=base_config.vocab_size,
+      output_vocab_size=base_config.vocab_size,
+      logits_via_embedding=base_config.logits_via_embedding,
+      dtype=jnp.bfloat16 if base_config.use_bfloat16 else jnp.float32,
+      emb_dim=base_config.emb_dim,
+      num_heads=base_config.num_heads,
+      num_layers=base_config.num_layers,
+      qkv_dim=base_config.qkv_dim,
+      mlp_dim=base_config.mlp_dim,
+      max_len=max(
+          base_config.max_target_length, base_config.max_eval_target_length
+      ),
+      dropout_rate=base_config.dropout_rate,
+      attention_dropout_rate=base_config.attention_dropout_rate,
+      kernel_init=nnx.initializers.xavier_uniform(),
+      bias_init=nnx.initializers.normal(stddev=1e-6),
+      **kwargs,
   )
   return base_config, config
 
@@ -93,9 +88,10 @@ def copy_var(nnx_name: str, linen_name: str):
         == flat_params_linen[linen_name].value.shape
       )
       flat_params_nnx[nnx_path].value = flat_params_linen[linen_name].value
-      assert flat_params_nnx[nnx_path].sharding == apply_rules(
-        flat_params_linen[linen_name].names
-      )
+      if not all(rule is None for rule in rules.values()):
+        assert flat_params_nnx[nnx_path].sharding == apply_rules(
+            flat_params_linen[linen_name].names
+        )
 
     copy_var('decoder/output_embed/embedding', 'decoder/Embed_0/embedding')
     copy_var(
@@ -196,14 +192,8 @@ def copy_var(nnx_name: str, linen_name: str):
 
   def test_forward_eval(self):
     _, config = get_transformer_config(
-      axis_rules=default.MeshRules(
-        embed='model',
-        mlp='data',
-        kv=None,
-        vocab=None,
-      ),
-      deterministic=True,
-      decode=False,
+        deterministic=True,
+        decode=False,
     )
     # Set dropout rates to avoid create dropout states
     config.dropout_rate = 0.0
@@ -233,14 +223,8 @@ def test_forward_decode(self):
     batch_size = 2
 
     _, config = get_transformer_config(
-      axis_rules=default.MeshRules(
-        embed='model',
-        mlp='data',
-        kv=None,
-        vocab=None,
-      ),
-      deterministic=True,
-      decode=True,
+        deterministic=True,
+        decode=True,
     )
     # Set dropout rates to avoid create dropout states
     config.dropout_rate = 0.0
@@ -293,14 +277,8 @@ def test_forward_decode(self):
 
   def test_forward_eval_set_mode(self):
     _, config = get_transformer_config(
-      axis_rules=default.MeshRules(
-        embed='model',
-        mlp='data',
-        kv=None,
-        vocab=None,
-      ),
-      deterministic=True,
-      decode=False,
+        deterministic=True,
+        decode=False,
     )
     # Set dropout rates to avoid create dropout states
     config.dropout_rate = 0.0
@@ -330,14 +308,8 @@ def test_forward_decode_set_mode(self):
     batch_size = 2
 
     _, config = get_transformer_config(
-      axis_rules=default.MeshRules(
-        embed='model',
-        mlp='data',
-        kv=None,
-        vocab=None,
-      ),
-      deterministic=True,
-      decode=True,
+        deterministic=True,
+        decode=True,
     )
     # Set dropout rates to avoid create dropout states
     config.dropout_rate = 0.0