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..17f2ab6591 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,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 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):