apple envy moe implementation by entrpn · Pull Request #4645 · AI-Hypercomputer/maxtext

entrpn · 2026-07-28T16:32:03Z

Description

Implements Apple Envy MoE model as described here

Tests

Envy MoE Parameter Equivalence Report

This report compares the parameter structures of Apple's Envy Switch-Base model between AXLearn and MaxText to verify architectural correctness.

1. Parameter Shape Mapping

Here is the direct mapping of parameter tensors for the 12-layer Switch-Base model (6 repeat blocks of Alternating Dense and Sparse MoE layers):

Component / Layer	AXLearn Parameter (Shape)	MaxText Parameter (Shape)	Equivalence Note
Token Embedding / Logits	`decoder/emb/token_emb/weight` `[32768, 1536]`	`token_embedder/embedding` `(32768, 1536)`	Exactly Shared & Identical (`logits_via_embedding: true`)
Final Decoder Norm	`decoder/output_norm/scale` `[1536]`	`decoder_norm/scale` `(1536,)`	Exactly Identical
Self-Attention Norm	`self_attention/norm/scale` `(6, 1536)`	`pre_self_attention_layer_norm/scale` `(1536, 6)`	Exactly Identical (Transposed layers axis)
Self-Attention QKV Proj	`self_attention/attention/i_proj/i_proj/qkv_proj/weight` `(6, 1536, 36, 128)`	Q, K, V are stored as 3 separate parameters: - `query/kernel` `(1536, 6, 12, 128)` - `key/kernel` `(1536, 6, 12, 128)` - `value/kernel` `(1536, 6, 12, 128)`	Exactly Identical. AXLearn groups QKV along the projection axis: $12 + 12 + 12 = 36$ heads total.
Self-Attention QK Norm	`attention/scale_query/norm/scale` `(6, 128)` `attention/scale_key/norm/scale` `(6, 128)`	`query_norm/scale` `(128, 6)` `key_norm/scale` `(128, 6)`	Exactly Identical (Transposed layers axis)
Self-Attention Output Proj	`self_attention/attention/o_proj/weight` `(6, 1536, 12, 128)`	`out/kernel` `(12, 6, 128, 1536)`	Exactly Identical
Dense FFN Norm	`feed_forward/prenorm/scale` `(6, 1536)`	`post_self_attention_layer_norm/scale` `(1536, 6)`	Exactly Identical (Transposed layers axis)
Dense FFN Post-Norm	`feed_forward/postnorm/scale` `(6, 1536)`	`post_ffw_norm/scale` `(1536, 6)`	Exactly Identical (`use_post_ffw_norm: true`)
Dense FFN Projections	`feed_forward/linear1_0/weight` `(6, 1536, 6144)` `feed_forward/linear1_1/weight` `(6, 1536, 6144)` `feed_forward/linear2/weight` `(6, 6144, 1536)`	`wi_0/kernel` `(1536, 6, 6144)` `wi_1/kernel` `(1536, 6, 6144)` `wo/kernel` `(6144, 6, 1536)`	Exactly Identical
Sparse MoE Norm	`feed_forward/prenorm/scale` `(6, 1536)`	`post_self_attention_layer_norm/scale` `(1536, 6)`	Exactly Identical (Transposed layers axis)
Sparse MoE Post-Norm	`feed_forward/postnorm/scale` `(6, 1536)`	`post_ffw_norm/scale` `(1536, 6)`	Exactly Identical (`use_post_ffw_norm: true`)
Sparse MoE Gate	`feed_forward/gate_weight` `(6, 1536, 128)`	`MoeBlock_0/gate/kernel` `(1536, 6, 128)`	Exactly Identical
Sparse MoE Experts Proj	`feed_forward/wi_0_weight` `(6, 128, 1536, 6144)` `feed_forward/wi_1_weight` `(6, 128, 1536, 6144)` `feed_forward/wo_weight` `(6, 128, 6144, 1536)`	`MoeBlock_0/wi_0` `(128, 6, 1536, 6144)` `MoeBlock_0/wi_1` `(128, 6, 1536, 6144)` `MoeBlock_0/wo` `(128, 6, 6144, 1536)`	Exactly Identical

