apple envy moe implementation

[entrpn]

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:

1. logits_via_embedding: true: Shares the token embedding parameters with the output logits projection.
2. 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]

Caution

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

[codecov]

Codecov Report

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

Files with missing lines	Patch %	Lines
src/maxtext/models/envy.py	75.21%	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!

[JamesDeng42]

LGTM

[hengtaoguo]

Could you update the PR description? Thanks!

[RissyRan]

Is this a new PR? Could you link to the old one for future reference for review comments?