From abf2b7e79525379b254e1c3d2defee23bc62748d Mon Sep 17 00:00:00 2001 From: Juan Acevedo Date: Mon, 20 Jul 2026 19:36:43 -0700 Subject: [PATCH] apple envy moe implementation --- src/maxtext/common/common_types.py | 2 + .../configs/models/envy-switch-base.yml | 38 +++ .../configs/models/envy-switch-large.yml | 38 +++ .../configs/models/envy-switch-xxl.yml | 38 +++ src/maxtext/configs/models/envy-test.yml | 38 +++ src/maxtext/configs/types.py | 10 + src/maxtext/layers/decoders.py | 15 + src/maxtext/layers/nnx_decoders.py | 13 + src/maxtext/models/envy.py | 320 ++++++++++++++++++ tests/unit/envy_test.py | 86 +++++ 10 files changed, 598 insertions(+) create mode 100644 src/maxtext/configs/models/envy-switch-base.yml create mode 100644 src/maxtext/configs/models/envy-switch-large.yml create mode 100644 src/maxtext/configs/models/envy-switch-xxl.yml create mode 100644 src/maxtext/configs/models/envy-test.yml create mode 100644 src/maxtext/models/envy.py create mode 100644 tests/unit/envy_test.py diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index 9dd4aff3dc..6cff94c49e 100644 --- a/src/maxtext/common/common_types.py +++ b/src/maxtext/common/common_types.py @@ -114,6 +114,8 @@ class DecoderBlockType(enum.Enum): LLAMA4 = "llama4" OLMO3 = "olmo3" DEEPSEEK4 = "deepseek4" + ENVY = "envy" + class VisionEncoderBlockType(enum.Enum): diff --git a/src/maxtext/configs/models/envy-switch-base.yml b/src/maxtext/configs/models/envy-switch-base.yml new file mode 100644 index 0000000000..1ea2c23836 --- /dev/null +++ b/src/maxtext/configs/models/envy-switch-base.yml @@ -0,0 +1,38 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Model configuration for Envy Switch-Base (30B) + +decoder_block: "envy" +mlp_activations: ["silu", "linear"] +enable_dropout: false +logits_via_embedding: false +tokenizer_type: "huggingface" + +base_emb_dim: 1536 +base_num_decoder_layers: 12 +base_num_query_heads: 12 +base_num_kv_heads: 12 +base_mlp_dim: 6144 +base_moe_mlp_dim: 6144 +vocab_size: 32768 +normalization_layer_epsilon: 1e-05 +rope_max_timescale: 500000 +num_experts: 128 +num_experts_per_tok: 1 +capacity_factor: 2.0 +use_qk_norm: true +interleave_moe_layer_step: 2 +inhomogeneous_layer_cycle_interval: 2 +max_target_length: 8192 diff --git a/src/maxtext/configs/models/envy-switch-large.yml b/src/maxtext/configs/models/envy-switch-large.yml new file mode 100644 index 0000000000..f84983c23c --- /dev/null +++ b/src/maxtext/configs/models/envy-switch-large.yml @@ -0,0 +1,38 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Model configuration for Envy Switch-Large (104B) + +decoder_block: "envy" +mlp_activations: ["silu", "linear"] +enable_dropout: false +logits_via_embedding: false +tokenizer_type: "huggingface" + +base_emb_dim: 2048 +base_num_decoder_layers: 24 +base_num_query_heads: 16 +base_num_kv_heads: 16 +base_mlp_dim: 8192 +base_moe_mlp_dim: 8192 +vocab_size: 32768 +normalization_layer_epsilon: 1e-05 +rope_max_timescale: 500000 +num_experts: 128 +num_experts_per_tok: 1 +capacity_factor: 2.0 +use_qk_norm: true +interleave_moe_layer_step: 2 +inhomogeneous_layer_cycle_interval: 2 +max_target_length: 8192 diff --git a/src/maxtext/configs/models/envy-switch-xxl.yml b/src/maxtext/configs/models/envy-switch-xxl.yml new file mode 100644 index 0000000000..903175a4b9 --- /dev/null +++ b/src/maxtext/configs/models/envy-switch-xxl.yml @@ -0,0 +1,38 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Model configuration for Envy Switch-XXL (520B) + +decoder_block: "envy" +mlp_activations: ["silu", "linear"] +enable_dropout: false +logits_via_embedding: false +tokenizer_type: "huggingface" + +base_emb_dim: 8192 +base_num_decoder_layers: 24 +base_num_query_heads: 64 +base_num_kv_heads: 8 +base_mlp_dim: 20480 +base_moe_mlp_dim: 20480 +vocab_size: 32768 +normalization_layer_epsilon: 1e-05 +rope_max_timescale: 500000 +num_experts: 64 +num_experts_per_tok: 1 +capacity_factor: 2.0 +use_qk_norm: true +interleave_moe_layer_step: 2 +inhomogeneous_layer_cycle_interval: 2 +max_target_length: 8192 diff --git a/src/maxtext/configs/models/envy-test.yml b/src/maxtext/configs/models/envy-test.yml new file mode 100644 index 0000000000..ca78a663b3 --- /dev/null +++ b/src/maxtext/configs/models/envy-test.yml @@ -0,0 +1,38 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Model configuration for Envy-test + +decoder_block: "envy" +mlp_activations: ["silu", "linear"] +enable_dropout: false +logits_via_embedding: false +tokenizer_type: "huggingface" + +base_emb_dim: 8 +base_num_decoder_layers: 4 +base_num_query_heads: 4 +base_num_kv_heads: 2 +base_mlp_dim: 16 +base_moe_mlp_dim: 16 +vocab_size: 32 +normalization_layer_epsilon: 1e-05 +rope_max_timescale: 500000 +num_experts: 8 +num_experts_per_tok: 1 +capacity_factor: 2.0 +use_qk_norm: true +interleave_moe_layer_step: 2 +inhomogeneous_layer_cycle_interval: 2 +max_target_length: 64 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 22a2a5ff0b..ca692ea995 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -280,6 +280,10 @@ class ProfilerType(str, Enum): "olmo3-7b", "olmo3-7b-pt", "olmo3-32b", + "envy-test", + "envy-switch-base", + "envy-switch-large", + "envy-switch-xxl", ] @@ -3470,6 +3474,12 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de f"The number of decoder layers ({self.base_num_decoder_layers}) must be divisible by interleave moe layer step " f"({self.interleave_moe_layer_step})" ) + if self.decoder_block == DecoderBlockType.ENVY: + if self.base_num_decoder_layers % self.interleave_moe_layer_step != 0: + raise ValueError( + f"The number of decoder layers ({self.base_num_decoder_layers}) must be divisible by interleave moe layer step " + f"({self.interleave_moe_layer_step})" + ) if self.decoder_block in ( DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 6ce5ffc114..d88ab56aa8 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -44,6 +44,7 @@ deepseek4, deepseek_batchsplit, deepseek_batchsplit_fp8, + envy, gemma, gemma2, gemma3, @@ -493,6 +494,9 @@ def get_decoder_layers(self): return [llama4.Llama4ScannableBlockToLinen] if self.config.scan_layers else [llama4.Llama4DecoderLayerToLinen] case DecoderBlockType.OLMO3: return [olmo3.Olmo3ScannableBlockToLinen] if self.config.scan_layers else [olmo3.Olmo3DecoderLayerToLinen] + case DecoderBlockType.ENVY: + return [envy.EnvyScannableBlockToLinen] if self.config.scan_layers else [envy.EnvyDecoderLayerToLinen] + case _: # Default case to handle any unknown decoder block types. @@ -528,6 +532,7 @@ def get_scannable(normal_cls, scannable_cls): DecoderBlockType.DEEPSEEK: [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer], DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), + DecoderBlockType.ENVY: get_scannable(envy.EnvyDecoderLayer, envy.EnvyScannableBlock), } if cfg.decoder_block not in layer_map: @@ -1078,6 +1083,11 @@ def __call__( "nope_layer_interval": self.config.nope_layer_interval, "interleave_moe_layer_step": self.config.interleave_moe_layer_step, } + if cfg.decoder_block == DecoderBlockType.ENVY: + layer_kwargs = { + "interleave_moe_layer_step": self.config.interleave_moe_layer_step, + } + # Update broadcast_args and in_axes_tuple for vLLM RPA in_axes_tuple = (nn.broadcast,) * len(broadcast_args) current_broadcast_args = list(broadcast_args) @@ -1203,6 +1213,11 @@ def __call__( "is_nope_layer": llama4.determine_is_nope_layer(lyr, self.config.nope_layer_interval), "is_moe_layer": llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step), } + if cfg.decoder_block == DecoderBlockType.ENVY: + layer_kwargs = { + "is_moe_layer": (lyr + 1) % self.config.interleave_moe_layer_step == 0, + } + if cfg.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, DecoderBlockType.DEEPSEEK4): layer_kwargs = {"layer_idx": lyr} if cfg.decoder_block == DecoderBlockType.DEEPSEEK4: diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index de33114911..3571081c8e 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -48,6 +48,7 @@ deepseek4, deepseek_batchsplit, deepseek_batchsplit_fp8, + envy, gemma, gemma2, gemma3, @@ -680,6 +681,11 @@ def _init_scanned_generic(self, decoder_block_classes, rngs): "nope_layer_interval": self.config.nope_layer_interval, "interleave_moe_layer_step": self.config.interleave_moe_layer_step, } + if config.decoder_block == DecoderBlockType.ENVY: + layer_kwargs = { + "interleave_moe_layer_step": self.config.interleave_moe_layer_step, + } + if num_layers > 0: self.layers = self._create_scanned_layers( @@ -726,6 +732,11 @@ def _init_sequential_generic(self, decoder_block_classes, rngs): "is_nope_layer": llama4.determine_is_nope_layer(lyr, self.config.nope_layer_interval), "is_moe_layer": llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step), } + elif config.decoder_block == DecoderBlockType.ENVY: + layer_kwargs = { + "is_moe_layer": (lyr + 1) % self.config.interleave_moe_layer_step == 0, + } + elif config.decoder_block in { DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, @@ -1053,6 +1064,7 @@ def get_deepseek(): DecoderBlockType.QWEN3_5: get_scannable(qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock), DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), + DecoderBlockType.ENVY: get_scannable(envy.EnvyDecoderLayer, envy.EnvyScannableBlock), } if cfg.decoder_block not in layer_map: @@ -1212,6 +1224,7 @@ def get_norm_layer(self, num_features: int, rngs: nnx.Rngs): DecoderBlockType.SIMPLE_MLP, DecoderBlockType.LLAMA4, DecoderBlockType.OLMO3, + DecoderBlockType.ENVY, }: return functools.partial( RMSNorm, diff --git a/src/maxtext/models/envy.py b/src/maxtext/models/envy.py new file mode 100644 index 0000000000..f6303ffb13 --- /dev/null +++ b/src/maxtext/models/envy.py @@ -0,0 +1,320 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Envy MoE decoder layer definition.""" +# pylint: disable=arguments-differ, disable=no-name-in-module, missing-function-docstring + +from flax import linen as nn +from flax import nnx +import jax +from jax.ad_checkpoint import checkpoint_name +import jax.numpy as jnp +from jax.sharding import Mesh +from maxtext.common.common_types import Config, MODEL_MODE_TRAIN, MODEL_MODE_PREFILL +from maxtext.layers import initializers +from maxtext.layers import nnx_wrappers +from maxtext.layers import quantizations +from maxtext.layers.attentions import Attention +from maxtext.layers.linears import Dropout +from maxtext.layers.linears import MlpBlock +from maxtext.layers import moe +from maxtext.layers.normalizations import RMSNorm +from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.utils import max_utils + + +class EnvyDecoderLayer(nnx.Module): + """Transformer decoder layer for Envy.""" + + def __init__( + self, + config: Config, + mesh: Mesh, + model_mode: str, + rngs: nnx.Rngs, + quant: None | Quant = None, + is_moe_layer: bool = False, + ): + """Initializes the Envy decoder layer. + + Args: + config: The main model configuration object. + mesh: The device mesh used for sharding parameters and activations. + model_mode: One of MODEL_MODE_TRAIN, MODEL_MODE_PREFILL, or MODEL_MODE_AUTOREGRESSIVE. + rngs: An `nnx.Rngs` object to provide random numbers. + quant: An optional configuration for quantization. Defaults to None. + is_moe_layer: If True, this layer will use a MoE block. Defaults to False as Dense. + """ + self.config = config + self.mesh = mesh + self.quant = quant + self.rngs = rngs + self.is_moe_layer = is_moe_layer + + batch_size, seq_len = max_utils.get_batch_seq_len_for_mode(config, model_mode) + dummy_inputs_shape = (batch_size, seq_len, config.emb_dim) + + self.pre_self_attention_layer_norm = RMSNorm( + num_features=config.emb_dim, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + kernel_axes=("norm",), + epsilon=config.normalization_layer_epsilon, + rngs=rngs, + ) + + query_pre_attn_scalar = config.head_dim**-0.5 + self.self_attention = Attention( + config=config, + num_query_heads=config.num_query_heads, + num_kv_heads=config.num_kv_heads, + head_dim=config.head_dim, + max_target_length=config.max_target_length, + max_prefill_predict_length=config.max_prefill_predict_length, + attention_kernel=config.attention, + inputs_q_shape=dummy_inputs_shape, + inputs_kv_shape=dummy_inputs_shape, + mesh=mesh, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + dropout_rate=config.dropout_rate, + float32_qk_product=config.float32_qk_product, + float32_logits=config.float32_logits, + quant=self.quant, + kv_quant=quantizations.configure_kv_quant(config), + prefill_cache_axis_order=tuple(map(int, config.prefill_cache_axis_order.split(","))), + ar_cache_axis_order=tuple(map(int, config.ar_cache_axis_order.split(","))), + compute_axis_order=tuple(map(int, config.compute_axis_order.split(","))), + reshape_q=config.reshape_q, + use_ragged_attention=config.use_ragged_attention, + ragged_block_size=config.ragged_block_size, + use_qk_norm=config.use_qk_norm, + query_pre_attn_scalar=query_pre_attn_scalar, + model_mode=model_mode, + rngs=rngs, + ) + + self.post_self_attention_layer_norm = RMSNorm( + num_features=config.emb_dim, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + kernel_axes=("norm",), + epsilon=config.normalization_layer_epsilon, + rngs=self.rngs, + ) + + if self.is_moe_layer: + self.MoeBlock_0 = moe.RoutedMoE( + config=config, + num_experts=config.num_experts, + num_experts_per_tok=config.num_experts_per_tok, + mesh=mesh, + kernel_init=initializers.nd_dense_init(config.dense_init_scale, "fan_in", "truncated_normal"), + kernel_axes=("embed", None), + intermediate_dim=config.base_moe_mlp_dim, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + quant=self.quant, + rngs=self.rngs, + ) + else: + self.mlp = MlpBlock( + mesh=self.mesh, + in_features=config.emb_dim, + intermediate_dim=config.mlp_dim, + activations=config.mlp_activations, + intermediate_dropout_rate=config.dropout_rate, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + config=config, + quant=self.quant, + model_mode=model_mode, + rngs=self.rngs, + ) + + self.dropout = Dropout(rate=config.dropout_rate, broadcast_dims=(-2,), rngs=self.rngs) + if model_mode == MODEL_MODE_PREFILL: + self.activation_axis_names = ("activation_batch", "prefill_activation_norm_length", "activation_embed") + else: + self.activation_axis_names = ("activation_batch", "activation_norm_length", "activation_embed") + + @property + def moe_block(self): + return self.MoeBlock_0 + + def __call__( + self, + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=None, + slot: None | int = None, + kv_cache=None, + attention_metadata=None, + ): + cfg = self.config + + # Unpack inputs if it's a tuple (e.g. from a scanned sequential block returning (hidden_states, stacked_kv_cache, layer_idx)) + is_scan_carry = False + if isinstance(inputs, tuple) and len(inputs) == 3: + hidden_states, stacked_kv_cache, layer_idx = inputs + kv_cache = stacked_kv_cache[layer_idx] + inputs = hidden_states + is_scan_carry = True + elif isinstance(inputs, tuple): + inputs = inputs[0] + inputs = nn.with_logical_constraint(inputs, self.activation_axis_names) + inputs = checkpoint_name(inputs, "decoder_layer_input") + + lnx = self.pre_self_attention_layer_norm(inputs) + lnx = nn.with_logical_constraint(lnx, self.activation_axis_names) + + # Self-attention block + attention_lnx, kv_cache = self.self_attention( + lnx, + lnx, + decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=deterministic, + model_mode=model_mode, + slot=slot, + previous_chunk=previous_chunk, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + ) + attention_lnx = nn.with_logical_constraint(attention_lnx, self.activation_axis_names) + intermediate_inputs = inputs + attention_lnx + + # Fully Connected / MLP block + hidden_states = self.post_self_attention_layer_norm(intermediate_inputs) + hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names) + + load_balance_loss = None + if self.is_moe_layer: + mlp_lnx, load_balance_loss, _ = self.moe_block(hidden_states) + else: + mlp_lnx = self.mlp(hidden_states, deterministic=deterministic) + mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names) + + layer_output = mlp_lnx + intermediate_inputs + layer_output = self.dropout(layer_output, deterministic=deterministic) + layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names) + + if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: + self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) + + if cfg.record_internal_nn_metrics: + self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output)) + self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output)) + self.sow( + nnx.Intermediate, + "activation_fraction_zero", + jnp.sum(layer_output == 0) / jnp.size(layer_output), + ) + + if is_scan_carry: + def update_cache(cache, val): + if jnp.size(val) > 0: + return cache.at[layer_idx].set(val) + return cache + + stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) + return (layer_output, stacked_kv_cache, layer_idx + 1), None + elif cfg.scan_layers: + return layer_output, None + else: + return layer_output, kv_cache + + +EnvyDecoderLayerToLinen = nnx_wrappers.to_linen_class( + EnvyDecoderLayer, + base_metadata_fn=initializers.variable_to_logically_partitioned, +) + + +class EnvyScannableBlock(nnx.Module): + """A repeatable block of inhomogeneous layers for Envy.""" + + def __init__( + self, + config: Config, + mesh: Mesh, + model_mode: str, + rngs: nnx.Rngs, + quant: None | Quant = None, + interleave_moe_layer_step: int = 2, + ): + self.config = config + self.mesh = mesh + self.model_mode = model_mode + self.quant = quant + self.rngs = rngs + self.interleave_moe_layer_step = interleave_moe_layer_step + + for layer_id in range(self.config.inhomogeneous_layer_cycle_interval): + moe_layer = (layer_id + 1) % self.interleave_moe_layer_step == 0 + layer_name = f"layers_{layer_id}" + layer = EnvyDecoderLayer( + config=self.config, + mesh=self.mesh, + model_mode=self.model_mode, + rngs=self.rngs, + quant=self.quant, + is_moe_layer=moe_layer, + ) + setattr(self, layer_name, layer) + + def __call__( + self, + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=None, + slot: None | int = None, + kv_cache=None, + attention_metadata=None, + ): + cfg = self.config + + inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + inputs = checkpoint_name(inputs, "decoder_layer_input") + y = inputs + for layer_id in range(cfg.inhomogeneous_layer_cycle_interval): + y = getattr(self, f"layers_{layer_id}")( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=previous_chunk, + slot=slot, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + ) + if cfg.scan_layers: + y = y[0] + if cfg.scan_layers: + return y, None + else: + return y + + +EnvyScannableBlockToLinen = nnx_wrappers.to_linen_class( + EnvyScannableBlock, + base_metadata_fn=initializers.variable_to_logically_partitioned, +) diff --git a/tests/unit/envy_test.py b/tests/unit/envy_test.py new file mode 100644 index 0000000000..b06d1b39e5 --- /dev/null +++ b/tests/unit/envy_test.py @@ -0,0 +1,86 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for Envy MoE.""" + +import sys +import unittest + +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +from maxtext.configs import pyconfig +from maxtext.common.common_types import MODEL_MODE_TRAIN, MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE +from maxtext.layers import quantizations +from maxtext.models import models +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path +import numpy as np + + +class EnvyTest(unittest.TestCase): + """Compile and verification tests for Envy MoE.""" + + def setUp(self): + super().setUp() + self.cfg = pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + run_name="test", + enable_checkpointing=False, + model_name="envy-test", + dtype="float32", + per_device_batch_size=1.0 / jax.device_count(), + ) + self.rng = jax.random.PRNGKey(123) + + devices_array = maxtext_utils.create_device_mesh(self.cfg, devices=[jax.devices()[0]]) + mesh = Mesh(devices_array, self.cfg.mesh_axes) + quant = quantizations.configure_quantization(self.cfg) + self.model = models.transformer_as_linen(config=self.cfg, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + self.example_batch = { + "inputs": jnp.array([[5, 12, 18, 3, 22]], dtype=jnp.int32), + "inputs_position": jnp.array([[0, 1, 2, 3, 4]], dtype=jnp.int32), + "inputs_segmentation": jnp.array([[1, 1, 1, 1, 1]], dtype=jnp.int32), + "targets": jnp.array([[12, 18, 3, 22, 1]], dtype=jnp.int32), + "targets_position": jnp.array([[0, 1, 2, 3, 4]], dtype=jnp.int32), + "targets_segmentation": jnp.array([[1, 1, 1, 1, 0]], dtype=jnp.int32), + } + + def test_model_initialization_and_forward(self): + """Verifies that the Envy model initializes and computes a training forward pass without error.""" + model_vars = self.model.init( + {"params": self.rng, "aqt": self.rng}, + self.example_batch["inputs"], + self.example_batch["inputs_position"], + enable_dropout=False, + ) + + logits, _ = self.model.apply( + model_vars, + self.example_batch["inputs"], + self.example_batch["inputs_position"], + decoder_segment_ids=self.example_batch["inputs_segmentation"], + enable_dropout=False, + rngs={"dropout": self.rng, "aqt": self.rng}, + mutable="intermediates", + ) + + self.assertEqual(logits.shape, (1, 5, self.cfg.vocab_size)) + # Make sure logits are not all zeros or NaNs + self.assertFalse(np.isnan(logits).any()) + self.assertFalse((logits == 0).all()) + + +if __name__ == "__main__": + unittest.main()