2. Parameter Counts Analysis

AXLearn Parameter Count: 22,077,958,656 (22.07B)
MaxText Parameter Count: 22,077,958,656 (22.07B)
Difference: 0 parameters (100% exact match)

Key Configurations Enabled:

logits_via_embedding: true: Shares the token embedding parameters with the output logits projection.
use_post_ffw_norm: true: Activates the post-FFN normalization layer (post_ffw_norm in MaxText, mapping to postnorm in AXLearn's hybridnorm FFN structure).

Ran test:

python -m pytest tests/unit/envy_test.py
=========================================================== test session starts ===========================================================
platform linux -- Python 3.12.11, pytest-8.4.2, pluggy-1.6.0
rootdir: /mnt/disks/external_disk/maxtext
configfile: pytest.ini
plugins: anyio-4.13.0, typeguard-2.13.3, jaxtyping-0.3.9, xdist-3.8.0, hypothesis-6.151.9, mock-3.15.1
collected 1 item                                                                                                                          

tests/unit/envy_test.py .                                                                                                           [100%]

============================================================ warnings summary =============================================================
maxtext_venv/lib/python3.12/site-packages/flax/nnx/variablelib.py:324
  /mnt/disks/external_disk/maxtext/maxtext_venv/lib/python3.12/site-packages/flax/nnx/variablelib.py:324: DeprecationWarning: jax.core.Effect is deprecated. Use jax.extend.core.Effect.
    class VariableEffect(jax.core.Effect): ...

tests/unit/envy_test.py::EnvyTest::test_model_initialization_and_forward
  /mnt/disks/external_disk/maxtext/maxtext_venv/lib/python3.12/site-packages/flax/core/tracers.py:30: DeprecationWarning: jax.core.get_opaque_trace_state is deprecated. Use jax.extend.core.get_opaque_trace_state.
    return jax.core.get_opaque_trace_state(convention="flax")

tests/unit/envy_test.py: 90 warnings
  /mnt/disks/external_disk/maxtext/src/maxtext/layers/nnx_scan.py:84: DeprecationWarning: '.value' access is now deprecated. For Variable[Array] instances use:
  
    variable[...]
  
  For other Variable types use:
  
    variable.get_value()
  
    is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"),

tests/unit/envy_test.py: 90 warnings
  /mnt/disks/external_disk/maxtext/src/maxtext/layers/nnx_scan.py:56: DeprecationWarning: '.value' access is now deprecated. For Variable[Array] instances use:
  
    variable[...]
  
  For other Variable types use:
  
    variable.get_value()
  
    if hasattr(leaf, "replace") and hasattr(leaf, "value"):

tests/unit/envy_test.py: 90 warnings
  /mnt/disks/external_disk/maxtext/src/maxtext/layers/nnx_scan.py:65: DeprecationWarning: 'sharding_names' is deprecated, use 'out_sharding' instead.
    value = getattr(leaf, key, None)

tests/unit/envy_test.py::EnvyTest::test_model_initialization_and_forward
tests/unit/envy_test.py::EnvyTest::test_model_initialization_and_forward
  /mnt/disks/external_disk/maxtext/maxtext_venv/lib/python3.12/site-packages/qwix/_src/qconfig.py:178: DeprecationWarning: The 'm.iter_modules()' method is deprecated; use the 'nnx.iter_modules(m)' function instead.
    for _, node in model.iter_modules():

tests/unit/envy_test.py::EnvyTest::test_model_initialization_and_forward
tests/unit/envy_test.py::EnvyTest::test_model_initialization_and_forward
  /mnt/disks/external_disk/maxtext/maxtext_venv/lib/python3.12/site-packages/flax/nnx/statelib.py:378: DeprecationWarning: `flax.nnx.State` will be deprecated and be replaced by the built-in Python dict. Please use the equivalent `nnx.split_state` instead.
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
==================================================== 1 passed, 276 warnings in 14.50s =====================================================
sys:1: DeprecationWarning: builtin type swigvarlink has no __module__ attribute

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

gemini-code-assist · 2026-07-28T16:32:06Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

codecov · 2026-07-28T16:42:32Z

Codecov Report

❌ Patch coverage is 70.40000% with 37 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/models/envy.py	74.56%	18 Missing and 11 partials ⚠️
src/maxtext/layers/decoders.py	0.00%	4 Missing and 2 partials ⚠️
src/maxtext/layers/nnx_decoders.py	50.00%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

github-actions · 2026-07-28T17:17:18Z

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions · 2026-07-28T17:21:24Z

🤖 I'm sorry @RissyRan, but I was unable to process your request. Please see the logs for more details.

JamesDeng42

left some comments, otherwise, it looks good to me.

github-actions · 2026-07-29T19:28:44Z

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions

## 📋 Review Summary

This Pull Request successfully implements the Apple Envy MoE model (specifically supporting the Switch-Base, Switch-Large, and Switch-XXL models) in MaxText. The alternating inhomogeneous dense and sparse layers, RMSNorm configurations, and parameters are precisely matched (reaching an exact 100% parameter count match of 22.07B parameters for Switch-Base). The changes are exceptionally well-structured, follow existing design conventions, and are verified by unit tests.

🔍 General Feedback

Rigorous Validation: The detailed parameter equivalence report comparing AXLearn and MaxText structure and shape mappings is exemplary and ensures high correctness.
Clean Architectural Integration: Adding Envy as an inhomogeneous decoder block option integrates perfectly across Linen (decoders.py) and NNX (nnx_decoders.py) layer builders.
Comprehensive Test Coverage: The unit test tests/unit/envy_test.py covers the compilation, initialization, and training forward pass.

github-actions · 2026-07-29T19:34:21Z

+    cfg = self.config
+
+    inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed"))
+    inputs = checkpoint_name(inputs, "decoder_layer_input")


🟡 Performance/Correctness: In EnvyScannableBlock.__call__, the logical constraint on inputs is hardcoded to "activation_norm_length". In MODEL_MODE_PREFILL mode, the sequence length dimension is named "prefill_activation_norm_length". Leaving this hardcoded can cause a sharding constraint conflict with EnvyDecoderLayer (which correctly selects the axis name based on model_mode), potentially leading to unnecessary compilation layout transitions on accelerators.

To align with EnvyDecoderLayer and ensure optimal sharding behavior across both prefill and decoding modes, dynamically select the axis names based on model_mode.

Suggested change

inputs = checkpoint_name(inputs, "decoder_layer_input")

activation_axis_names = (

("activation_batch", "prefill_activation_norm_length", "activation_embed")

if model_mode == MODEL_MODE_PREFILL

else ("activation_batch", "activation_norm_length", "activation_embed")

)

inputs = nn.with_logical_constraint(inputs, activation_axis_names)

github-actions · 2026-07-29T19:34:32Z

+  @property
+  def moe_block(self):
+    return self.MoeBlock_0


🟢 Maintainability/Robustness: In EnvyDecoderLayer, the moe_block property directly returns self.MoeBlock_0. However, when is_moe_layer is False (for dense layers), self.MoeBlock_0 is not initialized, so accessing self.moe_block will raise an AttributeError.

Although the current layer implementation only calls moe_block when is_moe_layer is True, other tooling/utilities (such as model visualization, parameter counting, weight inspection, or serialization) might inspect properties unconditionally. Making this safe with getattr prevents potential runtime crashes.

Suggested change

@property

def moe_block(self):

return self.MoeBlock_0

@property

def moe_block(self):

return getattr(self, "MoeBlock_0", None)

github-actions · 2026-07-29T19:34:41Z

                "nope_layer_interval": self.config.nope_layer_interval,
                "interleave_moe_layer_step": self.config.interleave_moe_layer_step,
            }
