From 009e288aa72c00f1dddad2612b5f30a113aecc09 Mon Sep 17 00:00:00 2001 From: Jacky Fang Date: Fri, 7 Nov 2025 10:23:17 +0000 Subject: [PATCH 1/2] feat: rewrite rngs reading logic and make get_abstract_state compatible to nnx model --- src/MaxText/layers/nnx_wrappers.py | 11 ++++- src/MaxText/maxtext_utils.py | 71 +++++++++++++++++------------ src/MaxText/model_creation_utils.py | 8 ++-- 3 files changed, 57 insertions(+), 33 deletions(-) diff --git a/src/MaxText/layers/nnx_wrappers.py b/src/MaxText/layers/nnx_wrappers.py index fd4bc1f6b6..a135501f36 100644 --- a/src/MaxText/layers/nnx_wrappers.py +++ b/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: diff --git a/src/MaxText/maxtext_utils.py b/src/MaxText/maxtext_utils.py index 6cb8b1e10d..7a8ce9fa8f 100644 --- a/src/MaxText/maxtext_utils.py +++ b/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,7 +883,24 @@ 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) + # For NNX models, get partition specs from the model's NNX parameters + # The model has nnx.Param(..., sharding=...) annotations that we need to preserve + graphdef, model_params_state, model_other_vars = nnx.split(model, nnx.Param, ...) + param_partition_specs = nnx.get_partition_spec(model_params_state) + + # Extract PartitionSpec values from the NNX State structure + param_logical_annotations = jax.tree.map( + lambda x: x.value if isinstance(x, nnx.VariableState) else x, + param_partition_specs, + is_leaf=lambda x: isinstance(x, nnx.VariableState), + ) + + # Get partition specs for the entire state (to handle opt_state, step, etc.) + # This will return None for fields without annotations + full_state_logical_annotations = nn.get_partition_spec(abstract_state) + + # Replace just the params field with NNX partition specs + state_logical_annotations = full_state_logical_annotations.replace(params=param_logical_annotations) 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: diff --git a/src/MaxText/model_creation_utils.py b/src/MaxText/model_creation_utils.py index 632eca0088..84aee4d836 100644 --- a/src/MaxText/model_creation_utils.py +++ b/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): From db316068d7d58b27a096ab2a7fe90e9560f1c7fe Mon Sep 17 00:00:00 2001 From: Jacky Fang Date: Tue, 11 Nov 2025 08:00:40 +0000 Subject: [PATCH 2/2] feat: Update get_abstract_state logic to use nnx --- src/MaxText/maxtext_utils.py | 20 ++------------------ 1 file changed, 2 insertions(+), 18 deletions(-) diff --git a/src/MaxText/maxtext_utils.py b/src/MaxText/maxtext_utils.py index 7a8ce9fa8f..17f2ab6591 100644 --- a/src/MaxText/maxtext_utils.py +++ b/src/MaxText/maxtext_utils.py @@ -883,25 +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) - # For NNX models, get partition specs from the model's NNX parameters - # The model has nnx.Param(..., sharding=...) annotations that we need to preserve - graphdef, model_params_state, model_other_vars = nnx.split(model, nnx.Param, ...) - param_partition_specs = nnx.get_partition_spec(model_params_state) - - # Extract PartitionSpec values from the NNX State structure - param_logical_annotations = jax.tree.map( - lambda x: x.value if isinstance(x, nnx.VariableState) else x, - param_partition_specs, - is_leaf=lambda x: isinstance(x, nnx.VariableState), - ) - - # Get partition specs for the entire state (to handle opt_state, step, etc.) - # This will return None for fields without annotations - full_state_logical_annotations = nn.get_partition_spec(abstract_state) - - # Replace just the params field with NNX partition specs - state_logical_annotations = full_state_logical_annotations.replace(params=param_logical_annotations) + 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