From 40bfec83064b715a3e1432a15af157e0610bd8fa Mon Sep 17 00:00:00 2001 From: Parsa Assadi <83480666+parsley9877@users.noreply.github.com> Date: Mon, 5 Oct 2026 16:37:40 +0000 Subject: [PATCH] Add WeaverOmniTransformer full 36-layer backbone architecture and tests --- src/maxtext/common/common_types.py | 1 + .../configs/models/weaver-mini-diffuser.yml | 46 + src/maxtext/configs/models/weaver.yml | 50 + src/maxtext/configs/types.py | 2 + src/maxtext/layers/nnx_decoders.py | 2 + src/maxtext/models/weaver.py | 891 +++++++++++++++++- tests/unit/weaver_layers_test.py | 458 ++++++++- 7 files changed, 1423 insertions(+), 27 deletions(-) create mode 100644 src/maxtext/configs/models/weaver-mini-diffuser.yml create mode 100644 src/maxtext/configs/models/weaver.yml diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index 04082beb93..a8aa4fe197 100644 --- a/src/maxtext/common/common_types.py +++ b/src/maxtext/common/common_types.py @@ -141,6 +141,7 @@ class DecoderBlockType(enum.Enum): OLMO3 = "olmo3" DEEPSEEK4 = "deepseek4" ENVY = "envy" + WEAVER = "weaver" class VisionEncoderBlockType(enum.Enum): diff --git a/src/maxtext/configs/models/weaver-mini-diffuser.yml b/src/maxtext/configs/models/weaver-mini-diffuser.yml new file mode 100644 index 0000000000..b085780005 --- /dev/null +++ b/src/maxtext/configs/models/weaver-mini-diffuser.yml @@ -0,0 +1,46 @@ +# Copyright 2023–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 config for Weaver Mini Diffuser (36-layer Mixture-of-Transformers Joint Backbone) + +# Core Architectural Parameters +decoder_block: "weaver" +base_emb_dim: 4096 +base_mlp_dim: 12288 +base_num_query_heads: 32 +base_num_kv_heads: 8 +base_num_decoder_layers: 36 +head_dim: 128 +mlp_activations: ["silu", "linear"] +vocab_size: 151936 +normalization_layer_epsilon: 1.0e-6 +use_qk_norm: true +logits_via_embedding: false + +# RoPE Settings +rope_max_timescale: 5000000 + +# General Model Settings +enable_dropout: false + +# Multimodal & 3D M-RoPE Settings +use_multimodal: true +use_mrope: true +mrope_section: [24, 20, 20] + +# Note: Weaver-specific diffusion & pathway parameters (latent_channels, patch_size, +# time_embed_in_channels, timestep_scale, timestep_max_period, qk_norm_for_text, +# qk_norm_for_diffusion, use_und_k_norm_for_gen, weaver_block_q, weaver_block_kv) +# are isolated in src/maxtext/configs/models/weaver.yml and loaded by WeaverConfig. + diff --git a/src/maxtext/configs/models/weaver.yml b/src/maxtext/configs/models/weaver.yml new file mode 100644 index 0000000000..5334f7ff5a --- /dev/null +++ b/src/maxtext/configs/models/weaver.yml @@ -0,0 +1,50 @@ +# Copyright 2023–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. + +# Isolated Weaver Mixture-of-Transformers diffusion & pathway configuration. +# Loaded directly by WeaverConfig in src/maxtext/models/weaver.py so that +# Weaver-specific diffusion parameters stay isolated from base.yml and types.py. + +# Core Architectural Parameters +decoder_block: "weaver" +base_emb_dim: 4096 +base_mlp_dim: 12288 +base_num_query_heads: 32 +base_num_kv_heads: 8 +base_num_decoder_layers: 36 +head_dim: 128 +mlp_activations: ["silu", "linear"] +vocab_size: 151936 +normalization_layer_epsilon: 1.0e-6 +use_qk_norm: true +logits_via_embedding: false + +# RoPE & 3D M-RoPE Settings +rope_max_timescale: 5000000 +enable_dropout: false +use_multimodal: true +use_mrope: true +mrope_section: [24, 20, 20] + +# Weaver Diffusion & Dual-Pathway Settings +latent_channels: 48 +patch_size: 2 +time_embed_in_channels: 256 +timestep_scale: 0.001 +timestep_max_period: 10000 +qk_norm_for_text: true +qk_norm_for_diffusion: true +use_und_k_norm_for_gen: false +weaver_block_q: 128 +weaver_block_kv: 128 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 235030f555..1b2bb58b1d 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -292,6 +292,7 @@ class ProfilerType(str, Enum): "qwen3-vl-30b-a3b", "weaver-mini", "weaver-max", + "weaver-mini-diffuser", "qwen3-next-80b-a3b", "qwen3-omni-30b-a3b", "qwen3-custom-30b-a3b", @@ -5407,6 +5408,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "maxtext-omni-gemma3-qwen3", "weaver-mini", "weaver-max", + "weaver-mini-diffuser", ) if self.model_name not in valid_mm_models and self.model_name != "default": raise ValueError(f"Multimodal is only supported for {valid_mm_models}, not {self.model_name}") diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index fab552bf2f..d6114305e9 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -69,6 +69,7 @@ qwen3_5, qwen3_custom, simple_layer, + weaver, ) from maxtext.multimodal import utils as mm_utils @@ -1240,6 +1241,7 @@ def get_deepseek(): DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), DecoderBlockType.ENVY: get_scannable(envy.EnvyDecoderLayer, envy.EnvyScannableBlock), + DecoderBlockType.WEAVER: [weaver.WeaverMoTDecoderLayer], } if cfg.decoder_block not in layer_map: diff --git a/src/maxtext/models/weaver.py b/src/maxtext/models/weaver.py index 0809b8bd48..439f5cb337 100644 --- a/src/maxtext/models/weaver.py +++ b/src/maxtext/models/weaver.py @@ -29,6 +29,10 @@ 6. ``WeaverMLP`` feed-forward block (``silu`` SwiGLU and ``relu2`` variants). 7. ``WeaverMoTDecoderLayer`` dual-pathway Mixture-of-Transformers decoder block with canonical weight alias properties. +8. ``WeaverOmniTransformer`` full 36-layer joint multimodal diffusion backbone + integrating ``embed_tokens``, ``proj_in``, ``time_embedder``, stacked + ``WeaverMoTDecoderLayer`` blocks (``scan_layers=False`` and + ``scan_layers=True``), ``norm``, ``norm_moe_gen``, and ``proj_out``. ======================================================================== WEAVER DECODER LAYER DATA FLOW @@ -59,7 +63,10 @@ from collections.abc import Callable, Sequence import dataclasses +import functools +import inspect import math +import pathlib from typing import Any, Literal, overload from flax import nnx @@ -71,23 +78,57 @@ from maxtext.kernels.tokamax_splash_attention import splash_attention_kernel from maxtext.kernels.tokamax_splash_attention import splash_attention_mask from maxtext.kernels.tokamax_splash_attention import splash_attention_mask_info +from maxtext.layers import nnx_scan import numpy as np +import yaml # ============================================================================== # 0. Configuration & Pure-NNX Leaf Layer Builders # ============================================================================== +_WEAVER_YAML_PATH = pathlib.Path(__file__).resolve().parents[1] / "configs" / "models" / "weaver.yml" + + +@functools.lru_cache(maxsize=1) +def _load_weaver_yaml_defaults() -> dict[str, Any]: + """Loads isolated Weaver-specific defaults from ``src/maxtext/configs/models/weaver.yml``. + + Keeping Weaver-specific diffusion and dual-pathway parameters in ``weaver.yml`` + avoids polluting core MaxText ``base.yml`` and ``types.py`` while providing a + single source of truth for ``WeaverConfig`` defaults. + """ + if _WEAVER_YAML_PATH.is_file(): + with _WEAVER_YAML_PATH.open("r", encoding="utf-8") as f: + loaded = yaml.safe_load(f) + if isinstance(loaded, dict): + return loaded + return {} + + +def _resolve_jax_dtype(dtype: Any) -> Any: + """Resolves a string or dtype object to a JAX dtype.""" + if isinstance(dtype, str): + return getattr(jnp, dtype, jnp.dtype(dtype)) + return dtype + @dataclasses.dataclass class WeaverConfig: - """Configuration for Weaver Mixture-of-Transformers blocks.""" + """Configuration for Weaver Mixture-of-Transformers blocks and backbone.""" hidden_size: int = 4096 head_dim: int = 128 num_attention_heads: int = 32 num_key_value_heads: int = 8 intermediate_size: int = 12288 + num_hidden_layers: int = 36 + vocab_size: int = 151936 + latent_channels: int = 48 + patch_size: int = 2 + time_embed_in_channels: int = 256 + timestep_scale: float = 0.001 + timestep_max_period: int = 10000 rms_norm_eps: float = 1e-6 hidden_act: str = "silu" attention_bias: bool = False @@ -96,8 +137,18 @@ class WeaverConfig: use_und_k_norm_for_gen: bool = False attention_kernel: str = "dot_product" use_ragged_ops: bool = False + block_q: int = 128 + block_kv: int = 128 mrope_section: tuple[int, int, int] = (24, 20, 20) - rope_theta: float = 1_000_000.0 + rope_theta: float = 5_000_000.0 + scan_layers: bool = False + param_scan_axis: int = 0 + remat_policy: str = "none" + embed_kernel_axes: tuple[None | str, ...] = ("vocab", "embed") + proj_in_kernel_axes: tuple[None | str, ...] = ("patch", "embed") + time_mlp_in_kernel_axes: tuple[None | str, ...] = ("time_embed", "embed") + time_mlp_out_kernel_axes: tuple[None | str, ...] = ("embed", "embed") + proj_out_kernel_axes: tuple[None | str, ...] = ("embed", "patch") q_kernel_axes: tuple[None | str, ...] = ("embed", "heads") kv_kernel_axes: tuple[None | str, ...] = ("embed", "kv_heads") o_kernel_axes: tuple[None | str, ...] = ("heads", "embed") @@ -114,6 +165,108 @@ def __post_init__(self) -> None: "num_attention_heads must be divisible by num_key_value_heads, " f"got {self.num_attention_heads=} {self.num_key_value_heads=}" ) + self.dtype = _resolve_jax_dtype(self.dtype) + self.weight_dtype = _resolve_jax_dtype(self.weight_dtype) + + @classmethod + def from_maxtext_config(cls, config: Any) -> WeaverConfig: + """Translates a MaxText ``pyconfig`` object into a ``WeaverConfig`` dataclass. + + Clarification on ``WeaverConfig.from_maxtext_config`` vs ``Module.from_config``: + * ``WeaverConfig.from_maxtext_config`` (this method) is a **dataclass config + adapter**: it converts a MaxText ``pyconfig.HyperParameters`` / ``MaxTextConfig`` + object (or a partial config stub in unit tests) into a self-contained + ``WeaverConfig`` dataclass instance. + * ``WeaverMoTDecoderLayer.from_config`` and ``WeaverOmniTransformer.from_config`` + are **Flax NNX module factory classmethods**: they instantiate and return the + neural network module itself from a ``WeaverConfig`` or MaxText ``pyconfig`` + object (delegating config normalization to ``from_maxtext_config``). + + Why ``getattr(config, ..., default)`` is used here: + 1. Bridges MaxText's attribute names (e.g. ``emb_dim`` / ``base_emb_dim``, + ``num_query_heads``, ``mlp_dim``, ``num_decoder_layers``, + ``normalization_layer_epsilon``, ``rope_max_timescale``) to the + self-contained ``WeaverConfig`` field names. + 2. Reads Weaver-specific diffusion and pathway defaults (``latent_channels``, + ``patch_size``, ``time_embed_in_channels``, ``timestep_scale``, + ``timestep_max_period``, ``qk_norm_for_text``, ``qk_norm_for_diffusion``, + ``use_und_k_norm_for_gen``, ``weaver_block_q``, ``weaver_block_kv``) from + the isolated ``src/maxtext/configs/models/weaver.yml`` file when not + present on ``config``, avoiding the need to register Weaver-only fields in + core ``base.yml`` or ``types.py``. + 3. Allows callers and unit tests to pass partial/mock config objects without + requiring every MaxText field to be populated. + """ + if isinstance(config, cls): + return config + + yml = _load_weaver_yaml_defaults() + + mlp_acts = getattr(config, "mlp_activations", yml.get("mlp_activations", ["silu", "linear"])) + hidden_act = "relu2" if "relu2" in mlp_acts else "silu" + + raw_attn = str(getattr(config, "attention", "dot_product")) + attn_kernel = raw_attn if raw_attn in ("dot_product", "flash", "splash") else "dot_product" + + mrope_raw = getattr(config, "mrope_section", yml.get("mrope_section", (24, 20, 20))) + mrope_tuple = (int(mrope_raw[0]), int(mrope_raw[1]), int(mrope_raw[2])) + + # Core architectural fields: check derived MaxText names (e.g. emb_dim) first, + # then base MaxText names (e.g. base_emb_dim), then isolated weaver.yml defaults. + return cls( + hidden_size=int(getattr(config, "emb_dim", getattr(config, "base_emb_dim", yml.get("base_emb_dim", 4096)))), + head_dim=int(getattr(config, "head_dim", yml.get("head_dim", 128))), + num_attention_heads=int( + getattr( + config, + "num_query_heads", + getattr(config, "base_num_query_heads", yml.get("base_num_query_heads", 32)), + ) + ), + num_key_value_heads=int( + getattr(config, "num_kv_heads", getattr(config, "base_num_kv_heads", yml.get("base_num_kv_heads", 8))) + ), + intermediate_size=int( + getattr(config, "mlp_dim", getattr(config, "base_mlp_dim", yml.get("base_mlp_dim", 12288))) + ), + num_hidden_layers=int( + getattr( + config, + "num_decoder_layers", + getattr(config, "base_num_decoder_layers", yml.get("base_num_decoder_layers", 36)), + ) + ), + vocab_size=int(getattr(config, "vocab_size", yml.get("vocab_size", 151936))), + # Weaver-specific diffusion & pathway fields default to src/maxtext/configs/models/weaver.yml. + latent_channels=int(getattr(config, "latent_channels", yml.get("latent_channels", 48))), + patch_size=int(getattr(config, "patch_size", yml.get("patch_size", 2))), + time_embed_in_channels=int(getattr(config, "time_embed_in_channels", yml.get("time_embed_in_channels", 256))), + timestep_scale=float(getattr(config, "timestep_scale", yml.get("timestep_scale", 0.001))), + timestep_max_period=int(getattr(config, "timestep_max_period", yml.get("timestep_max_period", 10000))), + rms_norm_eps=float(getattr(config, "normalization_layer_epsilon", yml.get("normalization_layer_epsilon", 1e-6))), + hidden_act=hidden_act, + attention_bias=bool(getattr(config, "attention_bias", False)), + qk_norm_for_text=bool( + getattr( + config, + "qk_norm_for_text", + getattr(config, "use_qk_norm", yml.get("qk_norm_for_text", True)), + ) + ), + qk_norm_for_diffusion=bool(getattr(config, "qk_norm_for_diffusion", yml.get("qk_norm_for_diffusion", True))), + use_und_k_norm_for_gen=bool(getattr(config, "use_und_k_norm_for_gen", yml.get("use_und_k_norm_for_gen", False))), + attention_kernel=attn_kernel, + use_ragged_ops=bool(getattr(config, "use_ragged_attention", False)), + block_q=int(getattr(config, "weaver_block_q", yml.get("weaver_block_q", 128))), + block_kv=int(getattr(config, "weaver_block_kv", yml.get("weaver_block_kv", 128))), + mrope_section=mrope_tuple, + rope_theta=float(getattr(config, "rope_max_timescale", yml.get("rope_max_timescale", 5_000_000.0))), + scan_layers=bool(getattr(config, "scan_layers", False)), + param_scan_axis=int(getattr(config, "param_scan_axis", 0)), + remat_policy=str(getattr(config, "remat_policy", "none")), + dtype=_resolve_jax_dtype(getattr(config, "dtype", jnp.float32)), + weight_dtype=_resolve_jax_dtype(getattr(config, "weight_dtype", jnp.float32)), + ) def _adapt_kernel_init( @@ -123,17 +276,52 @@ def _adapt_kernel_init( if kernel_init is None: return nnx.initializers.lecun_normal() + num_params = len(inspect.signature(kernel_init).parameters) + def _wrapped(key: jax.Array, shape: Sequence[int], dtype: Any = jnp.float32) -> jax.Array: - try: - return kernel_init(key, shape, dtype) - except TypeError: + if num_params >= 5: in_axis = (0,) out_axis = tuple(range(1, len(shape))) return kernel_init(key, shape, dtype, in_axis, out_axis) + return kernel_init(key, shape, dtype) return _wrapped +def _make_embed( + num_embeddings: int, + features: int, + *, + dtype: Any = jnp.float32, + weight_dtype: Any = jnp.float32, + embedding_init: Callable[..., jax.Array] | None = None, + kernel_axes: tuple[None | str, ...] = ("vocab", "embed"), + sharding_hook: Callable[[Any, str], Any] | None = None, + hook_name: str = "embed_kernel_init", + rngs: nnx.Rngs, +) -> nnx.Embed: + """Constructs a pure ``nnx.Embed`` with logical-axis annotations.""" + base_init = _adapt_kernel_init(embedding_init or nnx.initializers.normal(stddev=0.02)) + if sharding_hook is not None: + base_init = sharding_hook(base_init, hook_name) + meta_init = nnx.with_metadata( + base_init, + out_sharding=kernel_axes, + kernel_axes=kernel_axes, + eager_sharding=False, + ) + embed = nnx.Embed( + num_embeddings=num_embeddings, + features=features, + dtype=dtype, + param_dtype=weight_dtype, + embedding_init=meta_init, + rngs=rngs, + ) + embed.kernel_axes = kernel_axes + return embed + + def _make_linear( in_features: int, out_features: int, @@ -537,7 +725,7 @@ def compute_3d_mrope_cos_sin( position_ids_3d: jax.Array, head_dim: int, mrope_section: tuple[int, int, int] = (24, 20, 20), - rope_theta: float = 1000000.0, + rope_theta: float = 5_000_000.0, ) -> tuple[jax.Array, jax.Array]: """Computes interleaved 3D M-RoPE cosine and sine tensors from 3D coordinates.""" if sum(mrope_section) != head_dim // 2: @@ -578,6 +766,8 @@ def causal_understanding_attention( scale: float | None = None, mask_value: float = -1e9, attention_kernel: str = "dot_product", + block_q: int = 128, + block_kv: int = 128, ) -> jax.Array: """Kernel 1: Causal Understanding Attention with cross-sample fences.""" n_und, num_heads, head_dim = q_und.shape @@ -594,8 +784,6 @@ def causal_understanding_attention( scale = 1.0 / math.sqrt(head_dim) if attention_kernel == "flash": - block_q = 128 - block_kv = 128 if n_und % block_q != 0: raise ValueError( "FlashAttention requires sequence length divisible by block_q" @@ -621,8 +809,6 @@ def causal_understanding_attention( return out.astype(q_und.dtype) elif attention_kernel == "splash": - block_q = 128 - block_kv = 128 if n_und % block_q != 0: raise ValueError( "Splash Attention requires sequence length divisible by block size" @@ -670,6 +856,8 @@ def full_generative_attention( scale: float | None = None, mask_value: float = -1e9, attention_kernel: str = "dot_product", + block_q: int = 128, + block_kv: int = 128, ) -> jax.Array: """Kernel 2: Full Generative Attention with cross-sample fences.""" n_gen, num_heads, head_dim = q_gen.shape @@ -687,8 +875,6 @@ def full_generative_attention( scale = 1.0 / math.sqrt(head_dim) if attention_kernel == "flash": - block_q = 128 - block_kv = 128 if n_gen % block_q != 0 or total_tokens % block_kv != 0: raise ValueError( "FlashAttention requires sequence lengths divisible by block sizes" @@ -715,8 +901,6 @@ def full_generative_attention( return out.astype(q_gen.dtype) elif attention_kernel == "splash": - block_q = 128 - block_kv = 128 if n_gen % block_q != 0 or total_tokens % block_kv != 0: raise ValueError( "Splash Attention requires sequence lengths divisible by block sizes" @@ -782,6 +966,8 @@ def __init__( rms_norm_eps: float = 1e-6, attention_kernel: str = "dot_product", use_ragged_ops: bool = False, + block_q: int = 128, + block_kv: int = 128, q_kernel_axes: tuple[None | str, ...] = ("embed", "heads"), kv_kernel_axes: tuple[None | str, ...] = ("embed", "kv_heads"), o_kernel_axes: tuple[None | str, ...] = ("heads", "embed"), @@ -808,6 +994,8 @@ def __init__( rms_norm_eps = cfg.rms_norm_eps attention_kernel = cfg.attention_kernel use_ragged_ops = cfg.use_ragged_ops + block_q = cfg.block_q + block_kv = cfg.block_kv q_kernel_axes = cfg.q_kernel_axes kv_kernel_axes = cfg.kv_kernel_axes o_kernel_axes = cfg.o_kernel_axes @@ -826,6 +1014,8 @@ def __init__( self.attention_kernel = attention_kernel self.use_ragged_ops = use_ragged_ops + self.block_q = block_q + self.block_kv = block_kv self.dim = dim self.num_heads = num_heads self.num_kv_heads = num_kv_heads @@ -1187,6 +1377,8 @@ def __call__( causal_q_offsets=metadata.causal_q_offsets, scale=self.scaling, attention_kernel=self.attention_kernel, + block_q=self.block_q, + block_kv=self.block_kv, ) und_attn_context = checkpoint_name(hook(und_attn_context, "attn_out"), "attention_out") und_attn_flat = und_attn_context.reshape(n_und, self.num_heads * self.head_dim) @@ -1204,6 +1396,8 @@ def __call__( sample_kv_offsets=metadata.sample_kv_offsets, scale=self.scaling, attention_kernel=self.attention_kernel, + block_q=self.block_q, + block_kv=self.block_kv, ) gen_attn_context = checkpoint_name(hook(gen_attn_context, "gen_attn_out"), "gen_attention_out") gen_attn_flat = gen_attn_context.reshape(n_gen, self.num_heads * self.head_dim) @@ -1376,6 +1570,8 @@ def __init__( use_und_k_norm_for_gen: bool = False, attention_kernel: str = "dot_product", use_ragged_ops: bool = False, + block_q: int = 128, + block_kv: int = 128, q_kernel_axes: tuple[None | str, ...] = ("embed", "heads"), kv_kernel_axes: tuple[None | str, ...] = ("embed", "kv_heads"), o_kernel_axes: tuple[None | str, ...] = ("heads", "embed"), @@ -1390,10 +1586,15 @@ def __init__( kernel_init: Callable[..., jax.Array] | None = None, rngs: nnx.Rngs, ): - """Initializes the Weaver Mixture-of-Transformers decoder layer.""" + """Initializes the Weaver Mixture-of-Transformers decoder layer. + + Accepts either an integer ``hidden_size`` (with explicit architectural + arguments), a ``WeaverConfig`` dataclass, or a MaxText ``pyconfig`` object + (which is converted to ``WeaverConfig`` via ``WeaverConfig.from_maxtext_config``). + """ del mesh - if isinstance(hidden_size, WeaverConfig): - cfg = hidden_size + if not isinstance(hidden_size, int): + cfg = hidden_size if isinstance(hidden_size, WeaverConfig) else WeaverConfig.from_maxtext_config(hidden_size) hidden_size = cfg.hidden_size head_dim = cfg.head_dim num_attention_heads = cfg.num_attention_heads @@ -1407,6 +1608,8 @@ def __init__( use_und_k_norm_for_gen = cfg.use_und_k_norm_for_gen attention_kernel = cfg.attention_kernel use_ragged_ops = cfg.use_ragged_ops + block_q = cfg.block_q + block_kv = cfg.block_kv q_kernel_axes = cfg.q_kernel_axes kv_kernel_axes = cfg.kv_kernel_axes o_kernel_axes = cfg.o_kernel_axes @@ -1445,6 +1648,8 @@ def __init__( rms_norm_eps=rms_norm_eps, attention_kernel=attention_kernel, use_ragged_ops=use_ragged_ops, + block_q=block_q, + block_kv=block_kv, q_kernel_axes=q_kernel_axes, kv_kernel_axes=kv_kernel_axes, o_kernel_axes=o_kernel_axes, @@ -1495,13 +1700,19 @@ def _mlp() -> WeaverMLP: @classmethod def from_config( cls, - config: WeaverConfig, + config: WeaverConfig | Any, *, rngs: nnx.Rngs, sharding_hook: Callable[[Any, str], Any] | None = None, kernel_init: Callable[..., jax.Array] | None = None, ) -> WeaverMoTDecoderLayer: - """Constructs ``WeaverMoTDecoderLayer`` from a ``WeaverConfig``.""" + """Flax NNX module factory that instantiates ``WeaverMoTDecoderLayer`` from a config. + + Unlike ``WeaverConfig.from_maxtext_config`` (which converts a MaxText + ``pyconfig`` into a ``WeaverConfig`` dataclass), ``from_config`` is a module + factory classmethod that constructs and returns a ``WeaverMoTDecoderLayer`` + instance from either a ``WeaverConfig`` or a MaxText ``pyconfig`` object. + """ return cls(config, sharding_hook=sharding_hook, kernel_init=kernel_init, rngs=rngs) # --------------------------------------------------------------------- @@ -1628,3 +1839,645 @@ def __call__( mlp_out_gen = self.mlp_moe_gen(mlp_in_gen) return hook(residual_und + mlp_out_und, "post_mlp"), hook(residual_gen + mlp_out_gen, "gen_post_mlp") + + +# ============================================================================== +# 8. Weaver Omni-Transformer Full Backbone Architecture (Pure Flax NNX) +# ============================================================================== + + +def patchify_latents(latents: jax.Array, patch_size: int = 2) -> jax.Array: + """Patchifies 5D video/image latents ``[B, C, T, H, W]`` into ``[B, S_gen, p*p*C]``. + + Matches the spatial 2D patchification layout ``b c t (h p1) (w p2) -> b (t h w) (p1 p2 c)``. + + Args: + latents: Latent tensor of shape ``[B, C, T, H, W]``. + patch_size: Spatial patch size ``p`` along height and width. + + Returns: + Patchified token sequence of shape ``[B, T * (H // p) * (W // p), p * p * C]``. + """ + if latents.ndim != 5: + raise ValueError(f"latents must be a 5D array [B, C, T, H, W], got shape {latents.shape}.") + bsz, c_lat, t_lat, h_lat, w_lat = latents.shape + if h_lat % patch_size != 0 or w_lat % patch_size != 0: + raise ValueError(f"Latent spatial dimensions ({h_lat}, {w_lat}) must be divisible by patch_size ({patch_size}).") + h_patch = h_lat // patch_size + w_patch = w_lat // patch_size + s_gen = t_lat * h_patch * w_patch + patch_dim = patch_size * patch_size * c_lat + + reshaped = latents.reshape(bsz, c_lat, t_lat, h_patch, patch_size, w_patch, patch_size) + permuted = jnp.einsum("bcthpwq->bthwpqc", reshaped) + return permuted.reshape(bsz, s_gen, patch_dim) + + +def unpatchify_latents( + patches: jax.Array, + t_lat: int, + h_lat: int, + w_lat: int, + *, + patch_size: int = 2, + latent_channels: int = 48, +) -> jax.Array: + """Unpatchifies generation tokens ``[B, S_gen, p*p*C]`` or ``[B*S_gen, p*p*C]`` to ``[B, C, T, H, W]``. + + Matches the inverse spatial layout ``(b t h w) (p1 p2 c) -> b c t (h p1) (w p2)``. + + Args: + patches: Patch token tensor of shape ``[B, S_gen, p*p*C]`` or ``[B * S_gen, p*p*C]``. + t_lat: Number of latent temporal frames ``T``. + h_lat: Latent spatial height ``H``. + w_lat: Latent spatial width ``W``. + patch_size: Spatial patch size ``p`` along height and width. + latent_channels: Number of latent channels ``C``. + + Returns: + Unpatchified latent tensor of shape ``[B, C, T, H, W]``. + """ + if h_lat % patch_size != 0 or w_lat % patch_size != 0: + raise ValueError(f"Latent spatial dimensions ({h_lat}, {w_lat}) must be divisible by patch_size ({patch_size}).") + h_patch = h_lat // patch_size + w_patch = w_lat // patch_size + s_gen = t_lat * h_patch * w_patch + patch_dim = patch_size * patch_size * latent_channels + + if patches.shape[-1] != patch_dim: + raise ValueError(f"Expected patch feature dim {patch_dim}, got {patches.shape[-1]}.") + + if patches.ndim == 2: + if patches.shape[0] % s_gen != 0: + raise ValueError(f"Total patch tokens ({patches.shape[0]}) must be divisible by s_gen ({s_gen}).") + bsz = patches.shape[0] // s_gen + elif patches.ndim == 3: + bsz = patches.shape[0] + if patches.shape[1] != s_gen: + raise ValueError(f"Expected s_gen={s_gen} along axis 1, got {patches.shape[1]}.") + else: + raise ValueError(f"patches must be 2D or 3D, got shape {patches.shape}.") + + reshaped = patches.reshape(bsz, t_lat, h_patch, w_patch, patch_size, patch_size, latent_channels) + permuted = jnp.einsum("bthwpqc->bcthpwq", reshaped) + return permuted.reshape(bsz, latent_channels, t_lat, h_lat, w_lat) + + +def get_weaver_timestep_embedding( + timesteps: jax.Array, + num_channels: int = 256, + *, + flip_sin_to_cos: bool = True, + downscale_freq_shift: float = 0.0, + scale: float = 1.0, + max_period: int = 10000, +) -> jax.Array: + """Computes sinusoidal timestep embeddings in ``float32``. + + Args: + timesteps: Timestep tensor of shape ``[N]`` or ``[B, T]``. + num_channels: Number of sinusoidal embedding channels (default 256). + flip_sin_to_cos: Whether to order ``[cos, sin]`` instead of ``[sin, cos]``. + downscale_freq_shift: Frequency shift denominator offset. + scale: Multiplier applied to sinusoidal arguments. + max_period: Controls minimum frequency of the embeddings. + + Returns: + Float32 sinusoidal embeddings of shape ``(*timesteps.shape, num_channels)``. + """ + half_dim = num_channels // 2 + exponent = -math.log(float(max_period)) * jnp.arange(0, half_dim, dtype=jnp.float32) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = jnp.exp(exponent) + emb = timesteps[..., None].astype(jnp.float32) * emb + + emb = scale * emb + emb = jnp.concatenate([jnp.sin(emb), jnp.cos(emb)], axis=-1) + + if flip_sin_to_cos: + emb = jnp.concatenate([emb[..., half_dim:], emb[..., :half_dim]], axis=-1) + + if num_channels % 2 == 1: + emb = jnp.pad(emb, ((0, 0),) * (emb.ndim - 1) + ((0, 1),)) + + return emb + + +class WeaverTimeEmbedder(nnx.Module): + """Sinusoidal + 2-layer SiLU MLP diffusion timestep embedder using pure ``nnx.Linear``.""" + + def __init__( + self, + in_channels: int = 256, + time_embed_dim: int = 4096, + *, + timestep_scale: float = 0.001, + max_period: int = 10000, + use_bias: bool = True, + in_kernel_axes: tuple[None | str, ...] = ("time_embed", "embed"), + out_kernel_axes: tuple[None | str, ...] = ("embed", "embed"), + sharding_hook: Callable[[Any, str], Any] | None = None, + dtype: Any = jnp.float32, + weight_dtype: Any = jnp.float32, + kernel_init: Callable[..., jax.Array] | None = None, + rngs: nnx.Rngs, + ): + """Initializes the diffusion timestep embedder.""" + self.in_channels = in_channels + self.time_embed_dim = time_embed_dim + self.timestep_scale = timestep_scale + self.max_period = max_period + self.use_bias = use_bias + self.dtype = _resolve_jax_dtype(dtype) + self.weight_dtype = _resolve_jax_dtype(weight_dtype) + + self.mlp_0 = _make_linear( + in_channels, + time_embed_dim, + use_bias=use_bias, + dtype=jnp.float32, + weight_dtype=self.weight_dtype, + kernel_init=kernel_init, + kernel_axes=in_kernel_axes, + sharding_hook=sharding_hook, + hook_name="time_mlp_0_kernel_init", + rngs=rngs, + ) + self.mlp_2 = _make_linear( + time_embed_dim, + time_embed_dim, + use_bias=use_bias, + dtype=jnp.float32, + weight_dtype=self.weight_dtype, + kernel_init=kernel_init, + kernel_axes=out_kernel_axes, + sharding_hook=sharding_hook, + hook_name="time_mlp_2_kernel_init", + rngs=rngs, + ) + + @property + def linear_1(self) -> nnx.Linear: + """First linear projection in the timestep MLP (``mlp_0``).""" + return self.mlp_0 + + @property + def linear_2(self) -> nnx.Linear: + """Second linear projection in the timestep MLP (``mlp_2``).""" + return self.mlp_2 + + def __call__(self, timesteps: jax.Array) -> jax.Array: + """Embeds scalar/per-frame diffusion timesteps into ``time_embed_dim``.""" + scaled_t = timesteps.astype(jnp.float32) * self.timestep_scale + t_emb = get_weaver_timestep_embedding( + scaled_t, + num_channels=self.in_channels, + flip_sin_to_cos=True, + downscale_freq_shift=0.0, + max_period=self.max_period, + ) + hidden = jax.nn.silu(self.mlp_0(t_emb)) + return self.mlp_2(hidden).astype(self.dtype) + + +def build_weaver_3d_position_ids( + batch_size: int, + s_und: int, + t_lat: int, + h_patch: int, + w_patch: int, +) -> jax.Array: + """Constructs packed 3D M-RoPE coordinates ``[N_total, 3]`` for ``[text, vision]`` samples. + + For each sample: + * Understanding (text) tokens have 1D positions ``(i, i, i)`` for ``i in 0..s_und-1``. + * Generation (vision) tokens have 3D positions ``(s_und + t, h, w)`` across the + ``(t_lat, h_patch, w_patch)`` patch grid. + + Args: + batch_size: Number of samples ``B``. + s_und: Number of understanding (text) tokens per sample. + t_lat: Number of temporal latent frames ``T``. + h_patch: Number of spatial height patches ``H // p``. + w_patch: Number of spatial width patches ``W // p``. + + Returns: + Integer coordinate array of shape ``[B * (s_und + t_lat * h_patch * w_patch), 3]``. + """ + pos_und = jnp.broadcast_to(jnp.arange(s_und, dtype=jnp.int32)[:, None], (s_und, 3)) + + t_grid = jnp.repeat(jnp.arange(t_lat, dtype=jnp.int32) + s_und, h_patch * w_patch) + h_grid = jnp.tile(jnp.repeat(jnp.arange(h_patch, dtype=jnp.int32), w_patch), t_lat) + w_grid = jnp.tile(jnp.arange(w_patch, dtype=jnp.int32), t_lat * h_patch) + pos_gen = jnp.stack([t_grid, h_grid, w_grid], axis=-1) + + sample_pos = jnp.concatenate([pos_und, pos_gen], axis=0) + return jnp.tile(sample_pos, (batch_size, 1)) + + +class WeaverOmniTransformer(nnx.Module): + """Complete Weaver joint Mixture-of-Transformers diffusion backbone. + + Integrates: + * ``embed_tokens``: Text token embedding table (``nnx.Embed``). + * ``proj_in`` (alias ``vae2llm``): Patchified latent input projection (``nnx.Linear``). + * ``time_embedder``: Sinusoidal + 2-layer SiLU MLP timestep embedder (``WeaverTimeEmbedder``). + * ``layers`` / ``scanned_layers``: Stack of ``WeaverMoTDecoderLayer`` blocks supporting + both ``scan_layers=False`` and ``scan_layers=True``. + * ``norm``: Final understanding pathway RMSNorm (``nnx.RMSNorm``). + * ``norm_moe_gen``: Final generation pathway RMSNorm (``nnx.RMSNorm``). + * ``proj_out`` (alias ``llm2vae``): Output projection back to latent patch space (``nnx.Linear``). + """ + + def __init__( + self, + config: WeaverConfig | Any | int = 4096, + *, + head_dim: int | None = None, + num_attention_heads: int | None = None, + num_key_value_heads: int | None = None, + intermediate_size: int | None = None, + num_hidden_layers: int | None = None, + vocab_size: int | None = None, + latent_channels: int | None = None, + patch_size: int | None = None, + time_embed_in_channels: int | None = None, + timestep_scale: float | None = None, + timestep_max_period: int | None = None, + rms_norm_eps: float | None = None, + hidden_act: str | None = None, + attention_bias: bool | None = None, + qk_norm_for_text: bool | None = None, + qk_norm_for_diffusion: bool | None = None, + use_und_k_norm_for_gen: bool | None = None, + attention_kernel: str | None = None, + use_ragged_ops: bool | None = None, + block_q: int | None = None, + block_kv: int | None = None, + mrope_section: tuple[int, int, int] | None = None, + rope_theta: float | None = None, + scan_layers: bool | None = None, + param_scan_axis: int | None = None, + remat_policy: str | None = None, + dtype: Any | None = None, + weight_dtype: Any | None = None, + sharding_hook: Callable[[Any, str], Any] | None = None, + mesh: sharding.Mesh | None = None, + kernel_init: Callable[..., jax.Array] | None = None, + rngs: nnx.Rngs, + ): + """Initializes ``WeaverOmniTransformer`` from ``WeaverConfig``, MaxText ``Config``, or kwargs. + + Config resolution happens in two steps: + 1. Normalize ``config`` into a ``WeaverConfig`` dataclass (using + ``WeaverConfig.from_maxtext_config`` when a MaxText ``pyconfig`` object is + provided). + 2. Apply any non-``None`` keyword arguments via ``dataclasses.replace``. + """ + del mesh + if isinstance(config, WeaverConfig): + cfg = dataclasses.replace(config) + elif isinstance(config, int): + cfg = WeaverConfig(hidden_size=config) + else: + cfg = WeaverConfig.from_maxtext_config(config) + + local_vars = locals() + overrides = { + f.name: local_vars[f.name] + for f in dataclasses.fields(WeaverConfig) + if f.name in local_vars and local_vars[f.name] is not None + } + if overrides: + cfg = dataclasses.replace(cfg, **overrides) + + self.config = cfg + self.hidden_size = cfg.hidden_size + self.head_dim = cfg.head_dim + self.num_attention_heads = cfg.num_attention_heads + self.num_key_value_heads = cfg.num_key_value_heads + self.intermediate_size = cfg.intermediate_size + self.num_hidden_layers = cfg.num_hidden_layers + self.vocab_size = cfg.vocab_size + self.latent_channels = cfg.latent_channels + self.patch_size = cfg.patch_size + self.patch_dim = cfg.patch_size * cfg.patch_size * cfg.latent_channels + self.time_embed_in_channels = cfg.time_embed_in_channels + self.timestep_scale = cfg.timestep_scale + self.timestep_max_period = cfg.timestep_max_period + self.rms_norm_eps = cfg.rms_norm_eps + self.hidden_act = cfg.hidden_act + self.mrope_section = cfg.mrope_section + self.rope_theta = cfg.rope_theta + self.scan_layers = cfg.scan_layers + self.param_scan_axis = cfg.param_scan_axis + self.remat_policy = cfg.remat_policy + self.use_ragged_ops = cfg.use_ragged_ops + self.dtype = cfg.dtype + self.weight_dtype = cfg.weight_dtype + self.sharding_hook = sharding_hook or (lambda x, name: x) + + # 1. Text token embedding table. + self.embed_tokens = _make_embed( + num_embeddings=cfg.vocab_size, + features=cfg.hidden_size, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_axes=cfg.embed_kernel_axes, + sharding_hook=sharding_hook, + hook_name="embed_tokens_kernel_init", + rngs=rngs, + ) + + # 2. Latent patch input projection (vae2llm). + self.proj_in = _make_linear( + self.patch_dim, + cfg.hidden_size, + use_bias=True, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_init=kernel_init, + kernel_axes=cfg.proj_in_kernel_axes, + sharding_hook=sharding_hook, + hook_name="proj_in_kernel_init", + rngs=rngs, + ) + + # 3. Diffusion timestep embedder. + self.time_embedder = WeaverTimeEmbedder( + in_channels=cfg.time_embed_in_channels, + time_embed_dim=cfg.hidden_size, + timestep_scale=cfg.timestep_scale, + max_period=cfg.timestep_max_period, + in_kernel_axes=cfg.time_mlp_in_kernel_axes, + out_kernel_axes=cfg.time_mlp_out_kernel_axes, + sharding_hook=sharding_hook, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_init=kernel_init, + rngs=rngs, + ) + + # 4. Stacked Mixture-of-Transformers decoder blocks. + if cfg.scan_layers: + self.layers = None + self.scanned_layers = nnx_scan.create_scanned_layers( + lambda layer_rngs: WeaverMoTDecoderLayer( + cfg, + sharding_hook=sharding_hook, + kernel_init=kernel_init, + rngs=layer_rngs, + ), + length=cfg.num_hidden_layers, + param_scan_axis=cfg.param_scan_axis, + metadata_axis_name="layers", + rngs=rngs, + ) + else: + self.layers = nnx.List( + [ + WeaverMoTDecoderLayer( + cfg, + sharding_hook=sharding_hook, + kernel_init=kernel_init, + rngs=rngs, + ) + for _ in range(cfg.num_hidden_layers) + ] + ) + self.scanned_layers = None + + # 5. Final pathway RMSNorms. + self.norm = _make_rmsnorm( + cfg.hidden_size, + epsilon=cfg.rms_norm_eps, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_axes=cfg.norm_kernel_axes, + sharding_hook=sharding_hook, + hook_name="norm_scale_init", + rngs=rngs, + ) + self.norm_moe_gen = _make_rmsnorm( + cfg.hidden_size, + epsilon=cfg.rms_norm_eps, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_axes=cfg.norm_kernel_axes, + sharding_hook=sharding_hook, + hook_name="norm_moe_gen_scale_init", + rngs=rngs, + ) + + # 6. Output projection to latent patch space (llm2vae). + self.proj_out = _make_linear( + cfg.hidden_size, + self.patch_dim, + use_bias=True, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_init=kernel_init, + kernel_axes=cfg.proj_out_kernel_axes, + sharding_hook=sharding_hook, + hook_name="proj_out_kernel_init", + rngs=rngs, + ) + + @classmethod + def from_config( + cls, + config: WeaverConfig | Any, + *, + rngs: nnx.Rngs, + mesh: sharding.Mesh | None = None, + sharding_hook: Callable[[Any, str], Any] | None = None, + kernel_init: Callable[..., jax.Array] | None = None, + **overrides: Any, + ) -> WeaverOmniTransformer: + """Flax NNX module factory that instantiates ``WeaverOmniTransformer`` from a config. + + Unlike ``WeaverConfig.from_maxtext_config`` (which translates a MaxText + ``pyconfig`` into a ``WeaverConfig`` dataclass), ``from_config`` is a module + factory classmethod that constructs and returns a ``WeaverOmniTransformer`` + instance from either a ``WeaverConfig`` or a MaxText ``pyconfig`` object, + forwarding any keyword overrides to ``WeaverOmniTransformer.__init__``. + """ + return cls( + config, + mesh=mesh, + sharding_hook=sharding_hook, + kernel_init=kernel_init, + rngs=rngs, + **overrides, + ) + + @property + def vae2llm(self) -> nnx.Linear: + """Canonical alias for the latent patch input projection (``proj_in``).""" + return self.proj_in + + @property + def llm2vae(self) -> nnx.Linear: + """Canonical alias for the latent patch output projection (``proj_out``).""" + return self.proj_out + + @overload + def __call__( + self, + input_ids: jax.Array, + latents: jax.Array, + timesteps: jax.Array, + *, + position_ids: jax.Array | None = None, + metadata: WeaverPackingMetadata | None = None, + return_hidden_states: Literal[False] = False, + ) -> jax.Array: + ... + + @overload + def __call__( + self, + input_ids: jax.Array, + latents: jax.Array, + timesteps: jax.Array, + *, + position_ids: jax.Array | None = None, + metadata: WeaverPackingMetadata | None = None, + return_hidden_states: Literal[True], + ) -> tuple[jax.Array, jax.Array]: + ... + + def __call__( + self, + input_ids: jax.Array, + latents: jax.Array, + timesteps: jax.Array, + *, + position_ids: jax.Array | None = None, + metadata: WeaverPackingMetadata | None = None, + return_hidden_states: bool = False, + ) -> jax.Array | tuple[jax.Array, jax.Array]: + """Executes the full Weaver Omni-Transformer forward pass. + + Args: + input_ids: Understanding (text) token IDs of shape ``[B, S_und]``. + latents: Noisy video/image latent tensor of shape ``[B, C, T, H, W]``. + timesteps: Diffusion timesteps of shape ``[B]``, ``[B, 1]``, or ``[B, T]``. + position_ids: Optional 3D M-RoPE position IDs of shape ``[N_total, 3]`` + or ``[3, N_total]``. Computed automatically from ``(B, S_und, T, H//p, W//p)`` + when ``None``. + metadata: Optional pre-built ``WeaverPackingMetadata``. Constructed + automatically for ``[S_und] * B`` and ``[S_gen] * B`` when ``None``. + return_hidden_states: When ``True``, returns ``(preds_vision, last_hidden_state)`` + where ``last_hidden_state`` has shape ``[N_total, hidden_size]``. + + Returns: + Predicted velocity/noise latent tensor of shape ``[B, C, T, H, W]`` (or a + tuple ``(preds_vision, last_hidden_state)`` when ``return_hidden_states=True``). + """ + if input_ids.ndim != 2: + raise ValueError(f"input_ids must be 2D [B, S_und], got shape {input_ids.shape}.") + if latents.ndim != 5: + raise ValueError(f"latents must be 5D [B, C, T, H, W], got shape {latents.shape}.") + + bsz, s_und = input_ids.shape + bsz_lat, c_lat, t_lat, h_lat, w_lat = latents.shape + if bsz != bsz_lat: + raise ValueError(f"Batch size mismatch between input_ids ({bsz}) and latents ({bsz_lat}).") + if c_lat != self.latent_channels: + raise ValueError(f"Expected latent_channels={self.latent_channels}, got {c_lat}.") + + p = self.patch_size + h_patch = h_lat // p + w_patch = w_lat // p + s_gen = t_lat * h_patch * w_patch + + # 1. Embed understanding (text) tokens -> [B * S_und, hidden_size]. + text_embeds = self.embed_tokens(input_ids).astype(self.dtype) + und_seq = text_embeds.reshape(bsz * s_und, self.hidden_size) + + # 2. Patchify latents, project to hidden_size, and add timestep embeddings -> [B * S_gen, hidden_size]. + patches = patchify_latents(latents.astype(self.dtype), patch_size=p) + vision_embeds = self.proj_in(patches) + + time_embeds = self.time_embedder(timesteps).astype(self.dtype) + if time_embeds.ndim == 2: + time_embeds = time_embeds[:, None, :] + elif time_embeds.ndim == 3 and time_embeds.shape[1] == t_lat and s_gen != t_lat: + time_embeds = jnp.repeat(time_embeds, h_patch * w_patch, axis=1) + + gen_tokens = vision_embeds + time_embeds + gen_seq = gen_tokens.reshape(bsz * s_gen, self.hidden_size) + + # 3. Build packing metadata and 3D M-RoPE rotary tables. + if metadata is None: + metadata = build_weaver_packing_metadata([s_und] * bsz, [s_gen] * bsz) + + if position_ids is None: + pos_3d = build_weaver_3d_position_ids( + batch_size=bsz, + s_und=s_und, + t_lat=t_lat, + h_patch=h_patch, + w_patch=w_patch, + ) + else: + if position_ids.ndim == 2 and position_ids.shape[0] == 3 and position_ids.shape[1] == metadata.total_tokens: + pos_3d = position_ids.T + else: + pos_3d = position_ids + + cos, sin = compute_3d_mrope_cos_sin( + pos_3d, + head_dim=self.head_dim, + mrope_section=self.mrope_section, + rope_theta=self.rope_theta, + ) + cos = cos.astype(self.dtype) + sin = sin.astype(self.dtype) + + # 4. Apply stacked Mixture-of-Transformers decoder blocks. + if self.scan_layers: + assert self.scanned_layers is not None + remat = self.remat_policy not in ("none", "") + und_seq, gen_seq = nnx_scan.apply_scanned_layers( + self.scanned_layers, + (und_seq, gen_seq), + length=self.num_hidden_layers, + param_scan_axis=self.param_scan_axis, + apply_fn=lambda layer, carry: layer(carry[0], carry[1], metadata, cos, sin), + remat=remat, + ) + else: + assert self.layers is not None + for layer in self.layers: + und_seq, gen_seq = layer(und_seq, gen_seq, metadata, cos, sin) + + # 5. Final pathway RMSNorms. + und_seq = self.norm(und_seq) + gen_seq = self.norm_moe_gen(gen_seq) + + # 6. Project generation tokens to patch space and unpatchify to [B, C, T, H, W]. + gen_patches = self.proj_out(gen_seq) + preds_vision = unpatchify_latents( + gen_patches, + t_lat=t_lat, + h_lat=h_lat, + w_lat=w_lat, + patch_size=p, + latent_channels=c_lat, + ) + + if return_hidden_states: + last_hidden_state = reinterleave_streams( + und_tokens=und_seq, + gen_tokens=gen_seq, + packed_und_token_indexes=metadata.packed_und_token_indexes, + packed_gen_token_indexes=metadata.packed_gen_token_indexes, + total_tokens=metadata.total_tokens, + use_ragged_ops=self.use_ragged_ops, + ) + return preds_vision, last_hidden_state + + return preds_vision diff --git a/tests/unit/weaver_layers_test.py b/tests/unit/weaver_layers_test.py index 92a981f3ba..6ae289d690 100644 --- a/tests/unit/weaver_layers_test.py +++ b/tests/unit/weaver_layers_test.py @@ -12,12 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit and golden parity tests for Weaver layers. +"""Unit, golden parity, and TPU tests for Weaver layers and WeaverOmniTransformer. -Consolidates tests for attention, multi-layer perceptron (MLP), and -Mixture-of-Transformers decoder layer modules. +Consolidates tests for attention, multi-layer perceptron (MLP), +Mixture-of-Transformers decoder layer modules, and the full WeaverOmniTransformer +backbone. """ +import dataclasses import os import subprocess from typing import Any @@ -27,23 +29,25 @@ from flax import nnx import jax import jax.numpy as jnp +from maxtext.common.common_types import DecoderBlockType +from maxtext.configs import pyconfig from maxtext.models import weaver +from maxtext.utils import globals as maxtext_globals import numpy as np import pytest -pytestmark = [ - pytest.mark.cpu_only, -] - -# Module-level aliases for weaver attention and layer symbols. +# Module-level aliases for weaver attention, layer, and backbone symbols. WeaverConfig = weaver.WeaverConfig WeaverMLP = weaver.WeaverMLP WeaverMoTDecoderLayer = weaver.WeaverMoTDecoderLayer WeaverAttention = weaver.WeaverAttention WeaverDualAttention = weaver.WeaverDualAttention +WeaverOmniTransformer = weaver.WeaverOmniTransformer +WeaverTimeEmbedder = weaver.WeaverTimeEmbedder apply_rotary_pos_emb = weaver.apply_rotary_pos_emb build_causal_understanding_mask = weaver.build_causal_understanding_mask build_causal_understanding_splash_mask = weaver.build_causal_understanding_splash_mask +build_weaver_3d_position_ids = weaver.build_weaver_3d_position_ids build_weaver_packing_metadata = weaver.build_weaver_packing_metadata build_full_generative_mask = weaver.build_full_generative_mask build_full_generative_splash_mask = weaver.build_full_generative_splash_mask @@ -51,11 +55,16 @@ compile_weaver_splash_mask = weaver.compile_weaver_splash_mask compute_3d_mrope_cos_sin = weaver.compute_3d_mrope_cos_sin full_generative_attention = weaver.full_generative_attention +get_weaver_timestep_embedding = weaver.get_weaver_timestep_embedding +patchify_latents = weaver.patchify_latents reinterleave_streams = weaver.reinterleave_streams +unpatchify_latents = weaver.unpatchify_latents unpack_streams = weaver.unpack_streams _ATTENTION_GOLDEN_FILENAME = "weaver_attention_golden_data.npz" _MOT_GOLDEN_FILENAME = "weaver_mot_golden_data.npz" +_TRANSFORMER_GOLDEN_FILENAME = "weaver_transformer_golden_data.npz" +_BASE_CONFIG_PATH = os.path.join(maxtext_globals.MAXTEXT_CONFIGS_DIR, "base.yml") _TOLERANCES = { "float32": {"rtol": 1e-4, "atol": 1e-4}, @@ -94,6 +103,7 @@ def _get_golden_data_path(filename: str) -> str | None: return None +@pytest.mark.cpu_only class WeaverPackingLogicTest(parameterized.TestCase): """Tests for packing logic, metadata, unpacking, and re-interleaving.""" @@ -144,6 +154,7 @@ def test_unpack_and_reinterleave_roundtrip(self): np.testing.assert_allclose(packed_tokens, reconstructed, rtol=1e-6, atol=1e-6) +@pytest.mark.cpu_only class WeaverMaskingTest(parameterized.TestCase): """Tests for sample isolation and attention masks.""" @@ -208,6 +219,7 @@ def test_full_generative_mask(self): self.assertTrue(bool(mask[q, kv]), f"Expected True at ({q}, {kv})") +@pytest.mark.cpu_only class WeaverDualAttentionKernelTest(parameterized.TestCase): """Tests for the functional attention kernels.""" @@ -297,6 +309,7 @@ def test_kernel_2_cross_sample_isolation(self): np.testing.assert_allclose(out_base[:3], out_perturbed[:3], rtol=1e-6, atol=1e-6) +@pytest.mark.cpu_only class WeaverMRoPETest(parameterized.TestCase): """Tests for 3D M-RoPE position encoding.""" @@ -326,6 +339,7 @@ def test_3d_mrope_shape_and_norm(self): np.testing.assert_allclose(norm_x, norm_x_rot, rtol=1e-5, atol=1e-5) +@pytest.mark.cpu_only class WeaverDualAttentionModuleTest(parameterized.TestCase): """Tests for the WeaverDualAttention NNX module.""" @@ -629,6 +643,7 @@ def test_ragged_ops_parity(self): np.testing.assert_array_equal(np.array(reint_def), np.array(reint_rag)) +@pytest.mark.cpu_only class WeaverMLPTest(parameterized.TestCase): """Tests for the per-tower feed-forward block.""" @@ -665,6 +680,7 @@ def test_rejects_unknown_activation(self): ) +@pytest.mark.cpu_only class WeaverMoTDecoderLayerStructureTest(parameterized.TestCase): """Tests for layer composition, aliasing, and shape contracts.""" @@ -869,6 +885,7 @@ def test_config_construction(self): self.assertEqual(gen_out.shape, (6, 64)) +@pytest.mark.cpu_only @pytest.mark.scheduled_only class WeaverDualAttentionGoldenParityTest(parameterized.TestCase): """Tests verifying exact numerical parity against golden outputs.""" @@ -1028,6 +1045,7 @@ def test_case_block_aligned_no_rope_flash(self): self._run_golden_parity_check("case_block_aligned_no_rope", kernel="flash") +@pytest.mark.cpu_only @pytest.mark.scheduled_only class WeaverMoTDecoderLayerGoldenParityTest(parameterized.TestCase): """Numerical parity against the reference Mixture-of-Transformers decoder layer.""" @@ -1198,5 +1216,429 @@ def test_golden_parity(self, case_name, precision): self._run_parity_check(case_name, precision) +def _copy_unscanned_to_scanned_layers( + unscanned_layers: nnx.List, + scanned_layers: WeaverMoTDecoderLayer, +) -> None: + """Stacks per-layer weights from unscanned ``nnx.List`` into ``scanned_layers``.""" + num_layers = len(unscanned_layers) + per_layer_flat = [dict(nnx.to_flat_state(nnx.split(unscanned_layers[i], nnx.Param, ...)[1])) for i in range(num_layers)] + + _, scanned_params, _ = nnx.split(scanned_layers, nnx.Param, ...) + for path, scanned_leaf in nnx.to_flat_state(scanned_params): + stacked_val = jnp.stack([per_layer_flat[i][path][...] for i in range(num_layers)], axis=0) + scanned_leaf[...] = stacked_val + + +@pytest.mark.cpu_only +class WeaverOmniTransformerUnitTest(parameterized.TestCase): + """CPU unit tests for WeaverOmniTransformer components and configuration.""" + + def test_patchify_and_unpatchify_roundtrip(self): + """Verifies patchify_latents and unpatchify_latents form an exact roundtrip.""" + bsz, c_lat, t_lat, h_lat, w_lat = 2, 48, 3, 6, 8 + patch_size = 2 + latents = jax.random.normal(jax.random.PRNGKey(0), (bsz, c_lat, t_lat, h_lat, w_lat)) + + patches = patchify_latents(latents, patch_size=patch_size) + expected_s_gen = t_lat * (h_lat // patch_size) * (w_lat // patch_size) + expected_patch_dim = patch_size * patch_size * c_lat + self.assertEqual(patches.shape, (bsz, expected_s_gen, expected_patch_dim)) + + reconstructed_3d = unpatchify_latents( + patches, + t_lat=t_lat, + h_lat=h_lat, + w_lat=w_lat, + patch_size=patch_size, + latent_channels=c_lat, + ) + self.assertEqual(reconstructed_3d.shape, (bsz, c_lat, t_lat, h_lat, w_lat)) + np.testing.assert_allclose(np.array(latents), np.array(reconstructed_3d), rtol=1e-6, atol=1e-6) + + # Also verify 2D [B * S_gen, patch_dim] input to unpatchify_latents. + patches_2d = patches.reshape(bsz * expected_s_gen, expected_patch_dim) + reconstructed_2d = unpatchify_latents( + patches_2d, + t_lat=t_lat, + h_lat=h_lat, + w_lat=w_lat, + patch_size=patch_size, + latent_channels=c_lat, + ) + np.testing.assert_allclose(np.array(latents), np.array(reconstructed_2d), rtol=1e-6, atol=1e-6) + + def test_timestep_embedding_and_time_embedder(self): + """Verifies sinusoidal timestep embedding and WeaverTimeEmbedder shapes and aliases.""" + timesteps = jnp.array([250.0, 750.0], dtype=jnp.float32) + sin_emb = get_weaver_timestep_embedding(timesteps * 0.001, num_channels=256) + self.assertEqual(sin_emb.shape, (2, 256)) + self.assertEqual(sin_emb.dtype, jnp.float32) + # At t=0, cos is 1 and sin is 0 -> first half is 1, second half is 0. + zero_emb = get_weaver_timestep_embedding(jnp.array([0.0], dtype=jnp.float32), num_channels=256) + np.testing.assert_allclose(np.array(zero_emb[0, :128]), np.ones(128), rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(np.array(zero_emb[0, 128:]), np.zeros(128), rtol=1e-6, atol=1e-6) + + embedder = WeaverTimeEmbedder(in_channels=256, time_embed_dim=128, rngs=nnx.Rngs(0)) + self.assertIs(embedder.linear_1, embedder.mlp_0) + self.assertIs(embedder.linear_2, embedder.mlp_2) + out = embedder(timesteps) + self.assertEqual(out.shape, (2, 128)) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + + def test_build_weaver_3d_position_ids(self): + """Verifies 3D M-RoPE coordinates for packed [text, vision] samples.""" + pos_3d = build_weaver_3d_position_ids(batch_size=2, s_und=4, t_lat=2, h_patch=2, w_patch=3) + s_gen = 2 * 2 * 3 + sample_len = 4 + s_gen + self.assertEqual(pos_3d.shape, (2 * sample_len, 3)) + + # Text positions for sample 0 are (i, i, i) for i in 0..3. + expected_text = np.stack([np.arange(4)] * 3, axis=-1) + np.testing.assert_array_equal(np.array(pos_3d[:4]), expected_text) + + # First vision token of sample 0 has (t=4, h=0, w=0); last has (t=5, h=1, w=2). + np.testing.assert_array_equal(np.array(pos_3d[4]), np.array([4, 0, 0])) + np.testing.assert_array_equal(np.array(pos_3d[sample_len - 1]), np.array([5, 1, 2])) + # Sample 1 repeats the same per-sample coordinate layout. + np.testing.assert_array_equal(np.array(pos_3d[:sample_len]), np.array(pos_3d[sample_len:])) + + def test_weaver_mini_diffuser_config_loading(self): + """Verifies weaver-mini-diffuser.yml loads and populates WeaverConfig.""" + cfg = pyconfig.initialize( + [ + None, + _BASE_CONFIG_PATH, + "model_name=weaver-mini-diffuser", + "skip_jax_distributed_system=True", + ] + ) + self.assertEqual(cfg.decoder_block, DecoderBlockType.WEAVER) + self.assertEqual(cfg.emb_dim, 4096) + self.assertEqual(cfg.mlp_dim, 12288) + self.assertEqual(cfg.num_query_heads, 32) + self.assertEqual(cfg.num_kv_heads, 8) + self.assertEqual(cfg.num_decoder_layers, 36) + self.assertEqual(cfg.head_dim, 128) + self.assertEqual(cfg.vocab_size, 151936) + + weaver_cfg = WeaverConfig.from_maxtext_config(cfg) + self.assertEqual(weaver_cfg.hidden_size, 4096) + self.assertEqual(weaver_cfg.intermediate_size, 12288) + self.assertEqual(weaver_cfg.num_attention_heads, 32) + self.assertEqual(weaver_cfg.num_key_value_heads, 8) + self.assertEqual(weaver_cfg.num_hidden_layers, 36) + self.assertEqual(weaver_cfg.head_dim, 128) + self.assertEqual(weaver_cfg.vocab_size, 151936) + self.assertEqual(weaver_cfg.latent_channels, 48) + self.assertEqual(weaver_cfg.patch_size, 2) + self.assertEqual(weaver_cfg.time_embed_in_channels, 256) + self.assertEqual(weaver_cfg.timestep_scale, 0.001) + self.assertEqual(weaver_cfg.timestep_max_period, 10000) + self.assertTrue(weaver_cfg.qk_norm_for_text) + self.assertTrue(weaver_cfg.qk_norm_for_diffusion) + self.assertFalse(weaver_cfg.use_und_k_norm_for_gen) + self.assertEqual(weaver_cfg.block_q, 128) + self.assertEqual(weaver_cfg.block_kv, 128) + self.assertEqual(weaver_cfg.mrope_section, (24, 20, 20)) + self.assertEqual(weaver_cfg.rope_theta, 5_000_000.0) + + def test_scan_layers_false_and_true_equivalence(self): + """Verifies scan_layers=False and scan_layers=True produce identical outputs.""" + base_cfg = WeaverConfig( + hidden_size=32, + head_dim=16, + num_attention_heads=2, + num_key_value_heads=1, + intermediate_size=64, + num_hidden_layers=2, + vocab_size=64, + latent_channels=48, + patch_size=2, + mrope_section=(4, 2, 2), + scan_layers=False, + ) + model_unscanned = WeaverOmniTransformer.from_config(base_cfg, rngs=nnx.Rngs(42)) + model_scanned = WeaverOmniTransformer.from_config( + base_cfg, + scan_layers=True, + rngs=nnx.Rngs(99), + ) + + # Copy top-level weights and per-layer weights from model_unscanned to model_scanned. + model_scanned.embed_tokens.embedding[...] = model_unscanned.embed_tokens.embedding[...] + model_scanned.proj_in.kernel[...] = model_unscanned.proj_in.kernel[...] + model_scanned.proj_in.bias[...] = model_unscanned.proj_in.bias[...] + model_scanned.time_embedder.mlp_0.kernel[...] = model_unscanned.time_embedder.mlp_0.kernel[...] + model_scanned.time_embedder.mlp_0.bias[...] = model_unscanned.time_embedder.mlp_0.bias[...] + model_scanned.time_embedder.mlp_2.kernel[...] = model_unscanned.time_embedder.mlp_2.kernel[...] + model_scanned.time_embedder.mlp_2.bias[...] = model_unscanned.time_embedder.mlp_2.bias[...] + model_scanned.norm.scale[...] = model_unscanned.norm.scale[...] + model_scanned.norm_moe_gen.scale[...] = model_unscanned.norm_moe_gen.scale[...] + model_scanned.proj_out.kernel[...] = model_unscanned.proj_out.kernel[...] + model_scanned.proj_out.bias[...] = model_unscanned.proj_out.bias[...] + + assert model_unscanned.layers is not None + assert model_scanned.scanned_layers is not None + _copy_unscanned_to_scanned_layers(model_unscanned.layers, model_scanned.scanned_layers) + + self.assertIs(model_unscanned.vae2llm, model_unscanned.proj_in) + self.assertIs(model_unscanned.llm2vae, model_unscanned.proj_out) + + input_ids = jax.random.randint(jax.random.PRNGKey(1), (2, 4), 0, 64, dtype=jnp.int32) + latents = jax.random.normal(jax.random.PRNGKey(2), (2, 48, 2, 4, 4)) + timesteps = jnp.array([200.0, 800.0], dtype=jnp.float32) + + out_unscanned, hidden_unscanned = model_unscanned(input_ids, latents, timesteps, return_hidden_states=True) + out_scanned, hidden_scanned = model_scanned(input_ids, latents, timesteps, return_hidden_states=True) + + self.assertEqual(out_unscanned.shape, (2, 48, 2, 4, 4)) + self.assertEqual(out_scanned.shape, (2, 48, 2, 4, 4)) + np.testing.assert_allclose(np.array(out_unscanned), np.array(out_scanned), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(np.array(hidden_unscanned), np.array(hidden_scanned), rtol=1e-5, atol=1e-5) + + +@pytest.mark.cpu_only +@pytest.mark.scheduled_only +class WeaverOmniTransformerGoldenParityTest(parameterized.TestCase): + """End-to-end numerical parity tests against reference GPU golden outputs.""" + + golden_data: Any = None + + @classmethod + def setUpClass(cls): + super().setUpClass() + golden_path = _get_golden_data_path(_TRANSFORMER_GOLDEN_FILENAME) + if golden_path is not None: + cls.golden_data = np.load(golden_path) + + def setUp(self): + super().setUp() + if self.golden_data is None: + self.skipTest( + f"Golden test asset {_TRANSFORMER_GOLDEN_FILENAME} not found locally under /tmp " + "and could not be downloaded from gs://maxtext-test-assets/." + ) + + def _populate_unscanned_weights( + self, + model: WeaverOmniTransformer, + data: dict[str, np.ndarray], + prefix: str, + dtype: Any, + ) -> None: + """Loads reference weights from the golden archive into an unscanned WeaverOmniTransformer.""" + + def arr(key: str) -> jax.Array: + return jnp.array(data[f"{prefix}_{key}"], dtype=dtype) + + model.embed_tokens.embedding[...] = arr("embed_tokens_weight") + model.proj_in.kernel[...] = arr("proj_in_weight").T + model.proj_in.bias[...] = arr("proj_in_bias") + model.time_embedder.mlp_0.kernel[...] = arr("time_embedder_mlp_0_weight").T + model.time_embedder.mlp_0.bias[...] = arr("time_embedder_mlp_0_bias") + model.time_embedder.mlp_2.kernel[...] = arr("time_embedder_mlp_2_weight").T + model.time_embedder.mlp_2.bias[...] = arr("time_embedder_mlp_2_bias") + model.norm.scale[...] = arr("norm_weight") + model.norm_moe_gen.scale[...] = arr("norm_moe_gen_weight") + model.proj_out.kernel[...] = arr("proj_out_weight").T + model.proj_out.bias[...] = arr("proj_out_bias") + + assert model.layers is not None + for l_idx, layer in enumerate(model.layers): + lp = f"layer_{l_idx}" + layer.input_layernorm.scale[...] = arr(f"{lp}_input_layernorm_weight") + layer.input_layernorm_moe_gen.scale[...] = arr(f"{lp}_input_layernorm_moe_gen_weight") + layer.post_attention_layernorm.scale[...] = arr(f"{lp}_post_attention_layernorm_weight") + layer.post_attention_layernorm_moe_gen.scale[...] = arr(f"{lp}_post_attention_layernorm_moe_gen_weight") + + layer.self_attn.q_proj.kernel[...] = arr(f"{lp}_self_attn_q_proj_weight").T + layer.self_attn.k_proj.kernel[...] = arr(f"{lp}_self_attn_k_proj_weight").T + layer.self_attn.v_proj.kernel[...] = arr(f"{lp}_self_attn_v_proj_weight").T + layer.self_attn.o_proj.kernel[...] = arr(f"{lp}_self_attn_o_proj_weight").T + + layer.self_attn.q_proj_gen.kernel[...] = arr(f"{lp}_self_attn_q_proj_moe_gen_weight").T + layer.self_attn.k_proj_gen.kernel[...] = arr(f"{lp}_self_attn_k_proj_moe_gen_weight").T + layer.self_attn.v_proj_gen.kernel[...] = arr(f"{lp}_self_attn_v_proj_moe_gen_weight").T + layer.self_attn.o_proj_gen.kernel[...] = arr(f"{lp}_self_attn_o_proj_moe_gen_weight").T + + if layer.self_attn.q_norm is not None and layer.self_attn.k_norm is not None: + layer.self_attn.q_norm.scale[...] = arr(f"{lp}_self_attn_q_norm_weight") + layer.self_attn.k_norm.scale[...] = arr(f"{lp}_self_attn_k_norm_weight") + if layer.self_attn.q_norm_gen is not None and layer.self_attn.k_norm_gen is not None: + layer.self_attn.q_norm_gen.scale[...] = arr(f"{lp}_self_attn_q_norm_moe_gen_weight") + layer.self_attn.k_norm_gen.scale[...] = arr(f"{lp}_self_attn_k_norm_moe_gen_weight") + if layer.self_attn.k_norm_und_for_gen is not None: + layer.self_attn.k_norm_und_for_gen.scale[...] = arr(f"{lp}_self_attn_k_norm_und_for_gen_weight") + + for mlp_mod, mlp_name in ((layer.mlp, "mlp"), (layer.mlp_moe_gen, "mlp_moe_gen")): + if mlp_mod.gate_proj is not None: + mlp_mod.gate_proj.kernel[...] = arr(f"{lp}_{mlp_name}_gate_proj_weight").T + mlp_mod.up_proj.kernel[...] = arr(f"{lp}_{mlp_name}_up_proj_weight").T + mlp_mod.down_proj.kernel[...] = arr(f"{lp}_{mlp_name}_down_proj_weight").T + + @parameterized.named_parameters( + ("weaver_mini_unscanned", "case0", "weaver_mini", False), + ("weaver_mini_scanned", "case0", "weaver_mini", True), + ("weaver_max_unscanned", "case1", "weaver_max", False), + ("weaver_max_scanned", "case1", "weaver_max", True), + ) + def test_e2e_golden_parity(self, prefix: str, variant: str, scan_layers: bool): + """Verifies E2E forward pass parity against the reference GPU implementation.""" + assert self.golden_data is not None + data = dict(self.golden_data) + dtype = jnp.bfloat16 + + is_mini = variant == "weaver_mini" + cfg = WeaverConfig( + hidden_size=128, + head_dim=32, + num_attention_heads=4, + num_key_value_heads=2, + intermediate_size=256, + num_hidden_layers=2, + vocab_size=256, + latent_channels=48, + patch_size=2, + time_embed_in_channels=256, + timestep_scale=0.001, + rms_norm_eps=1e-6, + hidden_act="silu" if is_mini else "relu2", + qk_norm_for_text=is_mini, + qk_norm_for_diffusion=True, + use_und_k_norm_for_gen=not is_mini, + mrope_section=(6, 5, 5), + rope_theta=5_000_000.0, + scan_layers=False, + dtype=dtype, + weight_dtype=dtype, + ) + + ref_unscanned = WeaverOmniTransformer(cfg, rngs=nnx.Rngs(0)) + self._populate_unscanned_weights(ref_unscanned, data, prefix, dtype) + + if scan_layers: + scanned_cfg = dataclasses.replace(cfg, scan_layers=True) + model = WeaverOmniTransformer(scanned_cfg, rngs=nnx.Rngs(1)) + model.embed_tokens.embedding[...] = ref_unscanned.embed_tokens.embedding[...] + model.proj_in.kernel[...] = ref_unscanned.proj_in.kernel[...] + model.proj_in.bias[...] = ref_unscanned.proj_in.bias[...] + model.time_embedder.mlp_0.kernel[...] = ref_unscanned.time_embedder.mlp_0.kernel[...] + model.time_embedder.mlp_0.bias[...] = ref_unscanned.time_embedder.mlp_0.bias[...] + model.time_embedder.mlp_2.kernel[...] = ref_unscanned.time_embedder.mlp_2.kernel[...] + model.time_embedder.mlp_2.bias[...] = ref_unscanned.time_embedder.mlp_2.bias[...] + model.norm.scale[...] = ref_unscanned.norm.scale[...] + model.norm_moe_gen.scale[...] = ref_unscanned.norm_moe_gen.scale[...] + model.proj_out.kernel[...] = ref_unscanned.proj_out.kernel[...] + model.proj_out.bias[...] = ref_unscanned.proj_out.bias[...] + assert ref_unscanned.layers is not None + assert model.scanned_layers is not None + _copy_unscanned_to_scanned_layers(ref_unscanned.layers, model.scanned_layers) + else: + model = ref_unscanned + + input_ids = jnp.array(data[f"{prefix}_input_ids"], dtype=jnp.int32) + latents = jnp.array(data[f"{prefix}_latents"], dtype=dtype) + timesteps = jnp.array(data[f"{prefix}_timesteps"], dtype=jnp.float32) + golden_pos_ids = data[f"{prefix}_position_ids"] + golden_preds_vision = data[f"{prefix}_preds_vision"] + golden_last_hidden = data[f"{prefix}_last_hidden_state"] + + # Verify automatic 3D M-RoPE position ID builder matches the reference position_ids. + auto_pos_3d = build_weaver_3d_position_ids(batch_size=2, s_und=6, t_lat=2, h_patch=2, w_patch=2) + np.testing.assert_array_equal(np.array(auto_pos_3d), golden_pos_ids.T) + + preds_vision, last_hidden = model( + input_ids, + latents, + timesteps, + return_hidden_states=True, + ) + + self.assertEqual(preds_vision.shape, (2, 48, 2, 4, 4)) + self.assertEqual(last_hidden.shape, (28, 128)) + + preds_np = np.array(preds_vision, dtype=np.float32) + hidden_np = np.array(last_hidden, dtype=np.float32) + + preds_avg_abs = float(np.mean(np.abs(preds_np - golden_preds_vision))) + hidden_avg_abs = float(np.mean(np.abs(hidden_np - golden_last_hidden))) + print( + f"\n[WeaverOmniTransformer E2E Parity: {variant} | scan_layers={scan_layers}]" + f" preds_max_abs={np.max(np.abs(preds_np - golden_preds_vision)):.4e}," + f" preds_avg_abs={preds_avg_abs:.4e}," + f" hidden_max_abs={np.max(np.abs(hidden_np - golden_last_hidden)):.4e}," + f" hidden_avg_abs={hidden_avg_abs:.4e}" + ) + + self.assertLess(hidden_avg_abs, 1.2e-2) + self.assertLess(preds_avg_abs, 1.2e-2) + np.testing.assert_allclose( + hidden_np, + golden_last_hidden, + rtol=3e-2, + atol=5.5e-2, + err_msg=f"last_hidden_state mismatch for {variant} (scan_layers={scan_layers})", + ) + np.testing.assert_allclose( + preds_np, + golden_preds_vision, + rtol=3e-2, + atol=5.5e-2, + err_msg=f"preds_vision mismatch for {variant} (scan_layers={scan_layers})", + ) + + +@pytest.mark.tpu_only +@pytest.mark.scheduled_only +class WeaverOmniTransformerTPUTest(parameterized.TestCase): + """Scheduled nightly TPU v5 test for full 36-layer WeaverOmniTransformer.""" + + @parameterized.named_parameters( + ("unscanned", False), + ("scanned", True), + ) + def test_full_36_layer_backbone_forward_tpu(self, scan_layers: bool): + """Executes full 36-layer WeaverOmniTransformer on TPU with random weights.""" + cfg = pyconfig.initialize( + [ + None, + _BASE_CONFIG_PATH, + "model_name=weaver-mini-diffuser", + f"scan_layers={scan_layers}", + "dtype=bfloat16", + "weight_dtype=bfloat16", + "skip_jax_distributed_system=True", + ] + ) + model = WeaverOmniTransformer.from_config(cfg, rngs=nnx.Rngs(0)) + self.assertEqual(model.num_hidden_layers, 36) + self.assertEqual(model.hidden_size, 4096) + self.assertEqual(model.latent_channels, 48) + + batch_size = 2 + s_und = 8 + t_lat, h_lat, w_lat = 2, 4, 4 + + input_ids = jax.random.randint(jax.random.PRNGKey(10), (batch_size, s_und), 0, model.vocab_size, dtype=jnp.int32) + latents = jax.random.normal( + jax.random.PRNGKey(11), + (batch_size, 48, t_lat, h_lat, w_lat), + dtype=jnp.bfloat16, + ) + timesteps = jnp.array([250.0, 750.0], dtype=jnp.float32) + + graphdef, state = nnx.split(model) + + @jax.jit + def forward(model_state, ids, lats, ts): + merged = nnx.merge(graphdef, model_state) + return merged(ids, lats, ts) + + output = forward(state, input_ids, latents, timesteps) + self.assertEqual(output.shape, (batch_size, 48, t_lat, h_lat, w_lat)) + self.assertTrue(bool(jnp.all(jnp.isfinite(output)))) + + if __name__ == "__main__": absltest.main()