+          if cfg.decoder_block == DecoderBlockType.ENVY:


🟢 Style: Since cfg.decoder_block can only match one block type, use elif instead of a separate if statement to align with the rest of the conditional checks.

Suggested change

if cfg.decoder_block == DecoderBlockType.ENVY:

elif cfg.decoder_block == DecoderBlockType.ENVY:

github-actions · 2026-07-29T19:34:48Z

                  "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:


🟢 Style: Use elif here as well to make the conditional block evaluation more efficient and clean.

Suggested change

if cfg.decoder_block == DecoderBlockType.ENVY:

elif cfg.decoder_block == DecoderBlockType.ENVY:

RissyRan

LGTM at high level. Will you test logits alignment in next PR for the model?

RissyRan · 2026-07-29T19:31:08Z

+rope_max_timescale: 500000
+num_experts: 128
+num_experts_per_tok: 1
+capacity_factor: 2.0


For capacity_factor is a implementation details (dropping/dropless), not related to model config itself.

Recommend to remove from here. If you'd like to assert this always happen, you could add an assertion in types.py file and ensure envy code block is called with capacity_factor=2

Similar comments for other envy configs.

RissyRan · 2026-07-29T19:31:32Z

+load_balance_loss_weight: 0.01
+interleave_moe_layer_step: 2
+inhomogeneous_layer_cycle_interval: 2
+max_target_length: 8192


