feat: rewrite rngs reading logic and make get_abstract_state compatib…

src/MaxText/layers/nnx_wrappers.py

@@ -241,7 +241,16 @@ def __call__(
     if rngs is None:
       rngs = self.to_nnx__rngs
     if isinstance(rngs, nnx.Rngs):
-      _rngs = {name: stream() for name, stream in rngs.items()}
+      # Extract RNG keys without calling stream() to avoid TraceContextError
+      # nnx.Rngs.items() yields (name, RngStream) pairs
+      # Each RngStream has a .key.value attribute containing the JAX PRNGKey
+      _rngs = {}
+      for name, stream in rngs.items():
+        if hasattr(stream, "key") and hasattr(stream.key, "value"):
+          _rngs[name] = stream.key.value
+    elif isinstance(rngs, dict):
+      # Allow passing RNG dict directly
+      _rngs = rngs
     elif isinstance(rngs, jax.Array):
       _rngs = {"params": rngs}
     else:

src/MaxText/maxtext_utils.py

@@ -19,6 +19,7 @@
 import pickle
 
 from flax import linen as nn
+from flax import nnx
 from flax.linen import partitioning as nn_partitioning
 from flax.training import train_state
 
@@ -294,9 +295,7 @@ def calculate_llama4_attention_tflops(config):
   num_chunked_layers = num_layers - num_global_layers
 
   # FLOPs for a single global attention layer (full attention, non-causal)
-  global_attention_flops_per_layer = (
-      4 * config.per_device_batch_size * seq_len**2 * config.num_query_heads * config.head_dim
-  )
+  global_attention_flops_per_layer = 4 * config.per_device_batch_size * seq_len**2 * config.num_query_heads * config.head_dim
 
   # FLOPs for a single chunked attention layer (non-causal)
   chunked_attention_flops_per_layer = _calculate_chunked_attention_flops_per_layer(config, seq_len, chunk_size)
@@ -323,9 +322,7 @@ def calculate_mla_tflops_per_device(config):
   else:
     # calculate query down and up flops
     q_flops = (
-        2
-        * batch_len
-        * (config.emb_dim * config.q_lora_rank + config.q_lora_rank * config.num_query_heads * qk_head_dim_sum)
+        2 * batch_len * (config.emb_dim * config.q_lora_rank + config.q_lora_rank * config.num_query_heads * qk_head_dim_sum)
     )
   # calculate mla kv projection with down and up flops
   kv_flops = (
@@ -338,9 +335,7 @@ def calculate_mla_tflops_per_device(config):
   )
   qkv_flops = q_flops + kv_flops
 
-  attention_flops = (
-      2 * batch_len * config.max_target_length * config.num_query_heads * (qk_head_dim_sum + config.v_head_dim)
-  )
+  attention_flops = 2 * batch_len * config.max_target_length * config.num_query_heads * (qk_head_dim_sum + config.v_head_dim)
   projection_flops = 2 * batch_len * config.emb_dim * config.num_query_heads * config.v_head_dim
   return qkv_flops, attention_flops, projection_flops
 
@@ -721,33 +716,34 @@ def init_decode_state(apply_fn, params) -> train_state.TrainState:
 
 
 def init_training_state(apply_fn, params, tx):
-  """Init train state with null opt state for decode."""
+  """Init train state for training."""
   state = train_state.TrainState.create(apply_fn=apply_fn, params=params, tx=tx)
   return state
 
 
 def init_initial_state(model, tx, config, is_training, key):
-  """
-  We pass in "static" objects like model, tx, config as JAX compares them by
-  object hash, and instantiating them inside causes pjit top-level annotations
-  to fail to match as pytree prefixes if we re-instantiate.
+  """Initialize training or decode state from an NNX model.
+
+  Args:
+    model: NNX model (already initialized)
+    tx: Optax optimizer transformation
+    config: Configuration object
+    is_training: True for training, False for decode
+    key: PRNG key (unused, kept for API compatibility)
+
+  Returns:
+    TrainState with model parameters and optimizer state.
 
-  Args: model, tx, config, is_training, key
+  Note:
+    Extracts only trainable parameters from the NNX model. The model structure
+    and RNG state are discarded as they're managed by the model object itself.
   """
-  input_shape = (config.micro_batch_size_to_train_on, config.max_target_length)
-  image_shape = multimodal_utils.get_dummy_image_shape_for_init(
-      config.model_name, batch_size=config.micro_batch_size_to_train_on
-  )
-  model_vars = model.init(
-      {"params": key, "dropout": key, "aqt": key},
-      np.ones(input_shape, dtype=jnp.int32),
-      np.ones(input_shape, dtype=jnp.int32),
-      encoder_images=np.ones(image_shape, dtype=jnp.int32) if config.use_multimodal else None,
-      # nnx_method="no_op",
-  )
+  # Extract only trainable parameters; discard structure and RNG state
+  _, params, _ = nnx.split(model, nnx.Param, ...)
+
   if is_training:
-    return init_training_state(model.apply, model_vars, tx)
-  return init_decode_state(model.apply, model_vars)
+    return init_training_state(None, params, tx)
+  return init_decode_state(None, params)
 
 
 def setup_decode_state(model, config, rng, mesh, checkpoint_manager):
@@ -887,8 +883,9 @@ def get_abstract_state(model, tx, config, rng, mesh, is_training=True):
   with nn_partitioning.axis_rules(config.logical_axis_rules):
     abstract_state = jax.eval_shape(init_state_partial)
 
-  state_logical_annotations = nn.get_partition_spec(abstract_state)
+  state_logical_annotations = nnx.get_partition_spec(abstract_state)
 
+  # Convert logical axis names to physical mesh axes (framework-agnostic utility, no NNX equivalent)
   state_mesh_shardings = nn.logical_to_mesh_sharding(state_logical_annotations, mesh, config.logical_axis_rules)
   if is_training and config.shard_optimizer_over_data:
     # Add data to sharding for optimizer state

src/MaxText/model_creation_utils.py

@@ -90,18 +90,20 @@ def from_config(
   return model
 
 
-def get_transformer_model(config, mesh, quant, rngs: nnx.Rngs | None = None, model_mode: str = MODEL_MODE_TRAIN)->nn.Module:
+def get_transformer_model(
+    config, mesh, quant, rngs: nnx.Rngs | None = None, model_mode: str = MODEL_MODE_TRAIN
+) -> nn.Module:
   """Returns the transformer model based on the configuration."""
   # TODO: use nnx model instead of flax linen model
 
   if config.model_fsdp_ag_once:
     if rngs is not None:
       return models.ZeroOneTransformer(config, mesh, quant=quant, model_mode=model_mode, rngs=rngs)
-    return models.zero_one_transformer_as_linen(config, mesh, quant=quant, model_mode=model_mode,rngs=rngs)
+    return models.zero_one_transformer_as_linen(config, mesh, quant=quant, model_mode=model_mode, rngs=rngs)
   else:
     if rngs is not None:
       return models.Transformer(config, mesh, quant=quant, model_mode=model_mode, rngs=rngs)
-  return models.transformer_as_linen(config, mesh, quant=quant,model_mode=model_mode,rngs=rngs)
+  return models.transformer_as_linen(config, mesh, quant=quant, model_mode=model_mode, rngs=rngs)
 
 
 def create_model(config, mesh, model_mode: str = MODEL_MODE_TRAIN, rngs: nnx.Rngs | None = None):