Similar comment with capacity factor.

RissyRan · 2026-07-29T19:37:06Z

+import numpy as np
+
+
+class EnvyTest(unittest.TestCase):


Does this test functionality of this model? If so, you could write a compile test like this:

maxtext/tests/unit/train_compile_test.py

Line 866 in 025f83f

def test_deepseek4(self, scan_layers):

entrpn requested a review from parambole as a code owner July 28, 2026 16:32

RissyRan added the gemini-review label Jul 28, 2026

JamesDeng42 reviewed Jul 28, 2026

View reviewed changes

Comment thread src/maxtext/configs/models/envy-switch-base.yml

Comment thread src/maxtext/configs/models/envy-switch-large.yml

JamesDeng42 approved these changes Jul 28, 2026

View reviewed changes

apple envy moe implementation

9194371

entrpn force-pushed the jfacevedo_apple_envy_2 branch from d4cb0d7 to 9194371 Compare July 28, 2026 18:41

RissyRan added gemini-review and removed gemini-review labels Jul 29, 2026

github-actions Bot reviewed Jul 29, 2026

View reviewed changes

RissyRan reviewed Jul 29, 2026

View reviewed changes

-    inputs = checkpoint_name(inputs, "decoder_layer_input")
+    activation_axis_names = (
+        ("activation_batch", "prefill_activation_norm_length", "activation_embed")
+        if model_mode == MODEL_MODE_PREFILL
+        else ("activation_batch", "activation_norm_length", "activation_embed")
+    )
+    inputs = nn.with_logical_constraint(inputs, activation_axis_names)

	if cfg.decoder_block == DecoderBlockType.ENVY:
	elif cfg.decoder_block == DecoderBlockType.ENVY:

Uh oh!

Conversation

entrpn commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Envy MoE Parameter Equivalence Report

1. Parameter Shape Mapping

2. Parameter Counts Analysis

Key Configurations Enabled:

Checklist

Uh oh!

gemini-code-assist Bot commented Jul 28, 2026

Uh oh!

codecov Bot commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions Bot commented Jul 28, 2026

Uh oh!

github-actions Bot commented Jul 28, 2026

Uh oh!

Uh oh!

Uh oh!

JamesDeng42 left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jul 29, 2026

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🔍 General Feedback

Uh oh!

github-actions Bot Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan left a comment

Choose a reason for hiding this comment

Uh oh!

RissyRan Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

entrpn commented Jul 28, 2026 •

edited

Loading

codecov Bot commented Jul 28, 2026 •

edited

Loading