diff --git a/docs/guides/optimization/custom_mesh_and_rule.md b/docs/guides/optimization/custom_mesh_and_rule.md index fbb8f3adc2..4d30523ba7 100644 --- a/docs/guides/optimization/custom_mesh_and_rule.md +++ b/docs/guides/optimization/custom_mesh_and_rule.md @@ -51,6 +51,12 @@ This rule utilizes the `data`, `stage`, `fsdp`, and `expert` axes. Its defining Similar in philosophy to `ep-as-cp.yml`, this configuration explicitly includes the `context` axis in the mesh layout alongside `data`, `stage`, `fsdp`, and `expert`. While context sharding is mapped to the `context` axis globally, within MoE components, this `context` axis dynamically shifts to perform expert parallelism instead of FSDP. This custom rule supports using CP and EP together. +### `cp-as-ep-dense-no-cp.yml` and `cp-as-ep-dense-fsdp.yml` + +Variants of `cp-as-ep.yml` that take the dense weights (`embed`, `embed_attn`, `q_lora`, `kv_lora`: attention and shared-expert projections, the MoE router and the shared-expert gate) off the `context` axis. Under `cp-as-ep` they are sharded over `fsdp` x `context` x `expert`, so every all-gather of them spans all devices and also synchronizes the context ranks, which do unequal attention work under causal context parallelism. `cp-as-ep-dense-no-cp.yml` shards them over `fsdp` x `expert` and `cp-as-ep-dense-fsdp.yml` over `fsdp` only, so those all-gathers run in independent groups of `fsdp` x `expert` (or `fsdp`) devices, at the cost of replicating the weights and their gradients `context`-way (`context` x `expert`-way). + +Both rules also map `embed_grad` to the token axes. Left to GSPMD, the gradient of a weight replicated over `context` (a partial sum over every token axis) is all-reduced at full size over all devices and then sliced. With `embed_grad` mapped, `linears.route_dense_weight_grad` reduce-scatters it over all devices instead, as under `cp-as-ep`, and then all-gathers the shard over the replicated axes. The GatedDeltaNet projections (`embed_gdn`), activations, the vocabulary weights and the routed experts are sharded exactly as in `cp-as-ep.yml`, and `RoutedMoE` keeps `("context", "expert")` as its expert-parallel axes. + ### `ep-as-dp.yml` Different with the rule in `base.yml`, this rule configures expert physical axis to function as data parallelism rather than FSDP. This removes the constraint where FSDPxEP is limited by specific model dimensions, particularly for small tensors such as attention projections. Ultimately, this change benefits large-scale training. diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_deepseek_family_ckpt.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_deepseek_family_ckpt.py index 6e2a596c17..307ec84c0b 100644 --- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_deepseek_family_ckpt.py +++ b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_deepseek_family_ckpt.py @@ -601,8 +601,6 @@ def _normalize(raw_key): moe["shared_experts"]["wo"]["kernel"][layer_idx, ...] = shared_wo # re-order - if q_lora_rank != 0: - moe["MoeBlock_0"]["gate"]["bias"] = np.transpose(moe["MoeBlock_0"]["gate"]["bias"], axes=(1, 0)) moe["MoeBlock_0"]["gate"]["kernel"] = np.transpose(moe["MoeBlock_0"]["gate"]["kernel"], axes=(1, 0, 2)) moe["shared_experts"]["wi_0"]["kernel"] = np.transpose(moe["shared_experts"]["wi_0"]["kernel"], axes=(1, 0, 2)) moe["shared_experts"]["wi_1"]["kernel"] = np.transpose(moe["shared_experts"]["wi_1"]["kernel"], axes=(1, 0, 2)) diff --git a/src/maxtext/checkpoint_conversion/to_huggingface.py b/src/maxtext/checkpoint_conversion/to_huggingface.py index f3b28bb90a..cf98b4dc3d 100644 --- a/src/maxtext/checkpoint_conversion/to_huggingface.py +++ b/src/maxtext/checkpoint_conversion/to_huggingface.py @@ -238,6 +238,12 @@ def _validate_or_update_architecture(hf_config, max_config, override: bool): ("vocab_size", "vocab_size"), ("global_head_dim", "global_head_dim"), ("num_global_key_value_heads", "global_num_kv_heads"), + ("hc_mult", "mhc_expansion_rate"), + ("num_hash_layers", "first_num_hash_layers"), + ("index_n_heads", "indexer_n_heads"), + ("index_head_dim", "indexer_head_dim"), + ("o_lora_rank", "o_lora_rank"), + ("o_groups", "o_groups"), ] if max_config.attention_type == "mla": diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index eb2b134466..cf670c927c 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -59,20 +59,18 @@ from typing import Any, Callable, List, Sequence import absl import ml_dtypes -import flax.linen as nn +from flax import nnx +from flax.nnx import variablelib from huggingface_hub import hf_hub_download, list_repo_files import jax from maxtext.configs import pyconfig from maxtext.configs.types import DType -from maxtext.common.common_types import MODEL_MODE_TRAIN from maxtext.checkpoint_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS from maxtext.checkpoint_conversion.utils.param_mapping import HOOK_FNS, PARAM_MAPPING from maxtext.checkpoint_conversion.utils.tensor_handling import apply_hook_fns, nesting_depth, slice_shape, stacked_axes from maxtext.checkpoint_conversion.utils.utils import MemoryMonitorTqdm, load_hf_dict_from_transformers, load_hf_dict_from_safetensors, param_key_parts_from_path, print_peak_memory, print_ram_usage, save_weights_to_checkpoint, validate_and_filter_param_map_keys from maxtext.inference.inference_utils import str2bool -from maxtext.layers import quantizations -from maxtext.models import models -from maxtext.utils import max_logging, max_utils, maxtext_utils +from maxtext.utils import max_logging, max_utils, maxtext_utils, model_creation_utils from maxtext.utils.globals import HF_IDS import numpy as np from orbax.checkpoint import type_handlers @@ -389,31 +387,30 @@ def get_maxtext_model_info(config): mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) max_logging.log("Initializing MaxText abstract model...") - quant = quantizations.configure_quantization(config) - maxtext_model_flax = models.transformer_as_linen(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - - # Get abstract model structure (name, shape) without materializing the weights to save memory. - # Extract the 'params' collection from the abstract model state. This focuses checkpoint - # conversion on trainable model parameters; variables outside the 'params' collection - # (such as non-trainable state or optimizer buffers) are not included. - abstract_params_tree = maxtext_utils.get_abstract_param(maxtext_model_flax, config)["params"] - - abstract_params_flat, abstract_params_treedef = jax.tree_util.tree_flatten_with_path( - abstract_params_tree, - is_leaf=lambda x: isinstance(x, nn.LogicallyPartitioned), - ) + # Get abstract model structure (name, shape) without materializing the weights to save memory + _, abstract_model = model_creation_utils.create_nnx_abstract_model(config, mesh) + + # Lay the variables out the way the checkpoint stores them: one collection per variable type + # (e.g. 'params', 'Tid2EidVar'), then the module path with Linen-style names ("layers_0"). + abstract_params_tree = {} + for path, variable in nnx.to_flat_state(nnx.state(abstract_model)): + if isinstance(variable, nnx.RngState): + continue + node = abstract_params_tree.setdefault(variablelib.variable_name_from_type(type(variable), allow_register=True), {}) + *parents, name = param_key_parts_from_path(path) + for part in parents: + node = node.setdefault(part, {}) + node[name] = variable.get_value() + + abstract_params_flat, abstract_params_treedef = jax.tree_util.tree_flatten_with_path(abstract_params_tree) max_logging.log("MaxText abstract model and state initialized.") # preprocess state maxtext_abstract_dict = {} for mt_target_idx, (path_tuple, abstract_leaf_value) in enumerate(abstract_params_flat): - mt_param_key = "params-" + "-".join(param_key_parts_from_path(path_tuple)) - if isinstance(abstract_leaf_value, nn.LogicallyPartitioned): - mt_target_shape = abstract_leaf_value.value.shape - else: - mt_target_shape = abstract_leaf_value.shape - maxtext_abstract_dict[mt_param_key] = (mt_target_idx, mt_target_shape) + mt_param_key = "-".join(param_key_parts_from_path(path_tuple)) + maxtext_abstract_dict[mt_param_key] = (mt_target_idx, abstract_leaf_value.shape) return maxtext_abstract_dict, abstract_params_treedef @@ -481,6 +478,7 @@ def _build_single_axis_stacked_tensor( hook_fns: Any, target_shape: tuple, config, + mt_key: str = "", ) -> np.ndarray: """Builds a MaxText tensor by stacking HF weights along a single axis. @@ -499,12 +497,16 @@ def _build_single_axis_stacked_tensor( """ tensors_to_stack = [] + axis_to_stack = 0 if config.scan_layers: - # If the target tensor rank exceeds param_scan_axis (e.g., multidimensional weights or 2D block scales), - # stack along param_scan_axis (typically axis 1 in scanned layers). For 1D tensors (e.g., per-layer scalar - # quantization scales like Llama 3.1 FP8 `kernel_scale` with shape (num_layers,)), stack along axis 0 - # because axis 1 does not exist on 1D tensors and stacking scalars along axis 0 produces the 1D vector. - axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0 + is_var_collection = "MoEBiasVar" in mt_key or "Tid2EidVar" in mt_key + is_unscanned_moe = "MoeBlock" in mt_key and "scanned_blocks" not in mt_key + if not (is_var_collection or is_unscanned_moe): + # If the target tensor rank exceeds param_scan_axis (e.g., multidimensional weights or 2D block scales), + # stack along param_scan_axis (typically axis 1 in scanned layers). For 1D tensors (e.g., per-layer scalar + # quantization scales like Llama 3.1 FP8 `kernel_scale` with shape (num_layers,)), stack along axis 0 + # because axis 1 does not exist on 1D tensors and stacking scalars along axis 0 produces the 1D vector. + axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0 else: # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0). axis_to_stack = 0 @@ -545,6 +547,8 @@ def _get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_t if not isinstance(hf_source_keys_or_key, list): # Case 1: Single hf key (str) def _loader(getter, key, shape, hook): + if key is None: + return apply_hook_fns(None, shape, hook) if isinstance(key, (list, tuple)): tensors = tuple(getter(k) for k in key) return apply_hook_fns(tensors, shape, hook) @@ -567,6 +571,7 @@ def _loader(getter, key, shape, hook): hook_fn, mt_target_shape_or_shapes, config, + mt_key, ) else: # isinstance(hf_source_keys_or_key[0], list) @@ -1022,6 +1027,7 @@ def main( jax_weights = convert_lora_to_maxtext_adapter(config, lora_weights, save_dtype) adapter_name = os.path.basename(os.path.normpath(hf_lora_adapter_path)) output_directory = os.path.join(output_directory, model_name_for_path, adapter_name) + state_params = jax_weights else: if lazy_load_tensors and config.use_multimodal: @@ -1093,6 +1099,8 @@ def main( } def _eager_getter(key): + if key is None: + return None resolved_key = resolve_scale_key(key, hf_state_dict_numpy) if resolved_key not in hf_state_dict_numpy: raise ValueError(f"HuggingFace key {key} not found in state_dict.") @@ -1180,11 +1188,19 @@ def _eager_getter(key): max_logging.log(f"Elapse for transform: {(time.time() - start) / 60:.2f} min") print_ram_usage("Before creating full JAX tree") - # Create final MaxText parameters tree + max_logging.log(f"Length of final_mt_weights: {len(final_mt_weights)}") + max_logging.log(f"Treedef num_leaves: {abstract_params_treedef.num_leaves}") + # Create final MaxText parameters tree containing all collections jax_weights = jax.tree_util.tree_unflatten(abstract_params_treedef, final_mt_weights) del final_mt_weights, abstract_params_treedef + state_params = jax_weights + print_ram_usage("Before saving") + leaves = jax.tree_util.tree_leaves(state_params) + max_logging.log(f"Length of state_params leaves: {len(leaves)}") + if leaves: + max_logging.log(f"Type of first leaf before saving: {type(leaves[0])}") if lazy_load_tensors and not is_adapter_only: max_logging.log("Starting checkpoint save (loading weights just-in-time)...") else: @@ -1195,7 +1211,7 @@ def _eager_getter(key): # and sharded across virtual devices. save_weights_to_checkpoint( output_directory, - jax_weights, + state_params, simulated_cpu_devices_count, config.checkpoint_storage_use_ocdbt, config.checkpoint_storage_use_zarr3, diff --git a/src/maxtext/checkpoint_conversion/to_maxtext_diffusers.py b/src/maxtext/checkpoint_conversion/to_maxtext_diffusers.py new file mode 100644 index 0000000000..2529d65f9d --- /dev/null +++ b/src/maxtext/checkpoint_conversion/to_maxtext_diffusers.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Universal Hugging Face Diffusers checkpoint converter and manager.""" diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index cbe3095cc7..d1a9865cf8 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -1608,6 +1608,86 @@ def __init__(self, **kwargs): qwen3_5_35b_a3b_config = PTConfig(**qwen3_5_35b_a3b_dict) # pytype: disable=wrong-arg-types qwen3_5_397b_a17b_config = PTConfig(**qwen3_5_397b_a17b_dict) # pytype: disable=wrong-arg-types +qwen3_5_9b_dict = { + "architectures": ["Qwen3_5ForConditionalGeneration"], + "image_token_id": 248056, + "model_type": "qwen3_5", + "text_config": { + "attention_bias": False, + "attention_dropout": 0.0, + "attn_output_gate": True, + "bos_token_id": None, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 12288, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + ] + * 8, + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_value_head_dim": 128, + "mamba_ssm_dtype": "float32", + "max_position_embeddings": 262144, + "mlp_only_layers": [], + "model_type": "qwen3_5_text", + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": False, + "num_attention_heads": 16, + "num_hidden_layers": 32, + "num_key_value_heads": 4, + "pad_token_id": None, + "partial_rotary_factor": 0.25, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + "rope_type": "default", + }, + "tie_word_embeddings": False, + "use_cache": True, + "vocab_size": 248320, + }, + "tie_word_embeddings": False, + "transformers_version": "5.16.0.dev0", + "video_token_id": 248057, + "vision_config": { + "deepstack_visual_indexes": [], + "depth": 27, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1152, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4304, + "model_type": "qwen3_5_vision", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 4096, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + }, + "vision_end_token_id": 248054, + "vision_start_token_id": 248053, +} + +try: + qwen3_5_9b_config = transformers.Qwen3_5Config(**qwen3_5_9b_dict) # pyrefly: ignore[missing-attribute] +except AttributeError: + qwen3_5_9b_config = PTConfig(**qwen3_5_9b_dict) # pytype: disable=wrong-arg-types + # from https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1/blob/main/config.json mixtral_8x7b_dict = { @@ -1883,9 +1963,9 @@ def __init__(self, **kwargs): } qwen3_vl_30b_a3b_config = PTConfig(**qwen3_vl_30b_a3b_dict) -cosmos3_nano_reasoner_dict = { - "architectures": ["Cosmos3OmniForConditionalGeneration"], - "model_type": "cosmos3_omni", +weaver_mini_dict = { + "architectures": ["WeaverOmniForConditionalGeneration"], + "model_type": "weaver_omni", "text_config": { "attention_bias": False, "attention_dropout": 0.0, @@ -1931,11 +2011,11 @@ def __init__(self, **kwargs): "temporal_patch_size": 2, }, } -cosmos3_nano_reasoner_config = PTConfig(**cosmos3_nano_reasoner_dict) +weaver_mini_config = PTConfig(**weaver_mini_dict) -cosmos3_super_reasoner_dict = { - "architectures": ["Cosmos3ForConditionalGeneration"], - "model_type": "cosmos3_omni", +weaver_max_dict = { + "architectures": ["WeaverForConditionalGeneration"], + "model_type": "weaver_omni", "text_config": { "attention_bias": False, "attention_dropout": 0.0, @@ -1981,7 +2061,7 @@ def __init__(self, **kwargs): "temporal_patch_size": 2, }, } -cosmos3_super_reasoner_config = PTConfig(**cosmos3_super_reasoner_dict) +weaver_max_config = PTConfig(**weaver_max_dict) # {maxtext model name: hf model config} @@ -2013,8 +2093,8 @@ def __init__(self, **kwargs): "qwen3-vl-2b": qwen3_vl_2b_config, "qwen3-vl-4b": qwen3_vl_4b_config, "qwen3-vl-30b-a3b": qwen3_vl_30b_a3b_config, - "cosmos3-nano-reasoner": cosmos3_nano_reasoner_config, - "cosmos3-super-reasoner": cosmos3_super_reasoner_config, + "weaver-mini": weaver_mini_config, + "weaver-max": weaver_max_config, "llama3.1-8b": llama31_8b_config, "llama3.1-8b-Instruct": llama31_8b_config, "llama3.1-70b": llama31_70b_config, @@ -2036,6 +2116,7 @@ def __init__(self, **kwargs): "qwen3.5-35b-a3b": qwen3_5_35b_a3b_config, "qwen3.5-35b-a3b-fp8": qwen3_5_35b_a3b_config, "qwen3.5-35b-fp8": qwen3_5_35b_a3b_config, + "qwen3.5-9b": qwen3_5_9b_config, "mixtral-8x7b": mixtral_8x7b_config, "mixtral-8x22b": mixtral_8x22b_config, "olmo3-7b": olmo3_7b_config, diff --git a/src/maxtext/checkpoint_conversion/utils/hf_shape.py b/src/maxtext/checkpoint_conversion/utils/hf_shape.py index c80c643d24..b74939cba3 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_shape.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_shape.py @@ -704,15 +704,16 @@ def QWEN3_5_HF_WEIGHTS_TO_SHAPE(config): vocab_size = config["vocab_size"] num_attention_heads = config["num_attention_heads"] num_key_value_heads = config["num_key_value_heads"] - num_experts = config["num_experts"] + num_experts = config.get("num_experts", 0) or 0 head_dim = config["head_dim"] linear_conv_kernel_dim = config["linear_conv_kernel_dim"] linear_key_head_dim = config["linear_key_head_dim"] linear_value_head_dim = config["linear_value_head_dim"] linear_num_key_heads = config["linear_num_key_heads"] linear_num_value_heads = config["linear_num_value_heads"] - moe_intermediate_size = config["moe_intermediate_size"] - shared_expert_intermediate_size = config["shared_expert_intermediate_size"] + moe_intermediate_size = config.get("moe_intermediate_size", 0) + shared_expert_intermediate_size = config.get("shared_expert_intermediate_size", 0) + intermediate_size = config.get("intermediate_size", 0) cycle_interval = config["full_attention_interval"] # --- Calculated Values --- @@ -776,27 +777,36 @@ def QWEN3_5_HF_WEIGHTS_TO_SHAPE(config): } ) - # --- MLP Logic (MoE + Shared) --- - mapping.update( - { - # Router - f"{layer_prefix}.mlp.gate.weight": [num_experts, hidden_size], - # Shared Experts (SwiGLU - Separate Weights) - f"{layer_prefix}.mlp.shared_expert.gate_proj.weight": [shared_expert_intermediate_size, hidden_size], - f"{layer_prefix}.mlp.shared_expert.up_proj.weight": [shared_expert_intermediate_size, hidden_size], - f"{layer_prefix}.mlp.shared_expert.down_proj.weight": [hidden_size, shared_expert_intermediate_size], - # Shared Expert Gate (learned scaling factor) - f"{layer_prefix}.mlp.shared_expert_gate.weight": [1, hidden_size], - } - ) + if num_experts > 1: + # --- MLP Logic (MoE + Shared) --- + mapping.update( + { + # Router + f"{layer_prefix}.mlp.gate.weight": [num_experts, hidden_size], + # Shared Experts (SwiGLU - Separate Weights) + f"{layer_prefix}.mlp.shared_expert.gate_proj.weight": [shared_expert_intermediate_size, hidden_size], + f"{layer_prefix}.mlp.shared_expert.up_proj.weight": [shared_expert_intermediate_size, hidden_size], + f"{layer_prefix}.mlp.shared_expert.down_proj.weight": [hidden_size, shared_expert_intermediate_size], + # Shared Expert Gate (learned scaling factor) + f"{layer_prefix}.mlp.shared_expert_gate.weight": [1, hidden_size], + } + ) - # --- Vectorized & Fused Routed Experts (No loop, no .weight suffix) --- - mapping.update( - { - f"{layer_prefix}.mlp.experts.gate_up_proj": [num_experts, 2 * moe_intermediate_size, hidden_size], - f"{layer_prefix}.mlp.experts.down_proj": [num_experts, hidden_size, moe_intermediate_size], - } - ) + # --- Vectorized & Fused Routed Experts (No loop, no .weight suffix) --- + mapping.update( + { + f"{layer_prefix}.mlp.experts.gate_up_proj": [num_experts, 2 * moe_intermediate_size, hidden_size], + f"{layer_prefix}.mlp.experts.down_proj": [num_experts, hidden_size, moe_intermediate_size], + } + ) + else: + mapping.update( + { + f"{layer_prefix}.mlp.gate_proj.weight": [intermediate_size, hidden_size], + f"{layer_prefix}.mlp.up_proj.weight": [intermediate_size, hidden_size], + f"{layer_prefix}.mlp.down_proj.weight": [hidden_size, intermediate_size], + } + ) return mapping @@ -1174,106 +1184,137 @@ def QWEN3_VL_HF_WEIGHTS_TO_SHAPE(config): def DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE(config): - """Returns a dictionary mapping HuggingFace weight names to shapes for DeepSeek V4.""" + """Returns the HuggingFace parameter shapes for DeepSeek V4.""" hidden_size = config["hidden_size"] vocab_size = config["vocab_size"] num_hidden_layers = config["num_hidden_layers"] q_lora_rank = config.get("q_lora_rank", 1024) - kv_lora_rank = config.get("kv_lora_rank", 512) - # MaxText scales o_lora_rank by o_groups to get o_a_out_features (8192) o_lora_rank = config.get("o_lora_rank", 1024) * config.get("o_groups", 8) num_attention_heads = config["num_attention_heads"] head_dim = config["head_dim"] - moe_intermediate_size = config["moe_intermediate_size"] n_routed_experts = config["n_routed_experts"] - - # Dynamic hyperparameters replacing hardcoded literals: - hc_mult = config.get("hc_mult", 4) - hc_dim = (2 * hc_mult) + (hc_mult**2) # 2(4) + 16 = 24 - num_hash_layers = config.get("num_hash_layers", 3) # 3 - index_n_heads = config.get("index_n_heads", 64) # 64 - index_head_dim = config.get("index_head_dim", 128) # 128 + hc_mult = config.get("hc_mult", config.get("mhc_expansion_rate", 4)) + hc_dim = (2 * hc_mult) + (hc_mult**2) + num_hash_layers = config.get("num_hash_layers", config.get("first_num_hash_layers", 3)) + index_n_heads = config.get("index_n_heads", config.get("indexer_n_heads", 64)) + index_head_dim = config.get("index_head_dim", config.get("indexer_head_dim", 128)) mapping = { - "model.embed_tokens.weight": [vocab_size, hidden_size], - "model.norm.weight": [hidden_size], + "embed.weight": [vocab_size, hidden_size], + "norm.weight": [hidden_size], "head.weight": [vocab_size, hidden_size], - "model.hc_head.hc_fn": [hc_mult, hidden_size * hc_mult], - "model.hc_head.hc_base": [hc_mult], - "model.hc_head.hc_scale": [1], + "hc_head_fn": [hc_mult, hidden_size * hc_mult], + "hc_head_base": [hc_mult], + "hc_head_scale": [1], } - for layer_idx in range(num_hidden_layers): - layer_prefix = f"model.layers.{layer_idx}" - layer_mapping = { - f"{layer_prefix}.input_layernorm.weight": [hidden_size], - f"{layer_prefix}.post_attention_layernorm.weight": [hidden_size], - f"{layer_prefix}.self_attn.q_a_proj.weight": [q_lora_rank, hidden_size], - f"{layer_prefix}.self_attn.q_a_norm.weight": [q_lora_rank], - f"{layer_prefix}.self_attn.q_b_proj.weight": [num_attention_heads * head_dim, q_lora_rank], - f"{layer_prefix}.self_attn.kv_proj.weight": [kv_lora_rank, hidden_size], - f"{layer_prefix}.self_attn.kv_norm.weight": [kv_lora_rank], - f"{layer_prefix}.self_attn.sinks": [num_attention_heads], - f"{layer_prefix}.self_attn.o_a_proj.weight": [o_lora_rank, hidden_size], - f"{layer_prefix}.self_attn.o_b_proj.weight": [hidden_size, o_lora_rank], - # MHC - f"{layer_prefix}.attn_hc.fn": [hc_dim, hidden_size * hc_mult], - f"{layer_prefix}.attn_hc.base": [hc_dim], - f"{layer_prefix}.attn_hc.scale": [num_hash_layers], - f"{layer_prefix}.ffn_hc.fn": [hc_dim, hidden_size * hc_mult], - f"{layer_prefix}.ffn_hc.base": [hc_dim], - f"{layer_prefix}.ffn_hc.scale": [num_hash_layers], - # MLP / MoE Block - f"{layer_prefix}.mlp.gate.weight": [n_routed_experts, hidden_size], - f"{layer_prefix}.mlp.gate.e_score_correction_bias": [n_routed_experts], - f"{layer_prefix}.mlp.shared_experts.gate_proj.weight": [moe_intermediate_size, hidden_size], - f"{layer_prefix}.mlp.shared_experts.up_proj.weight": [moe_intermediate_size, hidden_size], - f"{layer_prefix}.mlp.shared_experts.down_proj.weight": [hidden_size, moe_intermediate_size], + def _compressor_shapes(comp_prefix, comp_dim, norm_dim, r): + return { + f"{comp_prefix}.wgate.weight": [comp_dim, hidden_size], + f"{comp_prefix}.wkv.weight": [comp_dim, hidden_size], + f"{comp_prefix}.norm.weight": [norm_dim], + f"{comp_prefix}.ape": [r, comp_dim], } - # Experts - for e in range(n_routed_experts): - layer_mapping[f"{layer_prefix}.mlp.experts.{e}.w1.weight"] = [moe_intermediate_size, hidden_size] - layer_mapping[f"{layer_prefix}.mlp.experts.{e}.w3.weight"] = [moe_intermediate_size, hidden_size] - layer_mapping[f"{layer_prefix}.mlp.experts.{e}.w2.weight"] = [hidden_size, moe_intermediate_size] + def _layer_shapes(layer_prefix, is_hash, ratio): + """Returns the HF shapes for one DeepSeek-V4 transformer layer under `layer_prefix`. - # Compressors - if layer_idx >= 2: - c_type = "csa" if (layer_idx % 2 == 0) else "hca" - if c_type == "csa": + Args: + layer_prefix: HF key prefix, e.g. `layers.3` or `mtp.0`. + is_hash: Whether the layer uses hash routing (emits `tid2eid` instead of a gate bias). + ratio: Compression ratio; 0 means the layer has no compressor weights. + """ + layer_mapping = { + f"{layer_prefix}.attn_norm.weight": [hidden_size], + f"{layer_prefix}.ffn_norm.weight": [hidden_size], + f"{layer_prefix}.attn.wq_a.weight": [q_lora_rank, hidden_size], + f"{layer_prefix}.attn.q_norm.weight": [q_lora_rank], + f"{layer_prefix}.attn.wq_b.weight": [num_attention_heads * head_dim, q_lora_rank], + f"{layer_prefix}.attn.wkv.weight": [ + config.get("num_key_value_heads", 1) * config.get("head_dim", 64), + hidden_size, + ], + f"{layer_prefix}.attn.kv_norm.weight": [head_dim], + f"{layer_prefix}.attn.attn_sink": [num_attention_heads], + f"{layer_prefix}.attn.wo_a.weight": [ + o_lora_rank, + (num_attention_heads * head_dim) // config.get("o_groups", 8), + ], + f"{layer_prefix}.attn.wo_b.weight": [hidden_size, o_lora_rank], + f"{layer_prefix}.hc_attn_fn": [hc_dim, hidden_size * hc_mult], + f"{layer_prefix}.hc_attn_base": [hc_dim], + f"{layer_prefix}.hc_attn_scale": [3], + f"{layer_prefix}.hc_ffn_fn": [hc_dim, hidden_size * hc_mult], + f"{layer_prefix}.hc_ffn_base": [hc_dim], + f"{layer_prefix}.hc_ffn_scale": [3], + f"{layer_prefix}.ffn.gate.weight": [n_routed_experts, hidden_size], + f"{layer_prefix}.ffn.gate.bias": [n_routed_experts], + } + if is_hash: + layer_mapping[f"{layer_prefix}.ffn.gate.tid2eid"] = [vocab_size, config.get("num_experts_per_tok", 2)] + + if ratio > 0: + comp_dim = 2 * head_dim if ratio == 4 else head_dim + layer_mapping.update(_compressor_shapes(f"{layer_prefix}.attn.compressor", comp_dim, head_dim, ratio)) + if ratio == 4: layer_mapping.update( - { - f"{layer_prefix}.self_attn.compressor.kv_proj.weight": [1024, hidden_size], - f"{layer_prefix}.self_attn.compressor.gate_proj.weight": [1024, hidden_size], - f"{layer_prefix}.self_attn.compressor.position_bias": [4, 1024], - f"{layer_prefix}.self_attn.compressor.kv_norm.weight": [512], - } + _compressor_shapes(f"{layer_prefix}.attn.indexer.compressor", 2 * index_head_dim, index_head_dim, ratio) ) - layer_mapping.update( - { - f"{layer_prefix}.self_attn.compressor.indexer.gate_proj.weight": [256, hidden_size], - f"{layer_prefix}.self_attn.compressor.indexer.kv_proj.weight": [256, hidden_size], - f"{layer_prefix}.self_attn.compressor.indexer.q_b_proj.weight": [ - index_n_heads * index_head_dim, - q_lora_rank, - ], - f"{layer_prefix}.self_attn.compressor.indexer.scorer.weights_proj.weight": [index_n_heads, hidden_size], - f"{layer_prefix}.self_attn.compressor.indexer.position_bias": [4, 256], - f"{layer_prefix}.self_attn.compressor.indexer.kv_norm.weight": [128], - } - ) - else: # hca - layer_mapping.update( - { - f"{layer_prefix}.self_attn.compressor.kv_proj.weight": [512, hidden_size], - f"{layer_prefix}.self_attn.compressor.gate_proj.weight": [512, hidden_size], - f"{layer_prefix}.self_attn.compressor.position_bias": [128, 512], - f"{layer_prefix}.self_attn.compressor.kv_norm.weight": [512], - } + layer_mapping[f"{layer_prefix}.attn.indexer.weights_proj.weight"] = [index_n_heads, hidden_size] + layer_mapping[f"{layer_prefix}.attn.indexer.wq_b.weight"] = [index_n_heads * index_head_dim, q_lora_rank] + + for exp_idx in range(n_routed_experts): + layer_mapping[f"{layer_prefix}.ffn.experts.{exp_idx}.w1.weight"] = [moe_intermediate_size, hidden_size] + layer_mapping[f"{layer_prefix}.ffn.experts.{exp_idx}.w2.weight"] = [hidden_size, moe_intermediate_size] + layer_mapping[f"{layer_prefix}.ffn.experts.{exp_idx}.w3.weight"] = [moe_intermediate_size, hidden_size] + + layer_mapping[f"{layer_prefix}.ffn.shared_experts.w1.weight"] = [moe_intermediate_size, hidden_size] + layer_mapping[f"{layer_prefix}.ffn.shared_experts.w3.weight"] = [moe_intermediate_size, hidden_size] + layer_mapping[f"{layer_prefix}.ffn.shared_experts.w2.weight"] = [hidden_size, moe_intermediate_size] + + return layer_mapping + + compress_ratios = config["compress_ratios"] + for layer_idx in range(num_hidden_layers): + mapping.update( + _layer_shapes( + f"layers.{layer_idx}", + is_hash=layer_idx < num_hash_layers, + ratio=compress_ratios[layer_idx], ) + ) + + # Multi-Token Prediction heads. Each MTP depth owns a full DeepSeek-V4 decoder layer, + # built as DeepSeek4DecoderLayer(compress_ratio=0, is_hash_routing=False), so it has + # neither compressor weights nor hash routing but does carry a MoE gate bias. On top of + # that it has its own norms, a fused input projection and a hyper-connection head. + mtp_num_layers = config.get( + "num_nextn_predict_layers", + config.get("num_mtp_layers", config.get("mtp_num_layers", 0)), + ) + for mtp_idx in range(mtp_num_layers): + mtp_prefix = f"mtp.{mtp_idx}" + mapping.update(_layer_shapes(mtp_prefix, is_hash=False, ratio=0)) + mapping.update( + { + f"{mtp_prefix}.enorm.weight": [hidden_size], + f"{mtp_prefix}.hnorm.weight": [hidden_size], + f"{mtp_prefix}.norm.weight": [hidden_size], + f"{mtp_prefix}.e_proj.weight": [hidden_size, hidden_size], + f"{mtp_prefix}.h_proj.weight": [hidden_size, hidden_size], + # MaxText fuses these two into a single [2 * emb, emb] kernel, so the export + # splits one MaxText param into two HF params via this composite key. + (f"{mtp_prefix}.e_proj.weight", f"{mtp_prefix}.h_proj.weight"): ( + [hidden_size, hidden_size], + [hidden_size, hidden_size], + ), + f"{mtp_prefix}.hc_head_fn": [hc_mult, hidden_size * hc_mult], + f"{mtp_prefix}.hc_head_base": [hc_mult], + f"{mtp_prefix}.hc_head_scale": [1], + } + ) - mapping.update(layer_mapping) return mapping @@ -1316,6 +1357,7 @@ def DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE(config): "gpt-oss-120b": GPT_OSS_HF_WEIGHTS_TO_SHAPE, "mixtral-8x7b": MIXTRAL_HF_WEIGHTS_TO_SHAPE, "mixtral-8x22b": MIXTRAL_HF_WEIGHTS_TO_SHAPE, + "qwen3.5-9b": QWEN3_5_HF_WEIGHTS_TO_SHAPE, "qwen3.5-35b-a3b": QWEN3_5_HF_WEIGHTS_TO_SHAPE, "qwen3.5-35b-a3b-fp8": QWEN3_5_HF_WEIGHTS_TO_SHAPE, "qwen3.5-35b-fp8": QWEN3_5_HF_WEIGHTS_TO_SHAPE, diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index bb9e09d927..01e632b631 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -870,9 +870,10 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals - attention-in_proj_qkvz-kernel: (linear_attn.in_proj_qkv.weight, linear_attn.in_proj_z.weight) - attention-in_proj_ba-kernel: (linear_attn.in_proj_b.weight, linear_attn.in_proj_a.weight") """ - num_main_layers = config["text_config"]["num_hidden_layers"] + num_main_layers = getattr(maxtext_config, "num_decoder_layers", None) or config["text_config"]["num_hidden_layers"] layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval is_quantized = getattr(maxtext_config, "weight_dtype", None) == "float8_e4m3fn" + num_experts = getattr(maxtext_config, "num_experts", None) or config.get("text_config", config).get("num_experts", 1) # 1. Non-layer specific weight mappings mapping = { @@ -989,81 +990,92 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals } ) - # 3. Handle MLP: Gates and Shared Experts - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-mlp-routed_experts-gate-kernel": [ - f"model.language_model.layers.{i}.mlp.gate.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wi_0-kernel": [ - f"model.language_model.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wi_1-kernel": [ - f"model.language_model.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wo-kernel": [ - f"model.language_model.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert_gate-kernel": [ - f"model.language_model.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices - ], - } - ) - if is_quantized: - mapping.update( - { - f"{prefix}-mlp-shared_expert-wi_0-kernel_scale": [ - f"model.language_model.layers.{i}.mlp.shared_expert.gate_proj.weight_scale_inv" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wi_1-kernel_scale": [ - f"model.language_model.layers.{i}.mlp.shared_expert.up_proj.weight_scale_inv" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wo-kernel_scale": [ - f"model.language_model.layers.{i}.mlp.shared_expert.down_proj.weight_scale_inv" for i in hf_indices - ], - } - ) - - # 4. Handle MoE Routed Experts - if is_quantized: - num_experts = config.get("text_config", config).get("num_experts", 256) - mapping.update( + # 3. Handle MLP: Gates and Shared Experts (MoE) or Dense MLP + if num_experts > 1: + mapping.update( # pyrefly: ignore[no-matching-overload] { - f"{prefix}-mlp-routed_experts-wi_0": [ - [f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices] - for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wi_0_scale": [ - [f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight_scale_inv" for i in hf_indices] - for e in range(num_experts) + f"{prefix}-mlp-routed_experts-gate-kernel": [ + f"model.language_model.layers.{i}.mlp.gate.weight" for i in hf_indices ], - f"{prefix}-mlp-routed_experts-wi_1": [ - [f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight" for i in hf_indices] - for e in range(num_experts) + f"{prefix}-mlp-shared_expert-wi_0-kernel": [ + f"model.language_model.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices ], - f"{prefix}-mlp-routed_experts-wi_1_scale": [ - [f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight_scale_inv" for i in hf_indices] - for e in range(num_experts) + f"{prefix}-mlp-shared_expert-wi_1-kernel": [ + f"model.language_model.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices ], - f"{prefix}-mlp-routed_experts-wo": [ - [f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight" for i in hf_indices] - for e in range(num_experts) + f"{prefix}-mlp-shared_expert-wo-kernel": [ + f"model.language_model.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices ], - f"{prefix}-mlp-routed_experts-wo_scale": [ - [f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight_scale_inv" for i in hf_indices] - for e in range(num_experts) + f"{prefix}-mlp-shared_expert_gate-kernel": [ + f"model.language_model.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices ], } ) + if is_quantized: + mapping.update( + { + f"{prefix}-mlp-shared_expert-wi_0-kernel_scale": [ + f"model.language_model.layers.{i}.mlp.shared_expert.gate_proj.weight_scale_inv" for i in hf_indices + ], + f"{prefix}-mlp-shared_expert-wi_1-kernel_scale": [ + f"model.language_model.layers.{i}.mlp.shared_expert.up_proj.weight_scale_inv" for i in hf_indices + ], + f"{prefix}-mlp-shared_expert-wo-kernel_scale": [ + f"model.language_model.layers.{i}.mlp.shared_expert.down_proj.weight_scale_inv" for i in hf_indices + ], + } + ) + + # 4. Handle MoE Routed Experts + if is_quantized: + num_experts = config.get("text_config", config).get("num_experts", 256) + mapping.update( + { + f"{prefix}-mlp-routed_experts-wi_0": [ + [f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices] + for e in range(num_experts) + ], + f"{prefix}-mlp-routed_experts-wi_0_scale": [ + [f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight_scale_inv" for i in hf_indices] + for e in range(num_experts) + ], + f"{prefix}-mlp-routed_experts-wi_1": [ + [f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight" for i in hf_indices] + for e in range(num_experts) + ], + f"{prefix}-mlp-routed_experts-wi_1_scale": [ + [f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight_scale_inv" for i in hf_indices] + for e in range(num_experts) + ], + f"{prefix}-mlp-routed_experts-wo": [ + [f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight" for i in hf_indices] + for e in range(num_experts) + ], + f"{prefix}-mlp-routed_experts-wo_scale": [ + [f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight_scale_inv" for i in hf_indices] + for e in range(num_experts) + ], + } + ) + else: + mapping.update( # pyrefly: ignore[no-matching-overload] + { + f"{prefix}-mlp-routed_experts-wo": [ + f"model.language_model.layers.{i}.mlp.experts.down_proj" for i in hf_indices + ], + (f"{prefix}-mlp-routed_experts-wi_0", f"{prefix}-mlp-routed_experts-wi_1"): [ + f"model.language_model.layers.{i}.mlp.experts.gate_up_proj" for i in hf_indices + ], + } + ) else: - mapping.update( # pyrefly: ignore[no-matching-overload] + mapping.update( { - f"{prefix}-mlp-routed_experts-wo": [ - f"model.language_model.layers.{i}.mlp.experts.down_proj" for i in hf_indices - ], - (f"{prefix}-mlp-routed_experts-wi_0", f"{prefix}-mlp-routed_experts-wi_1"): [ - f"model.language_model.layers.{i}.mlp.experts.gate_up_proj" for i in hf_indices + f"{prefix}-mlp-wi_0-kernel": [ + f"model.language_model.layers.{i}.mlp.gate_proj.weight" for i in hf_indices ], + f"{prefix}-mlp-wi_1-kernel": [f"model.language_model.layers.{i}.mlp.up_proj.weight" for i in hf_indices], + f"{prefix}-mlp-wo-kernel": [f"model.language_model.layers.{i}.mlp.down_proj.weight" for i in hf_indices], } ) else: @@ -1142,63 +1154,72 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals } ) - # MLP: Gates and Shared Experts + # MLP: Gates and Shared Experts (MoE) or Dense MLP hf_mlp = f"model.language_model.layers.{i}.mlp" - mapping.update( - { - f"{prefix}-mlp-routed_experts-gate-kernel": (f"{hf_mlp}.gate.weight"), - f"{prefix}-mlp-shared_expert-wi_0-kernel": (f"{hf_mlp}.shared_expert.gate_proj.weight"), - f"{prefix}-mlp-shared_expert-wi_1-kernel": (f"{hf_mlp}.shared_expert.up_proj.weight"), - f"{prefix}-mlp-shared_expert-wo-kernel": (f"{hf_mlp}.shared_expert.down_proj.weight"), - f"{prefix}-mlp-shared_expert_gate-kernel": (f"{hf_mlp}.shared_expert_gate.weight"), - } - ) - if is_quantized: + if num_experts > 1: mapping.update( { - f"{prefix}-mlp-shared_expert-wi_0-kernel_scale": (f"{hf_mlp}.shared_expert.gate_proj.weight_scale_inv"), - f"{prefix}-mlp-shared_expert-wi_1-kernel_scale": (f"{hf_mlp}.shared_expert.up_proj.weight_scale_inv"), - f"{prefix}-mlp-shared_expert-wo-kernel_scale": (f"{hf_mlp}.shared_expert.down_proj.weight_scale_inv"), + f"{prefix}-mlp-routed_experts-gate-kernel": (f"{hf_mlp}.gate.weight"), + f"{prefix}-mlp-shared_expert-wi_0-kernel": (f"{hf_mlp}.shared_expert.gate_proj.weight"), + f"{prefix}-mlp-shared_expert-wi_1-kernel": (f"{hf_mlp}.shared_expert.up_proj.weight"), + f"{prefix}-mlp-shared_expert-wo-kernel": (f"{hf_mlp}.shared_expert.down_proj.weight"), + f"{prefix}-mlp-shared_expert_gate-kernel": (f"{hf_mlp}.shared_expert_gate.weight"), } ) + if is_quantized: + mapping.update( + { + f"{prefix}-mlp-shared_expert-wi_0-kernel_scale": (f"{hf_mlp}.shared_expert.gate_proj.weight_scale_inv"), + f"{prefix}-mlp-shared_expert-wi_1-kernel_scale": (f"{hf_mlp}.shared_expert.up_proj.weight_scale_inv"), + f"{prefix}-mlp-shared_expert-wo-kernel_scale": (f"{hf_mlp}.shared_expert.down_proj.weight_scale_inv"), + } + ) - # MoE Routed Experts - if is_quantized: - num_experts = config.get("text_config", config).get("num_experts", 256) - mapping.update( - { - f"{prefix}-mlp-routed_experts-wi_0": [ - f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wi_0_scale": [ - f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight_scale_inv" - for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wi_1": [ - f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight" for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wi_1_scale": [ - f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight_scale_inv" - for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wo": [ - f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight" for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wo_scale": [ - f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight_scale_inv" - for e in range(num_experts) - ], - } - ) + # MoE Routed Experts + if is_quantized: + num_experts = config.get("text_config", config).get("num_experts", 256) + mapping.update( + { + f"{prefix}-mlp-routed_experts-wi_0": [ + f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for e in range(num_experts) + ], + f"{prefix}-mlp-routed_experts-wi_0_scale": [ + f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight_scale_inv" + for e in range(num_experts) + ], + f"{prefix}-mlp-routed_experts-wi_1": [ + f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight" for e in range(num_experts) + ], + f"{prefix}-mlp-routed_experts-wi_1_scale": [ + f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight_scale_inv" + for e in range(num_experts) + ], + f"{prefix}-mlp-routed_experts-wo": [ + f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight" for e in range(num_experts) + ], + f"{prefix}-mlp-routed_experts-wo_scale": [ + f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight_scale_inv" + for e in range(num_experts) + ], + } + ) + else: + mapping.update( # pyrefly: ignore[no-matching-overload] + { + f"{prefix}-mlp-routed_experts-wo": f"model.language_model.layers.{i}.mlp.experts.down_proj", + ( + f"{prefix}-mlp-routed_experts-wi_0", + f"{prefix}-mlp-routed_experts-wi_1", + ): f"model.language_model.layers.{i}.mlp.experts.gate_up_proj", + } + ) else: - mapping.update( # pyrefly: ignore[no-matching-overload] + mapping.update( { - f"{prefix}-mlp-routed_experts-wo": f"model.language_model.layers.{i}.mlp.experts.down_proj", - ( - f"{prefix}-mlp-routed_experts-wi_0", - f"{prefix}-mlp-routed_experts-wi_1", - ): f"model.language_model.layers.{i}.mlp.experts.gate_up_proj", + f"{prefix}-mlp-wi_0-kernel": f"{hf_mlp}.gate_proj.weight", + f"{prefix}-mlp-wi_1-kernel": f"{hf_mlp}.up_proj.weight", + f"{prefix}-mlp-wo-kernel": f"{hf_mlp}.down_proj.weight", } ) @@ -1440,7 +1461,8 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): } layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval - num_main_layers = config["text_config"]["num_hidden_layers"] + num_main_layers = getattr(maxtext_config, "num_decoder_layers", None) or config["text_config"]["num_hidden_layers"] + num_experts = getattr(maxtext_config, "num_experts", None) or config.get("text_config", config).get("num_experts", 1) loop_indices = range(layer_cycle_interval) if scan_layers else range(num_main_layers) for i in loop_indices: @@ -1467,27 +1489,32 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): hooks[f"{prefix}-attention-out_proj-kernel_scale"] = transpose mlp_prefix = f"{prefix}-mlp" - hooks[f"{mlp_prefix}-routed_experts-gate-kernel"] = transpose - hooks[f"{mlp_prefix}-shared_expert-wi_0-kernel"] = transpose - hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel"] = transpose - hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose - hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose - if is_quantized: - hooks[f"{mlp_prefix}-shared_expert-wi_0-kernel_scale"] = transpose - hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel_scale"] = transpose - hooks[f"{mlp_prefix}-shared_expert-wo-kernel_scale"] = transpose - hooks[f"{mlp_prefix}-routed_experts-wi_0"] = transpose - hooks[f"{mlp_prefix}-routed_experts-wi_0_scale"] = transpose - hooks[f"{mlp_prefix}-routed_experts-wi_1"] = transpose - hooks[f"{mlp_prefix}-routed_experts-wi_1_scale"] = transpose - hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose - hooks[f"{mlp_prefix}-routed_experts-wo_scale"] = transpose + if num_experts > 1: + hooks[f"{mlp_prefix}-routed_experts-gate-kernel"] = transpose + hooks[f"{mlp_prefix}-shared_expert-wi_0-kernel"] = transpose + hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel"] = transpose + hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose + hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose + if is_quantized: + hooks[f"{mlp_prefix}-shared_expert-wi_0-kernel_scale"] = transpose + hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel_scale"] = transpose + hooks[f"{mlp_prefix}-shared_expert-wo-kernel_scale"] = transpose + hooks[f"{mlp_prefix}-routed_experts-wi_0"] = transpose + hooks[f"{mlp_prefix}-routed_experts-wi_0_scale"] = transpose + hooks[f"{mlp_prefix}-routed_experts-wi_1"] = transpose + hooks[f"{mlp_prefix}-routed_experts-wi_1_scale"] = transpose + hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose + hooks[f"{mlp_prefix}-routed_experts-wo_scale"] = transpose + else: + # pyrefly: ignore[unsupported-operation] + hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = ( + process_wi_0_wi_1 # pyrefly: ignore[unsupported-operation] + ) + hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose_expert else: - # pyrefly: ignore[unsupported-operation] - hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = ( - process_wi_0_wi_1 # pyrefly: ignore[unsupported-operation] - ) - hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose_expert + hooks[f"{mlp_prefix}-wi_0-kernel"] = reshape_kernel + hooks[f"{mlp_prefix}-wi_1-kernel"] = reshape_kernel + hooks[f"{mlp_prefix}-wo-kernel"] = reshape_kernel # Vision hooks for Qwen3.5 vision_config = config.get("vision_config", None) @@ -3949,8 +3976,8 @@ def replace_prefix(val): return mapping -def COSMOS3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): - """Returns mapping from MaxText to HuggingFace Cosmos3-Nano Reasoner weight paths.""" +def WEAVER_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): + """Returns mapping from MaxText to HuggingFace Weaver weight paths.""" # 1. Reuse QWEN3_VL mapping qwen3_vl_mapping = QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers) @@ -3968,7 +3995,7 @@ def translate_hf_key(val): elif val.startswith("model."): val = val[len("model.") :] - # Apply Cosmos3 specific attention naming + # Apply Weaver specific attention naming val = val.replace("self_attn.q_proj", "self_attn.to_q") val = val.replace("self_attn.k_proj", "self_attn.to_k") val = val.replace("self_attn.v_proj", "self_attn.to_v") @@ -3984,8 +4011,8 @@ def translate_hf_key(val): return mapping -def COSMOS3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False): - """Creates parameter transformation functions for Cosmos3-Nano Reasoner.""" +def WEAVER_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False): + """Creates parameter transformation functions for Weaver.""" # Hooks operate on MaxText Parameter paths, which are identical to Qwen3-VL return QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers, saving_to_hf) @@ -4126,330 +4153,355 @@ def reshape_vision_attn_out(input_tensor, target_shape): # {maxtext model name: {maxtext weight name: hf weight name}} -def DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): +def DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): """Maps MaxText parameter keys to HuggingFace parameter keys for DeepSeek V4.""" - n_layers = config["num_hidden_layers"] - num_experts = config.get("n_routed_experts", 8) + + def _get(cfg, key, default=None): + if isinstance(cfg, dict): + return cfg.get(key, default) + return getattr(cfg, key, default) + + n_layers = _get(config, "num_hidden_layers", getattr(maxtext_config, "base_num_decoder_layers", 43)) + num_experts = ( + maxtext_config.num_experts + if getattr(maxtext_config, "num_experts", 0) > 0 + else _get(config, "n_routed_experts", 256) + ) + num_hash_layers = _get(config, "num_hash_layers", getattr(maxtext_config, "first_num_hash_layers", 3)) + compress_ratios = _get(config, "compress_ratios", getattr(maxtext_config, "compress_ratios", None)) + + def _get_compressor_type(layer_idx): + if compress_ratios and layer_idx < len(compress_ratios): + ratio = compress_ratios[layer_idx] + else: + ratio = 0 if layer_idx < 2 else (4 if layer_idx % 2 == 0 else 128) + if ratio == 4: + return "csa" + elif ratio > 4: + return "hca" + return None mapping = { - "params-token_embedder-embedding": "model.embed_tokens.weight", - "params-decoder-decoder_norm-scale": "model.norm.weight", + "params-token_embedder-embedding": "embed.weight", + "params-decoder-decoder_norm-scale": "norm.weight", "params-decoder-logits_dense-kernel": "head.weight", - "params-decoder-hc_head-hc_fn": "model.hc_head.hc_fn", - "params-decoder-hc_head-hc_base": "model.hc_head.hc_base", - "params-decoder-hc_head-hc_scale": "model.hc_head.hc_scale", + "params-decoder-hc_head-hc_fn": "hc_head_fn", + "params-decoder-hc_head-hc_base": "hc_head_base", + "params-decoder-hc_head-hc_scale": "hc_head_scale", } - def add_layer_mapping(mt_layer_path, hf_layer_indices): - is_list = isinstance(hf_layer_indices, list) + def _add_layer(prefix, get_hf_key, is_hash, compressor_type=None): + mapping[f"{prefix}-pre_self_attention_layer_norm-scale"] = get_hf_key("attn_norm.weight") + mapping[f"{prefix}-post_self_attention_layer_norm-scale"] = get_hf_key("ffn_norm.weight") + mapping[f"{prefix}-mhc_attention-mhc_norm-scale"] = None + mapping[f"{prefix}-mhc_mlp-mhc_norm-scale"] = None + + for mod, hf_attr in [("mhc_attention", "hc_attn"), ("mhc_mlp", "hc_ffn")]: + for part in ["pre_alpha", "post_alpha", "res_alpha"]: + mapping[f"{prefix}-{mod}-{part}"] = get_hf_key(f"{hf_attr}_fn") + for part in ["pre_beta", "post_beta", "res_beta"]: + mapping[f"{prefix}-{mod}-{part}"] = get_hf_key(f"{hf_attr}_base") + for part in ["pre_alpha_scale", "post_alpha_scale", "res_alpha_scale"]: + mapping[f"{prefix}-{mod}-{part}"] = get_hf_key(f"{hf_attr}_scale") + + # Attention Core + mapping[f"{prefix}-self_attention-q_norm-scale"] = get_hf_key("attn.q_norm.weight") + mapping[f"{prefix}-self_attention-kv_norm-scale"] = get_hf_key("attn.kv_norm.weight") + mapping[f"{prefix}-self_attention-wq_a-kernel"] = get_hf_key("attn.wq_a.weight") + mapping[f"{prefix}-self_attention-wq_b-kernel"] = get_hf_key("attn.wq_b.weight") + mapping[f"{prefix}-self_attention-wkv-kernel"] = get_hf_key("attn.wkv.weight") + mapping[f"{prefix}-self_attention-sinks"] = get_hf_key("attn.attn_sink") + mapping[f"{prefix}-self_attention-o_a_proj-kernel"] = get_hf_key("attn.wo_a.weight") + mapping[f"{prefix}-self_attention-o_b_proj-kernel"] = get_hf_key("attn.wo_b.weight") + + # Compressors + def _add_compressor_keys(mt_comp_prefix, hf_comp_prefix): + mapping[f"{mt_comp_prefix}-gate_proj-kernel"] = get_hf_key(f"{hf_comp_prefix}.wgate.weight") + mapping[f"{mt_comp_prefix}-kv_proj-kernel"] = get_hf_key(f"{hf_comp_prefix}.wkv.weight") + mapping[f"{mt_comp_prefix}-kv_norm-scale"] = get_hf_key(f"{hf_comp_prefix}.norm.weight") + mapping[f"{mt_comp_prefix}-position_bias"] = get_hf_key(f"{hf_comp_prefix}.ape") + + if compressor_type in ("csa", "both"): + _add_compressor_keys(f"{prefix}-self_attention-csa_compressor", "attn.compressor") + _add_compressor_keys(f"{prefix}-self_attention-csa_compressor-indexer", "attn.indexer.compressor") + mapping[f"{prefix}-self_attention-csa_compressor-indexer-weights_proj-kernel"] = get_hf_key( + "attn.indexer.weights_proj.weight" + ) + mapping[f"{prefix}-self_attention-csa_compressor-indexer-q_proj-kernel"] = get_hf_key("attn.indexer.wq_b.weight") - def get_hf_key(subpath): - if subpath is None: - return None - if is_list: - return [f"model.layers.{idx}.{subpath}" for idx in hf_layer_indices] - else: - return f"model.layers.{hf_layer_indices}.{subpath}" + if compressor_type in ("hca", "both"): + _add_compressor_keys(f"{prefix}-self_attention-hca_compressor", "attn.compressor") - def get_hf_expert_keys(expert_subpath_template): - if is_list: - return [ - [f"model.layers.{idx}.mlp.experts.{e}.{expert_subpath_template}" for idx in hf_layer_indices] - for e in range(num_experts) - ] - else: - return [f"model.layers.{hf_layer_indices}.mlp.experts.{e}.{expert_subpath_template}" for e in range(num_experts)] - - layer_map = { - f"{mt_layer_path}-pre_self_attention_layer_norm-scale": get_hf_key("input_layernorm.weight"), - f"{mt_layer_path}-post_self_attention_layer_norm-scale": get_hf_key("post_attention_layernorm.weight"), - # Attention - f"{mt_layer_path}-self_attention-wq_a-kernel": get_hf_key("self_attn.q_a_proj.weight"), - f"{mt_layer_path}-self_attention-q_norm-scale": get_hf_key("self_attn.q_a_norm.weight"), - f"{mt_layer_path}-self_attention-wq_b-kernel": get_hf_key("self_attn.q_b_proj.weight"), - f"{mt_layer_path}-self_attention-wkv-kernel": get_hf_key("self_attn.kv_proj.weight"), - f"{mt_layer_path}-self_attention-kv_norm-scale": get_hf_key("self_attn.kv_norm.weight"), - f"{mt_layer_path}-self_attention-sinks": get_hf_key("self_attn.sinks"), - f"{mt_layer_path}-self_attention-o_a_proj-kernel": get_hf_key("self_attn.o_a_proj.weight"), - f"{mt_layer_path}-self_attention-o_b_proj-kernel": get_hf_key("self_attn.o_b_proj.weight"), - # mHC Attention - f"{mt_layer_path}-mhc_attention-mhc_norm-scale": None, - f"{mt_layer_path}-mhc_attention-pre_alpha": get_hf_key("attn_hc.fn"), - f"{mt_layer_path}-mhc_attention-post_alpha": get_hf_key("attn_hc.fn"), - f"{mt_layer_path}-mhc_attention-res_alpha": get_hf_key("attn_hc.fn"), - f"{mt_layer_path}-mhc_attention-pre_beta": get_hf_key("attn_hc.base"), - f"{mt_layer_path}-mhc_attention-post_beta": get_hf_key("attn_hc.base"), - f"{mt_layer_path}-mhc_attention-res_beta": get_hf_key("attn_hc.base"), - f"{mt_layer_path}-mhc_attention-pre_alpha_scale": get_hf_key("attn_hc.scale"), - f"{mt_layer_path}-mhc_attention-post_alpha_scale": get_hf_key("attn_hc.scale"), - f"{mt_layer_path}-mhc_attention-res_alpha_scale": get_hf_key("attn_hc.scale"), - # mHC MLP - f"{mt_layer_path}-mhc_mlp-mhc_norm-scale": None, - f"{mt_layer_path}-mhc_mlp-pre_alpha": get_hf_key("ffn_hc.fn"), - f"{mt_layer_path}-mhc_mlp-post_alpha": get_hf_key("ffn_hc.fn"), - f"{mt_layer_path}-mhc_mlp-res_alpha": get_hf_key("ffn_hc.fn"), - f"{mt_layer_path}-mhc_mlp-pre_beta": get_hf_key("ffn_hc.base"), - f"{mt_layer_path}-mhc_mlp-post_beta": get_hf_key("ffn_hc.base"), - f"{mt_layer_path}-mhc_mlp-res_beta": get_hf_key("ffn_hc.base"), - f"{mt_layer_path}-mhc_mlp-pre_alpha_scale": get_hf_key("ffn_hc.scale"), - f"{mt_layer_path}-mhc_mlp-post_alpha_scale": get_hf_key("ffn_hc.scale"), - f"{mt_layer_path}-mhc_mlp-res_alpha_scale": get_hf_key("ffn_hc.scale"), - # MoE Block - f"{mt_layer_path}-mlp-MoeBlock_0-gate-kernel": get_hf_key("mlp.gate.weight"), - # Shared Experts - f"{mt_layer_path}-mlp-shared_experts-wi_0-kernel": get_hf_key("mlp.shared_experts.gate_proj.weight"), - f"{mt_layer_path}-mlp-shared_experts-wi_1-kernel": get_hf_key("mlp.shared_experts.up_proj.weight"), - f"{mt_layer_path}-mlp-shared_experts-wo-kernel": get_hf_key("mlp.shared_experts.down_proj.weight"), - # Stacked Experts - f"{mt_layer_path}-mlp-MoeBlock_0-wi_0": get_hf_expert_keys("w1.weight"), - f"{mt_layer_path}-mlp-MoeBlock_0-wi_1": get_hf_expert_keys("w3.weight"), - f"{mt_layer_path}-mlp-MoeBlock_0-wo": get_hf_expert_keys("w2.weight"), - } + # MoE + mapping[f"{prefix}-mlp-MoeBlock_0-gate-kernel"] = get_hf_key("ffn.gate.weight") + if is_hash: + mapping[f"{prefix.replace('params-', 'Tid2EidVar-')}-mlp-MoeBlock_0-tid2eid"] = get_hf_key("ffn.gate.tid2eid") + else: + mapping[f"{prefix.replace('params-', 'MoEBiasVar-')}-mlp-MoeBlock_0-gate-bias"] = get_hf_key("ffn.gate.bias") - if (is_list and hf_layer_indices[0] >= 3) or (not is_list and hf_layer_indices >= 3): - layer_map[f"{mt_layer_path}-mlp-MoeBlock_0-gate-bias"] = get_hf_key("mlp.gate.e_score_correction_bias") + mapping[f"{prefix}-mlp-shared_experts-wi_0-kernel"] = get_hf_key("ffn.shared_experts.w1.weight") + mapping[f"{prefix}-mlp-shared_experts-wi_1-kernel"] = get_hf_key("ffn.shared_experts.w3.weight") + mapping[f"{prefix}-mlp-shared_experts-wo-kernel"] = get_hf_key("ffn.shared_experts.w2.weight") - first_idx = hf_layer_indices[0] if is_list else hf_layer_indices - if first_idx >= 2: - if first_idx % 2 == 0: - layer_map.update( - { - f"{mt_layer_path}-self_attention-csa_compressor-kv_proj-kernel": get_hf_key( - "self_attn.compressor.kv_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-gate_proj-kernel": get_hf_key( - "self_attn.compressor.gate_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-position_bias": get_hf_key( - "self_attn.compressor.position_bias" - ), - f"{mt_layer_path}-self_attention-csa_compressor-kv_norm-scale": get_hf_key( - "self_attn.compressor.kv_norm.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-gate_proj-kernel": get_hf_key( - "self_attn.compressor.indexer.gate_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-kv_proj-kernel": get_hf_key( - "self_attn.compressor.indexer.kv_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-q_proj-kernel": get_hf_key( - "self_attn.compressor.indexer.q_b_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-weights_proj-kernel": get_hf_key( - "self_attn.compressor.indexer.scorer.weights_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-position_bias": get_hf_key( - "self_attn.compressor.indexer.position_bias" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-kv_norm-scale": get_hf_key( - "self_attn.compressor.indexer.kv_norm.weight" - ), - } - ) - else: - layer_map.update( - { - f"{mt_layer_path}-self_attention-hca_compressor-kv_proj-kernel": get_hf_key( - "self_attn.compressor.kv_proj.weight" - ), - f"{mt_layer_path}-self_attention-hca_compressor-gate_proj-kernel": get_hf_key( - "self_attn.compressor.gate_proj.weight" - ), - f"{mt_layer_path}-self_attention-hca_compressor-position_bias": get_hf_key( - "self_attn.compressor.position_bias" - ), - f"{mt_layer_path}-self_attention-hca_compressor-kv_norm-scale": get_hf_key( - "self_attn.compressor.kv_norm.weight" - ), - } - ) + mapping[f"{prefix}-mlp-MoeBlock_0-wi_0"] = get_hf_key("ffn.experts.{e}.w1.weight", per_expert=True) + mapping[f"{prefix}-mlp-MoeBlock_0-wi_1"] = get_hf_key("ffn.experts.{e}.w3.weight", per_expert=True) + mapping[f"{prefix}-mlp-MoeBlock_0-wo"] = get_hf_key("ffn.experts.{e}.w2.weight", per_expert=True) + + def _make_unrolled_hf_key_fn(layer_idx): + def get_hf_key(target, per_expert=False): + if per_expert: + return [f"layers.{layer_idx}.{target.replace('{e}', str(e))}" for e in range(num_experts)] + return f"layers.{layer_idx}.{target}" - mapping.update(layer_map) # pyrefly: ignore[no-matching-overload] + return get_hf_key + + def _make_scanned_hf_key_fn(layer_indices): + def get_hf_key(target, per_expert=False): + if per_expert: + return [[f"layers.{idx}.{target.replace('{e}', str(e))}" for idx in layer_indices] for e in range(num_experts)] + return [f"layers.{idx}.{target}" for idx in layer_indices] + + return get_hf_key if not scan_layers: for i in range(n_layers): - add_layer_mapping(f"params-decoder-layers_{i}", i) + _add_layer( + f"params-decoder-layers_{i}", + _make_unrolled_hf_key_fn(i), + is_hash=(i < num_hash_layers), + compressor_type=_get_compressor_type(i), + ) else: - for i in range(3): - add_layer_mapping(f"params-decoder-layers_{i}", i) - add_layer_mapping("params-decoder-scanned_blocks-layers_0", list(range(3, n_layers, 2))) - add_layer_mapping("params-decoder-scanned_blocks-layers_1", list(range(4, n_layers, 2))) + # 1. Unrolled Prefix Layers + for i in range(num_hash_layers): + _add_layer( + f"params-decoder-layers_{i}", + _make_unrolled_hf_key_fn(i), + is_hash=True, + compressor_type=_get_compressor_type(i), + ) + + # 2. Scanned Blocks + hca_layers = list(range(num_hash_layers, n_layers, 2)) + csa_layers = list(range(num_hash_layers + 1, n_layers, 2)) + + # Layer 0 in Scanned Blocks (HCA) + _add_layer( + "params-decoder-scanned_blocks-layers_0", + _make_scanned_hf_key_fn(hca_layers), + is_hash=False, + compressor_type="hca", + ) + + # Layer 1 in Scanned Blocks (CSA) + _add_layer( + "params-decoder-scanned_blocks-layers_1", + _make_scanned_hf_key_fn(csa_layers), + is_hash=False, + compressor_type="csa", + ) - for i in range(3): - mapping[f"Tid2EidVar-decoder-layers_{i}-mlp-MoeBlock_0-tid2eid"] = f"model.layers.{i}.mlp.gate.tid2eid" + # 3. MTP Layers (shared by both the scanned and unrolled paths above). + # Follow the MaxText config, not the HF one: the HF checkpoint always ships MTP weights + # (`num_nextn_predict_layers`), but the MaxText model only has MTP params when enabled. + mtp_num_layers = maxtext_config.mtp_num_layers + + def _make_mtp_hf_key_fn(mtp_idx): + def get_hf_key(target, per_expert=False): + if per_expert: + return [f"mtp.{mtp_idx}.{target.replace('{e}', str(e))}" for e in range(num_experts)] + return f"mtp.{mtp_idx}.{target}" + + return get_hf_key + + for j in range(mtp_num_layers): + mtp_prefix = f"params-mtp_block-mtp_layer_{j + 1}" + + # Note: MultiTokenPredictionLayer renames its submodules to `mtp_{k}_*` via property + # setters, except for `hc_head`, which keeps its bare name. + mapping[f"{mtp_prefix}-mtp_{j + 1}_embedding_norm-scale"] = f"mtp.{j}.enorm.weight" + mapping[f"{mtp_prefix}-mtp_{j + 1}_hidden_state_norm-scale"] = f"mtp.{j}.hnorm.weight" + mapping[f"{mtp_prefix}-mtp_{j + 1}_final_norm-scale"] = f"mtp.{j}.norm.weight" + # MaxText fuses the embedding and hidden-state projections into one [2 * emb, emb] kernel. + mapping[f"{mtp_prefix}-mtp_{j + 1}_projection-kernel"] = (f"mtp.{j}.e_proj.weight", f"mtp.{j}.h_proj.weight") + mapping[f"{mtp_prefix}-hc_head-hc_fn"] = f"mtp.{j}.hc_head_fn" + mapping[f"{mtp_prefix}-hc_head-hc_base"] = f"mtp.{j}.hc_head_base" + mapping[f"{mtp_prefix}-hc_head-hc_scale"] = f"mtp.{j}.hc_head_scale" + + # The MTP transformer layer is built as DeepSeek4DecoderLayer(compress_ratio=0, + # is_hash_routing=False), so it has neither a compressor nor hash routing, but it + # does carry a MoE gate bias. + _add_layer( + f"{mtp_prefix}-mtp_{j + 1}_transformer_layer", + _make_mtp_hf_key_fn(j), + is_hash=False, + compressor_type=None, + ) return mapping -def DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False): +def DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False): """Returns hook functions for transforming weights between MaxText and HuggingFace for DeepSeek V4.""" - def transpose(input_tensor, target_shape=None): - return np.transpose(input_tensor) - - def ones_norm(input_tensor, target_shape=None): - return np.ones(target_shape, dtype=np.float32) # pyrefly: ignore[no-matching-overload] + def _get(cfg, key, default=None): + if isinstance(cfg, dict): + return cfg.get(key, default) + return getattr(cfg, key, default) - def identity(input_tensor, target_shape=None): - return input_tensor + n_layers = _get(config, "num_hidden_layers", getattr(maxtext_config, "base_num_decoder_layers", 43)) + num_hash_layers = _get(config, "num_hash_layers", getattr(maxtext_config, "first_num_hash_layers", 3)) - # Reshaping functions for wq_b, wkv, o_a_proj - def reshape_transpose_wq_b(input_tensor, target_shape=None): - # HF: [n_heads * q_head_dim, kv_lora_rank] - # MaxText: [kv_lora_rank, n_heads, q_head_dim] - if saving_to_hf: - tensor = input_tensor.reshape((input_tensor.shape[0], -1)) - return np.transpose(tensor) - tensor = np.transpose(input_tensor) # [kv_lora_rank, n_heads * q_head_dim] - return tensor.reshape(target_shape) - - def reshape_transpose_wkv(input_tensor, target_shape=None): - # HF: [n_kv_heads * (q_head_dim + v_head_dim), kv_lora_rank] - # MaxText: [kv_lora_rank, n_kv_heads, q_head_dim + v_head_dim] - if saving_to_hf: - tensor = input_tensor.reshape((input_tensor.shape[0], -1)) - return np.transpose(tensor) - tensor = np.transpose(input_tensor) - return tensor.reshape(target_shape) - - def reshape_transpose_o_a(input_tensor, target_shape=None): - # HF: [n_heads * v_head_dim, kv_lora_rank] (e.g. [8192, 4096]) - # MaxText: [n_heads, v_head_dim, kv_lora_rank] (e.g. [8, 4096, 1024]) - # We must reshape first and then permute (transpose) to get correct ordering. + def reshape_kernel(input_tensor, target_shape): if saving_to_hf: - tensor = np.transpose(input_tensor, (0, 2, 1)) - return tensor.reshape(target_shape) - num_heads = target_shape[0] # pyrefly: ignore[unsupported-operation] - embed_dim = target_shape[1] # pyrefly: ignore[unsupported-operation] - kv_lora_rank = target_shape[2] # pyrefly: ignore[unsupported-operation] - tensor = input_tensor.reshape((num_heads, kv_lora_rank, embed_dim)) - return np.transpose(tensor, (0, 2, 1)) + flipped_target_shape = np.flip(np.array(target_shape)) + return input_tensor.reshape(flipped_target_shape).T + else: + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f"RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}") from e - # Functions for mHC split - def mhc_split_fn_pre(input_tensor, target_shape=None): - return np.transpose(input_tensor[0:4, :]) + def reshape_o_a_proj(input_tensor, target_shape): + if saving_to_hf: + heads, head_dim, in_dim = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] + if target_shape[0] == in_dim: + return input_tensor.transpose(2, 0, 1).reshape(target_shape) + else: + return input_tensor.transpose(0, 2, 1).reshape(target_shape) + else: + heads, head_dim, in_dim = target_shape[0], target_shape[1], target_shape[2] + if input_tensor.shape[0] == in_dim: + return input_tensor.reshape(in_dim, heads, head_dim).transpose(1, 2, 0) + else: + return input_tensor.reshape(heads, in_dim, head_dim).transpose(0, 2, 1) - def mhc_split_fn_post(input_tensor, target_shape=None): - return np.transpose(input_tensor[4:8, :]) + # --- MHC Hook Functions --- + mhc_composite_fns = { + "alpha": lambda w, *args, **kwargs: np.concatenate(w, axis=-1).T, + "beta": lambda w, *args, **kwargs: np.concatenate([w[0], w[1], w[2].flatten()], axis=-1), + "scale": lambda w, *args, **kwargs: np.concatenate(w, axis=-1), + } - def mhc_split_fn_res(input_tensor, target_shape=None): - return np.transpose(input_tensor[8:24, :]) + def mhc_slice_hook(mode: str, part_idx: int): + """Unified slicer for MHC parameters across alpha (0), beta (1), and scale (2).""" + + def hook(tensor: np.ndarray, target_shape=None): + del target_shape + if mode == "scale": + return tensor[part_idx : part_idx + 1] + t = tensor.T if mode == "alpha" else tensor + k = int(np.sqrt(1 + t.shape[-1]) - 1) + if part_idx == 0: + return t[..., :k] + elif part_idx == 1: + return t[..., k : 2 * k] + else: + res = t[..., 2 * k :] + return res.reshape(k, k) if mode == "beta" else res - def mhc_split_base_pre(input_tensor, target_shape=None): - return input_tensor[0:4] + return hook - def mhc_split_base_post(input_tensor, target_shape=None): - return input_tensor[4:8] + def mhc_dummy_norm(input_tensor, target_shape=None): + dtype = input_tensor.dtype if input_tensor is not None else np.float32 + return np.ones(target_shape, dtype=dtype) - def mhc_split_base_res(input_tensor, target_shape=None): - return input_tensor[8:24].reshape(target_shape) + def unpad_hf_embedding_layer(input_tensor, target_shape): + target_vocab_size = target_shape[0] + if input_tensor.shape[0] == target_vocab_size: + return input_tensor[:target_vocab_size, :] + else: + # MaxText is (emb_dim, vocab_size) + return input_tensor[:, :target_vocab_size].T - def mhc_split_scale_pre(input_tensor, target_shape=None): - return np.array([input_tensor[0]]).reshape(target_shape) + def unpad_logits_layer(input_tensor, target_shape): + if saving_to_hf: + return input_tensor[:, : target_shape[0]].T + else: + return input_tensor[: target_shape[1], :].T - def mhc_split_scale_post(input_tensor, target_shape=None): - return np.array([input_tensor[1]]).reshape(target_shape) + def reshape_mtp_projection(weights, target_shape=None): + """Splits/joins the fused MTP projection kernel against separate e_proj/h_proj tensors.""" + del target_shape + if saving_to_hf: + half = weights.shape[0] // 2 + e_proj = weights[:half, :].T + h_proj = weights[half:, :].T + return e_proj, h_proj + else: + e_proj, h_proj = weights + return np.concatenate([e_proj.T, h_proj.T], axis=0) - def mhc_split_scale_res(input_tensor, target_shape=None): - return np.array([input_tensor[2]]).reshape(target_shape) + mapping = { + "params-token_embedder-embedding": unpad_hf_embedding_layer, + "params-decoder-logits_dense-kernel": unpad_logits_layer, + "params-decoder-hc_head-hc_fn": reshape_kernel, + } - mapping = {} + def _attach_layer_hooks(prefix): + mapping[f"{prefix}-self_attention-o_a_proj-kernel"] = reshape_o_a_proj + mapping[f"{prefix}-mhc_attention-mhc_norm-scale"] = mhc_dummy_norm + mapping[f"{prefix}-mhc_mlp-mhc_norm-scale"] = mhc_dummy_norm + + for key in [ + f"{prefix}-self_attention-wq_a-kernel", + f"{prefix}-self_attention-wq_b-kernel", + f"{prefix}-self_attention-wkv-kernel", + f"{prefix}-self_attention-o_b_proj-kernel", + f"{prefix}-mlp-MoeBlock_0-gate-kernel", + f"{prefix}-mlp-MoeBlock_0-wi_0", + f"{prefix}-mlp-MoeBlock_0-wi_1", + f"{prefix}-mlp-MoeBlock_0-wo", + f"{prefix}-mlp-shared_experts-wi_0-kernel", + f"{prefix}-mlp-shared_experts-wi_1-kernel", + f"{prefix}-mlp-shared_experts-wo-kernel", + f"{prefix}-self_attention-csa_compressor-gate_proj-kernel", + f"{prefix}-self_attention-csa_compressor-kv_proj-kernel", + f"{prefix}-self_attention-csa_compressor-indexer-gate_proj-kernel", + f"{prefix}-self_attention-csa_compressor-indexer-kv_proj-kernel", + f"{prefix}-self_attention-csa_compressor-indexer-weights_proj-kernel", + f"{prefix}-self_attention-csa_compressor-indexer-q_proj-kernel", + f"{prefix}-self_attention-hca_compressor-gate_proj-kernel", + f"{prefix}-self_attention-hca_compressor-kv_proj-kernel", + ]: + mapping[key] = reshape_kernel + + for mod in ("mhc_attention", "mhc_mlp"): + for mode, suffix in (("alpha", "alpha"), ("beta", "beta"), ("scale", "alpha_scale")): + keys = ( + f"{prefix}-{mod}-pre_{suffix}", + f"{prefix}-{mod}-post_{suffix}", + f"{prefix}-{mod}-res_{suffix}", + ) + if saving_to_hf: + mapping[keys] = mhc_composite_fns[mode] + else: + for idx, k in enumerate(keys): + mapping[k] = mhc_slice_hook(mode, idx) - # Base mapping logic from original file - for key, hf_key in DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers).items(): - if hf_key is None: - mapping[key] = ones_norm - elif "token_embedder-embedding" in key: - mapping[key] = identity - elif "-wkv-kernel" in key: - mapping[key] = reshape_transpose_wkv - elif "-wq_b-kernel" in key: - mapping[key] = reshape_transpose_wq_b - elif "-o_a_proj-kernel" in key: - mapping[key] = reshape_transpose_o_a - elif "mhc" in key: - if "pre_alpha" in key and "scale" not in key: - mapping[key] = mhc_split_fn_pre - elif "post_alpha" in key and "scale" not in key: - mapping[key] = mhc_split_fn_post - elif "res_alpha" in key and "scale" not in key: - mapping[key] = mhc_split_fn_res - elif "pre_beta" in key: - mapping[key] = mhc_split_base_pre - elif "post_beta" in key: - mapping[key] = mhc_split_base_post - elif "res_beta" in key: - mapping[key] = mhc_split_base_res - elif "pre_alpha_scale" in key: - mapping[key] = mhc_split_scale_pre - elif "post_alpha_scale" in key: - mapping[key] = mhc_split_scale_post - elif "res_alpha_scale" in key: - mapping[key] = mhc_split_scale_res - elif "position_bias" in key: - mapping[key] = identity - elif "hc_head-hc_fn" in key: - mapping[key] = transpose - elif "hc_head-hc_base" in key or "hc_head-hc_scale" in key: - mapping[key] = identity - elif isinstance(hf_key, list): - mapping[key] = transpose - elif "-kernel" in key or "-embedding" in key or "-sinks" in key: - mapping[key] = transpose - - if saving_to_hf: - - def mhc_concat_fn(input_tensors, target_shape=None): - if len(input_tensors) != 3: - raise ValueError(f"mhc_concat_fn expected 3 tensors (pre, post, res), got {len(input_tensors)}") - tensors = [np.asarray(t) for t in input_tensors] - res = np.transpose(np.concatenate(tensors, axis=1)) - return res.reshape(target_shape) if target_shape is not None else res - - def mhc_concat_base(input_tensors, target_shape=None): - if len(input_tensors) != 3: - raise ValueError(f"mhc_concat_base expected 3 tensors (pre, post, res), got {len(input_tensors)}") - tensors = [np.asarray(t).ravel() for t in input_tensors] - res = np.concatenate(tensors, axis=0) - return res.reshape(target_shape) if target_shape is not None else res - - def mhc_concat_scale(input_tensors, target_shape=None): - if len(input_tensors) != 3: - raise ValueError(f"mhc_concat_scale expected 3 tensors (pre, post, res), got {len(input_tensors)}") - tensors = [np.asarray(t).ravel() for t in input_tensors] - res = np.concatenate(tensors, axis=0) - return res.reshape(target_shape) if target_shape is not None else res - - # Process composite mappings - keys_to_delete = [] - keys_to_add = {} - for key in list(mapping.keys()): - if "mhc" in key and "pre_alpha" in key and "scale" not in key: - post = key.replace("pre_alpha", "post_alpha") - res = key.replace("pre_alpha", "res_alpha") - keys_to_delete.extend([key, post, res]) - keys_to_add[(key, post, res)] = mhc_concat_fn - - if "mhc" in key and "pre_beta" in key: - post = key.replace("pre_beta", "post_beta") - res = key.replace("pre_beta", "res_beta") - keys_to_delete.extend([key, post, res]) - keys_to_add[(key, post, res)] = mhc_concat_base - - if "mhc" in key and "pre_alpha_scale" in key: - post = key.replace("pre_alpha_scale", "post_alpha_scale") - res = key.replace("pre_alpha_scale", "res_alpha_scale") - keys_to_delete.extend([key, post, res]) - keys_to_add[(key, post, res)] = mhc_concat_scale - - for k in set(keys_to_delete): - if k in mapping: - del mapping[k] - mapping.update(keys_to_add) + prefixes = [f"params-decoder-layers_{i}" for i in range(n_layers if not scan_layers else num_hash_layers)] + if scan_layers: + prefixes.extend(["params-decoder-scanned_blocks-layers_0", "params-decoder-scanned_blocks-layers_1"]) + for prefix in prefixes: + _attach_layer_hooks(prefix) + + # MTP layers reuse the same per-layer hooks, plus two MTP-specific transforms. + # Follow the MaxText config, not the HF one: the HF checkpoint always ships MTP weights + # (`num_nextn_predict_layers`), but the MaxText model only has MTP params when enabled. + mtp_num_layers = maxtext_config.mtp_num_layers + for j in range(mtp_num_layers): + mtp_prefix = f"params-mtp_block-mtp_layer_{j + 1}" + mapping[f"{mtp_prefix}-mtp_{j + 1}_projection-kernel"] = reshape_mtp_projection + mapping[f"{mtp_prefix}-hc_head-hc_fn"] = reshape_kernel + _attach_layer_hooks(f"{mtp_prefix}-mtp_{j + 1}_transformer_layer") return mapping +DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING = DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_MAPPING +DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN = DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN + + PARAM_MAPPING = { "gemma2-2b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING, "gemma2-9b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4478,8 +4530,8 @@ def mhc_concat_scale(input_tensors, target_shape=None): "qwen3-vl-2b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-vl-4b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-vl-30b-a3b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, - "cosmos3-nano-reasoner": COSMOS3_MAXTEXT_TO_HF_PARAM_MAPPING, - "cosmos3-super-reasoner": COSMOS3_MAXTEXT_TO_HF_PARAM_MAPPING, + "weaver-mini": WEAVER_MAXTEXT_TO_HF_PARAM_MAPPING, + "weaver-max": WEAVER_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-8b-Instruct": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4491,7 +4543,7 @@ def mhc_concat_scale(input_tensors, target_shape=None): "deepseek2-16b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "deepseek3.2-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, - "deepseek4-284b": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING, + "deepseek4-284b": DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_MAPPING, "gpt-oss-20b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING, "gpt-oss-120b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4501,6 +4553,7 @@ def mhc_concat_scale(input_tensors, target_shape=None): "qwen3.5-35b-a3b": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3.5-35b-a3b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3.5-35b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING, + "qwen3.5-9b": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING, "mixtral-8x7b": MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING, "mixtral-8x22b": MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING, "olmo3-7b": OLMO3_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4537,8 +4590,8 @@ def mhc_concat_scale(input_tensors, target_shape=None): "qwen3-vl-2b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3-vl-4b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3-vl-30b-a3b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, - "cosmos3-nano-reasoner": COSMOS3_MAXTEXT_TO_HF_PARAM_HOOK_FN, - "cosmos3-super-reasoner": COSMOS3_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "weaver-mini": WEAVER_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "weaver-max": WEAVER_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-8b-Instruct": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, @@ -4550,8 +4603,8 @@ def mhc_concat_scale(input_tensors, target_shape=None): "deepseek2-16b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek3.2-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, - "deepseek4-tiny": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN, - "deepseek4-284b": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "deepseek4-tiny": DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "deepseek4-284b": DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN, "gpt-oss-20b": GPT_OSS_TO_HF_PARAM_HOOK_FN, "gpt-oss-120b": GPT_OSS_TO_HF_PARAM_HOOK_FN, "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN, @@ -4560,6 +4613,7 @@ def mhc_concat_scale(input_tensors, target_shape=None): "qwen3.5-35b-a3b": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3.5-35b-a3b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3.5-35b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "qwen3.5-9b": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3-next-80b-a3b": QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_HOOK_FN, "mixtral-8x7b": MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN, "mixtral-8x22b": MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN, diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index 7008461ade..4bd92fb8ca 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/src/maxtext/checkpoint_conversion/utils/utils.py @@ -475,6 +475,17 @@ def save_safetensor_file( if "model.safetensors" in state_dict and isinstance(state_dict["model.safetensors"], dict): state_dict = state_dict["model.safetensors"] + # safetensors' numpy backend (0.8.0) serializes a non-contiguous array's + # raw buffer and ignores strides, writing the correct shape metadata with + # untransposed data -- silent corruption. Its torch backend raises on + # non-contiguous input instead; the numpy path does not. Any hook ending in + # `.T` (e.g. the DeepSeek-V4 MHC alpha composite) produces exactly such a + # view, so normalise every tensor to C-contiguous before writing. + for _key in list(state_dict.keys()): + _val = state_dict[_key] + if isinstance(_val, np.ndarray) and not _val.flags["C_CONTIGUOUS"]: + state_dict[_key] = np.ascontiguousarray(_val) + if output_dir_final.startswith("gs://"): local_path = os.path.join(local_dir_to_save_to, file_name) numpy_save_file(state_dict, local_path, metadata={"format": "pt"}) @@ -1029,7 +1040,7 @@ def extract_nnx_weights(weights_dict: dict) -> dict[str, np.ndarray]: return result -def extract_linen_weights(weights_dict: dict) -> dict[str, np.ndarray]: +def extract_linen_weights(weights_dict: dict, prefix: str = "params-") -> dict[str, np.ndarray]: """Extract weights from Linen checkpoint structure. Linen checkpoints have structure: {'params': {'decoder': {'decoder_norm': {'scale': array}}}} @@ -1037,6 +1048,7 @@ def extract_linen_weights(weights_dict: dict) -> dict[str, np.ndarray]: Args: weights_dict: Linen checkpoint weights dictionary + prefix: Prefix for parameter keys (defaults to "params-") Returns: Dictionary mapping parameter names to weight arrays @@ -1045,7 +1057,7 @@ def extract_linen_weights(weights_dict: dict) -> dict[str, np.ndarray]: leaves_with_paths = jax.tree_util.tree_leaves_with_path(weights_dict) for path_tuple, leaf_value in leaves_with_paths: path_keys = param_key_parts_from_path(path_tuple) - maxtext_param_key = "params-" + "-".join(path_keys) + maxtext_param_key = prefix + "-".join(path_keys) if not isinstance(leaf_value, (jax.Array, np.ndarray)): raise ValueError(f"Leaf value for {maxtext_param_key} is not an array. Type: {type(leaf_value)}.") result[maxtext_param_key] = leaf_value @@ -1087,8 +1099,18 @@ def detect_and_extract_checkpoint(checkpoint_dict: dict) -> dict[str, np.ndarray else: # Linen checkpoint: check if there's a nested 'params' key if isinstance(actual_weights_dict, dict) and "params" in actual_weights_dict: - actual_weights_dict = actual_weights_dict["params"] - max_logging.log("Detected Linen checkpoint structure") + extra_collections = [k for k in actual_weights_dict if k != "params"] + if extra_collections: + max_logging.log( + f"Detected multi-collection Linen checkpoint structure with collections: {list(actual_weights_dict.keys())}" + ) + result = {} + for col_name, col_dict in actual_weights_dict.items(): + result.update(extract_linen_weights(col_dict, prefix=f"{col_name}-")) + return result + else: + actual_weights_dict = actual_weights_dict["params"] + max_logging.log("Detected Linen checkpoint structure") else: max_logging.log("Detected Linen checkpoint structure (single params layer)") return extract_linen_weights(actual_weights_dict) @@ -1292,8 +1314,9 @@ def save_weights_to_checkpoint( if checkpoint_manager is None: raise RuntimeError("Failed to create Orbax checkpoint manager.") + params = jax_weights if (isinstance(jax_weights, dict) and "params" in jax_weights) else {"params": jax_weights} state_new = train_state.TrainState( - step=step_number_to_save_new_ckpt, apply_fn=None, params={"params": jax_weights}, tx=None, opt_state={} # type: ignore + step=step_number_to_save_new_ckpt, apply_fn=None, params=params, tx=None, opt_state={} # type: ignore ) logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3)) diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 94ba1617ab..508c516158 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -15,7 +15,9 @@ """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer.""" +import collections import contextlib +import dataclasses import importlib import os import time @@ -23,6 +25,7 @@ from etils import epath from flax import nnx +from flax.nnx import variablelib from flax import struct @@ -45,6 +48,7 @@ from maxtext.utils import globals as maxtext_globals from maxtext.utils import max_logging from maxtext.utils import sharding as maxtext_sharding +from orbax.checkpoint import logging as ocp_logging from orbax.checkpoint import v1 as ocp from orbax.checkpoint._src.arrays import sharding as sharding_utils @@ -63,6 +67,17 @@ CheckpointManager = ocp.training.Checkpointer | EmergencyCheckpointManager | EmergencyReplicatorCheckpointManager +def _deep_merge_dicts(target: dict[str, Any], source: dict[str, Any]) -> dict[str, Any]: + """Recursively merges source dict into target dict.""" + out = dict(target) + for k, v in source.items(): + if k in out and isinstance(out[k], dict) and isinstance(v, dict): + out[k] = _deep_merge_dicts(out[k], v) + else: + out[k] = v + return out + + def _weight_mismatches(want, have, path=(), check_missing: bool = True, is_quantized_param: bool = False): """Returns `(path, problem)` for each weight in `want` that `have` didn't restore. @@ -100,6 +115,134 @@ def _weight_mismatches(want, have, path=(), check_missing: bool = True, is_quant return [] +def _lookup_path(tree, path): + """Returns the leaf at `path` in a nested dict, or None if absent.""" + node = tree + for key in path: + if not isinstance(node, dict) or key not in node: + return None + node = node[key] + return node + + +def _alias_legacy_collections(params_collection, stored_collection_meta): + """Redirects weights whose Flax collection postdates the checkpoint into `params`. + + A weight's Linen collection name is derived from its NNX Variable *class* + (`variable_name_from_type`). So promoting a weight from `nnx.Param` to a custom + `nnx.Variable` subclass silently moves it from the `params` collection into a new + one -- and every checkpoint written before that promotion still stores it under + `params`. + + This bit us badly: DeepSeek-V4's routed gate bias became a `MoEBiasVar`, but the + converted checkpoints store it at `params/params/decoder/layers_N/mlp/MoeBlock_0/ + gate/bias`. The restore request simply skipped the unknown `MoEBiasVar` collection, + so all 40 bias tensors stayed at their zero initialiser with no error raised, which + silently corrupted expert selection in every MoE layer. + + Rather than drop such a collection, look each of its leaves up inside the + checkpoint's `params` collection and request it from there. A leaf that is genuinely + absent is left alone so the existing mismatch check can still report it. + + Args: + params_collection: mapping of collection name -> nested weight tree, as it will be + handed to Orbax. + stored_collection_meta: mapping of collection name -> nested metadata tree read off + the checkpoint, or None when metadata could not be read. + + Returns: + `(patched_collection, aliased)` where `aliased` maps collection name -> list of + leaf paths that were moved into `params`. Callers restoring into a structure that + is keyed by collection (i.e. Linen) must move those paths back afterwards; the NNX + path merges all collections by path, so it needs no un-aliasing. + """ + if not isinstance(stored_collection_meta, dict) or not stored_collection_meta: + return params_collection, {} + + stored_params = stored_collection_meta.get("params") + if not isinstance(stored_params, dict): + return params_collection, {} + + patched = dict(params_collection) + params_flat = dict(nnx.traversals.flatten_mapping(patched.get("params", {}))) + aliased = {} + + for col, tree in params_collection.items(): + if col == "params" or col in stored_collection_meta or not isinstance(tree, dict): + continue + tree_flat = dict(nnx.traversals.flatten_mapping(tree)) + moved = [] + for path, leaf in tree_flat.items(): + if _lookup_path(stored_params, path) is None: + continue # genuinely absent; let the normal mismatch check speak. + params_flat[path] = leaf + moved.append(path) + if moved: + aliased[col] = moved + moved_set = set(moved) + remaining = {p: v for p, v in tree_flat.items() if p not in moved_set} + if remaining: + patched[col] = nnx.traversals.unflatten_mapping(remaining) + else: + patched.pop(col, None) + max_logging.log( + f"Checkpoint predates the '{col}' collection; restoring its {len(moved)} weight(s) " + f"from the 'params' collection instead of skipping them." + ) + + if aliased: + patched["params"] = nnx.traversals.unflatten_mapping(params_flat) + return patched, aliased + + +def _unalias_legacy_collections(restored_collection, aliased): + """Moves aliased weights back from `params` into their own collection. + + Inverse of `_alias_legacy_collections`, needed only for the Linen path, whose + restored structure stays keyed by collection. + """ + if not aliased or not isinstance(restored_collection, dict): + return restored_collection + + out = dict(restored_collection) + params_flat = dict(nnx.traversals.flatten_mapping(out.get("params", {}))) + for col, paths in aliased.items(): + col_flat = dict(nnx.traversals.flatten_mapping(out.get(col, {}))) + for path in paths: + if path in params_flat: + col_flat[path] = params_flat.pop(path) + if col_flat: + out[col] = nnx.traversals.unflatten_mapping(col_flat) + out["params"] = nnx.traversals.unflatten_mapping(params_flat) + return out + + +def _log_unconsumed_checkpoint_weights(want, stored_collection): + """Warns about checkpoint tensors the model never asked for. + + `_weight_mismatches` only walks `want`, so it answers "did the checkpoint supply + everything the model needs?" and never "did the model consume everything the + checkpoint supplied?". The second question is the one that would have caught the + DeepSeek-V4 routed-bias regression immediately. This is a warning rather than an + error because legitimately partial loads (LoRA, SFT, params-only) are common. + """ + if not isinstance(want, dict) or not isinstance(stored_collection, dict): + return + try: + want_paths = set(nnx.traversals.flatten_mapping(want).keys()) + stored_paths = set(nnx.traversals.flatten_mapping(stored_collection).keys()) + except Exception: # pylint: disable=broad-except + return + unconsumed = sorted("/".join(str(p) for p in path) for path in stored_paths - want_paths) + if not unconsumed: + return + preview = ", ".join(unconsumed[:5]) + max_logging.log( + f"WARNING: {len(unconsumed)} checkpoint weight(s) were not requested by the model and " + f"will be ignored (the model will use its initialised values). First few: {preview}" + ) + + def _expected_and_restored_params(abstract_nnx_state, restored_linen): """Returns the model's expected weights and the checkpoint's restored weights, as pure dicts. @@ -176,6 +319,13 @@ def _linen_items_to_nnx(restored_linen, abstract_nnx_state): nnx_aux = restored_linen.get("nnx_aux") if nnx_aux: train_state_nnx.apply_checkpoint_aux(aux_state, nnx_aux) + elif isinstance(restored_linen.get("params"), dict): + custom_collections = {k: v for k, v in restored_linen["params"].items() if k != "params" and isinstance(v, dict)} + if custom_collections: + merged_custom = {} + for col_data in custom_collections.values(): + merged_custom = _deep_merge_dicts(merged_custom, col_data) + train_state_nnx.apply_checkpoint_aux(aux_state, {"model": merged_custom}) return nnx.merge_state(linen_state, aux_state, ephemeral) @@ -253,9 +403,17 @@ def _restored_linen_to_nnx(restored_linen, abstract_nnx_state, config=None): def _abstract_params(abstract_unboxed_pre_state): - """Returns the state's weights: the NNX Param subtree, or Linen's `params` collection.""" + """Returns the state's weights: persistent model variables (excluding transient state), or Linen's `params` collection.""" if isinstance(abstract_unboxed_pre_state, nnx.State): - return nnx.split_state(abstract_unboxed_pre_state.model, nnx.Param, ...)[0] + model_state = ( + abstract_unboxed_pre_state["model"] if "model" in abstract_unboxed_pre_state else abstract_unboxed_pre_state + ) + return model_state.filter( + lambda path, var: not isinstance(var, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat)) + ) + elif hasattr(abstract_unboxed_pre_state, "model") and isinstance(abstract_unboxed_pre_state.model, nnx.Module): + state = nnx.state(abstract_unboxed_pre_state.model) + return state.filter(lambda path, var: not isinstance(var, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat))) return abstract_unboxed_pre_state.params @@ -392,6 +550,7 @@ def create_orbax_checkpoint_manager( todelete_subdir: str | None = None, todelete_full_path: str | None = None, ocdbt_target_data_file_size_bytes: int | None = None, + colocated_python_checkpointing: bool = False, ): """Returns an Orbax v1 training ``Checkpointer``, or None if checkpointing is disabled.""" if not enable_checkpointing: @@ -399,12 +558,6 @@ def create_orbax_checkpoint_manager( return None # TODO: b/529622681 - Remove deprecated settings. - if orbax_logger is not None: - max_logging.warning( - "Cloud logging (enable_checkpoint_cloud_logger) is disabled because" - " Orbax v1 now configures its own logger internally. This config" - " setting is ignored and will be removed." - ) if dataset_type is not None: max_logging.warning( "Specifying dataset_type upon checkpointer creation is deprecated and" @@ -428,6 +581,7 @@ def create_orbax_checkpoint_manager( todelete_full_path=todelete_full_path, todelete_subdir=todelete_subdir, partial_load=True, + colocated_python_checkpointing=colocated_python_checkpointing, ) manager = ocp.training.Checkpointer( @@ -444,11 +598,83 @@ def create_orbax_checkpoint_manager( ) # Necessary bridge to support v0 backward compatibility. manager.use_async = use_async # pyrefly: ignore[missing-attribute] + if orbax_logger is not None: + _attach_checkpoint_logger(manager, orbax_logger) max_logging.log("Checkpoint manager created!") return manager +def _attach_checkpoint_logger(manager: ocp.training.Checkpointer, orbax_logger: Any) -> None: + """Routes Orbax checkpoint step statistics to `orbax_logger` (e.g. Goodput). + + The Orbax v1 `Checkpointer` does not expose a `logger` argument and builds its + internal v0 `CheckpointManager` with the default `StandardLogger`, so save + statistics never reach Cloud Logging. Attach the logger to the internal + manager so v0-compatible save step statistics are emitted again. + + TODO: b/529622681 - Pass the logger via the public v1 API once Orbax supports it. + + Args: + manager: The Orbax v1 training `Checkpointer`. + orbax_logger: An Orbax `AbstractLogger`, e.g. `ocp_logging.CloudLogger`. + """ + internal_manager = getattr(manager, "_manager", None) + if internal_manager is None or not hasattr(internal_manager, "_logger"): + max_logging.warning( + "Could not attach checkpoint logger: Orbax v1 Checkpointer internals changed." + " Checkpoint save/restore statistics will not be sent to Cloud Logging." + ) + return + try: + internal_manager._logger = orbax_logger # pylint: disable=protected-access + # v1 `load_checkpointables` bypasses the internal manager's restore logging, + # so keep a handle to log restore statistics from `load_state_if_possible`. + manager.orbax_logger = orbax_logger # pyrefly: ignore[missing-attribute] + except AttributeError as e: + max_logging.warning( + f"Could not attach checkpoint logger: {e}." + " Checkpoint save/restore statistics will not be sent to Cloud Logging." + ) + return + max_logging.log("Attached checkpoint logger to Orbax v1 Checkpointer.") + + +def _load_checkpointables_with_logging( + checkpoint_manager: ocp.training.Checkpointer, + step: int, + abstract_checkpointables: dict[str, Any], +) -> dict[str, Any]: + """Loads checkpointables and logs v0-compatible restore step statistics.""" + orbax_logger = getattr(checkpoint_manager, "orbax_logger", None) + if orbax_logger is None: + return checkpoint_manager.load_checkpointables(step, abstract_checkpointables) + + # Logging is auxiliary: failures here must never block the restore itself. + step_stats = None + try: + step_stats = ocp_logging.step_statistics.RestoreStepStatistics() + step_stats.step = step + step_stats.directory = str(checkpoint_manager.directory) + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.warning(f"Failed to initialize checkpoint restore statistics: {e}") + + start_time = time.time() + restored = checkpoint_manager.load_checkpointables(step, abstract_checkpointables) + end_time = time.time() + + if step_stats is not None: + try: + step_stats.checkpoint_manager_start_time = start_time + step_stats.checkpointer_start_time = start_time + step_stats.checkpointer_duration_secs = end_time - start_time + step_stats.checkpoint_manager_duration_secs = end_time - start_time + orbax_logger.log_entry(dataclasses.asdict(step_stats)) + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.warning(f"Failed to log checkpoint restore statistics: {e}") + return restored + + def print_save_message(step, async_checkpointing): if async_checkpointing: max_logging.log(f"Started an asynchronous checkpoint save for step {step}") @@ -604,7 +830,7 @@ def load_state_if_possible( abstract_checkpointables["iter"] = grain_utility.for_restore( checkpoint_manager, step, data_iterator, expansion_factor_real_data ) - restored = checkpoint_manager.load_checkpointables(step, abstract_checkpointables) + restored = _load_checkpointables_with_logging(checkpoint_manager, step, abstract_checkpointables) if is_diloco: restored_items = diloco_checkpoint_utils.from_diloco_checkpoint_dict( restored["items"], abstract_unboxed_pre_state, config=maxtext_config @@ -641,6 +867,8 @@ def load_state_if_possible( use_ocdbt=use_ocdbt, use_zarr3=use_zarr3, enable_single_replica_ckpt_restoring=bool(enable_single_replica_ckpt_restoring), + convert_flat_moe_weights=bool(getattr(maxtext_config, "moe_flat_fsdp_weights", False)), + colocated_python_checkpointing=bool(getattr(maxtext_config, "colocated_python_checkpointing", False)), ) return None, restored_params elif load_full_state_from_path != "": @@ -663,14 +891,36 @@ def load_state_if_possible( def setup_checkpoint_logger(config) -> Any | None: # pytype: disable=attribute-error - """DEPRECATED: Setup checkpoint logger.""" - # TODO: b/529622681 - Remove this config option entirely. - if config.enable_checkpoint_cloud_logger: + """Sets up the Orbax checkpoint Cloud Logger used by ML Goodput. + + Args: + config: MaxText config. + + Returns: + An Orbax `CloudLogger` writing to the `goodput_` log if + `enable_checkpoint_cloud_logger` is set, otherwise None. + """ + if not config.enable_checkpoint_cloud_logger: + return None + max_logging.log("Setting up checkpoint logger...") + cloud_logger_cls = getattr(ocp_logging, "CloudLogger", None) + cloud_logger_options_cls = getattr(ocp_logging, "CloudLoggerOptions", None) + if cloud_logger_cls is None or cloud_logger_options_cls is None: max_logging.warning( - "Cloud logging (enable_checkpoint_cloud_logger) is disabled because" - " Orbax v1 now configures its own logger internally. This config" - " setting is ignored and will be removed." + "enable_checkpoint_cloud_logger is set but Orbax CloudLogger is unavailable" + " (is google-cloud-logging installed?). Checkpoint logs are disabled." + ) + return None + logger_name = f"goodput_{config.run_name}" + try: + orbax_cloud_logger = cloud_logger_cls( + options=cloud_logger_options_cls(job_name=config.run_name, logger_name=logger_name) ) + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.warning(f"Failed to initialize Orbax CloudLogger: {e}. Checkpoint logs are disabled.") + return None + max_logging.log(f"Successfully set up checkpoint cloud logger: {logger_name}") + return orbax_cloud_logger def _scale_sharding(weight_sharding, scale_shape): @@ -788,6 +1038,107 @@ def maybe_dequantize_restored_params(restored_weights: Any, expected_param_pytre return out +_FLAT_MOE_LEAVES = ("wi_0", "wi_1", "wo") + + +def _flat_moe_restore_plan(want, stored_weights): + """Finds routed-expert weights the model wants flat (moe_flat_fsdp_weights) but the checkpoint stores 3D. + + Returns {path: (flat_target_sds, stored_shape_sds)} for every leaf named wi_0/wi_1/wo whose expected shape is + the flat [E * rows, (L,) cols] view of a stored [E, (L,) rows, cols] array. `stored_shape_sds` carries the stored + shape with the (expert | fsdp)-style sharding derived from the flat target's (expert, fsdp) dim-0 sharding, so + Orbax restores the 3D/4D array directly in a layout that `_flat_moe_convert` turns into the flat one. + """ + plan = {} + + def _walk(w, s, path): + if isinstance(w, dict): + for k, v in w.items(): + _walk(v, s.get(k) if isinstance(s, dict) else None, path + (k,)) + return + if not path or path[-1] not in _FLAT_MOE_LEAVES or s is None: + return + want_shape = tuple(getattr(w, "shape", ()) or ()) + stored_shape = tuple(getattr(s, "shape", ()) or ()) + if not want_shape or not stored_shape or want_shape == stored_shape or len(want_shape) != len(stored_shape) - 1: + return + if len(stored_shape) == 4: # scanned: [E, L, rows, cols] -> [E * rows, L, cols] + e, l, r, c = stored_shape + flat_shape = (e * r, l, c) + elif len(stored_shape) == 3: # unscanned (e.g. MTP): [E, rows, cols] -> [E * rows, cols] + e, r, c = stored_shape + flat_shape = (e * r, c) + else: + return + if want_shape != flat_shape: + return + sharding = getattr(w, "sharding", None) + mesh = getattr(sharding, "mesh", None) + spec = getattr(sharding, "spec", None) + if mesh is None or spec is None: + return + dim0 = spec[0] if len(spec) > 0 else None + dim0_axes = () if dim0 is None else ((dim0,) if isinstance(dim0, str) else tuple(dim0)) + # Dim 0 is (expert, fsdp...) minus size-one axes. With a single axis left we can't tell which it is, so shard + # the stored rows dim on it: rows (embed/mlp) divide any mesh axis, the expert count may not. + expert_axis = dim0_axes[0] if len(dim0_axes) > 1 else None + rest = dim0_axes[1:] if len(dim0_axes) > 1 else dim0_axes + rest_spec = None if not rest else (rest[0] if len(rest) == 1 else rest) + tail = tuple(spec)[1:] + (None,) * max(0, len(flat_shape) - len(spec)) + if len(stored_shape) == 4: + stored_spec = jax.sharding.PartitionSpec(expert_axis, tail[0], rest_spec, tail[1]) + else: + stored_spec = jax.sharding.PartitionSpec(expert_axis, rest_spec, tail[0]) + stored_sds = jax.ShapeDtypeStruct(stored_shape, w.dtype, sharding=jax.sharding.NamedSharding(mesh, stored_spec)) + plan[path] = (w, stored_sds) + + _walk(want, stored_weights, ()) + return plan + + +def _set_path(tree, path, value): + node = tree + for key in path[:-1]: + if not isinstance(node, dict) or key not in node: + return False + node = node[key] + if not isinstance(node, dict) or path[-1] not in node: + return False + node[path[-1]] = value + return True + + +def _flat_moe_convert(x, target): + """Reshards a restored [E, (L,) rows, cols] weight into the flat [E * rows, (L,) cols] target. + + The merge of the expert dim with the FSDP-sharded row dim is a real data movement (each device's flat block is a + contiguous run of one expert's rows), so it is done as a jit with `out_shardings`, one layer at a time to bound the + transient memory, and the layers are stacked back along the scan axis. + """ + out_sharding = target.sharding + mesh = out_sharding.mesh + # PartitionSpec may be shorter than the rank (e.g. P() when replicated); pad with None. + spec = tuple(out_sharding.spec) + (None,) * (len(target.shape) - len(out_sharding.spec)) + if x.ndim == 4: + e, l, r, c = x.shape + layer_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(spec[0], spec[2])) + + def _layer_fn(a, idx): + return jax.lax.dynamic_index_in_dim(a, idx, axis=1, keepdims=False).reshape(e * r, c).astype(target.dtype) + + _layer = jax.jit(_layer_fn, out_shardings=layer_sharding) + layers = [_layer(x, jnp.int32(i)) for i in range(l)] + x.delete() + stacked = jax.jit(lambda *ys: jnp.stack(ys, axis=1), out_shardings=out_sharding)(*layers) + for y in layers: + y.delete() + return stacked + e, r, c = x.shape + y = jax.jit(lambda a: a.reshape(e * r, c).astype(target.dtype), out_shardings=out_sharding)(x) + x.delete() + return y + + def load_params_from_path( load_parameters_from_path, abstract_unboxed_params, @@ -795,8 +1146,15 @@ def load_params_from_path( use_ocdbt=True, use_zarr3=True, enable_single_replica_ckpt_restoring: bool = False, + convert_flat_moe_weights: bool = False, + colocated_python_checkpointing: bool = False, ): - """Load decode params from checkpoint at specified path.""" + """Load decode params from checkpoint at specified path. + + With `convert_flat_moe_weights` (moe_flat_fsdp_weights), 3D/4D routed-expert weights in the checkpoint are converted + to the flat [E * rows, (L,) cols] layout the model expects. With `colocated_python_checkpointing` (Pathways single + controller only), the restore runs on the workers through the Orbax colocated-Python dispatcher. + """ assert load_parameters_from_path, "load_parameters_from_path is not defined." max_logging.log(f"restoring params from {load_parameters_from_path}") @@ -815,6 +1173,103 @@ def load_params_from_path( if restore_key not in ("model_params", "model"): restore_key = "params" + checkpointable_name = "items" if (path / "items").exists() else None + # Orbax v1 fails a mid-load shape mismatch itself, with an error that reports the + # shapes but not which weight; compare the stored metadata first so the error names + # it. A metadata read failure falls through to the load (worst case: Orbax's error). + # With colocated Python, read the metadata under the same checkpointing implementation + # as the load: outside a Context, Orbax on Pathways requires remote Python instead. + metadata_context = ( + checkpoint_context.build_context(colocated_python_checkpointing=True) + if colocated_python_checkpointing + else contextlib.nullcontext() + ) + try: + with metadata_context: + stored = ocp.metadata(path, checkpointable_name=checkpointable_name).metadata + except Exception as e: # pylint: disable=broad-except + max_logging.log(f"Skipping pre-load shape check, checkpoint metadata unreadable: {e}") + stored = None + + # Flat routed-expert weights (moe_flat_fsdp_weights): restore the checkpoint's 3D/4D arrays, convert after the load. + flat_moe_plan = {} + if convert_flat_moe_weights and isinstance(stored, dict): + stored_root = stored.get(restore_key) + # An NNX `want` is bare weights, so drop the stored `params` collection level; a Linen `want` keeps it. + stored_params = ( + stored_root.get("params") + if is_nnx and isinstance(stored_root, dict) and isinstance(stored_root.get("params"), dict) + else stored_root + ) + flat_moe_plan = _flat_moe_restore_plan(want, stored_params) + for wpath, (_, stored_sds) in flat_moe_plan.items(): + _set_path(want, wpath, stored_sds) + if flat_moe_plan: + max_logging.log( + f"moe_flat_fsdp_weights: {len(flat_moe_plan)} routed-expert weights will be restored in the checkpoint's" + f" 3D layout and converted to the flat layout: {['/'.join(map(str, p)) for p in flat_moe_plan]}" + ) + + aliased_collections = {} + if restore_key in ("model_params", "model"): + params_collection = want + elif is_nnx and restore_key == "params": + stored_params_meta = ( + stored.get("params") if isinstance(stored, dict) and isinstance(stored.get("params"), dict) else {} + ) + # `None` means "we could not read which collections the checkpoint holds", which is + # NOT the same as "it holds only `params`". Defaulting to {"params"} would drop every + # custom collection from the request below, turning a transient metadata read failure + # into a hard restore failure for a checkpoint that is perfectly valid. + stored_collections = set(stored_params_meta.keys()) if stored_params_meta else None + + collection_leaves = collections.defaultdict(dict) + for path_tuple, leaf in nnx.to_flat_state(abstract_unboxed_params): + if isinstance(leaf, nnx.Variable): + col = variablelib.variable_name_from_type(leaf.type, allow_register=True) + val = leaf.get_value() if hasattr(leaf, "get_value") else leaf + else: + col = "params" + val = leaf + collection_leaves[col][tuple(path_tuple)] = val + + params_collection = {col: nnx.traversals.unflatten_mapping(d) for col, d in collection_leaves.items()} + if "params" not in params_collection: + params_collection["params"] = want + for wpath, (_, stored_sds) in flat_moe_plan.items(): + for col_tree in params_collection.values(): + _set_path(col_tree, wpath, stored_sds) + + # Augment "params" collection with companion scales from checkpoint metadata if present (FP8 dequantize-on-load) + if "params" in params_collection: + params_collection["params"] = _augment_want_with_scales( + params_collection["params"], stored, is_nnx=True, restore_key="params" + ) + + # Rescue weights whose collection is newer than the checkpoint BEFORE discarding + # unknown collections below -- otherwise they are dropped with no error and stay + # at their initialiser values. See _alias_legacy_collections. + params_collection, aliased_collections = _alias_legacy_collections(params_collection, stored_params_meta) + + # Anything the checkpoint is KNOWN not to hold genuinely is not there; keep dropping it + # so Orbax's partial_load does not hand back unmaterialized leaves. When the collection + # set is unknown, request everything and let the post-load mismatch check adjudicate. + params_collection = { + col: tree + for col, tree in params_collection.items() + if stored_collections is None or col in stored_collections or col == "params" + } + else: + # Only reachable when `is_nnx` is False: `restore_key` is normalized to one of three + # values above, and the preceding branches claim both the non-"params" keys and every + # NNX "params" case. So this is the Linen path, whose `want` is already collection-keyed. + augmented_want = _augment_want_with_scales(want, stored, is_nnx=False, restore_key=restore_key) + params_collection = augmented_want + if isinstance(stored, dict): + # Linen keeps its structure keyed by collection, so the same rescue applies here; + # the restored tree is un-aliased again after the load. + params_collection, aliased_collections = _alias_legacy_collections(params_collection, stored.get(restore_key)) + # Memory optimization: restore only the "params" key (the checkpoint may also hold opt_state/step); # partial_load drops the rest. The abstract carries shape/dtype/sharding directly. context = checkpoint_context.build_context( @@ -823,31 +1278,29 @@ def load_params_from_path( checkpoint_storage_concurrent_gb=checkpoint_storage_concurrent_gb, partial_load=True, enable_single_replica_ckpt_restoring=enable_single_replica_ckpt_restoring, + colocated_python_checkpointing=colocated_python_checkpointing, ) # Dispatch on the on-disk layout instead of assuming a step root: callers pass step roots, # v0-style pytree dirs (normalized above), and v0 flat params-only checkpoints # (save_params_to_path wrote the pytree directly at the directory). + audit_stored = None with context: - checkpointable_name = "items" if (path / "items").exists() else None - # Orbax v1 fails a mid-load shape mismatch itself, with an error that reports the - # shapes but not which weight; compare the stored metadata first so the error names - # it. A metadata read failure falls through to the load (worst case: Orbax's error). - try: - stored = ocp.metadata(path, checkpointable_name=checkpointable_name).metadata - except Exception as e: # pylint: disable=broad-except - max_logging.log(f"Skipping pre-load shape check, checkpoint metadata unreadable: {e}") - stored = None - - augmented_want = _augment_want_with_scales(want, stored, is_nnx, restore_key) - if restore_key in ("model_params", "model"): - params_collection = augmented_want - else: - params_collection = {"params": augmented_want} if is_nnx else augmented_want - if isinstance(stored, dict): stored_collection = stored.get(restore_key) + audit_stored = stored_collection if restore_key == "params" and is_nnx and isinstance(stored_collection, dict): - stored_collection = stored_collection.get("params") + # Everything the checkpoint actually holds, captured BEFORE the request-time + # filtering below, which keeps only the collections the model asked for. The + # unconsumed audit needs the unfiltered view or it cannot see a dropped weight. + audit_stored = {} + for col_tree in stored_collection.values(): + if isinstance(col_tree, dict): + audit_stored = _deep_merge_dicts(audit_stored, col_tree) + merged_stored = {} + for col_name in params_collection.keys(): + if col_name in stored_collection and isinstance(stored_collection[col_name], dict): + merged_stored = _deep_merge_dicts(merged_stored, stored_collection[col_name]) + stored_collection = merged_stored _raise_weight_problems(_weight_mismatches(want, stored_collection, check_missing=False)) restored = ocp.load( path, @@ -855,6 +1308,11 @@ def load_params_from_path( checkpointable_name=checkpointable_name, # pyrefly: ignore[bad-argument-type] ) restored_collection = restored[restore_key] # pyrefly: ignore[bad-index] + # Weights rescued out of a newer collection came back inside `params`; put them back + # where the model expects them. The NNX branch below merges every collection by path, + # so it needs no un-aliasing. + if not is_nnx and aliased_collections: + restored_collection = _unalias_legacy_collections(restored_collection, aliased_collections) # partial_load lets Orbax return an unmaterialized leaf for a weight the checkpoint lacks, # and a stored array at its own shape rather than the target's. Either reaches the model and # fails much later without naming the weight, so check here -- the params-only load @@ -862,12 +1320,49 @@ def load_params_from_path( if restore_key in ("model_params", "model"): restored_weights = restored_collection + elif is_nnx: + restored_weights = {} + if isinstance(restored_collection, dict): + # Process "params" first so custom collection weights always overwrite stale legacy copies. + for col_key in sorted(restored_collection, key=lambda k: (k != "params", k)): + col_val = restored_collection[col_key] + if isinstance(col_val, dict): + restored_weights = _deep_merge_dicts(restored_weights, col_val) + else: + restored_weights[col_key] = col_val + else: + restored_weights = restored_collection else: - restored_weights = restored_collection["params"] if is_nnx else restored_collection + # Linen: `restored_collection` is already the collection mapping that Linen's + # `want` is keyed by (e.g. {"params": ..., "Tid2EidVar": ...}). Do NOT unwrap + # the "params" key here -- that drops every other collection and shifts the + # remaining weights up a level, so the mismatch check reports the entire model + # as missing. Only the NNX branch above needs per-collection flattening. + restored_weights = restored_collection + # Convert the restored 3D/4D routed-expert weights to the flat layout the model expects. + for wpath, (flat_sds, stored_sds) in flat_moe_plan.items(): + arr = _lookup_path(restored_weights, wpath) + if not isinstance(arr, jax.Array): + continue + t0 = time.time() + converted = _flat_moe_convert(arr, flat_sds) + assert _set_path(restored_weights, wpath, converted) + _set_path(want, wpath, flat_sds) + max_logging.log( + f"moe_flat_fsdp_weights: converted {'/'.join(map(str, wpath))} {tuple(stored_sds.shape)} ->" + f" {tuple(converted.shape)} in {time.time() - t0:.1f}s" + ) # Dequantize if checkpoint had companion kernel_scale and target want is unquantized. restored_weights = maybe_dequantize_restored_params(restored_weights, want) + # Validate against everything the model needs. Rebuilding the expectation from the + # post-filter request would excuse precisely the weight whose collection the checkpoint + # lacked -- which is the silent corruption this check exists to catch, not an exemption. _raise_on_weight_mismatch(want, restored_weights) + # Answer the other half of the question: did the model consume everything the checkpoint + # offered? Compare against what is ON DISK; comparing against the restored tree can never + # differ, because the restore only ever returns what was requested. + _log_unconsumed_checkpoint_weights(want, audit_stored) if is_nnx: nnx.replace_by_pure_dict(abstract_unboxed_params, restored_weights) return abstract_unboxed_params diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index c8147d36d2..a1617a5113 100644 --- a/src/maxtext/common/common_types.py +++ b/src/maxtext/common/common_types.py @@ -151,6 +151,7 @@ class DecoderBlockType(enum.Enum): OLMO3 = "olmo3" DEEPSEEK4 = "deepseek4" ENVY = "envy" + WEAVER = "weaver" class VisionEncoderBlockType(enum.Enum): @@ -202,9 +203,12 @@ class CustomRule(enum.Enum): DEFAULT = "" PURE_FSDP = "pure-fsdp" CP_AS_EP = "cp-as-ep" # Support CP and EP together + CP_AS_EP_DENSE_NO_CP = "cp-as-ep-dense-no-cp" # cp-as-ep with the dense weights on fsdp x expert, not context + CP_AS_EP_DENSE_FSDP = "cp-as-ep-dense-fsdp" # cp-as-ep with the dense weights on fsdp only EP_AS_CP = "ep-as-cp" # Support EP only PIPELINE_LARGE_MOE = "pipeline-large-moe" FSDP_2D = "2d-fsdp" EP_AS_DP = "ep-as-dp" SHARD_EXP_ON_FSDP = "shard-exp-on-fsdp" FSDP_AS_DP_FOR_ATTN = "fsdp-as-dp-for-attn" + FSDP_AS_DP_FOR_ATTN_RF = "fsdp-as-dp-for-attn-rf" # fsdp-as-dp-for-attn with the router weight sharded on fsdp diff --git a/src/maxtext/common/data_loader.py b/src/maxtext/common/data_loader.py index 410dad7ddf..cffc68d8e1 100644 --- a/src/maxtext/common/data_loader.py +++ b/src/maxtext/common/data_loader.py @@ -44,9 +44,13 @@ def loader_exception_guard(config): except Exception as e: # pylint: disable=broad-except elastic_utils.maybe_bubble_elastic_exception(config, e) if isinstance(e, StopIteration): - raise exceptions.StopTraining(f"You may have run out of training data. Received {type(e)}" f" exception: ({e})") + raise exceptions.StopTraining( + f"You may have run out of training data. Received {type(e)}" f" exception: ({e})" + ) from e else: - raise exceptions.StopTraining(f"`next(self.data_iterator)` failed with {type(e)} exception: ({e}).") + raise exceptions.StopTraining( + f"`next(self.data_iterator)` failed with {type(e)} exception: ({e}).", + ) from e class DataLoader: diff --git a/src/maxtext/common/gcloud_stub.py b/src/maxtext/common/gcloud_stub.py index 52c39a8abb..cec3ccb43f 100644 --- a/src/maxtext/common/gcloud_stub.py +++ b/src/maxtext/common/gcloud_stub.py @@ -314,6 +314,12 @@ def start_goodput_uploader(self): def start_step_deviation_uploader(self): print("[DECOUPLED NO-OP] goodput step deviation uploader skipped.") + def start_rolling_window_goodput_uploader(self, *_a, **_k): + print("[DECOUPLED NO-OP] rolling window goodput uploader skipped.") + + def stop_rolling_window_goodput_uploader(self, *_a, **_k): + print("[DECOUPLED NO-OP] rolling window goodput uploader stop skipped.") + monitoring_ns = SimpleNamespace(GCPOptions=_StubMonitoringOptions, GoodputMonitor=_StubGoodputMonitor) goodput_ns = SimpleNamespace(GoodputRecorder=_StubGoodputRecorder) return goodput_ns, monitoring_ns, True diff --git a/src/maxtext/common/goodput.py b/src/maxtext/common/goodput.py index cf4d8c0043..2ba5fd23ca 100644 --- a/src/maxtext/common/goodput.py +++ b/src/maxtext/common/goodput.py @@ -20,14 +20,18 @@ """ import contextlib -import jax from enum import Enum -from maxtext.utils import max_logging +import jax from maxtext.common.gcloud_stub import goodput_modules +from maxtext.utils import max_logging goodput, monitoring, _GOODPUT_STUB = goodput_modules() +def get_goodput_job_name(config) -> str: + return config.goodput_job_name or config.run_name + + class GoodputEvent(Enum): JOB = "job" TPU_INIT = "tpu_init" @@ -44,6 +48,7 @@ class GoodputEvent(Enum): def _construct_goodput_monitor(config, common_kwargs): """Constructs a GoodputMonitor, preferring the elastic-aware monitor when applicable.""" + job_name = common_kwargs["job_name"] if config.elastic_enabled: try: from maxtext.utils import elastic_utils # pylint: disable=import-outside-toplevel @@ -52,13 +57,13 @@ def _construct_goodput_monitor(config, common_kwargs): from ml_goodput_measurement import monitoring_elastic # pylint: disable=import-outside-toplevel monitor = monitoring_elastic.ElasticGoodputMonitor(include_slice_efficiency=True, **common_kwargs) - max_logging.log(f"Goodput: using ElasticGoodputMonitor for job: {config.run_name}") + max_logging.log(f"Goodput: using ElasticGoodputMonitor for job: {job_name}") return monitor except Exception as e: # pylint: disable=broad-exception-caught max_logging.log(f"Goodput: could not create elastic goodput monitor, falling back to base monitor: {e}") monitor = monitoring.GoodputMonitor(pathway_enabled=config.enable_pathways_goodput, **common_kwargs) - max_logging.log(f"Goodput: using GoodputMonitor for job: {config.run_name}") + max_logging.log(f"Goodput: using GoodputMonitor for job: {job_name}") return monitor @@ -87,9 +92,12 @@ def maybe_monitor_goodput(config): enable_gcp_step_deviation_metrics=config.enable_gcp_step_deviation_metrics, ) + job_name = get_goodput_job_name(config) + logger_name = f"goodput_{job_name}" + common_kwargs = { - "job_name": config.run_name, - "logger_name": f"goodput_{config.run_name}", + "job_name": job_name, + "logger_name": logger_name, "tensorboard_dir": config.tensorboard_dir, "upload_interval": config.goodput_upload_interval_seconds, "monitoring_enabled": True, @@ -101,12 +109,42 @@ def maybe_monitor_goodput(config): goodput_monitor = _construct_goodput_monitor(config, common_kwargs) goodput_monitor.start_goodput_uploader() - max_logging.log("Started Goodput upload to Tensorboard & GCM in the background!") - yield + max_logging.log("Started Cumulative Goodput upload to Tensorboard & GCM in the background!") + with maybe_monitor_rolling_window_goodput(goodput_monitor, config): + yield finally: if goodput_monitor: goodput_monitor.stop_goodput_uploader() - max_logging.log("Flushed final metrics and safe exited from Goodput monitoring.") + max_logging.log("Flushed final metrics and safe exited from cumulative Goodput monitoring.") + + +@contextlib.contextmanager +def maybe_monitor_rolling_window_goodput(goodput_monitor, config): + """Monitor rolling window goodput if enabled on the lead host. + + Failures to start or stop the rolling window uploader are logged and + swallowed so that monitoring degrades gracefully and never interrupts + training. Exceptions raised by the wrapped block are not caught here. + """ + if _GOODPUT_STUB or not goodput_monitor or not config.enable_rolling_window_goodput: + yield + return + started = False + try: + goodput_monitor.start_rolling_window_goodput_uploader(config.rolling_windows_seconds) + started = True + max_logging.log("Started Rolling Window Goodput upload to Tensorboard & GCM in the background!") + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.log(f"Goodput: failed to start rolling window goodput uploader, continuing without it: {e}") + try: + yield + finally: + if started: + try: + goodput_monitor.stop_rolling_window_goodput_uploader() + max_logging.log("Flushed final metrics and safe exited from rolling window Goodput monitoring.") + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.log(f"Goodput: failed to stop rolling window goodput uploader cleanly: {e}") @contextlib.contextmanager @@ -147,7 +185,8 @@ def create_goodput_recorder(config): if not config.enable_goodput_recording: return None - logger_name = f"goodput_{config.run_name}" + job_name = get_goodput_job_name(config) + logger_name = f"goodput_{job_name}" # Detect if we should use the elastic-aware recorder if config.elastic_enabled: @@ -157,13 +196,13 @@ def create_goodput_recorder(config): if elastic_utils.should_use_elastic(config): from ml_goodput_measurement import goodput_elastic # pylint: disable=import-outside-toplevel - recorder = goodput_elastic.ElasticGoodputRecorder(config.run_name, logger_name, jax.process_index() == 0) + recorder = goodput_elastic.ElasticGoodputRecorder(job_name, logger_name, jax.process_index() == 0) elastic_utils.record_slice_state(recorder) - max_logging.log(f"Goodput: created ElasticGoodputRecorder for job: {config.run_name}") + max_logging.log(f"Goodput: created ElasticGoodputRecorder for job: {job_name}") return recorder except Exception as e: # pylint: disable=broad-exception-caught max_logging.log(f"Goodput: could not create elastic goodput recorder, falling back to base recorder: {e}") - recorder = goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0) - max_logging.log(f"Goodput: created base GoodputRecorder for job: {config.run_name}") + recorder = goodput.GoodputRecorder(job_name, logger_name, jax.process_index() == 0) + max_logging.log(f"Goodput: created base GoodputRecorder for job: {job_name}") return recorder diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py index 2d6fa882f1..2aa399e5fb 100644 --- a/src/maxtext/common/metric_logger.py +++ b/src/maxtext/common/metric_logger.py @@ -56,6 +56,25 @@ } +STEP_DIAGNOSTICS_KEYS = ( + ("raw_grad_norm", "learning/raw_grad_norm"), + ("grad_norm", "learning/grad_norm"), + ("param_norm", "learning/param_norm"), + ("max_abs_grad", "diag/max_abs_grad"), + ("moe_bias_checksum", "diag/moe_bias_checksum"), + ("moe_bias_update_nonzero", "diag/moe_bias_update_nonzero"), + ("moe_overflow", "diag/moe_overflow"), +) + + +def step_diagnostics_text(scalars): + """The log_step_diagnostics part of the step line: 'diag: raw_grad_norm=... grad_norm=... ...' (%.9e; n/a if absent).""" + parts = [] + for name, key in STEP_DIAGNOSTICS_KEYS: + parts.append(f"{name}={float(scalars[key]):.9e}" if key in scalars else f"{name}=n/a") + return "diag: " + " ".join(parts) + + def _prepare_metrics_for_json(metrics, step, run_name): """Converts metric dictionary into json supported types (e.g. float)""" metrics_dict = {val: float(metrics["scalar"][val]) for val in metrics["scalar"]} @@ -227,6 +246,9 @@ def _log_training_metrics(self, metrics, step): if self.config.num_experts > 1: moe_lb_loss = scalars.get("learning/moe_lb_loss", 0.0) log_parts.append(f"moe_lb_loss: {moe_lb_loss:.6f}") + if "learning/moe_max_load_ratio" in scalars: + log_parts.append(f"moe_max_load_ratio: {scalars['learning/moe_max_load_ratio']:.4f}") + log_parts.append(f"moe_max_load_ratio_mean: {scalars['learning/moe_max_load_ratio_mean']:.4f}") if getattr(self.config, "mtp_num_layers", 0) > 0: mtp_loss = scalars.get("learning/mtp_loss", 0.0) @@ -237,8 +259,20 @@ def _log_training_metrics(self, metrics, step): indexer_l = scalars.get("learning/indexer_loss", 0.0) log_parts.append(f"indexer_loss: {float(indexer_l)}") + if getattr(self.config, "log_step_diagnostics", False): + log_parts.append(step_diagnostics_text(scalars)) + max_logging.log(", ".join(log_parts)) + # Logged here rather than in the train loop: metrics are flushed one step late, so reading + # the flag does not add a per-step host sync (which the in-layer fallback exists to avoid). + fallback_layers = scalars.get("learning/moe_dropless_fallback_layers") + if fallback_layers is not None and int(fallback_layers) > 0: + max_logging.log( + f"Step {step}: MoE ragged buffer overflow detected in {int(fallback_layers)} layer(s); " + f"in-layer dropless fallback was used (no tokens dropped)." + ) + def _log_eval_metrics(self, metrics, step): """Logs the final accumulated eval summary at the end of an eval run.""" scalars = metrics["scalar"] @@ -267,6 +301,13 @@ def _log_eval_metrics(self, metrics, step): max_logging.log(", ".join(log_parts)) + fallback_layers = int(scalars.get("eval/moe_dropless_fallback_layers", 0)) + if fallback_layers > 0: + max_logging.log( + f"Eval after train step {step}: MoE ragged buffer overflow detected in {fallback_layers} layer(s) " + f"summed over eval steps; in-layer dropless fallback was used (no tokens dropped)." + ) + def _log_running_eval_metrics(self, metrics, step): """Logs a per-eval-step running average (deferred by one eval step).""" scalars = metrics["scalar"] @@ -488,6 +529,10 @@ def _accumulate_eval_metrics(self, metrics): scalar.get("evaluation/mtp_acceptance_rate_percent", 0.0) ) self.cumulative_eval_metrics["scalar"]["eval/z_loss"] += float(scalar.get("evaluation/z_loss", 0.0)) + if "evaluation/moe_dropless_fallback_layers" in scalar: + self.cumulative_eval_metrics["scalar"]["eval/moe_dropless_fallback_layers"] += float( + scalar["evaluation/moe_dropless_fallback_layers"] + ) def record_train_metrics(self, metrics, step, step_time): """Records training metrics for the current step.""" diff --git a/src/maxtext/common/profiler.py b/src/maxtext/common/profiler.py index f6f56c372e..8fda452751 100644 --- a/src/maxtext/common/profiler.py +++ b/src/maxtext/common/profiler.py @@ -72,6 +72,9 @@ def __init__(self, config, offset_step=0): } ) + if self.mode == "xplane" and config.enable_continuous_profiling: + advanced_config["enable_continuous_profiling"] = True + if advanced_config: self.profiling_options.advanced_configuration = advanced_config diff --git a/src/maxtext/common/vertex_tensorboard.py b/src/maxtext/common/vertex_tensorboard.py index ce36efff1d..f098992c48 100644 --- a/src/maxtext/common/vertex_tensorboard.py +++ b/src/maxtext/common/vertex_tensorboard.py @@ -121,11 +121,11 @@ def configure_vertex_tensorboard(self, config): if not os.environ.get("EXPERIMENT_NAME"): os.environ["EXPERIMENT_NAME"] = config.run_name - if config.use_vertex_tensorboard: # running MaxText on GCE + if config.use_vertex_tensorboard: # MaxText creates/reuses the Tensorboard instance and Experiment tensorboard_url = self.setup() if tensorboard_url is None: raise ValueError("Unable to create Tensorboard and Experiment in Vertex AI.") max_logging.log(f"View your Vertex AI Tensorboard at: {tensorboard_url}") self.upload_data(config.tensorboard_dir) - elif os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"): # running MaxText via XPK + elif os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"): # instance and Experiment created outside MaxText self.upload_data(config.tensorboard_dir) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 08d79afd09..6a76b32413 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -21,6 +21,7 @@ run_name: "" model_name: "default" # override config settings to match a specific model. other than the override, nothing should use this! +use_m3_model: false # Use the m3 backend; requires a registered m3 model. override_model_config: false # When set to true allows overriding model parameters via CLI (or kwargs or env vars) for the purpose of debugging/testing. override_logical_axis_rules: false # When set overrides logical axis rules instead of merging them. @@ -184,6 +185,10 @@ kv_quant_axis: "heads_and_dkv" kv_quant_dtype: "int8" quantize_mtp: false # set to true to quantize mtp layers when quantization is enabled. requires mtp_num_layers > 0 and quantization=fp8_full. quantize_router_proj: true # set to false to exclude the moe router (gate) projection matmul from quantization when quantization is enabled. default is true for backward compatibility. applicable when use_qwix_quantization=true and quantization is set; ignored otherwise. +quantize_logits_proj: false # set to true to quantize output logits projection (logits_dense) when quantization is enabled. requires quantization=fp8_full, logits_via_embedding=false, num_vocab_tiling=1. also covers the mtp logits projection, which reuses the same head. +logits_proj_quant_calibration_method: "" # calibration method for logits_dense when quantize_logits_proj=true. if empty, inherits weight/act_quantization_calibration_method. e.g. "absmax" +logits_proj_bwd_quant_calibration_method: "" # bwd (gradient) calibration method for logits_dense when quantize_logits_proj=true, applied to both dlhs and drhs gradients. if empty, inherits bwd_quantization_calibration_method and drhs_grad_quantization_calibration_method. e.g. "absmax" +sa_qproj_bwd_bf16: false # keep self-attention query projection backward pass in bfloat16 instead of fp8 when fp8_full quantization is enabled. checkpoint_is_quantized: false # set to true if reading from a saved aqt quantized checkpoint # saves params quantized on fly at following path save_quantized_params_path: "" @@ -197,8 +202,15 @@ use_manual_quantization: false # a flag if to use manual quantization for batch weight_quantization_calibration_method: "absmax" act_quantization_calibration_method: "absmax" bwd_quantization_calibration_method: "absmax" -# shard the range finding operation for quantization. by default this is set to number of slices. +drhs_grad_quantization_calibration_method: null # weight-gradient-arm cotangent calibration; null = bwd_quantization_calibration_method +# AQT only (use_qwix_quantization=false, quantization=int8): number of local shards along the contraction +# axis of weight-gradient matmuls, each with its own scale. -1 = jax.local_device_count(); 0 = do not +# quantize the weight-gradient matmul. Qwix uses qwix_bwd_weight_grad_tile_count instead. quantization_local_shard_count: -1 +# Qwix only (use_qwix_quantization=true): number of tiles along the contraction axis of weight-gradient matmuls, each with its own scale. +# 1 = no tiling; >1 makes XLA all-reduce the full weight gradient under FSDP instead of reduce-scattering it. +# Applies to quantization=int8, fp8, etc.; ignored by fp8_full, which never tiles weight gradients (same as 1). +qwix_bwd_weight_grad_tile_count: 1 # The 'N' in N:M sparsity, representing the maximum number of non-zero values in each block. weight_sparsity_n: null # The 'M' in N:M sparsity, representing the number of values in each block. @@ -265,6 +277,8 @@ mtp_loss_scaling_factor: 0.1 # acceptance rate during evaluation. for example, a value of `1` targets the # first auxiliary prediction head. set to 0 to disable this specific evaluation mtp_eval_target_module: 0 +# reuse the main decoder's token embeddings instead of a second lookup +mtp_reuse_input_embedding: false # mixture of experts (moe) num_experts: 1 @@ -273,23 +287,51 @@ megablox: true sparse_matmul: true capacity_factor: -1.0 # a factor to decide expert capacity for token dropping, and no dropping by default num_moe_token_chunks: 1 # split routed-MoE tokens into N chunks (ring-of-experts path) so each chunk's EP all-gather / reduce-scatter overlaps the previous chunk's GMM compute. 1 = disabled (identical to baseline). +moe_accumulate_chunk_wgrad: false # with num_moe_token_chunks>1, gather/quantize MoE weights once and accumulate their gradients in place across chunks via tgmm_v2 partial_sum. requires use_tokamax_gmm=true, use_gmm_v2=true, num_moe_emb_chunks=0, prefuse_moe_weights=false. +moe_accumulate_wi_dlhs: false # accumulate the MoE wi_1 backward DLHS into the wi_0 DLHS in place (gmm_v2 partial_sum) and fold per-tensor DLHS scales into the kernel. Requires use_tokamax_gmm=True, use_gmm_v2=True, prefuse_moe_weights=False. +moe_gmm_v2_dlhs_transpose_rhs: false # use native transpose_rhs in GMM v2 backward DLHS; false keeps the explicit RHS transpose +moe_flat_fsdp_weights: false # store routed-expert weights as 2D [num_experts * rows, cols] sharded (expert, fsdp) on dim 0 to eliminate FSDP all-gather / reduce-scatter relayout copies. +moe_bwd_prequant_before_unsort: false # fuse TC unsort-reduce into the wo GMM custom VJP so backward quantization runs on the smaller [N_tokens, D // 128, 128] cotangent before the TC unsort-gather into the expert buffer. With quantization, needs fixed weight, act and bwd calibration. +moe_combine_bwd_direct_qarray: false # with moe_bwd_prequant_before_unsort, fuse the EP combine psum_scatter into the wo GMM custom VJP so backward quantizes the local cotangent once before the EP all-gather and TC unsort-gather. Incompatible with moe_quantize_combine_bwd_method. +moe_3d_dlhs_partial_sum: false # skip redundant zero-init in 3D DLHS gmm_v2 when accumulating onto an existing 3D partial_sum (e.g. wi_0 DLHS accumulating onto wi_1 DLHS). moe_chunk_barrier: false # wraps its iteration with optimization_barrier to chain the chunked ring-of-experts MoE loop so each chunk's input is fenced on the previous chunk's output, forcing sequential (no-interleave) chunk execution. Math unchanged (barrier is identity), loss bit-exact. Needs num_moe_token_chunks>1 + use_ring_of_experts=True to matter. moe_chunk_pipeline: false # software-pipeline the chunked ring-of-experts loop: chunk c+1's EP token all-gather runs under chunk c's GMMs and chunk c's gather-reduce + reduce-scatter under chunk c+1's, via optimization barriers (math unchanged; skipped with a warning when --xla_tpu_aggressive_opt_barrier_removal=true or ENABLED applies). Needs num_moe_token_chunks>1 + use_ring_of_experts=True. +moe_log_max_load_ratio: false # debug: log learning/moe_max_load_ratio{,_mean} (expert-shard load / balanced load) on the ring-of-experts ragged path; overflow iff max > ragged_buffer_factor. ragged_buffer_factor: -1.0 # a factor to determine the size of the ragged buffer for routed MoE activations. moe_exclude_padding_from_routing: false # train only: tokens with decoder_segment_ids == 0 (packing/padding) take no expert slots (zero group_sizes / ragged-buffer contribution, replayed ids ignored) and get a zero MoE output. Non-pad routing and outputs unchanged. # By default (-1), the routed buffer is worst case size to ensure no dropping. # When set to 1.0 this buffer if set to the size assuming perfectly balanced. If the routing dictates # a size larger than this then tokens are dropped. # In general if ragged_buffer_factor > 0, the ragged_buffer_size is balanced_size * ragged_buffer_factor. +eval_ragged_buffer_factor: -1.0 # evaluation override for ragged_buffer_factor. +# When the ragged buffer would drop tokens: None = drop them; step = replay the whole train step with a +# dropless buffer; layer = per MoE layer, lax.cond into a chunked dropless branch (no step replay). +# Valid values: None, step, layer +moe_dropless_fallback: None +retry_dropless_first_steps: 0 # with moe_dropless_fallback=step or retry_when_tokens_dropped=true, run steps with step - start_step < N directly with the dropless program (no normal attempt). 0 = off. +first_phase_ragged_buffer_factor: 0.0 # with moe_dropless_fallback=step or retry_when_tokens_dropped=true and retry_dropless_first_steps=N, the first N steps use this ragged buffer factor (drops replay dropless). 0 = off. +log_required_ragged_buffer_factor: false # probe: log per step and MoE layer the min ragged_buffer_factor avoiding drops (REQUIRED_RBF). +warmup_programs_in_init: false # before init_stop/run_start, run each precompiled train/eval program once on a synthetic batch (state unchanged), then a cross-host barrier. +warmup_skip_train_dropless: false # requires warmup_programs_in_init=true; do not warm the train_dropless program (its warmup holds a second copy of the train state in HBM and can OOM). Use only when that warmup OOMs; the first dropless replay then pays the program's first load. +log_step_diagnostics: false # append grad/param norms, max |grad|, routed-bias checksum, nonzero bias updates and the overflow flag to every training step log line retry_when_tokens_dropped: false # retry the entire train step with worst-case ragged buffer size if tokens would otherwise be dropped. retry_num_moe_token_chunks: 2 # token chunks of the dropless replay program, run one after another. -retry_dropless_first_steps: 0 # with retry_when_tokens_dropped=true, run the first N steps directly dropless (or with first_phase_ragged_buffer_factor). 0 = off. Training engine only. -first_phase_ragged_buffer_factor: 0.0 # with retry_when_tokens_dropped=true and retry_dropless_first_steps=N, the first N steps use this ragged buffer factor (drops replay dropless). 0 = off. Training engine only. retry_dropless_final_micro_batch: false # with retry_when_tokens_dropped=true, the final micro-batch of each step (apply_optimizer=True) runs the dropless (rbf=-1) program directly; overflowing earlier micro-batches still replay dropless. log_moe_drops: false # drop audit: sow moe_dropped_assignments / moe_max_load_ratio per MoE layer (ring-of-experts ragged path); the engine logs MOE_DROP_AUDIT per micro-batch. moe_expert_input_dim: -1 # feature dimension of the tokens entering the MoE expert blocks. base_moe_mlp_dim: -1 # intermediate dimension at MoE layer. For a fully MoE model, base_mlp_dim must be equal to base_moe_mlp_dim. prefuse_moe_weights: false # whether to pre-fuse MoE weights (wi_0 and wi_1 into wi) during initialization +# With prefuse_moe_weights on the sparse-matmul training path: keep wi = [w0 | w1] one tensor through the fp8 +# fake-quant, the FSDP all-gather and the prefetch barrier into the up-projection GEMM (no split + concatenate). +moe_unsplit_prefused_wi: false +# With prefuse_moe_weights on the sparse-matmul path: the activation on the whole [gate | up] up-projection output as +# one custom-VJP op whose backward writes d[gate | up] in one pass. Same values. +moe_prefused_swiglu_vjp: false +# The MoE SwiGLU in the wi GMM epilogue (megablox gmm_silu, tokamax gmm_v2): the forward kernel writes the +# [buffer, mlp] activation instead of the [buffer, 2 * mlp] GMM output (differentiated: act, gate and up, i.e. +# [buffer, 3 * mlp]). Numerics change: silu reads the float32 accumulator. +# Needs prefuse_moe_weights, use_tokamax_gmm and use_gmm_v2, silu+linear without clipping, no bias or quantization. +moe_wi_gmm_fused_swiglu: false fp8_moe: false # whether MoE weights are quantized to float8_e4m3fn in inference/rollout # Keep routed-expert weights in weight_dtype but round them in the forward pass to the per-channel FP8 grid # fp8_moe stores them on, with straight-through gradients: the trainer computes with the rollout's weights. @@ -297,19 +339,31 @@ fp8_moe_fake_quant: false # Trainer side: the rollout serves fp8_moe experts, so weight sync quantizes them (and emits their scales). rollout_fp8_moe: false load_balance_loss_weight: 0.0 # weight for the load balance loss +moe_use_megatron_seq_aux_loss: false # when true, sigmoid routers (routed_score_func="sigmoid") use Megatron-LM's sequence-wise aux loss summed over MoE layers; false (default) uses the Switch-style loss averaged over MoE layers. use_random_routing: false # whether to use random routing for debug/test purpose use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul +router_topk_matmul_vjp: false # deepseek routing: dense one-hot VJP for the top-k weight gather (no scatter-add in the backward) use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism ring_of_experts_local_routing: false # ring of experts: route each shard's own tokens, then all-gather only the top-k expert ids and weights (not the router logits) +# With ring_of_experts_local_routing + moe_chunk_pipeline: route every token chunk once and all-gather ids + weights as +# ONE float32 array per layer (tagged moe_routing_maps) instead of two gathers per chunk. Same routing. +moe_roe_packed_routing_gather: false +# Ring of experts with ragged sort: activation + down projection + unsort as one op whose backward reads the down +# projection's input, not its output (no recomputed wo GEMM under remat). Numerics change: gradients to bf16 rounding. +moe_router_grad_from_wo_input: false ring_of_experts_row_major_reduce_scatter: false # ring of experts: lay out the expert-parallel reduce-scatter input row-major in the forward pass, so XLA does not copy the combine output into another layout and the result back +moe_topk_before_ep_all_gather: false # ring of experts: run top-k on local tokens and all-gather top-k weights/ids instead of router logits te_moe_block: false # whether to use TransformerEngine's MoEBlock for routed MoE layers te_ep_overflow_check_every_n_steps: 20 # host check interval for TE EP overflow when ragged_buffer_factor limits receive capacity te_gmm_quantization: "" # quantization mode for TE GMM matmuls; available options are te_no_quant and te_mxfp8; must be specified when te_moe_block is true -# Whether to quantize token activations before the All-Gather across EP shards in Ring of Experts, reducing -# inter-chip token traffic proportionally to the quantized width (2x for FP8 vs BF16). +# Whether to quantize token activations before the All-Gather across EP shards in Ring-of-Experts forward pass +# Reduces inter-chip token traffic proportionally to the quantized width (2x for FP8 vs BF16). # Pipeline: quantize token -> EP all-gather -> sort -> GMM, so tokens stay quantized end to end. # Requires use_ring_of_experts=true, use_gmm_v2=true, and quantization with fixed act calibration. moe_quantize_token_all_gather: false +# Quantization method for output gradients before the All-Gather across EP shards in Ring-of-Experts backward pass (Token Combine) +# Value: "" = unquantized BF16; "rowwise" = dynamic token-wise FP8 E5M2; "fixed," = static per-tensor FP8 E5M2 +moe_quantize_combine_bwd_method: "" num_moe_emb_chunks: 0 # number of chunks for overlapping token all-gather and GMM computation along embedding dimension # If true, peel the 'expert' mesh axis off the MoE dispatch/MLP batch dim so the expert GEMM # stays expert-parallel (AllToAll); false keeps 'expert' on the batch dim (activation_batch_moe). @@ -318,6 +372,24 @@ moe_dispatch_no_expert_sharding: false use_ragged_sort: false # whether to use the Pallas ragged-sort kernels in the MoE permute path; valid both with and # without `use_ring_of_experts` (with EP > 1). When `use_ring_of_experts=True` the kernels run # inside `permute`/`unpermute`; otherwise they run inside `local_permute`/local-unpermute. +# Ring-of-experts ragged sort routing index maps: 'sort' = two stable argsorts (default); 'count' = one packed +# single-operand sort plus counting for the inverse (and argmax passes for the top-k of non-forced tokens under +# router replay / padding exclusion). Same permutation, same values. +# Needs sparse_matmul, use_ring_of_experts and use_ragged_sort (and te_moe_block=False). +moe_routing_index_math: 'sort' +# With moe_routing_index_math='sort': the ragged sort's argsorts (and the SparseCore gather-reduce validity sort) as +# unstable two-operand sorts of unique (key, index) pairs instead of stable argsorts (three operands). Same values. +moe_ragged_sort_unique_pair_argsort: false +moe_tc_ragged_sort: false # ring-of-experts ragged sort/unsort on TensorCore (Pallas DMA kernels) instead of the SparseCore + # kernels; requires use_ring_of_experts, use_ragged_sort and ragged_buffer_factor > 0. +moe_tc_ragged_gather_block_size: 1024 # moe_tc_ragged_sort: rows per grid step of the TC ragged gather. +moe_tc_ragged_reduce_block_size: 896 # moe_tc_ragged_sort: VMEM rows per block of the TC ragged gather-reduce. +moe_tc_ragged_mask_padding: true # moe_tc_ragged_sort: zero the gather buffer rows past the valid count (left uninitialized otherwise). +moe_tc_ragged_weights_on_activation: false # moe_tc_ragged_sort: apply routing weights on the (buffer, mlp) activation before wo instead of on the expert output before the unsort. +moe_tc_ragged_flatten_block_size: 0 # moe_tc_ragged_sort: if > 0, Pallas (N, D/128, 128) <-> (N, D) relayout over the valid rows, this many rows per grid step; 0 uses XLA reshapes. +moe_tc_ragged_3d_gmm: false # moe_tc_ragged_sort with tokamax gmm_v2: keep the sorted buffer in the (buffer, emb // 128, 128) layout through the GMMs. +moe_tc_ragged_3d_dispatch: false # moe_tc_ragged_sort + moe_tc_ragged_3d_gmm: keep tokens in the (batch, seq, emb // 128, 128) layout through the EP all-gather and combine. +moe_tc_window_routing: false # moe_tc_ragged_sort: keep only the shard's [start, start + cap) routing window, packed uint32 sorts. moe_use_direct_token_gather: false # whether to gather tokens directly in expert order instead of materializing Top-K copies use_gather_mosaic_kernel: false # whether to use a custom mosaic kernel for token gather ops ragged_gather_fallback: false # when true, unconditionally use the JAX reference implementation instead of the @@ -330,6 +402,13 @@ ragged_gather_cost_estimate_bytes_accessed: -1 # -1 means auto-compute, any > 0 ragged_gather_reduce_cost_estimate_bytes_accessed: -1 # -1 means auto-compute, any > 0 value overrides the bytes_accessed cost estimate for the ragged gather reduce kernel # Pin MoE FSDP and EP all-gathers to dedicated SparseCores using compute_on moe_pin_sparse_core_all_gathers: false +# Split controls (OR-ed with moe_pin_sparse_core_all_gathers): pin only the ring-of-experts EP all-gathers (dispatch +# tokens and routing inputs) or only the FSDP weight all-gathers. fwd_only keeps the pinned FSDP weight all-gather in the +# forward and runs its transpose (the weight-gradient reshard) outside compute_on, so the gradient's cross-slice +# all-reduce is not consumed by a SparseCore region. +moe_pin_sparse_core_ep_all_gathers: false +moe_pin_sparse_core_fsdp_all_gathers: false +moe_pin_sparse_core_fsdp_all_gathers_fwd_only: false moe_fsdp_all_gather_sparse_core_id: 0 moe_ep_all_gather_sparse_core_id: 1 # tunable tiling dimensions used for mlp gmm @@ -338,6 +417,9 @@ moe_ep_all_gather_sparse_core_id: 1 wi_tile_fwd_batch_seq: 512 wi_tile_fwd_embed_dim: 1024 wi_tile_fwd_mlp_dim: 1024 +eval_wi_tile_fwd_batch_seq: 512 +eval_wi_tile_fwd_embed_dim: 1024 +eval_wi_tile_fwd_mlp_dim: 1024 wi_tile_dlhs_batch_seq: 512 wi_tile_dlhs_embed_dim: 1024 wi_tile_dlhs_mlp_dim: 1024 @@ -348,6 +430,9 @@ wi_tile_drhs_mlp_dim: 1024 wo_tile_fwd_batch_seq: 512 wo_tile_fwd_embed_dim: 1024 wo_tile_fwd_mlp_dim: 1024 +eval_wo_tile_fwd_batch_seq: 512 +eval_wo_tile_fwd_embed_dim: 1024 +eval_wo_tile_fwd_mlp_dim: 1024 wo_tile_dlhs_batch_seq: 512 wo_tile_dlhs_embed_dim: 1024 wo_tile_dlhs_mlp_dim: 1024 @@ -399,6 +484,7 @@ routed_scaling_factor: 1.0 # scaling factor for routing scores routed_score_func: "" # scoring function for routing routed_bias: false # a flag if a learnable bias is added for routing routed_bias_update_rate: 0.0 # a flag indicate the update rate applied to the router bias term +defer_small_all_reduces: false # ring-of-experts only: reduce the per-layer routed-bias counts once after the layer loop rather than with one small all-reduce per layer mlp_bias: false # a flag if a learnable bias is added for MLP matmul, and originally implemented to support the GPT-OSS model architecture. n_routing_groups: -1 # number of groups for routing, disabled by default first_num_hash_layers: 0 # number of hash routing layers, used in DeepSeek V4 (0 means disabled) @@ -407,6 +493,9 @@ topk_routing_group: -1 # number of top groups to route inputs. For EP, # all-to-all communication with compute. Currently only implemented with DeepSeek sparse layers. use_batch_split_schedule: false # a flag if splitting batch into micro-batches to hide communications that yields performance benefits. batch_split_factor: 1 # the factor by which to split the batch. Only used if use_batch_split_schedule is true. +use_lineage: false # a flag to use Lineage DeepSeek-V3 execution. +lineage_quantization: "none" # "none" or "fp8_full": quantization of the Lineage routed experts. Only used if use_lineage is true. +lineage_expert_permutation: None # per-layer expert permutation for Lineage sparse/MTP layers: None (no shuffling) or "lpt" (on-device LPT load balancing from the previous step's expert counts). # For complex architectures like llama4 there are repeated sets of # inhomogeneous layers. E.g. maverick uses [dense+rope, moe+rope, dense+rope, moe+nope] @@ -480,6 +569,10 @@ moe_mlpwi_0: 'remat' moe_mlpwi_1: 'remat' moe_mlpwo: 'remat' moe_route: 'remat' # Ring-of-experts ragged sort integer results (sorted and revert indices, group sizes); saving them keeps both argsorts out of the bwd replay. Needs sparse_matmul, use_ring_of_experts and use_ragged_sort ('remat', 'device', or 'offload') +moe_x_sorted: 'remat' # routed (expert-sorted) MoE input plus its routing metadata; 'device' keeps it in HBM (tokens x top-k x emb per layer) so the backward skips the dispatch all-gather and ragged sort, 'offload' moves it to host, 'remat' recomputes it. Not supported with num_moe_emb_chunks > 0. +moe_routing_maps: 'remat' # MoE routing maps (top-k expert ids, ragged-sort permutations, sorted token indices and group sizes); with moe_tc_window_routing, also the TC group sizes +# TODO: Update to moe_tc_routing_weights. +moe_tc_routing_checkpoint: 'remat' # TC per-buffer-row routing weights (moe_tc_ragged_sort with moe_tc_ragged_weights_on_activation) query_proj: 'remat' key_proj: 'remat' value_proj: 'remat' @@ -498,6 +591,7 @@ parameter_memory_host_offload: false # Whether to use jax.lax.scan over layers (stacked/unstacked checkpoint). We recommend setting this to false when using pipeline parallelism, instead scanning the PP iterations. When resuming from a checkpoint, this flag is auto-determined from metadata. scan_layers: true param_scan_axis: 1 +gather_scanned_kernels_once: false # The attention parameter dictates the specific algorithm/methodology used to compute the attention scores # The attention_type parameter determines the variants of attention, e.g. global or local_sliding @@ -563,6 +657,7 @@ qk_nope_head_dim: 128 qk_rope_head_dim: 64 v_head_dim: 128 use_sliced_mla_proj: false # Whether to slice projection kernel weights before contraction in MLA instead of running full projection + jnp.split. +use_mla_absorbed_mqa: false # Whether to use Absorbed-Query Latent MQA mode in MLA (optimizes AR decode KV cache footprint/bandwidth at the cost of higher prefill/training FLOPs). # Compressed Attention parameters o_lora_rank: 0 # Output LoRA rank for Compressed Attention. @@ -709,6 +804,8 @@ logical_axis_rules: [ ['mlp_moe', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']], ['embed_moe', ['fsdp', 'fsdp_transpose', 'context', 'context_usp_ulysses']], ['embed_moe', ['fsdp', 'context', 'context_usp_ulysses']], + ['embed_router', ['fsdp', 'fsdp_transpose', 'context', 'context_usp_ulysses']], + ['embed_router', ['fsdp', 'context', 'context_usp_ulysses']], # ========================================== # Standard MLP / Dense Layers / Model Structure # ========================================== @@ -1034,6 +1131,7 @@ skip_first_n_steps_for_profiler: 1 # Profile for a small number of steps to avoid a large profile file size. profiler_steps: 5 hide_profiler_step_metric: false +enable_continuous_profiling: false # [Experimental] Enables continuous profiling (Allows saving >2GB traces) Requires JAX >= 0.11.2. profile_cleanly: true # If set to true, adds a block_until_ready on train state which aligns the profile for each step. profile_periodically_period: -1 # If set to a positive integer, profile every profile_periodically_period steps. # This is useful to debug scenarios where performance is changing. @@ -1076,6 +1174,14 @@ enable_data_shuffling: true data_shuffle_seed: 0 init_weights_seed: 0 +# Step-0 / MLPerf time-to-train options (all off by default). +block_state_before_run_start: false # before run_start: block_until_ready(state), warm the LR schedule, cross-host barrier. No data access. +warm_input_reshard_before_run_start: false # before run_start: device_put a synthetic zero batch to compile the input reshard program. No data access. +warm_eval_input_reshard_before_run_start: false # before run_start: device_put a synthetic zero EVAL batch to compile the eval input reshard program. No data access. +eval_cache_prefill_in_background: false # after run_start: background thread iterates the eval tf.data pipeline once so cache() is complete before the first eval. +train_shard_in_read: false # c4_mlperf train: shard by file in tfds instead of read-all-then-ds.shard (needs >= n_hosts files). +train_interleave_cycle_length: -1 # c4_mlperf train: number of files each host reads concurrently (tfds interleave_cycle_length). Applies with or without train_shard_in_read; -1 keeps the tfds default (16). With train_shard_in_read, set it below files-per-host so data_shuffle_seed changes which files are read first. + # DiLoCo params. enable_diloco: false diloco_sync_period: 36 @@ -1095,6 +1201,10 @@ dcn_bandwidth_interface: "eth0" enable_streaming_diloco: false num_diloco_fragments: null use_sequential_layers: false +# Split non-scanned matrices (embedding, output projection; any ndim >= 2 leaf, whatever its size) across the layer +# fragments, along their longest unsharded axis; matrices whose long axes are all sharded stay whole in fragment 0. +# Requires num_diloco_fragments >= 3. +diloco_bucketize_non_scanned: false num_communication_overlapping_steps: 0 communication_overlapping_alpha: 0.0 @@ -1178,6 +1288,10 @@ rope_attention_scaling: false # Scale the rotary embedding output # Keep rank-5 pair tensor intact [batch_size, sequence_length, num_heads, half_dim, 2] # where half_dim is the number of 2D rotation planes and 2 contains [real, imag] coordinates, returning interleaved RoPE rope_pairwise: false +# Compute RoPE angles as position * inv_freqs instead of gathering from a [max_position_embeddings, half_dim] table +rope_direct_position_freqs: false +# With rope_pairwise, stay in rank 4 and apply the pair swap as a block-skew matmul instead of a rank-5 reshape +rope_dot_pairwise: false # Ahead of time Compilation (aka AOT) # Only set these arguments if you are running train_compile or loading a compiled train step. @@ -1218,6 +1332,7 @@ abort_on_inf_loss: true # Check for Inf and abort if found in training loss # Goodput parameters enable_goodput_recording: false +goodput_job_name: "" monitor_goodput: false goodput_upload_interval_seconds: 30 enable_pathways_goodput: false @@ -1225,6 +1340,8 @@ monitor_step_time_deviation: true step_deviation_interval_seconds: 30 enable_gcp_goodput_metrics: true enable_gcp_step_deviation_metrics: true +rolling_windows_seconds: [3600, 86400, 259200, 432000] +enable_rolling_window_goodput: true # GCP workload monitoring report_heartbeat_metric_for_gcp_monitoring: false @@ -1234,12 +1351,14 @@ report_performance_metric_for_gcp_monitoring: false enable_tensorboard: true # Vertex AI Tensorboard Configurations - https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/use_vertex_ai_tensorboard.md -# Set to true for GCE, false if running via XPK +# Set to true to have MaxText create (or reuse) the Vertex AI Tensorboard instance and Experiment and upload logs. +# Leave false if they are created outside MaxText and UPLOAD_DATA_TO_TENSORBOARD is set in the environment +# (legacy XPK set this automatically; with Cluster Toolkit pass it via `gcluster job submit --env`). use_vertex_tensorboard: false -# Project to create Vertex AI Tensorboard in for GCE, blank if project is set using 'gcloud config set project' -# Set this to blank if running via XPK +# Project to create Vertex AI Tensorboard in, blank if project is set using 'gcloud config set project' +# Leave blank when use_vertex_tensorboard is false vertex_tensorboard_project: "" -# Region to create Vertex AI Tensorboard in for GCE, blank if running via XPK +# Region to create Vertex AI Tensorboard in, blank when use_vertex_tensorboard is false # Vertex AI supported regions: https://cloud.google.com/vertex-ai/docs/general/locations#available-regions vertex_tensorboard_region: "" @@ -1321,6 +1440,9 @@ ragged_block_size: 256 sa_block_q: 512 sa_block_kv: 512 sa_block_kv_compute: 512 +eval_sa_block_q: 512 +eval_sa_block_kv: 512 +eval_sa_block_kv_compute: 512 sa_block_q_dkv: 512 sa_block_kv_dkv: 512 sa_block_kv_dkv_compute: 512 @@ -1331,15 +1453,30 @@ sa_bwd_dkv_megacore: false # megacore-parallel kv-head groups in the static dkv sa_q_layout: "HEAD_DIM_MINOR" sa_k_layout: "HEAD_DIM_MINOR" sa_v_layout: "HEAD_DIM_MINOR" +eval_sa_q_layout: "HEAD_DIM_MINOR" +eval_sa_k_layout: "HEAD_DIM_MINOR" +eval_sa_v_layout: "HEAD_DIM_MINOR" use_splash_scheduler: false # to use tokamax splash attention scheduler. ring_scan_unroll: 1 # unroll factor for Tokamax ring attention scans; 0 or >= ring size fully unrolls. sa_fuse_reciprocal: true # defaults to true in Tokamax sa_use_base2_exp: true # defaults to true in Tokamax +# Diagonal sub-tile skipping (requires tokamax>=0.0.15 and use_tokamax_splash=true). Square blocks are required: +# sa_block_q == sa_block_kv == sa_block_kv_compute, sa_block_q_dkv == sa_block_kv_dkv == sa_block_kv_dkv_compute, +# and eval_sa_block_q == eval_sa_block_kv == eval_sa_block_kv_compute under custom eval rules. +# Only applies when the kernel gets a pure causal mask: global/mla attention layers with q_len == kv_len. Disabled +# for local_sliding/chunk/full/compressed/block_diffusion layers, decode, ring/USP context parallelism and +# load-balanced all_gather context parallelism (context_parallel_load_balance: true, the default). +sa_qk_diag_skip: false # skip strictly causal-masked sub-tiles in diagonal QK blocks +sa_sv_diag_skip: false # skip zeroed strictly causal-masked sub-tiles in diagonal SV blocks +sa_qk_diag_grid: 2 # sub-grid split factor for diagonal skip (power of 2 >= 2) # local_sa_* variants apply to local (sliding window) attention layers; # if None, each inherits from the corresponding global sa_* flag. local_sa_block_q: None # inherits from sa_block_q if None local_sa_block_kv: None # inherits from sa_block_kv if None local_sa_block_kv_compute: None # inherits from sa_block_kv_compute if None +eval_local_sa_block_q: None # inherits from eval_sa_block_q if None +eval_local_sa_block_kv: None # inherits from eval_sa_block_kv if None +eval_local_sa_block_kv_compute: None # inherits from eval_sa_block_kv_compute if None local_sa_block_q_dkv: None # inherits from sa_block_q_dkv if None local_sa_block_kv_dkv: None # inherits from sa_block_kv_dkv if None local_sa_block_kv_dkv_compute: None # inherits from sa_block_kv_dkv_compute if None @@ -1349,6 +1486,9 @@ local_sa_use_fused_bwd_kernel: None # inherits from sa_use_fused_bwd_kernel if local_sa_q_layout: None # inherits from sa_q_layout if None local_sa_k_layout: None # inherits from sa_k_layout if None local_sa_v_layout: None # inherits from sa_v_layout if None +eval_local_sa_q_layout: None # inherits from eval_sa_q_layout if None +eval_local_sa_k_layout: None # inherits from eval_sa_k_layout if None +eval_local_sa_v_layout: None # inherits from eval_sa_v_layout if None local_use_splash_scheduler: None # inherits from use_splash_scheduler if None local_sa_fuse_reciprocal: None # inherits from sa_fuse_reciprocal if None local_sa_use_base2_exp: None # inherits from sa_use_base2_exp if None @@ -1499,6 +1639,17 @@ gdn_cp_mode: "auto" # ('highest' = 6 bf16 passes on TPU, 'high' = 3 passes). Only used with use_gdn_kernel: true and # sequence-sharded GDN CP (gdn_cp_mode 'seq', or 'auto' when it does not pick head sharding); a no-op otherwise. gdn_cp_matmul_precision: "highest" +# GDN kernel path: build the Conv1D + SiLU halo-extended input [conv_state | qkv] in qkv's dtype and cast each tap +# to float32 where it is read, instead of materializing a float32 copy of the whole input. Same values: the outputs +# and the qkv, conv-state and bias gradients are bit-identical on CPU, the conv weight gradient matches up to float32 +# summation order. Requires use_gdn_kernel: true. +gdn_conv1d_input_in_qkv_dtype: false +# GDN kernel path: build the Pallas GDN forward kernel's Conv1D output from dense per-token [conv_dim / 1024, 8, 128] +# views instead of jnp.stack (kernels/gdn/compute_conv1d.py). Same values; the October 2026 TPU compiler lays the +# stacked value out with 7 of 8 sublanes as padding, so the SiLU after it runs on 8x the vector registers. Requires +# use_gdn_kernel: true and a GDN Conv1D width (2 * key heads * key head dim + value heads * value head dim) that is +# a multiple of 1024. +gdn_conv1d_dense_token_stack: false use_tokamax_splash: false # Setting this flag will use a non-pallas implementation. diff --git a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-fsdp.yml b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-fsdp.yml new file mode 100644 index 0000000000..c11f43cc36 --- /dev/null +++ b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-fsdp.yml @@ -0,0 +1,97 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Variant of cp-as-ep.yml whose only difference is that the dense weights are sharded on fsdp only. +# +# As cp-as-ep-dense-no-cp.yml, but the dense weights ('embed', 'embed_attn', 'q_lora', 'kv_lora') are also replicated +# over the expert axis: their all-gathers run in context x expert independent fsdp groups, and their gradients are +# reduce-scattered over all devices ('embed_grad') and then all-gathered over context x expert. Twice the HBM of +# cp-as-ep-dense-no-cp.yml for those weights. The GatedDeltaNet projections ('embed_gdn'), activations, the +# vocabulary weights and the routed experts are as in cp-as-ep, and RoutedMoE uses ("context", "expert") as its +# expert-parallel axes. +# +# This rule uses data, stage, FSDP, context, and expert. Context axis acts as expert +# parallelism in core MoE part (between EP all2all). +mesh_axes: ['data', 'stage', 'fsdp', 'context', 'expert'] +data_sharding: [['data', 'stage', 'fsdp', 'context', 'expert']] +logical_axis_rules: [ + # ========================================== + # Vocabulary Embedding + # ========================================== + # Vocab Activations + ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'expert']], + # Vocabulary tiling flattens [batch, length] into one dimension. It must also be sharded + # over the length axis (context), or every context rank computes the logits of all the + # context shards. context is last because it is the minor part of the flattened index. + ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'expert', 'context']], + # Vocab Weights + ['vocab', []], + ['embed_vocab', ['fsdp', 'context', 'expert']], + # ========================================== + # Attention + # ========================================== + # Attention Activations + ['activation_batch_attn', ['data', 'fsdp', 'expert']], + ['activation_heads', []], + ['activation_kv_heads', []], + ['activation_length_attn', ['context']], + ['activation_q_length', ['context']], + ['activation_kv_length', []], + ['activation_embed_attn', []], + ['activation_kv', []], + ['activation_kv_batch', ['data', 'fsdp', 'expert']], + ['activation_kv_head_dim', []], + # Attention Weights + ['heads', []], + ['q_heads', []], + ['kv_heads', []], + ['qkv', []], + ['kv', []], + ['kv_head_dim', []], + ['q_lora', ['fsdp']], + ["q_lora_up_proj", []], + ['kv_lora', ['fsdp']], + ["kv_lora_up_proj", []], + # ========================================== + # Mixture of Experts (MoE) + # ========================================== + # MoE Activations + ['activation_batch_moe', ['data', 'fsdp']], + ['activation_exp', ['context', 'expert']], + # MoE Weights + ['exp', ['context', 'expert']], + ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], + # ========================================== + # Standard MLP / Dense Layers / Model Structure + # ========================================== + # Dense Activations + ['activation_mlp', []], + ['activation_batch', ['data', 'fsdp', 'expert']], + ['activation_length', ['context']], + ['activation_norm_length', ['context']], + ['activation_embed', []], + ['activation_stage', 'stage'], + # General Weights + ['mlp', []], + ['layers', 'stage'], + # Dense weights: not sharded on context, see the header. + ['embed', ['fsdp']], + ['embed_attn', ['fsdp']], + # GatedDeltaNet projections (models/qwen3.py) stay sharded as in cp-as-ep. + ['embed_gdn', ['fsdp', 'context', 'expert']], + # The dense weights' gradients are reduce-scattered over every token axis, as in cp-as-ep, then + # gathered over the axes the weights are replicated on (linears.route_dense_weight_grad). + ['embed_grad', ['fsdp', 'context', 'expert']], + ] diff --git a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-no-cp.yml b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-no-cp.yml new file mode 100644 index 0000000000..6387338ae3 --- /dev/null +++ b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-no-cp.yml @@ -0,0 +1,100 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Variant of cp-as-ep.yml whose only difference is that the dense weights are not sharded on the context axis. +# +# cp-as-ep shards the dense weights ('embed', 'embed_attn', 'q_lora', 'kv_lora': attention and shared-expert +# projections, the MoE router and the shared-expert gate) over fsdp x context x expert, so every all-gather of them +# spans all devices and also synchronizes the context ranks, which do unequal attention work. Here they are sharded +# over fsdp x expert and replicated over context: those all-gathers run in context-many independent fsdp x expert +# groups, at the cost of a context-way replication of the weights and their gradients. Their gradients are still +# reduce-scattered over all devices, as in cp-as-ep ('embed_grad'), and the shard is then all-gathered over context; +# left to GSPMD, the gradient would be all-reduced at full size over all devices instead. The GatedDeltaNet +# projections ('embed_gdn'), activations, the vocabulary weights and the routed experts ('exp' over context x expert) +# are sharded as in cp-as-ep, and RoutedMoE uses ("context", "expert") as its expert-parallel axes. +# +# This rule uses data, stage, FSDP, context, and expert. Context axis acts as expert +# parallelism in core MoE part (between EP all2all). +mesh_axes: ['data', 'stage', 'fsdp', 'context', 'expert'] +data_sharding: [['data', 'stage', 'fsdp', 'context', 'expert']] +logical_axis_rules: [ + # ========================================== + # Vocabulary Embedding + # ========================================== + # Vocab Activations + ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'expert']], + # Vocabulary tiling flattens [batch, length] into one dimension. It must also be sharded + # over the length axis (context), or every context rank computes the logits of all the + # context shards. context is last because it is the minor part of the flattened index. + ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'expert', 'context']], + # Vocab Weights + ['vocab', []], + ['embed_vocab', ['fsdp', 'context', 'expert']], + # ========================================== + # Attention + # ========================================== + # Attention Activations + ['activation_batch_attn', ['data', 'fsdp', 'expert']], + ['activation_heads', []], + ['activation_kv_heads', []], + ['activation_length_attn', ['context']], + ['activation_q_length', ['context']], + ['activation_kv_length', []], + ['activation_embed_attn', []], + ['activation_kv', []], + ['activation_kv_batch', ['data', 'fsdp', 'expert']], + ['activation_kv_head_dim', []], + # Attention Weights + ['heads', []], + ['q_heads', []], + ['kv_heads', []], + ['qkv', []], + ['kv', []], + ['kv_head_dim', []], + ['q_lora', ['fsdp', 'expert']], + ["q_lora_up_proj", []], + ['kv_lora', ['fsdp', 'expert']], + ["kv_lora_up_proj", []], + # ========================================== + # Mixture of Experts (MoE) + # ========================================== + # MoE Activations + ['activation_batch_moe', ['data', 'fsdp']], + ['activation_exp', ['context', 'expert']], + # MoE Weights + ['exp', ['context', 'expert']], + ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], + # ========================================== + # Standard MLP / Dense Layers / Model Structure + # ========================================== + # Dense Activations + ['activation_mlp', []], + ['activation_batch', ['data', 'fsdp', 'expert']], + ['activation_length', ['context']], + ['activation_norm_length', ['context']], + ['activation_embed', []], + ['activation_stage', 'stage'], + # General Weights + ['mlp', []], + ['layers', 'stage'], + # Dense weights: not sharded on context, see the header. + ['embed', ['fsdp', 'expert']], + ['embed_attn', ['fsdp', 'expert']], + # GatedDeltaNet projections (models/qwen3.py) stay sharded as in cp-as-ep. + ['embed_gdn', ['fsdp', 'context', 'expert']], + # The dense weights' gradients are reduce-scattered over every token axis, as in cp-as-ep, then + # gathered over the axes the weights are replicated on (linears.route_dense_weight_grad). + ['embed_grad', ['fsdp', 'context', 'expert']], + ] diff --git a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml index bd3828825d..fdcff51c1a 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml +++ b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml @@ -63,6 +63,7 @@ logical_axis_rules: [ # MoE Weights ['exp', ['context', 'expert']], ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], # ========================================== # Standard MLP / Dense Layers / Model Structure # ========================================== diff --git a/src/maxtext/configs/custom_mesh_and_rule/ep-as-cp.yml b/src/maxtext/configs/custom_mesh_and_rule/ep-as-cp.yml index 142c441867..0dfdf75258 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/ep-as-cp.yml +++ b/src/maxtext/configs/custom_mesh_and_rule/ep-as-cp.yml @@ -60,6 +60,7 @@ logical_axis_rules: [ # MoE Weights ['exp', 'expert'], ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], # ========================================== # Standard MLP / Dense Layers / Model Structure # ========================================== diff --git a/src/maxtext/configs/custom_mesh_and_rule/ep-as-dp.yml b/src/maxtext/configs/custom_mesh_and_rule/ep-as-dp.yml index b6091c2d9b..f9e885dee2 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/ep-as-dp.yml +++ b/src/maxtext/configs/custom_mesh_and_rule/ep-as-dp.yml @@ -61,6 +61,7 @@ logical_axis_rules: [ ['exp', 'expert'], ['mlp_moe', ['fsdp_transpose']], ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], # ========================================== # Standard MLP / Dense Layers / Model Structure # ========================================== diff --git a/src/maxtext/configs/custom_mesh_and_rule/fsdp-as-dp-for-attn-rf.yml b/src/maxtext/configs/custom_mesh_and_rule/fsdp-as-dp-for-attn-rf.yml new file mode 100644 index 0000000000..27b6829dd7 --- /dev/null +++ b/src/maxtext/configs/custom_mesh_and_rule/fsdp-as-dp-for-attn-rf.yml @@ -0,0 +1,115 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Variant of fsdp-as-dp-for-attn whose only difference is that the MoE router weight ('embed_router') is +# sharded on fsdp instead of fully replicated. See the 'embed_router' rule below. + +# This rule uses data, FSDP, FSDP_transpose and expert, and assigns them as: +# +# FSDP acts like DP for attn, but FSDP for moe +# EP acts like FSDP for attn, and EP for moe +# +# That is, attention weights ('q_lora', 'kv_lora' and 'embed_attn') are sharded +# on the expert axis and replicated over fsdp / fsdp_transpose, the MoE router +# weight ('embed_router') is sharded on fsdp, while routed/shared MoE and +# dense MLP weights keep using FSDP as in ep-as-dp.yml. Outside the core MoE +# part (between EP all2all) the expert axis carries the batch, so sharding the +# attention weights on it behaves like FSDP over a much smaller and faster axis: +# it trades HBM for a cheaper all-gather. +# +# This is mainly useful for large scale training. The attention weights are +# small, so once the FSDP rank is large (larger than 512) each shard is tiny and +# their all-gather is dominated by latency rather than bandwidth, which makes +# the attention layer latency bound. Sharding the attention weights on the +# (much smaller) expert axis keeps the per-shard slice large enough that the +# collective stays bandwidth bound, at the cost of replicating those weights +# over fsdp / fsdp_transpose. MoE weights are large enough to keep scaling with +# FSDP, so they are left alone. DeepSeek V3 MLPerf training at 4k sequence +# length is one such case. + +mesh_axes: ['data', 'fsdp', 'fsdp_transpose', 'expert'] +data_sharding: [['data', 'fsdp', 'fsdp_transpose', 'expert']] +logical_axis_rules: [ + # ========================================== + # Vocabulary Embedding + # ========================================== + # Vocab Activations + ['activation_embed_and_logits_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']], + ['activation_embed_and_logits_batch_sequence', ['data', 'fsdp', 'fsdp_transpose', 'expert']], + ['activation_vocab', []], + # Vocab Weights + ['vocab', []], + ['embed_vocab', ['fsdp', 'fsdp_transpose']], + # ========================================== + # Attention + # ========================================== + # Attention Activations + ['activation_batch_attn', ['data', 'fsdp', 'fsdp_transpose', 'expert']], + ['activation_heads', []], + ['activation_kv_heads', []], + ['activation_length_attn', ['context']], + ['activation_q_length', ['context']], + ['activation_kv_length', []], + ['activation_embed_attn', ['tensor', 'tensor_transpose']], + ['activation_kv', []], + ['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']], + ['activation_kv_head_dim', []], + # Attention Weights, sharded on expert instead of fsdp: + # fsdp acts like DP here, expert acts like FSDP + ['q_lora', ['expert']], + ["q_lora_up_proj", []], + ['kv_lora', ['expert']], + ["kv_lora_up_proj", []], + ['embed_attn', ['expert']], + # ========================================== + # Mixture of Experts (MoE) + # ========================================== + # MoE Activations + ['activation_batch_moe', ['data', 'fsdp', 'fsdp_transpose', 'expert']], + ['activation_length_moe', []], + ['activation_norm_length_moe', []], + ['activation_embed_moe', []], + ['activation_mlp_moe', []], + ['activation_exp', ['expert']], + # MoE Weights + ['exp', 'expert'], + ['mlp_moe', ['fsdp_transpose']], + ['embed_moe', ['fsdp']], + # 'rf' variant: shard the router weight on fsdp (as in ep-as-dp) instead of replicating it. + # A replicated router makes GSPMD reduce the bf16[embed, num_experts] router gradient with a + # synchronous TensorCore all-reduce over every device inside the backward scan loop. The fsdp + # sharding does not remove the reduction (the gradient is still a partial sum over all batch + # axes, including expert): it becomes a reduce-scatter over fsdp plus an all-reduce of only the + # [embed / fsdp, num_experts] shard over data / fsdp_transpose / expert, at the cost of a small + # forward/remat all-gather. + ['embed_router', ['fsdp']], + # ========================================== + # Standard MLP / Dense Layers / Model Structure + # ========================================== + # Dense Activations + ['activation_mlp', []], + # Note activation batch and length also get used in vocab + ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']], + ['activation_length', []], + ['activation_norm_length', []], + ['activation_embed', []], + ['activation_stage', []], + # General Weights + ['mlp', ['fsdp_transpose']], + ['embed', ['fsdp', 'fsdp_transpose']], + ['embed', ['fsdp']], + # ========================================== + # Deprecated / Scheduled for Removal + # ========================================== + ['exp_with_fsdp', 'fsdp'], + ] diff --git a/src/maxtext/configs/custom_mesh_and_rule/fsdp-as-dp-for-attn.yml b/src/maxtext/configs/custom_mesh_and_rule/fsdp-as-dp-for-attn.yml index 193480bd10..6e663f763b 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/fsdp-as-dp-for-attn.yml +++ b/src/maxtext/configs/custom_mesh_and_rule/fsdp-as-dp-for-attn.yml @@ -18,7 +18,8 @@ # EP acts like FSDP for attn, and EP for moe # # That is, attention weights ('q_lora', 'kv_lora' and 'embed_attn') are sharded -# on the expert axis and replicated over fsdp / fsdp_transpose, while MoE and +# on the expert axis and replicated over fsdp / fsdp_transpose, the MoE router +# weight ('embed_router') is fully replicated, while routed/shared MoE and # dense MLP weights keep using FSDP as in ep-as-dp.yml. Outside the core MoE # part (between EP all2all) the expert axis carries the batch, so sharding the # attention weights on it behaves like FSDP over a much smaller and faster axis: @@ -82,6 +83,8 @@ logical_axis_rules: [ ['exp', 'expert'], ['mlp_moe', ['fsdp_transpose']], ['embed_moe', ['fsdp']], + # Replicate small router weight to eliminate forward/remat all-gather + ['embed_router', []], # ========================================== # Standard MLP / Dense Layers / Model Structure # ========================================== diff --git a/src/maxtext/configs/custom_mesh_and_rule/pipeline-large-moe.yml b/src/maxtext/configs/custom_mesh_and_rule/pipeline-large-moe.yml index fe076b526e..7c40592397 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/pipeline-large-moe.yml +++ b/src/maxtext/configs/custom_mesh_and_rule/pipeline-large-moe.yml @@ -63,6 +63,7 @@ logical_axis_rules: [ ['embed_attn', ['fsdp', 'expert']], # remove context from embed sharding ['embed_vocab', ['fsdp', 'expert']], ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], ['q_lora', ['fsdp']], ['kv_lora', ['fsdp']], ['norm', ['tensor']], diff --git a/src/maxtext/configs/custom_mesh_and_rule/pure-fsdp.yml b/src/maxtext/configs/custom_mesh_and_rule/pure-fsdp.yml index 02249602fc..53484dd37d 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/pure-fsdp.yml +++ b/src/maxtext/configs/custom_mesh_and_rule/pure-fsdp.yml @@ -29,6 +29,7 @@ logical_axis_rules: [ ['embed_attn', ['fsdp']], ['embed_vocab', ['fsdp']], ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], ['q_lora', ['fsdp']], ['kv_lora', ['fsdp']], ['exp_with_fsdp', 'fsdp'], diff --git a/src/maxtext/configs/custom_mesh_and_rule/shard-exp-on-fsdp.yml b/src/maxtext/configs/custom_mesh_and_rule/shard-exp-on-fsdp.yml index 41e9487852..807686953f 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/shard-exp-on-fsdp.yml +++ b/src/maxtext/configs/custom_mesh_and_rule/shard-exp-on-fsdp.yml @@ -55,6 +55,7 @@ logical_axis_rules: [ ['exp', ['expert', 'fsdp']], # fsdp shards exp instead embed ['mlp_moe', ['fsdp_transpose']], ['embed_moe', ['context']], + ['embed_router', ['context']], # ========================================== # Standard MLP / Dense Layers / Model Structure # ========================================== diff --git a/src/maxtext/configs/decoupled_base_test.yml b/src/maxtext/configs/decoupled_base_test.yml index 485bb4e8ae..0288a05788 100644 --- a/src/maxtext/configs/decoupled_base_test.yml +++ b/src/maxtext/configs/decoupled_base_test.yml @@ -4,6 +4,8 @@ base_config: base.yml # Output goes to a local relative directory so tests do not require GCS. +# Tests that call get_test_base_output_directory() write to a per-process +# subdirectory of this root instead; only tests that rely on this default share it. base_output_directory: ./maxtext_local_output/gcloud_decoupled_test_logs run_name: test_decoupled @@ -70,5 +72,3 @@ monitor_goodput: false goodput_upload_interval_seconds: 0 enable_pathways_goodput: false enable_gcp_goodput_metrics: false - -tensorboard_dir: "./maxtext_local_output/gcloud_decoupled_test_logs/tensorboard" diff --git a/src/maxtext/configs/gpu/a3/llama_2_7b/16vm.sh b/src/maxtext/configs/gpu/a3/llama_2_7b/16vm.sh index 1d193ae318..54b163b31f 100644 --- a/src/maxtext/configs/gpu/a3/llama_2_7b/16vm.sh +++ b/src/maxtext/configs/gpu/a3/llama_2_7b/16vm.sh @@ -1,8 +1,11 @@ echo "Running 16vm.sh" -# Example command to invoke this script via XPK -# python3 xpk/xpk.py workload create --cluster ${CLUSTER_NAME} \ -# --workload ${WORKLOAD_NAME} --docker-image=gcr.io/supercomputer-testing/${LOCAL_IMAGE_NAME} \ -# --device-type ${DEVICE_TYPE} --num-slices 16 --priority=high \ +# Example command to submit this script with Cluster Toolkit (gcluster). See +# docs/run_maxtext/run_maxtext_via_cluster_toolkit.md to install gcluster and set the default +# project, cluster, and location (`gcluster job config set ...`). COMPUTE_TYPE is the GPU +# type, e.g. h100-80gb-8 (a3-highgpu-8g). +# gcluster job submit --name=${WORKLOAD_NAME} \ +# --image=gcr.io/${PROJECT_ID}/${LOCAL_IMAGE_NAME} \ +# --compute-type=${COMPUTE_TYPE} --num-nodes=16 --priority=high \ # --command "bash src/maxtext/configs/gpu/a3/llama_2_7b/16vm.sh" # Stop execution if any command exits with error diff --git a/src/maxtext/configs/gpu/a3/llama_2_7b/1vm.sh b/src/maxtext/configs/gpu/a3/llama_2_7b/1vm.sh index 0958449bf8..459f5ab40b 100644 --- a/src/maxtext/configs/gpu/a3/llama_2_7b/1vm.sh +++ b/src/maxtext/configs/gpu/a3/llama_2_7b/1vm.sh @@ -1,9 +1,12 @@ echo "Running 1vm.sh" -# Example command to invoke this script via XPK -# python3 xpk/xpk.py workload create --cluster ${CLUSTER_NAME} \ -# --workload ${WORKLOAD_NAME} --docker-image=gcr.io/supercomputer-testing/${LOCAL_IMAGE_NAME} \ -# --device-type ${DEVICE_TYPE} --num-slices 1 \ +# Example command to submit this script with Cluster Toolkit (gcluster). See +# docs/run_maxtext/run_maxtext_via_cluster_toolkit.md to install gcluster and set the default +# project, cluster, and location (`gcluster job config set ...`). COMPUTE_TYPE is the GPU +# type, e.g. h100-80gb-8 (a3-highgpu-8g). +# gcluster job submit --name=${WORKLOAD_NAME} \ +# --image=gcr.io/${PROJECT_ID}/${LOCAL_IMAGE_NAME} \ +# --compute-type=${COMPUTE_TYPE} --num-nodes=1 \ # --command "bash src/maxtext/configs/gpu/a3/llama_2_7b/1vm.sh" # Stop execution if any command exits with error diff --git a/src/maxtext/configs/gpu/a3/llama_2_7b/2vm.sh b/src/maxtext/configs/gpu/a3/llama_2_7b/2vm.sh index 34b7f345a0..6e158393f0 100644 --- a/src/maxtext/configs/gpu/a3/llama_2_7b/2vm.sh +++ b/src/maxtext/configs/gpu/a3/llama_2_7b/2vm.sh @@ -1,9 +1,12 @@ echo "Running 2vm.sh" -# Example command to invoke this script via XPK -# python3 xpk/xpk.py workload create --cluster ${CLUSTER_NAME} \ -# --workload ${WORKLOAD_NAME} --docker-image=gcr.io/supercomputer-testing/${LOCAL_IMAGE_NAME} \ -# --device-type ${DEVICE_TYPE} --num-slices 2 \ +# Example command to submit this script with Cluster Toolkit (gcluster). See +# docs/run_maxtext/run_maxtext_via_cluster_toolkit.md to install gcluster and set the default +# project, cluster, and location (`gcluster job config set ...`). COMPUTE_TYPE is the GPU +# type, e.g. h100-80gb-8 (a3-highgpu-8g). +# gcluster job submit --name=${WORKLOAD_NAME} \ +# --image=gcr.io/${PROJECT_ID}/${LOCAL_IMAGE_NAME} \ +# --compute-type=${COMPUTE_TYPE} --num-nodes=2 \ # --command "bash src/maxtext/configs/gpu/a3/llama_2_7b/2vm.sh" # Stop execution if any command exits with error diff --git a/src/maxtext/configs/gpu/a3/llama_2_7b/4vm.sh b/src/maxtext/configs/gpu/a3/llama_2_7b/4vm.sh index 17bbba0b7c..164ca9cd1c 100644 --- a/src/maxtext/configs/gpu/a3/llama_2_7b/4vm.sh +++ b/src/maxtext/configs/gpu/a3/llama_2_7b/4vm.sh @@ -1,8 +1,11 @@ echo "Running 4vm.sh" -# Example command to invoke this script via XPK -# python3 xpk/xpk.py workload create --cluster ${CLUSTER_NAME} \ -# --workload ${WORKLOAD_NAME} --docker-image=gcr.io/supercomputer-testing/${LOCAL_IMAGE_NAME} \ -# --device-type ${DEVICE_TYPE} --num-slices 4 \ +# Example command to submit this script with Cluster Toolkit (gcluster). See +# docs/run_maxtext/run_maxtext_via_cluster_toolkit.md to install gcluster and set the default +# project, cluster, and location (`gcluster job config set ...`). COMPUTE_TYPE is the GPU +# type, e.g. h100-80gb-8 (a3-highgpu-8g). +# gcluster job submit --name=${WORKLOAD_NAME} \ +# --image=gcr.io/${PROJECT_ID}/${LOCAL_IMAGE_NAME} \ +# --compute-type=${COMPUTE_TYPE} --num-nodes=4 \ # --command "bash src/maxtext/configs/gpu/a3/llama_2_7b/4vm.sh" # Stop execution if any command exits with error diff --git a/src/maxtext/configs/gpu/a3/llama_2_7b/8vm.sh b/src/maxtext/configs/gpu/a3/llama_2_7b/8vm.sh index 187fd84177..192f9bc2bf 100644 --- a/src/maxtext/configs/gpu/a3/llama_2_7b/8vm.sh +++ b/src/maxtext/configs/gpu/a3/llama_2_7b/8vm.sh @@ -1,8 +1,11 @@ echo "Running 8vm.sh" -# Example command to invoke this script via XPK -# python3 xpk/xpk.py workload create --cluster ${CLUSTER_NAME} \ -# --workload ${WORKLOAD_NAME} --docker-image=gcr.io/supercomputer-testing/${LOCAL_IMAGE_NAME} \ -# --device-type ${DEVICE_TYPE} --num-slices 8 \ +# Example command to submit this script with Cluster Toolkit (gcluster). See +# docs/run_maxtext/run_maxtext_via_cluster_toolkit.md to install gcluster and set the default +# project, cluster, and location (`gcluster job config set ...`). COMPUTE_TYPE is the GPU +# type, e.g. h100-80gb-8 (a3-highgpu-8g). +# gcluster job submit --name=${WORKLOAD_NAME} \ +# --image=gcr.io/${PROJECT_ID}/${LOCAL_IMAGE_NAME} \ +# --compute-type=${COMPUTE_TYPE} --num-nodes=8 \ # --command "bash src/maxtext/configs/gpu/a3/llama_2_7b/8vm.sh" # Stop execution if any command exits with error diff --git a/src/maxtext/configs/gpu/a3/llama_3.1_405b/128vm.sh b/src/maxtext/configs/gpu/a3/llama_3.1_405b/128vm.sh index 89af9b655b..82d1a3f1c4 100644 --- a/src/maxtext/configs/gpu/a3/llama_3.1_405b/128vm.sh +++ b/src/maxtext/configs/gpu/a3/llama_3.1_405b/128vm.sh @@ -1,12 +1,15 @@ echo "Running 128vm.sh" -# Example command to invoke this script via XPK, assume you've installed xpk +# Example command to submit this script with Cluster Toolkit (gcluster). See +# docs/run_maxtext/run_maxtext_via_cluster_toolkit.md to install gcluster. LOCATION is the +# cluster's region (regional clusters) or zone (zonal clusters). COMPUTE_TYPE is the GPU type, +# e.g. h100-80gb-8 (a3-highgpu-8g). The script runs on 128 nodes (dcn_fsdp_parallelism=128). # COMMAND="bash src/maxtext/configs/gpu/a3/llama_3.1_405b/128vm.sh" # COMMAND='export LD_LIBRARY_PATH=/usr/local/cuda-12.6/compat:$LD_LIBRARY_PATH;'"${COMMAND}"; # -# xpk workload create --project=${PROJECT}--cluster=${CLUSTER_NAME} --zone=${ZONE} \ -# --workload=${WORKLOAD_NAME} --docker-image=gcr.io/supercomputer-testing/${LOCAL_IMAGE_NAME} \ -# --device-type=${DEVICE_TYPE} --num-nodes=2 --priority=high \ -# --command="$COMMAND" --env=XLA_FLAGS=$XLA_FLAGS +# gcluster job submit --project=${PROJECT_ID} --cluster=${CLUSTER_NAME} --location=${LOCATION} \ +# --name=${WORKLOAD_NAME} --image=gcr.io/${PROJECT_ID}/${LOCAL_IMAGE_NAME} \ +# --compute-type=${COMPUTE_TYPE} --num-nodes=128 --priority=high \ +# --command="$COMMAND" --env="XLA_FLAGS=$XLA_FLAGS" # Stop execution if any command exits with error set -e diff --git a/src/maxtext/configs/inference/inference.yml b/src/maxtext/configs/inference/inference.yml index 0099934825..b76e3e9d2c 100644 --- a/src/maxtext/configs/inference/inference.yml +++ b/src/maxtext/configs/inference/inference.yml @@ -3,8 +3,8 @@ base_config: "base.yml" logical_axis_rules: [ ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']], ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert']], - ['activation_heads', ['tensor', 'tensor_transpose', 'tensor_sequence']], - ['activation_kv_heads', ['tensor', 'tensor_transpose', 'tensor_sequence']], + ['activation_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']], + ['activation_kv_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']], ['activation_length', ['context_autoregressive']], ['activation_length', ['context_autoregressive']], ['activation_q_length', ['context_autoregressive']], @@ -48,6 +48,10 @@ logical_axis_rules: [ ['embed_moe', ['fsdp', 'context_autoregressive', 'tensor_transpose']], ['embed_moe', ['fsdp', 'fsdp_transpose', 'context_autoregressive']], ['embed_moe', ['fsdp', 'context_autoregressive']], + ['embed_router', ['fsdp', 'fsdp_transpose', 'context_autoregressive', 'tensor_transpose']], + ['embed_router', ['fsdp', 'context_autoregressive', 'tensor_transpose']], + ['embed_router', ['fsdp', 'fsdp_transpose', 'context_autoregressive']], + ['embed_router', ['fsdp', 'context_autoregressive']], ['norm', ['tensor', 'tensor_transpose', 'tensor_sequence']], ['layers', 'stage'], ['kv', []], diff --git a/src/maxtext/configs/inference/vllm.yml b/src/maxtext/configs/inference/vllm.yml index 04612cef50..89f2b53735 100644 --- a/src/maxtext/configs/inference/vllm.yml +++ b/src/maxtext/configs/inference/vllm.yml @@ -78,6 +78,7 @@ logical_axis_rules: [ ['exp', ['attn_dp_expert', 'expert', 'dcp']], ['mlp_moe', ['attn_dp', 'model']], ['embed_moe', []], + ['embed_router', []], # ========================================== # Standard MLP / Dense Layers / Model Structure # ========================================== diff --git a/src/maxtext/configs/models/deepseek3-671b-2dfsdp.yml b/src/maxtext/configs/models/deepseek3-671b-2dfsdp.yml index 6190c7eba6..2d9ba8c49f 100644 --- a/src/maxtext/configs/models/deepseek3-671b-2dfsdp.yml +++ b/src/maxtext/configs/models/deepseek3-671b-2dfsdp.yml @@ -72,6 +72,7 @@ logical_axis_rules: [ ['embed', ['fsdp']], ['embed_attn', ['fsdp']], ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], ['embed_vocab', ['fsdp', 'fsdp_transpose']], ['q_lora', ['fsdp']], ['kv_lora', ['fsdp']], diff --git a/src/maxtext/configs/models/deepseek3-671b-batchsplit.yml b/src/maxtext/configs/models/deepseek3-671b-batchsplit.yml index acfcb54f44..8a98e249da 100644 --- a/src/maxtext/configs/models/deepseek3-671b-batchsplit.yml +++ b/src/maxtext/configs/models/deepseek3-671b-batchsplit.yml @@ -75,6 +75,7 @@ logical_axis_rules: [ ['embed', ['fsdp']], ['embed_attn', ['fsdp']], ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], ['embed_vocab', ['fsdp']], ['embed_no_exp', ['fsdp']], ['embed_no_exp_moe', ['fsdp']], diff --git a/src/maxtext/configs/models/deepseek4-284b.yml b/src/maxtext/configs/models/deepseek4-284b.yml index 064fd23dc1..5623878b6c 100644 --- a/src/maxtext/configs/models/deepseek4-284b.yml +++ b/src/maxtext/configs/models/deepseek4-284b.yml @@ -59,6 +59,7 @@ routed_scaling_factor: 1.5 # --- Attention configuration --- +attention: "dot_product" # use_indexer requires dot_product, or flash with use_tokamax_splash attention_type: 'compressed' use_indexer: true q_lora_rank: 1024 @@ -67,7 +68,12 @@ o_lora_rank: 1024 sliding_window_size: 128 # --- RoPE --- -rope_type: "default" +# Values below mirror the official HF config.json `rope_scaling`: +# {"type": "yarn", "factor": 16, "beta_fast": 32, "beta_slow": 1, "original_max_position_embeddings": 65536}. +rope_type: "yarn" +rope_factor: 16 +beta_fast: 32 +beta_slow: 1 rope_max_timescale: 10000 # Main RoPE theta compressed_rope_max_timescale: 160000 # Compressed RoPE theta max_position_embeddings: 1048576 diff --git a/src/maxtext/configs/models/qwen3.5-9b.yml b/src/maxtext/configs/models/qwen3.5-9b.yml new file mode 100644 index 0000000000..0a2b0c43df --- /dev/null +++ b/src/maxtext/configs/models/qwen3.5-9b.yml @@ -0,0 +1,49 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# maxtext/configs/models/qwen3.5-9b.yml + +decoder_block: "qwen3_5" + +# Core Architectural Parameters +base_emb_dim: 4096 +base_num_decoder_layers: 32 +base_num_query_heads: 16 +base_num_kv_heads: 4 +head_dim: 256 +vocab_size: 248320 +normalization_layer_epsilon: 1.0e-6 + +# Dense MLP block (not MoE) +base_mlp_dim: 12288 +mlp_activations: ["silu", "linear"] + +# GatedDeltaNet Specific Parameters for Linear Attention (GDN) +inhomogeneous_layer_cycle_interval: 4 +gdn_conv_kernel_dim: 4 +gdn_key_head_dim: 128 +gdn_value_head_dim: 128 +gdn_num_key_heads: 16 +gdn_num_value_heads: 32 +gdn_chunk_size: 64 + +# RoPE Settings +rope_max_timescale: 10000000 +partial_rotary_factor: 0.25 + +# General Model Settings +enable_dropout: False +logits_via_embedding: False +normalize_embedding_logits: False +tokenizer_type: "huggingface" diff --git a/src/maxtext/configs/models/weaver-max.yml b/src/maxtext/configs/models/weaver-max.yml new file mode 100644 index 0000000000..f13b7047dd --- /dev/null +++ b/src/maxtext/configs/models/weaver-max.yml @@ -0,0 +1,41 @@ +# Model config for Weaver Max + +# Core Architectural Parameters +decoder_block: "qwen3" +base_emb_dim: 5120 +base_mlp_dim: 25600 +base_num_query_heads: 64 +base_num_kv_heads: 8 +base_num_decoder_layers: 64 +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 Settings +use_multimodal: true +use_mrope: true +mrope_section: [24, 20, 20] + +# Vision Encoder Configuration +image_size_for_vit: 768 +hidden_size_for_vit: 1152 +intermediate_size_for_vit: 4304 +num_attention_heads_for_vit: 16 +num_hidden_layers_for_vit: 27 +num_channels_for_vit: 3 +patch_size_for_vit: 16 +temporal_patch_size_for_vit: 2 +spatial_merge_size_for_vit: 2 +out_hidden_size_for_vit: 5120 +num_position_embeddings_for_vit: 2304 +deepstack_visual_indexes_for_vit: [8, 16, 24] +vision_encoder_block: "qwen3_vl" 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-mini.yml b/src/maxtext/configs/models/weaver-mini.yml new file mode 100644 index 0000000000..c36b572864 --- /dev/null +++ b/src/maxtext/configs/models/weaver-mini.yml @@ -0,0 +1,41 @@ +# Model config for Weaver Mini (Text-only mode) + +# Core Architectural Parameters +decoder_block: "qwen3" +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 Settings +use_multimodal: true +use_mrope: true +mrope_section: [24, 20, 20] + +# Vision Encoder Configuration +image_size_for_vit: 768 +hidden_size_for_vit: 1152 +intermediate_size_for_vit: 4304 +num_attention_heads_for_vit: 16 +num_hidden_layers_for_vit: 27 +num_channels_for_vit: 3 +patch_size_for_vit: 16 +temporal_patch_size_for_vit: 2 +spatial_merge_size_for_vit: 2 +out_hidden_size_for_vit: 4096 +num_position_embeddings_for_vit: 2304 +deepstack_visual_indexes_for_vit: [8, 16, 24] +vision_encoder_block: "qwen3_vl" 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/post_train/rl.yml b/src/maxtext/configs/post_train/rl.yml index f01b85b80e..1841cd5752 100644 --- a/src/maxtext/configs/post_train/rl.yml +++ b/src/maxtext/configs/post_train/rl.yml @@ -166,6 +166,15 @@ max_target_length: 1024 kv_cache_buffer: 256 hbm_utilization_vllm: 0.72 swap_space_vllm_gb: 2 +# Maps to tunix RolloutConfig.rollout_vllm_free_kv_cache_during_weight_sync -> VllmConfig.free_kv_cache_during_weight_sync +# (effective once the tunix pin declares it; until then the setting is ignored and the trainer logs that). +# The weight sync materializes a second copy of the sampler weights on HBM while the new values are resharded in; +# freeing the KV cache first makes room for that copy. false keeps the pool (hbm_utilization_vllm of each chip) +# allocated across the sync and only resets its prefix-cache entries, which saves the free + re-allocation +# (about 2 s per step on a colocated Qwen3-0.6B run at hbm_utilization_vllm=0.3). With false the sync runs out of +# HBM when that second copy does not fit next to the pool, so check the headroom before disabling: the default +# hbm_utilization_vllm above leaves far less room than 0.3, and colocated setups have the least. +free_kv_cache_during_weight_sync: true # Generation Configuration During Training # Important to keep a high-ish temperature for varied, diverse responses during # training. @@ -318,6 +327,8 @@ logical_axis_rules: [ ['mlp_moe', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']], ['embed_moe', ['fsdp', 'fsdp_transpose', 'context']], ['embed_moe', ['fsdp', 'context']], + ['embed_router', ['fsdp', 'fsdp_transpose', 'context']], + ['embed_router', ['fsdp', 'context']], # Standard MLP / Dense Layers / Model Structure ['activation_mlp', ['tensor', 'tensor_sequence']], diff --git a/src/maxtext/configs/post_train/rl_mt_jt.yml b/src/maxtext/configs/post_train/rl_mt_jt.yml index 43a37343fc..d44569b507 100644 --- a/src/maxtext/configs/post_train/rl_mt_jt.yml +++ b/src/maxtext/configs/post_train/rl_mt_jt.yml @@ -58,6 +58,10 @@ logical_axis_rules: [ ['embed_moe', ['fsdp', 'context_autoregressive', 'tensor_transpose']], ['embed_moe', ['fsdp', 'fsdp_transpose', 'context_autoregressive']], ['embed_moe', ['fsdp', 'context_autoregressive']], + ['embed_router', ['fsdp', 'fsdp_transpose', 'context_autoregressive', 'tensor_transpose']], + ['embed_router', ['fsdp', 'context_autoregressive', 'tensor_transpose']], + ['embed_router', ['fsdp', 'fsdp_transpose', 'context_autoregressive']], + ['embed_router', ['fsdp', 'context_autoregressive']], ['norm', ['tensor', 'tensor_transpose', 'tensor_sequence']], ['layers', 'stage'], ['kv', []], diff --git a/src/maxtext/configs/pyconfig.py b/src/maxtext/configs/pyconfig.py index a68c54cf8e..f71d1405a4 100644 --- a/src/maxtext/configs/pyconfig.py +++ b/src/maxtext/configs/pyconfig.py @@ -89,7 +89,7 @@ def _resolve_or_infer_config(argv: list[str] | None = None, **kwargs) -> tuple[s # if passing at least two arguments via list (no kwargs), then we have to specify # first one as either "" or python script like train_rl.py or train.py # the second argument is the yaml file - if len(argv) >= 2 and argv[1].endswith(".yml"): + if len(argv) >= 2 and argv[1].lower().endswith((".yml", ".yaml")): return resolve_config_path(argv[1]), argv[2:] module = _module_from_path(argv[0]) if len(argv) > 0 else None if module not in _CONFIG_FILE_MAPPING: @@ -349,6 +349,7 @@ def __init__(self, pydantic_config: types.MaxTextConfig): final_dict["weight_dtype"] if not final_dict["mu_dtype"] else jnp.dtype(final_dict["mu_dtype"]) ) final_dict["logical_axis_rules"] = _lists_to_tuples(final_dict["logical_axis_rules"]) + final_dict["logical_axis_rules_for_eval"] = _lists_to_tuples(final_dict["logical_axis_rules_for_eval"]) final_dict["data_sharding"] = _lists_to_tuples(final_dict["data_sharding"]) final_dict["decoder_block"] = DecoderBlockType(final_dict["decoder_block"]) diff --git a/src/maxtext/configs/pyconfig_deprecated.py b/src/maxtext/configs/pyconfig_deprecated.py index 51e9cc1679..c64ee66a00 100644 --- a/src/maxtext/configs/pyconfig_deprecated.py +++ b/src/maxtext/configs/pyconfig_deprecated.py @@ -705,7 +705,7 @@ def __init__(self, argv: list[str], **kwargs): def user_init(raw_keys): """Transformations between the config data and configs used at runtime""" if raw_keys["run_name"] == "": - raw_keys["run_name"] = os.environ.get("JOBSET_NAME") # using XPK default + raw_keys["run_name"] = os.environ.get("JOBSET_NAME") # set by some JobSet launchers, e.g. for Pathways if raw_keys["run_name"] == "": now = datetime.datetime.now() timestamp = now.strftime("%Y-%m-%d-%H-%M") diff --git a/src/maxtext/configs/tpu/v4/22b.sh b/src/maxtext/configs/tpu/v4/22b.sh index 83fad6c38e..fcdff96831 100644 --- a/src/maxtext/configs/tpu/v4/22b.sh +++ b/src/maxtext/configs/tpu/v4/22b.sh @@ -18,7 +18,7 @@ echo "Running 22b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script for training: # bash src/maxtext/configs/tpu/v4/22b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v4/52b.sh b/src/maxtext/configs/tpu/v4/52b.sh index f774679c50..07b67cab1a 100644 --- a/src/maxtext/configs/tpu/v4/52b.sh +++ b/src/maxtext/configs/tpu/v4/52b.sh @@ -18,7 +18,7 @@ echo "Running 52b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v4/52b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5e/128b.sh b/src/maxtext/configs/tpu/v5e/128b.sh index 2fdc307b25..7d82ef9d07 100644 --- a/src/maxtext/configs/tpu/v5e/128b.sh +++ b/src/maxtext/configs/tpu/v5e/128b.sh @@ -4,7 +4,7 @@ echo "Running 128b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5e/128b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5e/16b.sh b/src/maxtext/configs/tpu/v5e/16b.sh index da66a72adf..39ffb531ab 100644 --- a/src/maxtext/configs/tpu/v5e/16b.sh +++ b/src/maxtext/configs/tpu/v5e/16b.sh @@ -4,7 +4,7 @@ echo "Running 16b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5e/16b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5e/32b.sh b/src/maxtext/configs/tpu/v5e/32b.sh index ff870948c3..6884307779 100644 --- a/src/maxtext/configs/tpu/v5e/32b.sh +++ b/src/maxtext/configs/tpu/v5e/32b.sh @@ -4,7 +4,7 @@ echo "Running 32b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5e/32b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5e/64b.sh b/src/maxtext/configs/tpu/v5e/64b.sh index faec50b553..f6919f479f 100644 --- a/src/maxtext/configs/tpu/v5e/64b.sh +++ b/src/maxtext/configs/tpu/v5e/64b.sh @@ -4,7 +4,7 @@ echo "Running 64b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5e/64b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5e/gpt3_175b.sh b/src/maxtext/configs/tpu/v5e/gpt3_175b.sh index d9c95ce2c6..4643a4c2ff 100644 --- a/src/maxtext/configs/tpu/v5e/gpt3_175b.sh +++ b/src/maxtext/configs/tpu/v5e/gpt3_175b.sh @@ -3,7 +3,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5e/gpt3_175b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5e/llama2_13b.sh b/src/maxtext/configs/tpu/v5e/llama2_13b.sh index e79997da64..7f90b39578 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_13b.sh +++ b/src/maxtext/configs/tpu/v5e/llama2_13b.sh @@ -3,7 +3,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5e/llama2_13b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5e/llama2_70b.sh b/src/maxtext/configs/tpu/v5e/llama2_70b.sh index aee3babe44..341c014c70 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_70b.sh +++ b/src/maxtext/configs/tpu/v5e/llama2_70b.sh @@ -3,7 +3,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5e/llama2_70b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5e/llama2_7b.sh b/src/maxtext/configs/tpu/v5e/llama2_7b.sh index e4aa662b2c..d9a595b381 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_7b.sh +++ b/src/maxtext/configs/tpu/v5e/llama2_7b.sh @@ -3,7 +3,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5e/llama2_7b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5p/1024b.sh b/src/maxtext/configs/tpu/v5p/1024b.sh index d42bc8a70e..e3b91c8f15 100644 --- a/src/maxtext/configs/tpu/v5p/1024b.sh +++ b/src/maxtext/configs/tpu/v5p/1024b.sh @@ -4,7 +4,7 @@ echo "Running 1024b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5p/1024b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5p/128b.sh b/src/maxtext/configs/tpu/v5p/128b.sh index e53c6a9ede..517fa341f5 100644 --- a/src/maxtext/configs/tpu/v5p/128b.sh +++ b/src/maxtext/configs/tpu/v5p/128b.sh @@ -4,7 +4,7 @@ echo "Running 128b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5p/128b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5p/256b.sh b/src/maxtext/configs/tpu/v5p/256b.sh index 984f281e32..0687760d18 100644 --- a/src/maxtext/configs/tpu/v5p/256b.sh +++ b/src/maxtext/configs/tpu/v5p/256b.sh @@ -4,7 +4,7 @@ echo "Running 256b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # PLATFORM (Optional, can be "gke" or "gce", default is "gce") # # Example to invoke this script: diff --git a/src/maxtext/configs/tpu/v5p/32b.sh b/src/maxtext/configs/tpu/v5p/32b.sh index de129f0ec5..306d1be2d5 100644 --- a/src/maxtext/configs/tpu/v5p/32b.sh +++ b/src/maxtext/configs/tpu/v5p/32b.sh @@ -4,7 +4,7 @@ echo "Running 32b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5p/32b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5p/512b.sh b/src/maxtext/configs/tpu/v5p/512b.sh index 42e6c7cec7..8b2378f7bb 100644 --- a/src/maxtext/configs/tpu/v5p/512b.sh +++ b/src/maxtext/configs/tpu/v5p/512b.sh @@ -4,7 +4,7 @@ echo "Running 512b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # PLATFORM (Optional, can be "gke" or "gce", default is "gce") # # Example to invoke this script: diff --git a/src/maxtext/configs/tpu/v5p/64b.sh b/src/maxtext/configs/tpu/v5p/64b.sh index 90c01c2c74..2897f4ce68 100644 --- a/src/maxtext/configs/tpu/v5p/64b.sh +++ b/src/maxtext/configs/tpu/v5p/64b.sh @@ -4,7 +4,7 @@ echo "Running 64b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5p/64b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5p/llama2_70b.sh b/src/maxtext/configs/tpu/v5p/llama2_70b.sh index 2228794617..f6cf78823e 100644 --- a/src/maxtext/configs/tpu/v5p/llama2_70b.sh +++ b/src/maxtext/configs/tpu/v5p/llama2_70b.sh @@ -5,7 +5,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5p/llama2_70b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v5p/llama2_7b.sh b/src/maxtext/configs/tpu/v5p/llama2_7b.sh index 774b7ffe61..8194c47c42 100644 --- a/src/maxtext/configs/tpu/v5p/llama2_7b.sh +++ b/src/maxtext/configs/tpu/v5p/llama2_7b.sh @@ -5,7 +5,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v5p/llama2_7b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v6e/gemma2_27b.sh b/src/maxtext/configs/tpu/v6e/gemma2_27b.sh index 970722c6f8..ce1b477359 100644 --- a/src/maxtext/configs/tpu/v6e/gemma2_27b.sh +++ b/src/maxtext/configs/tpu/v6e/gemma2_27b.sh @@ -3,7 +3,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v6e/gemma2_27b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v6e/gemma2_9b.sh b/src/maxtext/configs/tpu/v6e/gemma2_9b.sh index 2728f46aed..0df42d5ba1 100644 --- a/src/maxtext/configs/tpu/v6e/gemma2_9b.sh +++ b/src/maxtext/configs/tpu/v6e/gemma2_9b.sh @@ -3,7 +3,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v6e/gemma2_9b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v6e/gemma3_27b.sh b/src/maxtext/configs/tpu/v6e/gemma3_27b.sh index 80824e9f46..4af86e4142 100644 --- a/src/maxtext/configs/tpu/v6e/gemma3_27b.sh +++ b/src/maxtext/configs/tpu/v6e/gemma3_27b.sh @@ -3,7 +3,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v6e/gemma3_27b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v6e/gpt3_175b.sh b/src/maxtext/configs/tpu/v6e/gpt3_175b.sh index b7f3b40454..95d0d67d16 100644 --- a/src/maxtext/configs/tpu/v6e/gpt3_175b.sh +++ b/src/maxtext/configs/tpu/v6e/gpt3_175b.sh @@ -3,7 +3,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v6e/gpt3_175b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh b/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh index 8279acede6..ca1dba6cec 100644 --- a/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh +++ b/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh @@ -3,7 +3,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh b/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh index c5bd75d94f..918e19f22f 100644 --- a/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh +++ b/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh @@ -3,7 +3,7 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) +# RUN_NAME (Required, unless run_name is already set in base.yml or the JOBSET_NAME environment variable is set) # # Example to invoke this script: # bash src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh RUN_NAME="" OUTPUT_PATH="gs://" diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index ba872c9a4b..166db8c136 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -249,6 +249,7 @@ class ProfilerType(str, Enum): "deepseek3-671b", "deepseek3-671b-2dfsdp", "deepseek3-671b-batchsplit", + "deepseek3-671b-lineage", "deepseek3-test", "deepseek3-tiny", "deepseek3.2-671b", @@ -290,11 +291,13 @@ class ProfilerType(str, Enum): "qwen3-vl-2b", "qwen3-vl-4b", "qwen3-vl-30b-a3b", - "cosmos3-nano-reasoner", - "cosmos3-super-reasoner", + "weaver-mini", + "weaver-max", + "weaver-mini-diffuser", "qwen3-next-80b-a3b", "qwen3-omni-30b-a3b", "qwen3-custom-30b-a3b", + "qwen3.5-9b", "qwen3.5-35b-a3b", "qwen3.5-35b-a3b-fp8", "qwen3.5-35b-fp8", @@ -330,6 +333,7 @@ class RunInfo(BaseModel): description="The name of the run. Checkpoints will be stored under this name.", ) model_name: ModelName = Field("default", description="The name of the model configuration to use.") + use_m3_model: bool = Field(False, description="Use the m3 backend for a registered model.") override_model_config: bool = Field(False, description="If True, allows overriding model parameters via CLI.") override_logical_axis_rules: bool = Field( False, @@ -620,9 +624,51 @@ class Quantization(BaseModel): " function) and does not control quantization." ), ) + quantize_logits_proj: bool = Field( + False, + description=( + "If True, quantizes the output logits (logits_dense) projection when quantization is enabled." + " Targets the `logits_dense` module (matching path regex r'decoder/logits_dense.*'). Default is False." + " Only supported with `quantization=fp8_full`, `logits_via_embedding=False`, and `num_vocab_tiling = 1`." + ), + ) + logits_proj_quant_calibration_method: str = Field( + "", + description=( + "Calibration method for the output logits (logits_dense) projection when `quantize_logits_proj=True`." + " If empty (default), inherits `weight_quantization_calibration_method` and" + " `act_quantization_calibration_method`. Set to e.g. 'absmax' to force absmax calibration." + ), + ) + logits_proj_bwd_quant_calibration_method: str = Field( + "", + description=( + "Backward (gradient) calibration method for the output logits (logits_dense) projection when" + " `quantize_logits_proj=True`. Applies to both the activation-gradient and weight-gradient arms. If empty" + " (default), inherits `bwd_quantization_calibration_method` and `drhs_grad_quantization_calibration_method`." + " Set to e.g. 'absmax' to force absmax calibration." + ), + ) + sa_qproj_bwd_bf16: bool = Field( + False, + description=( + "If True, keeps the backward pass in bfloat16 (unquantized) for self-attention query projection matmuls" + " when `quantization=fp8_full`." + ), + ) kv_quant_axis: KvQuantAxis = Field(KvQuantAxis.HEADS_AND_DKV, description="Axes to quantize over for the KV cache.") kv_quant_dtype: Literal["int8", "int4"] = Field("int8", description="Data type for KV cache quantization.") - quantization_local_shard_count: int = Field(-1, description="Shards the range finding operation for quantization.") + quantization_local_shard_count: int = Field( + -1, + description="AQT int8 only: number of local shards along the contraction axis of weight-gradient matmuls, each" + " with its own scale. -1 means jax.local_device_count(); 0 leaves the weight-gradient matmul unquantized.", + ) + qwix_bwd_weight_grad_tile_count: PositiveInt = Field( + 1, + description="Qwix: split the contraction axis of each weight-gradient matmul into this many tiles, each with" + " its own quantization scale. 1 disables tiling. Tiles are summed after the matmul, so under FSDP tiling" + " makes XLA all-reduce the full weight gradient instead of reduce-scattering it. fp8_full ignores it.", + ) use_qwix_quantization: bool = Field(False, description="Whether to use qwix for quantization.") use_manual_quantization: bool = Field( False, @@ -640,6 +686,15 @@ class Quantization(BaseModel): "absmax", description="Quantization calibration method used for gradients.", ) + drhs_grad_quantization_calibration_method: str | None = Field( + None, + description=( + "Calibration for the cotangent in the weight-gradient matmul only (qwix drhs_grad; dW = X^T dY). None " + "inherits bwd_quantization_calibration_method. With absmax, this arm's per-channel scale reduces over the " + "token axis and therefore across every data shard, including DCN slices; a fixed range (e.g. 'fixed,0.01') " + "removes that reduction. The activation-gradient arm (dX = dY W^T) is unaffected." + ), + ) weight_sparsity_n: int | None = Field( None, description=("The 'N' in N:M sparsity, representing the maximum number of non-zero" " values in each block."), @@ -724,6 +779,10 @@ class MTP(BaseModel): 0, description="Specifies which MTP layer is used to calculate metrics.", ) + mtp_reuse_input_embedding: bool = Field( + False, + description="Reuse the main decoder's token embeddings instead of a second lookup.", + ) class LogitsAndLoss(BaseModel): @@ -864,6 +923,13 @@ class MlaAttention(BaseModel): " instead of running full projection + jnp.split." ), ) + use_mla_absorbed_mqa: bool = Field( + False, + description=( + "Whether to use Absorbed-Query Latent MQA mode in MLA. Trades higher prefill/training FLOPs " + "for 57x smaller KV cache footprint and memory bandwidth during autoregressive decode." + ), + ) class CompressedAttention(BaseModel): @@ -929,6 +995,11 @@ class SplashAttention(BaseModel): sa_block_q: int = Field(512, description="Block size for Q in splash attention.") sa_block_kv: int = Field(512, description="Block size for KV in splash attention.") sa_block_kv_compute: int = Field(512, description="Block size for KV compute in splash attention.") + eval_sa_block_q: int = Field(512, description="Block size for Q in splash attention during evaluation.") + eval_sa_block_kv: int = Field(512, description="Block size for KV in splash attention during evaluation.") + eval_sa_block_kv_compute: int = Field( + 512, description="Block size for KV compute in splash attention during evaluation." + ) sa_block_q_dkv: int = Field(512, description="Block size for Q_dkv in splash attention.") sa_block_kv_dkv: int = Field(512, description="Block size for KV_dkv in splash attention.") sa_block_kv_dkv_compute: int = Field(512, description="Block size for KV_dkv compute in splash attention.") @@ -942,6 +1013,9 @@ class SplashAttention(BaseModel): sa_q_layout: str = Field("HEAD_DIM_MINOR", description="Layout for Q in splash attention.") sa_k_layout: str = Field("HEAD_DIM_MINOR", description="Layout for K in splash attention.") sa_v_layout: str = Field("HEAD_DIM_MINOR", description="Layout for V in splash attention.") + eval_sa_q_layout: str = Field("HEAD_DIM_MINOR", description="Layout for Q in splash attention during evaluation.") + eval_sa_k_layout: str = Field("HEAD_DIM_MINOR", description="Layout for K in splash attention during evaluation.") + eval_sa_v_layout: str = Field("HEAD_DIM_MINOR", description="Layout for V in splash attention during evaluation.") use_splash_scheduler: bool = Field(False, description="Use experimental splash attention scheduler.") ring_scan_unroll: NonNegativeInt = Field( 1, @@ -952,10 +1026,22 @@ class SplashAttention(BaseModel): ) sa_fuse_reciprocal: bool = Field(True, description="Maps to fuse_reciprocal in SplashConfig.") sa_use_base2_exp: bool = Field(True, description="Maps to use_base2_exp in SplashConfig.") + sa_qk_diag_skip: bool = Field(False, description="Maps to qk_diag_skip in SplashConfig.") + sa_sv_diag_skip: bool = Field(False, description="Maps to sv_diag_skip in SplashConfig.") + sa_qk_diag_grid: int = Field(2, description="Maps to qk_diag_grid in SplashConfig.") # If None, each local_sa_* flag inherits from the corresponding sa_* flag. local_sa_block_q: int | None = Field(None, description="Block size for Q in local splash attention.") local_sa_block_kv: int | None = Field(None, description="Block size for KV in local splash attention.") local_sa_block_kv_compute: int | None = Field(None, description="Block size for KV compute in local splash attention.") + eval_local_sa_block_q: int | None = Field( + None, description="Block size for Q in local splash attention during evaluation." + ) + eval_local_sa_block_kv: int | None = Field( + None, description="Block size for KV in local splash attention during evaluation." + ) + eval_local_sa_block_kv_compute: int | None = Field( + None, description="Block size for KV compute in local splash attention during evaluation." + ) local_sa_block_q_dkv: int | None = Field(None, description="Block size for Q_dkv in local splash attention.") local_sa_block_kv_dkv: int | None = Field(None, description="Block size for KV_dkv in local splash attention.") local_sa_block_kv_dkv_compute: int | None = Field( @@ -969,6 +1055,15 @@ class SplashAttention(BaseModel): local_sa_q_layout: str | None = Field(None, description="Layout for Q in local splash attention.") local_sa_k_layout: str | None = Field(None, description="Layout for K in local splash attention.") local_sa_v_layout: str | None = Field(None, description="Layout for V in local splash attention.") + eval_local_sa_q_layout: str | None = Field( + None, description="Layout for Q in local splash attention during evaluation." + ) + eval_local_sa_k_layout: str | None = Field( + None, description="Layout for K in local splash attention during evaluation." + ) + eval_local_sa_v_layout: str | None = Field( + None, description="Layout for V in local splash attention during evaluation." + ) local_use_splash_scheduler: bool | None = Field(None, description="Use experimental local splash attention scheduler.") local_sa_fuse_reciprocal: bool | None = Field(None, description="Maps to local fuse_reciprocal in SplashConfig.") local_sa_use_base2_exp: bool | None = Field(None, description="Maps to local use_base2_exp in SplashConfig.") @@ -1006,6 +1101,37 @@ class SplashAttention(BaseModel): ) +def _reject_unless(setting: str, requirements: list[tuple[bool, str]]) -> None: + """Raises a ValueError naming every requirement of `setting` that is not met. + + Args: + setting: the setting being validated, as the user wrote it (e.g. "moe_unsplit_prefused_wi=True"). + requirements: (met, requirement) pairs, `requirement` describing what `setting` needs. + """ + unmet = [requirement for met, requirement in requirements if not met] + if unmet: + raise ValueError(f"{setting} requires: {'; '.join(unmet)}.") + + +# The MLP checkpoint names that the remat_policy presets save or offload (Decoder.minimal_policy and get_remat_policy +# in nnx_decoders.py, which point back here: keep the two in sync). +# remat_policy=custom names them per tensor instead (mlpwi_0, mlpwo, moe_mlpwi_0, ... = device / offload). +_REMAT_PRESET_MLP_NAMES = { + "minimal": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"), + "minimal_with_context": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"), + "minimal_flash": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"), + "minimal_with_quantization": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"), + "minimal_with_context_and_quantization": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"), + "minimal_offloaded": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"), + "save_dot_except_mlpwi": ("mlpwo",), +} +# The checkpoint names of the routed experts' up-projection output in the sparse-matmul path: the gate half carries +# mlpwi_0 and moe_mlpwi_0, the up half moe_mlpwi_1 (moe.py gmm_up). +_MOE_UP_PROJECTION_NAMES = ("mlpwi_0", "moe_mlpwi_0", "moe_mlpwi_1") +# ... and of its down-projection output: mlpwo and moe_mlpwo. +_MOE_DOWN_PROJECTION_NAMES = ("mlpwo", "moe_mlpwo") + + class MoEGeneral(BaseModel): """General configuration for Mixture of Experts (MoE) layers.""" @@ -1016,6 +1142,63 @@ class MoEGeneral(BaseModel): -1.0, description="Ragged buffer factor. If < 0, ragged buffer is worst case size.", ) + eval_ragged_buffer_factor: float = Field( + -1.0, + description="Evaluation ragged buffer factor, applied under eval logical axis rules. If < 0, worst case size.", + ) + moe_dropless_fallback: Literal["step", "layer"] | None = Field( + None, + description=( + "What to do when the ragged buffer (ragged_buffer_factor > 0) would drop tokens. None: drop them." + " 'step': roll the step's update back in-graph and replay the whole train step with a dropless buffer." + " 'layer': per MoE layer, calculate the overflow from the router's top-k and take a chunked dropless" + " branch (lax.cond) for that layer only; no step replay." + ), + ) + retry_dropless_first_steps: int = Field( + 0, + description=( + "With moe_dropless_fallback='step', run the first N steps after the loop's start step (i.e. steps with " + "step - start_step < N; equal to step < N when training starts at step 0) directly with the dropless " + "program, skipping the attempt with the normal program. 0 = off." + ), + ) + first_phase_ragged_buffer_factor: float = Field( + 0.0, + description=( + "With retry_dropless_first_steps=N > 0, run the first N steps with a precompiled first-phase program whose " + "RoutedMoE modules use this ragged buffer factor (instead of the dropless program); a first-phase step that " + "still drops tokens is replayed with the dropless program. Must be >= ragged_buffer_factor. 0 = off." + ), + ) + log_required_ragged_buffer_factor: bool = Field( + False, + description=( + "Probe: log per train step and per MoE layer the minimum ragged_buffer_factor that would have avoided drops " + "(max over shards of the tokens routed to a shard / the per-shard buffer at factor 1). Adds one scalar " + "max all-reduce per MoE layer and token chunk, and a device-to-host fetch per step." + ), + ) + warmup_programs_in_init: bool = Field( + False, + description=( + "After the setup precompiles and before init_stop/run_start are logged, execute each precompiled train " + "program (normal, first-phase, dropless) and eval program (eval, eval dropless) once through its jit on a " + "synthetic batch (PRNG token ids, never the dataset), discard the outputs so parameters, optimizer state " + "and routed biases are unchanged, block on completion and run a cross-host barrier. Moves the first " + "execution (and the hosts' arrival skew) out of the scored window. Skipped when the programs are not " + "precompiled (compiled_trainstep_file or AutoPGLE)." + ), + ) + warmup_skip_train_dropless: bool = Field( + False, + description=( + "Requires warmup_programs_in_init=True. Skips the train_dropless program in that warmup (other programs are " + "still warmed). Its warmup runs on a copy of the donated train state, so HBM holds two state copies plus " + "the dropless program reservation, which can OOM. Use only when that warmup OOMs: skipping moves the " + "dropless program's first load and execution cost into the first dropless replay, if any." + ), + ) moe_exclude_padding_from_routing: bool = Field( False, description=( @@ -1036,14 +1219,6 @@ class MoEGeneral(BaseModel): " engine, each rematerialized, so its backward holds one chunk's worst-case buffers at a time." ), ) - retry_dropless_first_steps: int = Field( - 0, - description=( - "With retry_when_tokens_dropped=True, run the first N steps after the loop's start step directly " - "with the dropless program (skipping the normal attempt) or with first_phase_ragged_buffer_factor. 0 = off. " - "Training engine only (pre_train/train.py ignores it)." - ), - ) retry_dropless_final_micro_batch: bool = Field( False, description=( @@ -1063,15 +1238,6 @@ class MoEGeneral(BaseModel): "moe_dropped_assignments per step. Adds a scalar psum and pmax per MoE layer call." ), ) - first_phase_ragged_buffer_factor: float = Field( - 0.0, - description=( - "With retry_dropless_first_steps=N > 0, run the first N steps with a first-phase program whose " - "RoutedMoE modules use this ragged buffer factor (instead of the dropless program directly); a first-phase " - "step that still drops tokens is replayed with the dropless program. Must be >= ragged_buffer_factor. 0 = off. " - "Training engine only (pre_train/train.py ignores it)." - ), - ) num_moe_token_chunks: PositiveInt = Field( 1, description=( @@ -1082,6 +1248,67 @@ class MoEGeneral(BaseModel): " use_ring_of_experts=True." ), ) + moe_accumulate_chunk_wgrad: bool = Field( + False, + description=( + "With num_moe_token_chunks>1, chain the expert weights through the token chunks: chunk 0 gathers and" + " quantizes them once and later chunks reuse the result. In the backward pass each chunk's weight-gradient" + " tgmm accumulates in place onto the later chunks' gradient (tgmm_v2 partial_sum) instead of" + " materializing and summing one gradient per chunk. Requires use_tokamax_gmm=True, use_gmm_v2=True," + " num_moe_emb_chunks=0 and prefuse_moe_weights=False." + ), + ) + moe_accumulate_wi_dlhs: bool = Field( + False, + description=( + "Route the routed-MoE wi_1 GMM input through the wi_0 GMM so that in the backward pass the wi_1 DLHS" + " is accumulated in place into the wi_0 DLHS (gmm_v2 partial_sum) instead of summed by a separate add." + " Also folds per-tensor DLHS scales into the gmm_v2 kernel for all MoE GMMs. Requires" + " use_tokamax_gmm=True, use_gmm_v2=True and prefuse_moe_weights=False." + ), + ) + moe_flat_fsdp_weights: bool = Field( + False, + description=( + "Store every routed-expert weight as a 2D [num_experts * rows, cols] matrix whose first dim is" + " sharded over (expert, fsdp) instead of the 3D [E, rows, cols] layout sharded (expert | fsdp | -)." + " The per-device FSDP shard is then a contiguous, tile-aligned row block at every FSDP size, the fp8" + " weight all-gather / weight-grad reduce-scatter run on the major dim, and the gathered weight is a free" + " bitcast of the [E, rows, cols] operand the gmm kernels expect (no relayout copies, no 7/8 sublane" + " padding). Requires shard_embed_moe_on_fsdp with fixed weight calibration (explicit QAG) and the" + " tokamax gmm_v2 custom-VJP path. 3D checkpoints loaded via load_parameters_path are converted to the flat" + " layout on load; full-state resumes must come from a flat checkpoint." + ), + ) + moe_gmm_v2_dlhs_transpose_rhs: bool = Field( + False, + description="Use native transpose_rhs in GMM v2 backward DLHS instead of an explicit RHS transpose.", + ) + moe_bwd_prequant_before_unsort: bool = Field( + False, + description=( + "Fuse the TC unsort-reduce into the wo GMM custom VJP so backward quantization runs on the smaller" + " [N_tokens, D // 128, 128] cotangent before the TC unsort-gather into the expert buffer." + " With quantization, needs fixed weight, act and bwd calibration (per-tensor scales applied in token order)." + ), + ) + moe_combine_bwd_direct_qarray: bool = Field( + False, + description=( + "With moe_bwd_prequant_before_unsort=True, fuse the EP combine psum_scatter into the wo GMM custom VJP" + " so backward quantizes the local cotangent once before the EP all-gather and TC unsort-gather." + " Incompatible with moe_quantize_combine_bwd_method. Each EP shard scales its cotangent before the all-gather, so" + " the fixed weight, act and bwd calibration required by moe_bwd_prequant_before_unsort must match across" + " shards." + ), + ) + moe_3d_dlhs_partial_sum: bool = Field( + False, + description=( + "Skip redundant zero-initialization in 3D DLHS gmm_v2 when accumulating onto an existing 3D partial_sum" + " (e.g. wi_0 DLHS accumulating onto wi_1 DLHS)." + ), + ) moe_chunk_barrier: bool = Field( False, description=( @@ -1112,6 +1339,16 @@ class MoEGeneral(BaseModel): " use_ring_of_experts=True and num_moe_emb_chunks=0." ), ) + moe_log_max_load_ratio: bool = Field( + False, + description=( + "Debug only. Log learning/moe_max_load_ratio{,_mean}: the max (mean) over MoE layers of" + " the largest (tokens routed to an expert shard) / (perfectly balanced load) over token" + " chunks and expert shards. A step overflows the ragged buffer iff the max exceeds" + " ragged_buffer_factor (up to buffer rounding). Requires sparse_matmul=True," + " te_moe_block=False, use_ring_of_experts=True and use_ragged_sort=True." + ), + ) moe_expert_input_dim: int = Field( -1, @@ -1123,10 +1360,35 @@ class MoEGeneral(BaseModel): description="Padded intermediate dimension at MoE layer for efficient GMM_v2 kernel execution.", ) load_balance_loss_weight: NonNegativeFloat = Field(0.0, description="Weight for the load balancing auxiliary loss.") + moe_use_megatron_seq_aux_loss: bool = Field( + False, + description=( + "When True, sigmoid routers (routed_score_func='sigmoid', with te_moe_block=False) use" + " Megatron-LM's sequence-wise auxiliary load-balancing loss (fp32 per-token normalized sigmoid scores," + " token fractions from an unbiased, ungrouped top-k over the full sequence, summed over MoE layers)." + " When False (default), sigmoid routers use the Switch Transformer-style loss averaged over MoE layers." + ), + ) use_custom_sort_vjp: bool = Field( True, description="Whether to use a custom VJP sort for efficient backward pass processing in sparse matmul.", ) + router_topk_matmul_vjp: bool = Field( + False, + description=( + "DeepSeek routing: take the top-k routing weights with a custom VJP whose backward builds the dense" + " [..., num_experts] gradient as a one-hot compare-select-sum over the k slots, instead of the scatter-add" + " that jnp.take_along_axis transposes to. Forward and gradient values are unchanged." + ), + ) + moe_topk_before_ep_all_gather: bool = Field( + False, + description=( + "Ring of experts: run top-k on the local tokens and all-gather the top-k weights and expert ids" + " instead of the full router logits. The token sort still runs after the all-gather." + " Ignored with forced, hash or random routing." + ), + ) use_ring_of_experts: bool = Field( False, description="Whether to use Ring of Experts for sparse matmul expert parallelism.", @@ -1141,6 +1403,38 @@ class MoEGeneral(BaseModel): " Requires use_ring_of_experts=True." ), ) + moe_roe_packed_routing_gather: bool = Field( + False, + description=( + "With ring_of_experts_local_routing and moe_chunk_pipeline: route the layer's local tokens once (every" + " token chunk) and all-gather the top-k weights and expert ids as ONE float32 [batch, seq, 2*top_k] array" + " (ids are exact in float32), instead of two all-gathers (ids, weights) per token chunk: one SparseCore" + " offload per layer instead of 2 x num_moe_token_chunks, and one weight-gradient reduce-scatter instead of" + " one per chunk. Same routing; bit-identical with float32 routing weights (float32_weight_sum=True), while" + " bfloat16 weights get their cotangent reduce-scattered in float32. The per-chunk load-balance loss is" + " computed as without it. Tagged moe_routing_maps. Requires sparse_matmul=True, te_moe_block=False," + " ring_of_experts_local_routing=True, moe_chunk_pipeline=True and use_random_routing=False (random routing" + " draws per call, so one draw for the layer differs from one per chunk)." + ), + ) + moe_router_grad_from_wo_input: bool = Field( + False, + description=( + "Ring of Experts with the ragged sort: run the activation, the down projection (wo GEMM) and the unsort" + " as one op whose backward does not read the down projection's output. The routing-weight gradient" + " is computed as from the down projection's input and the routing weight is applied" + " on the [rows, mlp] side, so a rematerialized layer no longer recomputes the down-projection GEMM in the" + " backward, and the combine backward's pass over [rows, hidden] tensors moves to [rows, mlp] ones. The" + " forward is bit-identical. NUMERICS CHANGE: each gradient keeps one rounding to the activation dtype, at" + " a different point, so the gradients agree with the unfused ops to bf16 rounding, not bit for bit." + " Requires sparse_matmul=True, te_moe_block=False, use_ring_of_experts=True and use_ragged_sort=True, and" + " the down projection to be the bare GEMM of the sorted rows: no tensor parallelism, mlp_bias=False, no" + " weight quantization, moe_tc_ragged_sort, moe_accumulate_chunk_wgrad, moe_bwd_prequant_before_unsort and" + " moe_quantize_token_all_gather off, num_moe_emb_chunks=0; and a remat policy that does not save the routed" + " experts' down-projection output (mlpwo, moe_mlpwo; the minimal* and save_dot_except_mlpwi policies save" + " mlpwo): the fused op never forms it, so the policy would silently stop saving the MoE tensor." + ), + ) ring_of_experts_row_major_reduce_scatter: bool = Field( False, description=( @@ -1190,10 +1484,101 @@ class MoEGeneral(BaseModel): False, description="Whether to use ragged kernel for sorting, improve performance when EP is enabled.", ) + moe_routing_index_math: Literal["sort", "count"] = Field( + "sort", + description=( + "How the ring-of-experts ragged sort computes its routing index maps. 'sort' (default): the expert" + " permutation and its inverse by two stable argsorts. 'count': the permutation by one single-operand" + " packed uint32 sort (XLA lowers the stable argsort with an extra, duplicated iota operand); the sorted" + " token indices by division instead of a gather; the inverse permutation by counting, on this shard's" + " slots only, instead of a second argsort; the router-weight gradient back to slot order by a gather" + " instead of a scatter; the gather-reduce validity sort packed into one uint32 operand; and, for tokens" + " without forced ids under router replay or moe_exclude_padding_from_routing, the router's top-k indices" + " by k argmax passes instead of a sort. Exact: the same permutation, so the same outputs and gradients," + " dropped tokens included. With moe_tc_ragged_sort, which replaces the ragged sort of the dropping program," + " it applies to the dropless program's ragged sort only. Requires sparse_matmul=True, te_moe_block=False," + " use_ring_of_experts=True and use_ragged_sort=True." + ), + ) + moe_ragged_sort_unique_pair_argsort: bool = Field( + False, + description=( + "With moe_routing_index_math='sort': run the ring-of-experts ragged sort's argsorts (the expert" + " permutation, its inverse, and the SparseCore gather-reduce's validity sort) as unstable sorts of the" + " unique (key, index) pairs instead of stable argsorts. XLA lowers a stable jnp.argsort in the forward" + " pass with a third, duplicated iota operand; the pair sort has two operands and exactly one result, the" + " same permutation. Exact (same values). An alternative to moe_routing_index_math='count', which removes" + " these sorts instead. Requires sparse_matmul=True, te_moe_block=False, use_ring_of_experts=True," + " use_ragged_sort=True and moe_routing_index_math='sort'." + ), + ) + moe_quantize_combine_bwd_method: str | None = Field( + "", + description=( + "Quantization method for Token Combine backward All-Gather. " + "'' = unquantized BF16; 'rowwise' = dynamic token-wise FP8 E5M2; " + "'fixed,' = static per-tensor FP8 E5M2 (e.g. 'fixed,1.0', 'fixed,57344')." + ), + ) ragged_sort_use_single_sparsecore: bool = Field( False, description="Whether to run ragged sort kernels on 1 SparseCore instead of all SparseCores.", ) + moe_tc_ragged_sort: bool = Field( + False, + description=( + "Ring-of-experts ragged sort/unsort on TensorCore (DMA Pallas kernels ported from lineage)" + " instead of the SparseCore ragged_gather / ragged_gather_reduce kernels. Requires" + " use_ring_of_experts, use_ragged_sort and ragged_buffer_factor > 0." + ), + ) + moe_tc_ragged_gather_block_size: int = Field(1024, description="Rows per grid step of the TC ragged gather.") + moe_tc_ragged_reduce_block_size: int = Field(896, description="VMEM rows per block of the TC ragged gather-reduce.") + moe_tc_ragged_mask_padding: bool = Field( + True, description="Zero TC ragged-gather buffer rows past the valid count (the kernel leaves them uninitialized)." + ) + moe_tc_ragged_weights_on_activation: bool = Field( + False, + description=( + "With moe_tc_ragged_sort: multiply routing weights into the expert activation (buffer, mlp)" + " before the wo matmul instead of into the (buffer, emb) expert output before the unsort." + ), + ) + moe_tc_ragged_flatten_block_size: int = Field( + 0, + description=( + "If > 0, relayout (N, D/128, 128) <-> (N, D) around the TC ragged kernels with a Pallas kernel over" + " only the valid rows (lineage ragged_flatten), this many rows per grid step. 0 uses XLA reshapes." + ), + ) + moe_tc_ragged_3d_gmm: bool = Field( + False, + description=( + "With moe_tc_ragged_sort (tokamax gmm_v2): keep the sorted token buffer in the TC kernels' 3D layout" + " (buffer, emb // 128, 128). The wi gmm_v2 / tgmm_v2 kernels consume it and the wo gmm_v2 produces it" + " directly (relayout in VMEM), removing the (N, D/128, 128) <-> (N, D) relayouts around the sorted" + " buffer. Embed-dim GMM tiles must be multiples of 128 that divide the embed dim; tiles that are not" + " multiples of 1024 (e.g. 3584 / 1792) DMA whole (emb // 128, 128) slabs and relayout the window in VMEM." + ), + ) + moe_tc_ragged_3d_dispatch: bool = Field( + False, + description=( + "With moe_tc_ragged_sort and moe_tc_ragged_3d_gmm under ring of experts: reshape the local tokens to" + " (batch, seq, emb // 128, 128) before the expert-parallel all-gather, and keep the TC unsort output and" + " the combine reduce-scatter in that layout (only the local result is reshaped back). This removes the" + " full-size (N, D/128, 128) <-> (N, D) relayouts of the gathered tokens and of the combine output, in" + " both the forward and the backward pass. Raises if it cannot apply (e.g. emb not a multiple of 128)." + ), + ) + moe_tc_window_routing: bool = Field( + False, + description=( + "With moe_tc_ragged_sort: keep only this shard's [start, start + cap) window of the sorted slots (so the " + "combine metadata sorts cap instead of N * topk slots), skip the unused inverse permutation, use uint32 " + "bit-packed sorts, and gather buffer row weights with unique_indices=True (no backward scatter-sort)." + ), + ) moe_use_direct_token_gather: bool = Field( False, description="Whether to gather tokens directly in expert order instead of materializing Top-K copies.", @@ -1236,6 +1621,28 @@ class MoEGeneral(BaseModel): False, description="Pin FSDP and EP all-gathers in MoE to dedicated SparseCores using compute_on.", ) + gather_scanned_kernels_once: bool = Field( + False, + description=( + "Gather small FSDP-sharded MoE router and attention query kernels once before the scanned layer stack" + " instead of inside every scanned layer." + ), + ) + moe_pin_sparse_core_ep_all_gathers: bool = Field( + False, + description="Pin only the ring-of-experts EP all-gathers (dispatch tokens and routing inputs) to a SparseCore" + " (OR-ed with moe_pin_sparse_core_all_gathers).", + ) + moe_pin_sparse_core_fsdp_all_gathers: bool = Field( + False, + description="Pin only the MoE FSDP weight all-gathers to a SparseCore (OR-ed with" + " moe_pin_sparse_core_all_gathers).", + ) + moe_pin_sparse_core_fsdp_all_gathers_fwd_only: bool = Field( + False, + description="When the MoE FSDP weight all-gathers are pinned, pin only the forward all-gather; its transpose" + " (the weight-gradient reshard) runs outside compute_on, as with the pin off.", + ) moe_fsdp_all_gather_sparse_core_id: int = Field( 0, description="SparseCore ID to pin MoE FSDP all-gathers to when moe_pin_sparse_core_all_gathers is True.", @@ -1318,6 +1725,53 @@ class MoEGeneral(BaseModel): description="Whether to pre-fuse MoE weights (w0 and w1) during initialization. " "This enables a single FFN1 grouped GEMM in sparse MoE paths and passes fused weights directly in vllm_rpa mode.", ) + moe_unsplit_prefused_wi: bool = Field( + False, + description=( + "With prefuse_moe_weights on the sparse-matmul training path: keep the prefused wi = [w0 | w1] one tensor" + " from the parameter to the up-projection GEMM, instead of slicing it into w0 and w1 at the layer entry and" + " concatenating the gathered halves back before the GEMM. One fp8 fake-quant (per output channel, so the" + " same values as rounding the halves), one FSDP all-gather, one operand through the expert-weight" + " prefetch barrier, and no post-gather concatenate. Exact (same values). Requires sparse_matmul=True," + " te_moe_block=False, no weight quantization, an unsharded expert MLP dimension (tensor, tensor_sequence" + " and autoregressive parallelism 1), the default expert-weight sharding (not shard_exp_on_fsdp," + " use_2d_fsdp_sharding, shard_embed_moe_on_fsdp, moe_flat_fsdp_weights or" + " moe_fsdp_use_two_stage_all_gather), num_moe_emb_chunks=0 and moe_accumulate_chunk_wgrad=False." + ), + ) + moe_prefused_swiglu_vjp: bool = Field( + False, + description=( + "With prefuse_moe_weights on the sparse-matmul path: apply the MoE's gated activation (SwiGLU, or any" + " other mlp_activations pair, e.g. GPT-OSS's clipped form) to the whole [gate | up] up-projection output as" + " one custom-VJP op whose backward returns d[gate | up] as one array, computed by the same ops autodiff" + " emits. Split, the two [rows, mlp] cotangents leave their fusion separately and XLA concatenates them in" + " an extra full pass over the expert buffer. Exact (same values; the gradient bits can differ only in the" + " sign of an exact zero). An alternative to moe_wi_gmm_fused_swiglu for configurations that one does not" + " support (other GMM backends, other activations). Requires sparse_matmul=True, te_moe_block=False," + " mlp_bias=False, num_moe_emb_chunks=0, and a remat policy that does not save the routed experts' gate or" + " up half (mlpwi_0, moe_mlpwi_0, moe_mlpwi_1; the minimal* policies save mlpwi_0): the fused activation" + " never forms them, so the policy would silently stop saving the MoE tensor." + ), + ) + moe_wi_gmm_fused_swiglu: bool = Field( + False, + description=( + "Apply the MoE SwiGLU in the epilogue of the wi (up-projection) GMM (megablox ops.gmm_silu, the tokamax" + " gmm_v2 kernel with fuse_act): the forward kernel writes the [buffer, mlp] activation of the rows it" + " computes instead of the [buffer, 2 * mlp] GMM output that a separate elementwise pass reads back over every" + " buffer row, padding included. When differentiated it writes the activation and the gate and up halves" + " ([buffer, 3 * mlp] in all, gate and up rounded as the unfused GMM rounds its output), and the backward is" + " the unfused one. NUMERICS CHANGE: silu reads the float32 accumulator instead of its bfloat16 rounding (the" + " more precise value), so the activation differs by rounding. Requires prefuse_moe_weights=True," + " sparse_matmul=True, te_moe_block=False, use_tokamax_gmm=True, use_gmm_v2=True, no weight quantization," + " mlp_activations=['silu', 'linear'] without clipping, mlp_bias=False, num_moe_emb_chunks=0," + " moe_flat_fsdp_weights, moe_tc_ragged_3d_gmm, moe_quantize_token_all_gather and moe_accumulate_chunk_wgrad" + " off, and a remat policy that does not save the routed experts' gate or up half (mlpwi_0, moe_mlpwi_0," + " moe_mlpwi_1; the minimal* policies save mlpwi_0): the fused GMM never forms them, so the policy would" + " silently stop saving the MoE tensor." + ), + ) fp8_moe: bool = Field( False, description="Whether MoE weights are quantized to float8_e4m3fn in inference/rollout.", @@ -1367,7 +1821,9 @@ def validate_moe_chunks(self) -> "MoEGeneral": @model_validator(mode="after") def validate_moe_sharding_strategy(self) -> "MoEGeneral": """Ensure that only one MoE FSDP sharding strategy is active at a time.""" - if self.moe_pin_sparse_core_all_gathers and self.moe_fsdp_use_two_stage_all_gather: + if ( + self.moe_pin_sparse_core_all_gathers or self.moe_pin_sparse_core_fsdp_all_gathers + ) and self.moe_fsdp_use_two_stage_all_gather: raise ValueError( "SparseCore pinning for MoE all-gathers (`moe_pin_sparse_core_all_gathers=True`) " "is not supported with `moe_fsdp_use_two_stage_all_gather=True`." @@ -1380,6 +1836,206 @@ def validate_moe_sharding_strategy(self) -> "MoEGeneral": ) return self + # The validators of the MoE performance keys below share one form: each lists (met, requirement) pairs and + # _reject_unless raises naming every unmet one. They read the fields of the other config mixins (sparse_matmul, + # mlp_bias, remat_policy, the parallelism degrees, ...) directly: they run on the combined configs (MaxTextConfig + # and RLConfig), which have all of them, so a renamed field fails loudly instead of passing the check. + + def _prefused_sparse_matmul_requirements(self) -> list[tuple[bool, str]]: + """What the MoE layer needs to run the prefused [w0 | w1] sparse-matmul up projection in training.""" + return [ + (self.prefuse_moe_weights, "prefuse_moe_weights=True"), + (self.sparse_matmul, "sparse_matmul=True"), + (not self.te_moe_block, "te_moe_block=False"), + ( + self.attention not in ("vllm_rpa", "vllm_batched_rpa", "vllm_batched_rpa_long_ctx"), + "an attention other than vllm_rpa* (the fused-MoE inference path)", + ), + ] + + def _no_weight_quantization_requirement(self) -> tuple[bool, str]: + return (not self.quantization and not self.use_qwix_quantization, "no weight quantization") + + def _remat_keeps_moe_tensor_requirement(self, tensor: str, names: tuple[str, ...]) -> tuple[bool, str]: + """Requirement: the remat policy saves none of `names`, the checkpoint names of the routed experts' `tensor`. + + The keys that ask for this never form that tensor (it is fused into the op that reads it), so a policy that + saves or offloads it would silently stop saving the MoE tensor; the shared expert's dense MLP, which carries the + generic names (mlpwi_0, mlpwo) too, would still be saved. + """ + if self.remat_policy == "custom": + saved = [name for name in names if getattr(self, name) != RematLocation.REMAT] + else: + saved = [name for name in names if name in _REMAT_PRESET_MLP_NAMES.get(self.remat_policy, ())] + return ( + not saved, + f"a remat policy that does not save the routed experts' {tensor}, which this key never forms" + f" (remat_policy={self.remat_policy} saves it as {', '.join(saved)})", + ) + + def _expert_mlp_dim_shards(self) -> int: + """Devices that split the expert weights' MLP dimension inside the expert shard_map. + + By the default logical axis rules: 'mlp_no_fsdp' -> tensor, tensor_sequence, autoregressive. + """ + degrees = ( + self.ici_tensor_parallelism, + self.dcn_tensor_parallelism, + self.ici_tensor_sequence_parallelism, + self.dcn_tensor_sequence_parallelism, + self.ici_autoregressive_parallelism, + self.dcn_autoregressive_parallelism, + ) + return math.prod(degrees) + + def _ring_of_experts_ragged_sort_requirements(self) -> list[tuple[bool, str]]: + """What the MoE layer needs to run the ring-of-experts ragged sort.""" + return [ + (self.sparse_matmul, "sparse_matmul=True"), + (not self.te_moe_block, "te_moe_block=False"), + (self.use_ring_of_experts, "use_ring_of_experts=True"), + (self.use_ragged_sort, "use_ragged_sort=True"), + ] + + @model_validator(mode="after") + def validate_moe_roe_packed_routing_gather(self) -> "MoEGeneral": + """The packed gather replaces the per-chunk routing gathers of the pipelined local-routing chunk loop.""" + if self.moe_roe_packed_routing_gather: + _reject_unless( + "moe_roe_packed_routing_gather=True", + [ + (self.sparse_matmul, "sparse_matmul=True"), + (not self.te_moe_block, "te_moe_block=False"), + (self.ring_of_experts_local_routing, "ring_of_experts_local_routing=True"), + (self.moe_chunk_pipeline, "moe_chunk_pipeline=True (only the pipelined chunk loop gathers per chunk)"), + # Random routing draws per get_topk call: one draw for the layer is not one per chunk. + (not self.use_random_routing, "use_random_routing=False"), + ], + ) + return self + + @model_validator(mode="after") + def validate_moe_router_grad_from_wo_input(self) -> "MoEGeneral": + """The fused op needs the ring-of-experts ragged sort and a down projection that is the bare GEMM of its rows.""" + if self.moe_router_grad_from_wo_input: + _reject_unless( + "moe_router_grad_from_wo_input=True", + self._ring_of_experts_ragged_sort_requirements() + + [ + ( + self.ici_tensor_parallelism * self.dcn_tensor_parallelism == 1, + "no tensor parallelism (the layer checks the mesh again when it is traced)", + ), + (not self.mlp_bias, "mlp_bias=False"), + self._no_weight_quantization_requirement(), + (not self.moe_tc_ragged_sort, "moe_tc_ragged_sort=False"), + (not self.moe_accumulate_chunk_wgrad, "moe_accumulate_chunk_wgrad=False"), + (not self.moe_bwd_prequant_before_unsort, "moe_bwd_prequant_before_unsort=False"), + (not self.moe_quantize_token_all_gather, "moe_quantize_token_all_gather=False"), + (self.num_moe_emb_chunks <= 0, "num_moe_emb_chunks=0"), + self._remat_keeps_moe_tensor_requirement("down-projection output", _MOE_DOWN_PROJECTION_NAMES), + ], + ) + return self + + @model_validator(mode="after") + def validate_moe_routing_index_math(self) -> "MoEGeneral": + """moe_routing_index_math='count' needs the ring-of-experts ragged sort, the only user of these index maps.""" + if self.moe_routing_index_math != "sort": + _reject_unless( + f"moe_routing_index_math={self.moe_routing_index_math!r}", self._ring_of_experts_ragged_sort_requirements() + ) + return self + + @model_validator(mode="after") + def validate_moe_ragged_sort_unique_pair_argsort(self) -> "MoEGeneral": + """The pair argsort rewrites the ring-of-experts ragged sort's argsorts, which "count" removes.""" + if self.moe_ragged_sort_unique_pair_argsort: + _reject_unless( + "moe_ragged_sort_unique_pair_argsort=True", + self._ring_of_experts_ragged_sort_requirements() + + [ + ( + self.moe_routing_index_math == "sort", + "moe_routing_index_math='sort' (with 'count' there are no argsorts left to rewrite; set one of them)", + ) + ], + ) + return self + + @model_validator(mode="after") + def validate_moe_unsplit_prefused_wi(self) -> "MoEGeneral": + """The whole wi needs the prefused sparse-matmul path and readers of w0/w1 that all take the [w0 | w1] layout.""" + if self.moe_unsplit_prefused_wi: + _reject_unless( + "moe_unsplit_prefused_wi=True", + self._prefused_sparse_matmul_requirements() + + [ + self._no_weight_quantization_requirement(), + # The default expert-weight sharding: these split w0 and w1 differently from the whole [w0 | w1]. + (not self.moe_fsdp_use_two_stage_all_gather, "moe_fsdp_use_two_stage_all_gather=False"), + (not self.moe_flat_fsdp_weights, "moe_flat_fsdp_weights=False"), + (not self.shard_exp_on_fsdp, "shard_exp_on_fsdp=False"), + (not self.use_2d_fsdp_sharding, "use_2d_fsdp_sharding=False"), + (not self.shard_embed_moe_on_fsdp, "shard_embed_moe_on_fsdp=False"), + (not self.moe_accumulate_chunk_wgrad, "moe_accumulate_chunk_wgrad=False"), + (self.num_moe_emb_chunks <= 0, "num_moe_emb_chunks=0"), + # Sharded, a device would hold a slice of [w0 | w1] instead of matching slices of w0 and w1. The layer + # checks the actual partition spec again when it is traced (custom logical axis rules). + ( + self._expert_mlp_dim_shards() == 1, + "an unsharded expert MLP dimension (tensor, tensor_sequence and autoregressive parallelism 1)", + ), + ], + ) + return self + + @model_validator(mode="after") + def validate_moe_prefused_swiglu_vjp(self) -> "MoEGeneral": + """The fused activation needs the whole prefused [gate | up] output, unbiased and unchunked, never saved per half.""" + if self.moe_prefused_swiglu_vjp: + _reject_unless( + "moe_prefused_swiglu_vjp=True", + self._prefused_sparse_matmul_requirements() + + [ + (not self.mlp_bias, "mlp_bias=False (biases are added per half)"), + (self.num_moe_emb_chunks <= 0, "num_moe_emb_chunks=0 (partial sums per half)"), + self._remat_keeps_moe_tensor_requirement("gate and up halves", _MOE_UP_PROJECTION_NAMES), + ( + not self.moe_wi_gmm_fused_swiglu, + "moe_wi_gmm_fused_swiglu=False (the wi GMM already applies the activation)", + ), + ], + ) + return self + + @model_validator(mode="after") + def validate_moe_wi_gmm_fused_swiglu(self) -> "MoEGeneral": + """The fused gmm_v2 SwiGLU epilogue replaces the up projection only on the plain bf16 prefused gmm_v2 path.""" + if self.moe_wi_gmm_fused_swiglu: + _reject_unless( + "moe_wi_gmm_fused_swiglu=True", + self._prefused_sparse_matmul_requirements() + + [ + # gmm_silu is the bf16 tokamax gmm_v2 kernel only. + (self.use_tokamax_gmm and self.use_gmm_v2, "use_tokamax_gmm=True and use_gmm_v2=True"), + self._no_weight_quantization_requirement(), + # The epilogue computes silu(gate) * up; GPT-OSS and clipped activations differ. + (list(self.mlp_activations) == ["silu", "linear"], "mlp_activations=['silu', 'linear']"), + (self.decoder_block != DecoderBlockType.GPT_OSS, "a decoder_block other than gpt_oss"), + (self.mlp_activations_limit <= 0.0, "mlp_activations_limit <= 0 (no clipping)"), + (not self.mlp_bias, "mlp_bias=False"), + (self.num_moe_emb_chunks <= 0, "num_moe_emb_chunks=0"), + (not self.moe_flat_fsdp_weights, "moe_flat_fsdp_weights=False"), + # A 2D bf16 token buffer. + (not self.moe_tc_ragged_3d_gmm, "moe_tc_ragged_3d_gmm=False"), + (not self.moe_quantize_token_all_gather, "moe_quantize_token_all_gather=False"), + (not self.moe_accumulate_chunk_wgrad, "moe_accumulate_chunk_wgrad=False"), + self._remat_keeps_moe_tensor_requirement("gate and up halves", _MOE_UP_PROJECTION_NAMES), + ], + ) + return self + class MoEKernels(BaseModel): """Configuration for MoE-specific kernels like Megablox.""" @@ -1392,6 +2048,16 @@ class MoEKernels(BaseModel): ) wi_tile_fwd_embed_dim: int = Field(1024, description="forward pass tiling dimension for embedding in GMM for wi.") wi_tile_fwd_mlp_dim: int = Field(1024, description="forward pass tiling dimension for MLP in GMM for wi.") + eval_wi_tile_fwd_batch_seq: int = Field( + 512, + description="evaluation forward pass tiling dimension for batch/sequence in GMM for wi.", + ) + eval_wi_tile_fwd_embed_dim: int = Field( + 1024, description="evaluation forward pass tiling dimension for embedding in GMM for wi." + ) + eval_wi_tile_fwd_mlp_dim: int = Field( + 1024, description="evaluation forward pass tiling dimension for MLP in GMM for wi." + ) wi_tile_dlhs_batch_seq: int = Field( 512, description="bwd pass dlhs tiling dimension for batch/sequence in GMM for wi.", @@ -1410,6 +2076,16 @@ class MoEKernels(BaseModel): ) wo_tile_fwd_embed_dim: int = Field(1024, description="forward pass tiling dimension for embedding in GMM for wo.") wo_tile_fwd_mlp_dim: int = Field(1024, description="forward pass tiling dimension for MLP in GMM for wo.") + eval_wo_tile_fwd_batch_seq: int = Field( + 512, + description="evaluation forward pass tiling dimension for batch/sequence in GMM for wo.", + ) + eval_wo_tile_fwd_embed_dim: int = Field( + 1024, description="evaluation forward pass tiling dimension for embedding in GMM for wo." + ) + eval_wo_tile_fwd_mlp_dim: int = Field( + 1024, description="evaluation forward pass tiling dimension for MLP in GMM for wo." + ) wo_tile_dlhs_batch_seq: int = Field( 512, description="bwd pass dlhs tiling dimension for batch/sequence in GMM for wo.", @@ -1455,7 +2131,20 @@ class DeepSeekMoE(BaseModel): routed_score_func: str = Field("", description="Scoring function for routing (e.g., 'softmax', 'sigmoid').") routed_bias: bool = Field(False, description="Whether to add a bias term for routing.") routed_bias_update_rate: float = Field(0.0, description="Update rate applied to the router bias term.") + defer_small_all_reduces: bool = Field( + False, + description="Ring-of-experts sparse_matmul only: drop the per-layer all-reduces of the routed-bias expert counts" + " from the scanned layer. Each MoE layer emits its local partial counts, the scan stacks them, and loss_fn" + " reduces the stacked array once after the layer loop (one all-reduce). Same sums as the per-layer path. (The" + " token-overflow flag of moe_dropless_fallback='step' is already reduced once per step without this flag.)", + ) log_moe_bias_norms: bool = Field(False, description="Whether to log the norms of MoE router biases.") + log_step_diagnostics: bool = Field( + False, + description="Append to every training step's log line: pre-clip and post-clip global grad norm, param norm," + " max |grad|, the routed-bias checksum (sum over MoE layers of sum(bias)), the number of nonzero routed-bias" + " update entries and the MoE overflow flag.", + ) mlp_bias: bool = Field( False, description="Whether to add a learnable bias for MLP matmul, " @@ -1475,6 +2164,31 @@ class DeepSeekMoE(BaseModel): 1, description="Factor by which to split the batch into micro-batches. Only used if use_batch_split_schedule is True.", ) + use_lineage: bool = Field( + False, + description="Whether to use Lineage DeepSeek-V3 execution.", + ) + lineage_quantization: Literal["none", "fp8_full"] = Field( + "none", + description=("Quantization of the Lineage sparse-layer routed experts. Only used" " if use_lineage is True."), + ) + + @classmethod + def _lineage_quantization_none(cls, v: Any) -> Any: + """pyconfig converts the string "none" to None; map it back.""" + return "none" if v is None else v + + # Manually apply the field_validator decorator outside of the class definition to avoid pytype issues + _validate_lineage_quantization = field_validator("lineage_quantization", mode="before")(_lineage_quantization_none) + lineage_expert_permutation: None | Literal["lpt"] = Field( + None, + description=( + "Per-layer expert permutation algorithm for the Lineage sparse and MTP layers (see `dsv3.dsv3`'s" + " `expert_permutation`). None disables expert shuffling. 'lpt' computes each step's permutation on device" + " from the previous step's per-expert token counts with `dsv3_expert_shuffle.lpt_expert_permutation`; the" + " first step uses the identity." + ), + ) class Qwen3Next(BaseModel): @@ -1531,6 +2245,39 @@ class Qwen3Next(BaseModel): " HIGHEST) or with head-sharded context parallelism." ), ) + gdn_conv1d_input_in_qkv_dtype: bool = Field( + False, + description=( + "GDN kernel path: build the Conv1D + SiLU input, the halo-extended [conv_state | qkv] (or qkv left-padded" + " by gdn_conv_kernel_dim - 1 zeros), in qkv's dtype and cast each of the gdn_conv_kernel_dim taps to" + " float32 where it is read, instead of materializing a float32 copy of the whole input that every tap" + " then slices (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py). bfloat16 -> float32 is lossless and every" + " product and sum is formed in the same order, so the forward outputs and the qkv / conv-state / bias" + " gradients are bit-identical on CPU; the conv weight gradient, a float32 reduction over batch x sequence," + " may be summed in a different order once XLA fuses the cast into it (CPU: a few ulp). Only which tensors" + " live in HBM changes; no optimization barrier is involved. Applies to the Conv1D forward and backward that" + " gdn_decoupled_conv1d and its custom-VJP rules trace. Requires use_gdn_kernel=True (the only path that" + " runs that Conv1D)." + ), + ) + gdn_conv1d_dense_token_stack: bool = Field( + False, + description=( + "GDN kernel path: in the Pallas GDN forward kernel, build each chunk's Conv1D output from per-token" + " [conv_dim / 1024, 8, 128] float32 views concatenated along the leading axis, instead of jnp.stack(axis=1)" + " of the per-token [1, 1, conv_dim] rows (kernels/gdn/compute_conv1d.py stack_conv1d_tokens). The value is" + " the same, element for element; only the layout the TPU compiler gives it changes. The stacked value's" + " second-minor dimension is 1, and the compiler of the October 2026 Pathways build lays it out in (8, 128)" + " tiles padded from 1 to 8 sublanes, so the SiLU that follows runs its exp and reciprocal on 8x as many" + " vector registers as the data fills; the dense views keep the value in dense tiles, as older compilers" + " already do for the stack. Applies to every forward kernel call of gdn_decoupled_conv1d and its forward" + " rule (both sequence-sharded context-parallel passes, and the forward rule's replay under remat); the" + " backward kernel is unchanged. Requires use_gdn_kernel=True (the only path that runs that kernel) and a" + " Conv1D width, 2 * gdn_num_key_heads * gdn_key_head_dim + gdn_num_value_heads * gdn_value_head_dim (one" + " rank's heads under head-sharded GDN context parallelism), that is a multiple of 1024: one float32 vector" + " register of 8 x 128." + ), + ) # ---------------------------------------------------------------------------- @@ -1629,6 +2376,8 @@ class Qwen3Next(BaseModel): ["mlp_moe", ["fsdp_transpose", "tensor", "tensor_sequence", "autoregressive"]], ["embed_moe", ["fsdp", "fsdp_transpose", "context", "context_usp_ulysses"]], ["embed_moe", ["fsdp", "context", "context_usp_ulysses"]], + ["embed_router", ["fsdp", "fsdp_transpose", "context", "context_usp_ulysses"]], + ["embed_router", ["fsdp", "context", "context_usp_ulysses"]], # ========================================== # Standard MLP / Dense Layers / Model Structure # ========================================== @@ -1910,6 +2659,30 @@ class RematAndOffload(BaseModel): RematLocation.REMAT, description="Remat policy for the second part of a gated MLP's output.", ) + moe_x_sorted: RematLocation = Field( + RematLocation.REMAT, + description=( + "Remat policy for the routed (post-dispatch, expert-sorted) MoE input plus its small " + "routing/metadata bundle. 'device' saves them across the remat boundary so the backward " + "does not re-run the dispatch token all-gather and ragged sort; the expert GMMs re-run " + "from the saved tensor. Default 'remat' recomputes (existing behavior)." + ), + ) + moe_routing_maps: RematLocation = Field( + RematLocation.REMAT, + description=( + "Remat policy for the MoE routing maps (top-k expert ids, ragged-sort permutations, sorted token indices " + "and group sizes)." + ), + ) + # TODO: Update to moe_tc_routing_weights. + moe_tc_routing_checkpoint: RematLocation = Field( + RematLocation.REMAT, + description=( + "Remat policy for the TC per-buffer-row routing weights (moe_tc_ragged_sort with " + "moe_tc_ragged_weights_on_activation); 'device' saves them so the backward skips the weight gather." + ), + ) mlpwo: RematLocation = Field( RematLocation.REMAT, description="Remat policy for the second MLP layer's output.", @@ -2390,6 +3163,51 @@ class TrainingLoop(BaseModel): data_shuffle_seed: int = Field(0, description="Seed for data shuffling.") init_weights_seed: int = Field(0, description="Seed for model weight initialization.") max_inflight_computations: int = Field(2, description="Maximum number of inflight computations on device.") + block_state_before_run_start: bool = Field( + False, + description=( + "Before init_stop/run_start: jax.block_until_ready(state), evaluate the learning-rate schedule once " + "(compiles its small eager ops) and run a cross-host barrier. Does not touch the dataset." + ), + ) + warm_input_reshard_before_run_start: bool = Field( + False, + description=( + "Before init_stop/run_start: device_put an all-zero synthetic batch through the same host->global->input " + "sharding path as the data loader, so the reshard program is compiled outside the timer. No dataset access." + ), + ) + warm_eval_input_reshard_before_run_start: bool = Field( + False, + description=( + "Before init_stop/run_start: device_put an all-zero synthetic eval batch through the eval loop's " + "host->global->eval-input sharding path (the same jax.device_put the eval loop performs), so the eval " + "reshard program is compiled outside the timer. No dataset access." + ), + ) + eval_cache_prefill_in_background: bool = Field( + False, + description=( + "Right after run_start (inside the timed region): a daemon thread iterates the eval tf.data pipeline once " + "so its in-memory cache() is complete before the first eval; the eval file reads then overlap training " + "steps instead of being paid serially inside the first eval. tf.data (c4_mlperf) eval loaders only." + ), + ) + train_shard_in_read: bool = Field( + False, + description=( + "c4_mlperf train split: shard by file in tfds (host i opens files i, i+n, ...) instead of reading the full " + "stream and keeping every n-th example. Needs at least as many files as data-loading hosts." + ), + ) + train_interleave_cycle_length: int = Field( + -1, + description=( + "c4_mlperf train split: tfds interleave_cycle_length, i.e. number of files read concurrently per host " + "(-1 = tfds default of 16). With train_shard_in_read, keep this below the files per host so that " + "data_shuffle_seed changes which files are read first." + ), + ) class ManifoldConstrainedHyperConnections(BaseModel): @@ -2454,6 +3272,21 @@ def validate_streaming_diloco_params(self) -> "DilocoParams": f"num_diloco_fragments ({self.num_diloco_fragments}) must be at least 2 when enable_streaming_diloco " "is True (1 for non-scanned parameters, at least 1 for scanned layers)." ) + if self.diloco_sync_period % self.num_diloco_fragments != 0: + # Mirrors fragmenter.get_streaming_schedule, which rounds the per-fragment interval. + steps_between_syncs = max(1, int(round(self.diloco_sync_period / self.num_diloco_fragments))) + max_logging.warning( + f"diloco_sync_period ({self.diloco_sync_period}) is not a multiple of num_diloco_fragments" + f" ({self.num_diloco_fragments}): streaming DiLoCo syncs a fragment every {steps_between_syncs} steps, so" + f" each fragment is synced every {steps_between_syncs * self.num_diloco_fragments} steps." + ) + if self.diloco_bucketize_non_scanned and not self.enable_streaming_diloco: + raise ValueError("diloco_bucketize_non_scanned requires enable_streaming_diloco=True.") + if self.diloco_bucketize_non_scanned and self.num_diloco_fragments < 3: + raise ValueError( + f"diloco_bucketize_non_scanned requires num_diloco_fragments >= 3, got {self.num_diloco_fragments}: with a" + " single layer fragment, bucketized matrices would move whole into it." + ) return self diloco_outer_lr: float = Field(0.3, description="learning rate for outer optimizer.") @@ -2485,6 +3318,17 @@ def validate_streaming_diloco_params(self) -> "DilocoParams": ), ) use_sequential_layers: bool = Field(False, description="Whether to sync layers sequentially (or interleaved).") + diloco_bucketize_non_scanned: bool = Field( + False, + description=( + "Split non-scanned matrices (ndim >= 2, e.g. embedding table, output projection) across the layer" + " fragments instead of syncing them whole in fragment 0, so every streaming DiLoCo step moves a similar" + " number of bytes. Each matrix is split along its longest unsharded axis with at least" + " num_diloco_fragments - 1 indices (vocab under FSDP), whatever its size; a matrix whose long axes are all" + " sharded (e.g. vocab under tensor parallelism) stays whole in fragment 0. Requires" + " num_diloco_fragments >= 3." + ), + ) num_communication_overlapping_steps: NonNegativeInt = Field( 0, description="Steps of communication overlap with computation. \\tau from the paper." ) @@ -2685,11 +3529,29 @@ class YarnRope(BaseModel): "and 2 represents the real/imag coordinates, returning interleaved RoPE." ), ) + rope_direct_position_freqs: bool = Field( + False, + description=( + "Compute YaRN RoPE angles directly from position * inv_freqs " + "instead of indexing a [max_position_embeddings, half_dim] lookup table. " + "Outputs differ for positions >= max_position_embeddings: the table lookup clamps them " + "to the last row, while the direct computation uses the true angle." + ), + ) + rope_dot_pairwise: bool = Field( + False, + description=( + "Keep rank-4 [batch, seq, heads, dim] tensor intact during pairwise RoPE " + "via a block-skew matmul on the minor dimension instead of reshaping to rank-5." + ), + ) @model_validator(mode="after") def validate_rope_pairwise(self) -> "YarnRope": if self.rope_pairwise and not self.rope_interleave: raise ValueError("rope_pairwise=True requires rope_interleave=True.") + if self.rope_dot_pairwise and not self.rope_pairwise: + raise ValueError("rope_dot_pairwise=True requires rope_pairwise=True.") return self @@ -2898,6 +3760,7 @@ class Profiling(BaseModel): profile_cleanly: bool = Field(True, description="Add block_until_ready to align profile for each step.") profile_periodically_period: int = Field(-1, description="If positive, profile every N steps.") hide_profiler_step_metric: bool = Field(False, description="Whether to enable profiler step metric.") + enable_continuous_profiling: bool = Field(False, description="If true, it will support saving profile > 2GB.") enable_jax_profiler: bool = Field(False, description="Enable the JAX live profiler.") jax_profiler_port: int = Field(9999, description="Port for the JAX profiler.") enable_tpu_profiling_options: bool = Field(False, description="Enable TPU advanced profiling options.") @@ -2979,6 +3842,9 @@ class Goodput(BaseModel): """Configuration for goodput monitoring.""" enable_goodput_recording: bool = Field(False, description="Enable goodput recording.") + goodput_job_name: str = Field( + "", description="Optional job name override for goodput recording and monitoring. Defaults to run_name when empty." + ) monitor_goodput: bool = Field(False, description="Monitor goodput.") goodput_upload_interval_seconds: int = Field(30, description="Interval to upload goodput metrics.") enable_pathways_goodput: bool = Field(False, description="Enable goodput monitoring for Pathways.") @@ -2986,6 +3852,12 @@ class Goodput(BaseModel): step_deviation_interval_seconds: int = Field(30, description="Interval to check step time deviation.") enable_gcp_goodput_metrics: bool = Field(True, description="Enable GCP goodput metrics.") enable_gcp_step_deviation_metrics: bool = Field(True, description="Enable GCP step deviation metrics.") + rolling_windows_seconds: list[PositiveInt] = Field( + [3600, 86400, 259200, 432000], + min_length=1, + description="Rolling window intervals in seconds. Must contain at least one window.", + ) + enable_rolling_window_goodput: bool = Field(True, description="Enable rolling window goodput.") class ElasticTraining(BaseModel): @@ -3032,7 +3904,14 @@ class Tensorboard(BaseModel): """Configuration for Tensorboard logging.""" enable_tensorboard: bool = Field(True, description="Enable Tensorboard logging.") - use_vertex_tensorboard: bool = Field(False, description="Set to True for GCE, False if running via XPK.") + use_vertex_tensorboard: bool = Field( + False, + description=( + "Set to True to have MaxText create (or reuse) the Vertex AI Tensorboard instance and Experiment and upload" + " logs. Leave False if they are created outside MaxText and UPLOAD_DATA_TO_TENSORBOARD is set in the" + " environment." + ), + ) vertex_tensorboard_project: Optional[str] = Field("", description="GCP project for Vertex AI Tensorboard.") vertex_tensorboard_region: Optional[str] = Field("", description="Region for Vertex AI Tensorboard.") @@ -3209,6 +4088,19 @@ class VLLM(BaseModel): kv_cache_buffer: int = Field(256, gt=0, description="Buffer for KV cache.") hbm_utilization_vllm: float = Field(0.72, gt=0.0, le=1.0, description="Target HBM utilization for vLLM.") swap_space_vllm_gb: int = Field(2, ge=0, description="Swap space in GB for vLLM.") + free_kv_cache_during_weight_sync: bool = Field( + True, + description=( + "Maps to tunix RolloutConfig.rollout_vllm_free_kv_cache_during_weight_sync -> " + "VllmConfig.free_kv_cache_during_weight_sync (ignored, with a log line, while the installed tunix " + "predates it). The weight sync materializes a second copy of the sampler weights on HBM; freeing the KV " + "cache first makes room for it. False keeps the pool (hbm_utilization_vllm of each chip) allocated across " + "the sync and only resets its prefix-cache entries, saving the free + re-allocation (about 2 s per step " + "on a colocated Qwen3-0.6B run at hbm_utilization_vllm=0.3). With False the sync runs out of HBM when " + "that second copy does not fit next to the pool: check the headroom before disabling, colocated setups " + "and a large hbm_utilization_vllm have the least." + ), + ) enable_dp_attention: bool = Field(False, description="Enable the attn_dp mesh axis in vLLM.") enable_expert_parallel: bool = Field(False, description="Enable expert parallelism in vLLM.") async_scheduling: bool = Field(False, description="Enable asynchronous scheduling in vLLM.") @@ -3732,6 +4624,13 @@ def get_individual_scales(scale: int) -> tuple[int, int, int, int]: return emb_scale, num_head_scale, mlp_dim_scale, layer_scale +def _resolve_parallelism(configured: list[int] | None, axis_map: dict[str, int], mesh_axes: list[str]) -> list[int]: + """Per-axis parallelism list for `mesh_axes` (preserves explicit list for custom physical mesh axes).""" + if configured and len(configured) == len(mesh_axes) and any(axis not in axis_map for axis in mesh_axes): + return list(configured) + return [axis_map.get(axis, 1) for axis in mesh_axes] + + def _resolved_fsdp_size(mesh_axes: list[str], parallelism: list[int], num_devices: int) -> int: """Combined size of the fsdp mesh axes, resolving a `-1` the way mesh creation will. @@ -3880,19 +4779,17 @@ def _validate_check_vma_is_supported(self): def validate_retry_when_tokens_dropped(self): """Validates prerequisites for the step-level dropless fallback retry.""" - if self.retry_dropless_first_steps < 0: - raise ValueError(f"retry_dropless_first_steps must be >= 0 (got {self.retry_dropless_first_steps}).") + if self.log_moe_drops and self.moe_log_max_load_ratio: + raise ValueError("log_moe_drops and moe_log_max_load_ratio both sow moe_max_load_ratio; enable only one.") if self.first_phase_ragged_buffer_factor < 0: raise ValueError(f"first_phase_ragged_buffer_factor must be >= 0 (got {self.first_phase_ragged_buffer_factor}).") - if self.retry_dropless_first_steps > 0 and not self.retry_when_tokens_dropped: - raise ValueError("retry_dropless_first_steps > 0 requires retry_when_tokens_dropped=True.") - if self.first_phase_ragged_buffer_factor > 0 and self.retry_dropless_first_steps == 0: - raise ValueError("first_phase_ragged_buffer_factor > 0 requires retry_dropless_first_steps > 0.") if self.retry_dropless_final_micro_batch and not self.retry_when_tokens_dropped: raise ValueError("retry_dropless_final_micro_batch=True requires retry_when_tokens_dropped=True.") if self.retry_when_tokens_dropped: + if self.moe_dropless_fallback is not None: + raise ValueError("retry_when_tokens_dropped=True cannot be combined with moe_dropless_fallback (both handle drops).") if self.num_experts <= 1: raise ValueError("retry_when_tokens_dropped=True requires num_experts > 1.") if self.ragged_buffer_factor == -1: @@ -3919,14 +4816,134 @@ def validate_retry_when_tokens_dropped(self): if self.te_moe_block: raise ValueError("first_phase_ragged_buffer_factor > 0 is not supported with te_moe_block=True.") + def validate_moe_dropless_fallback(self): + """Validates prerequisites for the step-level ('step') and in-layer ('layer') dropless fallback.""" + mode = self.moe_dropless_fallback + if mode is None: + return + prefix = f"moe_dropless_fallback='{mode}'" + if self.num_experts <= 1: + raise ValueError(f"{prefix} requires num_experts > 1.") + if self.ragged_buffer_factor <= 0: + raise ValueError(f"{prefix} requires ragged_buffer_factor > 0.0.") + if not self.use_ring_of_experts: + raise ValueError(f"{prefix} is currently only supported with use_ring_of_experts=True.") + if not self.use_ragged_sort: + raise ValueError(f"{prefix} requires use_ragged_sort=True.") + if self.num_moe_emb_chunks > 0: + raise ValueError(f"{prefix} does not support num_moe_emb_chunks > 0.") + if mode == "layer" and self.use_random_routing: + # The overflow check re-derives top-k outside permute(); random routing would not reproduce it. + raise ValueError(f"{prefix} does not support use_random_routing=True.") + if mode == "step": + if self.enable_diloco: + raise ValueError(f"{prefix} is not supported with enable_diloco=True.") + if self.compiled_trainstep_file: + raise ValueError(f"{prefix} is not supported with compiled_trainstep_file.") + if self.optimizer_memory_host_offload or self.parameter_memory_host_offload: + raise ValueError( + f"{prefix} is not supported with optimizer_memory_host_offload or parameter_memory_host_offload." + ) + + def validate_retry_dropless_first_steps_and_first_phase_buffer(self): + """Validates retry_dropless_first_steps, first_phase_ragged_buffer_factor and the REQUIRED_RBF probe.""" + if self.retry_dropless_first_steps < 0: + raise ValueError(f"retry_dropless_first_steps must be >= 0 (got {self.retry_dropless_first_steps}).") + if self.retry_dropless_first_steps > 0 and self.moe_dropless_fallback != "step" and not self.retry_when_tokens_dropped: + raise ValueError( + "retry_dropless_first_steps > 0 requires moe_dropless_fallback='step' or retry_when_tokens_dropped=True." + ) + if self.first_phase_ragged_buffer_factor < 0: + raise ValueError( + f"first_phase_ragged_buffer_factor must be 0 (off) or > 0 (got {self.first_phase_ragged_buffer_factor})." + ) + if self.first_phase_ragged_buffer_factor > 0: + if self.retry_dropless_first_steps <= 0: + raise ValueError("first_phase_ragged_buffer_factor > 0 requires retry_dropless_first_steps > 0.") + if self.moe_dropless_fallback != "step" and not self.retry_when_tokens_dropped: + raise ValueError( + "first_phase_ragged_buffer_factor > 0 requires moe_dropless_fallback='step' or retry_when_tokens_dropped=True." + ) + if self.ragged_buffer_factor <= 0: + raise ValueError("first_phase_ragged_buffer_factor > 0 requires ragged_buffer_factor > 0.") + if self.first_phase_ragged_buffer_factor < self.ragged_buffer_factor: + raise ValueError( + f"first_phase_ragged_buffer_factor ({self.first_phase_ragged_buffer_factor}) must be >= " + f"ragged_buffer_factor ({self.ragged_buffer_factor})." + ) + if self.te_moe_block: + raise ValueError("first_phase_ragged_buffer_factor > 0 is not supported with te_moe_block=True.") + if self.log_required_ragged_buffer_factor: + if self.te_moe_block: + raise ValueError("log_required_ragged_buffer_factor=True is not supported with te_moe_block=True.") + if not (self.use_ring_of_experts and self.use_ragged_sort): + raise ValueError( + "log_required_ragged_buffer_factor=True requires use_ring_of_experts=True and use_ragged_sort=True." + ) + if self.num_moe_emb_chunks > 0: + raise ValueError("log_required_ragged_buffer_factor=True does not support num_moe_emb_chunks > 0.") + + def validate_moe_log_max_load_ratio(self): + """Validates that moe_log_max_load_ratio is used with the ring-of-experts ragged path.""" + if self.moe_log_max_load_ratio and not ( + self.sparse_matmul and not self.te_moe_block and self.use_ring_of_experts and self.use_ragged_sort + ): + raise ValueError( + "moe_log_max_load_ratio=True requires sparse_matmul=True, te_moe_block=False," + " use_ring_of_experts=True and use_ragged_sort=True." + ) + + def validate_moe_tc_ragged_sort(self): + """Validates that moe_tc_ragged_sort is used with the truncated-buffer ring-of-experts ragged sort.""" + if self.moe_tc_ragged_3d_dispatch and not (self.moe_tc_ragged_sort and self.moe_tc_ragged_3d_gmm): + # The remaining requirements of the 3D dispatch are those of moe_tc_ragged_sort / moe_tc_ragged_3d_gmm below. + raise ValueError("moe_tc_ragged_3d_dispatch=True requires moe_tc_ragged_sort=True and moe_tc_ragged_3d_gmm=True.") + # The per-buffer-row routing weights only exist (and are only tagged) on the TC path with weights on activation; + # otherwise a non-remat moe_tc_routing_checkpoint would be silently ignored. + if self.moe_tc_routing_checkpoint != RematLocation.REMAT and not ( + self.moe_tc_ragged_sort and self.moe_tc_ragged_weights_on_activation + ): + raise ValueError( + f"moe_tc_routing_checkpoint={RematLocation(self.moe_tc_routing_checkpoint).value} requires" + " moe_tc_ragged_sort=True and moe_tc_ragged_weights_on_activation=True; use moe_tc_routing_checkpoint=remat." + ) + if self.moe_tc_window_routing and not self.moe_tc_ragged_sort: + raise ValueError("moe_tc_window_routing=True requires moe_tc_ragged_sort=True.") + if not self.moe_tc_ragged_sort: + return + if not self.use_ring_of_experts: + raise ValueError("moe_tc_ragged_sort=True requires use_ring_of_experts=True.") + if not self.use_ragged_sort: + raise ValueError("moe_tc_ragged_sort=True requires use_ragged_sort=True.") + if self.ragged_buffer_factor <= 0: + raise ValueError("moe_tc_ragged_sort=True requires ragged_buffer_factor > 0.0.") + if self.moe_tc_ragged_3d_gmm: + if not (self.use_tokamax_gmm and self.use_gmm_v2): + raise ValueError("moe_tc_ragged_3d_gmm=True requires use_tokamax_gmm=True and use_gmm_v2=True.") + unsupported = ["mlp_bias"] if self.mlp_bias else [] + if self.ici_tensor_parallelism > 1 or self.dcn_tensor_parallelism > 1: + unsupported.append("tensor parallelism") + if self.num_moe_emb_chunks > 0: + unsupported.append("num_moe_emb_chunks > 0") + if unsupported: + raise ValueError(f"moe_tc_ragged_3d_gmm=True does not support: {', '.join(unsupported)}.") + + def validate_moe_topk_before_ep_all_gather(self): + """Validates that moe_topk_before_ep_all_gather is used with ring of experts.""" + if self.moe_topk_before_ep_all_gather and not (self.sparse_matmul and self.use_ring_of_experts): + raise ValueError("moe_topk_before_ep_all_gather=True requires sparse_matmul=True and use_ring_of_experts=True.") + def validate_ragged_buffer_factor(self): - """Validates that ragged_buffer_factor is used with supported settings.""" + """Validates that ragged_buffer_factor and eval_ragged_buffer_factor are used with supported settings.""" if self.te_moe_block: - if 0 < self.ragged_buffer_factor < 1.0: - raise ValueError("te_moe_block=True requires ragged_buffer_factor >= 1.0, or <= 0 for worst-case capacity.") + if 0 < self.ragged_buffer_factor < 1.0 or self.eval_ragged_buffer_factor > 0: + raise ValueError( + "te_moe_block=True requires ragged_buffer_factor >= 1.0, or <= 0 for worst-case capacity, " + "and does not support eval_ragged_buffer_factor > 0." + ) return - if self.ragged_buffer_factor <= 0: + if self.ragged_buffer_factor <= 0 and self.eval_ragged_buffer_factor <= 0: return # Not using a ragged buffer factor if self.use_ring_of_experts and not self.use_ragged_sort: @@ -4113,6 +5130,39 @@ def validate_mllog(self): "at least one eval_accuracy event)." ) + def validate_splash_diag_skip(self): + """Validates Tokamax Splash diagonal sub-tile skipping and warns when it cannot take effect.""" + if self.sa_qk_diag_grid < 2 or (self.sa_qk_diag_grid & (self.sa_qk_diag_grid - 1)) != 0: + raise ValueError(f"sa_qk_diag_grid must be a power of 2 and >= 2, got {self.sa_qk_diag_grid}.") + if not (self.use_tokamax_splash and (self.sa_qk_diag_skip or self.sa_sv_diag_skip)): + return + # Tokamax checks square training blocks when the kernel is built at startup, but `eval_sa_*` blocks are + # only used under custom eval rules and first traced at the first eval step, so check them here. + uses_eval_sa_blocks = bool(self.logical_axis_rules_for_eval) and ( + self.logical_axis_rules_for_eval != self.logical_axis_rules + ) + if uses_eval_sa_blocks and not self.eval_sa_block_q == self.eval_sa_block_kv == self.eval_sa_block_kv_compute: + raise ValueError( + "sa_qk_diag_skip/sa_sv_diag_skip with custom eval logical axis rules require square eval blocks " + "(eval_sa_block_q == eval_sa_block_kv == eval_sa_block_kv_compute); got " + f"{self.eval_sa_block_q}/{self.eval_sa_block_kv}/{self.eval_sa_block_kv_compute}." + ) + # The skip only runs when the Splash kernel gets a pure `CausalMask`; otherwise it is silently disabled. + cp_size = self.ici_context_parallelism * self.dcn_context_parallelism + strategy = self.context_parallel_strategy.lower() + if self.attention_type not in (AttentionType.GLOBAL.value, AttentionType.MLA.value): + reason = f"attention_type={self.attention_type!r} (only 'global' and 'mla' are supported)" + elif cp_size > 1 and strategy in ("ring", "usp"): + reason = f"context_parallel_strategy={strategy!r} (ring/USP attention does not use the flags)" + elif cp_size > 1 and strategy == "all_gather" and self.context_parallel_load_balance: + reason = "load-balanced all_gather context parallelism (it uses LoadBalancedCausalMask)" + else: + return + max_logging.warning( + f"sa_qk_diag_skip={self.sa_qk_diag_skip}, sa_sv_diag_skip={self.sa_sv_diag_skip} have no effect with " + f"{reason}. Diagonal sub-tile skipping requires a pure causal mask." + ) + def validate_num_moe_emb_chunks(self): """ Validates that num_moe_emb_chunks is used with supported settings. @@ -4123,6 +5173,13 @@ def validate_num_moe_emb_chunks(self): f"num_moe_emb_chunks > 0 requires use_gmm_v2=True and use_ring_of_experts=True. " f"Got use_gmm_v2={self.use_gmm_v2}, use_ring_of_experts={self.use_ring_of_experts}." ) + # The emb-chunking path (moe_emb_chunking in moe.py) routes per chunk and does not tag its routed input with + # checkpoint_name("moe_x_sorted"), so a non-remat moe_x_sorted would be silently ignored there. + if self.moe_x_sorted != RematLocation.REMAT: + raise ValueError( + f"moe_x_sorted={RematLocation(self.moe_x_sorted).value} is not supported with num_moe_emb_chunks > 0; " + "use moe_x_sorted=remat." + ) def validate_moe_expert_weight_prefetch(self): """Validates that moe_expert_weight_prefetch is set only on the path that places the barrier.""" @@ -4157,6 +5214,20 @@ def validate_moe_quantize_token_all_gather(self): " (static scaling mode)." ) + def validate_moe_quantize_combine_bwd_method(self): + """Validates that moe_quantize_combine_bwd_method is only used with ring of experts.""" + if self.moe_quantize_combine_bwd_method: + if not self.use_ring_of_experts: + raise ValueError( + f"moe_quantize_combine_bwd_method='{self.moe_quantize_combine_bwd_method}' requires use_ring_of_experts=True." + ) + method = self.moe_quantize_combine_bwd_method + if not (method == "rowwise" or method.startswith("fixed")): + raise ValueError( + f"Unsupported moe_quantize_combine_bwd_method: '{method}'. " + "Supported options: '', 'rowwise', 'fixed,'." + ) + @staticmethod def _load_mesh_config_from_yaml(rule_value: str) -> dict: """Helper to load and parse custom mesh YAML configurations.""" @@ -4248,7 +5319,13 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None: """ if is_rl_config: MaxTextConfig.validate_gdn_remat_requires_kernel(cfg) + dense_token_stack = getattr(cfg, "gdn_conv1d_dense_token_stack", False) if cfg.decoder_block not in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5): + if dense_token_stack: + raise ValueError( + "gdn_conv1d_dense_token_stack=True requires a decoder with Gated DeltaNet layers (qwen3_next or" + f" qwen3_5), got decoder_block={cfg.decoder_block.value}." + ) return cp_size = math.prod( max(1, getattr(cfg, f"{mesh}_{axis}_parallelism")) @@ -4264,6 +5341,16 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None: raise ValueError("GatedDeltaNet context parallelism requires context_parallel_load_balance=False.") if not cfg.use_gdn_kernel: + if cfg.gdn_conv1d_input_in_qkv_dtype: + raise ValueError( + "gdn_conv1d_input_in_qkv_dtype=True requires use_gdn_kernel=True: only the GDN kernel path runs the" + " Conv1D + SiLU (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py) whose input it keeps in qkv's dtype." + ) + if dense_token_stack: + raise ValueError( + "gdn_conv1d_dense_token_stack=True requires use_gdn_kernel=True: only the Pallas GDN forward kernel" + " runs the in-VMEM Conv1D (kernels/gdn/compute_conv1d.py) whose output it assembles." + ) return if cfg.ici_tensor_parallelism > 1 or cfg.dcn_tensor_parallelism > 1: raise ValueError( @@ -4271,8 +5358,8 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None: f"'tensor' (got ici_tensor_parallelism={cfg.ici_tensor_parallelism}, " f"dcn_tensor_parallelism={cfg.dcn_tensor_parallelism})." ) - if cfg.gdn_chunk_size != 64: - raise ValueError(f"use_gdn_kernel=True requires gdn_chunk_size=64, got {cfg.gdn_chunk_size}.") + if cfg.gdn_chunk_size not in (64, 128): + raise ValueError(f"use_gdn_kernel=True requires gdn_chunk_size in (64, 128), got {cfg.gdn_chunk_size}.") if cfg.gdn_key_head_dim % 128 or cfg.gdn_value_head_dim % 128: raise ValueError( "use_gdn_kernel=True requires gdn_key_head_dim and gdn_value_head_dim to be multiples of 128, " @@ -4280,6 +5367,25 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None: ) if not cfg.use_qk_norm_in_gdn: raise ValueError("use_gdn_kernel=True requires use_qk_norm_in_gdn=True.") + # With packing=False the kernel skips its segment-aware path in training, so reject the inputs that still + # put several documents in one row. + if not cfg.packing: + if is_rl_config and cfg.max_seq_token_per_tpu > 0: + raise ValueError( + "use_gdn_kernel=True with max_seq_token_per_tpu > 0 requires packing=True: packed rows hold several " + "sequences, and with packing=False the kernel ignores their segment IDs in training." + ) + if ( + not is_rl_config + and cfg.dataset_type == DatasetType.GRAIN + and cfg.grain_file_type in ("mmap", "mmap_npy") + and cfg.reset_attention_mask + ): + raise ValueError( + f"use_gdn_kernel=True with grain_file_type={cfg.grain_file_type!r} and reset_attention_mask=True " + "requires packing=True: those rows hold several documents split at EOD, and with packing=False the " + "kernel ignores their segment IDs in training." + ) seq_lens = {"max_target_length": cfg.max_target_length} if is_rl_config and cfg.max_seq_token_per_tpu > 0: seq_lens["max_seq_token_per_tpu"] = cfg.max_seq_token_per_tpu @@ -4290,6 +5396,40 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None: f"degree ({cp_size}) times gdn_chunk_size ({cfg.gdn_chunk_size}), so that every per-rank sequence " "is a whole number of kernel chunks." ) + if dense_token_stack: + MaxTextConfig.validate_gdn_dense_token_stack_width(cfg, cp_size) + + @staticmethod + def validate_gdn_dense_token_stack_width(cfg: Any, cp_size: int) -> None: + """Fails if gdn_conv1d_dense_token_stack cannot regroup the Conv1D rows the GDN forward kernel convolves. + + The option regroups each token's float32 Conv1D row in whole vector registers of 8 x 128 + (kernels/gdn/compute_conv1d.py DENSE_TOKEN_STACK_ALIGNMENT), so the width each forward-kernel call convolves must + be a multiple of 1024: the whole Conv1D, or one rank's heads under head-sharded GDN context parallelism. + + `cp_size` counts a context parallelism left at -1 (filled in when the mesh is built) as 1, as validate_gdn_config's + sequence-length check does. If the filled-in degree head-shards the heads to a width that is not a multiple of + 1024, wrapper.fused_conv1d_gdn's own width check raises when the layer is traced. + """ + num_k_heads, num_v_heads = cfg.gdn_num_key_heads, cfg.gdn_num_value_heads + # Head-sharded as model_runner._resolve_gdn_cp_mode decides it for training: gdn_cp_mode 'head', or 'auto' at a + # degree of 2 that divides the key heads. (Head counts the degree does not divide fail there on their own.) + cp_mode = getattr(cfg, "gdn_cp_mode", "auto") + head_sharded = cp_size > 1 and ( + cp_mode == "head" or (cp_mode == "auto" and cp_size <= min(2, num_k_heads) and num_k_heads % cp_size == 0) + ) + ranks = cp_size if head_sharded and num_k_heads % cp_size == 0 and num_v_heads % cp_size == 0 else 1 + width = 2 * (num_k_heads // ranks) * cfg.gdn_key_head_dim + (num_v_heads // ranks) * cfg.gdn_value_head_dim + if width % 1024: + if ranks > 1: + where = f"one rank's heads under head-sharded GDN context parallelism of degree {cp_size}" + else: + where = "the whole Conv1D" + raise ValueError( + f"gdn_conv1d_dense_token_stack=True requires the width each GDN forward kernel convolves ({where}: 2 * key" + f" heads * gdn_key_head_dim + value heads * gdn_value_head_dim = {width}) to be a multiple of 1024: it" + " regroups each token's float32 row in whole vector registers of 8 x 128." + ) @model_validator(mode="after") def set_derived_and_validate_values(self) -> "MaxTextConfig": @@ -4319,12 +5459,13 @@ def set_derived_and_validate_values(self) -> "MaxTextConfig": _ep_disabled_flags = { "use_random_routing": False, "use_ragged_sort": False, - "retry_when_tokens_dropped": False, + "moe_dropless_fallback": None, "use_ring_of_experts": False, "num_moe_emb_chunks": 0, } if not self.te_moe_block: _ep_disabled_flags["ragged_buffer_factor"] = -1.0 + _ep_disabled_flags["eval_ragged_buffer_factor"] = -1.0 for flag_name, disabled_value in _ep_disabled_flags.items(): current = getattr(self, flag_name) if current != disabled_value: @@ -4698,6 +5839,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "gdn_cp_state", "mlpwi", "moe_mlpwi_0", + "moe_x_sorted", + "moe_routing_maps", + "moe_tc_routing_checkpoint", "moe_mlpwi_1", "moe_mlpwo", "moe_route", @@ -4815,13 +5959,19 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de ): self.logical_axis_rules.append(["aqt_amax_history", ("stage",)]) - # H. RESOLVE local_sa_* FLAGS: inherit from global sa_* if not explicitly set. + # H. RESOLVE local_sa_* and eval_local_sa_* FLAGS: inherit from global sa_* / eval_sa_* if not explicitly set. if self.local_sa_block_q is None: self.local_sa_block_q = self.sa_block_q if self.local_sa_block_kv is None: self.local_sa_block_kv = self.sa_block_kv if self.local_sa_block_kv_compute is None: self.local_sa_block_kv_compute = self.sa_block_kv_compute + if self.eval_local_sa_block_q is None: + self.eval_local_sa_block_q = self.eval_sa_block_q + if self.eval_local_sa_block_kv is None: + self.eval_local_sa_block_kv = self.eval_sa_block_kv + if self.eval_local_sa_block_kv_compute is None: + self.eval_local_sa_block_kv_compute = self.eval_sa_block_kv_compute if self.local_sa_block_q_dkv is None: self.local_sa_block_q_dkv = self.sa_block_q_dkv if self.local_sa_block_kv_dkv is None: @@ -4840,6 +5990,12 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de self.local_sa_k_layout = self.sa_k_layout if self.local_sa_v_layout is None: self.local_sa_v_layout = self.sa_v_layout + if self.eval_local_sa_q_layout is None: + self.eval_local_sa_q_layout = self.eval_sa_q_layout + if self.eval_local_sa_k_layout is None: + self.eval_local_sa_k_layout = self.eval_sa_k_layout + if self.eval_local_sa_v_layout is None: + self.eval_local_sa_v_layout = self.eval_sa_v_layout if self.local_use_splash_scheduler is None: self.local_use_splash_scheduler = self.use_splash_scheduler if self.local_sa_fuse_reciprocal is None: @@ -5014,6 +6170,29 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de f"Setting `indexer_cutoff_threshold='{self.indexer_cutoff_threshold}'` is only valid when " "`use_indexer=True` (DeepSeek Sparse Attention / MLA Indexer)." ) + if self.use_mla_absorbed_mqa: + if self.attention_type != AttentionType.MLA.value: + raise ValueError( + f"`use_mla_absorbed_mqa=True` requires `attention_type='{AttentionType.MLA.value}'`, " + f"but found attention_type='{self.attention_type}'." + ) + if self.attention not in ("dot_product", "autoselected"): + raise ValueError( + f"`use_mla_absorbed_mqa=True` requires `attention` to be 'dot_product' or 'autoselected', " + f"but found attention='{self.attention}'." + ) + if self.mla_naive_kvcache: + raise ValueError("`use_mla_absorbed_mqa=True` is incompatible with `mla_naive_kvcache=True`.") + has_quant = ( + self.quantize_kvcache + or self.use_qwix_quantization + or self.use_manual_quantization + or self.experimental_sa_quant_q_fp8 + or self.experimental_sa_quant_k_fp8 + or (self.quantization and self.quantization not in (QuantizationType.NONE, QuantizationType.TE_NO_QUANT)) + ) + if has_quant: + raise ValueError("`use_mla_absorbed_mqa=True` is incompatible with quantization or FP8 options.") if self.attention_type == AttentionType.CHUNK.value and ( not isinstance(self.chunk_attn_window_size, int) or self.chunk_attn_window_size <= 0 ): @@ -5054,8 +6233,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError("`training_objective='block_diffusion'` currently requires `dataset_type='hf'`.") if self.use_dpo: raise ValueError("`training_objective='block_diffusion'` is not compatible with DPO.") - if self.use_sft: - raise ValueError("`training_objective='block_diffusion'` currently supports pre-training only.") if self.use_multimodal or self.use_audio: raise ValueError("`training_objective='block_diffusion'` currently supports text-only training.") valid_model_contracts = { @@ -5072,11 +6249,36 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError("`block_diffusion_canvas_policy='seed_and_mask'` requires `causal_block_size >= 2`.") if self.quantize_kvcache and not self.kv_quant_axis: raise ValueError("`kv_quant_axis` cannot be empty when quantize_kvcache is True.") + if self.mtp_reuse_input_embedding and self.use_multimodal: + raise ValueError("`mtp_reuse_input_embedding` does not support multimodal inputs.") if self.quantize_mtp: if self.mtp_num_layers <= 0: raise ValueError("`quantize_mtp` can only be enabled when `mtp_num_layers > 0`.") if self.quantization != "fp8_full": raise ValueError("`quantize_mtp` can only be enabled when `quantization='fp8_full'`.") + if self.quantize_logits_proj: + if self.logits_via_embedding: + raise ValueError( + "`quantize_logits_proj` cannot be enabled when input and output embeddings are tied" + " (`logits_via_embedding=True`)." + ) + if self.quantization != "fp8_full": + raise ValueError("`quantize_logits_proj` can only be enabled when `quantization='fp8_full'`.") + if self.num_vocab_tiling > 1: + raise ValueError( + "`quantize_logits_proj` is not supported with `num_vocab_tiling > 1`: under vocab tiling the" + " decoder skips `apply_output_head` in train mode, so `logits_dense` is absent from the forward" + " pass that `qwix.quantize_model` traces, and the projection runs later inside" + " `vocab_tiling_nnx_loss` on a merged model copy. Whether interception survives that path is" + " unverified; this combination is rejected until it is tested." + ) + if self.logits_dot_in_fp32: + raise ValueError( + "`logits_dot_in_fp32=True` is rejected with `quantize_logits_proj=True`: the fp32 cast on" + " the logits operands is undone by requantization at the projection matmul, so the projection" + " remains quantized while believing you configured fp32. Set `quantize_logits_proj=False` to keep" + " the projection in fp32." + ) if self.quantization and self.use_qwix_quantization and self.quantize_router_proj and self.float32_gate_logits: raise ValueError( "`float32_gate_logits=True` is rejected with `quantize_router_proj=True`: the fp32 cast on" @@ -5149,11 +6351,20 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de if self.model_name.startswith("deepseek4") and self.first_num_hash_layers > 0 and self.use_ring_of_experts: raise ValueError("DeepSeek V4 hash routing is currently not supported with ring of experts.") self.validate_ragged_buffer_factor() + self.validate_moe_dropless_fallback() self.validate_retry_when_tokens_dropped() + self.validate_moe_log_max_load_ratio() + self.validate_moe_tc_ragged_sort() self.validate_num_moe_emb_chunks() self.validate_moe_quantize_token_all_gather() self.validate_moe_expert_weight_prefetch() + self.validate_moe_topk_before_ep_all_gather() + self.validate_moe_quantize_combine_bwd_method() self.validate_mllog() + self.validate_splash_diag_skip() + self.validate_retry_dropless_first_steps_and_first_phase_buffer() + if self.warmup_skip_train_dropless and not self.warmup_programs_in_init: + raise ValueError("warmup_skip_train_dropless=True requires warmup_programs_in_init=True.") if self.enable_streaming_diloco: if not self.scan_layers: @@ -5165,11 +6376,15 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de f"The number of decoder layers ({self.num_decoder_layers}) must be divisible by " f"(num_diloco_fragments - 1) ({num_transformer_fragments}) when enable_streaming_diloco is True." ) + if self.diloco_bucketize_non_scanned and self.shard_optimizer_over_data: + # train_compile.py derives the outer param shardings after the Zero-1 overlay, the runtime path before it, + # so the bucket split axes could differ between AOT and training. + raise ValueError("diloco_bucketize_non_scanned does not support shard_optimizer_over_data (Zero-1).") - # Gemma 4 small (E2B / E4B) uses per-layer KV sharing, which is incompatible with nn.scan. + # Gemma 4 small (E2B / E4B) uses per-layer KV sharing, which is incompatible with scanned layers. if self.model_name in ("gemma4-e2b", "gemma4-e4b") and self.scan_layers: raise ValueError( - f"{self.model_name} requires scan_layers=False (per-layer KV sharing is incompatible with nn.scan)." + f"{self.model_name} requires scan_layers=False (per-layer KV sharing is incompatible with scanned layers)." ) if self.use_multimodal: # Gemma 4 small (E2B / E4B) only supports text for now; multimodal @@ -5191,8 +6406,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", "maxtext-omni-gemma3-qwen3", - "cosmos3-nano-reasoner", - "cosmos3-super-reasoner", + "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}") @@ -5247,7 +6463,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de DecoderBlockType.QWEN3_NEXT, ): decoder_name = self.decoder_block.value - if not self.sparse_matmul: + if self.num_experts > 1 and not self.sparse_matmul: raise ValueError( f"'explicit' sharding with the '{decoder_name}' decoder requires" " `sparse_matmul=True`; the dense matmul MoE path has not been" @@ -5612,8 +6828,14 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError("`share_kv_projections` is not compatible with `fused_qkv`.") if self.share_kv_projections and self.attention_type == "mla": raise ValueError("`share_kv_projections` is not compatible with `attention_type='mla'`.") - if self.use_sliced_mla_proj and (self.quantization or self.use_qwix_quantization): - raise ValueError("`use_sliced_mla_proj` is not supported with quantization.") + if self.use_sliced_mla_proj and self.quantization: + # Qwix interception still quantizes the sliced dot_general; module-level quant + # (AQT / Fp8 / TE / batch split) would silently fall back to the unsliced path. + if not self.use_qwix_quantization or self.use_batch_split_schedule: + raise ValueError( + "`use_sliced_mla_proj` with quantization is only supported via Qwix interception " + "(use_qwix_quantization=True, use_batch_split_schedule=False)." + ) if self.use_manual_quantization and not self.use_batch_split_schedule: raise ValueError("manual quantization is only used when `use_batch_split_schedule=True`.") @@ -5642,9 +6864,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de if val != 0 and val < 4: raise ValueError(f"compress_ratio must be 0 (disabled) or >= 4, got {val}") - if self.decoder_block == DecoderBlockType.DEEPSEEK4 and self.mtp_num_layers > 0: - raise ValueError("DeepSeek4 decoder block currently does not support MTP layers.") - if self.num_kv_shared_layers > 0: if self.fused_qkv: raise ValueError("`num_kv_shared_layers > 0` is not compatible with `fused_qkv`.") @@ -5663,6 +6882,97 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de f"max_target_length={self.max_target_length}." ) + if self.moe_accumulate_chunk_wgrad and not ( + self.use_tokamax_gmm and self.use_gmm_v2 and self.num_moe_emb_chunks == 0 and not self.prefuse_moe_weights + ): + raise ValueError( + "moe_accumulate_chunk_wgrad=True requires use_tokamax_gmm=True, use_gmm_v2=True, num_moe_emb_chunks=0" + " and prefuse_moe_weights=False." + ) + + if self.moe_accumulate_wi_dlhs and not (self.use_tokamax_gmm and self.use_gmm_v2 and not self.prefuse_moe_weights): + raise ValueError( + "moe_accumulate_wi_dlhs=True requires use_tokamax_gmm=True, use_gmm_v2=True and prefuse_moe_weights=False." + ) + + if self.moe_bwd_prequant_before_unsort and not ( + self.moe_tc_ragged_3d_dispatch and self.moe_tc_ragged_weights_on_activation + ): + raise ValueError( + "moe_bwd_prequant_before_unsort=True requires moe_tc_ragged_3d_dispatch=True and" + " moe_tc_ragged_weights_on_activation=True." + ) + if self.moe_combine_bwd_direct_qarray and not self.moe_bwd_prequant_before_unsort: + raise ValueError("moe_combine_bwd_direct_qarray=True requires moe_bwd_prequant_before_unsort=True.") + if self.moe_bwd_prequant_before_unsort and self.quantization: + # The fused unsort applies the weight / act / bwd scales to the token-order cotangent, so they must be + # per-tensor; the config's qwix rule keeps channelwise axes, so only fixed calibration gives that. With + # moe_combine_bwd_direct_qarray, each EP shard also scales its cotangent before the all-gather, so a dynamic + # (shard-local) scale would differ across shards. + dynamic = [ + f"{name}_quantization_calibration_method={method!r}" + for name in ("weight", "act", "bwd") + if not (method := getattr(self, f"{name}_quantization_calibration_method")).lower().startswith("fixed") + ] + if dynamic: + raise ValueError( + "moe_bwd_prequant_before_unsort=True with quantization requires fixed weight, act and bwd calibration," + f" got {', '.join(dynamic)}." + ) + if self.moe_combine_bwd_direct_qarray and self.moe_quantize_combine_bwd_method: + raise ValueError( + f"moe_quantize_combine_bwd_method={self.moe_quantize_combine_bwd_method!r} has no effect with" + " moe_combine_bwd_direct_qarray=True: the fused EP combine quantizes the cotangent with the gmm bwd rule." + " Unset moe_quantize_combine_bwd_method." + ) + if self.moe_3d_dlhs_partial_sum and not (self.moe_accumulate_wi_dlhs and self.moe_tc_ragged_3d_gmm): + raise ValueError("moe_3d_dlhs_partial_sum=True requires moe_accumulate_wi_dlhs=True and moe_tc_ragged_3d_gmm=True.") + + if self.moe_flat_fsdp_weights: + if not ( + self.sparse_matmul + and self.shard_embed_moe_on_fsdp + and self.use_tokamax_gmm + and self.use_gmm_v2 + and not self.prefuse_moe_weights + ): + raise ValueError( + "moe_flat_fsdp_weights=True requires sparse_matmul=True, shard_embed_moe_on_fsdp=True," + " use_tokamax_gmm=True, use_gmm_v2=True and prefuse_moe_weights=False." + ) + # These either pick their own 3D expert weight layout or carry per-expert params that assume it. + if ( + self.shard_exp_on_fsdp + or self.use_2d_fsdp_sharding + or self.use_batch_split_schedule + or self.weight_dtype in (DType.FLOAT8_E4M3FN, DType.FLOAT8_E5M2) + or self.decoder_block == DecoderBlockType.GEMMA4 + ): + raise ValueError( + "moe_flat_fsdp_weights=True is incompatible with shard_exp_on_fsdp, use_2d_fsdp_sharding," + " use_batch_split_schedule, fp8 weight_dtype and the gemma4 decoder block." + ) + + if self.use_lineage: + if not self.scan_layers: + raise ValueError("use_lineage=True requires scan_layers=True.") + if self.decoder_block != DecoderBlockType.DEEPSEEK: + raise ValueError(f"use_lineage=True requires decoder_block='deepseek', got decoder_block={self.decoder_block!r}.") + if self.attention_type != "mla": + raise ValueError(f"use_lineage=True requires attention_type='mla', got attention_type={self.attention_type!r}.") + if self.rope_type != RopeType.YARN: + raise ValueError(f"use_lineage=True requires rope_type='yarn', got rope_type={self.rope_type!r}.") + if self.capacity_factor <= 0: + raise ValueError(f"use_lineage=True requires capacity_factor > 0, got capacity_factor={self.capacity_factor}.") + + # The freqs_cis table lookup clamps positions >= max_position_embeddings to the last row, while + # rope_direct_position_freqs uses the true angle; keep positions in range so both settings agree. + if self.rope_direct_position_freqs and self.max_target_length > self.max_position_embeddings: + raise ValueError( + "rope_direct_position_freqs=True requires max_target_length <= max_position_embeddings, got " + f"max_target_length={self.max_target_length}, max_position_embeddings={self.max_position_embeddings}." + ) + # I. FINAL TYPE CONVERSIONS AND DERIVED LISTS ici_map = { "diloco": self.ici_diloco_parallelism, @@ -5684,7 +6994,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "dcp": (1), "pcp": (1), } - self.ici_parallelism = [ici_map[axis] for axis in self.mesh_axes] + self.ici_parallelism = _resolve_parallelism(self.ici_parallelism, ici_map, self.mesh_axes) dcn_map = { "diloco": self.dcn_diloco_parallelism, @@ -5706,7 +7016,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "dcp": (1), "pcp": (1), } - self.dcn_parallelism = [dcn_map[axis] for axis in self.mesh_axes] + self.dcn_parallelism = _resolve_parallelism(self.dcn_parallelism, dcn_map, self.mesh_axes) # Zero-1 (`shard_optimizer_over_data`) shards the optimizer moments over the "data" # axis on top of whatever layout the parameters already have. FSDP shards the @@ -5755,6 +7065,17 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de ) self._validate_check_vma_is_supported() + if self.defer_small_all_reduces: + if not (self.use_ring_of_experts and self.sparse_matmul): + raise ValueError( + "defer_small_all_reduces requires use_ring_of_experts=True and sparse_matmul=True (the deferral is" + " implemented on the ring-of-experts sparse_matmul path only)." + ) + if self.moe_dropless_fallback == "layer": + raise ValueError( + "defer_small_all_reduces does not support moe_dropless_fallback='layer' (the in-layer fallback's two" + " branches would emit per-chunk counts of different chunkings)." + ) # Final string-to-enum conversions if they haven't been coerced by pydantic yet. if isinstance(self.decoder_block, str): @@ -6084,6 +7405,9 @@ def set_derived_values_and_validate(self) -> "RLConfig": "gdn_cp_state", "mlpwi", "moe_mlpwi_0", + "moe_x_sorted", + "moe_routing_maps", + "moe_tc_routing_checkpoint", "moe_mlpwi_1", "moe_mlpwo", "moe_route", @@ -6133,10 +7457,10 @@ def get_parallelism_map(prefix: str) -> dict[str, int]: } ici_map = get_parallelism_map("ici") - self.ici_parallelism = [ici_map[axis] for axis in self.mesh_axes] + self.ici_parallelism = _resolve_parallelism(self.ici_parallelism, ici_map, self.mesh_axes) dcn_map = get_parallelism_map("dcn") - self.dcn_parallelism = [dcn_map[axis] for axis in self.mesh_axes] + self.dcn_parallelism = _resolve_parallelism(self.dcn_parallelism, dcn_map, self.mesh_axes) MaxTextConfig.validate_gdn_config(self, is_rl_config=True) MaxTextConfig._validate_colocated_checkpointing(self) diff --git a/src/maxtext/examples/lora_llama3_demo.ipynb b/src/maxtext/examples/lora_llama3_demo.ipynb index 16b098060a..cc31ae9f2d 100644 --- a/src/maxtext/examples/lora_llama3_demo.ipynb +++ b/src/maxtext/examples/lora_llama3_demo.ipynb @@ -33,9 +33,9 @@ "source": [ "## Prerequisites\n", "\n", - "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", + "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", "\n", - "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." + "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." ] }, { @@ -61,7 +61,7 @@ "source": [ "## Installation: MaxText and Post training Dependencies\n", "\n", - "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." + "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/page/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." ] }, { diff --git a/src/maxtext/examples/maxtext_with_gepa.ipynb b/src/maxtext/examples/maxtext_with_gepa.ipynb index f3230d3f03..b24b086578 100644 --- a/src/maxtext/examples/maxtext_with_gepa.ipynb +++ b/src/maxtext/examples/maxtext_with_gepa.ipynb @@ -14,7 +14,7 @@ "\n", "## Installation: MaxText and Post training Dependencies\n", "\n", - "Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." + "Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/page/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." ] }, { diff --git a/src/maxtext/examples/native_lora_demo.ipynb b/src/maxtext/examples/native_lora_demo.ipynb index cce8529307..479ba05a47 100644 --- a/src/maxtext/examples/native_lora_demo.ipynb +++ b/src/maxtext/examples/native_lora_demo.ipynb @@ -28,9 +28,9 @@ "source": [ "## Prerequisites\n", "\n", - "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", + "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", "\n", - "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." + "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." ] }, { @@ -54,7 +54,7 @@ "source": [ "## Installation: MaxText and Post-Training Dependencies\n", "\n", - "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." + "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/page/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." ] }, { diff --git a/src/maxtext/examples/rl_llama3_demo.ipynb b/src/maxtext/examples/rl_llama3_demo.ipynb index 68a9ff95eb..9d51954b23 100644 --- a/src/maxtext/examples/rl_llama3_demo.ipynb +++ b/src/maxtext/examples/rl_llama3_demo.ipynb @@ -26,9 +26,9 @@ "source": [ "## Prerequisites\n", "\n", - "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", + "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", "\n", - "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." + "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." ] }, { @@ -52,7 +52,7 @@ "source": [ "## Installation: MaxText and Post training Dependencies\n", "\n", - "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." + "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/page/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." ] }, { @@ -278,7 +278,7 @@ " print(\"❌ Training Failed!\")\n", " print(\"=\" * 80)\n", " print(\"\\nFor troubleshooting, refer to the Common Pitfalls & Debugging section:\")\n", - " print(\"https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html#common-pitfalls-debugging\")\n", + " print(\"https://maxtext.readthedocs.io/page/guides/run_python_notebook.html#common-pitfalls-debugging\")\n", " raise\n" ] }, @@ -341,7 +341,7 @@ "source": [ "## 📚 Learn More\n", "\n", - "- **CLI Usage**: https://maxtext.readthedocs.io/en/latest/tutorials/rl.html\n", + "- **CLI Usage**: https://maxtext.readthedocs.io/page/tutorials/posttraining/rl.html\n", "- **Configuration**: See `src/maxtext/configs/rl.yml` for all available options\n", "- **Documentation**: Check `src/maxtext/trainers/post_train/rl/train_rl.py` for the `rl_train` function implementation" ] diff --git a/src/maxtext/examples/sft_llama3_demo_gpu.ipynb b/src/maxtext/examples/sft_llama3_demo_gpu.ipynb index cb293cd088..caa7be8a4f 100644 --- a/src/maxtext/examples/sft_llama3_demo_gpu.ipynb +++ b/src/maxtext/examples/sft_llama3_demo_gpu.ipynb @@ -573,7 +573,7 @@ "source": [ "## Learn more\n", "\n", - "- **CLI Usage**: https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/sft.html\n", + "- **CLI Usage**: https://maxtext.readthedocs.io/page/tutorials/posttraining/sft.html\n", "- **Configuration**: See `src/maxtext/configs/post_train/sft.yml` for all available options\n", "- **Documentation**: Check `src/maxtext/trainers/post_train/sft/train_sft.py` for the `train` function implementation" ] diff --git a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb index 216769ac48..ac4ebd1a02 100644 --- a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb +++ b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb @@ -32,9 +32,9 @@ "source": [ "## Prerequisites\n", "\n", - "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", + "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", "\n", - "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." + "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." ] }, { @@ -62,7 +62,7 @@ "source": [ "## Installation: MaxText and Post training Dependencies\n", "\n", - "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." + "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/page/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." ] }, { @@ -315,7 +315,7 @@ " print(\"❌ Training Failed!\")\n", " print(\"=\" * 60)\n", " print(\"\\nFor troubleshooting, refer to the Common Pitfalls & Debugging section:\")\n", - " print(\"https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html#common-pitfalls-debugging\")\n", + " print(\"https://maxtext.readthedocs.io/page/guides/run_python_notebook.html#common-pitfalls-debugging\")\n", " raise\n" ] }, @@ -378,7 +378,7 @@ "source": [ "## 📚 Learn More\n", "\n", - "- **CLI Usage**: https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/sft.html\n", + "- **CLI Usage**: https://maxtext.readthedocs.io/page/tutorials/posttraining/sft.html\n", "- **Configuration**: See `src/maxtext/configs/post_train/sft.yml` for all available options\n", "- **Documentation**: Check `src/maxtext/trainers/post_train/sft/train_sft.py` for the `train` function implementation" ] diff --git a/src/maxtext/examples/sft_multimodal_gemma3_demo.ipynb b/src/maxtext/examples/sft_multimodal_gemma3_demo.ipynb index c50c62cc28..983f76833e 100644 --- a/src/maxtext/examples/sft_multimodal_gemma3_demo.ipynb +++ b/src/maxtext/examples/sft_multimodal_gemma3_demo.ipynb @@ -31,9 +31,9 @@ "source": [ "## Prerequisites\n", "\n", - "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", + "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", "\n", - "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." + "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." ] }, { @@ -59,7 +59,7 @@ "source": [ "## Installation: MaxText and Post training Dependencies\n", "\n", - "**Running the notebook on Visual Studio or JupyterLab**: Before proceeding, create a virtual environment and install the required post-training dependencies by following Option 3: Installing [tpu-post-train] in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." + "**Running the notebook on Visual Studio or JupyterLab**: Before proceeding, create a virtual environment and install the required post-training dependencies by following Option 3: Installing [tpu-post-train] in the [MaxText installation guide](https://maxtext.readthedocs.io/page/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." ] }, { @@ -283,7 +283,7 @@ " print(\"❌ Training Failed!\")\n", " print(\"=\" * 60)\n", " print(\"\\nFor troubleshooting, refer to the Common Pitfalls & Debugging section:\")\n", - " print(\"https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html#common-pitfalls-debugging\")\n", + " print(\"https://maxtext.readthedocs.io/page/guides/run_python_notebook.html#common-pitfalls-debugging\")\n", " raise\n" ] }, diff --git a/src/maxtext/examples/sft_qwen3_demo.ipynb b/src/maxtext/examples/sft_qwen3_demo.ipynb index 0e0e838951..bdb8355748 100644 --- a/src/maxtext/examples/sft_qwen3_demo.ipynb +++ b/src/maxtext/examples/sft_qwen3_demo.ipynb @@ -36,9 +36,9 @@ "source": [ "## Prerequisites\n", "\n", - "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", + "Before running this notebook, make sure your environment is set up for the method you are using. Follow the [Run MaxText Python Notebooks on TPUs](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html) guide and complete all steps for your chosen method (Google Colab, VS Code, or Local Jupyter Lab) before proceeding.\n", "\n", - "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/en/latest/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." + "If you run into issues, refer to the [Common Pitfalls & Debugging](https://maxtext.readthedocs.io/page/guides/run_python_notebook.html#common-pitfalls-debugging) section of the guide." ] }, { @@ -66,7 +66,7 @@ "source": [ "## Installation: MaxText and Post training Dependencies\n", "\n", - "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." + "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/page/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it." ] }, { diff --git a/src/maxtext/experimental/rl/README.md b/src/maxtext/experimental/rl/README.md index 30941bc449..cf7d4c42f3 100644 --- a/src/maxtext/experimental/rl/README.md +++ b/src/maxtext/experimental/rl/README.md @@ -4,29 +4,29 @@ This directory contains code and documentation for **GRPO**, a reinforcement lea ## Key Concepts -* **GRPO (Group Relative Policy Optimization):** A policy optimization algorithm that uses reinforcement learning principles to improve the quality of a language model's responses based on a provided reward signal. -* **FSDP (Fully Sharded Data Parallelism):** A technique for distributing model parameters across multiple devices during training, allowing for efficient training of large models. -* **DP (Data Parallelism):** A technique for distributing data across multiple devices during inference, allowing for efficient inference of large models. -* **Sharding:** For more information on sharding, refer to [getting_started/Sharding.md](https://github.com/AI-Hypercomputer/maxtext/blob/main/getting_started/Sharding.md) -* **[Pathways](https://cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro):** A system that allows to orchestrate computations across multiple TPU slices, enabling training and inference on very large TPU topologies. -* **inference_rollouts:** This is a config parameter to control the frequency in training steps to reshard and copy training parameters to the inference mesh. Must be an integer between 1 and `steps` (total number of training steps). +- **GRPO (Group Relative Policy Optimization):** A policy optimization algorithm that uses reinforcement learning principles to improve the quality of a language model's responses based on a provided reward signal. +- **FSDP (Fully Sharded Data Parallelism):** A technique for distributing model parameters across multiple devices during training, allowing for efficient training of large models. +- **DP (Data Parallelism):** A technique for distributing data across multiple devices during inference, allowing for efficient inference of large models. +- **Sharding:** For more information on sharding, refer to [docs/guides/optimization/sharding.md](https://maxtext.readthedocs.io/page/guides/optimization/sharding.html) +- **[Pathways](https://cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro):** A system that allows to orchestrate computations across multiple TPU slices, enabling training and inference on very large TPU topologies. +- **inference_rollouts:** This is a config parameter to control the frequency in training steps to reshard and copy training parameters to the inference mesh. Must be an integer between 1 and `steps` (total number of training steps). ## Getting Started ### Prerequisites -1. **Google Cloud Platform (GCP):** You need a GCP project with sufficient TPU quota. -2. **GKE Cluster with Pathways:** - * Follow the instructions to create a GKE cluster with Pathways support: [Create a GKE Cluster with Pathways](https://cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster) - * This involves setting up node pools for both training and inference. -3. **Pathways Container Images:** Your GCP service account must be allowlisted to access the Pathways container images. Contact your Google Cloud representative for assistance. -4. **XPK:** (Recommended) Install and configure XPK, the Accelerated Processing Kit, to simplify the management of your GKE cluster and Pathways environment. +1. **Google Cloud Platform (GCP):** You need a GCP project with sufficient TPU quota. +2. **GKE Cluster with Pathways:** + - Follow the instructions to create a GKE cluster with Pathways support: [Create a GKE Cluster with Pathways](https://cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster) + - This involves setting up node pools for both training and inference. +3. **Pathways Container Images:** Your GCP service account must be allowlisted to access the Pathways container images. Contact your Google Cloud representative for assistance. +4. **Cluster Management Tools:** [Cluster Toolkit](../../../../docs/run_maxtext/run_maxtext_via_cluster_toolkit.md) (`gcluster`) is recommended for managing GKE clusters and submitting Pathways RL workloads. (Note: the experimental `tests/end_to_end/tpu/test_grpo.sh` script below still uses legacy XPK commands.) 5. **HF_TOKEN:** You need to have a Hugging Face token to access the models. ### Setup -* Make sure your GKE cluster is properly configured with Pathways. This includes having the correct node pools for training and inference, and that the Pathways Resource Manager, Proxy, and Worker are correctly deployed. -* Set the environment variables for your cluster, project, and zone. +- Make sure your GKE cluster is properly configured with Pathways. This includes having the correct node pools for training and inference, and that the Pathways Resource Manager, Proxy, and Worker are correctly deployed. +- Set the environment variables for your cluster, project, and zone. ## Running GRPO @@ -34,20 +34,23 @@ This repository includes a shell script, `tests/end_to_end/tpu/test_grpo.sh`, th **How it works:** -* **TPU Topology:** This script targets a v5p-256 cluster, which consists of 256 TPU cores. -* **Device Allocation:** - * **Training:** 64 devices are reserved for training. The model is sharded using FSDP across these devices. - * **Inference (Sampling):** The other 64 devices are used for inference. The model is sharded using DP across samplers and TP within each sampler. -* **Samplers:** The inference devices are grouped into samplers. -* **Example Command:** +- **TPU Topology:** This script targets a v5p-256 cluster, which consists of 256 TPU cores. +- **Device Allocation:** + - **Training:** 64 devices are reserved for training. The model is sharded using FSDP across these devices. + - **Inference (Sampling):** The other 64 devices are used for inference. The model is sharded using DP across samplers and TP within each sampler. +- **Samplers:** The inference devices are grouped into samplers. +- **Example Command:** ```bash -HF_TOKEN=${HF_TOKEN} \ +HF_TOKEN=${HF_TOKEN?} \ MODEL=llama3.3-70b \ TOKENIZER=meta-llama/Llama-3.3-70B-Instruct \ NUM_SAMPLERS=8 \ DEVICES_PER_SAMPLER=8 \ +TRAINING_SUBSLICE=4,16 \ +INFERENCE_SUBSLICE=4,16 \ TRAINING_PER_DEVICE_BATCH_SIZE=1 \ INFERENCE_PER_DEVICE_BATCH_SIZE=8 \ STEPS=20 \ bash tests/end_to_end/tpu/test_grpo.sh +``` diff --git a/src/maxtext/inference/decode.py b/src/maxtext/inference/decode.py index 23061e7ba3..90e095defb 100644 --- a/src/maxtext/inference/decode.py +++ b/src/maxtext/inference/decode.py @@ -131,7 +131,7 @@ def main(argv: Sequence[str]) -> None: try: # TODO: update jetstream.engine.tokenizer_api.Tokenizer to maintain tokenizer state. - has_chat_template = getattr(tokenizer_model.tokenizer, "chat_template", False) # pytype: disable=attribute-error + has_chat_template = getattr(tokenizer_model.tokenizer, "chat_template", False) except AttributeError as _: has_chat_template = False is_bos = config.add_bos and not has_chat_template @@ -149,13 +149,13 @@ def main(argv: Sequence[str]) -> None: position_ids, mrope_position_deltas = processor_qwen3_omni.get_rope_index( input_ids=tokens[np.newaxis, :], # Add batch dimension for processing - image_grid_thw=processor_outputs.pixel_grid_thw, # pytype: disable=attribute-error - video_grid_thw=processor_outputs.video_grid_thw, # pytype: disable=attribute-error + image_grid_thw=processor_outputs.pixel_grid_thw, # pyrefly: ignore[missing-attribute] + video_grid_thw=processor_outputs.video_grid_thw, # pyrefly: ignore[missing-attribute] attention_mask=np.ones_like(tokens)[np.newaxis, :], use_audio_in_video=config.use_audio and getattr(processor_outputs, "num_videos", 0) > 0, - audio_lengths=processor_outputs.audio_lengths, # pytype: disable=attribute-error - second_per_grids=processor_outputs.video_second_per_grid, # pytype: disable=attribute-error - spatial_merge_size=config.spatial_merge_size_for_vit, # pytype: disable=attribute-error + audio_lengths=processor_outputs.audio_lengths, # pyrefly: ignore[missing-attribute] + second_per_grids=processor_outputs.video_second_per_grid, # pyrefly: ignore[missing-attribute] + spatial_merge_size=config.spatial_merge_size_for_vit, position_id_per_seconds=config.position_id_per_seconds, config=config, ) @@ -245,8 +245,8 @@ def main(argv: Sequence[str]) -> None: output = tokenizer_model.decode(results) print(f"Input `{text}` -> `{output}`") - assert output.startswith( - config.autoregressive_decode_assert + assert " ".join(output.split()).startswith( + " ".join(config.autoregressive_decode_assert.split()) ), f"generated text mismatch {output=}, {config.autoregressive_decode_assert=}" # Deactivate profiler diff --git a/src/maxtext/inference/diffuser_pipelines/__init__.py b/src/maxtext/inference/diffuser_pipelines/__init__.py new file mode 100644 index 0000000000..2221e69d05 --- /dev/null +++ b/src/maxtext/inference/diffuser_pipelines/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Diffusion inference pipelines package.""" diff --git a/src/maxtext/inference/diffuser_pipelines/pipeline_utils.py b/src/maxtext/inference/diffuser_pipelines/pipeline_utils.py new file mode 100644 index 0000000000..dfc93955ae --- /dev/null +++ b/src/maxtext/inference/diffuser_pipelines/pipeline_utils.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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-agnostic diffusion pipeline utilities and base classes.""" diff --git a/src/maxtext/inference/diffuser_pipelines/wan_pipeline.py b/src/maxtext/inference/diffuser_pipelines/wan_pipeline.py new file mode 100644 index 0000000000..1380b1a566 --- /dev/null +++ b/src/maxtext/inference/diffuser_pipelines/wan_pipeline.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Wan text-to-video and image-to-video inference pipeline.""" diff --git a/src/maxtext/inference/kvcache.py b/src/maxtext/inference/kvcache.py index df0dec0bcc..bc524fb1c1 100644 --- a/src/maxtext/inference/kvcache.py +++ b/src/maxtext/inference/kvcache.py @@ -19,7 +19,6 @@ import jax import jax.numpy as jnp -from flax import linen as nn from flax import nnx from aqt.jax.v2 import config as aqt_config @@ -29,6 +28,7 @@ from maxtext.common.common_types import Array, AxisNames, AxisIdxes, Config, CACHE_BATCH_PREFILL, DType, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, MODEL_MODE_AUTOREGRESSIVE, CACHE_HEADS_NONE, DECODING_ACTIVE_SEQUENCE_INDICATOR from maxtext.common.common_types import CACHE_BATCH, CACHE_SEQUENCE, CACHE_HEADS, CACHE_KV, CACHE_SCALE_BATCH, CACHE_SCALE_SEQUENCE, CACHE_SCALE_HEADS, CACHE_SCALE_KV +from maxtext.utils.sharding import with_logical_constraint MAX_INT8 = 127.5 @@ -433,7 +433,7 @@ def _initialize_ar_cache_vars(self, model_mode): jnp.zeros(cache_shape_key, dtype=dtype), out_sharding=cache_axis_names, ) - self.cached_ar_key[...] = nn.with_logical_constraint( + self.cached_ar_key[...] = with_logical_constraint( self.cached_ar_key[...], cache_axis_names, ) @@ -442,7 +442,7 @@ def _initialize_ar_cache_vars(self, model_mode): jnp.zeros(cache_shape_value, dtype=dtype), out_sharding=cache_axis_names, ) - self.cached_ar_value[...] = nn.with_logical_constraint( + self.cached_ar_value[...] = with_logical_constraint( self.cached_ar_value[...], cache_axis_names, ) @@ -503,9 +503,9 @@ def update_gdn_states(self, new_recurrent_state: Array, new_conv_state: Array) - # not just during MODEL_MODE_PREFILL. This ensures AR updates persist. if getattr(self, "cached_prefill_key", None) is not None: self.cached_prefill_key.set_value( - nn.with_logical_constraint(new_recurrent_state, (cache_batch_axis_name, CACHE_HEADS, None, None)) + with_logical_constraint(new_recurrent_state, (cache_batch_axis_name, CACHE_HEADS, None, None)) ) - self.cached_prefill_value.set_value(nn.with_logical_constraint(new_conv_state, (cache_batch_axis_name, None, None))) + self.cached_prefill_value.set_value(with_logical_constraint(new_conv_state, (cache_batch_axis_name, None, None))) def _get_ar_cache_vars(self): return self.ar_key_vars, self.ar_value_vars, self.cache_ar_segment_id, self.cache_ar_index, self.cached_ar_lengths @@ -747,8 +747,8 @@ def value_body(i, val): cached_value.get_value(), one_token_value_shaped_for_cache, ar_cache_update_idx, ar_cache_update_axis ) ) - cached_key.set_value(nn.with_logical_constraint(cached_key.get_value(), ar_cache_axis_names)) - cached_value.set_value(nn.with_logical_constraint(cached_value.get_value(), ar_cache_axis_names)) + cached_key.set_value(with_logical_constraint(cached_key.get_value(), ar_cache_axis_names)) + cached_value.set_value(with_logical_constraint(cached_value.get_value(), ar_cache_axis_names)) if self.kv_quant: ar_cache_scale_axis_names = transpose_tuple(self.cache_scale_logical_axis_names, self.ar_cache_axis_order) diff --git a/src/maxtext/inference/mlperf/README.md b/src/maxtext/inference/mlperf/README.md index c02b6e4a3c..1d42602e65 100644 --- a/src/maxtext/inference/mlperf/README.md +++ b/src/maxtext/inference/mlperf/README.md @@ -59,7 +59,7 @@ mv 09292024_mixtral_15k_mintoken2_v1.pkl mixtral-processed-data.pkl ``` ### Install Maxtext -For instructions on installing MaxText on your VM, please refer to the [official documentation](https://maxtext.readthedocs.io/en/latest/install_maxtext.html). Then, run the following commands to install the additional dependencies: +For instructions on installing MaxText on your VM, please refer to the [official documentation](https://maxtext.readthedocs.io/page/install_maxtext.html). Then, run the following commands to install the additional dependencies: ``` cd ~ git clone https://github.com/AI-Hypercomputer/maxtext.git diff --git a/src/maxtext/inference/sample.py b/src/maxtext/inference/sample.py new file mode 100644 index 0000000000..7e02b8cd6f --- /dev/null +++ b/src/maxtext/inference/sample.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Canonical diffusion sampling and generation CLI entrypoint.""" diff --git a/src/maxtext/inference/vllm_decode.py b/src/maxtext/inference/vllm_decode.py index 501823baa6..e6fbf20958 100644 --- a/src/maxtext/inference/vllm_decode.py +++ b/src/maxtext/inference/vllm_decode.py @@ -89,11 +89,12 @@ def build_chat_messages(config: Config) -> list[dict[str, Any]]: return messages -def decode_with_vllm(config: Config) -> None: - """Decode using vLLM with a MaxText model implementation. +def create_vllm_llm(config: Config, **overrides: Any) -> LLM: + """Create a vLLM engine with optional keyword overrides. Args: - config: MaxText config. + config: MaxText config; the adapter must already be registered. + **overrides: Additional or overriding keyword arguments for ``vllm.LLM``. """ # Prepare vLLM Arguments vllm_args = { @@ -138,6 +139,8 @@ def decode_with_vllm(config: Config) -> None: if config.max_num_seqs is not None: vllm_args["max_num_seqs"] = config.max_num_seqs + vllm_args.update(overrides) + global _USE_MROPE _USE_MROPE = config.use_multimodal and config.use_mrope @@ -151,7 +154,16 @@ def decode_with_vllm(config: Config) -> None: vllm_config = pyconfig.initialize(argv_list) with logical_axis_rules(vllm_config.logical_axis_rules): - llm = LLM(**vllm_args) + return LLM(**vllm_args) + + +def decode_with_vllm(config: Config) -> None: + """Decode using vLLM with a MaxText model implementation. + + Args: + config: MaxText config. + """ + llm = create_vllm_llm(config) max_logging.log("Generating output...") tokenizer = transformers.AutoTokenizer.from_pretrained( diff --git a/src/maxtext/inference/vllm_score.py b/src/maxtext/inference/vllm_score.py new file mode 100644 index 0000000000..385b1d4c7d --- /dev/null +++ b/src/maxtext/inference/vllm_score.py @@ -0,0 +1,532 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Next-token scoring for any MaxText model served through vLLM on TPU. + +This is a thin shim over ``maxtext.inference.vllm_decode``: the engine is built with the exact +same arguments (``create_vllm_llm``) and chat prompts with the same ``build_chat_messages``; on +top of that it adds *scoring* instead of sampling - the next-token distribution (top-k +log-probs) and the probabilities of user supplied candidate continuations. This is the +primitive needed by "decision" / bounded classification models, multiple-choice evals, +reward scoring, etc. It works for every model supported by the MaxText<->vLLM adapter. + +Example usage (single prompt, score option letters A/B):: + + python3 -m maxtext.inference.vllm_score src/maxtext/configs/base.yml \ + model_name=qwen3.5-9b \ + tokenizer_path=Qwen/Qwen3.5-9B \ + load_parameters_path= \ + vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ + ici_tensor_parallelism=4 \ + hbm_utilization_vllm=0.6 \ + max_target_length=2048 \ + use_chat_template=true \ + system_prompt="Answer with only the option letter A or B." \ + "prompt='What is the capital of Germany? A: Berlin, B: Paris'" \ + --candidates=A,B --top_k=10 --method=auto \ + --enable_thinking=false --enable_prefix_caching=false + +Add ``use_chat_template=true`` (and optionally ``system_prompt=...``) to wrap ``prompt`` with the +tokenizer chat template exactly as ``vllm_decode`` does. + +Batch mode: ``--prompts_file=`` where every line is +``{"prompt": "...", "candidates": ["A", "B"]}`` (``candidates`` optional; strings or token +ids). Batch prompts are used as supplied; chat-template and thinking flags do not apply. +Results are printed and optionally written to ``--output_file`` as JSONL. + +Notes on log-prob support in vLLM on TPU (``tpu_inference`` JAX backend): + +* ``SamplingParams(logprobs=K)`` returns the top-``K`` entries of ``log_softmax(raw_logits)`` + for every generated token (``logprobs_mode`` only switches raw vs. *processed* logprobs; + raw logits are not materialized on TPU). +* ``SamplingParams(prompt_logprobs=k)`` always includes the log-prob of the *actual* prompt + token regardless of rank. Appending a candidate to the prompt and reading its prompt + log-prob is therefore an exact way to score an arbitrary candidate token. +* ``SamplingParams.logprob_token_ids`` / ``allowed_token_ids`` exist upstream but are not + wired into the TPU sampler yet, so they are not relied upon here. + +Because ``softmax`` restricted to a candidate set is invariant to the constant shift between +raw logits and log-probabilities, candidate probabilities computed from vLLM log-probs are +identical to the ones computed from raw logits. +""" + +import json +import os +from typing import Any, Sequence + +from absl import app +from absl import flags +import jax +import numpy as np + +from maxtext.common.common_types import Config +from maxtext.configs import pyconfig +from maxtext.inference import vllm_decode +import maxtext.integration.vllm.maxtext_vllm_adapter as adapter +from maxtext.utils import max_logging +import transformers +from vllm import LLM +from vllm.inputs import TokensPrompt +from vllm.sampling_params import SamplingParams + +DEFAULT_TOP_K = 20 + +FLAGS = flags.FLAGS +flags.DEFINE_string( + "candidates", "", "Comma separated candidate text continuations. Use JSONL integer values for raw token IDs." +) +flags.DEFINE_string( + "prompts_file", + "", + "JSONL with {prompt, candidates} per line. Prompts must already be formatted; chat-template flags are ignored.", +) +flags.DEFINE_string("output_file", "", "Optional JSONL output path.") +flags.DEFINE_integer("top_k", DEFAULT_TOP_K, "Number of next-token log-probs to request / print per prompt.") +flags.DEFINE_enum("method", "auto", ["auto", "top_k", "exact"], "Candidate scoring method (see VllmScorer).") +flags.DEFINE_bool("add_special_tokens", False, "Let the tokenizer add BOS/special tokens to text prompts.") +flags.DEFINE_bool( + "enable_prefix_caching", + False, + "Enable vLLM prefix caching. Disabled by default for scoring: on hybrid (GDN/mamba) models the TPU " + "prefix-cache path was observed to change next-token log-probs for identical prompts.", +) +flags.DEFINE_integer( + "repeat", + 1, + "Run the scoring pass this many times (determinism check); results of the last run are reported.", +) +flags.DEFINE_bool( + "enable_thinking", + False, + "Passed to tokenizer.apply_chat_template when use_chat_template=true. Off by default for scoring: Qwen3-style " + "templates otherwise leave the prompt inside an open block, so the next token is reasoning text rather " + "than the answer.", +) + +flags.register_validator("top_k", lambda value: value >= 1, message="top_k must be at least 1.") +flags.register_validator("repeat", lambda value: value >= 1, message="repeat must be at least 1.") + +TokenIds = Sequence[int] +PromptLike = str | TokenIds + + +def setup_scoring_environment() -> None: + """Set the TPU environment for scoring with an in-process vLLM engine. + + ``VLLM_ENABLE_V1_MULTIPROCESSING=0`` runs the vLLM ``EngineCore`` in-process. This is + required because MaxText's ``pyconfig.initialize`` already initializes the JAX TPU backend + in the calling process; a forked ``EngineCore`` would otherwise deadlock on its first XLA + compilation (same approach as ``maxtext.eval.runner.server_manager``). + """ + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + os.environ["SKIP_JAX_PRECOMPILE"] = "1" + os.environ["NEW_MODEL_DESIGN"] = "1" + + jax.config.update("jax_default_prng_impl", "unsafe_rbg") + os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0" + if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""): + os.environ["LIBTPU_INIT_ARGS"] = ( + os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true" + ) + + +class VllmScorer: + """Model-agnostic next-token scoring on top of a ``vllm.LLM`` engine. + + Typical usage:: + + adapter.register(config) + llm = vllm_decode.create_vllm_llm(config, max_logprobs=DEFAULT_TOP_K, enable_prefix_caching=False) + scorer = VllmScorer(llm) + top = scorer.next_token_logprobs(["Paris is the"]) # {token_id: logprob} + lps = scorer.candidate_logprobs([prompt_ids], [[tok_a, tok_b]]) # exact candidate scores + + Prompts may be text or pre-tokenized ids. Text prompts are tokenized with + ``add_special_tokens=False`` so chat-templated prompts are not altered. + """ + + def __init__(self, llm: LLM, default_top_k: int = DEFAULT_TOP_K): + self.llm = llm + self.tokenizer = llm.get_tokenizer() + if default_top_k < 1: + raise ValueError("default_top_k must be at least 1") + self.default_top_k = default_top_k + max_logprobs = llm.llm_engine.model_config.max_logprobs + if max_logprobs != -1 and default_top_k > max_logprobs: + max_logging.log( + f"VllmScorer: default_top_k={default_top_k} exceeds engine max_logprobs={max_logprobs}; " + f"clamping. Pass max_logprobs= to create_vllm_llm to raise the limit." + ) + self.default_top_k = max_logprobs + + def to_token_ids(self, prompt: PromptLike, add_special_tokens: bool = False) -> list[int]: + """Tokenize text or copy an existing token sequence.""" + if isinstance(prompt, str): + return list(self.tokenizer.encode(prompt, add_special_tokens=add_special_tokens)) + return [int(t) for t in prompt] + + def _tokens_prompts(self, prompts: Sequence[PromptLike], add_special_tokens: bool) -> list[TokensPrompt]: + return [TokensPrompt(prompt_token_ids=self.to_token_ids(p, add_special_tokens)) for p in prompts] + + def next_token_logprobs( + self, + prompts: Sequence[PromptLike], + top_k: int | None = None, + add_special_tokens: bool = False, + ) -> list[dict[int, float]]: + """Returns the top-``k`` next-token log-probs (``log_softmax`` of raw logits) per prompt. + + One single-token generation request (``max_tokens=1``, greedy) is issued per prompt; all prompts are + batched by the vLLM scheduler. + + Returns: + A list (one per prompt) of ``{token_id: logprob}`` dicts containing the top-``k`` tokens + (the greedy token is always among them). + """ + k = self.default_top_k if top_k is None else top_k + max_logprobs = self.llm.llm_engine.model_config.max_logprobs + if k < 1 or (max_logprobs != -1 and k > max_logprobs): + raise ValueError(f"top_k must be positive and within the engine max_logprobs limit ({max_logprobs})") + sampling_params = SamplingParams(max_tokens=1, temperature=0.0, logprobs=k, detokenize=False) + outputs = self.llm.generate(self._tokens_prompts(prompts, add_special_tokens), sampling_params, use_tqdm=False) + results: list[dict[int, float]] = [] + for output in outputs: + step_logprobs = output.outputs[0].logprobs + if not step_logprobs: + raise RuntimeError("vLLM did not return logprobs; make sure the engine was created with max_logprobs>=k.") + results.append({int(tid): float(lp.logprob) for tid, lp in step_logprobs[0].items()}) + return results + + def greedy_next_token(self, prompts: Sequence[PromptLike], add_special_tokens: bool = False) -> list[int]: + """Returns the arg-max next token id for each prompt.""" + sampling_params = SamplingParams(max_tokens=1, temperature=0.0, detokenize=False) + outputs = self.llm.generate(self._tokens_prompts(prompts, add_special_tokens), sampling_params, use_tqdm=False) + return [int(o.outputs[0].token_ids[0]) for o in outputs] + + def exact_candidate_logprobs( + self, + prompts: Sequence[PromptLike], + candidate_token_ids: Sequence[TokenIds], + add_special_tokens: bool = False, + ) -> list[np.ndarray]: + """Exact next-token log-probs of arbitrary candidate tokens via ``prompt_logprobs``. + + For every (prompt, candidate) pair one request ``prompt + [candidate]`` is issued and the + log-prob of the appended candidate (conditioned on the prompt) is read back from vLLM's + prompt log-probs. Never misses a candidate; costs ``sum(len(cands))`` prefill requests + (batched by the scheduler). + """ + if len(prompts) != len(candidate_token_ids): + raise ValueError("prompts and candidate_token_ids must have the same length") + for candidates in candidate_token_ids: + for candidate in candidates: + _validate_token_id(self.tokenizer, candidate) + requests: list[TokensPrompt] = [] + owners: list[tuple[int, int]] = [] + for i, (prompt, cands) in enumerate(zip(prompts, candidate_token_ids)): + ids = self.to_token_ids(prompt, add_special_tokens) + for j, cand in enumerate(cands): + requests.append(TokensPrompt(prompt_token_ids=ids + [int(cand)])) + owners.append((i, j)) + if not requests: + return [np.zeros((0,), dtype=np.float64) for _ in prompts] + + # Request a positive count to avoid backends treating zero as disabled. + # The actual candidate token is included regardless of its rank. + sampling_params = SamplingParams(max_tokens=1, temperature=0.0, prompt_logprobs=1, detokenize=False) + outputs = self.llm.generate(requests, sampling_params, use_tqdm=False) + + results = [np.full((len(c),), np.nan, dtype=np.float64) for c in candidate_token_ids] + for (i, j), output, request in zip(owners, outputs, requests): + cand = request["prompt_token_ids"][-1] + prompt_logprobs = output.prompt_logprobs + if not prompt_logprobs or prompt_logprobs[-1] is None or cand not in prompt_logprobs[-1]: + raise RuntimeError(f"vLLM did not return a prompt logprob for candidate token {cand} (prompt {i}).") + results[i][j] = float(prompt_logprobs[-1][cand].logprob) + return results + + def candidate_logprobs( + self, + prompts: Sequence[PromptLike], + candidate_token_ids: Sequence[TokenIds], + method: str = "auto", + top_k: int | None = None, + add_special_tokens: bool = False, + *, + top_k_logprobs: Sequence[dict[int, float]] | None = None, + ) -> list[np.ndarray]: + """Returns ``log_softmax`` scores of the given candidate next tokens for each prompt. + + Args: + prompts: Batch of prompts (text or token ids). + candidate_token_ids: Per-prompt list of candidate token ids to score. + method: ``"top_k"`` (one request per prompt; raises if a candidate is outside the top-k), + ``"exact"`` (``prompt_logprobs`` based, one request per candidate) or ``"auto"`` + (top-k first, exact fallback only for prompts with missing candidates). + top_k: Number of logprobs to request in the top-k pass (defaults to ``default_top_k``). + add_special_tokens: Whether to let the tokenizer add special tokens to text prompts. + top_k_logprobs: Optional precomputed distributions aligned with prompts, using the + same tokenization and top_k. Ignored for the exact method. + + Returns: + A list of ``np.ndarray`` (float64) aligned with ``candidate_token_ids``. + """ + if method not in ("auto", "top_k", "exact"): + raise ValueError(f"Unknown method {method!r}; expected 'auto', 'top_k' or 'exact'.") + if len(prompts) != len(candidate_token_ids): + raise ValueError("prompts and candidate_token_ids must have the same length") + for candidates in candidate_token_ids: + for candidate in candidates: + _validate_token_id(self.tokenizer, candidate) + + if method == "exact": + return self.exact_candidate_logprobs(prompts, candidate_token_ids, add_special_tokens) + + if top_k_logprobs is not None and len(top_k_logprobs) != len(prompts): + raise ValueError("prompts and top_k_logprobs must have the same length") + top = ( + self.next_token_logprobs(prompts, top_k=top_k, add_special_tokens=add_special_tokens) + if top_k_logprobs is None + else top_k_logprobs + ) + results: list[np.ndarray] = [] + missing: list[int] = [] + for i, (dist, cands) in enumerate(zip(top, candidate_token_ids)): + scores = np.array([dist.get(int(c), np.nan) for c in cands], dtype=np.float64) + if np.isnan(scores).any(): + missing.append(i) + results.append(scores) + + if missing: + if method == "top_k": + raise RuntimeError( + f"Candidates outside top-{top_k or self.default_top_k} for prompts {missing}; " + "increase top_k / max_logprobs or use method='exact' or 'auto'." + ) + max_logging.log(f"VllmScorer: falling back to exact prompt_logprobs scoring for prompts {missing}.") + exact = self.exact_candidate_logprobs( + [prompts[i] for i in missing], [candidate_token_ids[i] for i in missing], add_special_tokens + ) + for i, scores in zip(missing, exact): + results[i] = scores + return results + + def candidate_probabilities( + self, + prompts: Sequence[PromptLike], + candidate_token_ids: Sequence[TokenIds], + **kwargs: Any, + ) -> list[np.ndarray]: + """Softmax over the candidate set only (identical to ``softmax(raw_logits[cands])``).""" + return [ + _candidate_probabilities(scores) for scores in self.candidate_logprobs(prompts, candidate_token_ids, **kwargs) + ] + + +def _parse_candidates(raw: Any) -> list[Any]: + """Parse comma-separated options or copy a JSON candidate list.""" + if raw is None or raw == "": + return [] + if isinstance(raw, str): + raw = [c for c in raw.split(",") if c != ""] + return list(raw) + + +def _validate_token_id(tokenizer: Any, token_id: int) -> None: + """Reject non-integer and out-of-vocabulary candidate token IDs.""" + if isinstance(token_id, (bool, np.bool_)) or not isinstance(token_id, (int, np.integer)): + raise ValueError("Candidate token IDs must be integers") + if not 0 <= token_id < len(tokenizer): + raise ValueError(f"Candidate token ID {token_id} is outside the tokenizer vocabulary") + + +def candidate_token_id( + tokenizer: Any, + prompt_text: str, + prompt_ids: Sequence[int], + candidate: Any, + add_special_tokens: bool = False, +) -> int: + """Resolves a candidate (token id or string) to the single token id that follows ``prompt_ids``. + + Only integer candidates are raw token IDs; numeric strings are candidate text. + Strings are tokenized *in context* using the original prompt and its special-token + setting (prompt + candidate) so that BPE merges across the + boundary are detected. Multi-token candidates are rejected: next-token scoring is only well + defined for single-token continuations. + """ + if add_special_tokens and prompt_ids and prompt_ids[-1] == tokenizer.eos_token_id: + raise ValueError("The tokenizer appended EOS to the prompt; use add_special_tokens=false for next-token scoring.") + if isinstance(candidate, bool) or not isinstance(candidate, (int, str)): + raise ValueError("Candidates must be text strings or integer token IDs") + if isinstance(candidate, int): + _validate_token_id(tokenizer, candidate) + return candidate + combined = tokenizer.encode(prompt_text + candidate, add_special_tokens=add_special_tokens) + suffix = combined[len(prompt_ids) :] + if list(combined[: len(prompt_ids)]) != list(prompt_ids) or len(suffix) != 1: + raise ValueError( + f"Candidate {candidate!r} is not a single token at the prompt boundary (got {len(suffix)} tokens: {suffix})." + ) + return int(suffix[0]) + + +def load_requests(config: Config, tokenizer: Any) -> list[dict[str, Any]]: + """Builds the list of {prompt, prompt_ids, candidates, candidate_ids} requests from flags/config.""" + requests: list[dict[str, Any]] = [] + if FLAGS.prompts_file: + with open(FLAGS.prompts_file, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + item = json.loads(line) + requests.append({"prompt": item["prompt"], "candidates": _parse_candidates(item.get("candidates"))}) + else: + prompt = config.prompt + if config.use_chat_template: + # Same chat formatting as vllm_decode (system_prompt / multimodal placeholders included), + # plus enable_thinking so scoring happens at the answer position. + prompt = tokenizer.apply_chat_template( + vllm_decode.build_chat_messages(config), + tokenize=False, + add_generation_prompt=True, + add_special_tokens=False, + enable_thinking=FLAGS.enable_thinking, + ) + requests.append({"prompt": prompt, "candidates": _parse_candidates(FLAGS.candidates)}) + + for req in requests: + req["prompt_ids"] = tokenizer.encode(req["prompt"], add_special_tokens=FLAGS.add_special_tokens) + req["candidate_ids"] = [ + candidate_token_id(tokenizer, req["prompt"], req["prompt_ids"], candidate, FLAGS.add_special_tokens) + for candidate in req["candidates"] + ] + return requests + + +def score_requests(scorer: VllmScorer, requests: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Runs top-k next-token scoring and candidate scoring for all requests.""" + tokenizer = scorer.tokenizer + prompt_ids = [req["prompt_ids"] for req in requests] + top = scorer.next_token_logprobs(prompt_ids, top_k=FLAGS.top_k) + + with_cands = [i for i, req in enumerate(requests) if req["candidate_ids"]] + cand_scores: dict[int, np.ndarray] = {} + if with_cands: + scores = scorer.candidate_logprobs( + [prompt_ids[i] for i in with_cands], + [requests[i]["candidate_ids"] for i in with_cands], + method=FLAGS.method, + top_k=FLAGS.top_k, + top_k_logprobs=[top[i] for i in with_cands], + ) + cand_scores = dict(zip(with_cands, scores)) + + results = [] + for i, req in enumerate(requests): + dist = top[i] + greedy_id = max(dist, key=dist.get) + result: dict[str, Any] = { + "prompt": req["prompt"], + "num_prompt_tokens": len(req["prompt_ids"]), + "greedy_token": {"id": greedy_id, "text": tokenizer.decode([greedy_id]), "logprob": dist[greedy_id]}, + "top_k": [ + {"id": tid, "text": tokenizer.decode([tid]), "logprob": lp} + for tid, lp in sorted(dist.items(), key=lambda kv: kv[1], reverse=True)[: FLAGS.top_k] + ], + } + if i in cand_scores: + lps = cand_scores[i] + probs = _candidate_probabilities(lps) + result["candidates"] = [ + {"candidate": str(c), "id": tid, "logprob": float(lp), "prob_among_candidates": float(p)} + for c, tid, lp, p in zip(req["candidates"], req["candidate_ids"], lps, probs) + ] + result["prediction"] = str(req["candidates"][int(np.argmax(lps))]) + results.append(result) + return results + + +def _candidate_probabilities(logprobs: np.ndarray) -> np.ndarray: + """Normalize log-probabilities over the supplied candidate set.""" + if len(logprobs) == 0: + return np.zeros((0,), dtype=np.float64) + weights = np.exp(logprobs - np.max(logprobs)) + return weights / np.sum(weights) + + +def log_results(results: Sequence[dict[str, Any]]) -> None: + """Log each prediction with options ranked by candidate probability.""" + for index, result in enumerate(results, start=1): + heading = f"Result {index}" + if "prediction" in result: + heading += f" — prediction: {result['prediction']}" + lines = [heading, f"Prompt ({result['num_prompt_tokens']} tokens): {result['prompt']!r}"] + candidates = result.get("candidates", []) + if candidates: + ranked = sorted(candidates, key=lambda option: option["prob_among_candidates"], reverse=True) + option_width = max(len("Option"), *(len(repr(option["candidate"])) for option in ranked)) + lines.append(f"{'Option':<{option_width}} {'Option-prob':>12} {'Log-prob':>10}") + for option in ranked: + lines.append( + f"{option['candidate']!r:<{option_width}} {option['prob_among_candidates']:>12.2%} " + f"{option['logprob']:>10.4f}" + ) + greedy = result["greedy_token"] + lines.append(f"Greedy next token: {greedy['text']!r} (id={greedy['id']}, log-prob={greedy['logprob']:.4f})") + if result["top_k"]: + lines.append("Top-k next tokens:") + for token in result["top_k"]: + lines.append(f" {token['text']!r} (id={token['id']}, log-prob={token['logprob']:.4f})") + max_logging.log("\n".join(lines)) + + +def main(argv: Sequence[str]) -> None: + """Load the model, score requests, and report the option outcomes.""" + setup_scoring_environment() + + config = pyconfig.initialize(argv) + adapter.register(config) + + # Resolve prompts/candidates first (fail fast on multi-token candidates) - same tokenizer + # source as vllm_decode; the engine's tokenizer is identical since vLLM loads tokenizer_path. + tokenizer = transformers.AutoTokenizer.from_pretrained(config.tokenizer_path, token=config.hf_access_token) + requests = load_requests(config, tokenizer) + + llm = vllm_decode.create_vllm_llm(config, max_logprobs=FLAGS.top_k, enable_prefix_caching=FLAGS.enable_prefix_caching) + scorer = VllmScorer(llm, default_top_k=FLAGS.top_k) + + for i in range(FLAGS.repeat): + results = score_requests(scorer, requests) + if FLAGS.repeat > 1: + summary = [ + [round(c["logprob"], 6) for c in r["candidates"]] + if "candidates" in r + else round(r["greedy_token"]["logprob"], 6) + for r in results + ] + max_logging.log(f"[repeat {i + 1}/{FLAGS.repeat}] candidate/greedy logprobs per prompt: {summary}") + + log_results(results) + + if FLAGS.output_file: + with open(FLAGS.output_file, "w", encoding="utf-8") as f: + for result in results: + f.write(json.dumps(result, ensure_ascii=False) + "\n") + max_logging.log(f"Wrote {len(results)} results to {FLAGS.output_file}") + + +if __name__ == "__main__": + app.run(main) diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py index a70a91a6a2..8421d6500d 100644 --- a/src/maxtext/input_pipeline/hf_data_processing.py +++ b/src/maxtext/input_pipeline/hf_data_processing.py @@ -52,7 +52,7 @@ def _get_training_objective_transform( pad_id: int, bos_token_id: int | None, ) -> input_pipeline_utils.ShiftData | input_pipeline_utils.BlockDiffusionCorruption | None: - """Selects target preparation for causal or block-diffusion pre-training. + """Selects target preparation for causal or block-diffusion training. Args: config: Training configuration containing the objective-specific settings. @@ -68,16 +68,15 @@ def _get_training_objective_transform( Raises: ValueError: If the objective is unsupported or block diffusion is combined with - an incompatible post-training or packing mode. + DPO or packing. """ objective = getattr(config, "training_objective", "causal_lm") if objective == "block_diffusion": - if use_sft: - raise ValueError("This block-diffusion integration currently supports pre-training only.") + completion_only = bool(use_sft and getattr(config, "sft_train_on_completion_only", False)) if use_dpo: - raise ValueError("Block-diffusion pre-training is not compatible with DPO.") + raise ValueError("Block-diffusion training is not compatible with DPO.") if packing: - raise ValueError("Block-diffusion pre-training requires packing=False.") + raise ValueError("Block-diffusion training requires packing=False.") return input_pipeline_utils.BlockDiffusionCorruption( block_size=config.causal_block_size, mask_id=config.block_diffusion_mask_id, @@ -85,6 +84,7 @@ def _get_training_objective_transform( logit_alignment=config.block_diffusion_logit_alignment, canvas_policy=config.block_diffusion_canvas_policy, axis=1, + completion_only=completion_only, ) if objective != "causal_lm": raise ValueError(f"Unsupported training objective: {objective}") @@ -431,6 +431,7 @@ def preprocessing_pipeline( completion_only=sft_train_on_completion_only, max_target_length=max_target_length, unk_id=pad_id, + training_objective=getattr(config, "training_objective", "causal_lm"), ) ) data_column_names = ("inputs", "targets") diff --git a/src/maxtext/input_pipeline/input_pipeline_interface.py b/src/maxtext/input_pipeline/input_pipeline_interface.py index 4a0b37532f..034b6ca4b8 100644 --- a/src/maxtext/input_pipeline/input_pipeline_interface.py +++ b/src/maxtext/input_pipeline/input_pipeline_interface.py @@ -46,7 +46,9 @@ def get_process_loading_real_data( return list(process_loading_real_data) -def create_process_specific_iterator(config: pyconfig.HyperParameters, mesh, process_indices, input_iterator): +def create_process_specific_iterator( + config: pyconfig.HyperParameters, mesh, process_indices, input_iterator, is_training: bool = True +): """ If the current process's index is among the `process_indices`, a real data iterator is created. Otherwise, a placeholder iterator is returned. @@ -55,7 +57,7 @@ def create_process_specific_iterator(config: pyconfig.HyperParameters, mesh, pro iterator_fn = functools.partial(input_iterator, config, mesh, process_indices) output_iterator = iterator_fn() else: - output_iterator = PlaceHolderDataIterator(config, mesh) + output_iterator = PlaceHolderDataIterator(config, mesh, is_training=is_training) return output_iterator @@ -116,5 +118,7 @@ def create_data_iterator(config: pyconfig.HyperParameters, mesh): if config.expansion_factor_real_data > 1: assert len(process_indices_eval) == jax.process_count() // config.expansion_factor_real_data - output_eval_iterator = create_process_specific_iterator(config, mesh, process_indices_eval, eval_iterator) + output_eval_iterator = create_process_specific_iterator( + config, mesh, process_indices_eval, eval_iterator, is_training=False + ) return output_train_iterator, output_eval_iterator diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index 01d228df16..c8b7075c28 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -360,33 +360,51 @@ def tokenization(example, hf_tokenizer, truncation, max_length, column_names): @dataclasses.dataclass class SFTPromptMasking(grain.MapTransform): """Construct inputs and targets for SFT training. Concat prompt and completion to generate inputs. - For targets, if train on completion only, the prompt will be masked by unk_id. Otherwise the same as inputs. + Causal SFT can mask prompt targets with ``unk_id``. Block diffusion keeps + clean same-position targets and emits role-derived completion metadata. """ - def __init__(self, text_column_name, completion_only, max_target_length, unk_id=0): + def __init__( + self, + text_column_name, + completion_only, + max_target_length, + unk_id=0, + training_objective="causal_lm", + ): + if training_objective not in ("causal_lm", "block_diffusion"): + raise ValueError(f"Unsupported training objective: {training_objective}") self.text_column_name = text_column_name self.completion_only = completion_only self.max_target_length = max_target_length self.unk_id = unk_id + self.training_objective = training_objective def map(self, element): """ Maps a single dataset element to an SFT training instance. It concatenates the prompt and completion to form the `inputs` sequence. For the `targets` sequence: - - If `self.completion_only` is `True`, the prompt portion of the - concatenated sequence is masked using `self.unk_id`. + - If `self.completion_only` is `True`, causal SFT masks the prompt portion + with `self.unk_id`; block-diffusion SFT keeps clean targets and records + completion eligibility separately. - If `self.completion_only` is `False`, the target sequence is identical to the input sequence. """ - inputs, targets = [], [] + is_block_diffusion = self.training_objective == "block_diffusion" + inputs, targets, completion_mask = [], [], [] for i, text in enumerate(element[self.text_column_name]): inputs += text - targets += [self.unk_id] * len(text) if self.completion_only and element["is_prompt"][i] else text - return { + is_prompt = element["is_prompt"][i] + targets += [self.unk_id] * len(text) if self.completion_only and is_prompt and not is_block_diffusion else text + completion_mask += [not is_prompt] * len(text) + output = { "inputs": np.asarray(inputs[: self.max_target_length], dtype=np.int32), "targets": np.asarray(targets[: self.max_target_length], dtype=np.int32), } + if is_block_diffusion: + output["completion_mask"] = np.asarray(completion_mask[: self.max_target_length], dtype=np.int32) + return output @dataclasses.dataclass @@ -914,7 +932,7 @@ def map( self.config is not None and getattr(self.config, "training_objective", "causal_lm") == "block_diffusion" ) for data_column in data_columns: - if data_column != "images": + if data_column not in ("images", "completion_mask"): if isinstance(element[data_column], mm_utils.PreprocessorOutput): raise TypeError("Only 'images' column can be of type PreprocessorOutput.") @@ -949,7 +967,7 @@ def map( element["images"] = self._pad_image_and_mask(element["images"]) # pyrefly: ignore[bad-argument-type] - elif preserve_pad_valued_tokens and key.endswith(("_segmentation", "_position")): + elif key == "completion_mask" or (preserve_pad_valued_tokens and key.endswith(("_segmentation", "_position"))): element[key] = self._pad_text(element[key], self.max_length, 0) # pyrefly: ignore[bad-argument-type] elif "true_length" not in key: element[key] = self._pad_text(element[key], self.max_length, self.pad_id) # pyrefly: ignore[bad-argument-type] @@ -1380,6 +1398,7 @@ class BlockDiffusionCorruption(grain.RandomMapTransform): logit_alignment: str = "same_position" canvas_policy: str = "all_masked" axis: int = 1 + completion_only: bool = False def random_map(self, element, rng: np.random.Generator): """Corrupts inputs while preserving clean targets and input metadata.""" @@ -1391,9 +1410,51 @@ def random_map(self, element, rng: np.random.Generator): "inputs, targets, and targets_segmentation must have identical shapes, got " f"{inputs.shape}, {targets.shape}, and {targets_segmentation.shape}" ) + validity_mask = targets_segmentation != 0 + raw_completion_mask = element.get("completion_mask") + if raw_completion_mask is None: + if self.completion_only: + if "completion_mask" not in element: + raise ValueError("completion-only block diffusion requires an explicit completion_mask") + raise ValueError("completion-only block diffusion requires a non-None completion_mask") + completion_mask = validity_mask + else: + completion_mask = np.asarray(raw_completion_mask) != 0 + if completion_mask.shape != inputs.shape: + raise ValueError(f"completion_mask must match inputs shape; received {completion_mask.shape} and {inputs.shape}") + if np.any(completion_mask & ~validity_mask): + raise ValueError("completion_mask must be a subset of valid target positions") + can_validate_block_order = ( + self.completion_only + and self.block_size > 0 + and inputs.ndim > 0 + and -inputs.ndim <= self.axis < inputs.ndim + and inputs.shape[self.axis % inputs.ndim] > 0 + ) + if can_validate_block_order: + axis = self.axis % inputs.ndim + sequence_length = inputs.shape[axis] + block_count = (sequence_length + self.block_size - 1) // self.block_size + padded_length = block_count * self.block_size + + def _as_blocks(mask): + rows = np.moveaxis(mask, axis, -1).reshape(-1, sequence_length) + padded = np.pad(rows, ((0, 0), (0, padded_length - sequence_length))) + return padded.reshape(rows.shape[0], block_count, self.block_size) + + validity_blocks = _as_blocks(validity_mask) + completion_blocks = _as_blocks(completion_mask) + seen_completion = np.maximum.accumulate(completion_blocks, axis=-1) + if np.any(seen_completion & validity_blocks & ~completion_blocks): + raise ValueError( + "completion-only block diffusion cannot place a prompt token after " + "a completion token in the same diffusion block; split the example " + "or align the next prompt to a block boundary" + ) + supervision_mask = completion_mask if self.completion_only else validity_mask result = block_diffusion_corruption.corrupt_tokens( inputs, - targets_segmentation != 0, + supervision_mask, rng, block_size=self.block_size, mask_id=self.mask_id, diff --git a/src/maxtext/input_pipeline/multihost_dataloading.py b/src/maxtext/input_pipeline/multihost_dataloading.py index 07237349c8..2cfc82ffc5 100644 --- a/src/maxtext/input_pipeline/multihost_dataloading.py +++ b/src/maxtext/input_pipeline/multihost_dataloading.py @@ -233,6 +233,7 @@ def __init__( raise ValueError("Type error: dataloader should be either tf.data.Dataset or Iterable.") self.out_of_data = False self.last_local_data = None + self._reset_and_unread = False self.generate_padding_batch = generate_padding_batch self.expansion_loading_factor_for_grain = expansion_loading_factor_for_grain _GLOBAL_MULTIHOST_ITERATORS.add(self) @@ -246,6 +247,7 @@ def reset(self): raise ValueError("Type error: dataloader should be either tf.data.Dataset or Iterable.") self.out_of_data = False self.last_local_data = None + self._reset_and_unread = True def close(self): if hasattr(self, "local_iterator"): @@ -260,10 +262,16 @@ def __del__(self): self.close() def __iter__(self): - self.reset() + # A for loop starts from the first batch. Skip the rebuild when `reset()` was just + # called and nothing has been read since (e.g. the eval loop in train.py resets the + # iterator explicitly, because wrappers such as `_ReorderedDataIterator` do not reset + # in `__iter__`); building a second tf.data iterator would discard the first one. + if not self._reset_and_unread: + self.reset() return self def __next__(self): + self._reset_and_unread = False return self._get_next_batch_sharded() def _get_next_batch_sharded(self) -> jax.Array: diff --git a/src/maxtext/input_pipeline/synthetic_data_processing.py b/src/maxtext/input_pipeline/synthetic_data_processing.py index 7c79eebfcb..964397ec2b 100644 --- a/src/maxtext/input_pipeline/synthetic_data_processing.py +++ b/src/maxtext/input_pipeline/synthetic_data_processing.py @@ -135,9 +135,9 @@ def raw_generate_synthetic_data(config: pyconfig.HyperParameters, data): class PlaceHolderDataIterator: """Creates a place holder synthetic data iterator for loading on subset of hosts""" - def __init__(self, config: pyconfig.HyperParameters, mesh): + def __init__(self, config: pyconfig.HyperParameters, mesh, is_training: bool = True): self.mesh = mesh - dataset = PlaceHolderDataIterator.get_place_holder_synthetic_data(config) + dataset = PlaceHolderDataIterator.get_place_holder_synthetic_data(config, is_training=is_training) self.data_generator = multihost_dataloading.MultiHostDataLoadIterator(dataset, self.mesh) def __iter__(self): @@ -150,9 +150,17 @@ def reset(self): pass @staticmethod - def get_place_holder_synthetic_data(config: pyconfig.HyperParameters): - """fill negative value in synthetic data""" - batch_size = config.global_batch_size_to_load // jax.process_count() + def get_place_holder_synthetic_data(config: pyconfig.HyperParameters, is_training: bool = True): + """fill negative value in synthetic data + + Args: + config: hyperparameters, used for the loaded batch size and the sequence length. + is_training: whether this placeholder stands in for the train iterator. The eval + iterator loads `global_batch_size_to_load_eval` rows instead, and every host has + to contribute the same number of rows to the global batch. + """ + global_batch_size_to_load = config.global_batch_size_to_load if is_training else config.global_batch_size_to_load_eval + batch_size = global_batch_size_to_load // jax.process_count() neg_ones = np.full((batch_size, config.max_target_length), -1, dtype=np.int32) batch = { "inputs": neg_ones, diff --git a/src/maxtext/input_pipeline/tfds_data_processing_c4_mlperf.py b/src/maxtext/input_pipeline/tfds_data_processing_c4_mlperf.py index 991439dbb1..ba57153b7f 100644 --- a/src/maxtext/input_pipeline/tfds_data_processing_c4_mlperf.py +++ b/src/maxtext/input_pipeline/tfds_data_processing_c4_mlperf.py @@ -14,6 +14,7 @@ """DEPRECATED: Input pipeline for c4 mlperf dataset.""" +from collections.abc import Sequence import functools import numpy as np @@ -32,12 +33,14 @@ import jax import jax.numpy as jnp from jax.experimental import multihost_utils +from jax.sharding import PartitionSpec as P from maxtext.input_pipeline import multihost_dataloading from maxtext.input_pipeline.packing import sequence_packing from maxtext.input_pipeline.input_pipeline_utils import get_tokenizer from maxtext.input_pipeline.input_pipeline_utils import TokenizeOp from maxtext.utils import max_logging +from maxtext.utils.sharding import remove_size_one_mesh_axis AUTOTUNE = tf.data.experimental.AUTOTUNE @@ -167,8 +170,24 @@ def _pad_to_batch_size( ds: tf.data.Dataset, batch_size: int, num_examples: int = 0, + num_batches: int = 0, + dataloading_host_count: int = 0, ) -> tf.data.Dataset: - """Pad unevenly distributed eval data in each shard with new entries to multiples of batch size.""" + """Pad unevenly distributed eval data in each shard with new entries to multiples of batch size. + + Args: + ds: the unbatched eval dataset of the current host. + batch_size: the number of real examples in each local batch. + num_examples: the number of examples in `ds`, when it is already known. Avoids a + full scan of the dataset. + num_batches: the number of batches every data loading host has to produce. When + positive it is used as is, which keeps the hosts in sync without communication. + dataloading_host_count: the number of JAX processes loading real data. Defaults to + `jax.process_count()`. + + Returns: + The dataset padded with entries whose `targets_segmentation` is 0. + """ # local_num represents the total number of examples in eval dataset, if num_examples > 0: @@ -185,9 +204,22 @@ def _get_num_examples(ds: tf.data.Dataset) -> int: local_num = _get_num_examples(ds) local_num_batches = (local_num + batch_size - 1) // batch_size - # Find the max number of batches required across all Jax processes. - num_batches_all = multihost_utils.process_allgather(jnp.array([local_num_batches]), tiled=False) - num_batches = np.max(num_batches_all) + if dataloading_host_count <= 0: + dataloading_host_count = jax.process_count() + if num_batches <= 0: + if dataloading_host_count < jax.process_count(): + # multihost_utils.process_allgather is a collective over *all* processes, but + # only the data loading hosts run this code, so calling it would hang. Every + # loading host reads an equally sized shard, so use the local batch count. + max_logging.log( + f"Only {dataloading_host_count} of {jax.process_count()} processes load eval data, " + "skipping the cross-host batch count allgather." + ) + num_batches = local_num_batches + else: + # Find the max number of batches required across all Jax processes. + num_batches_all = multihost_utils.process_allgather(jnp.array([local_num_batches]), tiled=False) + num_batches = int(np.max(num_batches_all)) pad_num = num_batches * batch_size - local_num assert pad_num >= 0 @@ -205,6 +237,62 @@ def _add_pad(x): return ds.concatenate(pad_ds) +def get_local_rows_loading_real_data( + data_sharding, global_batch_size_to_load, global_batch_size_to_train_on, max_target_length, mesh +): + """Get the rows of this host's batch that are not decimated away by the train/eval step. + + `MultiHostDataLoadIterator` splits the array loaded by a host into + `len(mesh.local_devices)` equally sized row blocks and hands block `i` to + `mesh.local_devices[i]`. When `global_batch_size_to_train_on` is smaller than + `global_batch_size_to_load`, the train/eval step keeps only the first + `global_batch_size_to_train_on` rows of the global batch (see the decimation in + `loss_fn`), so only the local rows placed on a device that holds one of those global + rows are actually used. This returns the sorted indices of those local rows, which is + the same mechanism `input_pipeline_interface.get_process_loading_real_data` uses to + select the loading hosts. + """ + local_batch_size = global_batch_size_to_load // jax.process_count() + num_local_devices = len(mesh.local_devices) + if num_local_devices == 0 or local_batch_size % num_local_devices != 0: + # The loader cannot split the local batch evenly over the local devices either, + # so there is no row-to-device mapping to reason about. + return list(range(local_batch_size)) + rows_per_device = local_batch_size // num_local_devices + + data_sharding_pspec = remove_size_one_mesh_axis(P(*data_sharding), mesh) + sharding = jax.sharding.NamedSharding(mesh, data_sharding_pspec) + devices_indices_map = sharding.devices_indices_map((global_batch_size_to_load, max_target_length)) + batch_cutoff = global_batch_size_to_train_on + local_rows = [] + for device_position, device in enumerate(mesh.local_devices): + indices = devices_indices_map[device] + if not indices[0].stop or indices[0].stop <= batch_cutoff: + local_rows.extend(range(device_position * rows_per_device, (device_position + 1) * rows_per_device)) + return local_rows + + +def _place_real_rows_in_local_batch(batch, local_batch_size: int, real_data_rows: Sequence[int]): + """Scatter the real examples of a batch into the local rows that are evaluated. + + When `global_batch_size_to_eval_on` is smaller than `global_batch_size_to_load_eval` + the eval step keeps only the first `global_batch_size_to_eval_on` rows of the global + batch. Only some of the rows this host loads end up in that prefix, so the real + examples are placed in those rows and the remaining rows are filled with dummy + examples (all zeros, hence `targets_segmentation == 0`) which carry no loss weight. + """ + rows = tf.constant(list(real_data_rows), dtype=tf.int32) + output = {} + for key, value in batch.items(): + # The final batch may hold fewer examples than `len(real_data_rows)`; the + # remaining real rows then stay dummy rows. + indices = rows[: tf.shape(value)[0], tf.newaxis] + shape = tf.concat([[local_batch_size], tf.shape(value)[1:]], axis=0) + scattered = tf.scatter_nd(indices, value, shape) + output[key] = tf.ensure_shape(scattered, [local_batch_size] + value.shape.as_list()[1:]) + return output + + def get_dataset( dataset_name: str, split: str, @@ -214,8 +302,12 @@ def get_dataset( enable_data_shuffling: bool = False, data_shuffle_seed: int = 0, shard_in_read: bool = False, + interleave_cycle_length: int = -1, ) -> tf.data.Dataset: """Load and return a dataset of examples.""" + extra_read_config = {} + if interleave_cycle_length > 0: + extra_read_config["interleave_cycle_length"] = interleave_cycle_length if shard_in_read: # shard dataset in reading read_config = tfds.ReadConfig( @@ -224,19 +316,34 @@ def get_dataset( input_pipeline_id=dataloading_host_index, num_input_pipelines=dataloading_host_count, ), + **extra_read_config, ) ds_builder = tfds.builder(dataset_name, data_dir=data_dir) ds_builder.download_and_prepare() ds = ds_builder.as_dataset(split=split, read_config=read_config, shuffle_files=enable_data_shuffling) else: # shard dataset after reading - read_config = tfds.ReadConfig(shuffle_seed=data_shuffle_seed) + read_config = tfds.ReadConfig(shuffle_seed=data_shuffle_seed, **extra_read_config) ds_builder = tfds.builder(dataset_name, data_dir=data_dir) ds = ds_builder.as_dataset(split=split, read_config=read_config, shuffle_files=enable_data_shuffling) ds = ds.shard(num_shards=dataloading_host_count, index=dataloading_host_index) return ds +def _files_hold_one_example_each( + dataset_name: str, split: str, data_dir: str | None, dataloading_host_count: int +) -> bool: + """Whether sharding `split` by file gives every host the same examples as sharding by example. + + With `shard_in_read=True` host i opens only files i, i + n, i + 2n, ... (n loading hosts), + while `ds.shard` after reading makes host i read files 0..i to find its first example. + When every file holds exactly one example the two select the same examples in the same + order, so the eval batches are unchanged. Only the dataset metadata is read here. + """ + file_instructions = tfds.builder(dataset_name, data_dir=data_dir).info.splits[split].file_instructions + return len(file_instructions) >= dataloading_host_count and all(f.skip == 0 and f.take == 1 for f in file_instructions) + + def format_fn(x, eos_id: int = 1, pad_id: int = 0): """Format function for c4_mlperf.""" x["inputs"] = x["targets"] @@ -353,8 +460,30 @@ def preprocess_eval_dataset( num_examples: int = 0, is_tokenized_dataset: bool = True, use_stream_chunking: bool | None = None, + real_data_rows: Sequence[int] | None = None, + num_batches: int = 0, + dataloading_host_count: int = 0, ) -> tf.data.Dataset: - """Preprocess the evaluation dataset.""" + """Preprocess the evaluation dataset. + + Args: + eval_ds: the raw eval dataset shard of the current host. + sp_tokenizer: tokenizer, used for its pad and eos ids. + eval_global_batch_size_to_load: cluster-wide number of rows loaded per eval step. + max_target_length: sequence length. + num_examples: number of examples in `eval_ds`, when already known. + is_tokenized_dataset: whether `eval_ds` already holds token ids. + use_stream_chunking: whether to chunk a continuous token stream, defaults to + `is_tokenized_dataset`. + real_data_rows: rows of the per-host batch that are not decimated away by the eval + step. Defaults to every row. The other rows are filled with dummy examples. + num_batches: number of batches every data loading host has to produce. + dataloading_host_count: number of JAX processes loading real data. + + Returns: + The batched eval dataset, with one batch of `eval_global_batch_size_to_load // + jax.process_count()` rows per eval step. + """ pad_id = sp_tokenizer.pad_id eod_id = sp_tokenizer.eos_id @@ -386,10 +515,35 @@ def preprocess_eval_dataset( eval_ds = sequence_packing.pack_dataset(eval_ds, max_target_length, pad_id=pad_id) eval_ds = eval_ds.map(lambda x: format_fn(x, pad_id=pad_id), num_parallel_calls=AUTOTUNE) + # Every host contributes the same number of rows to the global batch, but only the + # rows in `real_data_rows` survive the decimation in the eval step, so only those are + # filled with real examples. + local_batch_size = eval_global_batch_size_to_load // jax.process_count() + if real_data_rows is None: + real_data_rows = range(local_batch_size) + real_data_rows = list(real_data_rows) + real_batch_size = len(real_data_rows) + # ensure array split in an equal division for each device # pad zeros up to the same batch_size among all processes - eval_ds = _pad_to_batch_size(eval_ds, eval_global_batch_size_to_load // jax.process_count(), num_examples) - eval_ds = eval_ds.batch(eval_global_batch_size_to_load // jax.process_count(), drop_remainder=False) + eval_ds = _pad_to_batch_size(eval_ds, real_batch_size, num_examples, num_batches, dataloading_host_count) + eval_ds = eval_ds.batch(real_batch_size, drop_remainder=False) + if real_data_rows != list(range(local_batch_size)): + eval_ds = eval_ds.map( + functools.partial( + _place_real_rows_in_local_batch, + local_batch_size=local_batch_size, + real_data_rows=real_data_rows, + ), + num_parallel_calls=AUTOTUNE, + ) + + # The eval loop consumes at most `num_batches` batches per eval. Truncate before the + # cache so that the first eval completes it: `cache()` only commits after a full pass + # over its input, and a partially read cache is discarded when the iterator is reset, + # which would make every eval read the eval files again. + if num_batches > 0: + eval_ds = eval_ds.take(num_batches) # We are running eval over exactly one epoch. # We explicitly cache the entire epoch (in memory) to ensure that it is the @@ -420,6 +574,8 @@ def make_c4_mlperf_train_iterator( data_dir=train_data_dir, enable_data_shuffling=config.enable_data_shuffling, data_shuffle_seed=config.data_shuffle_seed, + shard_in_read=config.train_shard_in_read, + interleave_cycle_length=config.train_interleave_cycle_length, ) train_ds = rekey(train_ds, {"inputs": None, "targets": train_col}) @@ -477,15 +633,50 @@ def make_c4_mlperf_eval_iterator( dataloading_host_count=len(process_indices), data_dir=eval_data_dir, enable_data_shuffling=False, + shard_in_read=_files_hold_one_example_each( + config.eval_dataset_name, eval_split, eval_data_dir, len(process_indices) + ), ) eval_ds = rekey(eval_ds, {"inputs": None, "targets": eval_col}) sp_tokenizer = get_tokenizer( config.tokenizer_path, config.tokenizer_type, config.add_bos, config.add_eos, config.hf_access_token ) + + # Only the first `global_batch_size_to_eval_on` rows of the loaded global batch are + # evaluated (see the decimation in `loss_fn`), so those are the only rows that need + # real examples. + global_batch_size_to_load_eval = config.global_batch_size_to_load_eval + global_batch_size_to_eval_on = getattr(config, "global_batch_size_to_eval_on", 0) or global_batch_size_to_load_eval + local_batch_size = global_batch_size_to_load_eval // jax.process_count() + real_data_rows = list(range(local_batch_size)) + if global_batch_size_to_eval_on < global_batch_size_to_load_eval: + rows = get_local_rows_loading_real_data( + config.data_sharding, + global_batch_size_to_load_eval, + global_batch_size_to_eval_on, + config.max_target_length, + global_mesh, + ) + if rows: + real_data_rows = rows + else: + # This host was selected to load real data, so it should hold at least one + # evaluated row. Fall back to filling every row rather than dropping the shard. + max_logging.log( + f"Host {jax.process_index()} loads real eval data but holds no evaluated row, " + "filling the whole local batch with real examples." + ) + max_logging.log( + f"Eval data loading host {jax.process_index()} fills {len(real_data_rows)} of " + f"{local_batch_size} local rows with real examples." + ) + if num_examples <= 0: eval_steps = getattr(config, "eval_steps", -1) - eval_batch_size = getattr(config, "global_batch_size_to_load_eval", 0) + # Each eval step only evaluates `global_batch_size_to_eval_on` examples, the rest of + # the loaded batch is decimated before the loss is computed. + eval_batch_size = global_batch_size_to_eval_on # When eval_steps > 0, derive total cluster-wide eval examples to avoid # the slow linear scan of the entire dataset during startup # (_get_num_examples). @@ -499,20 +690,29 @@ def make_c4_mlperf_eval_iterator( # calculates the correct local batch count and padding entries without # linear counting. local_num_examples = 0 + num_batches = 0 if num_examples > 0: host_idx = process_indices.index(jax.process_index()) if jax.process_index() in process_indices else 0 per_host, remainder = divmod(num_examples, len(process_indices)) local_num_examples = per_host + (1 if host_idx < remainder else 0) + # Every loading host has to produce the same number of batches. Derive it from the + # largest local share instead of an allgather, which would hang when only a subset + # of the processes loads real data. + max_local_num_examples = per_host + (1 if remainder else 0) + num_batches = (max_local_num_examples + len(real_data_rows) - 1) // len(real_data_rows) use_stream_chunking = getattr(config, "use_stream_chunking", None) eval_ds = preprocess_eval_dataset( eval_ds, sp_tokenizer=sp_tokenizer, - eval_global_batch_size_to_load=config.global_batch_size_to_load_eval, + eval_global_batch_size_to_load=global_batch_size_to_load_eval, max_target_length=config.max_target_length, num_examples=local_num_examples, is_tokenized_dataset=is_tokenized_dataset, use_stream_chunking=use_stream_chunking, + real_data_rows=real_data_rows, + num_batches=num_batches, + dataloading_host_count=len(process_indices), ) eval_multihost_gen = multihost_dataloading.MultiHostDataLoadIterator(eval_ds, global_mesh) diff --git a/src/maxtext/integration/tunix/diffusion_sft.py b/src/maxtext/integration/tunix/diffusion_sft.py new file mode 100644 index 0000000000..c53f8ab055 --- /dev/null +++ b/src/maxtext/integration/tunix/diffusion_sft.py @@ -0,0 +1,182 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://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. + +"""MaxText adapters for target-aligned Tunix diffusion SFT.""" + +from collections.abc import Mapping +from typing import Any + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion.block_diffusion import target_alignment +from tunix.diffusion import types as diffusion_types + + +_REQUIRED_FIELDS = ( + "inputs", + "inputs_position", + "inputs_segmentation", + "targets", + "targets_position", + "targets_segmentation", + "completion_mask", + "corruption_mask", + "targets_loss_mask", +) + + +def _concrete_numpy(value): + """Returns a NumPy view only for values already resident on the host.""" + if isinstance(value, (jax.Array, jax.core.Tracer, jax.ShapeDtypeStruct)): + return None + return np.asarray(value) + + +def _validate_batch_masks( + positions, + targets_segmentation, + completion_mask, + corruption_mask, + loss_weights, + *, + alignment, + block_size, + completion_only, +): + """Validates the prepared block-diffusion supervision contract.""" + shapes = { + "positions": tuple(positions.shape), + "targets_segmentation": tuple(targets_segmentation.shape), + "completion_mask": tuple(completion_mask.shape), + "corruption_mask": tuple(corruption_mask.shape), + "targets_loss_mask": tuple(loss_weights.shape), + } + if len(set(shapes.values())) != 1: + raise ValueError(f"diffusion SFT masks must have identical shapes; received {shapes}") + concrete = [ + _concrete_numpy(value) + for value in ( + positions, + targets_segmentation, + completion_mask, + corruption_mask, + loss_weights, + ) + ] + if any(value is None for value in concrete): + return + concrete_positions = np.asarray(concrete[0], dtype=np.int32) + validity = np.asarray(concrete[1]) != 0 + completion, corruption = (np.asarray(value, dtype=bool) for value in concrete[2:4]) + weights = np.asarray(concrete[4], dtype=np.float32) + weighted = weights != 0 + if not np.all(np.isfinite(weights)) or np.any(weights < 0): + raise ValueError("targets_loss_mask must contain finite nonnegative weights") + if np.any(completion & ~validity): + raise ValueError("completion_mask must be a subset of valid target positions") + if completion_only: + block_ids = concrete_positions // block_size + unseen_block = np.iinfo(block_ids.dtype).min + seen_completion_blocks = np.maximum.accumulate(np.where(validity & completion, block_ids, unseen_block), axis=-1) + if np.any(validity & ~completion & (seen_completion_blocks == block_ids)): + raise ValueError( + "completion-only block diffusion cannot place a prompt token after " + "a completion token in the same diffusion block; split the example " + "or align the next prompt to a block boundary" + ) + supervision = completion if completion_only else validity + if np.any(corruption & ~supervision): + raise ValueError("corruption_mask must be a subset of the configured supervision scope") + if np.any(weighted & ~supervision): + raise ValueError("targets_loss_mask must be a subset of the configured supervision scope") + allowed = corruption.copy() + if alignment == "shifted": + allowed |= (concrete_positions > 0) & (concrete_positions % block_size == 0) + if np.any(weighted & ~allowed): + raise ValueError("diffusion SFT loss weights must own corrupted targets or shifted block anchors") + + +def create_batch_adapter(config): + """Builds a raw MaxText batch adapter for Tunix diffusion SFT.""" + alignment = config.block_diffusion_logit_alignment + block_size = int(config.causal_block_size) + completion_only = bool(config.sft_train_on_completion_only) + + def adapt(raw_batch: Mapping[str, Any]) -> diffusion_types.DiffusionTokenBatch: + unavailable = sorted(name for name in _REQUIRED_FIELDS if name not in raw_batch or raw_batch[name] is None) + if unavailable: + raise ValueError(f"block-diffusion SFT requires non-None batch fields; unavailable {unavailable}") + _validate_batch_masks( + raw_batch["targets_position"], + raw_batch["targets_segmentation"], + raw_batch["completion_mask"], + raw_batch["corruption_mask"], + raw_batch["targets_loss_mask"], + alignment=alignment, + block_size=block_size, + completion_only=completion_only, + ) + targets = jnp.asarray(raw_batch["targets"]) + positions = jnp.asarray(raw_batch["targets_position"], dtype=jnp.int32) + validity_mask = jnp.asarray(raw_batch["targets_segmentation"]) != 0 + completion_mask = jnp.asarray(raw_batch["completion_mask"], dtype=jnp.bool_) + corruption_mask = jnp.asarray(raw_batch["corruption_mask"], dtype=jnp.bool_) + raw_loss_weights = jnp.asarray(raw_batch["targets_loss_mask"], dtype=jnp.float32) + supervision_mask = completion_mask if completion_only else validity_mask + allowed = corruption_mask + if alignment == "shifted": + allowed |= (positions > 0) & (positions % block_size == 0) + loss_weights = jnp.where(validity_mask & supervision_mask & allowed, raw_loss_weights, 0.0) + return diffusion_types.DiffusionTokenBatch.create( + model_inputs={ + "input_tokens": jnp.asarray(raw_batch["inputs"]), + "input_positions": jnp.asarray(raw_batch["inputs_position"], dtype=jnp.int32), + "input_segmentation": jnp.asarray(raw_batch["inputs_segmentation"]), + "targets": targets, + "target_positions": positions, + "target_segmentation": jnp.asarray(raw_batch["targets_segmentation"]), + }, + target_ids=targets, + loss_weights=loss_weights, + ) + + return adapt + + +def create_target_aligned_logits_fn(config): + """Builds a MaxText scorer satisfying Tunix's diffusion logits contract.""" + alignment = config.block_diffusion_logit_alignment + enable_dropout = bool(config.enable_dropout) + + def logits_fn(model: nnx.Module, model_inputs: diffusion_types.ModelInputs): + base_model = getattr(model, "base", model) + logits = base_model( + decoder_input_tokens=model_inputs["input_tokens"], + decoder_positions=model_inputs["input_positions"], + decoder_segment_ids=model_inputs["input_segmentation"], + enable_dropout=enable_dropout, + decoder_target_tokens=model_inputs["targets"], + decoder_target_mask=model_inputs["target_segmentation"], + ) + return target_alignment.align_logits_to_targets( + logits, + alignment, + model_inputs["target_positions"], + model_inputs["target_segmentation"] != 0, + ) + + return logits_fn diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py index 2a95fb46ab..9a62537a21 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py @@ -25,6 +25,7 @@ import ast import contextlib import copy +import inspect import json import logging import re @@ -487,6 +488,15 @@ def unroll_gemma_scanned_weights(weights): return unflatten_dict(new_flat_w) +def _accepts_kwarg(cls: Any, name: str) -> bool: + """Whether `cls(...)` takes the keyword `name` (or `**kwargs`); True when its signature cannot be inspected.""" + try: + params = inspect.signature(cls).parameters + except (ValueError, TypeError): + return True + return name in params or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) + + def _log_and_flush_traceback(msg: str) -> None: """Logs an error with formatted traceback and flushes all logging handlers.""" logging.error("%s:\n%s", msg, traceback.format_exc()) @@ -625,6 +635,19 @@ def __init__( ) self._maxtext_config = maxtext_config + # Forwarded only when the installed tunix declares the option on both + # configs; an older tunix always frees the KV cache around a weight sync. + optional_vllm_config = {} + if hasattr(rollout_config, "rollout_vllm_free_kv_cache_during_weight_sync"): + if _accepts_kwarg(VllmConfig, "free_kv_cache_during_weight_sync"): + optional_vllm_config["free_kv_cache_during_weight_sync"] = ( + rollout_config.rollout_vllm_free_kv_cache_during_weight_sync + ) + else: + logging.warning( + "Installed tunix VllmConfig does not support free_kv_cache_during_weight_sync; the setting is ignored." + ) + self._sampler = MaxTextVllmSampler( tokenizer=tokenizer, config=VllmConfig( # pylint: disable=unexpected-keyword-arg,no-value-for-parameter @@ -647,6 +670,7 @@ def __init__( engine_kwargs=engine_kwargs, additional_config=rollout_additional_config, sampling_kwargs=rollout_config.rollout_vllm_sampling_kwargs, + **optional_vllm_config, ), direct_maxtext_sync=direct_maxtext_sync, model_name=getattr(maxtext_config, "model_name", ""), diff --git a/src/maxtext/kernels/attention/jax_flash_attention.py b/src/maxtext/kernels/attention/jax_flash_attention.py index 0c4e100ea5..b626549931 100644 --- a/src/maxtext/kernels/attention/jax_flash_attention.py +++ b/src/maxtext/kernels/attention/jax_flash_attention.py @@ -13,109 +13,60 @@ # limitations under the License. """JAX implementation without using Pallas for Flash Attention.""" -from typing import Optional, Tuple, Union +from __future__ import annotations + +import functools +from typing import Any, Callable, Dict, Optional, Tuple, Union import jax +from jax import lax from jax.experimental import layout from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as mask_lib from jax.experimental.xla_metadata import must_fuse_call import jax.numpy as jnp -from maxtext.kernels.attention import splash_attention_kernel - +from maxtext.kernels.attention.splash_attention_kernel import SegmentIds DLL = layout.Layout Layout = layout.Format -SegmentIds = splash_attention_kernel.SegmentIds +_LONG_SEQUENCE_THRESHOLD = 4096 -# This function computes masked flash attention using a block-sparse approach. -# This implementation keeps the full batch and number of heads dimensions -# throughout the attention computation while iterating through blocks of the -# key/value sequence and, within each, iterates through blocks of the query -# sequence. The `mask_blocked` is used to skip computations for blocks where all -# attention scores are masked out, improving efficiency for sparse masks. -def flash_attention_block_masked( - q: jnp.ndarray, - k: jnp.ndarray, - v: jnp.ndarray, - segment_ids: SegmentIds | None, - block_kv: int, - block_q: int, - mask: mask_lib.Mask | jax.Array, - mask_value: float, - cap: Optional[float] = None, - save_residuals: bool = False, - logits_dtype: jax.typing.DTypeLike = jnp.bfloat16, - loop_unroll: int | bool = True, - fuse_logits: bool = True, -) -> Union[jnp.ndarray, Tuple[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]]: - """Computes masked flash attention using block-sparse masking. +def mask_blocker(mask: jnp.ndarray, block_q: int, block_kv: int) -> jnp.ndarray: + """Creates a blocked mask from a full mask. Args: - q: Query tensor with shape (batch_size, num_kv_heads, - num_q_heads_per_kv_head, q_seq_len, head_dim). - k: Key tensor with shape (batch_size, num_kv_heads, kv_seq_len, head_dim). - v: Value tensor with shape (batch_size, num_kv_heads, kv_seq_len, - v_head_dim). - segment_ids: SegmentIds are a mechanism to ensure that there is no - cross-attention between segments (fraction of a sequence) that have been - concatenated together into a sequence. Each array is a list of ids - (integers). Only tokens with the same id are allowed to attend to each - other. It stores the segment ids of the query and key/value sequences. - block_kv: Block size for the key/value sequence dimension. + mask: The attention mask with shape of (batch_size, q_seq_len, kv_seq_len). block_q: Block size for the query sequence dimension. - mask: The full attention mask. A rank-2 mask with shape - (q_seq_len, kv_seq_len) is shared by every batch item. A rank-3 mask with - shape (batch_size, q_seq_len, kv_seq_len) supplies an independent mask - for each batch item. - mask_value: The value to use for masked-out attention scores. - cap: Optional cap for attention logits. This helps to prevent extremely - large logits: capped_logits = jnp.tanh(logits / attn_logits_soft_cap) * - attn_logits_soft_cap - logits_dtype: Preferred element type for the fused query-key product that - forms the attention probabilities; pass the query dtype when - full-precision logits are required. - loop_unroll: Unroll setting for the block loops; True fully unrolls, which - is prohibitive for long sequences with a runtime mask. - fuse_logits: Wrap the logits computation in a must-fuse XLA metadata call. - The wrapper does not support reverse-mode differentiation inside a - sharded computation; pass False when gradients are required. - save_residuals: Whether to save residuals. If True, returns a tuple of - (output, dict=(logsumexp, max_logits)). Both `logsumexp` and `max_logits` - are of shape (batch_size, num_kv_heads, num_q_heads // num_kv_heads, - q_seq_len). + block_kv: Block size for the key/value sequence dimension. Returns: - If save_residuals is True, returns a tuple containing: - - The output of the attention computation. - - A dict of (logsumexp, max_logits) - Otherwise, returns the output of the attention computation. + A blocked mask where each element indicates the number of non-zero + elements in the corresponding block of the original mask. """ - batch_size, num_q_heads, q_seq_len, qk_head_dim_size = q.shape - _, num_kv_heads, kv_seq_len, _ = k.shape - v_head_dim_size = v.shape[-1] - data_type = q.dtype - q_groups = num_q_heads // num_kv_heads - q = q.reshape( - ( - batch_size, - num_kv_heads, - q_groups, - q_seq_len, - qk_head_dim_size, - ) - ) + batch_size, q_seq_len, kv_seq_len = mask.shape if q_seq_len % block_q != 0: raise ValueError(f"q_seq_len {q_seq_len} must be divisible by block_q {block_q}") if kv_seq_len % block_kv != 0: raise ValueError(f"kv_seq_len {kv_seq_len} must be divisible by block_kv {block_kv}") + q_blocks = q_seq_len // block_q + kv_blocks = kv_seq_len // block_kv + + blocked_mask = mask.reshape(batch_size, q_blocks, block_q, kv_blocks, block_kv) + return jnp.count_nonzero(blocked_mask, axis=(2, 4)).astype(jnp.int32) - # Calculate the number of key/value and query blocks. - num_kv_blocks = kv_seq_len // block_kv - num_q_blocks = q_seq_len // block_q +def _prepare_mask_and_block_predicate( + mask: mask_lib.Mask | jax.Array, + batch_size: int, + q_seq_len: int, + kv_seq_len: int, + block_q: int, + block_kv: int, + segment_ids: SegmentIds | None, +) -> Tuple[jax.Array, Callable[[int, int], Any]]: + """Extracts the full mask and constructs the block computation predicate.""" mask_array = mask[:, :] if mask_array.ndim == 2: # A rank-2 mask is shared by every item in the batch. @@ -123,13 +74,13 @@ def flash_attention_block_masked( else: if mask_array.shape != (batch_size, q_seq_len, kv_seq_len): raise ValueError( - f"Batched attention mask must have shape {(batch_size, q_seq_len, kv_seq_len)}, got {mask_array.shape}." + "Batched attention mask must have shape" f" {(batch_size, q_seq_len, kv_seq_len)}, got {mask_array.shape}." ) - mask_full = mask_array + mask_full = jnp.asarray(mask_array) if segment_ids is not None: - segment_ids_q = segment_ids.q[:, :, None] - segment_ids_kv = segment_ids.kv[:, None, :] + segment_ids_q = segment_ids.q[:, :, None] if segment_ids.q.ndim == 2 else segment_ids.q[None, :, None] + segment_ids_kv = segment_ids.kv[:, None, :] if segment_ids.kv.ndim == 2 else segment_ids.kv[None, None, :] mask_full = jnp.logical_and(mask_full, segment_ids_q == segment_ids_kv) if isinstance(mask, mask_lib.CausalMask): @@ -141,19 +92,76 @@ def flash_attention_block_masked( # Min kv_idx in block j: j * block_kv # Condition: (i + 1) * block_q - 1 >= j * block_kv # Which simplifies to: (i + 1) * block_q > j * block_kv - def should_compute_block(i, j): + def should_compute_block(i: int, j: int) -> bool | jax.Array: return (i + 1) * block_q > j * block_kv else: mask_blocked = jax.jit(mask_blocker, static_argnums=[1, 2])(mask_full, block_q, block_kv) - def should_compute_block(i, j): + def should_compute_block(i: int, j: int) -> bool | jax.Array: mask_i_j_slice = jax.lax.dynamic_slice(mask_blocked, (0, i, j), (batch_size, 1, 1)) # The compute_attention_block should be executed if at least one element # in the slice is non-zero, meaning at least one batch requires work for # this block. return jnp.any(jnp.not_equal(mask_i_j_slice, 0)) + return mask_full, should_compute_block + + +def _compute_residual_stats( + l: jnp.ndarray, + m: jnp.ndarray, +) -> Dict[str, jnp.ndarray]: + """Computes logsumexp and max_logits residual statistics with stopped gradients.""" + stats = {"logsumexp": m + jnp.log(l), "max_logits": m} + return jax.tree.map(jax.lax.stop_gradient, stats) + + +def _flash_attention_block_masked_short_seq( + q: jnp.ndarray, + k: jnp.ndarray, + v: jnp.ndarray, + segment_ids: SegmentIds | None, + block_kv: int, + block_q: int, + mask: mask_lib.Mask | jax.Array, + mask_value: float, + cap: Optional[float] = None, + save_residuals: bool = False, + logits_dtype: jax.typing.DTypeLike = jnp.bfloat16, + loop_unroll: int | bool = True, + fuse_logits: bool = True, +) -> Union[jnp.ndarray, Tuple[jnp.ndarray, Dict[str, jnp.ndarray]]]: + """Computes masked flash attention for short sequences (q_seq_len <= 4096).""" + batch_size, num_q_heads, q_seq_len, qk_head_dim_size = q.shape + _, num_kv_heads, kv_seq_len, _ = k.shape + v_head_dim_size = v.shape[-1] + data_type = q.dtype + q_groups = num_q_heads // num_kv_heads + q = q.reshape( + ( + batch_size, + num_kv_heads, + q_groups, + q_seq_len, + qk_head_dim_size, + ) + ) + + # Calculate the number of key/value and query blocks. + num_kv_blocks = kv_seq_len // block_kv + num_q_blocks = q_seq_len // block_q + + mask_full, should_compute_block = _prepare_mask_and_block_predicate( + mask=mask, + batch_size=batch_size, + q_seq_len=q_seq_len, + kv_seq_len=kv_seq_len, + block_q=block_q, + block_kv=block_kv, + segment_ids=segment_ids, + ) + # Initialize `l` (logsumexp) and `m` (max_logits) for the online softmax. # `l` is initialized to 0 since no blocks have been processed yet and the sum # is 0. @@ -287,42 +295,1244 @@ def identity(output, l, m): output, l, m = jax.lax.fori_loop(0, num_kv_blocks, outer_loop_body, (output, l, m), unroll=loop_unroll) - # Reshape the output to drop the size one dimension at index 2, - # which corresponds to `num_q_heads // num_kv_heads` when - # num_q_heads == num_kv_heads. - output = output.squeeze(axis=2) + # Fold q_groups (dimension 2) back into the query heads dimension correctly + # for both MHA (q_groups=1) and GQA (q_groups > 1). + output = output.reshape((batch_size, num_q_heads, q_seq_len, v_head_dim_size)) if not save_residuals: # To avoid remat of the output, we can use context=hbm remat policy as in # maxtext/configs/types.py return output - l = l.squeeze(axis=2) - m = m.squeeze(axis=2) - stats = {"logsumexp": m + jnp.log(l), "max_logits": m} - stats = jax.tree.map(jax.lax.stop_gradient, stats) - return output, stats + l = l.reshape((batch_size, num_q_heads, q_seq_len)) + m = m.reshape((batch_size, num_q_heads, q_seq_len)) + return output, _compute_residual_stats(l, m) -def mask_blocker(mask: jnp.ndarray, block_q: int, block_kv: int) -> jnp.ndarray: - """Creates a blocked mask from a full mask. +def _flash_attention_block_masked_long_seq( + q: jnp.ndarray, + k: jnp.ndarray, + v: jnp.ndarray, + segment_ids: SegmentIds | None, + block_kv: int, + block_q: int, + mask: mask_lib.Mask | jax.Array, + mask_value: float, + cap: Optional[float] = None, + save_residuals: bool = False, + logits_dtype: jax.typing.DTypeLike = jnp.bfloat16, + loop_unroll: int | bool = True, + fuse_logits: bool = True, +) -> Union[jnp.ndarray, Tuple[jnp.ndarray, Dict[str, jnp.ndarray]]]: + """Computes masked flash attention for long sequences (q_seq_len > 4096).""" + batch_size, num_q_heads, q_seq_len, qk_head_dim_size = q.shape + _, num_kv_heads, kv_seq_len, _ = k.shape + v_head_dim_size = v.shape[-1] + data_type = q.dtype + q_groups = num_q_heads // num_kv_heads + q = q.reshape( + ( + batch_size, + num_kv_heads, + q_groups, + q_seq_len, + qk_head_dim_size, + ) + ) + + # Calculate the number of key/value and query blocks. + num_kv_blocks = kv_seq_len // block_kv + num_q_blocks = q_seq_len // block_q + + mask_full, should_compute_block = _prepare_mask_and_block_predicate( + mask=mask, + batch_size=batch_size, + q_seq_len=q_seq_len, + kv_seq_len=kv_seq_len, + block_q=block_q, + block_kv=block_kv, + segment_ids=segment_ids, + ) + + l = jnp.zeros((batch_size, num_kv_heads, q_groups, q_seq_len), dtype=data_type) + m = jnp.full( + (batch_size, num_kv_heads, q_groups, q_seq_len), + mask_value, + dtype=data_type, + ) + output = jnp.zeros( + ( + batch_size, + num_kv_heads, + q_groups, + q_seq_len, + v_head_dim_size, + ), + dtype=data_type, + ) + + # Augment V with a column of ones before the loop so that l_i_j = p_i_j.sum(-1) + # is computed as the last column of the P @ V MXU contraction rather than a + # separate reduce instruction that cannot fuse with DUS. With this we get rid + # of one time-consuming loop fusion on long sequence lengths. + v_ones = jnp.ones(v.shape[:-1] + (1,), dtype=v.dtype) + v_to_slice = jnp.concatenate([v, v_ones], axis=-1) + + # Outer loop over the key/value blocks. + def outer_loop_body(j, carried): + output, l, m = carried + + # Inner loop over the query blocks. + def inner_loop_body(i, carried_inner): + output, l, m = carried_inner + + def compute_attention_block(output, l, m): + output_i_slice = jax.lax.dynamic_slice_in_dim(output, i * block_q, block_q, axis=-2) + l_i_slice = jax.lax.dynamic_slice_in_dim(l, i * block_q, block_q, axis=-1) + m_i_slice = jax.lax.dynamic_slice_in_dim(m, i * block_q, block_q, axis=-1) + full_mask_i_j_slice = jax.lax.dynamic_slice( + mask_full, + (0, i * block_q, j * block_kv), + (batch_size, block_q, block_kv), + ) + broadcasted_mask = jnp.broadcast_to( + full_mask_i_j_slice[:, None, None, :, :], + (batch_size, num_kv_heads, q_groups, block_q, block_kv), + ) + + k_j_slice = jax.lax.dynamic_slice_in_dim(k, j * block_kv, block_kv, axis=-2) + v_j_slice = jax.lax.dynamic_slice_in_dim(v_to_slice, j * block_kv, block_kv, axis=-2) + q_i_slice = jax.lax.dynamic_slice_in_dim(q, i * block_q, block_q, axis=-2) + + def fusion_qk(q_i_slice, k_j_slice, broadcasted_mask): + s_i_j = jnp.einsum( + "bxhqc,bxkc->bxhqk", + q_i_slice, + k_j_slice, + preferred_element_type=logits_dtype, + ) + if cap is not None: + s_i_j = jnp.tanh(s_i_j / cap) + s_i_j = s_i_j * cap + s_i_j = jnp.where(broadcasted_mask, s_i_j, mask_value) + m_i_j = s_i_j.max(axis=-1) + return s_i_j, m_i_j + + if fuse_logits: + s_i_j, m_i_j = must_fuse_call("qk")(fusion_qk)(q_i_slice, k_j_slice, broadcasted_mask) + else: + s_i_j, m_i_j = fusion_qk(q_i_slice, k_j_slice, broadcasted_mask) + assert m_i_j.shape == m_i_slice.shape + + def fusion_pv( + s_i_j, + v_j_slice, + l_i_slice, + m_i_slice, + output_i_slice, + m_i_j, + ): + p_i_j = jnp.exp(s_i_j - m_i_j[..., None]) + m_i_new = jnp.maximum(m_i_slice, m_i_j) + m_i_difference = jnp.exp(m_i_slice - m_i_new) + m_i_j_difference = jnp.exp(m_i_j - m_i_new) + + pv_aug = jnp.einsum( + "bxhqk,bxkc->bxhqc", + p_i_j, + v_j_slice, + preferred_element_type=data_type, + ) + # This forces the layout of the final @V to have better utilization + # by avoiding using XLU: + # Changes conv emitter type from + # EmitAllInputFeatureInSublanesOutputBatchInSublanesXposeReuse to + # EmitInputBatchInLanes + pv_aug = layout.with_layout_constraint(pv_aug, DLL(major_to_minor=(0, 1, 2, 4, 3))) + pv = pv_aug[..., :-1] + l_i_j = pv_aug[..., -1] + + l_i_new = m_i_difference * l_i_slice + m_i_j_difference * l_i_j + output_i_slice_new = m_i_difference[..., None] * output_i_slice + m_i_j_difference[..., None] * pv + return output_i_slice_new, l_i_new, m_i_new + + if fuse_logits: + output_i_slice_new, l_i_new, m_i_new = must_fuse_call("pv")(fusion_pv)( + s_i_j, v_j_slice, l_i_slice, m_i_slice, output_i_slice, m_i_j + ) + else: + output_i_slice_new, l_i_new, m_i_new = fusion_pv(s_i_j, v_j_slice, l_i_slice, m_i_slice, output_i_slice, m_i_j) + + output = jax.lax.dynamic_update_index_in_dim(output, output_i_slice_new, i * block_q, axis=-2) + l = jax.lax.dynamic_update_index_in_dim(l, l_i_new, i * block_q, axis=-1) + m = jax.lax.dynamic_update_index_in_dim(m, m_i_new, i * block_q, axis=-1) + return output, l, m + + def identity(output, l, m): + """A no-op identity function.""" + return output, l, m + + output, l, m = jax.lax.cond( + should_compute_block(i, j), + compute_attention_block, + identity, + output, + l, + m, + ) + return output, l, m + + output, l, m = jax.lax.fori_loop(0, num_q_blocks, inner_loop_body, (output, l, m), unroll=loop_unroll) + return (output, l, m) + + output, l, m = jax.lax.fori_loop(0, num_kv_blocks, outer_loop_body, (output, l, m), unroll=loop_unroll) + + # Fold q_groups (dimension 2) back into the query heads dimension correctly + # for both MHA (q_groups=1) and GQA (q_groups > 1). + output = output.reshape((batch_size, num_q_heads, q_seq_len, v_head_dim_size)) + l = l.reshape((batch_size, num_q_heads, q_seq_len)) + output = output / l[..., None] + if not save_residuals: + # To avoid remat of the output, we can use context=hbm remat policy as in + # maxtext/configs/types.py + return output + + m = m.reshape((batch_size, num_q_heads, q_seq_len)) + return output, _compute_residual_stats(l, m) + + +# This function computes masked flash attention using a block-sparse approach. +# This implementation keeps the full batch and number of heads dimensions +# throughout the attention computation while iterating through blocks of the +# key/value sequence and, within each, iterates through blocks of the query +# sequence. The `mask_blocked` is used to skip computations for blocks where all +# attention scores are masked out, improving efficiency for sparse masks. +def flash_attention_block_masked( + q: jnp.ndarray, + k: jnp.ndarray, + v: jnp.ndarray, + segment_ids: SegmentIds | None, + block_kv: int, + block_q: int, + mask: mask_lib.Mask | jax.Array, + mask_value: float, + cap: Optional[float] = None, + save_residuals: bool = False, + logits_dtype: jax.typing.DTypeLike = jnp.bfloat16, + loop_unroll: int | bool = True, + fuse_logits: bool = True, +) -> Union[jnp.ndarray, Tuple[jnp.ndarray, Dict[str, jnp.ndarray]]]: + """Computes masked flash attention using block-sparse masking. + + Dispatches dynamically to the short-sequence implementation (q_seq_len <= + 4096) + or the long-sequence implementation (q_seq_len > 4096). Args: - mask: The attention mask with shape of (batch_size, q_seq_len, kv_seq_len). - block_q: Block size for the query sequence dimension. + q: Query tensor with shape (batch_size, num_kv_heads, + num_q_heads_per_kv_head, q_seq_len, head_dim). + k: Key tensor with shape (batch_size, num_kv_heads, kv_seq_len, head_dim). + v: Value tensor with shape (batch_size, num_kv_heads, kv_seq_len, + v_head_dim). + segment_ids: SegmentIds are a mechanism to ensure that there is no + cross-attention between segments (fraction of a sequence) that have been + concatenated together into a sequence. Each array is a list of ids + (integers). Only tokens with the same id are allowed to attend to each + other. It stores the segment ids of the query and key/value sequences. block_kv: Block size for the key/value sequence dimension. + block_q: Block size for the query sequence dimension. + mask: The full attention mask. A rank-2 mask with shape (q_seq_len, + kv_seq_len) is shared by every batch item. A rank-3 mask with shape + (batch_size, q_seq_len, kv_seq_len) supplies an independent mask for each + batch item. + mask_value: The value to use for masked-out attention scores. + cap: Optional cap for attention logits. This helps to prevent extremely + large logits: capped_logits = jnp.tanh(logits / attn_logits_soft_cap) * + attn_logits_soft_cap + save_residuals: Whether to save residuals. If True, returns a tuple of + (output, dict=(logsumexp, max_logits)). Both `logsumexp` and `max_logits` + are of shape (batch_size, num_kv_heads, num_q_heads // num_kv_heads, + q_seq_len). + logits_dtype: Preferred element type for the fused query-key product that + forms the attention probabilities; pass the query dtype when + full-precision logits are required. + loop_unroll: Unroll setting for the block loops; True fully unrolls, which + is prohibitive for long sequences with a runtime mask. + fuse_logits: Wrap the logits computation in a must-fuse XLA metadata call. + The wrapper does not support reverse-mode differentiation inside a sharded + computation; pass False when gradients are required. Returns: - A blocked mask where each element indicates the number of non-zero - elements in the corresponding block of the original mask. + If save_residuals is True, returns a tuple containing: + - The output of the attention computation. + - A dict of (logsumexp, max_logits) + Otherwise, returns the output of the attention computation. """ - batch_size, q_seq_len, kv_seq_len = mask.shape + q_seq_len = q.shape[-2] + kv_seq_len = k.shape[-2] if q_seq_len % block_q != 0: raise ValueError(f"q_seq_len {q_seq_len} must be divisible by block_q {block_q}") if kv_seq_len % block_kv != 0: raise ValueError(f"kv_seq_len {kv_seq_len} must be divisible by block_kv {block_kv}") - q_blocks = q_seq_len // block_q - kv_blocks = kv_seq_len // block_kv - blocked_mask = mask.reshape(batch_size, q_blocks, block_q, kv_blocks, block_kv) - return jnp.count_nonzero(blocked_mask, axis=(2, 4)).astype(jnp.int32) + if q_seq_len <= _LONG_SEQUENCE_THRESHOLD: + return _flash_attention_block_masked_short_seq( + q=q, + k=k, + v=v, + segment_ids=segment_ids, + block_kv=block_kv, + block_q=block_q, + mask=mask, + mask_value=mask_value, + cap=cap, + save_residuals=save_residuals, + logits_dtype=logits_dtype, + loop_unroll=loop_unroll, + fuse_logits=fuse_logits, + ) + return _flash_attention_block_masked_long_seq( + q=q, + k=k, + v=v, + segment_ids=segment_ids, + block_kv=block_kv, + block_q=block_q, + mask=mask, + mask_value=mask_value, + cap=cap, + save_residuals=save_residuals, + logits_dtype=logits_dtype, + loop_unroll=loop_unroll, + fuse_logits=fuse_logits, + ) + + +def _apply_mask_and_soft_cap_bwd( + qk: jax.Array, + mask: jax.Array | None, + q_segment_ids: jax.Array | None, + kv_segment_ids: jax.Array | None, + mask_value: float, + attn_logits_soft_cap: float | None, +) -> tuple[jax.Array, jax.Array]: + """Applies mask and soft cap to qk logits.""" + qk_uncapped = qk + if attn_logits_soft_cap is not None: + qk = jnp.tanh(qk_uncapped / attn_logits_soft_cap) + qk = qk * attn_logits_soft_cap + + if q_segment_ids is not None and kv_segment_ids is not None: + seg_mask = q_segment_ids[..., :, None] == kv_segment_ids[..., None, :] + if mask is not None: + mask = jnp.logical_and(mask, seg_mask) + else: + mask = seg_mask + + if mask is not None: + qk = jnp.where(mask, qk, mask_value) + + return qk, qk_uncapped + + +@functools.partial( + jax.jit, + static_argnames=[ + "block_q", + "block_kv", + "block_head", + "dtype", + "is_causal", + "use_log2", + "fuse", + "unroll_loops", + ], +) +def _flash_attention_mha_bwd_jax( + q: jax.Array, + k: jax.Array, + v: jax.Array, + o: jax.Array, + do: jax.Array, + logsumexp: jax.Array, + segment_ids: SegmentIds | None = None, + mask: jax.Array | None = None, + mask_value: float = -1e30, + attn_logits_soft_cap: float | None = None, + block_q: int = 1024, + block_kv: int = 1024, + block_head: int = 1, + dtype: jax.typing.DTypeLike = jnp.float32, + is_causal: bool = False, + use_log2: bool = False, + di: jax.Array | None = None, + fuse: bool = True, + unroll_loops: bool = False, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Memory efficient JAX implementation of Flash Attention backward pass for MHA.""" + dp_dup_must_fuse = False + dp_dedup_must_fuse = fuse + q_seq_len = q.shape[-2] + kv_seq_len = k.shape[-2] + + if di is None: + di_arr = jnp.sum(o * do, axis=-1) + else: + di_arr = di + + def should_compute_block(i, j): + return not is_causal or (j + 1) * block_q > i * block_kv + + num_q_heads = q.shape[-3] + eff_block_head = min(block_head, num_q_heads) + if num_q_heads % eff_block_head != 0: + raise ValueError(f"num_q_heads {num_q_heads} must be divisible by block_head" f" {block_head}") + if q_seq_len % block_q != 0: + raise ValueError(f"q_seq_len {q_seq_len} must be divisible by block_q {block_q}") + if kv_seq_len % block_kv != 0: + raise ValueError(f"kv_seq_len {kv_seq_len} must be divisible by block_kv {block_kv}") + + slice_sizes_k = k.shape[:-3] + (eff_block_head, block_kv, k.shape[-1]) + slice_sizes_v = v.shape[:-3] + (eff_block_head, block_kv, v.shape[-1]) + slice_sizes_q = q.shape[:-3] + (eff_block_head, block_q, q.shape[-1]) + slice_sizes_do = do.shape[:-3] + (eff_block_head, block_q, do.shape[-1]) + slice_sizes_log = logsumexp.shape[:-2] + (eff_block_head, block_q) + + if mask is None: + mask_arr = None + elif isinstance(mask, mask_lib.Mask): + mask_arr = mask[:, :] + else: + mask_arr = mask + + def head_kernel(head_idx, head_carry): + head_start = head_idx * eff_block_head + + def k_block_kernel(i, carry): + dq_carry, dk_carry, dv_carry = carry + kv_slice_start = i * block_kv + + start_indices_k = (0,) * (k.ndim - 3) + (head_start, kv_slice_start, 0) + start_indices_v = (0,) * (v.ndim - 3) + (head_start, kv_slice_start, 0) + + k_block = lax.dynamic_slice(k, start_indices_k, slice_sizes_k) + v_block = lax.dynamic_slice(v, start_indices_v, slice_sizes_v) + kv_segment_ids_block = ( + lax.dynamic_slice_in_dim(segment_ids.kv, kv_slice_start, block_kv, axis=-1) if segment_ids is not None else None + ) + + k_block_prec = k_block.astype(dtype) + v_block_prec = v_block.astype(dtype) + + dk_block = jnp.zeros_like(k_block_prec) + dv_block = jnp.zeros_like(v_block_prec) + + def q_block_kernel(j, inner_carry): + dk_block_carry, dv_block_carry, dq_carry_inner = inner_carry + q_slice_start = j * block_q + + start_indices_q = (0,) * (q.ndim - 3) + (head_start, q_slice_start, 0) + start_indices_do = (0,) * (do.ndim - 3) + (head_start, q_slice_start, 0) + start_indices_log = (0,) * (logsumexp.ndim - 2) + ( + head_start, + q_slice_start, + ) + + q_block = lax.dynamic_slice(q, start_indices_q, slice_sizes_q) + do_block = lax.dynamic_slice(do, start_indices_do, slice_sizes_do) + logsumexp_block = lax.dynamic_slice(logsumexp, start_indices_log, slice_sizes_log).astype(dtype) + di_block = lax.dynamic_slice(di_arr, start_indices_log, slice_sizes_log) + q_segment_ids_block = ( + lax.dynamic_slice_in_dim(segment_ids.q, q_slice_start, block_q, axis=-1) if segment_ids is not None else None + ) + + if mask_arr is None: + mask_block_qkv = None + elif mask_arr.ndim == 2: + mask_block_qkv = lax.dynamic_slice( + mask_arr, + (q_slice_start, kv_slice_start), + (block_q, block_kv), + ) + elif mask_arr.ndim == 3: + if mask_arr.shape[0] == q.shape[0] and q.ndim == 4 and mask_arr.shape[0] != num_q_heads: + mask_slice = lax.dynamic_slice( + mask_arr, + (0, q_slice_start, kv_slice_start), + (mask_arr.shape[0], block_q, block_kv), + ) + mask_block_qkv = mask_slice[:, None, :, :] + else: + mask_block_qkv = lax.dynamic_slice( + mask_arr, + (head_start, q_slice_start, kv_slice_start), + (eff_block_head, block_q, block_kv), + ) + else: + mask_block_qkv = lax.dynamic_slice( + mask_arr, + (0, head_start, q_slice_start, kv_slice_start), + (mask_arr.shape[0], eff_block_head, block_q, block_kv), + ) + + q_block_prec = q_block.astype(dtype) + do_block_prec = do_block.astype(dtype) + + def compute_updates(): + qk = jnp.einsum("...hd,...kd->...hk", q_block_prec, k_block_prec) + qk_masked, qk_uncapped = _apply_mask_and_soft_cap_bwd( + qk, + mask_block_qkv, + q_segment_ids_block, + kv_segment_ids_block, + mask_value, + attn_logits_soft_cap, + ) + + if use_log2: + p = 2 ** (qk_masked - logsumexp_block[..., None]) + else: + p = jnp.exp(qk_masked - logsumexp_block[..., None]) + + dv_block_update = jnp.einsum("...hk,...hd->...kd", p, do_block_prec) + + def fuse_dk_compute( + do_block_prec, + v_block_prec, + di_block, + p, + attn_logits_soft_cap, + q_block_prec, + ): + dp = jnp.einsum("...hd,...kd->...hk", do_block_prec, v_block_prec) + ds = p * (dp - di_block[..., None]) + if attn_logits_soft_cap is not None: + normalized = qk_uncapped / attn_logits_soft_cap + d = jnp.tanh(normalized) + g = ds * (1 - d) + ds = g + g * d + dk_block_update = jnp.einsum("...hk,...hd->...kd", ds, q_block_prec) + return dk_block_update + + def fuse_dq_compute( + do_block_prec, + v_block_prec, + di_block, + p, + attn_logits_soft_cap, + k_block_prec, + ): + dp = jnp.einsum("...hd,...kd->...hk", do_block_prec, v_block_prec) + ds = p * (dp - di_block[..., None]) + if attn_logits_soft_cap is not None: + normalized = qk_uncapped / attn_logits_soft_cap + d = jnp.tanh(normalized) + g = ds * (1 - d) + ds = g + g * d + dq_block_update = jnp.einsum("...hk,...kd->...hd", ds, k_block_prec) + return dq_block_update + + def fuse_dq_dk_compute( + do_block_prec, + v_block_prec, + di_block, + p, + attn_logits_soft_cap, + q_block_prec, + k_block_prec, + ): + dp = jnp.einsum("...hd,...kd->...hk", do_block_prec, v_block_prec) + ds = p * (dp - di_block[..., None]) + if attn_logits_soft_cap is not None: + normalized = qk_uncapped / attn_logits_soft_cap + d = jnp.tanh(normalized) + g = ds * (1 - d) + ds = g + g * d + dk_block_update = jnp.einsum("...hk,...hd->...kd", ds, q_block_prec) + dq_block_update = jnp.einsum("...hk,...kd->...hd", ds, k_block_prec) + return dk_block_update, dq_block_update + + if dp_dup_must_fuse: + dk_block_update = must_fuse_call("1")(fuse_dk_compute)( + do_block_prec, + v_block_prec, + di_block, + p, + attn_logits_soft_cap, + q_block_prec, + ) + dq_block = must_fuse_call("2")(fuse_dq_compute)( + do_block_prec, + v_block_prec, + di_block, + p, + attn_logits_soft_cap, + k_block_prec, + ) + elif dp_dedup_must_fuse: + dk_block_update, dq_block = must_fuse_call("3")(fuse_dq_dk_compute)( + do_block_prec, + v_block_prec, + di_block, + p, + attn_logits_soft_cap, + q_block_prec, + k_block_prec, + ) + else: + dp = jnp.einsum("...hd,...kd->...hk", do_block_prec, v_block_prec) + ds = (dp - di_block[..., None]) * p + if attn_logits_soft_cap is not None: + normalized = qk_uncapped / attn_logits_soft_cap + d = jnp.tanh(normalized) + g = ds * (1 - d) + ds = g + g * d + dk_block_update = jnp.einsum("...hk,...hd->...kd", ds, q_block_prec) + dq_block = jnp.einsum("...hk,...kd->...hd", ds, k_block_prec) + + return dk_block_update, dv_block_update, dq_block + + def skip_updates(): + return ( + jnp.zeros_like(k_block_prec), + jnp.zeros_like(v_block_prec), + jnp.zeros_like(q_block_prec), + ) + + if is_causal: + dk_block_update, dv_block_update, dq_block = lax.cond( + should_compute_block(i, j), + compute_updates, + skip_updates, + ) + else: + dk_block_update, dv_block_update, dq_block = compute_updates() + + dk_block_carry = dk_block_carry + dk_block_update + dv_block_carry = dv_block_carry + dv_block_update + dq_carry_inner = lax.dynamic_update_slice( + dq_carry_inner, + dq_block + lax.dynamic_slice(dq_carry_inner, start_indices_q, slice_sizes_q), + start_indices_q, + ) + return dk_block_carry, dv_block_carry, dq_carry_inner + + dk_block, dv_block, dq_carry = lax.fori_loop( + 0, + q_seq_len // block_q, + q_block_kernel, + (dk_block, dv_block, dq_carry), + unroll=unroll_loops, + ) + + dk_carry = lax.dynamic_update_slice(dk_carry, dk_block, start_indices_k) + dv_carry = lax.dynamic_update_slice(dv_carry, dv_block, start_indices_v) + return dq_carry, dk_carry, dv_carry + + num_kv_blocks = kv_seq_len // block_kv + dq_out, dk_out, dv_out = lax.fori_loop( + 0, + num_kv_blocks, + k_block_kernel, + head_carry, + unroll=unroll_loops, + ) + return dq_out, dk_out, dv_out + + dk = jnp.zeros_like(k, dtype=dtype) + dv = jnp.zeros_like(v, dtype=dtype) + dq = jnp.zeros_like(q, dtype=dtype) + num_head_blocks = num_q_heads // eff_block_head + dq, dk, dv = lax.fori_loop( + 0, + num_head_blocks, + head_kernel, + (dq, dk, dv), + unroll=unroll_loops, + ) + return dq.astype(q.dtype), dk.astype(k.dtype), dv.astype(v.dtype) + + +def flash_attention_mha_bwd_jax( + q: jax.Array, + k: jax.Array, + v: jax.Array, + o: jax.Array, + do: jax.Array, + logsumexp: jax.Array, + segment_ids: SegmentIds | None = None, + mask: jax.Array | mask_lib.Mask | None = None, + mask_value: float = -1e30, + attn_logits_soft_cap: float | None = None, + block_q: int = 1024, + block_kv: int = 1024, + block_head: int = 1, + dtype: jax.typing.DTypeLike = jnp.float32, + is_causal: bool = False, + use_log2: bool = False, + di: jax.Array | None = None, + fuse: bool = True, + unroll_loops: bool = False, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Memory efficient JAX implementation of Flash Attention backward pass for MHA.""" + if mask is not None: + mask_arr = jnp.asarray(mask[:, :]) if isinstance(mask, mask_lib.Mask) else jnp.asarray(mask) + else: + mask_arr = None + return _flash_attention_mha_bwd_jax( + q=q, + k=k, + v=v, + o=o, + do=do, + logsumexp=logsumexp, + segment_ids=segment_ids, + mask=mask_arr, + mask_value=mask_value, + attn_logits_soft_cap=attn_logits_soft_cap, + block_q=block_q, + block_kv=block_kv, + block_head=block_head, + dtype=dtype, + is_causal=is_causal, + use_log2=use_log2, + di=di, + fuse=fuse, + unroll_loops=unroll_loops, + ) + + +@functools.partial( + jax.jit, + static_argnames=[ + "block_q", + "block_kv", + "block_q_head", + "block_kv_head", + "dtype", + "is_causal", + "use_log2", + "fuse", + ], +) +def _flash_attention_gqa_bwd_jax( + q: jax.Array, + k: jax.Array, + v: jax.Array, + o: jax.Array, + do: jax.Array, + logsumexp: jax.Array, + segment_ids: SegmentIds | None = None, + mask: jax.Array | None = None, + block_q: int = 1024, + block_kv: int = 1024, + block_q_head: int = 1, + block_kv_head: int = 1, + mask_value: float = -1e30, + attn_logits_soft_cap: float | None = None, + dtype: jax.typing.DTypeLike = jnp.float32, + is_causal: bool = False, + use_log2: bool = False, + di: jax.Array | None = None, + fuse: bool = True, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Memory efficient JAX implementation of Flash Attention backward pass for GQA.""" + dp_dup_must_fuse = fuse + q_seq_len = q.shape[-2] + kv_seq_len = k.shape[-2] + num_q_heads = q.shape[-3] + num_kv_heads = k.shape[-3] + head_multiplier = num_q_heads // num_kv_heads + + eff_block_kv_head = min(block_kv_head, num_kv_heads) + eff_block_q_head = min(block_q_head, head_multiplier) + num_kv_head_blocks = num_kv_heads // eff_block_kv_head + num_q_head_blocks = head_multiplier // eff_block_q_head + + if q_seq_len % block_q != 0: + raise ValueError(f"q_seq_len {q_seq_len} must be divisible by block_q {block_q}") + if kv_seq_len % block_kv != 0: + raise ValueError(f"kv_seq_len {kv_seq_len} must be divisible by block_kv {block_kv}") + if num_kv_heads % eff_block_kv_head != 0: + raise ValueError(f"num_kv_heads {num_kv_heads} must be divisible by block_kv_head" f" {block_kv_head}") + if head_multiplier % eff_block_q_head != 0: + raise ValueError(f"head_multiplier {head_multiplier} must be divisible by block_q_head" f" {block_q_head}") + + if mask is None: + mask_arr = None + elif isinstance(mask, mask_lib.Mask): + mask_arr = mask[:, :] + else: + mask_arr = mask + + if mask_arr is not None: + if mask_arr.ndim == 2: + mask_arr = jnp.broadcast_to(mask_arr[None, :, :], (num_q_heads, *mask_arr.shape)) + elif mask_arr.ndim == 3 and mask_arr.shape[0] != num_q_heads: + mask_arr = jnp.broadcast_to(mask_arr, (num_q_heads, *mask_arr.shape[1:])) + + q_gqa = q.reshape(q.shape[:-3] + (num_kv_heads, head_multiplier, q_seq_len, q.shape[-1])) + do_gqa = do.reshape(do.shape[:-3] + (num_kv_heads, head_multiplier, q_seq_len, do.shape[-1])) + logsumexp_gqa = logsumexp.reshape(logsumexp.shape[:-2] + (num_kv_heads, head_multiplier, q_seq_len)) + + di_arr = jnp.sum(o * do, axis=-1, dtype=dtype) if di is None else di + di_gqa = di_arr.reshape(di_arr.shape[:-2] + (num_kv_heads, head_multiplier, q_seq_len)) + + def should_compute_block(i, j): + return not is_causal or (j + 1) * block_q > i * block_kv + + slice_sizes_k = k.shape[:-3] + (eff_block_kv_head, block_kv, k.shape[-1]) + slice_sizes_v = v.shape[:-3] + (eff_block_kv_head, block_kv, v.shape[-1]) + slice_sizes_q = q_gqa.shape[:-4] + ( + eff_block_kv_head, + eff_block_q_head, + block_q, + q_gqa.shape[-1], + ) + slice_sizes_do = do_gqa.shape[:-4] + ( + eff_block_kv_head, + eff_block_q_head, + block_q, + do_gqa.shape[-1], + ) + slice_sizes_log = logsumexp_gqa.shape[:-3] + ( + eff_block_kv_head, + eff_block_q_head, + block_q, + ) + + def kv_head_kernel(kv_head_idx, kv_head_carry): + dk_kv, dv_kv, dq_kv = kv_head_carry + kv_head_start = kv_head_idx * eff_block_kv_head + + def q_head_kernel(q_head_idx, q_head_carry): + dk_q, dv_q, dq_q = q_head_carry + q_head_start = q_head_idx * eff_block_q_head + + def k_block_kernel(i, k_carry): + dk_k, dv_k, dq_k = k_carry + kv_slice_start = i * block_kv + + start_indices_k = (0,) * (k.ndim - 3) + ( + kv_head_start, + kv_slice_start, + 0, + ) + start_indices_v = (0,) * (v.ndim - 3) + ( + kv_head_start, + kv_slice_start, + 0, + ) + + k_block = lax.dynamic_slice(k, start_indices_k, slice_sizes_k) + v_block = lax.dynamic_slice(v, start_indices_v, slice_sizes_v) + kv_segment_ids_block = ( + lax.dynamic_slice_in_dim(segment_ids.kv, kv_slice_start, block_kv, axis=-1) + if segment_ids is not None + else None + ) + + k_block_prec = k_block.astype(dtype) + v_block_prec = v_block.astype(dtype) + + dk_block = jnp.zeros_like(k_block_prec) + dv_block = jnp.zeros_like(v_block_prec) + + def q_block_kernel(j, inner_carry): + dk_inner, dv_inner, dq_inner = inner_carry + q_slice_start = j * block_q + + start_indices_q = (0,) * (q_gqa.ndim - 4) + ( + kv_head_start, + q_head_start, + q_slice_start, + 0, + ) + start_indices_do = (0,) * (do_gqa.ndim - 4) + ( + kv_head_start, + q_head_start, + q_slice_start, + 0, + ) + start_indices_log = (0,) * (logsumexp_gqa.ndim - 3) + ( + kv_head_start, + q_head_start, + q_slice_start, + ) + + q_block = lax.dynamic_slice(q_gqa, start_indices_q, slice_sizes_q) + do_block = lax.dynamic_slice(do_gqa, start_indices_do, slice_sizes_do) + logsumexp_block = lax.dynamic_slice(logsumexp_gqa, start_indices_log, slice_sizes_log).astype(dtype) + di_block = lax.dynamic_slice(di_gqa, start_indices_log, slice_sizes_log) + + q_segment_ids_block = ( + lax.dynamic_slice_in_dim(segment_ids.q, q_slice_start, block_q, axis=-1) + if segment_ids is not None + else None + ) + + if mask_arr is None: + mask_block_qkv = None + else: + mask_block_qkv = lax.dynamic_slice( + mask_arr, + ( + kv_head_start * head_multiplier + q_head_start, + q_slice_start, + kv_slice_start, + ), + (eff_block_kv_head * eff_block_q_head, block_q, block_kv), + ) + mask_block_qkv = mask_block_qkv.reshape(eff_block_kv_head, eff_block_q_head, block_q, block_kv) + + q_block_prec = q_block.astype(dtype) + do_block_prec = do_block.astype(dtype) + + def compute_updates(): + qk = jnp.einsum("...hqmd,...hkd->...hqmk", q_block_prec, k_block_prec) + qk_masked, qk_uncapped = _apply_mask_and_soft_cap_bwd( + qk, + mask_block_qkv, + q_segment_ids_block, + kv_segment_ids_block, + mask_value, + attn_logits_soft_cap, + ) + + if use_log2: + p = 2 ** (qk_masked - logsumexp_block[..., None]) + else: + p = jnp.exp(qk_masked - logsumexp_block[..., None]) + + dv_block_update = jnp.einsum("...hqmk,...hqmd->...hkd", p, do_block_prec) + + def fuse_dk_compute( + do_block_prec, + v_block_prec, + di_block, + p, + attn_logits_soft_cap, + q_block_prec, + ): + dp = jnp.einsum("...hqmd,...hkd->...hqmk", do_block_prec, v_block_prec) + ds = p * (dp - di_block[..., None]) + if attn_logits_soft_cap is not None: + normalized = qk_uncapped / attn_logits_soft_cap + d = jnp.tanh(normalized) + g = ds * (1 - d) + ds = g + g * d + + dk_block_update = jnp.einsum("...hqmk,...hqmd->...hkd", ds, q_block_prec) + return dk_block_update + + def fuse_dq_compute( + do_block_prec, + v_block_prec, + di_block, + p, + attn_logits_soft_cap, + k_block_prec, + ): + dp = jnp.einsum("...hqmd,...hkd->...hqmk", do_block_prec, v_block_prec) + ds = p * (dp - di_block[..., None]) + if attn_logits_soft_cap is not None: + normalized = qk_uncapped / attn_logits_soft_cap + d = jnp.tanh(normalized) + g = ds * (1 - d) + ds = g + g * d + + dq_block_update = jnp.einsum("...hqmk,...hkd->...hqmd", ds, k_block_prec) + return dq_block_update + + if dp_dup_must_fuse: + dk_block_update = must_fuse_call("1")(fuse_dk_compute)( + do_block_prec, + v_block_prec, + di_block, + p, + attn_logits_soft_cap, + q_block_prec, + ) + dq_block = must_fuse_call("2")(fuse_dq_compute)( + do_block_prec, + v_block_prec, + di_block, + p, + attn_logits_soft_cap, + k_block_prec, + ) + else: + dp = jnp.einsum("...hqmd,...hkd->...hqmk", do_block_prec, v_block_prec) + ds = p * (dp - di_block[..., None]) + if attn_logits_soft_cap is not None: + normalized = qk_uncapped / attn_logits_soft_cap + d = jnp.tanh(normalized) + g = ds * (1 - d) + ds = g + g * d + dk_block_update = jnp.einsum("...hqmk,...hqmd->...hkd", ds, q_block_prec) + dq_block = jnp.einsum("...hqmk,...hkd->...hqmd", ds, k_block_prec) + + return dk_block_update, dv_block_update, dq_block + + def skip_updates(): + return ( + jnp.zeros_like(k_block_prec), + jnp.zeros_like(v_block_prec), + jnp.zeros_like(q_block_prec), + ) + + if is_causal: + dk_block_update, dv_block_update, dq_block = lax.cond( + should_compute_block(i, j), + compute_updates, + skip_updates, + ) + else: + dk_block_update, dv_block_update, dq_block = compute_updates() + + dk_inner = dk_inner + dk_block_update + dv_inner = dv_inner + dv_block_update + dq_inner = lax.dynamic_update_slice( + dq_inner, + dq_block + lax.dynamic_slice(dq_inner, start_indices_q, slice_sizes_q), + start_indices_q, + ) + return dk_inner, dv_inner, dq_inner + + dk_block, dv_block, dq_k = lax.fori_loop( + 0, + q_seq_len // block_q, + q_block_kernel, + (dk_block, dv_block, dq_k), + ) + + dk_k = lax.dynamic_update_slice( + dk_k, + dk_block + lax.dynamic_slice(dk_k, start_indices_k, slice_sizes_k), + start_indices_k, + ) + dv_k = lax.dynamic_update_slice( + dv_k, + dv_block + lax.dynamic_slice(dv_k, start_indices_v, slice_sizes_v), + start_indices_v, + ) + return dk_k, dv_k, dq_k + + num_kv_blocks = kv_seq_len // block_kv + dk_q, dv_q, dq_q = lax.fori_loop( + 0, + num_kv_blocks, + k_block_kernel, + (dk_q, dv_q, dq_q), + ) + return dk_q, dv_q, dq_q + + dk_kv, dv_kv, dq_kv = lax.fori_loop( + 0, + num_q_head_blocks, + q_head_kernel, + (dk_kv, dv_kv, dq_kv), + ) + return dk_kv, dv_kv, dq_kv + + dk = jnp.zeros_like(k, dtype=dtype) + dv = jnp.zeros_like(v, dtype=dtype) + dq_gqa = jnp.zeros_like(q_gqa, dtype=dtype) + + dk, dv, dq_gqa = lax.fori_loop( + 0, + num_kv_head_blocks, + kv_head_kernel, + (dk, dv, dq_gqa), + ) + + dq = dq_gqa.reshape(q.shape) + return dq.astype(q.dtype), dk.astype(k.dtype), dv.astype(v.dtype) + + +def flash_attention_gqa_bwd_jax( + q: jax.Array, + k: jax.Array, + v: jax.Array, + o: jax.Array, + do: jax.Array, + logsumexp: jax.Array, + segment_ids: SegmentIds | None = None, + mask: jax.Array | mask_lib.Mask | None = None, + block_q: int = 1024, + block_kv: int = 1024, + block_q_head: int = 1, + block_kv_head: int = 1, + mask_value: float = -1e30, + attn_logits_soft_cap: float | None = None, + dtype: jax.typing.DTypeLike = jnp.float32, + is_causal: bool = False, + use_log2: bool = False, + di: jax.Array | None = None, + fuse: bool = True, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Memory efficient JAX implementation of Flash Attention backward pass for GQA.""" + if mask is not None: + mask_arr = jnp.asarray(mask[:, :]) if isinstance(mask, mask_lib.Mask) else jnp.asarray(mask) + else: + mask_arr = None + return _flash_attention_gqa_bwd_jax( + q=q, + k=k, + v=v, + o=o, + do=do, + logsumexp=logsumexp, + segment_ids=segment_ids, + mask=mask_arr, + block_q=block_q, + block_kv=block_kv, + block_q_head=block_q_head, + block_kv_head=block_kv_head, + mask_value=mask_value, + attn_logits_soft_cap=attn_logits_soft_cap, + dtype=dtype, + is_causal=is_causal, + use_log2=use_log2, + di=di, + fuse=fuse, + ) + + +def flash_attention_bwd( + q: jnp.ndarray, + k: jnp.ndarray, + v: jnp.ndarray, + o: jnp.ndarray, + do: jnp.ndarray, + logsumexp: jnp.ndarray, + mask: jnp.ndarray | mask_lib.Mask | None = None, + segment_ids: SegmentIds | None = None, + mask_value: float = -1e30, + cap: Optional[float] = None, + block_q: int = 1024, + block_kv: int = 1024, + block_q_head: Optional[int] = None, + block_kv_head: Optional[int] = None, + block_head: Optional[int] = None, + dtype: jax.typing.DTypeLike | None = None, + is_causal: bool = False, + use_log2: bool = False, + di: Optional[jnp.ndarray] = None, + fuse: bool = False, +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Computes backward pass for Flash Attention supporting MHA, GQA, and MQA.""" + if mask is not None: + mask_arr = jnp.asarray(mask[:, :]) if isinstance(mask, mask_lib.Mask) else jnp.asarray(mask) + else: + mask_arr = None + q_is_3d = q.ndim == 3 + k_was_2d = False + k_was_3d_in_4d = False + + if q_is_3d: + q = jnp.expand_dims(q, axis=0) + o = jnp.expand_dims(o, axis=0) + do = jnp.expand_dims(do, axis=0) + logsumexp = jnp.expand_dims(logsumexp, axis=0) + if di is not None: + di = jnp.expand_dims(di, axis=0) + if k.ndim == 3: + k = jnp.expand_dims(k, axis=0) + v = jnp.expand_dims(v, axis=0) + elif k.ndim == 2: + k_was_2d = True + k = jnp.expand_dims(k, axis=(0, 1)) + v = jnp.expand_dims(v, axis=(0, 1)) + else: + raise ValueError(f"Unsupported k ndim {k.ndim} for 3D q") + else: + if k.ndim == 3: + k_was_3d_in_4d = True + k = jnp.expand_dims(k, axis=1) + v = jnp.expand_dims(v, axis=1) + + num_q_heads = q.shape[-3] + num_kv_heads = k.shape[-3] + is_mqa = num_kv_heads == 1 and num_q_heads > 1 + use_gqa = (num_kv_heads < num_q_heads) or is_mqa + + comp_dtype = dtype if dtype is not None else q.dtype + eff_block_q = min(block_q, q.shape[-2]) + eff_block_kv = min(block_kv, k.shape[-2]) + + if use_gqa: + head_multiplier = num_q_heads // num_kv_heads + eff_block_kv_head = min(block_kv_head, num_kv_heads) if block_kv_head is not None else num_kv_heads + eff_block_q_head = min(block_q_head, head_multiplier) if block_q_head is not None else head_multiplier + + dq, dk, dv = flash_attention_gqa_bwd_jax( + q=q, + k=k, + v=v, + o=o, + do=do, + logsumexp=logsumexp, + segment_ids=segment_ids, + mask=mask_arr, + block_q=eff_block_q, + block_kv=eff_block_kv, + block_q_head=eff_block_q_head, + block_kv_head=eff_block_kv_head, + mask_value=mask_value, + attn_logits_soft_cap=cap, + dtype=comp_dtype, + is_causal=is_causal, + use_log2=use_log2, + di=di, + fuse=fuse, + ) + else: + eff_block_head = ( + min(block_head, num_q_heads) + if block_head is not None + else ( + min(block_q_head, num_q_heads) + if block_q_head is not None + else (min(block_kv_head, num_q_heads) if block_kv_head is not None else num_q_heads) + ) + ) + dq, dk, dv = flash_attention_mha_bwd_jax( + q=q, + k=k, + v=v, + o=o, + do=do, + logsumexp=logsumexp, + segment_ids=segment_ids, + mask=mask_arr, + mask_value=mask_value, + attn_logits_soft_cap=cap, + block_q=eff_block_q, + block_kv=eff_block_kv, + block_head=eff_block_head, + dtype=comp_dtype, + is_causal=is_causal, + use_log2=use_log2, + di=di, + fuse=fuse, + ) + + if q_is_3d: + dq = jnp.squeeze(dq, axis=0) + if k_was_2d: + dk = jnp.squeeze(dk, axis=(0, 1)) + dv = jnp.squeeze(dv, axis=(0, 1)) + else: + dk = jnp.squeeze(dk, axis=0) + dv = jnp.squeeze(dv, axis=0) + elif k_was_3d_in_4d: + dk = jnp.squeeze(dk, axis=1) + dv = jnp.squeeze(dv, axis=1) + + return dq, dk, dv + + +splash_attention_mha_bwd_jax = flash_attention_mha_bwd_jax +splash_attention_gqa_bwd_jax = flash_attention_gqa_bwd_jax +flash_attention_block_masked_bwd = flash_attention_bwd diff --git a/src/maxtext/kernels/attention/splash_attention_kernel.py b/src/maxtext/kernels/attention/splash_attention_kernel.py index 7f803c1811..e3520bfca1 100644 --- a/src/maxtext/kernels/attention/splash_attention_kernel.py +++ b/src/maxtext/kernels/attention/splash_attention_kernel.py @@ -156,7 +156,7 @@ def _attention_reference( attn_logits_soft_cap: float | None, ): """Reference attention implementation.""" - return _attention_reference_default( # pytype: disable=bad-return-type + return _attention_reference_default( mask, q, k, @@ -218,7 +218,7 @@ def attention_reference( attn_logits_soft_cap: float | None = None, ) -> SplashCustomReturnType: """Reference attention implementation.""" - return _attention_reference( # pytype: disable=wrong-arg-types + return _attention_reference( mask, q, k, @@ -285,7 +285,7 @@ def _attention_reference_custom_bwd( logits = jnp.where(mask, logits, mask_value) p = jnp.exp(logits - logsumexp[..., None]) - do = do.astype(jnp.float32) # pytype: disable=attribute-error + do = do.astype(jnp.float32) dv = jnp.einsum("pt,pd->td", p, do).astype(v.dtype) dp = jnp.einsum("pd,td->pt", do, v.astype(jnp.float32)) @@ -425,13 +425,13 @@ def _wrapped( if is_grouped: def reshape_activations(activations): - if activations.ndim == 4: # pytype: disable=attribute-error - kv_heads, q_heads_per_kv_head, q_seq_len, head_dim = activations.shape # pytype: disable=attribute-error + if activations.ndim == 4: + kv_heads, q_heads_per_kv_head, q_seq_len, head_dim = activations.shape return activations.reshape( kv_heads * q_heads_per_kv_head, q_seq_len, head_dim, - ) # pytype: disable=attribute-error + ) return activations def reshape_residuals(residuals): @@ -626,7 +626,7 @@ def _apply_mask_and_soft_cap( assert rem == 0 q_sequence = jnp.tile( q_sequence_ref[...], (1, repeats) # pyrefly: ignore[unsupported-operation] - ) # [bq, k_slice.size] # pyrefly: ignore[unsupported-operation] + ) # [bq, k_slice.size] else: assert q_sequence_ref.shape == (NUM_SUBLANES, bq) # pyrefly: ignore[missing-attribute] @@ -635,7 +635,7 @@ def _apply_mask_and_soft_cap( q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq)) assert q_sequence.shape == k_sequence.shape - computed_mask = mask_function(q_sequence, k_sequence) # pytype: disable=wrong-arg-count + computed_mask = mask_function(q_sequence, k_sequence) if computed_mask.dtype != jnp.dtype(jnp.bool_): raise ValueError("Mask function must return a boolean-valued array, but got:" f" {computed_mask.dtype}") masks.append(computed_mask) @@ -766,7 +766,7 @@ def body(kv_compute_index, _): qk = apply_mask_and_soft_cap() - m_curr = qk.max(axis=-1)[:, None] # pytype: disable=attribute-error + m_curr = qk.max(axis=-1)[:, None] # pyrefly: ignore[missing-attribute] assert m_curr.shape == (bq, 1) m_next = jnp.maximum(m_prev, m_curr) assert m_next.shape == (bq, NUM_LANES) @@ -1137,7 +1137,7 @@ def _splash_attention_custom( # device. del dq_mask_info, dkv_mask_info - return _splash_attention_forward( # pytype: disable=wrong-arg-types + return _splash_attention_forward( # pyrefly: ignore[no-matching-overload] fwd_mask_info, q, k, @@ -1178,7 +1178,7 @@ def _splash_attention_fwd( if save_residuals: raise NotImplementedError("Higher-order AD not supported") - out, (logsumexp,) = _splash_attention_forward( # pytype: disable=wrong-arg-types + out, (logsumexp,) = _splash_attention_forward( # pyrefly: ignore[no-matching-overload] fwd_mask_info, q, k, @@ -1269,7 +1269,7 @@ def run(): q_segment_ids_ref, kv_segment_ids_ref, attn_logits_soft_cap=attn_logits_soft_cap, # pyrefly: ignore[bad-argument-type] - k_slice=pl.ds(0, bkv), # pyrefly: ignore[bad-argument-type] + k_slice=pl.ds(0, bkv), # When the iteration space is shrunk (for local attention for example), # the kv_index program_id does not correspond to the actual coordinates # of the KV data. Make sure to use the 'unshrunk' index (coming from the @@ -1636,7 +1636,7 @@ def _load_kv(ref, layout): q_segment_ids_ref, kv_segment_ids_ref, attn_logits_soft_cap=attn_logits_soft_cap, # pyrefly: ignore[bad-argument-type] - k_slice=slice_k, # pyrefly: ignore[bad-argument-type] + k_slice=slice_k, k_offset=kv_index * bkv + i * bkv_compute, bq=bq, k_in_lanes=False, @@ -1698,7 +1698,7 @@ def run(): elif num_kv_heads < num_q_heads: should_write = jnp.logical_and( should_write, q_head_index_per_kv_head == q_heads_per_kv_heads - 1 # pyrefly: ignore[unsupported-operation] - ) # pyrefly: ignore[unsupported-operation] + ) @pl.when(should_write) def end(): @@ -2112,7 +2112,7 @@ def _splash_attention_bwd( ) = res # di: [num_heads, q_seq_len] - di = jnp.einsum("hsd,hsd->hs", o.astype(jnp.float32), do.astype(jnp.float32)) # pytype: disable=attribute-error + di = jnp.einsum("hsd,hsd->hs", o.astype(jnp.float32), do.astype(jnp.float32)) dq, dk, dv = _splash_attention_bwd_dkv( q, k, @@ -2296,7 +2296,7 @@ def _collapse_partial_mask_blocks(mask_info: mask_info_lib.MaskInfo | None): dkv_mask_info = _collapse_partial_mask_blocks(dkv_mask_info) del dq_mask_info, dkv_mask_info - out, (logsumexp,) = _splash_attention_forward( # pytype: disable=wrong-arg-types + out, (logsumexp,) = _splash_attention_forward( # pyrefly: ignore[no-matching-overload] fwd_mask_info, q, k, @@ -2425,7 +2425,7 @@ def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding): partial_mask_blocks_spec = spec if self.fwd_mask_info.is_dynamic_mask else replicated # Shard q_sequence over the sequence dimension only. q_sequence_spec = jax.sharding.PartitionSpec(spec[1]) - mask_info_specs = mask_info_lib.MaskInfo( # pytype: disable=wrong-arg-types + mask_info_specs = mask_info_lib.MaskInfo( data_next=spec if self.fwd_mask_info.data_next is not None else None, # pyrefly: ignore[bad-argument-type] mask_next=spec if self.fwd_mask_info.mask_next is not None else None, # pyrefly: ignore[bad-argument-type] block_mask=spec if self.fwd_mask_info.block_mask is not None else None, # pyrefly: ignore[bad-argument-type] diff --git a/src/maxtext/kernels/attention/tokamax_ring_attention.py b/src/maxtext/kernels/attention/tokamax_ring_attention.py index ed581bb7ef..9fc7d5c865 100644 --- a/src/maxtext/kernels/attention/tokamax_ring_attention.py +++ b/src/maxtext/kernels/attention/tokamax_ring_attention.py @@ -224,9 +224,24 @@ def build_splash_config( dq_reduction_steps = config.dq_reduction_steps q_seq_len_per_shard = q_seq_len // context_parallel_size kv_seq_len_per_shard = kv_seq_len // context_parallel_size - block_q = min(config.sa_block_q, q_seq_len_per_shard) - block_kv = min(config.sa_block_kv, kv_seq_len_per_shard) - block_kv_compute = min(config.sa_block_kv_compute, kv_seq_len_per_shard) + is_eval = max_utils.is_eval(config) + if is_eval: + sa_block_q = config.eval_sa_block_q + sa_block_kv = config.eval_sa_block_kv + sa_block_kv_compute = config.eval_sa_block_kv_compute + sa_q_layout = config.eval_sa_q_layout + sa_k_layout = config.eval_sa_k_layout + sa_v_layout = config.eval_sa_v_layout + else: + sa_block_q = config.sa_block_q + sa_block_kv = config.sa_block_kv + sa_block_kv_compute = config.sa_block_kv_compute + sa_q_layout = config.sa_q_layout + sa_k_layout = config.sa_k_layout + sa_v_layout = config.sa_v_layout + block_q = min(sa_block_q, q_seq_len_per_shard) + block_kv = min(sa_block_kv, kv_seq_len_per_shard) + block_kv_compute = min(sa_block_kv_compute, kv_seq_len_per_shard) block_q_dkv = min(config.sa_block_q_dkv, q_seq_len_per_shard) block_kv_dkv = min(config.sa_block_kv_dkv, kv_seq_len_per_shard) block_kv_dkv_compute = min(config.sa_block_kv_dkv_compute, kv_seq_len_per_shard) @@ -246,16 +261,24 @@ def build_splash_config( block_kv_dkv=block_kv_dkv, block_kv_dkv_compute=block_kv_dkv_compute, use_fused_bwd_kernel=True, - q_layout=tokamax_splash_kernel.QKVLayout[config.sa_q_layout], - k_layout=tokamax_splash_kernel.QKVLayout[config.sa_k_layout], - v_layout=tokamax_splash_kernel.QKVLayout[config.sa_v_layout], + q_layout=tokamax_splash_kernel.QKVLayout[sa_q_layout], + k_layout=tokamax_splash_kernel.QKVLayout[sa_k_layout], + v_layout=tokamax_splash_kernel.QKVLayout[sa_v_layout], attn_logits_soft_cap=attn_logits_soft_cap, residual_checkpoint_name="context", use_base2_exp=False, - fwd_cost_estimate=pl.CostEstimate(flops=config.cost_estimate_flops_fwd, transcendentals=0, bytes_accessed=0) + fwd_cost_estimate=pl.CostEstimate( + flops=config.cost_estimate_flops_fwd, + transcendentals=0, + bytes_accessed=0, + ) if config.cost_estimate_flops_fwd >= 0 else None, - bwd_cost_estimate=pl.CostEstimate(flops=config.cost_estimate_flops_bwd, transcendentals=0, bytes_accessed=0) + bwd_cost_estimate=pl.CostEstimate( + flops=config.cost_estimate_flops_bwd, + transcendentals=0, + bytes_accessed=0, + ) if config.cost_estimate_flops_bwd >= 0 else None, dq_reduction_steps=dq_reduction_steps if dq_reduction_steps > 0 else None, diff --git a/src/maxtext/kernels/gather_reduce_sc.py b/src/maxtext/kernels/gather_reduce_sc.py index 0beab57e02..2f49ef8949 100644 --- a/src/maxtext/kernels/gather_reduce_sc.py +++ b/src/maxtext/kernels/gather_reduce_sc.py @@ -573,8 +573,8 @@ def get_row_val(row_idx): arith.cmpf( arith.CmpFPredicate.OEQ, weights_vecs[0], # pyrefly: ignore[unsupported-operation] - zero_vec_f32, # pyrefly: ignore[unbound-name, unsupported-operation] - ), # pyrefly: ignore[unbound-name, unsupported-operation] + zero_vec_f32, # pyrefly: ignore[unbound-name] + ), zero_vec_f32, row0, ) @@ -586,7 +586,7 @@ def get_row_val(row_idx): row8 = arith.select( arith.cmpf( arith.CmpFPredicate.OEQ, weights_vecs[8], zero_vec_f32 # pyrefly: ignore[unsupported-operation] - ), # pyrefly: ignore[unsupported-operation] + ), zero_vec_f32, row8, ) diff --git a/src/maxtext/kernels/gdn/compute_conv1d.py b/src/maxtext/kernels/gdn/compute_conv1d.py index 05e40e3c69..bdcba5efb6 100644 --- a/src/maxtext/kernels/gdn/compute_conv1d.py +++ b/src/maxtext/kernels/gdn/compute_conv1d.py @@ -20,6 +20,38 @@ from . import config +# One float32 vector register is 8 sublanes x 128 lanes. gdn_conv1d_dense_token_stack regroups causal_conv1d's +# per-token rows in units of one register (see stack_conv1d_tokens), so it needs rows of a multiple of this size. +_F32_VREG_SUBLANES = 8 +_VREG_LANES = 128 +DENSE_TOKEN_STACK_ALIGNMENT = _F32_VREG_SUBLANES * _VREG_LANES + + +def stack_conv1d_tokens(out_list: list[jax.Array], dense: bool) -> jax.Array: + """Stacks causal_conv1d's per-token rows, each [seq, 1, dim], into its [seq, chunk, 1, dim] output. + + dense=False is jnp.stack(out_list, axis=1). The stacked value's second-minor dimension is 1, and the TPU compiler + of the October 2026 Pathways build (and of the libtpu 0.0.51 nightlies) lays it out in (8, 128) tiles padded from + 1 to 8 sublanes. The SiLU that wrapper.inner_kernel applies next inherits that layout, so its exp and reciprocal + run on 8x as many vector registers as the data fills. Older compilers keep the value in dense (1, 128) tiles. + + dense=True (config gdn_conv1d_dense_token_stack) views each [1, 1, dim] float32 row as [dim / 1024, 8, 128], which + is free because the row already fills dim / 1024 registers, concatenates the views along the leading axis and + reshapes the result to [1, chunk, 1, dim]. The values and their order are the same; only the layout the compiler + chooses changes, and it stays dense on old and new compilers alike. Requires seq == 1 (GDNMode.PER_SEQ), float32 + rows and dim a multiple of DENSE_TOKEN_STACK_ALIGNMENT. + """ + if not dense: + return jnp.stack(out_list, axis=1) + row = out_list[0] + if row.shape[0] != 1 or row.size % DENSE_TOKEN_STACK_ALIGNMENT or jnp.dtype(row.dtype).itemsize != 4: + raise ValueError( + "gdn_conv1d_dense_token_stack needs one sequence per tile and float32 rows whose size is a multiple of" + f" {DENSE_TOKEN_STACK_ALIGNMENT}, got a {jnp.dtype(row.dtype)} row of shape {row.shape}." + ) + dense_rows = jnp.concatenate([o.reshape(-1, _F32_VREG_SUBLANES, _VREG_LANES) for o in out_list], axis=0) + return dense_rows.reshape(row.shape[0], len(out_list), *row.shape[1:]) + def encode_segment_ids( segment_ids: jax.Array, @@ -296,4 +328,4 @@ def causal_conv1d( new_conv_state, ) - return jnp.stack(out_list, axis=1), new_conv_state + return stack_conv1d_tokens(out_list, cfg.conv_dense_token_stack), new_conv_state diff --git a/src/maxtext/kernels/gdn/compute_gdn.py b/src/maxtext/kernels/gdn/compute_gdn.py index 36d8c7bf62..a19b43790a 100644 --- a/src/maxtext/kernels/gdn/compute_gdn.py +++ b/src/maxtext/kernels/gdn/compute_gdn.py @@ -15,11 +15,58 @@ """Core GDN forward computation with triangular inverse matrix caching.""" +import dataclasses + import jax import jax.numpy as jnp from . import config +# ============================================================================== +# Triangular-inverse algorithm selection. +# +# `xor_dnc` is the default algorithm (pure-MXU divide-and-conquer with a +# half-row final merge). `baseline_b16` is the legacy 16-row block forward +# substitution algorithm. Both compute (I + L)^-1 for a batch of unit +# lower-triangular matrices `t: [batch, chunk, chunk]`. +# ============================================================================== + +DEFAULT_TRI_INV_ALGO = "xor_dnc" +TRI_INV_ALGOS = ("xor_dnc", "baseline_b16") + + +def validate_tri_inv_algo(algo: str) -> str: + if algo not in TRI_INV_ALGOS: + raise ValueError(f"Unknown triangular-inverse algorithm {algo!r}; expected one of {TRI_INV_ALGOS}.") + return algo + + +def default_head_group_size( + chunk_size: int, + num_kq_heads: int, + num_v_heads: int, + act_dtype: jnp.dtype | None = None, +) -> int | None: + """Returns the default KQ-head group size for `chunk_size`, or None for all heads at once. + + At `chunk_size > 64`, processing all heads at once can exceed VMEM with `[num_v_heads, + chunk_size, chunk_size]` intermediates. Bounding each group to at most 32 value heads + for `bfloat16` activations (or 8 value heads for `float32` activations) at `chunk_size = 128` + keeps peak VMEM comfortably within budget while preserving full MXU utilization. + """ + if chunk_size <= 64 or num_kq_heads <= 0 or num_v_heads <= 0: + return None + v_per_kq = max(1, num_v_heads // num_kq_heads) + base_v_heads = 8 if (act_dtype is not None and jnp.dtype(act_dtype) == jnp.float32) else 32 + max_v_heads = max(v_per_kq, (base_v_heads * 128) // chunk_size) + max_kq = max(1, max_v_heads // v_per_kq) + if num_kq_heads <= max_kq: + return None + for g in range(max_kq, 0, -1): + if num_kq_heads % g == 0: + return g + return 1 + def l2_norm(x: jax.Array, eps: float = 1e-6) -> jax.Array: norm = jnp.sqrt(jnp.sum(x * x, axis=-1, keepdims=True, dtype=x.dtype) + eps) @@ -37,12 +84,12 @@ def get_mask_dtype(dtype: jnp.dtype) -> jnp.dtype: # NOTE: Fork of recurrent_scan_v2.py but applied various optimizations. -def invert_triangular_matrix( +def _tri_inv_baseline_b16( t: jax.Array, block_size: int = 16, precision: jax.lax.Precision | None = None, ) -> jax.Array: - """Compute invert matrix of a given triauglar matrix.""" + """Compute inverse matrix of a given triangular matrix using 16-row block forward substitution.""" # NOTE: if chunk_size=1, compiler will perform DCE. out_dtype = t.dtype @@ -94,6 +141,107 @@ def local_forward_sub(t_mat: jax.Array, b_mat: jax.Array) -> jax.Array: return jnp.concatenate(x_blocks, axis=1) +def _bmm(lhs: jax.Array, rhs: jax.Array, precision: jax.lax.Precision | None) -> jax.Array: + """Batched [n, a, b] @ [n, b, c] -> [n, a, c] with f32 accumulation.""" + return jax.lax.dot_general( + lhs, + rhs, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + precision=precision, + preferred_element_type=jnp.float32, + ) + + +def _row_col_iotas(shape: tuple[int, ...]) -> tuple[jax.Array, jax.Array]: + """Broadcastable [1, R, C] row and column int32 iotas for the last two dims of `shape`.""" + mask_shape = (1, shape[-2], shape[-1]) + return ( + jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1), + jax.lax.broadcasted_iota(jnp.int32, mask_shape, 2), + ) + + +def _dnc_merge_levels( + x: jax.Array, + lower_part: jax.Array, + first_block: int, + precision: jax.lax.Precision | None, + xor_rc: jax.Array | None = None, +) -> jax.Array: + """Merges a block-diagonal inverse with `first_block`-sized blocks up to full size. + + At level b every 2b-window [[A, 0], [B, C]] gets its lower-left block + -C^-1 B A^-1, i.e. X <- X - X ((M_b * L) X) at full width, where M_b selects the + lower-left bxb block of each 2b-window ((r ^ c) in [b, 2b) and r > c). + + Args: + x: [n, c, c] f32 block-diagonal inverse with `first_block`-sized blocks. + lower_part: [n, c, c] f32 lower-triangular matrix (zeros above the diagonal) + whose strictly lower part is L of T = I + L. + first_block: Size of the diagonal blocks already inverted in `x`. + precision: Matmul precision. + xor_rc: Optional precomputed `bitwise_xor(iota_r, iota_c)` of `[1, c, c]`. + + Returns: + [n, c, c] f32 inverse of T. + """ + chunk = x.shape[-1] + if xor_rc is None: + iota_r, iota_c = _row_col_iotas(x.shape) + xor_rc = jnp.bitwise_xor(iota_r, iota_c) + b = first_block + while b < chunk: + if 2 * b == chunk and b >= 8: + lower_bot = jnp.where(xor_rc[:, b:, :] >= b, lower_part[:, b:, :], 0.0) + mid = _bmm(lower_bot, x, precision) + mid_padded = jnp.concatenate([jnp.zeros_like(mid), mid], axis=1) + bot = x[:, b:, :] - _bmm(x[:, b:, :], mid_padded, precision) + x = jnp.concatenate([x[:, :b, :], bot], axis=1) + else: + level_mask = (xor_rc >= b) & (xor_rc < 2 * b) + lower_b = jnp.where(level_mask, lower_part, 0.0) + x = x - _bmm(x, _bmm(lower_b, x, precision), precision) + b *= 2 + return x + + +def _tri_inv_xor_dnc(t: jax.Array, precision: jax.lax.Precision | None = None) -> jax.Array: + """Pure-MXU divide and conquer (gdm/universal_ai/clup/delta_rule.py), 2(log2 C - 1) matmuls.""" + chunk = t.shape[-1] + if chunk <= 1 or (chunk & (chunk - 1)) != 0: + return _tri_inv_baseline_b16(t, precision=precision) + out_dtype = t.dtype + t32 = t.astype(jnp.float32) + iota_r, iota_c = _row_col_iotas(t.shape) + xor_rc = jnp.bitwise_xor(iota_r, iota_c) + x = jnp.where(xor_rc == 0, 1.0, jnp.where(xor_rc == 1, -t32, 0.0)) + return _dnc_merge_levels(x, t32, 2, precision, xor_rc).astype(out_dtype) + + +def invert_triangular_matrix( + t: jax.Array, + block_size: int = 16, + precision: jax.lax.Precision | None = None, + *, + algo: str = DEFAULT_TRI_INV_ALGO, +) -> jax.Array: + """Computes (I + L)^-1 for unit lower-triangular `t: [batch, chunk, chunk]`. + + Args: + t: Unit lower-triangular matrices (ones on the diagonal, zeros above it). + block_size: Block size of the `baseline_b16` algorithm. + precision: Matmul precision. + algo: One of `TRI_INV_ALGOS`; defaults to `"xor_dnc"`. + + Returns: + The inverses, in `t.dtype`. + """ + validate_tri_inv_algo(algo) + if algo == "xor_dnc": + return _tri_inv_xor_dnc(t, precision) + return _tri_inv_baseline_b16(t, block_size, precision) + + def fused_transpose_broadcast(x: jax.Array, src_dim: int, dst_dim: int) -> jax.Array: """Perform 1D transpose where results are broadcasted along src_dim.""" assert x.shape[dst_dim] == 1 @@ -114,33 +262,76 @@ def chunked_gdn_per_seq( cfg: config.GDNConfig, seg_c_col: jax.Array | None = None, # [1, chunk, 1] seg_aux_col: jax.Array | None = None, # [1, chunk, 1] -) -> tuple[jax.Array, jax.Array, jax.Array]: - """Perform chunked GDN over input [num_heads, chunk, head_dim].""" - # bf16 activations keep single-pass MXU matmuls even with an f32 compute dtype - # (HIGHEST would triple the forward MXU cost); HIGHEST needs f32 activations and compute. - dot_prec = ( - jax.lax.Precision.HIGHEST - if jnp.dtype(cfg.dtypes.act_in) == jnp.float32 and jnp.dtype(cfg.dtypes.compute) == jnp.float32 - else jax.lax.Precision.DEFAULT - ) +) -> tuple[jax.Array, jax.Array, jax.Array | None]: + """Perform chunked GDN over input [num_heads, chunk, head_dim]. + + With `cfg.head_group_size` set, kq heads (and their value heads) are processed + `head_group_size` at a time, which bounds the live [heads, chunk, chunk] + intermediates at large chunk sizes. + """ + group = cfg.head_group_size + if group is None or group >= cfg.num_kq_heads: + return _chunked_gdn_per_seq_heads( + q_large, k_large, v_large, gating_log, beta, state_prev, cfg, seg_c_col=seg_c_col, seg_aux_col=seg_aux_col + ) + if cfg.num_kq_heads % group != 0: + raise ValueError(f"head_group_size={group} must divide num_kq_heads={cfg.num_kq_heads}.") + group_v = group * cfg.v_per_kq_head + group_cfg = dataclasses.replace(cfg, num_kq_heads=group, num_v_heads=group_v, head_group_size=None) + # The decay cumsum is a chunk-long serial chain on one lane-dense vreg row that + # every head shares: compute it once instead of once per head group. The lane -> + # head transposes stay per group, so only that group's lane-padded + # [heads, chunk, 1] columns are live at a time. + valid_c, active_c, _ = _segment_columns(cfg, seg_c_col, seg_aux_col) + if valid_c is not None: + gating_log = jnp.where(valid_c, gating_log, 0.0) + beta = jnp.where(valid_c, beta, 0.0) + # [1, chunk, aligned_num_v_heads] + g_cum_sum_all = _decay_cumsum(gating_log, cfg, active_c) + outs, states, t_invs = [], [], [] + for g in range(cfg.num_kq_heads // group): + kq_sl = slice(g * group, (g + 1) * group) + v_sl = slice(g * group_v, (g + 1) * group_v) + out, state, t_inv = _chunked_gdn_per_seq_heads( + q_large[kq_sl], + k_large[kq_sl], + v_large[v_sl], + None, + beta[:, :, v_sl], + state_prev[v_sl], + group_cfg, + seg_c_col=seg_c_col, + seg_aux_col=seg_aux_col, + decay_cumsum=g_cum_sum_all[:, :, v_sl], + ) + outs.append(out) + states.append(state) + if cfg.emit_t_inv: + t_invs.append(t_inv) + t_inv_cat = jnp.concatenate(t_invs, axis=0) if cfg.emit_t_inv else None + return jnp.concatenate(outs, axis=0), jnp.concatenate(states, axis=0), t_inv_cat - valid_c = None - active_c = None - seg_prev = None + +def _segment_columns( + cfg: config.GDNConfig, + seg_c_col: jax.Array | None, + seg_aux_col: jax.Array | None, +) -> tuple[jax.Array | None, jax.Array | None, jax.Array | None]: + """(valid_c, active_c, seg_prev) [1, chunk, 1] columns for packed sequences, else Nones.""" if cfg.has_seg_ids and seg_c_col is not None and seg_aux_col is not None: valid_c = seg_c_col > 0.5 active_c = jnp.abs(seg_c_col) seg_prev = jnp.abs(seg_aux_col[:, cfg.prev_kernel_size : cfg.prev_kernel_size + 1, :]) - q_large = jnp.where(valid_c, q_large, 0.0) - k_large = jnp.where(valid_c, k_large, 0.0) - v_large = jnp.where(valid_c, v_large, 0.0) - gating_log = jnp.where(valid_c, gating_log, 0.0) - beta = jnp.where(valid_c, beta, 0.0) + return valid_c, active_c, seg_prev + return None, None, None - # NOTE: Repeat along non lane/sublane dim is free. - q_repeat = jnp.repeat(q_large, cfg.v_per_kq_head, axis=0) - k_repeat = jnp.repeat(k_large, cfg.v_per_kq_head, axis=0) +def _decay_cumsum( + gating_log: jax.Array, # [1, chunk, lanes], masked for packed sequences + cfg: config.GDNConfig, + active_c: jax.Array | None, +) -> jax.Array: + """[1, chunk, lanes] f32 cumulative log-decay (restarting at document starts when packed).""" # Compute cumulative sum of decay in float32 to prevent bf16 cumsum drift. gating_log = gating_log.astype(jnp.float32) # [1, 1, num_v_heads] @@ -153,7 +344,50 @@ def chunked_gdn_per_seq( else: g_cum_sum_list.append(g_cum_sum_list[-1] + gating_log[:, row : row + 1]) # [1, chunk, num_v_heads] - g_cum_sum_log = jnp.concat(g_cum_sum_list, axis=1) + return jnp.concat(g_cum_sum_list, axis=1) + + +def _chunked_gdn_per_seq_heads( + q_large: jax.Array, # [num_kq_heads, chunk, kq_head_dim] + k_large: jax.Array, # [num_kq_heads, chunk, kq_head_dim] + v_large: jax.Array, # [num_v_heads, chunk, v_head_dim] + gating_log: jax.Array | None, # [1, chunk, >= num_v_heads]; None with decay_cumsum + beta: jax.Array, # [1, chunk, >= num_v_heads]; already masked with decay_cumsum + state_prev: jax.Array, # [num_v_heads, kq_head_dim, v_head_dim] + cfg: config.GDNConfig, + seg_c_col: jax.Array | None = None, # [1, chunk, 1] + seg_aux_col: jax.Array | None = None, # [1, chunk, 1] + decay_cumsum: jax.Array | None = None, # [1, chunk, >= num_v_heads] +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Perform chunked GDN over all heads of input [num_heads, chunk, head_dim]. + + `decay_cumsum` is the precomputed (`_decay_cumsum`) cumulative log-decay of these + heads; it replaces `gating_log`, and `beta` must then already be masked. + """ + # bf16 activations keep single-pass MXU matmuls even with an f32 compute dtype + # (HIGHEST would triple the forward MXU cost); HIGHEST needs f32 activations and compute. + dot_prec = ( + jax.lax.Precision.HIGHEST + if jnp.dtype(cfg.dtypes.act_in) == jnp.float32 and jnp.dtype(cfg.dtypes.compute) == jnp.float32 + else jax.lax.Precision.DEFAULT + ) + + valid_c, active_c, seg_prev = _segment_columns(cfg, seg_c_col, seg_aux_col) + if valid_c is not None: + q_large = jnp.where(valid_c, q_large, 0.0) + k_large = jnp.where(valid_c, k_large, 0.0) + v_large = jnp.where(valid_c, v_large, 0.0) + if decay_cumsum is None: + gating_log = jnp.where(valid_c, gating_log, 0.0) + beta = jnp.where(valid_c, beta, 0.0) + + # NOTE: Repeat along non lane/sublane dim is free. + # [num_v_heads, chunk, kq_head_dim] + q_repeat = jnp.repeat(q_large, cfg.v_per_kq_head, axis=0) + k_repeat = jnp.repeat(k_large, cfg.v_per_kq_head, axis=0) + + # [1, chunk, num_v_heads] + g_cum_sum_log = _decay_cumsum(gating_log, cfg, active_c) if decay_cumsum is None else decay_cumsum # [num_v_heads, chunk, 1] g_cum_sum_log = fused_transpose_broadcast(g_cum_sum_log, src_dim=2, dst_dim=0) @@ -177,49 +411,44 @@ def chunked_gdn_per_seq( m_out = ((jnp.abs(active_c - active_last) < 0.5) & (active_last > 0.5)).astype(g_cum_sum_log.dtype) m_keep = ((jnp.abs(active_last - seg_prev) < 0.5) & (seg_prev > 0.5)).astype(g_cum_sum_log.dtype) - bwd_diff_safe = jnp.where(m_out > 0.5, -g_cum_sum_diff_log[..., -1:], -1e4) + bwd_diff_safe = jnp.where(m_out > 0.5, g_cum_sum_log[:, -1:] - g_cum_sum_log, -1e4) gating_backward = (jnp.exp(bwd_diff_safe) * m_out).astype(cfg.dtypes.compute) gating_forward = (jnp.exp(g_cum_sum_log) * m_in).astype(cfg.dtypes.compute) gating_last = (jnp.exp(g_cum_sum_log[:, -1:]) * m_keep).astype(cfg.dtypes.compute) else: gating_map = jnp.exp(g_cum_sum_diff_log).astype(cfg.dtypes.compute) # [num_v_heads, chunk, 1] - gating_backward = jnp.exp(-g_cum_sum_diff_log[..., -1:]).astype(cfg.dtypes.compute) + gating_backward = jnp.exp(g_cum_sum_log[:, -1:] - g_cum_sum_log).astype(cfg.dtypes.compute) # [num_v_heads, chunk, 1] gating_forward = jnp.exp(g_cum_sum_log).astype(cfg.dtypes.compute) # [num_v_heads, 1, 1] gating_last = gating_forward[:, -1:] mask_dtype = get_mask_dtype(cfg.dtypes.compute) - iota_r = jax.lax.broadcasted_iota(mask_dtype, gating_map.shape, 1) - iota_c = jax.lax.broadcasted_iota(mask_dtype, gating_map.shape, 2) - identity_mask = iota_r == iota_c - strictly_lower_mask = iota_r > iota_c - lower_mask = iota_r >= iota_c - # [num_v_heads, chunk, chunk] - gating_map_masked = jnp.where(strictly_lower_mask, gating_map, 0) - - # [num_v_heads, chunk, kq_head_dim] - k_beta_repeat = k_repeat * beta_large + mask_shape = (1, cfg.chunk_size, cfg.chunk_size) + iota_r = jax.lax.broadcasted_iota(mask_dtype, mask_shape, 1) + iota_c = jax.lax.broadcasted_iota(mask_dtype, mask_shape, 2) + gating_map = jnp.where(iota_r >= iota_c, gating_map, 0) - # [num_v_heads, chunk, chunk] - beta_k_k_t = jax.lax.dot( - k_beta_repeat, - k_repeat, + # Compute k @ k^T at num_kq_heads (4x fewer matmuls) before repeating to num_v_heads. + k_k_t = jax.lax.dot( + k_large, + k_large, dimension_numbers=(((2,), (2,)), ((0,), (0,))), precision=dot_prec, preferred_element_type=jnp.float32, ).astype(cfg.dtypes.compute) - gating_beta_k_k_t = gating_map_masked * beta_k_k_t - t = jnp.where(identity_mask, 1, gating_beta_k_k_t) + beta_k_k_t = jnp.repeat(k_k_t, cfg.v_per_kq_head, axis=0) * beta_large + gating_beta_k_k_t = jnp.where(iota_r > iota_c, gating_map * beta_k_k_t, 0) + t = jnp.where(iota_r == iota_c, 1, gating_beta_k_k_t) # [num_v_heads, chunk, chunk] - t_inv = invert_triangular_matrix(t, precision=dot_prec) + t_inv = invert_triangular_matrix(t, precision=dot_prec, algo=cfg.tri_inv_algo) # [num_v_heads, chunk, v_head_dim] v_beta_large = v_large * beta_large # [num_v_heads, chunk, kv_head_dim] - k_beta_gating = k_beta_repeat * gating_forward + k_beta_gating = k_repeat * (beta_large * gating_forward) # NOTE: If v_head_dim < mxu size, concatenating them will help increase mxu # utilization. Also, if v_head_dim is multiple of lane size, concat / split # along lane dim is free - making this optimization strictly beneficial. @@ -286,7 +515,6 @@ def chunked_gdn_per_seq( # [num_v_heads, chunk, chunk] out_qk = jnp.repeat(out_qk, cfg.v_per_kq_head, axis=0) out_qk *= gating_map - out_qk = jnp.where(lower_mask, out_qk, 0) # [num_v_heads, chunk, v_head_dim] out_new = jax.lax.dot( @@ -314,7 +542,7 @@ def chunked_gdn( a_log: jax.Array, dt_bias: jax.Array, cfg: config.GDNConfig, -) -> tuple[jax.Array, jax.Array, jax.Array]: +) -> tuple[jax.Array, jax.Array, jax.Array | None]: """Perform chunked GDN over input [seq, num_heads, chunk, head_dim].""" mask_dtype = get_mask_dtype(cfg.dtypes.compute) @@ -342,6 +570,7 @@ def chunked_gdn( dt_bias_f32 = dt_bias.reshape(1, 1, 1, -1).astype(jnp.float32) # NOTE: Any element-wise computations should occur before repeat. + # [seqs, num_kq_heads, chunk, kq_head_dim] if cfg.use_qk_norm_in_gdn: q_large = l2_norm(q_large) k_large = l2_norm(k_large) @@ -375,10 +604,11 @@ def chunked_gdn( ) out_list.append(out.swapaxes(0, 1)) state_list.append(state) - t_inv_list.append(t_inv) + if cfg.emit_t_inv: + t_inv_list.append(t_inv) out = jnp.stack(out_list, axis=0) state = jnp.stack(state_list, axis=0) - t_inv = jnp.stack(t_inv_list, axis=0) + t_inv = jnp.stack(t_inv_list, axis=0) if cfg.emit_t_inv else None return out, state, t_inv diff --git a/src/maxtext/kernels/gdn/config.py b/src/maxtext/kernels/gdn/config.py index 8ad3933295..a61ccf964c 100644 --- a/src/maxtext/kernels/gdn/config.py +++ b/src/maxtext/kernels/gdn/config.py @@ -26,6 +26,19 @@ DEFAULT_VMEM_LIMIT_FACTOR: float = 0.90 +# Optional GDNConfig fields and their defaults. +_TRI_INV_FIELD_DEFAULTS = { + "tri_inv_algo": "xor_dnc", + "head_group_size": None, + "emit_t_inv": True, +} + +# GDNConfig fields that change how the kernel lays out a value it computes either way, not what it computes or +# returns. They are left out of the kernel name and metadata, so the lowered program outside the kernel body is the +# same with and without them, and a profile shows the same kernel name for both variants (the run's config tells +# them apart). +_LAYOUT_ONLY_FIELDS = ("conv_dense_token_stack",) + class GDNMode(enum.StrEnum): """Execution mode for GDN kernel.""" @@ -82,6 +95,12 @@ class GDNConfig: num_buffers: int = 2 has_seg_ids: bool = False use_qk_norm_in_gdn: bool = True + tri_inv_algo: str = "xor_dnc" + head_group_size: int | None = None + emit_t_inv: bool = True + # Build the Conv1D output from dense per-token views (compute_conv1d.stack_conv1d_tokens; config + # gdn_conv1d_dense_token_stack). GDNMode.PER_SEQ only. + conv_dense_token_stack: bool = False @property def chunk_size(self) -> int: @@ -115,10 +134,24 @@ def aligned_num_v_heads(self) -> int: return pl.cdiv(self.num_v_heads + extra_lanes, num_lanes) * num_lanes def get_kernel_name(self) -> str: - return f"fused_conv1d_gdn_{self.mode.value}_b{self.seq_tile_size}" f"_c{self.chunk_size}" + """Pallas kernel name; non-default triangular-inverse options are appended as suffixes.""" + name = f"fused_conv1d_gdn_{self.mode.value}_b{self.seq_tile_size}" f"_c{self.chunk_size}" + if self.tri_inv_algo != "xor_dnc": + name += f"_{self.tri_inv_algo}" + if self.head_group_size is not None: + name += f"_hg{self.head_group_size}" + if not self.emit_t_inv: + name += "_notinv" + return name def get_metadata(self) -> dict[str, str | int | float]: + """Flattened config fields for the kernel's metadata (knobs at their defaults and layout-only fields are left out).""" cfgs_dict = dataclasses.asdict(self) + for field_name, default in _TRI_INV_FIELD_DEFAULTS.items(): + if cfgs_dict.get(field_name) == default: + cfgs_dict.pop(field_name) + for field_name in _LAYOUT_ONLY_FIELDS: + cfgs_dict.pop(field_name) ret = {} for path, val in jax.tree_util.tree_leaves_with_path(cfgs_dict): key = jax.tree_util.keystr(path, simple=True, separator=".") diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py index 4ea358228b..d563c647e6 100644 --- a/src/maxtext/kernels/gdn/gdn_bwd/api.py +++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py @@ -15,6 +15,7 @@ """Model-facing API and custom VJP boundary for GDN backward kernel.""" import functools +import inspect from typing import Any, Optional, Tuple import jax @@ -23,8 +24,10 @@ import jax.numpy as jnp from .. import compute_conv1d as local_compute_conv1d +from .. import compute_gdn as local_compute_gdn from .. import wrapper as local_gdn_wrapper from . import cp_gdn +from . import compute_conv1d_bwd from .compute_conv1d_bwd import conv1d_silu_bwd from .compute_conv1d_bwd import conv1d_silu_fwd from .jax_compute_gdn_states import _compute_forward_conv_and_states @@ -61,15 +64,7 @@ def decoupled_conv1d_gdn_bwd_kernel( conv_halo_seg: Optional[jax.Array] = None, init_seg: Optional[jax.Array] = None, interpret: bool | pltpu.InterpretParams | None = None, -) -> Tuple[ - jax.Array, - jax.Array, - jax.Array, - jax.Array, - Optional[jax.Array], - jax.Array, - jax.Array, -]: +) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array, Optional[jax.Array], jax.Array, jax.Array,]: """Decoupled Conv1D + GDN backward combining Pallas GDN bwd and JAX Conv1D bwd.""" del seq_lens, qkv if segment_ids is not None and conv_halo_seg is None and init_seg is None: @@ -104,6 +99,7 @@ def decoupled_conv1d_gdn_bwd_kernel( segment_ids=segment_ids, init_seg=init_seg, interpret=interpret, + recompute_t_inv=t_inv is None, ) d_pre_conv_qkv, d_conv_weight, d_conv_bias = conv1d_silu_bwd( @@ -152,12 +148,13 @@ def _run_local_gdn_decoupled_fwd( segment_ids: Optional[jax.Array] = None, conv_halo_seg: Optional[jax.Array] = None, init_seg: Optional[jax.Array] = None, -) -> Tuple[ - Tuple[jax.Array, Tuple[jax.Array, jax.Array]], - Optional[jax.Array], - Optional[jax.Array], -]: - """Runs local GDN forward pass on TPU returning (t_inv, chunk_states), or pure JAX on CPU.""" + conv_dense_token_stack: bool = False, +) -> Tuple[Tuple[jax.Array, Tuple[jax.Array, jax.Array]], Optional[jax.Array], Optional[jax.Array],]: + """Runs local GDN forward pass on TPU returning (t_inv, chunk_states), or pure JAX on CPU. + + `conv_dense_token_stack` goes to the Pallas forward kernel (wrapper.fused_conv1d_gdn); the CPU fallback has no + in-VMEM Conv1D and does not use it. + """ batch_size, seq_len, dim_size = qkv.shape if seq_len % chunk_size != 0: raise ValueError( @@ -168,6 +165,8 @@ def _run_local_gdn_decoupled_fwd( fallback_reason = pallas_unsupported_reason(head_k_dim=head_k_dim, head_v_dim=head_v_dim, chunk_size=chunk_size) if not on_cpu and fallback_reason is not None: raise ValueError(f"GDN Pallas TPU forward kernel does not support {fallback_reason}.") + emit_t_inv = chunk_size <= 64 + head_group_size = local_compute_gdn.default_head_group_size(chunk_size, num_k_heads, num_v_heads, qkv.dtype) if on_cpu: out, states = pure_jax_decoupled_conv1d_gdn( qkv=qkv, @@ -213,6 +212,8 @@ def _run_local_gdn_decoupled_fwd( conv_halo_seg=conv_halo_seg, init_seg=init_seg, ) + if not emit_t_inv: + t_inv = None return (out, states), t_inv, chunk_states num_seqs = batch_size @@ -267,6 +268,9 @@ def _run_local_gdn_decoupled_fwd( d_k=head_k_dim, d_v=head_v_dim, kernel_size=conv_kernel_size, + # seq_len is a multiple of chunk_size (checked above), so the prefill kernel writes every + # output row and the output buffer needs no zero fill. + zero_initialize_out=False, compute_precision=jnp.dtype(compute_dtype), mixed_tile_size=chunk_size, is_prefill_only=True, @@ -274,10 +278,17 @@ def _run_local_gdn_decoupled_fwd( segment_ids=segment_ids, conv_halo_seg=conv_halo_seg, init_seg=init_seg, + tri_inv_algo=local_compute_gdn.DEFAULT_TRI_INV_ALGO, + head_group_size=head_group_size, + emit_t_inv=emit_t_inv, + conv_dense_token_stack=conv_dense_token_stack, ) core_attn_out = core_attn_out_flat.reshape(batch_size, seq_len, num_v_heads, head_v_dim) - t_inv = t_inv_raw.astype(jnp.float32).reshape(batch_size, num_chunks, num_v_heads, chunk_size, chunk_size) + if t_inv_raw is None: + t_inv = None + else: + t_inv = t_inv_raw.astype(jnp.float32).reshape(batch_size, num_chunks, num_v_heads, chunk_size, chunk_size) chunk_states = chunk_states_raw.astype(jnp.float32).reshape(batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim) return ( ( @@ -346,8 +357,9 @@ def _run_cp_gdn_decoupled_fwd_impl( segment_ids: Optional[jax.Array] = None, seg_metadata: Optional[Tuple[Optional[jax.Array], Optional[jax.Array], Optional[jax.Array]]] = None, cp_matmul_precision: jax.lax.Precision = cp_gdn.DEFAULT_PRECISION, + conv_dense_token_stack: bool = False, ): - """Runs 2-pass sequence-sharded CP forward for GDN.""" + """Runs 2-pass sequence-sharded CP forward for GDN. Both passes get `conv_dense_token_stack`.""" batch_size = qkv.shape[0] if seg_metadata is not None: s_enc_local, conv_halo_seg, init_seg = seg_metadata @@ -377,7 +389,7 @@ def _run_cp_gdn_decoupled_fwd_impl( conv_halo = checkpoint_name(conv_halo, "gdn_cp_conv_halo") zero_rs = jnp.zeros((batch_size, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) - # Pass 1: Local GDN with zero initial recurrent state -> yields t_inv and S_ext_local + # Pass 1: Local GDN with zero initial recurrent state -> yields t_inv (when chunk_size <= 64) and S_ext_local (_, states_pass1), t_inv, _ = _run_local_gdn_decoupled_fwd( qkv, b, @@ -399,6 +411,7 @@ def _run_cp_gdn_decoupled_fwd_impl( segment_ids=s_enc_local, conv_halo_seg=conv_halo_seg, init_seg=init_seg, + conv_dense_token_stack=conv_dense_token_stack, ) # Convolve only the K channel slice (1/6th of qkv) needed by compose_local_from_t_inv @@ -460,6 +473,7 @@ def _run_cp_gdn_decoupled_fwd_impl( segment_ids=s_enc_local, conv_halo_seg=conv_halo_seg, init_seg=init_seg, + conv_dense_token_stack=conv_dense_token_stack, ) if s_enc_local is not None: # Owner = last rank with a valid local token (rank 0 if none). A rank whose @@ -473,7 +487,22 @@ def _run_cp_gdn_decoupled_fwd_impl( return (out, states), t_inv_2, chunk_states, conv_halo, s_in_r, m_local -@functools.partial(jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16, 17, 19)) +def _conv_input_in_qkv_dtype_scoped(fn): + """Traces fn inside compute_conv1d_bwd.conv_input_in_qkv_dtype_scope().""" + signature = inspect.signature(fn) + + @functools.wraps(fn) # keeps fn's signature, which jax.custom_vjp resolves arguments against + def scoped(*args, **kwargs): + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(bound.arguments["conv_input_in_qkv_dtype"]): + return fn(*args, **kwargs) + + return scoped + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16, 17, 19, 20, 21)) +@_conv_input_in_qkv_dtype_scoped def gdn_decoupled_conv1d( qkv: jax.Array, b: jax.Array, @@ -495,11 +524,23 @@ def gdn_decoupled_conv1d( cp_axis_name: str | tuple[str, ...] | None = None, segment_ids: Optional[jax.Array] = None, cp_matmul_precision: jax.lax.Precision = cp_gdn.DEFAULT_PRECISION, + conv_input_in_qkv_dtype: bool = False, + conv_dense_token_stack: bool = False, ) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]: """Decoupled Conv1D + GDN with Pallas backward pass and optional sequence-sharded CP. `cp_matmul_precision` sets the precision of the f32 matmuls that compose the recurrent state across sequence-sharded CP ranks (`cp_gdn`); it is unused without CP. + + `conv_input_in_qkv_dtype` (config gdn_conv1d_input_in_qkv_dtype) makes every Conv1D + SiLU forward and backward + traced by this op and its forward and backward rules keep the halo-extended conv input in qkv's dtype instead of + materializing it in float32 (compute_conv1d_bwd.py). Same values (the conv weight gradient up to float32 + summation order); only HBM traffic changes. + + `conv_dense_token_stack` (config gdn_conv1d_dense_token_stack) makes every Pallas forward kernel this op and its + forward rule run (both sequence-sharded CP passes, and the forward rule's replay under remat) build its Conv1D + output from dense per-token views (compute_conv1d.stack_conv1d_tokens). Same values; only the layout the TPU + compiler gives that value changes. The backward rule never runs the forward kernel. """ if _is_cp_active(cp_axis_name): (out, states), _, _, _, _, _ = _run_cp_gdn_decoupled_fwd_impl( @@ -523,6 +564,7 @@ def gdn_decoupled_conv1d( cp_axis_name=cp_axis_name, segment_ids=segment_ids, cp_matmul_precision=cp_matmul_precision, + conv_dense_token_stack=conv_dense_token_stack, ) return out, states @@ -550,6 +592,7 @@ def gdn_decoupled_conv1d( segment_ids=segment_ids, conv_halo_seg=conv_halo_seg, init_seg=init_seg, + conv_dense_token_stack=conv_dense_token_stack, ) return out, states @@ -564,6 +607,7 @@ def _unwrap_cotangent(g, like_zero: bool = False): return g +@_conv_input_in_qkv_dtype_scoped def _gdn_decoupled_conv1d_fwd( qkv: jax.Array, b: jax.Array, @@ -585,6 +629,8 @@ def _gdn_decoupled_conv1d_fwd( cp_axis_name: str | tuple[str, ...] | None = None, segment_ids: Optional[jax.Array] = None, cp_matmul_precision: jax.lax.Precision = cp_gdn.DEFAULT_PRECISION, + conv_input_in_qkv_dtype: bool = False, + conv_dense_token_stack: bool = False, ): """Forward rule for custom_vjp registration of gdn_decoupled_conv1d.""" ( @@ -598,20 +644,18 @@ def _gdn_decoupled_conv1d_fwd( conv_state, recurrent_state, segment_ids, - ) = _unwrap_primal( - ( - qkv, - b, - a, - conv_weight, - conv_bias, - a_log, - dt_bias, - conv_state, - recurrent_state, - segment_ids, - ) - ) + ) = _unwrap_primal(( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + segment_ids, + )) qkv = checkpoint_name(qkv, "gdn_fwd_conv") qkv = checkpoint_name(qkv, "gdn_conv_out") has_initial_state = (conv_state is not None) or (recurrent_state is not None) @@ -644,6 +688,7 @@ def _gdn_decoupled_conv1d_fwd( segment_ids=segment_ids, seg_metadata=seg_metadata, cp_matmul_precision=cp_matmul_precision, + conv_dense_token_stack=conv_dense_token_stack, ) out = checkpoint_name(out, "gdn_core_attn_out") residuals = ( @@ -687,6 +732,7 @@ def _gdn_decoupled_conv1d_fwd( segment_ids=segment_ids, conv_halo_seg=conv_halo_seg, init_seg=init_seg, + conv_dense_token_stack=conv_dense_token_stack, ) out = checkpoint_name(out, "gdn_core_attn_out") seg_bundle = (segment_ids, conv_halo_seg, init_seg) if segment_ids is not None else None @@ -718,19 +764,74 @@ def _gdn_decoupled_conv1d_bwd( compute_dtype: jnp.dtype, cp_axis_name: str | tuple[str, ...] | None = None, cp_matmul_precision: Any = cp_gdn.DEFAULT_PRECISION, + conv_input_in_qkv_dtype: bool = False, + conv_dense_token_stack: bool = False, residuals: tuple[Any, ...] | None = None, cotangents: tuple[Any, ...] | None = None, ): """Backward rule for custom_vjp registration of gdn_decoupled_conv1d.""" # Support positional calls that omit the trailing nondiff arguments: 10 args (no cp_axis_name, - # no cp_matmul_precision) or 11 args (no cp_matmul_precision). + # no cp_matmul_precision, no conv_input_in_qkv_dtype, no conv_dense_token_stack), 11 args (no + # cp_matmul_precision, no conv_input_in_qkv_dtype, no conv_dense_token_stack), 12 args (no + # conv_input_in_qkv_dtype, no conv_dense_token_stack) or 13 args (no conv_dense_token_stack). + # jax.custom_vjp itself always passes all 12 nondiff arguments. if residuals is None and cotangents is None and isinstance(cp_axis_name, tuple) and len(cp_axis_name) >= 10: residuals, cotangents = cp_axis_name, cp_matmul_precision cp_axis_name, cp_matmul_precision = None, cp_gdn.DEFAULT_PRECISION - elif cotangents is None and isinstance(cp_matmul_precision, tuple) and len(cp_matmul_precision) >= 10: - residuals, cotangents = cp_matmul_precision, residuals - cp_matmul_precision = cp_gdn.DEFAULT_PRECISION + elif ( + residuals is None + and cotangents is None + and isinstance(cp_matmul_precision, tuple) + and len(cp_matmul_precision) >= 10 + ): + residuals, cotangents = cp_matmul_precision, conv_input_in_qkv_dtype + cp_matmul_precision, conv_input_in_qkv_dtype = cp_gdn.DEFAULT_PRECISION, False + elif ( + residuals is None + and cotangents is None + and isinstance(conv_input_in_qkv_dtype, tuple) + and len(conv_input_in_qkv_dtype) >= 10 + ): + residuals, cotangents = conv_input_in_qkv_dtype, conv_dense_token_stack + conv_input_in_qkv_dtype, conv_dense_token_stack = False, False + elif cotangents is None and isinstance(conv_dense_token_stack, tuple) and len(conv_dense_token_stack) >= 10: + residuals, cotangents = conv_dense_token_stack, residuals + conv_dense_token_stack = False + # The backward rule recomputes the Conv1D with JAX ops; it never runs the Pallas forward kernel. + del conv_dense_token_stack + + with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(conv_input_in_qkv_dtype): + return _gdn_decoupled_conv1d_bwd_impl( + num_k_heads, + num_v_heads, + head_k_dim, + head_v_dim, + conv_kernel_size, + chunk_size, + use_qk_norm_in_gdn, + compute_dtype, + cp_axis_name, + cp_matmul_precision, + residuals, + cotangents, + ) + +def _gdn_decoupled_conv1d_bwd_impl( + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, + cp_axis_name: str | tuple[str, ...] | None, + cp_matmul_precision: Any, + residuals: tuple[Any, ...], + cotangents: tuple[Any, ...], +): + """_gdn_decoupled_conv1d_bwd after its argument shim, traced inside its conv_input_in_qkv_dtype scope.""" m_local_fwd = None seg_slot = None include_segment_ids_grad = False @@ -848,9 +949,12 @@ def _gdn_decoupled_conv1d_bwd( has_initial_state=bool(has_user_conv_state) or bool(has_user_recurrent_state), ) - # Recompute forward chunk states and t_inv if not cached in residuals + # When chunk_states is cached in residuals and t_inv was omitted (chunk_size > 64), + # the Pallas backward kernel rebuilds T^-1 per chunk in VMEM using xor_dnc instead + # of reading it from HBM or recomputing it with XLA. + recompute_t_inv_in_kernel = chunk_states is not None and t_inv_fwd is None conv_out_cached = None - if chunk_states is None or t_inv_fwd is None: + if not recompute_t_inv_in_kernel and (chunk_states is None or t_inv_fwd is None): qkv_conv, chunk_states_recomputed, t_inv_recomputed = _compute_forward_conv_and_states( qkv=pre_conv_qkv, b=b, @@ -949,6 +1053,7 @@ def _gdn_decoupled_conv1d_bwd( init_seg=init_seg, d_recurrent_state=dht_local, return_dh0=need_dh0, + recompute_t_inv=recompute_t_inv_in_kernel, ) if need_dh0: dy_conv, d_b, d_a, d_a_log, d_dt_bias, dh0_local = bwd_out @@ -1034,6 +1139,7 @@ def _gdn_decoupled_conv1d_bwd( init_seg=init_seg, d_recurrent_state=d_recurrent_state, return_dh0=need_dh0, + recompute_t_inv=recompute_t_inv_in_kernel, ) # pylint: disable=unbalanced-tuple-unpacking if need_dh0: diff --git a/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py b/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py index 3d379e5cf7..ef29e24b9e 100644 --- a/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py +++ b/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py @@ -29,8 +29,12 @@ def make_bwd_block_specs( padded_num_v_heads: int | None = None, has_dht: bool = False, has_dh0: bool = False, + has_t_inv: bool = True, ) -> Tuple[list[pl.BlockSpec], list[pl.BlockSpec], int, int]: - """Constructs reverse-scan Pallas emit_pipeline in_specs and out_specs for GDN backward.""" + """Constructs reverse-scan Pallas emit_pipeline in_specs and out_specs for GDN backward. + + With `has_t_inv=False` (t_inv recomputed in-kernel) the t_inv input spec is omitted. + """ if padded_num_v_heads is None: padded_num_v_heads = ((num_v_heads + 127) // 128) * 128 @@ -63,11 +67,16 @@ def rc(c): (None, None, None, num_v_heads, kq_head_dim, v_head_dim), lambda b, g, c: (b, g, rc(c), 0, 0, 0), ), - # 5: t_inv [B, G, C, num_v_heads, chunk_size, chunk_size] - pl.BlockSpec( - (None, None, None, num_v_heads, chunk_size, chunk_size), - lambda b, g, c: (b, g, rc(c), 0, 0, 0), - ), + ] + if has_t_inv: + in_specs.append( + # 5: t_inv [B, G, C, num_v_heads, chunk_size, chunk_size] + pl.BlockSpec( + (None, None, None, num_v_heads, chunk_size, chunk_size), + lambda b, g, c: (b, g, rc(c), 0, 0, 0), + ) + ) + in_specs += [ # 6: a_log [B, G, 1, padded_num_v_heads] pl.BlockSpec( (None, None, 1, padded_num_v_heads), diff --git a/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py b/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py index e30fc6bf28..a17489e967 100644 --- a/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py +++ b/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py @@ -14,11 +14,40 @@ """Conv1D + SiLU forward and unrolled shift-multiply backward primitives.""" -from typing import Optional, Tuple +import contextlib +import contextvars +from typing import Iterator, Optional, Tuple import jax import jax.numpy as jnp +# gdn_conv1d_input_in_qkv_dtype (see the comment above _conv1d_silu_fwd_narrow_input at the end of this file). +# conv1d_silu_fwd / conv1d_silu_bwd take it as an argument; left at None they read this trace-time setting, which +# gdn_decoupled_conv1d (gdn_bwd/api.py) scopes around its primal, forward and backward rules from its own static +# argument. A scoped setting rather than an argument on every internal helper because the backward rule is traced +# when the cotangent is pulled back, after the model's forward has returned, so only the custom_vjp's static +# arguments reach it; and the rules reach these functions through several helpers (jax_compute_gdn_states, the CP +# forward) that do not otherwise need it. +_CONV_INPUT_IN_QKV_DTYPE: contextvars.ContextVar[bool] = contextvars.ContextVar( + "gdn_conv1d_input_in_qkv_dtype", default=False +) + + +@contextlib.contextmanager +def conv_input_in_qkv_dtype_scope(enabled: bool) -> Iterator[None]: + """Scopes the default of conv1d_silu_fwd / conv1d_silu_bwd's `conv_input_in_qkv_dtype` while tracing.""" + token = _CONV_INPUT_IN_QKV_DTYPE.set(bool(enabled)) + try: + yield + finally: + _CONV_INPUT_IN_QKV_DTYPE.reset(token) + + +def _resolve_conv_input_in_qkv_dtype(conv_input_in_qkv_dtype: Optional[bool]) -> bool: + if conv_input_in_qkv_dtype is None: + return _CONV_INPUT_IN_QKV_DTYPE.get() + return bool(conv_input_in_qkv_dtype) + def _build_conv_tap_masks( segment_ids: jax.Array, @@ -53,8 +82,17 @@ def conv1d_silu_fwd( segment_ids: Optional[jax.Array] = None, conv_halo_seg: Optional[jax.Array] = None, output_dtype: Optional[jnp.dtype] = None, + conv_input_in_qkv_dtype: Optional[bool] = None, ) -> Tuple[jax.Array, jax.Array]: - """Forward Conv1D + SiLU returning (conv_out, qkv_conv).""" + """Forward Conv1D + SiLU returning (conv_out, qkv_conv). + + `conv_input_in_qkv_dtype` (None: the conv_input_in_qkv_dtype_scope() setting, default False) keeps the halo-extended input + in qkv's dtype and casts each tap to float32 where it is read. Same values; see _conv1d_silu_fwd_narrow_input. + """ + if _resolve_conv_input_in_qkv_dtype(conv_input_in_qkv_dtype): + return _conv1d_silu_fwd_narrow_input( + qkv, conv_weight, conv_bias, kernel_size, conv_state, segment_ids, conv_halo_seg, output_dtype + ) _, seq_len, _ = qkv.shape if conv_weight.ndim == 3: conv_weight_3d = conv_weight.astype(jnp.float32) @@ -92,8 +130,26 @@ def conv1d_silu_bwd( return_d_conv_state: bool = False, segment_ids: Optional[jax.Array] = None, conv_halo_seg: Optional[jax.Array] = None, + conv_input_in_qkv_dtype: Optional[bool] = None, ): - """Dedicated Conv1D + SiLU backward pass using JAX primitives.""" + """Dedicated Conv1D + SiLU backward pass using JAX primitives. + + `conv_input_in_qkv_dtype` (None: the conv_input_in_qkv_dtype_scope() setting, default False) keeps the halo-extended input + in qkv's dtype and casts each tap to float32 where it is read. Same values; see _conv1d_silu_bwd_narrow_input. + """ + if _resolve_conv_input_in_qkv_dtype(conv_input_in_qkv_dtype): + return _conv1d_silu_bwd_narrow_input( + qkv, + conv_weight, + conv_bias, + dy, + kernel_size, + conv_out, + conv_state, + return_d_conv_state, + segment_ids, + conv_halo_seg, + ) _, seq_len, _ = qkv.shape if conv_weight.ndim == 3: conv_weight_3d = conv_weight.astype(jnp.float32) @@ -179,3 +235,129 @@ def conv1d_silu_bwd( return dx, dw, db, d_conv_state return dx, dw, db + + +# -------------------------------------------------------------------------------------------------------------------- +# gdn_conv1d_input_in_qkv_dtype: the halo-extended conv input in qkv's dtype. +# +# The default path above builds conv_input = [conv_state | qkv] (or qkv left-padded by kernel_size - 1 zeros) in +# float32 and slices the kernel_size taps out of it. That materializes a float32 copy of the whole input, one buffer +# of batch x (sequence + kernel_size - 1) x channels per GDN layer in the backward pass, written once and read +# kernel_size times. +# +# The functions below build conv_input in qkv's dtype (bf16 for bf16 qkv and conv_state; the promoted dtype if they +# differ) and cast each tap slice to float32 where it is read, so XLA fuses the cast into the consumer and the +# float32 copy is never written. Every product and sum is formed in the same order, on the same float32 values, as in +# the default path (bf16 -> float32 is exact), so op by op every output is bit-identical; only which tensors are +# materialized in HBM changes. Compiled (measured on CPU), the forward outputs and dx / d_conv_state / db stay +# bit-identical, while the weight gradient dw, a float32 reduction over batch x sequence, can come out in a different +# summation order when XLA fuses the tap's cast into the reduction (CPU: a few ulp, only when the pre-activation is +# recomputed without packing or a conv state). No jax.lax.optimization_barrier is used: the narrow buffer alone is +# what removes the copy. +# Tests: tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py. + + +def _conv_input_narrow(qkv, conv_state, kernel_size): + """Halo-extended conv input in the narrowest exact dtype (bf16 for bf16 qkv and conv_state).""" + if conv_state is not None: + cdt = jnp.promote_types(qkv.dtype, conv_state.dtype) + return jnp.concatenate([conv_state.astype(cdt), qkv.astype(cdt)], axis=1) + return jnp.pad(qkv, ((0, 0), (kernel_size - 1, 0), (0, 0))) + + +def _conv_weight_3d_f32(conv_weight): + return conv_weight.astype(jnp.float32) if conv_weight.ndim == 3 else conv_weight[:, None, :].astype(jnp.float32) + + +def _preact_from_narrow_input(conv_input, w3d, conv_bias, kernel_size, seq_len, tap_masks, valid_mask): + """f32 pre-activation z with the default path's op order: sum over taps k of (x_k * mask_k) * w_k, + b, * valid.""" + if tap_masks is not None: + z = sum( + (conv_input[:, k : k + seq_len, :].astype(jnp.float32) * tap_masks[k]) * w3d[k, 0, :] for k in range(kernel_size) + ) + if conv_bias is not None: + z = z + conv_bias.astype(jnp.float32) + return z * valid_mask + z = sum(conv_input[:, k : k + seq_len, :].astype(jnp.float32) * w3d[k, 0, :] for k in range(kernel_size)) + if conv_bias is not None: + z = z + conv_bias.astype(jnp.float32) + return z + + +def _conv1d_silu_fwd_narrow_input( + qkv, conv_weight, conv_bias, kernel_size, conv_state, segment_ids, conv_halo_seg, output_dtype +): + """conv1d_silu_fwd reading the conv input in qkv's dtype. Returns (f32 conv_out, qkv_conv) like the default path.""" + _, seq_len, _ = qkv.shape + w3d = _conv_weight_3d_f32(conv_weight) + conv_input = _conv_input_narrow(qkv, conv_state, kernel_size) + valid_mask = tap_masks = None + if segment_ids is not None: + valid_mask, tap_masks = _build_conv_tap_masks(segment_ids, kernel_size, conv_halo_seg) + conv_out = _preact_from_narrow_input(conv_input, w3d, conv_bias, kernel_size, seq_len, tap_masks, valid_mask) + qkv_conv = jax.nn.silu(conv_out) + return conv_out, qkv_conv.astype(output_dtype or qkv.dtype) + + +def _conv1d_silu_bwd_narrow_input( + qkv, conv_weight, conv_bias, dy, kernel_size, conv_out, conv_state, return_d_conv_state, segment_ids, conv_halo_seg +): + """conv1d_silu_bwd reading the conv input in qkv's dtype; the same steps 1-4 as the default path.""" + _, seq_len, _ = qkv.shape + w3d = _conv_weight_3d_f32(conv_weight) + conv_input = _conv_input_narrow(qkv, conv_state, kernel_size) + valid_mask = tap_masks = None + if segment_ids is not None: + valid_mask, tap_masks = _build_conv_tap_masks(segment_ids, kernel_size, conv_halo_seg) + if conv_out is None: + conv_out = _preact_from_narrow_input(conv_input, w3d, conv_bias, kernel_size, seq_len, tap_masks, valid_mask) + z = conv_out.astype(jnp.float32) + + # dz = dy * SiLU'(z), as in the default path. + sig_z = jax.nn.sigmoid(z) + silu_prime = sig_z * (1.0 + z * (1.0 - sig_z)) + dz = dy.astype(jnp.float32) * silu_prime + if valid_mask is not None: + dz = dz * valid_mask + + if conv_bias is not None: + db = jnp.sum(dz, axis=(0, 1)).astype(conv_bias.dtype) + if conv_bias.ndim != 1: + db = db.reshape(conv_bias.shape) + else: + db = None + + # dw[k] = sum_{b,t} dz[b,t] * conv_input[b, t+k], each tap cast to float32 where it is read. + dw_rows = [] + for k in range(kernel_size): + x_k = conv_input[:, k : k + seq_len, :].astype(jnp.float32) + dz_k = dz * tap_masks[k] if tap_masks is not None else dz + dw_rows.append(jnp.sum(dz_k * x_k, axis=(0, 1))) + dw = jnp.stack(dw_rows, axis=0) + dw = dw[:, None, :].astype(conv_weight.dtype) if conv_weight.ndim == 3 else dw.astype(conv_weight.dtype) + + # dx: transposed convolution of dz with the reversed weights (does not read conv_input). + dz_pad = jnp.pad(dz, ((0, 0), (0, kernel_size - 1), (0, 0))) + w_rev = w3d[::-1] + if tap_masks is not None: + tap_masks_pad = [jnp.pad(tm, ((0, 0), (0, kernel_size - 1), (0, 0))) for tm in tap_masks] + dx = sum( + (dz_pad[:, j : j + seq_len, :] * tap_masks_pad[kernel_size - 1 - j][:, j : j + seq_len, :]) * w_rev[j, 0, :] + for j in range(kernel_size) + ) + else: + dx = sum(dz_pad[:, k : k + seq_len, :] * w_rev[k, 0, :] for k in range(kernel_size)) + dx = dx.astype(qkv.dtype) + + if return_d_conv_state: + d_cs_list = [] + for m in range(kernel_size - 1): + val = jnp.zeros((qkv.shape[0], qkv.shape[2]), dtype=jnp.float32) + for t in range(m + 1): + k_idx = m - t + dz_t = dz[:, t, :] * tap_masks[k_idx][:, t, :] if tap_masks is not None else dz[:, t, :] + val = val + dz_t * w3d[k_idx, 0, :] + d_cs_list.append(val) + d_conv_state = jnp.stack(d_cs_list, axis=1).astype(qkv.dtype) + return dx, dw, db, d_conv_state + return dx, dw, db diff --git a/src/maxtext/kernels/gdn/gdn_bwd/cp_gdn.py b/src/maxtext/kernels/gdn/gdn_bwd/cp_gdn.py index 4958bfccde..95358a2da4 100644 --- a/src/maxtext/kernels/gdn/gdn_bwd/cp_gdn.py +++ b/src/maxtext/kernels/gdn/gdn_bwd/cp_gdn.py @@ -36,12 +36,49 @@ import jax.numpy as jnp from .. import compute_conv1d as local_compute_conv1d +from .. import compute_gdn as local_compute_gdn # Precision of the f32 composition matmuls below: 6 bf16 passes on TPU. `gdn_cp_matmul_precision=high` # passes Precision.HIGH (3 passes) to them through their `precision` argument instead. DEFAULT_PRECISION = jax.lax.Precision.HIGHEST +def _compute_t_inv_batched( + k_orig: jax.Array, # [B, N_c, C, H_k, d_k] f32, already qk-normed if enabled + beta_h: jax.Array, # [N_c, B, H_v, C] f32, already masked if segment_ids + cumsum_h: jax.Array, # [N_c, B, H_v, C] f32 + seg_c: Optional[jax.Array], # [N_c, B, 1, C] f32 or None + valid_c: Optional[jax.Array], # [N_c, B, 1, C] bool or None + *, + repeats: int, + precision: jax.lax.Precision = DEFAULT_PRECISION, +) -> jax.Array: + """Computes `t_inv_c: [N_c, B, H_v, C, C]` across all local chunks in one batched call.""" + num_chunks, batch, num_v_heads, chunk_size = beta_h.shape + k_kq = jnp.transpose(k_orig, (1, 0, 3, 2, 4)) # [N_c, B, H_k, C, d_k] + if valid_c is not None: + k_kq = jnp.where(valid_c[..., None], k_kq, 0.0) + kk_kq = jnp.matmul(k_kq, jnp.swapaxes(k_kq, -1, -2), precision=precision) # [N_c, B, H_k, C, C] + kk_h = jnp.repeat(kk_kq, repeats, axis=2) # [N_c, B, H_v, C, C] + + mask_strict_base = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32), k=-1) + if seg_c is not None and valid_c is not None: + same_valid = (jnp.abs(seg_c[..., :, None] - seg_c[..., None, :]) < 0.5) & valid_c[..., :, None] + mask_strict = mask_strict_base[None, None, None, :, :] * same_valid.astype(jnp.float32) + else: + mask_strict = mask_strict_base + diff = cumsum_h[..., :, None] - cumsum_h[..., None, :] + safe_diff = jnp.where(mask_strict > 0.5, diff, -1e4) + g_mat_strict = jnp.exp(safe_diff) * mask_strict + + t_lower = kk_h * (beta_h[..., :, None] * g_mat_strict) + eye = jnp.eye(chunk_size, dtype=jnp.float32) + t_full = eye + t_lower + t_flat = t_full.reshape(num_chunks * batch * num_v_heads, chunk_size, chunk_size) + t_inv_flat = local_compute_gdn.invert_triangular_matrix(t_flat, precision=precision) + return t_inv_flat.reshape(num_chunks, batch, num_v_heads, chunk_size, chunk_size) + + def gather_cp_segment_metadata( segment_ids: Optional[jax.Array], cp_axis: str | tuple[str, ...], @@ -279,7 +316,7 @@ def compose_local_from_t_inv( a: jax.Array, a_log: jax.Array, dt_bias: jax.Array, - t_inv: jax.Array, + t_inv: Optional[jax.Array] = None, *, num_k_heads: int, num_v_heads: int, @@ -324,6 +361,7 @@ def compose_local_from_t_inv( dt_bias_f32 = dt_bias.astype(jnp.float32)[None, None, :, None] log_g = -jnp.exp(a_log_f32) * jax.nn.softplus(a_h + dt_bias_f32) + seg_c = None valid_c = None if segment_ids is not None: s_enc = local_compute_conv1d.encode_segment_ids(segment_ids.reshape(batch, seq_len), init_seg=init_seg) @@ -364,7 +402,18 @@ def compose_local_from_t_inv( k_beta_g = k_h * beta_h[..., None] * gating_forward k_scaled_bwd = k_h * gating_backward - t_inv_c = jnp.transpose(t_inv.astype(jnp.float32), (1, 0, 2, 3, 4)) # [N_c, B, H_v, C, C] + if t_inv is None: + t_inv_c = _compute_t_inv_batched( + k_orig, + beta_h, + cumsum_h, + seg_c, + valid_c, + repeats=repeats, + precision=precision, + ) + else: + t_inv_c = jnp.transpose(t_inv.astype(jnp.float32), (1, 0, 2, 3, 4)) # [N_c, B, H_v, C, C] # Pre-batched GEMMs across all chunks outside any sequential loop w_all = jnp.matmul(t_inv_c, k_beta_g, precision=precision) # [N_c, B, H_v, C, d_k] @@ -422,7 +471,7 @@ def compose_bwd_local_from_t_inv( a_log: jax.Array, dt_bias: jax.Array, do: jax.Array, - t_inv: jax.Array, + t_inv: Optional[jax.Array] = None, *, num_k_heads: int, num_v_heads: int, @@ -478,6 +527,8 @@ def compose_bwd_local_from_t_inv( dt_bias_f32 = dt_bias.astype(jnp.float32)[None, None, :, None] log_g = -jnp.exp(a_log_f32) * jax.nn.softplus(a_h + dt_bias_f32) + seg_c = None + valid_c = None mask_causal_base = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32), k=0) if segment_ids is not None: s_enc = local_compute_conv1d.encode_segment_ids(segment_ids.reshape(batch, seq_len), init_seg=init_seg) @@ -523,7 +574,18 @@ def compose_bwd_local_from_t_inv( q_g = q_h * gating_forward k_beta_g = k_h * beta_h[..., None] * gating_forward k_scaled_bwd = k_h * gating_backward - t_inv_c = jnp.transpose(t_inv.astype(jnp.float32), (1, 0, 2, 3, 4)) + if t_inv is None: + t_inv_c = _compute_t_inv_batched( + k_orig, + beta_h, + cumsum_h, + seg_c, + valid_c, + repeats=repeats, + precision=precision, + ) + else: + t_inv_c = jnp.transpose(t_inv.astype(jnp.float32), (1, 0, 2, 3, 4)) diff = cumsum_h[..., :, None] - cumsum_h[..., None, :] safe_diff_causal = jnp.where(mask_causal > 0.5, diff, -1e4) diff --git a/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py b/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py index f4c490af8b..b5b97e0454 100644 --- a/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py +++ b/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py @@ -24,6 +24,7 @@ import jax.numpy as jnp from .. import compute_conv1d as local_compute_conv1d +from .. import compute_gdn as local_compute_gdn from .bwd_memory_ref import make_bwd_block_specs from .runtime_utils import ensure_cpu_interpret_registered from .runtime_utils import pallas_unsupported_reason @@ -121,6 +122,10 @@ class GDNBackwardConfig: vmem_limit_mb: Optional[int] = None use_qk_norm_in_gdn: bool = False has_seg_ids: bool = False + # When True, t_inv is not an input: each chunk rebuilds T and inverts it in VMEM + # with `tri_inv_algo`. + recompute_t_inv: bool = False + tri_inv_algo: str = local_compute_gdn.DEFAULT_TRI_INV_ALGO @property def repeats(self) -> int: @@ -139,23 +144,37 @@ def _bwd_gdn_pipeline_body( has_dh0: bool = False, ) -> None: """Inner kernel executed per (batch, group, chunk) by emit_pipeline with manual GDN backward.""" - if len(refs) == 17: + if len(refs) == 17 and not cfg.recompute_t_inv: has_dht = True has_dh0 = True idx = 0 # pylint: disable=unbalanced-tuple-unpacking - ( - qkv_conv_ref, - b_ref, - a_ref, - do_ref, - chunk_states_ref, - t_inv_ref, - a_log_ref, - dt_bias_ref, - reset_ref, - ) = refs[idx : idx + 9] - idx += 9 + if cfg.recompute_t_inv: + ( + qkv_conv_ref, + b_ref, + a_ref, + do_ref, + chunk_states_ref, + a_log_ref, + dt_bias_ref, + reset_ref, + ) = refs[idx : idx + 8] + t_inv_ref = None + idx += 8 + else: + ( + qkv_conv_ref, + b_ref, + a_ref, + do_ref, + chunk_states_ref, + t_inv_ref, + a_log_ref, + dt_bias_ref, + reset_ref, + ) = refs[idx : idx + 9] + idx += 9 if has_dht: dht_ref = refs[idx] idx += 1 @@ -214,7 +233,7 @@ def _init(): a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32) do_val = do_ref[...].astype(jnp.float32) state_prev = chunk_states_ref[...].astype(jnp.float32) - t_inv_val = t_inv_ref[...].astype(jnp.float32) + t_inv_val = t_inv_ref[...].astype(jnp.float32) if t_inv_ref is not None else None a_log_val = a_log_ref[...][0, :num_v_heads].astype(jnp.float32) dt_bias_val = dt_bias_ref[...][0, :num_v_heads].astype(jnp.float32) @@ -266,8 +285,10 @@ def _init(): q_scaled = jnp.where(valid_c[:, :, None], q_scaled, 0.0) k_scaled_val = jnp.where(valid_c[:, :, None], k_scaled_val, 0.0) - q_rep = jnp.repeat(q_scaled, repeats, axis=1) - k_rep = jnp.repeat(k_scaled_val, repeats, axis=1) + q_kq = jnp.transpose(q_scaled, (1, 0, 2)) + k_kq = jnp.transpose(k_scaled_val, (1, 0, 2)) + q_h = jnp.repeat(q_kq, repeats, axis=0) + k_h = jnp.repeat(k_kq, repeats, axis=0) beta = jax.nn.sigmoid(b_val) sp_input = a_val + dt_bias_val @@ -281,19 +302,15 @@ def _init(): cumsum_log_g = jnp.dot(mask_cumsum, log_g, precision=jax.lax.Precision.HIGHEST) - q_h = jnp.transpose(q_rep, (1, 0, 2)) - k_h = jnp.transpose(k_rep, (1, 0, 2)) v_h = jnp.transpose(v, (1, 0, 2)) beta_h = jnp.transpose(beta, (1, 0)) cumsum_h = jnp.transpose(cumsum_log_g, (1, 0)) do_h = jnp.transpose(do_val, (1, 0, 2)) diff = cumsum_h[:, :, None] - cumsum_h[:, None, :] - safe_diff_strict = jnp.where(mask_strict[None, :, :] > 0.5, diff, -1e4) - g_mat_strict = jnp.exp(safe_diff_strict) * mask_strict[None, :, :] - safe_diff_causal = jnp.where(mask_causal[None, :, :] > 0.5, diff, -1e4) g_mat_causal = jnp.exp(safe_diff_causal) * mask_causal[None, :, :] + g_mat_strict = g_mat_causal * mask_strict[None, :, :] gating_forward = jnp.exp(cumsum_h)[:, :, None] gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None] @@ -306,9 +323,21 @@ def _init(): gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None] k_beta = k_h * beta_h[:, :, None] - k_h_T = jnp.swapaxes(k_h, -1, -2) - - A = t_inv_val + k_kq_T = jnp.swapaxes(k_kq, -1, -2) + + if cfg.recompute_t_inv: + # Rebuild T = I + strictly_lower(beta_i k_i.k_j exp(g_i - g_j)) and invert it here + # instead of reading the forward kernel's t_inv from HBM. Compute both q @ k^T and + # k @ k^T at num_kq_heads (4x fewer 3-pass bf16 matmuls) before repeating to num_v_heads. + kk_kq = _gdn_matmul(k_kq, k_kq_T, compute_dtype, precise=True) + t_lower = jnp.repeat(kk_kq, repeats, axis=0) * (beta_h[:, :, None] * g_mat_strict) + t_rows = jax.lax.broadcasted_iota(jnp.int32, t_lower.shape, 1) + t_cols = jax.lax.broadcasted_iota(jnp.int32, t_lower.shape, 2) + t_full = jnp.where(t_rows == t_cols, 1.0, t_lower) + inv_prec = jax.lax.Precision.HIGHEST if jnp.dtype(compute_dtype) == jnp.float32 else jax.lax.Precision.DEFAULT + A = local_compute_gdn.invert_triangular_matrix(t_full, precision=inv_prec, algo=cfg.tri_inv_algo).astype(jnp.float32) + else: + A = t_inv_val v_beta = v_h * beta_h[:, :, None] k_beta_g = k_beta * gating_forward @@ -331,7 +360,8 @@ def _init(): ) q_g = q_h * gating_forward - attn_unmasked = _gdn_matmul(q_h, k_h_T, compute_dtype, precise=True) + attn_kq = _gdn_matmul(q_kq, k_kq_T, compute_dtype, precise=True) + attn_unmasked = jnp.repeat(attn_kq, repeats, axis=0) attn = attn_unmasked * g_mat_causal k_scaled_bwd = k_h * gating_backward @@ -484,7 +514,7 @@ def _pallas_gdn_bwd_kernel_single_group( dt_bias: jax.Array, do: jax.Array, chunk_states: jax.Array, - t_inv: jax.Array, + t_inv: Optional[jax.Array], *, cfg: GDNBackwardConfig, segment_ids: Optional[jax.Array] = None, @@ -517,6 +547,8 @@ def _pallas_gdn_bwd_kernel_single_group( return_dh0=return_dh0, interpret=interpret, name=name, + recompute_t_inv=cfg.recompute_t_inv, + tri_inv_algo=cfg.tri_inv_algo, ) @@ -528,7 +560,7 @@ def pallas_gdn_bwd_kernel( dt_bias: jax.Array, do: jax.Array, chunk_states: jax.Array, - t_inv: jax.Array, + t_inv: Optional[jax.Array], *, num_v_heads: int, kq_head_dim: int, @@ -543,12 +575,18 @@ def pallas_gdn_bwd_kernel( return_dh0: bool = False, interpret: bool | pltpu.InterpretParams | None = None, name: str = "gdn_bwd_kernel", + recompute_t_inv: bool = False, + tri_inv_algo: str = local_compute_gdn.DEFAULT_TRI_INV_ALGO, ): """Executes the Pallas reverse-chunk GDNv3 backward kernel using emit_pipeline. Dispatches a single kernel call with 3D grid=(batch_size, num_groups, - num_chunks). + num_chunks). With `recompute_t_inv=True`, `t_inv` is ignored (may be None) and + each chunk rebuilds T^-1 in VMEM with `tri_inv_algo`. """ + if t_inv is None and not recompute_t_inv: + raise ValueError("pallas_gdn_bwd_kernel needs t_inv unless recompute_t_inv=True.") + local_compute_gdn.validate_tri_inv_algo(tri_inv_algo) if interpret is None: on_cpu = target_platform(qkv_conv) == "cpu" fallback_reason = pallas_unsupported_reason(head_k_dim=kq_head_dim, head_v_dim=v_head_dim, chunk_size=chunk_size) @@ -579,10 +617,14 @@ def pallas_gdn_bwd_kernel( _, init_seg = local_compute_conv1d.initial_state_segment_metadata(segment_ids, 4) if head_tile is not None: target_tile = head_tile - elif jnp.dtype(qkv_conv.dtype) == jnp.float32: - target_tile = _F32_HEAD_TILE else: - target_tile = _BF16_HEAD_TILE + if jnp.dtype(qkv_conv.dtype) == jnp.float32: + target_tile = _F32_HEAD_TILE + else: + target_tile = _BF16_HEAD_TILE + if chunk_size > 64: + # ~10 [tile, C, C] f32 intermediates are live per step; keep them at the C=64 budget. + target_tile = max(1, target_tile * 64 // chunk_size) max_possible = min(num_v_heads, target_tile) tile_v_heads = None @@ -615,6 +657,8 @@ def pallas_gdn_bwd_kernel( vmem_limit_mb=vmem_limit_mb, use_qk_norm_in_gdn=use_qk_norm_in_gdn, has_seg_ids=has_seg_ids, + recompute_t_inv=recompute_t_inv, + tri_inv_algo=tri_inv_algo, ) padded_tile_v_heads = cfg.padded_num_v_heads @@ -680,18 +724,20 @@ def pallas_gdn_bwd_kernel( ).swapaxes(1, 2) # 5. Prepare t_inv into (B, G, C, tile_v_heads, chunk_size, chunk_size) - t_inv_6d = ( - t_inv.reshape( - batch_size, - num_chunks, - num_groups, - tile_v_heads, - chunk_size, - chunk_size, - ) - .swapaxes(1, 2) - .astype(jnp.float32) - ) + t_inv_6d = None + if not recompute_t_inv: + t_inv_6d = ( + t_inv.reshape( + batch_size, + num_chunks, + num_groups, + tile_v_heads, + chunk_size, + chunk_size, + ) + .swapaxes(1, 2) + .astype(jnp.float32) + ) # 6. Prepare a_log and dt_bias into (B, G, 1, padded_tile_v_heads) a_log_2d = a_log.reshape(num_groups, tile_v_heads) @@ -744,7 +790,10 @@ def pallas_gdn_bwd_kernel( a_5d, do_6d, chunk_states_6d, - t_inv_6d, + ] + if t_inv_6d is not None: + in_args.append(t_inv_6d) + in_args += [ a_log_4d, dt_bias_4d, reset_hbm_5d, @@ -774,7 +823,10 @@ def pallas_gdn_bwd_kernel( padded_num_v_heads=padded_tile_v_heads, has_dht=has_dht, has_dh0=has_dh0, + has_t_inv=not recompute_t_inv, ) + if recompute_t_inv and name == "gdn_bwd_kernel": + name = f"gdn_bwd_kernel_recompute_{tri_inv_algo}" out_shapes_list = [ jax.ShapeDtypeStruct( diff --git a/src/maxtext/kernels/gdn/gdn_bwd/runtime_utils.py b/src/maxtext/kernels/gdn/gdn_bwd/runtime_utils.py index b70a0bf07f..6bdaaba0e1 100644 --- a/src/maxtext/kernels/gdn/gdn_bwd/runtime_utils.py +++ b/src/maxtext/kernels/gdn/gdn_bwd/runtime_utils.py @@ -21,6 +21,9 @@ from .. import compute_gdn as local_compute_gdn +# Chunk sizes supported by the Pallas TPU GDN forward and backward kernels. +_SUPPORTED_PALLAS_CHUNK_SIZES = (64, 128) + def target_platform(*arrays: Any) -> str: """Returns the platform ('tpu', 'cpu', 'gpu') the computation is being compiled for. @@ -43,8 +46,8 @@ def pallas_unsupported_reason(*, head_k_dim: int, head_v_dim: int, chunk_size: i return f"head_k_dim={head_k_dim} is not a multiple of 128" if head_v_dim % 128 != 0: return f"head_v_dim={head_v_dim} is not a multiple of 128" - if chunk_size != 64: - return f"chunk_size={chunk_size} != 64" + if chunk_size not in _SUPPORTED_PALLAS_CHUNK_SIZES: + return f"chunk_size={chunk_size} not in {_SUPPORTED_PALLAS_CHUNK_SIZES}" return None @@ -67,7 +70,7 @@ def ensure_cpu_interpret_registered() -> None: @jax.custom_vjp def invert_triangular_matrix(t: jax.Array) -> jax.Array: - """Computes inverse of unit lower-triangular matrix using Tokamax block forward substitution.""" + """Computes inverse of unit lower-triangular matrix using `compute_gdn.invert_triangular_matrix`.""" return local_compute_gdn.invert_triangular_matrix(t, block_size=16, precision=jax.lax.Precision.HIGHEST) diff --git a/src/maxtext/kernels/gdn/memory_ref.py b/src/maxtext/kernels/gdn/memory_ref.py index f29d79d81d..6d9ce6492e 100644 --- a/src/maxtext/kernels/gdn/memory_ref.py +++ b/src/maxtext/kernels/gdn/memory_ref.py @@ -366,6 +366,7 @@ class ChunkStatesBufferedRef(BaseBufferedRef): """DMA buffer for caching intermediate recurrent chunk states.""" def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Initiates asynchronous chunk-state copy from current output VMEM slot.""" assert self.sem_sends is not None assert self.window_ref is not None slot = self.current_copy_out_slot @@ -379,8 +380,18 @@ def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): dst_ref.at[p_id * self.cfg.seq_tile_size + idx], sem, ).start() + if self.buffer_count == 1: + for idx in range(self.cfg.seq_tile_size): + pltpu.make_async_copy( + vmem_ref.at[idx], + vmem_ref.at[idx], + sem, + ).wait() def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Waits for asynchronous chunk-state output copy to complete.""" + if self.buffer_count == 1: + return assert self.sem_sends is not None assert self.window_ref is not None slot = self.current_wait_out_slot @@ -457,21 +468,31 @@ def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): vmem_ref = self.window_ref.at[slot] p_id = grid_indices[0] + total_dma_size = 0 for idx in range(self.cfg.seq_tile_size): record = self.metadata_ref.get_record(p_id, idx) is_last_tile = record.is_last_tile s_idx = record.s_idx state_idx = self.metadata_ref.s_idx_to_state_indices[s_idx] dma_size = jnp.where(is_last_tile, 1, 0) + total_dma_size += dma_size pltpu.make_async_copy( vmem_ref.at[pl.ds(idx, dma_size)], # pyrefly: ignore[missing-attribute] dst_ref.at[pl.ds(state_idx, dma_size)], sem, ).start() + if self.buffer_count == 1: + pltpu.make_async_copy( + vmem_ref.at[pl.ds(0, total_dma_size)], # pyrefly: ignore[missing-attribute] + vmem_ref.at[pl.ds(0, total_dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).wait() def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): """Waits for asynchronous state output copy to complete.""" + if self.buffer_count == 1: + return assert self.sem_sends is not None assert self.window_ref is not None slot = self.current_wait_out_slot @@ -530,7 +551,18 @@ def create_allocs( pipeline_mode=pipeline_mode, ) - qkv_spec = block_spec_partial(block_shape=qkv_shape) + large_fp32_buffer_count = ( + 1 + if (cfg.chunk_size > 64 and cfg.num_v_heads >= 64 and jnp.dtype(cfg.dtypes.act_in) == jnp.float32) + else pipeline_mode.buffer_count + ) + large_fp32_spec_partial = functools.partial( + pl.BlockSpec, + memory_space=pltpu.VMEM, + index_map=lambda i: (i,), + pipeline_mode=pl.Buffered(buffer_count=large_fp32_buffer_count, use_lookahead=False), + ) + qkv_spec = large_fp32_spec_partial(block_shape=qkv_shape) ba_spec = block_spec_partial(block_shape=ba_shape) in_buffered_partial = functools.partial( InBufferedRef.input, @@ -539,7 +571,14 @@ def create_allocs( cfg=cfg, metadata_ref=metadata_ref, ) - qkv_alloc = in_buffered_partial(spec=qkv_spec, dtype_or_type=qkv_ref) + qkv_alloc = InBufferedRef.input( + spec=qkv_spec, + dtype_or_type=qkv_ref, + buffer_count=large_fp32_buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) b_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=b_ref) a_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=a_ref) @@ -553,7 +592,7 @@ def create_allocs( ) conv_spec = block_spec_partial(block_shape=conv_shape) - recurrent_spec = block_spec_partial(block_shape=recurrent_shape) + recurrent_spec = large_fp32_spec_partial(block_shape=recurrent_shape) state_buffered_partial = functools.partial( StateBufferedRef.input_output, buffer_count=pipeline_mode.buffer_count, @@ -562,7 +601,14 @@ def create_allocs( metadata_ref=metadata_ref, ) conv_alloc = state_buffered_partial(spec=conv_spec, dtype_or_type=conv_state_ref) - recurrent_alloc = state_buffered_partial(spec=recurrent_spec, dtype_or_type=recurrent_state_ref) + recurrent_alloc = StateBufferedRef.input_output( + spec=recurrent_spec, + dtype_or_type=recurrent_state_ref, + buffer_count=large_fp32_buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) allocs = [ qkv_alloc, @@ -598,9 +644,9 @@ def create_allocs( cfg.v_head_dim, ) chunk_states_alloc = ChunkStatesBufferedRef.output( - spec=block_spec_partial(block_shape=chunk_states_shape), + spec=large_fp32_spec_partial(block_shape=chunk_states_shape), dtype_or_type=chunk_states_ref, - buffer_count=pipeline_mode.buffer_count, + buffer_count=large_fp32_buffer_count, use_lookahead=pipeline_mode.use_lookahead, cfg=cfg, metadata_ref=metadata_ref, diff --git a/src/maxtext/kernels/gdn/model_runner.py b/src/maxtext/kernels/gdn/model_runner.py index a1cb7e0b88..9de1221497 100644 --- a/src/maxtext/kernels/gdn/model_runner.py +++ b/src/maxtext/kernels/gdn/model_runner.py @@ -617,6 +617,10 @@ def run_gdn_kernel_layer( cfg, cp_size, layer.num_k_heads, model_mode, use_gdn_kernel=True ) cp_matmul_precision = jax.lax.Precision[getattr(cfg, "gdn_cp_matmul_precision", "highest").upper()] + # Conv1D halo input in qkv's dtype instead of a float32 copy (compute_conv1d_bwd.py). Same values. + conv_input_in_qkv_dtype = cfg.gdn_conv1d_input_in_qkv_dtype + # Forward kernel's Conv1D output built from dense per-token views (compute_conv1d.py). Same values. + conv_dense_token_stack = bool(getattr(cfg, "gdn_conv1d_dense_token_stack", False)) if use_head_sharded_cp: qkv = jnp.concatenate([query, key, value_raw], axis=3) @@ -641,6 +645,11 @@ def run_gdn_kernel_layer( qkv = jnp.where(mask.reshape(mask.shape + (1,) * (qkv.ndim - mask.ndim)), qkv, 0.0) a = jnp.where(mask[..., None], a, jnp.asarray(-1e4, dtype=a.dtype)) b = jnp.where(mask[..., None], b, jnp.asarray(-1e4, dtype=b.dtype)) + if model_mode == MODEL_MODE_TRAIN and not getattr(cfg, "packing", True): + # packing=False means one sequence per row, whose padding the mask above already makes a no-op, so the + # kernel's segment path adds work and changes nothing. Other modes keep it: it also ends the carried conv + # state at the last valid token. + decoder_segment_ids = None batch, seq_len = qkv.shape[:2] conv_kernel_size = cfg.gdn_conv_kernel_dim @@ -887,6 +896,8 @@ def shard_mapped_gdn( # pylint: disable=too-many-positional-arguments use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, compute_dtype=state_dtype, segment_ids=seg_val, + conv_input_in_qkv_dtype=conv_input_in_qkv_dtype, + conv_dense_token_stack=conv_dense_token_stack, ) if next_cs is not None: @@ -926,6 +937,8 @@ def shard_mapped_gdn( # pylint: disable=too-many-positional-arguments cp_axis_name=cp_axis_name if use_seq_sharded_cp else None, segment_ids=seg_val, cp_matmul_precision=cp_matmul_precision, + conv_input_in_qkv_dtype=conv_input_in_qkv_dtype, + conv_dense_token_stack=conv_dense_token_stack, ) if not isinstance(qkv, jax.core.Tracer): @@ -968,6 +981,8 @@ def shard_mapped_gdn( # pylint: disable=too-many-positional-arguments use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, compute_dtype=state_dtype, segment_ids=decoder_segment_ids, + conv_input_in_qkv_dtype=conv_input_in_qkv_dtype, + conv_dense_token_stack=conv_dense_token_stack, ) core_attn_out = checkpoint_name(core_attn_out, "gdn_core_attn_out") diff --git a/src/maxtext/kernels/gdn/wrapper.py b/src/maxtext/kernels/gdn/wrapper.py index 7b951c7cee..952aac211a 100644 --- a/src/maxtext/kernels/gdn/wrapper.py +++ b/src/maxtext/kernels/gdn/wrapper.py @@ -139,6 +139,8 @@ def inner_kernel( carry_recurrent_scratch_ref=carry_recurrent_scratch_ref, cfg=cfg, ) + if chunk_states_slot_ref is not None: + chunk_states_slot_ref[...] = prev_recurrent.astype(chunk_states_slot_ref.dtype) # Step 1: Conv1D. qkv_in_compact = qkv_slot_ref[...].astype(jnp.float32) @@ -218,14 +220,39 @@ def inner_kernel( # Store output, recurrent, and t_inv to vmem. out_slot_ref[...] = out.astype(out_slot_ref.dtype) recurrent_slot_ref[...] = new_recurrent_state.astype(recurrent_slot_ref.dtype) - t_inv_slot_ref[...] = t_inv.astype(t_inv_slot_ref.dtype) - if chunk_states_slot_ref is not None: - chunk_states_slot_ref[...] = prev_recurrent.astype(chunk_states_slot_ref.dtype) + if t_inv_slot_ref is not None: + t_inv_slot_ref[...] = t_inv.astype(t_inv_slot_ref.dtype) if carry_recurrent_scratch_ref is not None: carry_recurrent_scratch_ref[...] = new_recurrent_state +def inner_kernel_without_t_inv( + qkv_slot_ref: jax.Ref, + b_slot_ref: jax.Ref, + a_slot_ref: jax.Ref, + conv_state_slot_ref: jax.Ref, + recurrent_slot_ref: jax.Ref, + out_slot_ref: jax.Array, + *args, + cfg: config.GDNConfig, + **kwargs, +) -> None: + """`inner_kernel` for pipelines without a t_inv output (cfg.emit_t_inv=False).""" + inner_kernel( + qkv_slot_ref, + b_slot_ref, + a_slot_ref, + conv_state_slot_ref, + recurrent_slot_ref, + out_slot_ref, + None, + *args, + cfg=cfg, + **kwargs, + ) + + def outer_kernel( # Inputs. metadata_ref: memory_ref.MetadataRef, @@ -251,6 +278,7 @@ def outer_kernel( del conv_state_out_ref, recurrent_state_out_ref chunk_states_ref = args[0] if (len(args) > 0 and cfg.mode == config.GDNMode.PER_SEQ) else None + emit_t_inv = cfg.emit_t_inv allocs = memory_ref.create_allocs( metadata_ref=metadata_ref, @@ -261,7 +289,7 @@ def outer_kernel( conv_state_ref=conv_state_ref, recurrent_state_ref=recurrent_state_ref, cfg=cfg, - t_inv_ref=t_inv_ref, + t_inv_ref=t_inv_ref if emit_t_inv else None, chunk_states_ref=chunk_states_ref, ) qkv_alloc = allocs[0] @@ -270,18 +298,24 @@ def outer_kernel( conv_alloc = allocs[3] recurrent_alloc = allocs[4] out_alloc = allocs[5] - t_inv_alloc = allocs[6] - chunk_states_alloc = allocs[7] if len(allocs) > 7 else None + next_alloc = 6 + t_inv_alloc = None + if emit_t_inv: + t_inv_alloc = allocs[next_alloc] + next_alloc += 1 + chunk_states_alloc = allocs[next_alloc] if len(allocs) > next_alloc else None num_tiles = metadata_ref.num_tiles[...] - out_specs = [out_alloc.spec, t_inv_alloc.spec] + out_specs = [out_alloc.spec] + if t_inv_alloc is not None: + out_specs.append(t_inv_alloc.spec) if chunk_states_alloc is not None: out_specs.append(chunk_states_alloc.spec) pipeline_func = pltpu.emit_pipeline( body=functools.partial( - inner_kernel, + inner_kernel if emit_t_inv else inner_kernel_without_t_inv, cfg=cfg, ), grid=(num_tiles,), @@ -297,7 +331,7 @@ def outer_kernel( @pl.with_scoped(allocations=allocs) def _run(allocations): - out_args = [out_ref, t_inv_ref] + out_args = [out_ref, t_inv_ref] if emit_t_inv else [out_ref] if chunk_states_ref is not None: out_args.append(chunk_states_ref) pipeline_func( @@ -334,6 +368,10 @@ def _run(allocations): "compute_precision", "is_prefill_only", "use_qk_norm_in_gdn", + "tri_inv_algo", + "head_group_size", + "emit_t_inv", + "conv_dense_token_stack", ), ) def fused_conv1d_gdn( @@ -365,8 +403,24 @@ def fused_conv1d_gdn( segment_ids: jax.Array | None = None, conv_halo_seg: jax.Array | None = None, init_seg: jax.Array | None = None, -) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array, jax.Array]: - """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states).""" + tri_inv_algo: str = compute_gdn.DEFAULT_TRI_INV_ALGO, + head_group_size: int | None = None, + emit_t_inv: bool = True, + conv_dense_token_stack: bool = False, +) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array | None, jax.Array]: + """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states). + + When `head_group_size` is None, `compute_gdn.default_head_group_size` selects an + appropriate KQ-head group size for `chunk_size > 64`. With `emit_t_inv=False`, + the returned `t_inv` is None. + + `conv_dense_token_stack` (config gdn_conv1d_dense_token_stack) makes the per-sequence kernel build its Conv1D + output from dense per-token views instead of jnp.stack (compute_conv1d.stack_conv1d_tokens): the same values, in + a layout the TPU compiler keeps dense. It requires `dim` to be a multiple of + compute_conv1d.DENSE_TOKEN_STACK_ALIGNMENT. The batched (decode) kernel, which convolves one token per sequence, + does not take the option and keeps jnp.stack for its one-token rows; MaxText's GDN layer never runs that kernel, + since gdn_bwd/api.py calls this function prefill-only. + """ act_in_dtype = qkv.dtype act_out_dtype = qkv.dtype conv_out_dtype = conv_state.dtype @@ -390,6 +444,12 @@ def fused_conv1d_gdn( assert query_start_loc.shape == (num_seqs + 1,) assert state_indices.shape == (num_seqs,) assert distribution.shape == (3,) + if conv_dense_token_stack and dim % compute_conv1d.DENSE_TOKEN_STACK_ALIGNMENT: + raise ValueError( + f"conv_dense_token_stack (gdn_conv1d_dense_token_stack) requires the Conv1D width ({dim}) to be a multiple of" + f" {compute_conv1d.DENSE_TOKEN_STACK_ALIGNMENT}; with head-sharded GDN context parallelism this is the width" + " of one rank's heads." + ) num_lanes = pltpu.get_tpu_info().num_lanes packing = 4 // act_in_dtype.itemsize @@ -485,8 +545,20 @@ def call_kernel( ) -> tuple[jax.Array, ...]: if mode == config.GDNMode.BATCHED: tile_size = decode_tile_size + per_seq_kwargs = {} else: tile_size = mixed_tile_size + resolved_head_group_size = ( + compute_gdn.default_head_group_size(tile_size, n_kq, n_v, act_in_dtype) + if head_group_size is None + else head_group_size + ) + per_seq_kwargs = { + "tri_inv_algo": compute_gdn.validate_tri_inv_algo(tri_inv_algo), + "head_group_size": resolved_head_group_size, + "emit_t_inv": emit_t_inv, + "conv_dense_token_stack": conv_dense_token_stack, + } cfg = config.GDNConfig( mode=mode, @@ -507,6 +579,7 @@ def call_kernel( recurrent_state=in_recurrent_state.dtype, conv_state=in_conv_state.dtype, ), + **per_seq_kwargs, ) if mode == config.GDNMode.BATCHED: @@ -542,10 +615,14 @@ def call_kernel( input_output_aliases[len(metadata_obj) + 5] = 0 num_chunks = cfg.batch_size // cfg.chunk_size - t_inv_shape = jax.ShapeDtypeStruct( - (num_chunks, cfg.num_v_heads, cfg.chunk_size, cfg.chunk_size), - cfg.dtypes.compute, - ) + if cfg.emit_t_inv: + t_inv_shape = jax.ShapeDtypeStruct( + (num_chunks, cfg.num_v_heads, cfg.chunk_size, cfg.chunk_size), + cfg.dtypes.compute, + ) + else: + # Never written; keeps the pallas_call output arity unchanged. + t_inv_shape = jax.ShapeDtypeStruct((1, 1, 8, 128), cfg.dtypes.compute) if mode == config.GDNMode.PER_SEQ: chunk_states_shape = jax.ShapeDtypeStruct( @@ -628,5 +705,7 @@ def call_kernel( ).astype(conv_out_dtype) out_conv_state = out_conv_state.at[state_indices].set(masked_cs) out_recurrent_state = out_recurrent_state.astype(recurrent_out_dtype) + if not emit_t_inv: + t_inv = None return out_act, (out_conv_state, out_recurrent_state), t_inv, chunk_states diff --git a/src/maxtext/kernels/megablox/__init__.py b/src/maxtext/kernels/megablox/__init__.py index f02632f3cb..ff2a18776e 100644 --- a/src/maxtext/kernels/megablox/__init__.py +++ b/src/maxtext/kernels/megablox/__init__.py @@ -13,4 +13,4 @@ # limitations under the License. """Megablox kernel""" -from maxtext.kernels.megablox.ops import gmm +from maxtext.kernels.megablox.ops import TcUnsortCfg, gmm, gmm_silu diff --git a/src/maxtext/kernels/megablox/backend.py b/src/maxtext/kernels/megablox/backend.py index 1418e7bcea..97c201ae25 100644 --- a/src/maxtext/kernels/megablox/backend.py +++ b/src/maxtext/kernels/megablox/backend.py @@ -433,13 +433,13 @@ def _accum(is_last_k_tile): loaded_lhs = lhs[...] loaded_lhs = dataclasses.replace(loaded_lhs, qvalue=mask_k_rem_lhs(loaded_lhs.qvalue)) else: - loaded_lhs = mask_k_rem_lhs(lhs[...]) # pytype: disable=wrong-arg-types + loaded_lhs = mask_k_rem_lhs(lhs[...]) if isinstance(rhs, qpl.QArray): loaded_rhs = rhs[...] loaded_rhs = dataclasses.replace(loaded_rhs, qvalue=mask_k_rem_rhs(loaded_rhs.qvalue)) else: - loaded_rhs = mask_k_rem_rhs(rhs[...]) # pytype: disable=wrong-arg-types + loaded_rhs = mask_k_rem_rhs(rhs[...]) if transpose_rhs: dot_general_dims = (((1,), (1,)), ((), ())) @@ -609,7 +609,7 @@ def tgmm( group_offset = group_offset[None] group_sizes = _validate_args( lhs=lhs, rhs=rhs, group_sizes=group_sizes, expected_rhs_dims=2 # pyrefly: ignore[bad-argument-type] - ) # pyrefly: ignore[bad-argument-type] + ) # Gather shape information. k, m, n = (lhs.shape[0], lhs.shape[1], rhs.shape[1]) diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index 2ef92ace29..2ca87711fe 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -18,12 +18,13 @@ import dataclasses import functools -from typing import List, Literal, Tuple +from typing import List, Literal, NamedTuple, Tuple import jax import jax.numpy as jnp from maxtext.kernels.megablox import backend from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_gmm_kernel as gmm_v2 from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_tgmm_kernel as tgmm_v2 +from maxtext.kernels.ragged import ragged_sort_tc from maxtext.layers import quantizations import qwix import qwix.pallas as qpl @@ -43,6 +44,21 @@ ) +class TcUnsortCfg(NamedTuple): + """Static config to fuse the TC unsort-reduce (and optionally the EP combine) into the GMM custom VJP. + + The forward reduces the GMM's expert-slot rows back to token order; the backward quantizes the smaller + token-order cotangent first and gathers it into expert slots afterwards. + """ + + unpadded_cap: int # GMM output rows that hold expert slots; the rest is tile padding. + num_out_tokens: int + topk: int + blocks: tuple # ragged_sort_tc blocks: (gather_block_size, flatten_block_size, keep_3d, tokens_3d). + mask_padding: bool + combine_rs_cfg: tuple[str, int] | None = None # (ep_axis_name, seq_len): also fuse the EP psum_scatter. + + def gmm( lhs: jnp.ndarray, rhs: jnp.ndarray, @@ -76,8 +92,28 @@ def gmm( use_gmm_v2: bool = False, use_gmm_v2_heuristic_tiling: bool = False, partial_sum: jnp.ndarray | None = None, + return_rhs: bool = False, + return_lhs: bool = False, + fuse_dlhs_scale: bool = False, + use_dlhs_transpose_rhs: bool = False, + flat_fsdp_shard: tuple[int, int, int] | None = None, + out_is_3d: bool = False, + fuse_3d_dlhs_ps: bool = False, + tc_routing: object = None, + tc_unsort_cfg: TcUnsortCfg | None = None, ): - """Grouped matrix multiplication operation.""" + """Grouped matrix multiplication operation. + + A 3D lhs ([m, k // 128, 128], e.g. tokens from the TensorCore ragged sort) is + detected from its rank; `out_is_3d` writes the output as [m, n // 128, 128]. + Both are only supported by the tokamax gmm_v2 backend; the backward pass then + consumes / produces the matching 3D cotangents without relayout copies. With + `out_is_3d`, a forward `partial_sum` must also be 3D; with `return_lhs`, the + cotangent of the returned (3D) lhs is accumulated into the 3D dlhs in place. + """ + lhs_ndim = (lhs.qvalue if isinstance(lhs, qpl.QArray) else lhs).ndim + if (lhs_ndim == 3 or out_is_3d) and not (use_tokamax_backend and use_gmm_v2): + raise NotImplementedError("3D lhs / out requires use_tokamax_backend=True and use_gmm_v2=True.") if interpret is None: # Default to native (TPU) lowering. `jax.devices()[0]` is NOT the compile TARGET: # during train_compile the local backend is CPU (JAX_PLATFORMS=cpu) while the mesh @@ -107,13 +143,25 @@ def gmm( act_calibration_method="absmax", ) + if return_rhs and not (use_tokamax_backend and use_gmm_v2): + # The cotangent of the returned rhs is accumulated into drhs via tgmm_v2 partial_sum, + # which other backends do not support; they would silently drop it. + raise ValueError("return_rhs=True requires use_tokamax_backend=True and use_gmm_v2=True.") + + if return_lhs and not (use_tokamax_backend and use_gmm_v2): + # The cotangent of the returned lhs is accumulated into dlhs via gmm_v2 partial_sum, + # which other backends do not support; they would silently drop it. + raise ValueError("return_lhs=True requires use_tokamax_backend=True and use_gmm_v2=True.") + + lhs_dtype = preferred_element_type if isinstance(lhs, qpl.QArray) else lhs.dtype + rhs_dtype = preferred_element_type if isinstance(rhs, qpl.QArray) else rhs.dtype gmm_fwd_bwd = lambda *args: _gmm_fwd(*args)[0] # pylint: disable=C3001 gmm_fwd_bwd = jax.custom_vjp( gmm_fwd_bwd, - nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16), + nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 26), ) - gmm_fwd_bwd.defvjp(_gmm_fwd, functools.partial(_gmm_bwd, lhs.dtype, rhs.dtype)) - return gmm_fwd_bwd( + gmm_fwd_bwd.defvjp(_gmm_fwd, functools.partial(_gmm_bwd, lhs_dtype, rhs_dtype)) + out, lhs_out, rhs_out = gmm_fwd_bwd( lhs, rhs, group_sizes, @@ -132,7 +180,33 @@ def gmm( use_gmm_v2, use_gmm_v2_heuristic_tiling, partial_sum, + return_rhs, + return_lhs, + fuse_dlhs_scale, + use_dlhs_transpose_rhs, + flat_fsdp_shard, + out_is_3d, + fuse_3d_dlhs_ps, + tc_routing, + tc_unsort_cfg, ) + if return_lhs and return_rhs: + return out, lhs_out, rhs_out + if return_lhs: + return out, lhs_out + if return_rhs: + return out, rhs_out + return out + + +def _scale_like(scale: jnp.ndarray, x: jnp.ndarray, dtype) -> jnp.ndarray: + """Per-tensor scale -> scalar; per-row scale [m, 1, ...] -> [m, 1, ..., 1] broadcastable to x.""" + if scale.size == 1: + return scale.reshape(()).astype(dtype) + if x.ndim == 3 or (scale.ndim != x.ndim and scale.size == scale.shape[0]): + # Per-row scale whose rank may differ from x (3D [m, 1, 1] vs 2D [m, k] or vice versa). + return scale.reshape(scale.shape[0], *([1] * (x.ndim - 1))).astype(dtype) + return scale.astype(dtype) # ============================================================================== @@ -169,8 +243,17 @@ def _gmm_fwd( use_gmm_v2: bool = False, use_gmm_v2_heuristic_tiling: bool = False, partial_sum: jnp.ndarray | None = None, + return_rhs: bool = False, + return_lhs: bool = False, + fuse_dlhs_scale: bool = False, + use_dlhs_transpose_rhs: bool = False, + flat_fsdp_shard: tuple[int, int, int] | None = None, + out_is_3d: bool = False, + fuse_3d_dlhs_ps: bool = False, + tc_routing: object = None, + tc_unsort_cfg: TcUnsortCfg | None = None, ) -> tuple[ - jnp.ndarray, + tuple[jnp.ndarray, jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray], tuple[ jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray, @@ -179,24 +262,49 @@ def _gmm_fwd( jnp.ndarray | None, bool, bool, + object, ], ]: """Forward function for GMM VJP. - - lhs: [m, k] + - lhs: [m, k] (or [m, k // 128, 128]) - rhs: [g, k, n] if transpose_rhs=False. [g, n, k] if transpose_rhs=True + - out: [m, n], or [m, n // 128, 128] if out_is_3d + + Always returns (out, lhs, rhs) so the output structure does not depend on return_lhs/return_rhs; + gmm() drops lhs unless return_lhs=True and rhs unless return_rhs=True. """ + del use_dlhs_transpose_rhs # Only affects the custom backward pass. + # Track whether operands arrived as QArray (e.g. from token all-gather quantization # or gathered weights) to return matching cotangent containers in backward pass. lhs_is_qarray = isinstance(lhs, qpl.QArray) rhs_is_qarray = isinstance(rhs, qpl.QArray) + if flat_fsdp_shard is not None: + if rhs_is_qarray or transpose_rhs or rhs.ndim != 2 or not weight_gather_axes: + raise NotImplementedError("flat_fsdp_shard needs a plain 2D rhs block, transpose_rhs=False and weight_gather_axes.") + if any(axis_idx != 0 for _, axis_idx in weight_gather_axes): + raise ValueError("flat_fsdp_shard: the flat rhs is gathered on dim 0 only.") + if not ( + use_tokamax_backend + and quantization_rule + and quantization_rule.bwd_qtype + and quantization_rule.weight_calibration_method.startswith("fixed") + ): + raise NotImplementedError("flat_fsdp_shard requires the tokamax QAG path (fixed weight calibration, bwd_qtype).") # Quantize activation and weight if quantization_rule: # pyrefly: ignore[bad-assignment] lhs, rhs = _fwd_quantize_activation_and_weight( - lhs, rhs, quantization_rule, use_gmm_v2, use_manual_quantization, transpose_rhs + lhs, + rhs, + quantization_rule, + use_gmm_v2, + use_manual_quantization, + transpose_rhs, + rhs_is_flat2d=flat_fsdp_shard is not None, ) # Quantization All-Gather (QAG) for weight: only supported for following conditions @@ -209,7 +317,7 @@ def _gmm_fwd( and weight_gather_axes ): # pyrefly: ignore[bad-assignment] - rhs = _fwd_gather_weight(rhs, weight_gather_axes) + rhs = _fwd_gather_weight(rhs, weight_gather_axes, gathered_shape=flat_fsdp_shard) # Backend Execution Routing if use_tokamax_backend and not use_gmm_v2: @@ -226,6 +334,7 @@ def _gmm_fwd( partial_sum, transpose_rhs, quantization_rule, + out_is_3d=out_is_3d, ) else: out = _fwd_run_megablox( @@ -241,7 +350,20 @@ def _gmm_fwd( lhs_vma_axes, ) - return out, ( + if tc_unsort_cfg is not None: + cfg = tc_unsort_cfg + out = ragged_sort_tc._gather_reduce( # pylint: disable=protected-access + out[: cfg.unpadded_cap], tc_routing, cfg.num_out_tokens, cfg.topk, cfg.blocks + ) + if cfg.combine_rs_cfg is not None: + # Reduce-scatter the [tokens, D // 128, 128] combine over EP into the local [batch, seq, D]. + ep_axis_name, seq_len = cfg.combine_rs_cfg + out = out.reshape(-1, seq_len, *out.shape[1:]).astype(preferred_element_type) + out = jax.lax.psum_scatter(out, ep_axis_name, scatter_dimension=0, tiled=True) + out = out.reshape(*out.shape[:2], -1) + + del return_rhs, return_lhs, fuse_dlhs_scale, fuse_3d_dlhs_ps # Only used by gmm() and _gmm_bwd. + return (out, lhs, rhs), ( lhs, rhs, group_sizes, @@ -249,6 +371,7 @@ def _gmm_fwd( partial_sum, lhs_is_qarray, rhs_is_qarray, + tc_routing, ) @@ -259,10 +382,11 @@ def _fwd_quantize_activation_and_weight( use_gmm_v2: bool, use_manual_quantization: bool, transpose_rhs: bool, + rhs_is_flat2d: bool = False, ) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray]: """Handles act and weight quantization for GMM forward inputs.""" if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) and not use_gmm_v2: - lhs = qpl.quantize( # pyrefly: ignore[bad-assignment] + lhs = qpl.quantize( lhs, quantization_rule.act_qtype, channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0], @@ -272,17 +396,21 @@ def _fwd_quantize_activation_and_weight( if quantization_rule.weight_qtype and not isinstance(rhs, qpl.QArray): if not use_manual_quantization: - rhs = qpl.quantize( # pyrefly: ignore[bad-assignment] + rhs = qpl.quantize( rhs, quantization_rule.weight_qtype, # If only considering the fwd pass, we could also enable channelwise # axes for the group axis, i.e., [0, 1 or 2]. However, this makes the # bwd pass unable to reuse the scale easily. - channelwise_axes=([] if quantization_rule.disable_channelwise_axes else ([1] if transpose_rhs else [2])), + channelwise_axes=( + [] + if quantization_rule.disable_channelwise_axes + else ([1] if (transpose_rhs or rhs_is_flat2d) else [2]) # flat 2D [rows, cols]: cols is the N axis + ), calibration_method=quantization_rule.weight_calibration_method, ) else: - rhs = quantizations.manual_quantize( # pyrefly: ignore[bad-assignment] + rhs = quantizations.manual_quantize( rhs, quantization_rule.weight_qtype, calibration_method=quantization_rule.weight_calibration_method, @@ -290,13 +418,34 @@ def _fwd_quantize_activation_and_weight( return lhs, rhs -def _fwd_gather_weight(rhs: qpl.QArray, weight_gather_axes: List[Tuple[str, int]]) -> qpl.QArray: - """Applies QAG (Quantization All-Gather) to RHS weights during forward pass.""" +def _fwd_gather_weight( + rhs: qpl.QArray, + weight_gather_axes: List[Tuple[str, int]], + gathered_shape: tuple[int, ...] | None = None, +) -> qpl.QArray: + """Applies QAG (Quantization All-Gather) to RHS weights during forward pass. + + With `gathered_shape` (moe_flat_fsdp_weights) the gathered 2D [G * rows, cols] + qvalue is bitcast to (G, rows, cols) and a per-channel [1, cols] scale to [1, 1, cols]. + """ + if gathered_shape is None: + for axis_name, axis_idx in weight_gather_axes: + rhs_qvalue = jax.lax.all_gather(rhs.qvalue, axis_name, axis=axis_idx, tiled=True) + # replace the qvalue with the gathered qvalue in the QArray + rhs = dataclasses.replace(rhs, qvalue=rhs_qvalue) + return rhs + + qvalue = rhs.qvalue for axis_name, axis_idx in weight_gather_axes: - rhs_qvalue = jax.lax.all_gather(rhs.qvalue, axis_name, axis=axis_idx, tiled=True) - # replace the qvalue with the gathered qvalue in the QArray - rhs = dataclasses.replace(rhs, qvalue=rhs_qvalue) - return rhs + qvalue = jax.lax.all_gather(qvalue, axis_name, axis=axis_idx, tiled=True) + qvalue = qvalue.reshape(gathered_shape) + + def _expand(s): + if s is None or s.ndim == 0 or s.ndim == len(gathered_shape): + return s + return s.reshape((1,) * (len(gathered_shape) - s.ndim) + tuple(s.shape)) + + return dataclasses.replace(rhs, qvalue=qvalue, scale=_expand(rhs.scale), zero_point=_expand(rhs.zero_point)) def _fwd_run_tokamax_v1( @@ -394,6 +543,7 @@ def _fwd_run_tokamax_v2( partial_sum: jnp.ndarray | None, transpose_rhs: bool, quantization_rule: qwix.QtRule | None = None, + out_is_3d: bool = False, ) -> jnp.ndarray: """Executes the Tokamax GMM V2 backend for forward pass OUT = LHS @ RHS.""" # if transpose_rhs=False, rhs is [g, k, n], remain unchanged @@ -428,8 +578,8 @@ def _fwd_run_tokamax_v2( fwd_tiling = gmm_v2.TileSizes(tile_m=tiling[0], tile_k=tiling[1], tile_n=tiling[2]) out = gmm_v2.gmm_v2( - lhs=lhs_operand, # pyrefly: ignore[bad-argument-type] - rhs=rhs_operand, # pyrefly: ignore[bad-argument-type] + lhs=lhs_operand, + rhs=rhs_operand, group_sizes=group_sizes, rhs_scale=rhs_scale, tile_info=fwd_tiling, @@ -438,12 +588,13 @@ def _fwd_run_tokamax_v2( group_offset=group_offset, lhs_scale=lhs_scale, maybe_quantize_lhs=maybe_quantize_lhs, + out_is_3d=out_is_3d, ) # gmm_v2 only rescales output when it quantizes lhs internally; for pre-quantized QArray # inputs, apply lhs.scale to the accumulated output here. if isinstance(lhs, qpl.QArray): - out = out * (lhs.scale.squeeze() if lhs.scale.size == 1 else lhs.scale).astype(out.dtype) + out = out * _scale_like(lhs.scale, out, out.dtype) return out @@ -497,6 +648,14 @@ def _gmm_bwd( rhs_vma_axes: tuple, use_gmm_v2: bool, use_gmm_v2_heuristic_tiling: bool, + return_rhs: bool, + return_lhs: bool, + fuse_dlhs_scale: bool, + use_dlhs_transpose_rhs: bool, + flat_fsdp_shard: tuple[int, int, int] | None, + out_is_3d: bool, + fuse_3d_dlhs_ps: bool, + tc_unsort_cfg: TcUnsortCfg | None, residual: tuple[ jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray, @@ -505,6 +664,7 @@ def _gmm_bwd( jnp.ndarray | None, bool, bool, + object, ], grad: jnp.ndarray, ) -> tuple[ @@ -514,6 +674,7 @@ def _gmm_bwd( None, jnp.ndarray | None, jnp.ndarray | None, + None, ]: """Backward function for throughput GMM VJP.""" ( @@ -524,8 +685,34 @@ def _gmm_bwd( partial_sum_fwd, lhs_is_qarray, rhs_is_qarray, + tc_routing, ) = residual num_actual_groups = residual_rhs.shape[0] + grad, dlhs_partial_sum, drhs_partial_sum = grad + # With return_lhs, the cotangent of the returned lhs is the input gradient already computed by its + # other consumers (e.g. the MoE wi_1 GMM). The DLHS GMM adds this dlhs onto it in place. + if not return_lhs: + dlhs_partial_sum = None + elif isinstance(dlhs_partial_sum, qpl.QArray): + dlhs_partial_sum = dlhs_partial_sum.qvalue + # With return_rhs, the cotangent of the returned rhs is the weight gradient already accumulated + # by its downstream consumers (e.g. later MoE token chunks). tgmm adds this drhs onto it in place. + if not return_rhs: + drhs_partial_sum = None + else: + if isinstance(drhs_partial_sum, qpl.QArray): + drhs_partial_sum = drhs_partial_sum.qvalue + if transpose_rhs: + drhs_partial_sum = drhs_partial_sum.swapaxes(1, 2) + # 3D layouts: the incoming grad matches the (possibly 3D) forward output and dlhs is + # produced in the layout of the forward lhs. + lhs_is_3d = (residual_lhs.qvalue if isinstance(residual_lhs, qpl.QArray) else residual_lhs).ndim == 3 + if tc_unsort_cfg is not None: + _check_tc_unsort_bwd_supported(residual_lhs, residual_rhs, quantization_rule, tc_unsort_cfg) + if tc_unsort_cfg.combine_rs_cfg is not None: + # The fused forward returned the local [batch, seq, D] combine; go back to the TC token layout. + grad = grad.reshape(-1, grad.shape[-1] // 128, 128) + del out_is_3d # grad.ndim carries it. # Jargon used here: # - lhs: input activation in forward pass, possibly quantized. @@ -551,8 +738,32 @@ def _gmm_bwd( if quantization_rule: dlhs_dout, drhs_dout = _bwd_quantize_gradient(dlhs_dout, drhs_dout, quantization_rule) + # 2b. Fused TC unsort: gather the (already scaled / quantized) token-order cotangent into expert slots. + if tc_unsort_cfg is not None: + cfg = tc_unsort_cfg + padded_cap = (residual_lhs.qvalue if isinstance(residual_lhs, qpl.QArray) else residual_lhs).shape[0] + + def _unsort_gather_arr(arr): + if cfg.combine_rs_cfg is not None: + arr = jax.lax.all_gather(arr, cfg.combine_rs_cfg[0], axis=0, tiled=True) + arr = ragged_sort_tc._gather(arr, tc_routing, cfg.unpadded_cap, cfg.blocks, cfg.mask_padding) # pylint: disable=protected-access + return jnp.pad(arr, [(0, padded_cap - arr.shape[0])] + [(0, 0)] * (arr.ndim - 1)) + + def _unsort_gather_maybe_qa(x): + if isinstance(x, qpl.QArray): + return dataclasses.replace(x, qvalue=_unsort_gather_arr(x.qvalue)) + return _unsort_gather_arr(x) + + same_dout = dlhs_dout is drhs_dout + dlhs_dout = _unsort_gather_maybe_qa(dlhs_dout) + drhs_dout = dlhs_dout if same_dout else _unsort_gather_maybe_qa(drhs_dout) + + lhs_scale_for_dlhs = None + if lhs_is_qarray and isinstance(residual_lhs, qpl.QArray): + lhs_scale_for_dlhs = residual_lhs.scale + # 3. DLHS Gradient Execution - dlhs = _compute_dlhs( + dlhs, lhs_scale_applied = _compute_dlhs( dlhs_dout, rhs, group_sizes, @@ -566,6 +777,12 @@ def _gmm_bwd( interpret, lhs_vma_axes, use_gmm_v2_heuristic_tiling, + partial_sum=dlhs_partial_sum, + lhs_scale=lhs_scale_for_dlhs, + fuse_dlhs_scale=fuse_dlhs_scale, + use_dlhs_transpose_rhs=use_dlhs_transpose_rhs, + dlhs_is_3d=lhs_is_3d, + fuse_3d_dlhs_ps=fuse_3d_dlhs_ps, ) # 4. DRHS Gradient Execution @@ -585,6 +802,8 @@ def _gmm_bwd( rhs_vma_axes, quantization_rule, use_gmm_v2_heuristic_tiling, + partial_sum=drhs_partial_sum, + flat_rhs_shard=flat_fsdp_shard is not None, ) # 5. Output Formatting @@ -599,10 +818,11 @@ def _gmm_bwd( # If forward inputs were QArray, wrap cotangents in QArray to match JAX pytree container # structure expected by autodiff. if lhs_is_qarray and isinstance(residual_lhs, qpl.QArray): - # Scale dlhs by lhs_scale to propagate chain rule through pre-quantized QArray input. - lhs_scale = (residual_lhs.scale.squeeze() if residual_lhs.scale.size == 1 else residual_lhs.scale).astype(dlhs.dtype) + if not lhs_scale_applied: + # Scale dlhs by lhs_scale to propagate chain rule through pre-quantized QArray input. + dlhs = dlhs * _scale_like(residual_lhs.scale, dlhs, dlhs.dtype) dlhs = qpl.QArray( - qvalue=dlhs * lhs_scale, + qvalue=dlhs, scale=jnp.zeros_like(residual_lhs.scale), zero_point=jnp.zeros_like(residual_lhs.zero_point) if residual_lhs.zero_point is not None else None, qtype=residual_lhs.qtype, @@ -617,7 +837,70 @@ def _gmm_bwd( qtype=residual_rhs.qtype, ) - return dlhs, drhs, None, None, d_existing_out, dpartial_sum + return dlhs, drhs, None, None, d_existing_out, dpartial_sum, None + + +def _check_tc_unsort_bwd_supported( + lhs: jnp.ndarray | qpl.QArray, + rhs: jnp.ndarray | qpl.QArray, + quantization_rule: qwix.QtRule | None, + tc_unsort_cfg: TcUnsortCfg, +): + """The fused TC unsort scales and quantizes the cotangent in token order, before the gather into expert slots.""" + # Per-row lhs and per-expert rhs scales are laid out by expert slot, so they can't be applied in token order. + for name, x in (("lhs", lhs), ("rhs", rhs)): + if isinstance(x, qpl.QArray) and x.scale.size != 1: + raise NotImplementedError(f"tc_unsort_cfg requires a per-tensor {name} scale, got shape {x.scale.shape}.") + if quantization_rule is None: + return + if ( + quantization_rule.act_qtype + and not isinstance(lhs, qpl.QArray) + and not quantization_rule.disable_channelwise_axes + and not quantizations.is_static_calibration(quantization_rule.act_calibration_method) + ): + raise NotImplementedError("tc_unsort_cfg requires a per-tensor act scale (fixed or disable_channelwise_axes).") + if tc_unsort_cfg.combine_rs_cfg is not None: + # Each EP shard applies the rhs / lhs scales to its local cotangent and quantizes it before the all-gather, so + # every scale must be the same on all shards; a dynamic one is computed from the shard's own experts and rows. + for name, qtype, method in ( + ("weight", quantization_rule.weight_qtype, quantization_rule.weight_calibration_method), + ("act", quantization_rule.act_qtype, quantization_rule.act_calibration_method), + ("bwd", quantization_rule.bwd_qtype, quantization_rule.bwd_calibration_method), + ): + if qtype and not quantizations.is_static_calibration(method): + raise NotImplementedError( + f"tc_unsort_cfg with combine_rs_cfg requires a fixed {name}_calibration_method, got {method!r}." + ) + if not quantization_rule.bwd_qtype or quantizations.is_static_calibration(quantization_rule.bwd_calibration_method): + return + # A dynamic per-row cotangent scale would be in token order, so it can't be gathered like the qvalue. + if not quantization_rule.disable_channelwise_axes: + raise NotImplementedError("tc_unsort_cfg requires a per-tensor bwd scale (fixed or disable_channelwise_axes).") + + +def _lhs_rhs_share_fixed_scale( + quantization_rule: qwix.QtRule | None, + lhs: qpl.QArray, + rhs: jnp.ndarray | qpl.QArray, +) -> bool: + """True when lhs and rhs carry the same fixed per-tensor scale, so grad * rhs.scale == grad * lhs.scale. + + Scale equality is inferred from the rule, not the (traced) arrays: qwix builds a fixed scale from the + calibration bound in the quantized array's dtype, so the scales match only if both operands were quantized + under this rule's calibration at the same dtype. A pre-quantized lhs (moe_quantize_token_all_gather) is + checked to use the rule's act_calibration_method; a scale dtype mismatch falls back to applying lhs.scale. + """ + return ( + isinstance(rhs, qpl.QArray) + and quantization_rule is not None + and quantization_rule.weight_qtype == quantization_rule.act_qtype + and quantizations.is_static_calibration(quantization_rule.weight_calibration_method) + and quantization_rule.weight_calibration_method == quantization_rule.act_calibration_method + and lhs.scale.size == 1 + and rhs.scale.size == 1 + and lhs.scale.dtype == rhs.scale.dtype + ) def _bwd_prepare_inputs( @@ -634,6 +917,7 @@ def _bwd_prepare_inputs( # dlhs_dout and drhs_dout can be different when quantization is enabled. dlhs_dout = grad drhs_dout = grad + orig_rhs = rhs # Apply rhs.scale to dlhs_dout, dlhs_dout[m, n] @ rhs_transpose[g, n, k] = dlhs[m, k] # Assume channelwise scale on rhs n. @@ -654,10 +938,9 @@ def _bwd_prepare_inputs( # GMM2 FWD performs lhs quantization inside kernel, lhs is stored as unquantized dtype # in the residual tuple. In BWD, we explicitly quantize lhs. if quantization_rule and quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray): - lhs = qpl.quantize( # pyrefly: ignore[bad-assignment] + lhs = qpl.quantize( lhs, quantization_rule.act_qtype, - # pyrefly: ignore[bad-argument-type] channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0], # pyrefly: ignore[bad-argument-type] calibration_method=quantization_rule.act_calibration_method, @@ -667,7 +950,10 @@ def _bwd_prepare_inputs( # Apply lhs.scale to drhs_dout, as axis m will disappear in drhs. if isinstance(lhs, qpl.QArray): # lhs - qvalue: [m, k] scale: [m, 1] - drhs_dout = drhs_dout * (lhs.scale.squeeze() if lhs.scale.size == 1 else lhs.scale).astype(grad.dtype) + if _lhs_rhs_share_fixed_scale(quantization_rule, lhs, orig_rhs): + drhs_dout = dlhs_dout + else: + drhs_dout = drhs_dout * _scale_like(lhs.scale, drhs_dout, grad.dtype) lhs = lhs.qvalue return dlhs_dout, drhs_dout, lhs, rhs @@ -680,20 +966,25 @@ def _bwd_quantize_gradient( ) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray]: """Applies backward quantization to incoming gradients.""" if quantization_rule.bwd_qtype: + bwd_cal = quantization_rule.bwd_calibration_method + share_quantized_grad = dlhs_dout is drhs_dout and quantizations.is_static_calibration(bwd_cal) dlhs_dout = qpl.quantize( # pyrefly: ignore[bad-argument-type] dlhs_dout, quantization_rule.bwd_qtype, channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0], - calibration_method=quantization_rule.bwd_calibration_method, - ) - drhs_dout = qpl.quantize( - # pyrefly: ignore[bad-argument-type] - drhs_dout, - quantization_rule.bwd_qtype, - channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [1], - calibration_method=quantization_rule.bwd_calibration_method, + calibration_method=bwd_cal, ) + if share_quantized_grad: + drhs_dout = dlhs_dout + else: + drhs_dout = qpl.quantize( + # pyrefly: ignore[bad-argument-type] + drhs_dout, + quantization_rule.bwd_qtype, + channelwise_axes=[] if quantization_rule.disable_channelwise_axes else list(range(1, drhs_dout.ndim)), + calibration_method=bwd_cal, + ) return dlhs_dout, drhs_dout @@ -716,10 +1007,20 @@ def _compute_dlhs( interpret: bool, lhs_vma_axes: tuple, use_gmm_v2_heuristic_tiling: bool, -) -> jnp.ndarray: - """Routes execution of DLHS based on backend choices.""" + partial_sum: jnp.ndarray | None = None, + lhs_scale: jnp.ndarray | None = None, + fuse_dlhs_scale: bool = False, + use_dlhs_transpose_rhs: bool = False, + dlhs_is_3d: bool = False, + fuse_3d_dlhs_ps: bool = False, +) -> tuple[jnp.ndarray, bool]: + """Routes execution of DLHS based on backend choices. + + Returns (dlhs, lhs_scale_applied). partial_sum (only set with return_lhs, i.e. tokamax gmm_v2) is + added to dlhs; lhs_scale_applied tells the caller whether lhs_scale was already folded into dlhs. + """ if use_tokamax_backend and not use_gmm_v2: - return _dlhs_run_tokamax_v1( + dlhs = _dlhs_run_tokamax_v1( dlhs_dout, rhs, group_sizes, @@ -727,14 +1028,29 @@ def _compute_dlhs( transpose_rhs, use_manual_quantization, ) + return dlhs, False elif use_tokamax_backend and use_gmm_v2: return _dlhs_run_tokamax_v2( - dlhs_dout, rhs, group_sizes, group_offset, lhs_dtype, tiling, use_gmm_v2_heuristic_tiling, transpose_rhs + dlhs_dout, + rhs, + group_sizes, + group_offset, + lhs_dtype, + tiling, + use_gmm_v2_heuristic_tiling, + transpose_rhs, + partial_sum=partial_sum, + lhs_scale=lhs_scale, + fuse_dlhs_scale=fuse_dlhs_scale, + use_dlhs_transpose_rhs=use_dlhs_transpose_rhs, + out_is_3d=dlhs_is_3d, + fuse_3d_dlhs_ps=fuse_3d_dlhs_ps, ) else: - return _dlhs_run_megablox( + dlhs = _dlhs_run_megablox( dlhs_dout, rhs, group_sizes, group_offset, lhs_dtype, tiling, transpose_rhs, interpret, lhs_vma_axes ) + return dlhs, False def _dlhs_run_tokamax_v1( @@ -784,16 +1100,20 @@ def _dlhs_scale_grad_by_rhs_scale( rhs_scale = rhs_scale.squeeze(axis=squeeze_axis) # 2. Apply scale (handle shared vs per-expert scales) + if rhs_scale.size == 1: + return grad * rhs_scale.reshape(()).astype(grad.dtype) if rhs_scale.shape[0] == 1: - return grad * rhs_scale.astype(grad.dtype) + scale = rhs_scale.astype(grad.dtype) else: - repeated_scale = jnp.repeat( + scale = jnp.repeat( rhs_scale.astype(grad.dtype), group_sizes, axis=0, total_repeat_length=grad.shape[0], ) - return grad * repeated_scale + if grad.ndim == 3: + scale = scale.reshape(scale.shape[0], *grad.shape[1:]) + return grad * scale def _dlhs_run_tokamax_v2( @@ -805,35 +1125,78 @@ def _dlhs_run_tokamax_v2( tiling: tuple, use_gmm_v2_heuristic_tiling: bool, transpose_rhs: bool, -) -> jnp.ndarray: + partial_sum: jnp.ndarray | None = None, + lhs_scale: jnp.ndarray | None = None, + fuse_dlhs_scale: bool = False, + use_dlhs_transpose_rhs: bool = False, + out_is_3d: bool = False, + fuse_3d_dlhs_ps: bool = False, +) -> tuple[jnp.ndarray, bool]: """Executes Tokamax GMM V2 backend for DLHS = DLHS_dout @ RHS^T.""" - # NOTE: We manually transpose RHS here because gmm_v2 lacks native transpose_rhs support. - dlhs_rhs = rhs if transpose_rhs else rhs.swapaxes(1, 2) - dlhs_lhs = dlhs_dout.qvalue if isinstance(dlhs_dout, qpl.QArray) else dlhs_dout + is_qarray_dout = isinstance(dlhs_dout, qpl.QArray) + dlhs_lhs = dlhs_dout.qvalue if is_qarray_dout else dlhs_dout if use_gmm_v2_heuristic_tiling: dlhs_tiling = gmm_v2.calculate_tiling else: dlhs_tiling = gmm_v2.TileSizes(tile_m=tiling[3], tile_k=tiling[4], tile_n=tiling[5]) + if use_dlhs_transpose_rhs: + kernel_transpose_rhs = not transpose_rhs + dlhs_rhs = rhs + else: + kernel_transpose_rhs = False + dlhs_rhs = rhs if transpose_rhs else rhs.swapaxes(1, 2) + # Sub-byte packing only supports the standard [group, k, n] RHS layout. + if kernel_transpose_rhs and jax.dtypes.itemsize_bits(rhs.dtype) < 8: + dlhs_rhs = rhs.swapaxes(1, 2) + kernel_transpose_rhs = False + + # The kernel adds partial_sum after applying its lhs_scale, so partial_sum can only go into the kernel + # when every scale on dlhs (dlhs_dout.scale and, with partial_sum, the pre-quantized lhs scale) is + # folded into that per-tensor lhs_scale. Otherwise scale and accumulate after the kernel. + scales = [s for s in (dlhs_dout.scale if is_qarray_dout else None, lhs_scale) if s is not None] + can_fuse_scale = (fuse_dlhs_scale or partial_sum is not None) and all(s.size == 1 for s in scales) + kernel_lhs_scale = None + kernel_partial_sum = None + lhs_scale_applied = False + if can_fuse_scale: + for scale in scales: + scale = scale.astype(jnp.float32).reshape(1, 1) + kernel_lhs_scale = scale if kernel_lhs_scale is None else kernel_lhs_scale * scale + lhs_scale_applied = lhs_scale is not None + kernel_partial_sum = partial_sum + dlhs = gmm_v2.gmm_v2( lhs=dlhs_lhs, rhs=dlhs_rhs, group_sizes=group_sizes, # rhs scale is already applied to dlhs_lhs rhs_scale=None, + partial_sum=kernel_partial_sum, tile_info=dlhs_tiling, - preferred_element_type=lhs_dtype, # pyrefly: ignore[bad-argument-type] + preferred_element_type=lhs_dtype, group_offset=group_offset, - # Bypass internal quantization if incoming dlhs_dout is already a QArray. - maybe_quantize_lhs=not isinstance(dlhs_dout, qpl.QArray), + lhs_scale=kernel_lhs_scale, + transpose_rhs=kernel_transpose_rhs, + # Bypass internal quantization if incoming dlhs_dout is already a QArray, or when lhs_scale is a + # post-matmul scale rather than a quantization scale. + maybe_quantize_lhs=not is_qarray_dout and kernel_lhs_scale is None, + zero_initialize=not (fuse_3d_dlhs_ps and out_is_3d and kernel_partial_sum is not None), + out_is_3d=out_is_3d, ) - # Rescale dlhs by dlhs_dout.scale when incoming gradient was pre-quantized QArray. - if isinstance(dlhs_dout, qpl.QArray): - dlhs = dlhs * (dlhs_dout.scale.squeeze() if dlhs_dout.scale.size == 1 else dlhs_dout.scale).astype(dlhs.dtype) + if not can_fuse_scale: + # Rescale dlhs by dlhs_dout.scale when incoming gradient was pre-quantized QArray. + if is_qarray_dout: + dlhs = dlhs * _scale_like(dlhs_dout.scale, dlhs, dlhs.dtype) + if partial_sum is not None: + if lhs_scale is not None: + dlhs = dlhs * _scale_like(lhs_scale, dlhs, dlhs.dtype) + lhs_scale_applied = True + dlhs = dlhs + partial_sum - return dlhs + return dlhs, lhs_scale_applied def _dlhs_run_megablox( @@ -882,13 +1245,23 @@ def _compute_drhs( rhs_vma_axes: tuple, quantization_rule: qwix.QtRule | None, use_gmm_v2_heuristic_tiling: bool, + partial_sum: jnp.ndarray | None = None, + flat_rhs_shard: bool = False, ) -> jnp.ndarray: """Routes execution of DRHS based on backend choices.""" if use_tokamax_backend and not use_gmm_v2: drhs = _drhs_run_tokamax_v1(drhs_dout, lhs, group_sizes, rhs_dtype, use_manual_quantization) elif use_tokamax_backend and use_gmm_v2: drhs = _drhs_run_tokamax_v2( - drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling, use_gmm_v2_heuristic_tiling + drhs_dout, + lhs, + group_sizes, + group_offset, + num_actual_groups, + rhs_dtype, + tiling, + use_gmm_v2_heuristic_tiling, + partial_sum=partial_sum, ) else: drhs = _drhs_run_megablox( @@ -896,13 +1269,23 @@ def _compute_drhs( ) if use_tokamax_backend and quantization_rule and quantization_rule.bwd_qtype and weight_gather_axes: - drhs = _drhs_scatter_weight(drhs, weight_gather_axes) + drhs = _drhs_scatter_weight(drhs, weight_gather_axes, flat_rhs_shard=flat_rhs_shard) return drhs -def _drhs_scatter_weight(drhs: jnp.ndarray, weight_gather_axes: List[Tuple[str, int]]) -> jnp.ndarray: - """Scatters the DRHS output back in the reverse order of the forward gather.""" +def _drhs_scatter_weight( + drhs: jnp.ndarray, + weight_gather_axes: List[Tuple[str, int]], + flat_rhs_shard: bool = False, +) -> jnp.ndarray: + """Scatters the DRHS output back in the reverse order of the forward gather. + + With `flat_rhs_shard` (moe_flat_fsdp_weights) the (G, rows, cols) grad is first + bitcast to the flat [G * rows, cols] matrix so the reduce-scatter runs on the major dim. + """ + if flat_rhs_shard: + drhs = drhs.reshape(-1, drhs.shape[-1]) for axis_name, axis_idx in reversed(weight_gather_axes): drhs = jax.lax.psum_scatter(drhs, axis_name, scatter_dimension=axis_idx, tiled=True) return drhs @@ -938,7 +1321,10 @@ def _drhs_run_tokamax_v1( def _drhs_prepare_bwd_scale(drhs_dout: qpl.QArray) -> jnp.ndarray: """Formats and broadcasts drhs_dout scale to (1, 1, size_n) for V2 TGMM kernel.""" scale = drhs_dout.scale - size_n = drhs_dout.shape[1] + size_n = drhs_dout.qvalue.size // drhs_dout.qvalue.shape[0] + if drhs_dout.qvalue.ndim == 3 or scale.size == 1: + # 3D grad [m, n // 128, 128]: scale is (1, 1, 1) per-tensor or (1, n // 128, 128) per-channel. + return jnp.broadcast_to(scale.reshape(1, 1, -1), (1, 1, size_n)) # per channel: (1, n) -> (1, 1, n) # per tensor: (1, 1) -> (1, 1, 1) rhs_scale = jnp.expand_dims(scale, axis=1) @@ -957,6 +1343,7 @@ def _drhs_run_tokamax_v2( rhs_dtype: jax.typing.DTypeLike, tiling: tuple, use_gmm_v2_heuristic_tiling: bool, + partial_sum: jnp.ndarray | None = None, ) -> jnp.ndarray: """Executes Tokamax TGMM V2 backend for DRHS = LHS^T @ DRHS_dout.""" drhs_rhs = drhs_dout.qvalue if isinstance(drhs_dout, qpl.QArray) else drhs_dout @@ -977,8 +1364,9 @@ def _drhs_run_tokamax_v2( group_sizes=group_sizes, num_actual_groups=num_actual_groups, rhs_scale=rhs_scale, + partial_sum=partial_sum, precision=jax.lax.Precision.DEFAULT, - preferred_element_type=rhs_dtype, # pyrefly: ignore[bad-argument-type] + preferred_element_type=rhs_dtype, group_offset=group_offset, tile_info=drhs_tiling, ) @@ -1007,3 +1395,140 @@ def _drhs_run_megablox( interpret=interpret, varying_axes=rhs_vma_axes, ) + + +# ============================================================================== +# GMM + SwiGLU: the activation in the gmm_v2 epilogue +# ============================================================================== + + +def gmm_silu( + lhs: jnp.ndarray, + rhs: jnp.ndarray, + group_sizes: jnp.ndarray, + preferred_element_type: jnp.dtype, + tiling: tuple[int, int, int, int, int, int, int, int, int], + group_offset: jnp.ndarray | int | None = None, + use_gmm_v2_heuristic_tiling: bool = False, + use_dlhs_transpose_rhs: bool = False, + fuse_dlhs_scale: bool = False, + fuse_3d_dlhs_ps: bool = False, +) -> jnp.ndarray: + """silu(lhs @ gate) * (lhs @ up) per group, with the activation in the gmm_v2 epilogue. + + `rhs` is the prefused [g, k, 2n] weight, gate = rhs[..., :n] and up = rhs[..., n:]. The unfused path writes + the [m, 2n] GMM output and runs the activation over every row in a separate elementwise pass; here the kernel + applies it to the f32 accumulator of the rows it computes and writes [m, n]. + + Numerics: the activation reads the f32 accumulator instead of its bf16 rounding, so `act` differs from the + unfused value by rounding (it is the more precise of the two). The backward is the unfused one: it takes + gate and up rounded exactly as the unfused GMM rounds its output (the kernel writes them next to act when + differentiated), the SwiGLU VJP of `jax.nn.silu(gate) * up`, and the unchanged tokamax gmm_v2 DLHS / TGMM + DRHS of `gmm`. + + Only the bf16 tokamax gmm_v2 path: no quantization, no partial_sum, no return_lhs / return_rhs, 2D lhs and + output, rhs not transposed. tiling is gmm's 9-tuple; the fused kernel's tile_n counts the columns of each + half, so tiling[2] is halved and every grid step does the unfused kernel's MXU work in the same VMEM. + """ + if isinstance(lhs, qpl.QArray) or isinstance(rhs, qpl.QArray) or lhs.ndim != 2 or rhs.ndim != 3: + raise NotImplementedError("gmm_silu: plain 2D lhs and [g, k, 2n] rhs only.") + return _gmm_silu( + lhs, + rhs, + group_sizes, + group_offset, + preferred_element_type, + tuple(tiling), + use_gmm_v2_heuristic_tiling, + use_dlhs_transpose_rhs, + fuse_dlhs_scale, + fuse_3d_dlhs_ps, + ) + + +def _gmm_silu_kernel(lhs, rhs, group_sizes, group_offset, preferred_element_type, tiling, heuristic, return_pre_act): + """The fused forward: act, or (act, gate, up) with return_pre_act.""" + if heuristic: + tile_info = gmm_v2.calculate_tiling + else: + half_n = rhs.shape[-1] // 2 + tile_n = min(tiling[2], rhs.shape[-1]) // 2 + if tiling[2] % 2 or (tile_n % 128 and tile_n != half_n): + raise ValueError(f"gmm_silu: fwd tile_n {tiling[2]} must halve to a multiple of 128 (or to n={half_n}).") + tile_info = gmm_v2.TileSizes(tile_m=tiling[0], tile_k=tiling[1], tile_n=tile_n) + return gmm_v2.gmm_v2( + lhs=lhs, + rhs=rhs, + group_sizes=group_sizes, + tile_info=tile_info, + preferred_element_type=preferred_element_type, + group_offset=group_offset, + fuse_act="silu", + return_pre_act=return_pre_act, + ) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(4, 5, 6, 7, 8, 9)) +def _gmm_silu( + lhs, rhs, group_sizes, group_offset, preferred_element_type, tiling, heuristic, dlhs_transpose, fuse_scale, fuse_3d +): + # Not differentiated: the activation alone, nothing else is written. + del dlhs_transpose, fuse_scale, fuse_3d + return _gmm_silu_kernel(lhs, rhs, group_sizes, group_offset, preferred_element_type, tiling, heuristic, False) + + +def _gmm_silu_fwd( + lhs, rhs, group_sizes, group_offset, preferred_element_type, tiling, heuristic, dlhs_transpose, fuse_scale, fuse_3d +): + """Differentiated: the kernel also writes gate and up, the residuals of the unfused backward.""" + del dlhs_transpose, fuse_scale, fuse_3d + act, gate, up = _gmm_silu_kernel(lhs, rhs, group_sizes, group_offset, preferred_element_type, tiling, heuristic, True) + return act, (lhs, rhs, group_sizes, group_offset, gate, up) + + +def _swiglu(gate, up): + """MaxText's MoE activation on the GMM output (`apply_ffn_activation`, mlp_activations silu + linear).""" + return (jax.nn.silu(gate) * up).astype(gate.dtype) + + +def _gmm_silu_bwd(preferred_element_type, tiling, heuristic, dlhs_transpose, fuse_scale, fuse_3d, residual, dact): + """The unfused backward from the saved bf16 gate and up: the SwiGLU VJP, then gmm's DLHS / TGMM.""" + lhs, rhs, group_sizes, group_offset, gate, up = residual + # The unfused backward: the SwiGLU VJP on the bf16 GMM output, then gmm's own backward on [d gate | d up]. + _, swiglu_vjp = jax.vjp(_swiglu, gate, up) + dgate, dup = swiglu_vjp(dact.astype(gate.dtype)) + grad = jnp.concatenate([dgate, dup], axis=-1) + # gmm's own backward, as gmm's custom VJP calls it for a plain bf16 tokamax gmm_v2 GMM (no quantization, no + # partial sum, no returned operands), with this op's DLHS options. + dlhs, drhs, *_ = _gmm_bwd( + lhs_dtype=lhs.dtype, + rhs_dtype=rhs.dtype, + preferred_element_type=preferred_element_type, + tiling=tiling, + transpose_rhs=False, + interpret=False, # the megablox backend only + quantization_rule=None, + use_tokamax_backend=True, + weight_gather_axes=None, + use_manual_quantization=False, + lhs_vma_axes=(), + rhs_vma_axes=(), + use_gmm_v2=True, + use_gmm_v2_heuristic_tiling=heuristic, + return_rhs=False, + return_lhs=False, + fuse_dlhs_scale=fuse_scale, + use_dlhs_transpose_rhs=dlhs_transpose, + flat_fsdp_shard=None, + out_is_3d=False, + fuse_3d_dlhs_ps=fuse_3d, + tc_unsort_cfg=None, + # (lhs, rhs, group_sizes, group_offset, partial_sum_fwd, lhs_is_qarray, rhs_is_qarray, tc_routing) + residual=(lhs, rhs, group_sizes, group_offset, None, False, False, None), + # (d out, d returned lhs, d returned rhs) + grad=(grad, None, None), + ) + return dlhs, drhs, None, None + + +_gmm_silu.defvjp(_gmm_silu_fwd, _gmm_silu_bwd) diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 15e30f0c17..8168559c99 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -27,6 +27,8 @@ from jax.experimental.pallas import tpu as pltpu import jax.numpy as jnp +from maxtext.kernels.megablox import relayout + # Util. @@ -84,7 +86,7 @@ class RhsRef(ABC): """Abstract class that defines interfaces for rhs values.""" @abstractmethod - def get_weight(self) -> jax.Array: + def get_weight(self, transpose_rhs: bool = False) -> jax.Array: ... @abstractmethod @@ -105,7 +107,8 @@ class WeightsRef(RhsRef): scale: Any | None bias: Any | None - def get_weight(self) -> jax.Array: + def get_weight(self, transpose_rhs: bool = False) -> jax.Array: + del transpose_rhs return self.weight[...] def get_scale(self) -> jax.Array: @@ -125,10 +128,11 @@ class FusedWeightsRef(RhsRef): gate: WeightsRef up: WeightsRef - def get_weight(self) -> jax.Array: + def get_weight(self, transpose_rhs: bool = False) -> jax.Array: w_gate = self.gate.get_weight() w_up = self.up.get_weight() - return jnp.concatenate([w_gate, w_up], axis=-1) + concat_axis = -2 if transpose_rhs else -1 + return jnp.concatenate([w_gate, w_up], axis=concat_axis) def get_scale(self) -> jax.Array: s_gate = self.gate.get_scale() @@ -201,6 +205,10 @@ class InputConfigs: # (dtype == quant_dtype) the scale dequantizes it (rhs); when it arrives # unquantized (dtype != quant_dtype) the scale quantizes it online (lhs). has_scale: bool = False + # The operand is stored as [size_m, size // 128, 128] (one contiguous slab per + # row) instead of [size_m, size]. The kernel DMAs 4D blocks and relayouts them + # to 2D tiles inside VMEM (see relayout.py). + is_3d: bool = False @property def should_use_external_scale(self) -> bool: @@ -242,6 +250,71 @@ class GmmConfigs: has_partial_sum: bool zero_init: bool fuse_act: str | None + transpose_rhs: bool + # Output is stored as [size_m, out_size_n // 128, 128] (one contiguous slab + # per row); tiles are relayouted from [tile_m, tile_n] inside VMEM. + out_is_3d: bool = False + # With fuse_act: also write the two halves of the accumulator before the + # activation (gate and up), each [size_m, out_size_n] in out_dtype, next to + # the activation output. The backward of the activation needs them. + return_pre_act: bool = False + + @property + def lhs_is_3d(self) -> bool: + return self.lhs_cfgs.is_3d + + @property + def lhs_tile_d0(self) -> int: + assert self.lhs_is_3d + return self.tiles.tile_k // pltpu.get_tpu_info().num_lanes + + @property + def out_tile_d0(self) -> int: + assert self.out_is_3d + return self.tiles.tile_n // pltpu.get_tpu_info().num_lanes + + @property + def lhs_window(self) -> bool: + """3D lhs whose K tile is not a legal D0 block: DMA full-D0 slabs, relayout the K window in VMEM.""" + return self.lhs_is_3d and relayout.needs_window(self.lhs_tile_d0, self.dims.size_k // pltpu.get_tpu_info().num_lanes) + + @property + def out_window(self) -> bool: + """3D out whose N tile is not a legal D0 block: a full-D0 out slab stays resident across N. + + The grid is then (gm, n, k) instead of (n, gm, k), so every N tile of a gm + tile is stored into the same VMEM slab before it is written back once. + """ + return self.out_is_3d and relayout.needs_window(self.out_tile_d0, self.out_size_n // pltpu.get_tpu_info().num_lanes) + + @property + def gm_axis(self) -> int: + return 0 if self.out_window else 1 + + @property + def n_axis(self) -> int: + return 1 if self.out_window else 0 + + @property + def lhs_chunk_d0(self) -> int | None: + """D0 rows per in-VMEM relayout chunk for a large 3D lhs K tile (None: relayout the whole tile). + + Relayouting e.g. a [512, 7168] fp8 tile at once materializes it in VMEM; chunks + of 1024 keep the relayout buffer small (the matmul accumulates across chunks). + """ + if not self.lhs_is_3d or self.lhs_cfgs.quant_dtype is not None: + return None + if self.num_quant_blocks_per_tile_k != 1 or self.rhs_cfgs.should_dequantize_before_matmul: + return None + # 8-bit only: there the chunked accumulation matches one K-wide matmul (bit-exact + # on some TPU generations, within ~1 bf16 ulp on others); bf16 matmuls round + # noticeably differently across chunk boundaries. + if jnp.dtype(self.lhs_cfgs.dtype).itemsize != 1 or jnp.dtype(self.rhs_cfgs.dtype).itemsize != 1: + return None + d0 = self.lhs_tile_d0 + if d0 <= 16 or d0 % 8: + return None + return 8 @property def num_quant_blocks_per_tile_k(self) -> int: @@ -266,6 +339,7 @@ def __init__(self, metadata_ref: MetadataRef, cfgs: GmmConfigs): self.cfgs = cfgs def lhs_index_map(self, _: jax.Array, gm_id: jax.Array, k_id: jax.Array): + """Block index of the lhs tile (rows of the group tile, K block; 3D lhs blocks K along D0).""" m_start = self.metadata_ref.gm_id_to_m_offset[gm_id] m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1] @@ -273,6 +347,10 @@ def lhs_index_map(self, _: jax.Array, gm_id: jax.Array, k_id: jax.Array): row_end = pl.cdiv(m_end, self.cfgs.dims.size_lhs_sublane) row_size = row_end - row_start + if self.cfgs.lhs_is_3d: + # Block is [rows, sublane, tile_d0, 128]: K is blocked along D0. A window + # block covers the full D0 (the K window is relayouted from VMEM). + return (pl.ds(row_start, row_size), 0, 0 if self.cfgs.lhs_window else k_id, 0) return (pl.ds(row_start, row_size), 0, k_id) def lhs_scale_index_map(self, _: jax.Array, gm_id: jax.Array, k_id: jax.Array): @@ -285,6 +363,8 @@ def lhs_scale_index_map(self, _: jax.Array, gm_id: jax.Array, k_id: jax.Array): def rhs_weight_index_map(self, n_id: jax.Array, gm_id: jax.Array, k_id: jax.Array): group_id = self.metadata_ref.gm_id_to_group_id[gm_id] + if self.cfgs.transpose_rhs: + return (group_id, n_id, k_id) return (group_id, k_id, n_id) def rhs_bias_index_map(self, n_id: jax.Array, gm_id: jax.Array, _: jax.Array): @@ -302,7 +382,7 @@ def rhs_scale_index_map(self, n_id: jax.Array, gm_id: jax.Array, k_id: jax.Array def out_index_map(self, n_id: jax.Array, gm_id: jax.Array, _: jax.Array): """Calculates index map for the output tensor.""" - is_last_gm = gm_id == (pl.num_programs(1) - 1) + is_last_gm = gm_id == (pl.num_programs(self.cfgs.gm_axis) - 1) m_start = self.metadata_ref.gm_id_to_m_offset[gm_id] m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1] @@ -312,9 +392,12 @@ def out_index_map(self, n_id: jax.Array, gm_id: jax.Array, _: jax.Array): row_end = jnp.where(is_last_gm, last_row_end, capped_row_end) row_size = row_end - row_start + if self.cfgs.out_is_3d: + return (pl.ds(row_start, row_size), 0, 0 if self.cfgs.out_window else n_id, 0) return (pl.ds(row_start, row_size), 0, n_id) def ps_index_map(self, n_id: jax.Array, gm_id: jax.Array, _: jax.Array): + """Block index of the partial_sum tile; a 3D partial_sum (with a 3D out) is blocked like the out.""" m_start = self.metadata_ref.gm_id_to_m_offset[gm_id] m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1] @@ -322,6 +405,8 @@ def ps_index_map(self, n_id: jax.Array, gm_id: jax.Array, _: jax.Array): row_end = pl.cdiv(m_end, self.cfgs.dims.size_lhs_sublane) row_size = row_end - row_start + if self.cfgs.out_is_3d: + return (pl.ds(row_start, row_size), 0, 0 if self.cfgs.out_window else n_id, 0) return (pl.ds(row_start, row_size), 0, n_id) @@ -333,15 +418,30 @@ def generate_block_specs( index_map = IndexMaps(metadata_ref, cfgs) bounded_slice_gm = pl.BoundedSlice(cfgs.tiles.tile_m // cfgs.dims.size_lhs_sublane) - lhs_value_spec = pl.BlockSpec( - (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_k), - index_map.lhs_index_map, - ) + num_lanes = pltpu.get_tpu_info().num_lanes + + # Index maps take (n_id, gm_id, k_id); with out_window the grid is (gm, n, k). + def _ordered(fn): + if not cfgs.out_window: + return fn + return lambda gm_id, n_id, k_id: fn(n_id, gm_id, k_id) + + if cfgs.lhs_is_3d: + lhs_d0 = cfgs.dims.size_k // num_lanes if cfgs.lhs_window else cfgs.lhs_tile_d0 + lhs_value_spec = pl.BlockSpec( + (bounded_slice_gm, cfgs.dims.size_lhs_sublane, lhs_d0, num_lanes), + _ordered(index_map.lhs_index_map), + ) + else: + lhs_value_spec = pl.BlockSpec( + (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_k), + _ordered(index_map.lhs_index_map), + ) lhs_scale_spec = None if cfgs.lhs_cfgs.has_scale: lhs_scale_spec = pl.BlockSpec( (1, 1), - index_map.lhs_scale_index_map, + _ordered(index_map.lhs_scale_index_map), ) lhs_block_spec = LhsRef(value=lhs_value_spec, scale=lhs_scale_spec) @@ -350,27 +450,36 @@ def generate_block_specs( packing = pl.cdiv(32, jax.dtypes.itemsize_bits(cfgs.rhs_cfgs.dtype)) tile_k_rhs //= packing + rhs_tile_shape = (None, cfgs.tiles.tile_n, tile_k_rhs) if cfgs.transpose_rhs else (None, tile_k_rhs, cfgs.tiles.tile_n) rhs_weight_spec = pl.BlockSpec( - (None, tile_k_rhs, cfgs.tiles.tile_n), - index_map.rhs_weight_index_map, - pipeline_mode=pl.Buffered(buffer_count=3), + rhs_tile_shape, + _ordered(index_map.rhs_weight_index_map), + # A full-K 3D lhs tile (e.g. a 512 x 7168 fp8 wi forward tile) needs VMEM for + # the relayout; double-buffer the (large) rhs tile there. + pipeline_mode=pl.Buffered(buffer_count=2 if cfgs.lhs_chunk_d0 is not None else 3), ) rhs_scale_block_spec = rhs_bias_block_spec = ps_block_spec = None if cfgs.rhs_cfgs.has_bias: rhs_bias_block_spec = pl.BlockSpec( (None, 1, cfgs.tiles.tile_n), - index_map.rhs_bias_index_map, + _ordered(index_map.rhs_bias_index_map), ) if cfgs.rhs_cfgs.has_scale: rhs_scale_block_spec = pl.BlockSpec( (None, cfgs.num_quant_blocks_per_tile_k, 1, cfgs.tiles.tile_n), - index_map.rhs_scale_index_map, + _ordered(index_map.rhs_scale_index_map), ) - if cfgs.has_partial_sum: + if cfgs.has_partial_sum and cfgs.out_is_3d: + ps_d0 = cfgs.out_size_n // num_lanes if cfgs.out_window else cfgs.out_tile_d0 + ps_block_spec = pl.BlockSpec( + (bounded_slice_gm, cfgs.dims.size_lhs_sublane, ps_d0, num_lanes), + _ordered(index_map.ps_index_map), + ) + elif cfgs.has_partial_sum: ps_block_spec = pl.BlockSpec( (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_n), - index_map.ps_index_map, + _ordered(index_map.ps_index_map), ) rhs_block_spec = WeightsRef( @@ -379,10 +488,17 @@ def generate_block_specs( bias=rhs_bias_block_spec, ) - out_block_spec = pl.BlockSpec( - (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_n), - index_map.out_index_map, - ) + if cfgs.out_is_3d: + out_d0 = cfgs.out_size_n // num_lanes if cfgs.out_window else cfgs.out_tile_d0 + out_block_spec = pl.BlockSpec( + (bounded_slice_gm, cfgs.dims.size_lhs_sublane, out_d0, num_lanes), + _ordered(index_map.out_index_map), + ) + else: + out_block_spec = pl.BlockSpec( + (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_n), + _ordered(index_map.out_index_map), + ) return (lhs_block_spec, rhs_block_spec, ps_block_spec), out_block_spec @@ -394,7 +510,7 @@ def inner_kernel( # In tiled_lhs_ref: LhsRef, # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_k] - tiled_rhs_ref: RhsRef, # [tile_k, tile_n] + tiled_rhs_ref: RhsRef, # [tile_k, tile_n] or [tile_n, tile_k] # Partial Sum tiled_ps_ref: jax.Array | None, # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_n] @@ -418,29 +534,52 @@ def inner_kernel( Args: tiled_lhs_ref: Contains value lhs[m_start:m_end, k_start:k_end] - tiled_rhs_ref: Contains value rhs[g_id, k_start:k_end, n_start:n_end]. where - g_id is the group associated with lhs[m_start:m_end, :] + tiled_rhs_ref: Contains the rhs tile for g_id. The tile is [tile_k, tile_n] + normally and [tile_n, tile_k] when transpose_rhs is enabled. tiled_out_ref: Contains value out[m_start:m_end, n_start:n_end] partial_out_ref: Contains last size_lhs_sublane rows of the previous output. Will be initialized to zero if this is first tile for grid[n_id, :, :]. acc_ref: Reference to the accumulator. metadata_ref: Reference to the metadata. cfgs: GmmConfigs. + + With cfgs.return_pre_act, tiled_out_ref and partial_out_ref are (act, gate, + up) triples. The activation takes the regular store path; gate and up are + stored from the same accumulator before the activation is applied. """ + pre_act_refs = () + if cfgs.return_pre_act: + (tiled_out_ref, *tiled_pre_act_refs), (partial_out_ref, *partial_pre_act_refs) = tiled_out_ref, partial_out_ref + pre_act_refs = tuple(zip(tiled_pre_act_refs, partial_pre_act_refs)) + def _matmul(is_first_k_step: bool, is_last_k_step: bool): tpu_info = pltpu.get_tpu_info() mxu_size = tpu_info.mxu_column_size # Step 1: Input pre-processing. - tiled_lhs = tiled_lhs_ref.get_value().reshape(-1, cfgs.tiles.tile_k)[...] - tiled_rhs = tiled_rhs_ref.get_weight() + lhs_chunk_d0 = cfgs.lhs_chunk_d0 + if cfgs.lhs_window: + # Block is [tile_m // sublane, sublane, D0, 128]; relayout the K window + # k_id of the full-D0 slab into [tile_m, tile_k]. + tiled_lhs = relayout.load_3d_window_as_2d( + tiled_lhs_ref.value, cfgs.tiles.tile_m, cfgs.lhs_tile_d0, pl.program_id(2) + ) + elif lhs_chunk_d0 is not None: + tiled_lhs: Any = None # Relayouted chunk by chunk in the matmul loop below. + elif cfgs.lhs_is_3d: + # Block is [tile_m // sublane, sublane, tile_d0, 128]; relayout into + # [tile_m, tile_k] with tokens on sublanes. + tiled_lhs = relayout.load_3d_as_2d(tiled_lhs_ref.value, cfgs.tiles.tile_m, cfgs.lhs_tile_d0) + else: + tiled_lhs = tiled_lhs_ref.get_value().reshape(-1, cfgs.tiles.tile_k)[...] + tiled_rhs = tiled_rhs_ref.get_weight(transpose_rhs=cfgs.transpose_rhs) # When rhs is packed (quantized dtype packed into uint32), unpack it # back to the original dtype using pltpu.bitcast which operates on K # axis. This expands the K dimension back to tile_k. if cfgs.rhs_cfgs.should_bitcast: tiled_rhs = pltpu.bitcast(tiled_rhs, cfgs.rhs_cfgs.dtype) - rhs_tile_n = tiled_rhs.shape[1] + rhs_tile_n = tiled_rhs.shape[0 if cfgs.transpose_rhs else 1] # This should only be taken in the case where we don't requantize # the scales and thus we need to dequantize inside VMEM to avoid small @@ -449,22 +588,61 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): rhs_qbs = cfgs.rhs_cfgs.quant_block_size tiled_rhs_scale = tiled_rhs_ref.get_scale().astype(acc_ref.dtype) num_blocks = cfgs.num_quant_blocks_per_tile_k - tiled_rhs_dequant = tiled_rhs.astype(acc_ref.dtype).reshape(num_blocks, rhs_qbs, rhs_tile_n) - tiled_rhs_dequant = tiled_rhs_dequant * tiled_rhs_scale - tiled_rhs = tiled_rhs_dequant.reshape(cfgs.tiles.tile_k, rhs_tile_n) + if cfgs.transpose_rhs: + tiled_rhs_dequant = tiled_rhs.astype(acc_ref.dtype).reshape(rhs_tile_n, num_blocks, rhs_qbs) + tiled_rhs_dequant = tiled_rhs_dequant * jnp.moveaxis(tiled_rhs_scale, -1, 0) + tiled_rhs = tiled_rhs_dequant.reshape(rhs_tile_n, cfgs.tiles.tile_k) + else: + tiled_rhs_dequant = tiled_rhs.astype(acc_ref.dtype).reshape(num_blocks, rhs_qbs, rhs_tile_n) + tiled_rhs_dequant = tiled_rhs_dequant * tiled_rhs_scale + tiled_rhs = tiled_rhs_dequant.reshape(cfgs.tiles.tile_k, rhs_tile_n) valid_k = cfgs.dims.size_k % cfgs.tiles.tile_k if is_last_k_step and valid_k != 0: - mask_rhs = lax.broadcasted_iota(jnp.int32, tiled_rhs.shape, 0) < valid_k + k_axis = 1 if cfgs.transpose_rhs else 0 + mask_rhs = lax.broadcasted_iota(jnp.int32, tiled_rhs.shape, k_axis) < valid_k tiled_rhs = jnp.where(mask_rhs, tiled_rhs, 0) - # The lhs K-tail is read out of bounds and may hold stale NaN/Inf, and - # NaN * 0 is NaN, so masking rhs alone is not enough. - mask_lhs = lax.broadcasted_iota(jnp.int32, tiled_lhs.shape, 1) < valid_k - tiled_lhs = jnp.where(mask_lhs, tiled_lhs, 0) + # The lhs block past size_k is stale VMEM too: zeroing only rhs still lets a NaN / Inf there poison the + # accumulator (NaN * 0 = NaN) and the in-kernel lhs quantization's block absmax. Only the 2D lhs path gets + # here: 3D lhs requires size_k % tile_k == 0 (make_gmm_configs). + if not cfgs.lhs_is_3d: + mask_lhs = lax.broadcasted_iota(jnp.int32, tiled_lhs.shape, 1) < valid_k + tiled_lhs = jnp.where(mask_lhs, tiled_lhs, 0) + + def rhs_block(start_k: int, end_k: int, start_n: int, end_n: int): + if cfgs.transpose_rhs: + return tiled_rhs[start_n:end_n, start_k:end_k] + return tiled_rhs[start_k:end_k, start_n:end_n] # Step 2: Matmul. acc_list = [] - if cfgs.lhs_cfgs.quant_dtype is None: + if lhs_chunk_d0 is not None: + # Unquantized, single quant block: relayout the 3D lhs K tile in chunks and + # accumulate each [tile_m, mxu] output column block across the chunks. + num_lanes = pltpu.get_tpu_info().num_lanes + n_starts = list(range(0, rhs_tile_n, mxu_size)) + acc_ns: list[Any] = [None] * len(n_starts) + for c0 in range(0, cfgs.lhs_tile_d0, lhs_chunk_d0): + lhs_c = relayout.load_3d_as_2d(tiled_lhs_ref.value, cfgs.tiles.tile_m, lhs_chunk_d0, c0) + start_k, end_k = c0 * num_lanes, (c0 + lhs_chunk_d0) * num_lanes + for i, start_n in enumerate(n_starts): + end_n = min(rhs_tile_n, start_n + mxu_size) + rhs_b = rhs_block(start_k, end_k, start_n, end_n) + if cfgs.transpose_rhs: + part = lax.dot_general(lhs_c, rhs_b, (((1,), (1,)), ((), ())), preferred_element_type=jnp.float32) + else: + part = jnp.matmul(lhs_c, rhs_b, preferred_element_type=jnp.float32) + acc_ns[i] = part if acc_ns[i] is None else acc_ns[i] + part + for i, start_n in enumerate(n_starts): + end_n = min(rhs_tile_n, start_n + mxu_size) + block_acc = acc_ns[i].astype(acc_ref.dtype) + if cfgs.rhs_cfgs.should_dequantize_after_matmul: + tiled_rhs_scale = tiled_rhs_ref.get_scale() + block_acc *= tiled_rhs_scale[0, :, start_n:end_n].astype(acc_ref.dtype) + acc_n = jnp.zeros((cfgs.tiles.tile_m, end_n - start_n), dtype=acc_ref.dtype) + acc_n += block_acc + acc_list.append(acc_n) + elif cfgs.lhs_cfgs.quant_dtype is None: # Unquantized matmul path. rhs_qbs = cfgs.rhs_cfgs.quant_block_size @@ -477,11 +655,21 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): start_k = b_id * rhs_qbs # pyrefly: ignore[unsupported-operation] end_k = start_k + rhs_qbs # pyrefly: ignore[unsupported-operation] - block_acc = jnp.matmul( - tiled_lhs[:, start_k:end_k], - tiled_rhs[start_k:end_k, start_n:end_n], - preferred_element_type=jnp.float32, - ).astype(acc_ref.dtype) + block_lhs = tiled_lhs[:, start_k:end_k] + block_rhs = rhs_block(start_k, end_k, start_n, end_n) + if cfgs.transpose_rhs: + block_acc = lax.dot_general( + block_lhs, + block_rhs, + (((1,), (1,)), ((), ())), + preferred_element_type=jnp.float32, + ).astype(acc_ref.dtype) + else: + block_acc = jnp.matmul( + block_lhs, + block_rhs, + preferred_element_type=jnp.float32, + ).astype(acc_ref.dtype) if cfgs.rhs_cfgs.should_dequantize_after_matmul: tiled_rhs_scale = tiled_rhs_ref.get_scale() @@ -524,7 +712,7 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): end_k = min(cfgs.tiles.tile_k, start_k + q_block_size) # pyrefly: ignore[unsupported-operation] block_lhs = tiled_lhs[:, start_k:end_k] - block_rhs = tiled_rhs[start_k:end_k, start_n:end_n] + block_rhs = rhs_block(start_k, end_k, start_n, end_n) # Perform lhs quantization. Note that for every block_lhs, # same computation will be performed tiles_n//mxu_size times. @@ -551,11 +739,19 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): if not tpu_info.is_matmul_supported(lhs_q_dtype, block_rhs.dtype): block_rhs = block_rhs.astype(lhs_q_dtype) - block_acc = jnp.matmul( - block_lhs_q, - block_rhs, - preferred_element_type=preferred_element_type, - ).astype(acc_ref.dtype) + if cfgs.transpose_rhs: + block_acc = lax.dot_general( + block_lhs_q, + block_rhs, + (((1,), (1,)), ((), ())), + preferred_element_type=preferred_element_type, + ).astype(acc_ref.dtype) + else: + block_acc = jnp.matmul( + block_lhs_q, + block_rhs, + preferred_element_type=preferred_element_type, + ).astype(acc_ref.dtype) block_acc *= block_scale.astype(acc_ref.dtype) @@ -574,56 +770,57 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): acc += acc_ref[...] if is_last_k_step: + if cfgs.lhs_cfgs.has_scale and cfgs.lhs_cfgs.quant_dtype is None: + acc *= tiled_lhs_ref.get_scale().astype(acc.dtype) if cfgs.rhs_cfgs.has_bias: tiled_rhs_bias = tiled_rhs_ref.get_bias() acc += tiled_rhs_bias.astype(acc.dtype) if cfgs.has_partial_sum: - ps_tile = tiled_ps_ref[...].reshape(acc.shape) # pyrefly: ignore[unsupported-operation] + if cfgs.out_window: + ps_tile = relayout.load_3d_window_as_2d( + tiled_ps_ref, cfgs.tiles.tile_m, cfgs.out_tile_d0, pl.program_id(cfgs.n_axis) + ) + elif cfgs.out_is_3d: + ps_tile = relayout.load_3d_as_2d(tiled_ps_ref, cfgs.tiles.tile_m, cfgs.out_tile_d0) + else: + ps_tile = tiled_ps_ref[...].reshape(acc.shape) # pyrefly: ignore[unsupported-operation] acc += ps_tile.astype(acc.dtype) - acc = apply_act_fn(acc, cfgs.fuse_act) - - gm_id = pl.program_id(1) - - # Mask out rows that does not belong to the current group. - m_start = metadata_ref.gm_id_to_m_offset[gm_id] - m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1] - m_offset = m_start - m_start % cfgs.dims.size_lhs_sublane - - m_start_local = m_start - m_offset - m_end_local = m_end - m_offset + if pre_act_refs: + # acc is [tile_m, 2 * tile_n] = gate | up. Store both halves before the activation, rounded to the + # output dtype exactly as the unfused kernel rounds its output; the activation below stays in f32. + for (tiled_pre_ref, partial_pre_ref), acc_half in zip(pre_act_refs, jnp.split(acc, 2, -1)): + _store_group_rows(tiled_pre_ref, partial_pre_ref, acc_half, metadata_ref, cfgs) - iota = lax.broadcasted_iota(jnp.int32, acc.shape, 0) - mask = jnp.logical_and(m_start_local <= iota, iota < m_end_local) - acc_masked = jnp.where(mask, acc, 0).reshape(tiled_out_ref.shape) + acc = apply_act_fn(acc, cfgs.fuse_act) # Write the final output to the output ref. - tiled_out_ref[...] = acc_masked.astype(tiled_out_ref.dtype) - - # If this is the first tile for grid[n_id, :, :], we initialize the - # partial out to zeros. Otherwise, partial out from last tile of - # grid[n_id-1, :, :] can be used and cause numeric issues. - partial_out_zeros = jnp.zeros_like(partial_out_ref) - - # Accumulate the partial output from the previous step. - tiled_out_ref[0] += jnp.where(gm_id == 0, partial_out_zeros, partial_out_ref[...]) - - # Consider following case where size_lhs_sublane = 4, number denotes group - # id and | denotes boundaries between sublanes: - # | 0 0 1 2 | 2 2 2 2 | 3 3 4 4 | - # - # Assuming group id of current step is 1, current step will not completely - # fill size_lhs_sublane rows and will be revisited at the next step. By - # storing the partial rows into the partial_out_ref, the next step can - # read them and accumulate to them. Additionally, for group id of 2, - # since it completely fills the size_lhs_sublane rows, we need to zero out - # partial_out_ref to avoid numeric error for group 3. - last_row = m_end_local // cfgs.dims.size_lhs_sublane - partial_out_ref[...] = jnp.where( - m_end_local % cfgs.dims.size_lhs_sublane == 0, - partial_out_zeros, - tiled_out_ref[last_row], + if not cfgs.out_is_3d: + _store_group_rows(tiled_out_ref, partial_out_ref, acc, metadata_ref, cfgs) + return + acc_masked, gm_id, m_end_local = _mask_group_rows(acc, metadata_ref, cfgs) + if cfgs.out_window: + # Full-D0 out slab resident across the N tiles of this gm tile. The + # cross-gm partial-row carry is kept per N tile, in the 2D domain. + sub = cfgs.dims.size_lhs_sublane + n_id = pl.program_id(cfgs.n_axis) + out2d = acc_masked.astype(tiled_out_ref.dtype) + prev = partial_out_ref[n_id] + prev = jnp.where(gm_id == 0, jnp.zeros_like(prev), prev) + out2d = jnp.concatenate([out2d[:sub] + prev, out2d[sub:]], axis=0) + # Carry the (possibly shared) last sublane row to the next gm tile. + acc_ref[...] = out2d.astype(acc_ref.dtype) + last_row = m_end_local // sub + last = acc_ref[pl.ds(pl.multiple_of(last_row * sub, sub), sub), :].astype(out2d.dtype) + partial_out_ref[n_id] = jnp.where(m_end_local % sub == 0, jnp.zeros_like(last), last) + relayout.store_2d_as_3d_window(tiled_out_ref, out2d, out2d.shape[0], cfgs.out_tile_d0, n_id) + return + # Block is [tile_m // sublane, sublane, tile_d0, 128]; scatter the + # [tile_m, tile_n] rows back into per-token [tile_d0, 128] slabs. + relayout.store_2d_as_3d( + tiled_out_ref, acc_masked.astype(tiled_out_ref.dtype), acc_masked.shape[0], cfgs.out_tile_d0 ) + _carry_partial_rows(tiled_out_ref, partial_out_ref, gm_id, m_end_local, cfgs) else: acc_ref[...] = acc @@ -666,6 +863,80 @@ def matmul_last(): ) +def _mask_group_rows(acc: jax.Array, metadata_ref: MetadataRef, cfgs: GmmConfigs): + """Zeroes the rows of `acc` that do not belong to the current gm tile's group. + + Returns the masked `acc`, the gm tile id and the group's end row relative to + the tile's first sublane row (what _carry_partial_rows needs). + """ + gm_id = pl.program_id(cfgs.gm_axis) + + # Mask out rows that does not belong to the current group. + m_start = metadata_ref.gm_id_to_m_offset[gm_id] + m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1] + m_offset = m_start - m_start % cfgs.dims.size_lhs_sublane + + m_start_local = m_start - m_offset + m_end_local = m_end - m_offset + + iota = lax.broadcasted_iota(jnp.int32, acc.shape, 0) + mask = jnp.logical_and(m_start_local <= iota, iota < m_end_local) + return jnp.where(mask, acc, 0), gm_id, m_end_local + + +def _carry_partial_rows( + tiled_out_ref: jax.Array, + partial_out_ref: jax.Array, + gm_id: jax.Array, + m_end_local: jax.Array, + cfgs: GmmConfigs, +): + """Adds the sublane rows the previous gm tile left in `partial_out_ref`; leaves this tile's unfinished ones there.""" + # If this is the first tile for grid[n_id, :, :], we initialize the + # partial out to zeros. Otherwise, partial out from last tile of + # grid[n_id-1, :, :] can be used and cause numeric issues. + partial_out_zeros = jnp.zeros_like(partial_out_ref) + + # Accumulate the partial output from the previous step. + tiled_out_ref[0] += jnp.where(gm_id == 0, partial_out_zeros, partial_out_ref[...]) + + # Consider following case where size_lhs_sublane = 4, number denotes group + # id and | denotes boundaries between sublanes: + # | 0 0 1 2 | 2 2 2 2 | 3 3 4 4 | + # + # Assuming group id of current step is 1, current step will not completely + # fill size_lhs_sublane rows and will be revisited at the next step. By + # storing the partial rows into the partial_out_ref, the next step can + # read them and accumulate to them. Additionally, for group id of 2, + # since it completely fills the size_lhs_sublane rows, we need to zero out + # partial_out_ref to avoid numeric error for group 3. + last_row = m_end_local // cfgs.dims.size_lhs_sublane + partial_out_ref[...] = jnp.where( + m_end_local % cfgs.dims.size_lhs_sublane == 0, + partial_out_zeros, + tiled_out_ref[last_row], + ) + + +def _store_group_rows( + tiled_ref: jax.Array, # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_n] + partial_ref: jax.Array, # [size_lhs_sublane, tile_n] + acc: jax.Array, # [tile_m, tile_n] + metadata_ref: MetadataRef, + cfgs: GmmConfigs, +): + """The 2D store of inner_kernel's last k step: the current group's rows of `acc`, with the partial-row carry. + + Zeroes the rows of other groups, rounds to the output dtype, stores the tile + and carries its unfinished last sublane rows to the next gm tile. Used for the + plain 2D output and for the extra (return_pre_act) outputs written from the + same accumulator. + """ + acc_masked, gm_id, m_end_local = _mask_group_rows(acc, metadata_ref, cfgs) + tiled_ref[...] = acc_masked.reshape(tiled_ref.shape).astype(tiled_ref.dtype) + _carry_partial_rows(tiled_ref, partial_ref, gm_id, m_end_local, cfgs) + + def fill_metadata( lhs_group_sizes_ref: jax.Array, # int32[size_lhs_group] group_offset_ref: jax.Array, # int32[1] @@ -769,8 +1040,15 @@ def zero_out_start( assert num_lanes == zero_ref.shape[-1] zero_ref[...] = jnp.zeros_like(zero_ref) - zero_dma = zero_ref.reshape(-1, dims.size_lhs_sublane, num_lanes) - out_dma = out_ref.reshape(-1, dims.size_lhs_sublane, out_ref.shape[-1]) + out_is_3d = out_ref.ndim == 3 + if out_is_3d: + # [size_m, D0, 128] -> [size_m // sublane, sublane, D0, 128]. Each DMA + # zeroes whole [D0, 128] slabs. + zero_dma = zero_ref.reshape(-1, dims.size_lhs_sublane, *zero_ref.shape[1:]) + out_dma = out_ref.reshape(-1, dims.size_lhs_sublane, *out_ref.shape[1:]) + else: + zero_dma = zero_ref.reshape(-1, dims.size_lhs_sublane, num_lanes) + out_dma = out_ref.reshape(-1, dims.size_lhs_sublane, out_ref.shape[-1]) row_size = zero_dma.shape[0] compute_start = metadata_ref.gm_id_to_m_offset[0] @@ -791,14 +1069,21 @@ def fill_zero(i, zero_size, *, start, end): dma_end = jnp.minimum(dma_start + row_size, end) dma_size = dma_end - dma_start - # Static loop. Will be unrolled during compile time. - for n_start in range(0, out_dma.shape[-1], num_lanes): - n_end = n_start + num_lanes + if out_is_3d: pltpu.make_async_copy( src_ref=zero_dma.at[pl.ds(0, dma_size)], - dst_ref=out_dma.at[pl.ds(dma_start, dma_size), :, n_start:n_end], + dst_ref=out_dma.at[pl.ds(dma_start, dma_size)], sem=semaphore_ref.at[0], ).start(priority=1) + else: + # Static loop. Will be unrolled during compile time. + for n_start in range(0, out_dma.shape[-1], num_lanes): + n_end = n_start + num_lanes + pltpu.make_async_copy( + src_ref=zero_dma.at[pl.ds(0, dma_size)], + dst_ref=out_dma.at[pl.ds(dma_start, dma_size), :, n_start:n_end], + sem=semaphore_ref.at[0], + ).start(priority=1) return zero_size + dma_size @@ -822,7 +1107,7 @@ def zero_out_end( *, dims: Dimensions, ): - out_dma = out_ref.reshape(-1, dims.size_lhs_sublane, out_ref.shape[-1]) + out_dma = out_ref.reshape(-1, dims.size_lhs_sublane, *out_ref.shape[1:]) pltpu.make_async_copy( src_ref=out_dma.at[pl.ds(0, zero_size)], dst_ref=out_dma.at[pl.ds(0, zero_size)], @@ -848,6 +1133,8 @@ def kernel_main( semaphore_ref: jax.Array | None, # [1] *, cfgs: GmmConfigs, + pre_act_out_refs: tuple[jax.Array, ...] = (), # return_pre_act: (gate, up) outputs, each shaped like out_ref + pre_act_partial_refs: tuple[jax.Array, ...] = (), # return_pre_act: their partial-row scratch ): """Entry point for GMM kernel. @@ -903,11 +1190,31 @@ def kernel_main( num_gm, dims=cfgs.dims, ) + # The pre-activation outputs have the act output's shape, so they zero the same rows with DMAs of the + # same size on the same semaphore; zero_out_end below waits for them one output at a time. + for pre_act_out_ref in pre_act_out_refs: + zero_out_start( + pre_act_out_ref, + zero_ref, # pyrefly: ignore[bad-argument-type] + semaphore_ref, # pyrefly: ignore[bad-argument-type] + metadata_ref, + num_gm, + dims=cfgs.dims, + ) (lhs_spec, rhs_spec, ps_spec), out_spec = generate_block_specs(metadata_ref, cfgs) + if pre_act_out_refs: + # (act, gate, up) tiles share one index map; inner_kernel receives them as one tuple. + out_spec = ((out_spec,) * (1 + len(pre_act_out_refs)),) if cfgs.fuse_act is not None: - rhs_up_ref = jax.tree.map(lambda x: x.at[..., cfgs.out_size_n :], rhs_ref) + if cfgs.transpose_rhs: + rhs_up_weight = rhs_ref.weight.at[:, cfgs.out_size_n :, :] + else: + rhs_up_weight = rhs_ref.weight.at[..., cfgs.out_size_n :] + rhs_up_scale = None if rhs_ref.scale is None else rhs_ref.scale.at[..., cfgs.out_size_n :] + rhs_up_bias = None if rhs_ref.bias is None else rhs_ref.bias.at[..., cfgs.out_size_n :] + rhs_up_ref = WeightsRef(weight=rhs_up_weight, scale=rhs_up_scale, bias=rhs_up_bias) rhs_ref = FusedWeightsRef(gate=rhs_ref, up=rhs_up_ref) # pyrefly: ignore[bad-assignment] rhs_spec = FusedWeightsRef( @@ -918,7 +1225,7 @@ def kernel_main( # Execute the inner kernel. pipeline_fn = pltpu.emit_pipeline( functools.partial(inner_kernel, cfgs=cfgs), - grid=(num_n, num_gm, num_k), + grid=(num_gm, num_n, num_k) if cfgs.out_window else (num_n, num_gm, num_k), in_specs=(lhs_spec, rhs_spec, ps_spec), out_specs=out_spec, ) @@ -926,17 +1233,67 @@ def kernel_main( # Bounded slice requires second last dim to be aligned to the sublane size. # rhs_ref uses static tiling thus reshape is not needed. The lhs quant scale # (when present) is small and statically tiled, so it is passed through as-is. - lhs_value_in = lhs_ref.value.reshape(-1, cfgs.dims.size_lhs_sublane, lhs_ref.value.shape[-1]) + if cfgs.lhs_is_3d: + # [size_m, D0, 128] -> [size_m // sublane, sublane, D0, 128]. + lhs_value_in = lhs_ref.value.reshape(-1, cfgs.dims.size_lhs_sublane, *lhs_ref.value.shape[-2:]) + else: + lhs_value_in = lhs_ref.value.reshape(-1, cfgs.dims.size_lhs_sublane, lhs_ref.value.shape[-1]) lhs_in = LhsRef(value=lhs_value_in, scale=lhs_ref.scale) ps_in = None if cfgs.has_partial_sum: - ps_in = partial_sum_ref.reshape(-1, cfgs.dims.size_lhs_sublane, partial_sum_ref.shape[-1]) - out_in = out_ref.reshape(-1, cfgs.dims.size_lhs_sublane, out_ref.shape[-1]) + # Works for both [size_m, size_n] and the 3D [size_m, D0, 128] partial_sum. + ps_in = partial_sum_ref.reshape(-1, cfgs.dims.size_lhs_sublane, *partial_sum_ref.shape[1:]) + # Works for both [size_m, size_n] and the 3D [size_m, D0, 128] output. + out_in = out_ref.reshape(-1, cfgs.dims.size_lhs_sublane, *out_ref.shape[1:]) scratches = [partial_out_ref, acc_ref, metadata_ref] + if pre_act_out_refs: + out_in = (out_in, *(r.reshape(-1, cfgs.dims.size_lhs_sublane, *r.shape[1:]) for r in pre_act_out_refs)) + scratches[0] = (partial_out_ref, *pre_act_partial_refs) pipeline_fn(lhs_in, rhs_ref, ps_in, out_in, scratches=scratches) if cfgs.zero_init: zero_out_end(out_ref, semaphore_ref, zero_size, dims=cfgs.dims) # pyrefly: ignore[bad-argument-type, unbound-name] + for pre_act_out_ref in pre_act_out_refs: + # pyrefly: ignore[bad-argument-type, unbound-name] + zero_out_end(pre_act_out_ref, semaphore_ref, zero_size, dims=cfgs.dims) + + +def kernel_main_pre_act( + lhs_group_sizes_ref: jax.Array, + group_offset_ref: jax.Array, + lhs_ref: LhsRef, + rhs_ref: WeightsRef, + partial_sum_ref: jax.Array, + out_ref: jax.Array, # act: [size_m, size_n // 2] + gate_out_ref: jax.Array, # gate before the activation: [size_m, size_n // 2] + up_out_ref: jax.Array, # up before the activation: [size_m, size_n // 2] + partial_out_ref: jax.Array, + acc_ref: jax.Array, + metadata_ref: MetadataRef, + zero_ref: jax.Array | None, + semaphore_ref: jax.Array | None, + partial_gate_ref: jax.Array, # [size_lhs_sublane, tile_n] + partial_up_ref: jax.Array, # [size_lhs_sublane, tile_n] + *, + cfgs: GmmConfigs, +): + """kernel_main for return_pre_act: pallas_call passes the gate and up outputs right after act.""" + kernel_main( + lhs_group_sizes_ref, + group_offset_ref, + lhs_ref, + rhs_ref, + partial_sum_ref, + out_ref, + partial_out_ref, + acc_ref, + metadata_ref, + zero_ref, + semaphore_ref, + cfgs=cfgs, + pre_act_out_refs=(gate_out_ref, up_out_ref), + pre_act_partial_refs=(partial_gate_ref, partial_up_ref), + ) def calculate_tiling( @@ -983,7 +1340,7 @@ def _is_tile_k_quant_block_compatible(tk: int) -> bool: # pyrefly: ignore[unsupported-operation] tk % rhs_cfgs.quant_block_size != 0 and rhs_cfgs.quant_block_size % tk != 0 # pyrefly: ignore[unsupported-operation] - ): # pyrefly: ignore[unsupported-operation] + ): return False return True @@ -1062,20 +1419,36 @@ def validate_inputs( fuse_act: str | None = None, maybe_quantize_lhs: bool = True, lhs_scale: jax.Array | None = None, + transpose_rhs: bool = False, ) -> Dimensions: """Validates the inputs for the GMM kernel.""" size_m = lhs.shape[0] - size_group, size_k, size_n = rhs.shape + if transpose_rhs: + size_group, size_n, size_k = rhs.shape + else: + size_group, size_k, size_n = rhs.shape size_lhs_group = group_sizes.shape[0] assert size_group <= size_lhs_group - assert lhs.shape == (size_m, size_k) - assert rhs.shape == (size_group, size_k, size_n) + if lhs.ndim == 3: + num_lanes = pltpu.get_tpu_info().num_lanes + assert lhs.shape == (size_m, size_k // num_lanes, num_lanes), (lhs.shape, size_k) + else: + assert lhs.shape == (size_m, size_k) + expected_rhs_shape = (size_group, size_n, size_k) if transpose_rhs else (size_group, size_k, size_n) + assert rhs.shape == expected_rhs_shape + if transpose_rhs and jax.dtypes.itemsize_bits(rhs.dtype) < 8: + raise ValueError("transpose_rhs is not supported for sub-byte rhs dtypes.") if rhs_bias is not None: assert rhs_bias.shape == (size_group, 1, size_n) if partial_sum is not None: - assert partial_sum.shape[-1] == size_n + if partial_sum.ndim == 3: + # 3D partial_sum [m, size_n // 128, 128] (requires a 3D out; checked in make_gmm_configs). + num_lanes = pltpu.get_tpu_info().num_lanes + assert partial_sum.shape[1:] == (size_n // num_lanes, num_lanes), (partial_sum.shape, size_n) + else: + assert partial_sum.shape[-1] == size_n # lhs's m dimension can sometimes be padded to wi_tile_fwd_batch_seq assert partial_sum.shape[0] <= size_m if rhs_scale is not None: @@ -1086,7 +1459,6 @@ def validate_inputs( assert size_k % num_quant_blocks == 0 if lhs_scale is not None: - assert maybe_quantize_lhs, "lhs_scale requires maybe_quantize_lhs=True." # Only per-tensor scales are supported for now. The current implementation generalizes to per-channel [M, 1] and # sub-channel [M, num_k_blocks]; extend the validation and the block spec / # index map together when adding those. @@ -1139,6 +1511,8 @@ def get_cost_estimate(cfgs: GmmConfigs): rhs_bytes += dims.size_group * dims.size_n * fp32_bytes out_bytes = dims.size_m * cfgs.out_size_n * cfgs.out_dtype.itemsize + if cfgs.return_pre_act: + out_bytes *= 3 # act, gate and up total_bytes = lhs_bytes + rhs_bytes + out_bytes @@ -1150,12 +1524,18 @@ def get_cost_estimate(cfgs: GmmConfigs): def get_scope_name(cfgs: GmmConfigs) -> str: + """The kernel's named scope: its dims, tiles and the flags that change the kernel.""" dims = cfgs.dims tiles = cfgs.tiles - return ( + scope_name = ( f"gmm_v2-g_{dims.size_group}-m_{dims.size_m}-k_{dims.size_k}-act_{cfgs.fuse_act}" f"-n_{dims.size_n}-tm_{tiles.tile_m}-tk_{tiles.tile_k}-tn_{tiles.tile_n}" ) + if cfgs.transpose_rhs: + scope_name += "-transpose_rhs" + if cfgs.return_pre_act: + scope_name += "-pre_act" + return scope_name def make_gmm_configs( @@ -1175,6 +1555,9 @@ def make_gmm_configs( zero_initialize: bool, fuse_act: str | None = None, lhs_scale: jax.Array | None = None, + transpose_rhs: bool = False, + out_is_3d: bool = False, + return_pre_act: bool = False, ): """Fills the GMM config for the GMM kernel.""" @@ -1186,9 +1569,10 @@ def make_gmm_configs( partial_sum, group_sizes, group_offset, - fuse_act, - maybe_quantize_lhs, - lhs_scale, + fuse_act=fuse_act, + maybe_quantize_lhs=maybe_quantize_lhs, + lhs_scale=lhs_scale, + transpose_rhs=transpose_rhs, ) if rhs_scale is not None: @@ -1226,13 +1610,13 @@ def make_gmm_configs( if not is_rhs_float: lhs_q_dtype = jnp.int8.dtype - if lhs_scale is not None: + if lhs_scale is not None and maybe_quantize_lhs: assert lhs_q_dtype is not None, ( "lhs_scale requires lhs quantization to engage, but no lhs quant " "dtype was selected. Ensure rhs is quantized and the hardware supports " "fp8/int8 matmul." ) - has_lhs_scale = lhs_scale is not None and lhs_q_dtype is not None + has_lhs_scale = lhs_scale is not None lhs_cfgs = InputConfigs( quant_dtype=lhs_q_dtype, @@ -1243,6 +1627,7 @@ def make_gmm_configs( quant_block_size=512, dtype=lhs.dtype, has_scale=has_lhs_scale, + is_3d=lhs.ndim == 3, ) if out_dtype is None: @@ -1262,6 +1647,46 @@ def make_gmm_configs( else: tiles = tile_info(dims, lhs_cfgs, rhs_cfgs, vmem_limit_bytes, fuse_act, partial_sum is not None) + num_lanes = pltpu.get_tpu_info().num_lanes + if lhs_cfgs.is_3d: + if tiles.tile_k % num_lanes or dims.size_k % tiles.tile_k: + raise ValueError(f"3D lhs requires tile_k % 128 == 0 and size_k % tile_k == 0; got {tiles=}, {dims=}.") + lhs_tile_d0, lhs_full_d0 = tiles.tile_k // num_lanes, dims.size_k // num_lanes + if relayout.needs_window(lhs_tile_d0, lhs_full_d0): + relayout.check_window(lhs_tile_d0, lhs_full_d0, lhs.dtype, store=False) + else: + relayout.check_tile_d0(lhs_tile_d0, full_d0=lhs_full_d0, dtype=lhs.dtype) + if tiles.tile_m % dims.size_lhs_sublane: + raise ValueError(f"3D lhs requires tile_m % sublane == 0; got {tiles=}, {dims=}.") + if out_is_3d: + out_size_n = dims.size_n // 2 if fuse_act is not None else dims.size_n + if tiles.tile_n % num_lanes or out_size_n % tiles.tile_n: + raise ValueError(f"3D out requires tile_n % 128 == 0 and out_size_n % tile_n == 0; got {tiles=}, {dims=}.") + out_tile_d0, out_full_d0 = tiles.tile_n // num_lanes, out_size_n // num_lanes + if relayout.needs_window(out_tile_d0, out_full_d0): + relayout.check_window(out_tile_d0, out_full_d0, out_dtype, store=True) + if fuse_act is not None: + raise NotImplementedError("A 3D out window does not support fuse_act.") + else: + relayout.check_tile_d0(out_tile_d0, full_d0=out_full_d0, dtype=out_dtype) + if partial_sum is not None: + if partial_sum.ndim != 3 or partial_sum.shape[0] != dims.size_m: + raise ValueError( + f"A 3D output requires a 3D partial_sum of shape [size_m, out_size_n // 128, 128]; got {partial_sum.shape}." + ) + if fuse_act is not None: + raise NotImplementedError("A 3D partial_sum does not support fuse_act.") + ps_d0 = out_full_d0 if relayout.needs_window(out_tile_d0, out_full_d0) else out_tile_d0 + relayout.check_tile_d0(ps_d0, full_d0=out_full_d0, dtype=partial_sum.dtype) + if relayout.needs_window(out_tile_d0, out_full_d0): + relayout.check_window(out_tile_d0, out_full_d0, partial_sum.dtype, store=False) + if tiles.tile_m % 2: + raise ValueError(f"3D out requires an even tile_m; got {tiles=}.") + elif partial_sum is not None and partial_sum.ndim == 3: + raise ValueError("A 3D partial_sum requires out_is_3d=True.") + if return_pre_act and (fuse_act is None or out_is_3d or partial_sum is not None): + raise NotImplementedError("return_pre_act needs fuse_act and a 2D output, and no partial_sum.") + return GmmConfigs( dims=dims, tiles=tiles, @@ -1272,11 +1697,17 @@ def make_gmm_configs( has_partial_sum=partial_sum is not None, zero_init=zero_initialize, fuse_act=fuse_act, + transpose_rhs=transpose_rhs, + out_is_3d=out_is_3d, + return_pre_act=return_pre_act, ) def get_metadata(cfgs: GmmConfigs) -> dict[str, str | int | float]: + """The kernel config flattened to the pallas_call metadata dict.""" cfgs_dict = dataclasses.asdict(cfgs) + if not cfgs.return_pre_act: + del cfgs_dict["return_pre_act"] # keeps the metadata of every other gmm_v2 call unchanged ret = {} for path, val in jax.tree_util.tree_leaves_with_path(cfgs_dict): key = jax.tree_util.keystr(path, simple=True, separator=".") @@ -1296,11 +1727,14 @@ def get_metadata(cfgs: GmmConfigs) -> dict[str, str | int | float]: "maybe_quantize_lhs", "zero_initialize", "fuse_act", + "transpose_rhs", + "out_is_3d", + "return_pre_act", ] ) def gmm_v2( lhs: jax.Array, # [size_m, size_k] - rhs: jax.Array, # [size_group, size_k, size_n] + rhs: jax.Array, # [size_group, size_k, size_n], or [size_group, size_n, size_k] if transposed group_sizes: jax.Array, # int32[size_lhs_group] rhs_scale: jax.Array | None = None, # [size_group, num_blocks, 1, out_size] rhs_bias: jax.Array | None = None, # [size_group, 1, out_size] @@ -1316,7 +1750,10 @@ def gmm_v2( maybe_quantize_lhs: bool = True, zero_initialize: bool = True, fuse_act: str | None = None, -) -> jax.Array: + transpose_rhs: bool = False, + out_is_3d: bool = False, + return_pre_act: bool = False, +) -> jax.Array | tuple[jax.Array, jax.Array, jax.Array]: """GMM kernel implemented with emit_pipeline. Dynamically calculate offset lhs/out tiles to reduce redundant computations. @@ -1325,18 +1762,19 @@ def gmm_v2( Args: lhs: lhs with shape [size_m, size_k]. - rhs: rhs with shape [size_group, size_k, size_n]. + rhs: rhs with shape [size_group, size_k, size_n], or + [size_group, size_n, size_k] when transpose_rhs=True. group_sizes: The group sizes of lhs rows of shape [size_lhs_group,]. rhs_scale: The rhs scale of shape [size_group, num_blocks, 1, out_size]. rhs_bias: The rhs bias of shape [size_group, 1, out_size]. partial_sum: Optional. Per-token partial sums of shape [size_m, size_n]. group_offset: Optional. The group offset of shape [1,]. - lhs_scale: Optional scale used to quantize the (unquantized) lhs - inside the kernel and the result is multiplied back by `scale`. The shape - encodes granularity; currently only per-tensor `[1, 1]` is supported. When - None, a quantized lhs uses the default dynamic per-block absmax - calibration. Only takes effect when maybe_quantize_lhs is True and rhs is - quantized. + lhs_scale: Optional per-tensor `[1, 1]` lhs scale. With maybe_quantize_lhs + and a quantized rhs, it is used to quantize the (unquantized) lhs inside + the kernel and the result is multiplied back by `scale`; when None, a + quantized lhs uses the default dynamic per-block absmax calibration. + Otherwise the accumulator is multiplied by `scale` after the matmul, + before rhs_bias and partial_sum are added. tile_info: The tile sizes or tile function to use. vmem_limit_bytes: Optional vmem limit in bytes. precision: Unused. Exists for compatibility reasons. @@ -1345,9 +1783,24 @@ def gmm_v2( maybe_quantize_lhs: Quantize lhs if set to True and rhs is quantized. zero_initialize: Whether to initialize unvisited output elements to zero. fuse_act: Activation function to fuse with GMM, None if no fusion. + transpose_rhs: Interpret rhs as [size_group, size_n, size_k] and contract + its last dimension without materializing a full rhs transpose. + out_is_3d: If True, the output is written directly in the 3D layout + [size_m, out_size_n // 128, 128], relayouted from [tile_m, tile_n] inside + VMEM. A 3D lhs ([size_m, size_k // 128, 128]) is detected from lhs.ndim. + Both require tile_k / tile_n to be a multiple of 128 that divides the + dim. Tiles whose D0 (tile // 128) is not a multiple of 8 use full-D0 + blocks and relayout the tile window in VMEM. + return_pre_act: With fuse_act, also return the gate and up halves of the + matmul before the activation (what the activation's backward needs), each + rounded to the output dtype as the unfused kernel rounds them. The + activation itself is computed from the f32 accumulator. 2D output only, + no partial_sum. Returns: - Output of shape [size_m, size_n]. + Output of shape [size_m, size_n] (or [size_m, size_n // 128, 128]); with + fuse_act the output width is size_n // 2. With return_pre_act, the tuple + (act, gate, up), each of shape [size_m, size_n // 2]. """ del precision @@ -1377,6 +1830,9 @@ def gmm_v2( zero_initialize=zero_initialize, fuse_act=fuse_act, lhs_scale=lhs_scale, + transpose_rhs=transpose_rhs, + out_is_3d=out_is_3d, + return_pre_act=return_pre_act, ) dims = cfgs.dims tiles = cfgs.tiles @@ -1401,9 +1857,16 @@ def gmm_v2( # Initialize scratch shapes. max_num_gm = dims.size_group + pl.cdiv(dims.size_m, tiles.tile_m) - 1 acc_cols = 2 * tiles.tile_n if cfgs.fuse_act is not None else tiles.tile_n + num_lanes = pltpu.get_tpu_info().num_lanes + if cfgs.out_window: + partial_out_shape = (pl.cdiv(cfgs.out_size_n, tiles.tile_n), dims.size_lhs_sublane, tiles.tile_n) + elif cfgs.out_is_3d: + partial_out_shape = (dims.size_lhs_sublane, cfgs.out_tile_d0, num_lanes) + else: + partial_out_shape = (dims.size_lhs_sublane, tiles.tile_n) scratch_shapes = [ # partial_out_ref - pltpu.VMEM((dims.size_lhs_sublane, tiles.tile_n), cfgs.out_dtype), + pltpu.VMEM(partial_out_shape, cfgs.out_dtype), # acc_ref pltpu.VMEM((tiles.tile_m, acc_cols), cfgs.acc_dtype), # metadata_ref @@ -1413,7 +1876,6 @@ def gmm_v2( ), ] - num_lanes = pltpu.get_tpu_info().num_lanes if cfgs.zero_init: # TODO: Create better heuristics for determining this value. target_zero_ref_bytes = 2 * 1024 * 1024 @@ -1429,18 +1891,37 @@ def gmm_v2( # size_n//num_lanes times in a single tile, we can significantly increase # tile_zero_m without triggering OOM. out_bytes = jnp.dtype(cfgs.out_dtype).itemsize - tile_zero_m = target_zero_ref_bytes // num_lanes // out_bytes - tile_zero_m = min(tile_zero_m, dims.size_m) + if cfgs.out_is_3d: + # A row is a whole [D0, 128] slab; the zero buffer holds full slabs, + # tile_zero_m a multiple of the sublane block. VMEM pads D0 up to the + # sublane tiling, so budget with the padded size. + out_d0 = align_to(cfgs.out_size_n, num_lanes) // num_lanes + sublane = dims.size_lhs_sublane + sublane_tiling = pltpu.get_tpu_info().get_sublane_tiling(cfgs.out_dtype) + row_bytes = align_to(out_d0, sublane_tiling) * num_lanes * out_bytes + tile_zero_m = target_zero_ref_bytes // row_bytes // sublane * sublane + tile_zero_m = max(sublane, min(tile_zero_m, dims.size_m)) + zero_shape = (tile_zero_m, out_d0, num_lanes) + else: + tile_zero_m = target_zero_ref_bytes // num_lanes // out_bytes + tile_zero_m = min(tile_zero_m, dims.size_m) + zero_shape = (tile_zero_m, num_lanes) scratch_shapes += [ - pltpu.VMEM((tile_zero_m, num_lanes), cfgs.out_dtype), + pltpu.VMEM(zero_shape, cfgs.out_dtype), pltpu.SemaphoreType.DMA((1,)), ] else: scratch_shapes += [None, None] + if cfgs.return_pre_act: + # Partial-row carries of the gate and up outputs (kernel_main_pre_act takes them last). + scratch_shapes += [pltpu.VMEM(partial_out_shape, cfgs.out_dtype) for _ in range(2)] aligned_n = align_to(cfgs.out_size_n, num_lanes) - out_init = jax.ShapeDtypeStruct((dims.size_m, aligned_n), cfgs.out_dtype) + if cfgs.out_is_3d: + out_init = jax.ShapeDtypeStruct((dims.size_m, aligned_n // num_lanes, num_lanes), cfgs.out_dtype) + else: + out_init = jax.ShapeDtypeStruct((dims.size_m, aligned_n), cfgs.out_dtype) lhs_in = LhsRef(value=lhs, scale=lhs_scale) rhs_weights = WeightsRef(weight=rhs, scale=rhs_scale, bias=rhs_bias) partial_sum_spec = None @@ -1454,9 +1935,11 @@ def gmm_v2( partial_sum_idx = sum(1 for x in leaves if x is not None) input_output_aliases = {partial_sum_idx: 0} - return pl.pallas_call( - functools.partial(kernel_main, cfgs=cfgs), - out_shape=out_init, + out_spec = pl.BlockSpec(memory_space=pltpu.HBM) + out = pl.pallas_call( + functools.partial(kernel_main_pre_act if cfgs.return_pre_act else kernel_main, cfgs=cfgs), + # return_pre_act: (act, gate, up), all of act's shape. + out_shape=(out_init,) * 3 if cfgs.return_pre_act else out_init, grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=2, in_specs=[ @@ -1471,7 +1954,7 @@ def gmm_v2( ), partial_sum_spec, ], - out_specs=pl.BlockSpec(memory_space=pltpu.HBM), + out_specs=(out_spec,) * 3 if cfgs.return_pre_act else out_spec, scratch_shapes=scratch_shapes, # pyrefly: ignore[bad-argument-type] ), compiler_params=pltpu.CompilerParams( @@ -1482,4 +1965,9 @@ def gmm_v2( cost_estimate=get_cost_estimate(cfgs), metadata=get_metadata(cfgs), # pyrefly: ignore[bad-argument-type] input_output_aliases=input_output_aliases, - )(group_sizes, group_offset, lhs_in, rhs_weights, partial_sum)[:, : cfgs.out_size_n] + )(group_sizes, group_offset, lhs_in, rhs_weights, partial_sum) + if cfgs.return_pre_act: + return tuple(o[:, : cfgs.out_size_n] for o in out) + if cfgs.out_is_3d: + return out + return out[:, : cfgs.out_size_n] diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py index a3238184b1..c241289443 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py @@ -27,6 +27,7 @@ import jax.numpy as jnp from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_gmm_kernel as gmm_v2 +from maxtext.kernels.megablox import relayout @jax.tree_util.register_dataclass @@ -178,6 +179,29 @@ def within_vmem_limit(tile_m, tile_k, tile_n): return gmm_v2.TileSizes(tile_m=tile_m, tile_k=tile_k, tile_n=tile_n) +def _operand_size(x: jax.Array, name: str) -> int: + """Returns the flattened minor size of a 2D `[M, S]` or 3D `[M, S//128, 128]` operand.""" + if x.ndim == 2: + return x.shape[1] + num_lanes = pltpu.get_tpu_info().num_lanes + if x.ndim != 3 or x.shape[-1] != num_lanes: + raise ValueError(f"tgmm {name} must be [M, S] or [M, S // {num_lanes}, {num_lanes}]; got {x.shape}.") + return x.shape[1] * num_lanes + + +def _check_3d_tile(tile_d0: int, full_d0: int, dtype) -> None: + if relayout.needs_window(tile_d0, full_d0): + relayout.check_window(tile_d0, full_d0, dtype, store=False) + else: + relayout.check_tile_d0(tile_d0, full_d0=full_d0, dtype=dtype) + + +def _window(tile: int, full: int) -> bool: + """3D operand tile that is not a legal D0 block (see relayout.needs_window).""" + num_lanes = pltpu.get_tpu_info().num_lanes + return relayout.needs_window(tile // num_lanes, full // num_lanes) + + def make_tgmm_configs( lhs: jax.Array, # [m, k] rhs: jax.Array, # [m, n] @@ -197,8 +221,9 @@ def make_tgmm_configs( assert lhs.shape[0] == rhs.shape[0], ( f"lhs and rhs m-dim mismatch: {lhs.shape[0]}!={rhs.shape[0]} {lhs.shape}" f" vs {rhs.shape}" ) - size_m, size_k = lhs.shape - _, size_n = rhs.shape + size_m = lhs.shape[0] + size_k = _operand_size(lhs, "lhs") + size_n = _operand_size(rhs, "rhs") if rhs_scale is not None: # rhs_scale.shape[0] is the number of quant blocks along the m (reduction) # dimension. tgmm_v2 only implements per-N (per-output-channel) scaling, @@ -242,11 +267,13 @@ def make_tgmm_configs( quant_block_size=rhs_quant_block_size_m, dtype=rhs.dtype, has_scale=(rhs_scale is not None), + is_3d=rhs.ndim == 3, ) lhs_cfgs = gmm_v2.InputConfigs( quant_dtype=None, quant_block_size=-1, dtype=lhs.dtype, + is_3d=lhs.ndim == 3, ) fuse_act = None # fuse_act has to be None in tgmm. @@ -266,6 +293,18 @@ def make_tgmm_configs( partial_sum is not None, ) + num_lanes = pltpu.get_tpu_info().num_lanes + if lhs_cfgs.is_3d: + if tiles.tile_k % num_lanes or size_k % tiles.tile_k: + raise ValueError(f"3D tgmm lhs requires tile_k % 128 == 0 and size_k % tile_k == 0; got {tiles=}, {size_k=}.") + _check_3d_tile(tiles.tile_k // num_lanes, size_k // num_lanes, lhs.dtype) + if rhs_cfgs.is_3d: + if tiles.tile_n % num_lanes or size_n % tiles.tile_n: + raise ValueError(f"3D tgmm rhs requires tile_n % 128 == 0 and size_n % tile_n == 0; got {tiles=}, {size_n=}.") + _check_3d_tile(tiles.tile_n // num_lanes, size_n // num_lanes, rhs.dtype) + if (lhs_cfgs.is_3d or rhs_cfgs.is_3d) and tiles.tile_m % size_lhs_sublane: + raise ValueError(f"3D tgmm requires tile_m % sublane == 0; got {tiles=}, {size_lhs_sublane=}.") + return gmm_v2.GmmConfigs( dims=dims, tiles=tiles, @@ -279,6 +318,7 @@ def make_tgmm_configs( # tgmm accumulation happens at the 'pallas_call' level. zero_init=False, fuse_act=fuse_act, + transpose_rhs=False, ) @@ -313,8 +353,49 @@ def tgmm_inner_kernel( # NB: grid=(num_n, num_k, num_gm) tiled_rhs_scale_ref = tiled_rhs_ref.scale - tiled_lhs_ref = tiled_lhs_ref.reshape(-1, tiled_lhs_ref.shape[-1]) - tiled_rhs_ref = tiled_rhs_ref.value.reshape(-1, tiled_rhs_ref.value.shape[-1]) + num_lanes = pltpu.get_tpu_info().num_lanes + tile_m = cfgs.tiles.tile_m + if cfgs.lhs_cfgs.is_3d: + # [tile_m // sublane, sublane, tile_d0, 128] -> [tile_m, tile_k] + lhs_3d_ref = tiled_lhs_ref + if _window(cfgs.tiles.tile_k, cfgs.dims.size_k): + # Full-D0 slab block; relayout the K window k_id (grid axis 1) from VMEM. + def load_lhs(): + return relayout.load_3d_window_as_2d(lhs_3d_ref, tile_m, cfgs.tiles.tile_k // num_lanes, pl.program_id(1)) + + else: + + def load_lhs(): + return relayout.load_3d_as_2d(lhs_3d_ref, tile_m, cfgs.tiles.tile_k // num_lanes) + + lhs_shape = (tile_m, cfgs.tiles.tile_k) + else: + tiled_lhs_ref = tiled_lhs_ref.reshape(-1, tiled_lhs_ref.shape[-1]) + + def load_lhs(): + return tiled_lhs_ref[...] + + lhs_shape = tiled_lhs_ref.shape + if cfgs.rhs_cfgs.is_3d: + rhs_3d_ref = tiled_rhs_ref.value + if _window(cfgs.tiles.tile_n, cfgs.dims.size_n): + # Full-D0 slab block; relayout the N window n_id (grid axis 0) from VMEM. + def load_rhs(): + return relayout.load_3d_window_as_2d(rhs_3d_ref, tile_m, cfgs.tiles.tile_n // num_lanes, pl.program_id(0)) + + else: + + def load_rhs(): + return relayout.load_3d_as_2d(rhs_3d_ref, tile_m, cfgs.tiles.tile_n // num_lanes) + + rhs_shape = (tile_m, cfgs.tiles.tile_n) + else: + tiled_rhs_ref = tiled_rhs_ref.value.reshape(-1, tiled_rhs_ref.value.shape[-1]) + + def load_rhs(): + return tiled_rhs_ref[...] # pyrefly: ignore[bad-index] + + rhs_shape = tiled_rhs_ref.shape # pyrefly: ignore[missing-attribute] gm_id = pl.program_id(2) def _matmul(is_new_group: bool, is_group_changing: bool): @@ -327,15 +408,15 @@ def _matmul(is_new_group: bool, is_group_changing: bool): m_offset = m_start - m_start % cfgs.dims.size_lhs_sublane m_start_local = m_start - m_offset m_end_local = m_end - m_offset - lhs_iota = lax.broadcasted_iota(jnp.int32, tiled_lhs_ref.shape, 0) + lhs_iota = lax.broadcasted_iota(jnp.int32, lhs_shape, 0) lhs_mask = jnp.logical_and(m_start_local <= lhs_iota, lhs_iota < m_end_local) - lhs_masked = jnp.where(lhs_mask, tiled_lhs_ref[...], 0) + lhs_masked = jnp.where(lhs_mask, load_lhs(), 0) # If there are no NaNs, masking both lhs and rhs shouldn't be necessary. # But without masking both, we sometimes see the result contain NaNs so we # decide to mask both to be safe. - rhs_iota = lax.broadcasted_iota(jnp.int32, tiled_rhs_ref.shape, 0) + rhs_iota = lax.broadcasted_iota(jnp.int32, rhs_shape, 0) rhs_mask = jnp.logical_and(m_start_local <= rhs_iota, rhs_iota < m_end_local) - rhs_masked = jnp.where(rhs_mask, tiled_rhs_ref[...], 0) + rhs_masked = jnp.where(rhs_mask, load_rhs(), 0) acc = jax.lax.dot_general( lhs_masked, @@ -414,21 +495,31 @@ def __init__(self, metadata_ref: gmm_v2.MetadataRef, cfgs: gmm_v2.GmmConfigs): self.cfgs = cfgs def lhs_index_map(self, n_id: jax.Array, k_id: jax.Array, gm_id: jax.Array): + """Block index of the lhs tile (group rows, K block; 3D lhs blocks K along D0).""" m_start = self.metadata_ref.gm_id_to_m_offset[gm_id] m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1] row_start = m_start // self.cfgs.dims.size_lhs_sublane row_end = pl.cdiv(m_end, self.cfgs.dims.size_lhs_sublane) row_size = row_end - row_start + if self.cfgs.lhs_cfgs.is_3d: + # Block is [rows, sublane, tile_d0, 128]; K is blocked along D0 (a window + # block covers the full D0). + k_blk = 0 if _window(self.cfgs.tiles.tile_k, self.cfgs.dims.size_k) else k_id + return (pl.ds(row_start, row_size), 0, k_blk, 0) return (pl.ds(row_start, row_size), 0, k_id) def rhs_index_map(self, n_id: jax.Array, k_id: jax.Array, gm_id: jax.Array): + """Block index of the rhs tile (group rows, N block; 3D rhs blocks N along D0).""" m_start = self.metadata_ref.gm_id_to_m_offset[gm_id] m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1] row_start = m_start // self.cfgs.dims.size_lhs_sublane row_end = pl.cdiv(m_end, self.cfgs.dims.size_lhs_sublane) row_size = row_end - row_start + if self.cfgs.rhs_cfgs.is_3d: + n_blk = 0 if _window(self.cfgs.tiles.tile_n, self.cfgs.dims.size_n) else n_id + return (pl.ds(row_start, row_size), 0, n_blk, 0) return (pl.ds(row_start, row_size), 0, n_id) def rhs_scale_index_map(self, n_id: jax.Array, k_id: jax.Array, gm_id: jax.Array): @@ -449,14 +540,20 @@ def generate_tgmm_block_specs( # after this reshape has size tile_m // size_lhs_sublane — i.e., the number of # "sublane-rows" in a tile. bounded_slice_gm = pl.BoundedSlice(cfgs.tiles.tile_m // cfgs.dims.size_lhs_sublane) - lhs_block_spec = pl.BlockSpec( - (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_k), - index_map.lhs_index_map, - ) - rhs_block_spec = pl.BlockSpec( - (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_n), - index_map.rhs_index_map, - ) + num_lanes = pltpu.get_tpu_info().num_lanes + sublane = cfgs.dims.size_lhs_sublane + if cfgs.lhs_cfgs.is_3d: + lhs_tile = cfgs.dims.size_k if _window(cfgs.tiles.tile_k, cfgs.dims.size_k) else cfgs.tiles.tile_k + lhs_block_shape = (bounded_slice_gm, sublane, lhs_tile // num_lanes, num_lanes) + else: + lhs_block_shape = (bounded_slice_gm, sublane, cfgs.tiles.tile_k) + lhs_block_spec = pl.BlockSpec(lhs_block_shape, index_map.lhs_index_map) + if cfgs.rhs_cfgs.is_3d: + rhs_tile = cfgs.dims.size_n if _window(cfgs.tiles.tile_n, cfgs.dims.size_n) else cfgs.tiles.tile_n + rhs_block_shape = (bounded_slice_gm, sublane, rhs_tile // num_lanes, num_lanes) + else: + rhs_block_shape = (bounded_slice_gm, sublane, cfgs.tiles.tile_n) + rhs_block_spec = pl.BlockSpec(rhs_block_shape, index_map.rhs_index_map) rhs_scale_block_spec = None if cfgs.rhs_cfgs.has_scale: rhs_scale_block_spec = pl.BlockSpec( @@ -568,13 +665,15 @@ def tgmm_kernel_main( semaphore_ref: DMA semaphore for the zeroing copies. cfgs: GmmConfigs object containing kernel configurations. """ - num_groups_to_zero = zero_out_start( - lhs_group_sizes_ref, - group_offset_ref, - out_ref, - zero_ref, - semaphore_ref, - ) + num_groups_to_zero = 0 + if not cfgs.has_partial_sum: + num_groups_to_zero = zero_out_start( + lhs_group_sizes_ref, + group_offset_ref, + out_ref, + zero_ref, + semaphore_ref, + ) num_k = pl.cdiv(cfgs.dims.size_k, cfgs.tiles.tile_k) num_n = pl.cdiv(cfgs.dims.size_n, cfgs.tiles.tile_n) @@ -592,9 +691,10 @@ def tgmm_kernel_main( in_specs=in_specs, out_specs=out_specs, ) - lhs_in = lhs_ref.reshape(-1, cfgs.dims.size_lhs_sublane, lhs_ref.shape[-1]) + # [M, S] -> [M // sublane, sublane, S]; 3D [M, D0, 128] -> [M // sublane, sublane, D0, 128]. + lhs_in = lhs_ref.reshape(-1, cfgs.dims.size_lhs_sublane, *lhs_ref.shape[1:]) rhs_value = rhs_ref.value - rhs_in = rhs_value.reshape(-1, cfgs.dims.size_lhs_sublane, rhs_value.shape[-1]) + rhs_in = rhs_value.reshape(-1, cfgs.dims.size_lhs_sublane, *rhs_value.shape[1:]) rhs_operand = OperandRef(value=rhs_in, scale=rhs_ref.scale) ps_in = None if cfgs.has_partial_sum: @@ -602,11 +702,12 @@ def tgmm_kernel_main( scratches = [acc_ref, metadata_ref] pipeline_fn(lhs_in, rhs_operand, ps_in, out_ref, scratches=scratches) - zero_out_end( - num_groups_to_zero, - out_ref, - semaphore_ref, - ) + if not cfgs.has_partial_sum: + zero_out_end( + num_groups_to_zero, + out_ref, + semaphore_ref, + ) def validate_tgmm_inputs( @@ -648,8 +749,8 @@ def validate_tgmm_inputs( ], ) def tgmm_v2( - lhs: jax.Array, # [size_m, size_k] - rhs: jax.Array, # [size_m, size_n] + lhs: jax.Array, # [size_m, size_k] or [size_m, size_k // 128, 128] + rhs: jax.Array, # [size_m, size_n] or [size_m, size_n // 128, 128] group_sizes: jax.Array, num_actual_groups: int, rhs_scale: jax.Array | None = None, # [1, 1, size_n] (per-N scale) @@ -756,6 +857,10 @@ def tgmm_v2( hbm_spec = pl.BlockSpec(memory_space=pltpu.HBM) partial_sum_spec = None if partial_sum is not None: + pad_k = aligned_k - dims.size_k + pad_n = aligned_n - dims.size_n + if pad_k > 0 or pad_n > 0: + partial_sum = jnp.pad(partial_sum, ((0, 0), (0, pad_k), (0, pad_n))) partial_sum_spec = hbm_spec in_specs = [ hbm_spec, # lhs @@ -793,8 +898,4 @@ def tgmm_v2( input_output_aliases=input_output_aliases, )(group_sizes, group_offset, lhs, rhs, partial_sum)[:, : dims.size_k, : dims.size_n] - if partial_sum is not None: - local_group_sizes = lax.dynamic_slice(group_sizes, (group_offset[0],), (num_actual_groups,)) - empty_mask = (local_group_sizes == 0).reshape(num_actual_groups, 1, 1) - return jnp.where(empty_mask, partial_sum, raw_out) return raw_out diff --git a/src/maxtext/kernels/megablox/relayout.py b/src/maxtext/kernels/megablox/relayout.py new file mode 100644 index 0000000000..ecd06517f0 --- /dev/null +++ b/src/maxtext/kernels/megablox/relayout.py @@ -0,0 +1,237 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""In-VMEM relayout between [tokens, D0, 128] and [tokens, D0 * 128]. + +Ported from Lineage's gmmv2/relayout.py (bf16 only) and +extended to 8-bit dtypes (fp8 / int8) and 32-bit dtypes (e.g. f32 partial sums). + +gmm_v2 feeds the MXU tiles whose sublanes index tokens and whose lanes index +the contracting dim. Tokens stored as ``[tokens, D0, 128]`` (one contiguous +slab per token) have the sublane axis of each VMEM tile indexing ``D0`` +instead. The helpers below convert between the layouts inside VMEM/vregs using +32-bit sublane-strided loads/stores (the only strided access Mosaic supports) +plus a few VALUE bit ops to (un)pack the sub-32-bit values. + +For a VMEM tile ``[tm, d0, 128]`` of a dtype with ``p`` values per 32-bit word +(``p = 2`` for bf16, ``p = 4`` for 8-bit): + +* Viewed as uint32 it is ``[tm, d0 // p, 128]``: rows ``p*i .. p*i + p - 1`` of + a token share one word (row ``p*i + j`` in bits ``[j*b, (j+1)*b)``). +* Viewed flat it is ``[tm * d0 // p, 128]`` and token ``t`` owns rows + ``[t * d0 // p, (t + 1) * d0 // p)``. +* A sublane-strided load with stride ``p * d0 // p`` starting at row + ``j * d0 // p + w`` returns word ``w`` of tokens ``j, j + p, j + 2p, ...``. + ``p`` such loads followed by a ``p x p`` sub-word transpose give ``p`` + ``[tm, 128]`` columns (``d0`` rows ``p*w .. p*w + p - 1``) with tokens in + sublanes, i.e. in the native packed layout of a ``[tm, 128]`` value. +""" + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +LANES = 128 + + +def packing(dtype) -> int: + bits = jax.dtypes.itemsize_bits(jnp.dtype(dtype)) + if bits not in (8, 16, 32): + raise ValueError(f"3D relayout supports 8/16/32-bit dtypes, got {dtype}.") + return 32 // bits + + +def check_tile_d0(tile_d0: int, full_d0: int | None = None, dtype=jnp.bfloat16) -> None: + """Raises if ``tile_d0`` cannot use the strided-load relayout. + + Args: + tile_d0: second-minor block dim (``tile_k // 128`` or ``tile_n // 128``). + full_d0: the full second-minor array dim, if known. A block that covers the + whole dim is always accepted (Pallas' "block == full dim" rule). + dtype: element dtype of the 3D operand. + """ + p = packing(dtype) + if tile_d0 % p != 0: + raise ValueError(f"3D relayout packs {p} rows per word; tile_d0 must be a multiple of {p}, got {tile_d0}.") + if tile_d0 % 8 != 0 and tile_d0 != full_d0: + raise ValueError( + "3D relayout requires tile_d0 % 8 == 0 (tile_k / tile_n a multiple" + f" of {8 * LANES}) or tile_d0 == full D0 ({full_d0}); got" + f" tile_d0={tile_d0}." + ) + + +def needs_window(tile_d0: int, full_d0: int) -> bool: + """Whether a ``tile_d0`` block of a ``full_d0`` 3D operand must use a full-D0 window block. + + Pallas TPU blocks need a second-minor dim that is a multiple of 8 or the full + dim, and sub-tile DMA windows are not supported. Such tiles (e.g. + 3584 / 1792 of 7168: tile_d0 = 28 / 14) instead DMA the whole ``[.., full_d0, + 128]`` slab and relayout the ``[d0_start, d0_start + tile_d0)`` window from VMEM. + """ + return tile_d0 != full_d0 and tile_d0 % 8 != 0 + + +def check_window(tile_d0: int, full_d0: int, dtype, *, store: bool) -> None: + """Raises if a full-D0 window block cannot serve ``tile_d0`` tiles.""" + p = packing(dtype) + if full_d0 % tile_d0 != 0: + raise ValueError(f"3D window tiles need tile_d0 ({tile_d0}) to divide D0 ({full_d0}).") + if full_d0 % p != 0: + raise ValueError(f"3D D0 ({full_d0}) must be a multiple of the packing {p}.") + # Loads handle windows that start / end inside a packed word; stores do not. + if (store or p == 2) and tile_d0 % p != 0: + raise ValueError(f"3D window {'stores' if store else 'loads'} need tile_d0 % {p} == 0, got {tile_d0}.") + + +def _flat_u32_view(ref): + """Returns ``ref`` (``[..., d0, 128]``) as uint32 ``[-1, 128]`` and the u32 rows per token.""" + full_d0 = ref.shape[-2] + p = packing(ref.dtype) + assert full_d0 % p == 0, (full_d0, p) + rows_per_tok = full_d0 // p + ref_u32 = ref.bitcast(jnp.uint32) # [..., d0 // p, 128] + return ref_u32.reshape(-1, LANES), rows_per_tok + + +def _u32(v): + return jnp.uint32(v) + + +def _transpose4x4_bytes(a0, a1, a2, a3): + """4x4 byte transpose: out_b byte j = a_j byte b.""" + lo16, hi16 = _u32(0x0000FFFF), _u32(0xFFFF0000) + # Stage 1 (16-bit halves): pair token phases (0, 2) and (1, 3). + b0 = (a0 & lo16) | (a2 << 16) # [a0.b0, a0.b1, a2.b0, a2.b1] + b2 = (a0 >> 16) | (a2 & hi16) # [a0.b2, a0.b3, a2.b2, a2.b3] + b1 = (a1 & lo16) | (a3 << 16) # [a1.b0, a1.b1, a3.b0, a3.b1] + b3 = (a1 >> 16) | (a3 & hi16) # [a1.b2, a1.b3, a3.b2, a3.b3] + # Stage 2 (bytes). + ev, od = _u32(0x00FF00FF), _u32(0xFF00FF00) + c0 = (b0 & ev) | ((b1 & ev) << 8) # [a0.b0, a1.b0, a2.b0, a3.b0] + c1 = ((b0 >> 8) & ev) | (b1 & od) # [a0.b1, a1.b1, a2.b1, a3.b1] + c2 = (b2 & ev) | ((b3 & ev) << 8) + c3 = ((b2 >> 8) & ev) | (b3 & od) + return c0, c1, c2, c3 + + +def load_3d_as_2d(ref, tm: int, tile_d0: int, d0_start: int = 0) -> jax.Array: + """Loads rows ``[d0_start, d0_start + tile_d0)`` of VMEM ref ``[..., D0, 128]`` as ``[tm, tile_d0 * 128]``. + + Args: + ref: VMEM ref whose trailing two dims are ``(D0, 128)`` (bf16 or 8-bit) and + whose leading dims multiply to at least ``tm``. ``D0 == tile_d0`` for a + plain block; a full-D0 window block (``needs_window``) has ``D0 > tile_d0``. + tm: number of leading rows (tokens) to load. Must be a multiple of the + packing (2 for bf16, 4 for 8-bit). + tile_d0: number of D0 rows to load. + d0_start: static first D0 row. For bf16 it must be even; for 8-bit it may + start / end inside a packed word. + + Returns: + Array ``[tm, tile_d0 * 128]`` of ``ref.dtype`` with tokens in rows, + identical to ``ref[:tm, d0_start:d0_start + tile_d0].reshape(tm, -1)``. + """ + dtype = ref.dtype + p = packing(dtype) + assert tm % p == 0, (tm, p) + assert 0 <= d0_start and d0_start + tile_d0 <= ref.shape[-2], (d0_start, tile_d0, ref.shape) + if p == 1: + # 32-bit: one sublane-strided load per D0 row, no unpacking. + x2d, rows_per_tok = ref.reshape(-1, LANES), ref.shape[-2] + cols = [x2d[pl.ds(c, tm, stride=rows_per_tok), :] for c in range(d0_start, d0_start + tile_d0)] + return jnp.concatenate(cols, axis=-1) + x2d, rows_per_tok = _flat_u32_view(ref) + stride = p * rows_per_tok # p tokens, in uint32 rows + cols = [] + if p == 2: + assert d0_start % 2 == 0 and tile_d0 % 2 == 0, (d0_start, tile_d0) + lo_mask, hi_mask = _u32(0xFFFF), _u32(0xFFFF0000) + for c in range(d0_start, d0_start + tile_d0, 2): + even = x2d[pl.ds(c // 2, tm // 2, stride=stride), :] + odd = x2d[pl.ds(rows_per_tok + c // 2, tm // 2, stride=stride), :] + lo = (even & lo_mask) | (odd << 16) # d0 == c + hi = (even >> 16) | (odd & hi_mask) # d0 == c + 1 + cols.append(pltpu.bitcast(lo, dtype)) # [tm, 128] + cols.append(pltpu.bitcast(hi, dtype)) + else: + d0_end = d0_start + tile_d0 + for w in range(d0_start // 4, pl.cdiv(d0_end, 4)): + a = [x2d[pl.ds(j * rows_per_tok + w, tm // 4, stride=stride), :] for j in range(4)] + for b, c in enumerate(_transpose4x4_bytes(*a)): + if d0_start <= 4 * w + b < d0_end: + cols.append(pltpu.bitcast(c, dtype)) # [tm, 128], d0 == 4 * w + b + return jnp.concatenate(cols, axis=-1) + + +def store_2d_as_3d(ref, x2d: jax.Array, tm: int, tile_d0: int, d0_start: int = 0) -> None: + """Stores ``x2d`` ``[tm, tile_d0 * 128]`` into rows ``[d0_start, d0_start + tile_d0)`` of ``[..., D0, 128]``. + + Inverse of :func:`load_3d_as_2d`; writes the first ``tm`` tokens of ``ref`` + and leaves its other D0 rows untouched. 8/16/32-bit dtypes are supported + (``x2d.dtype == ref.dtype``); ``d0_start`` and ``tile_d0`` must be multiples + of the packing. + """ + dtype = ref.dtype + assert x2d.dtype == dtype, (x2d.dtype, dtype) + assert x2d.shape == (tm, tile_d0 * LANES), (x2d.shape, tm, tile_d0) + p = packing(dtype) + assert tm % p == 0, (tm, p) + assert d0_start % p == 0 and tile_d0 % p == 0, (d0_start, tile_d0, p) + assert d0_start + tile_d0 <= ref.shape[-2], (d0_start, tile_d0, ref.shape) + if p == 1: + o2d, rows_per_tok = ref.reshape(-1, LANES), ref.shape[-2] + for c in range(tile_d0): + o2d[pl.ds(d0_start + c, tm, stride=rows_per_tok), :] = x2d[:, c * LANES : (c + 1) * LANES] + return + o2d, rows_per_tok = _flat_u32_view(ref) + stride = p * rows_per_tok + if p == 2: + lo_mask, hi_mask = _u32(0xFFFF), _u32(0xFFFF0000) + for c in range(0, tile_d0, 2): + lo = pltpu.bitcast(x2d[:, c * LANES : (c + 1) * LANES], jnp.uint32) + hi = pltpu.bitcast(x2d[:, (c + 1) * LANES : (c + 2) * LANES], jnp.uint32) + even = (lo & lo_mask) | (hi << 16) # even tokens: (c, c+1) + odd = (lo >> 16) | (hi & hi_mask) # odd tokens + r = (d0_start + c) // 2 + o2d[pl.ds(r, tm // 2, stride=stride), :] = even + o2d[pl.ds(rows_per_tok + r, tm // 2, stride=stride), :] = odd + else: + for w in range(tile_d0 // 4): + c = [pltpu.bitcast(x2d[:, (4 * w + b) * LANES : (4 * w + b + 1) * LANES], jnp.uint32) for b in range(4)] + # The 4x4 byte transpose is an involution: word j byte b <- column b byte j. + a = _transpose4x4_bytes(*c) + r = d0_start // 4 + w + for j in range(4): + o2d[pl.ds(j * rows_per_tok + r, tm // 4, stride=stride), :] = a[j] + + +def load_3d_window_as_2d(ref, tm: int, tile_d0: int, window_id) -> jax.Array: + """``load_3d_as_2d`` of window ``window_id`` (traced) of a full-D0 block: a switch over static starts.""" + num = ref.shape[-2] // tile_d0 + if num == 1: + return load_3d_as_2d(ref, tm, tile_d0) + branches = [lambda i=i: load_3d_as_2d(ref, tm, tile_d0, i * tile_d0) for i in range(num)] + return jax.lax.switch(window_id, branches) + + +def store_2d_as_3d_window(ref, x2d: jax.Array, tm: int, tile_d0: int, window_id) -> None: + """``store_2d_as_3d`` into window ``window_id`` (traced) of a full-D0 block.""" + num = ref.shape[-2] // tile_d0 + if num == 1: + store_2d_as_3d(ref, x2d, tm, tile_d0) + return + branches = [lambda i=i: store_2d_as_3d(ref, x2d, tm, tile_d0, i * tile_d0) for i in range(num)] + jax.lax.switch(window_id, branches) diff --git a/src/maxtext/kernels/ragged/ragged_gather.py b/src/maxtext/kernels/ragged/ragged_gather.py index b90bbcb411..d662c21b7d 100644 --- a/src/maxtext/kernels/ragged/ragged_gather.py +++ b/src/maxtext/kernels/ragged/ragged_gather.py @@ -445,7 +445,7 @@ def ragged_gather( core_axis_name="core", subcore_axis_name="subcore", ) - return pl.kernel( # pytype: disable=wrong-keyword-args + return pl.kernel( functools.partial( main_kernel, core_axis_name=vector_mesh.core_axis_name, diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index 1392149ac5..8256f85cfd 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -544,7 +544,7 @@ def ragged_gather_reduce( ) # Each output row from `main_kernel` will be of type float32, and then # casted to the input dtype when doing the filter operation. - out = pl.kernel( # pytype: disable=wrong-keyword-args + out = pl.kernel( functools.partial( main_kernel, core_axis_name=vector_mesh.core_axis_name, @@ -559,7 +559,6 @@ def ragged_gather_reduce( compiler_params=pltpu.CompilerParams( use_tc_tiling_on_sc=True, disable_bounds_checks=True, - needs_layout_passes=False, ), cost_estimate=get_cost_estimate( padded_input_size=padded_input_size, diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py index 571200ffb7..0b8b3fa101 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py @@ -263,15 +263,41 @@ def _calculate_col_chunk_size(col_size: int, num_simd_lanes: int) -> int: return 128 +def argsort_unique_pairs(keys: jax.Array, axis: int = -1) -> jax.Array: + """``jnp.argsort(keys, axis=axis, stable=True)``, computed as an unstable sort of the pairs ``(key, index)``. + + The pairs ``(keys[i], i)`` along ``axis`` are all distinct, so sorting them lexicographically has exactly one + result, the stable argsort's permutation: equal keys come out in increasing index order. Marked unstable, the + sort keeps two operands (keys, indices) and a two-key comparator. ``jnp.argsort`` emits a stable + ``sort_key_val(keys, iota)``, which XLA lowers in the forward pass with a THIRD operand (a second copy of the iota + as the tie-breaker). Measured on TPU v7x for an s32[655360] argsort: 0.61-0.63 ms for the three-operand sort, + 0.42 ms for the two-operand one (moe_ragged_sort_unique_pair_argsort). The index dtype and sharding are the ones + ``jnp.argsort`` uses. + """ + dim = axis % keys.ndim + # jnp.argsort's index dtype: the default integer type (int32 unless x64 is enabled). + idx_dtype = jax.dtypes.canonicalize_dtype(jnp.int64) + iota = jax.lax.broadcasted_iota(idx_dtype, keys.shape, dim, out_sharding=jax.typeof(keys).sharding) + _, indices = jax.lax.sort((keys, iota), dimension=dim, is_stable=False, num_keys=2) + return indices + + def _preprocess( valid_rows_mask: jax.Array, reduce_group_size: int, num_row_partitions: int, num_simd_lanes: int, row_chunk_size: int, + packed: bool = False, + unique_pair_argsort: bool = False, ) -> tuple[jax.Array, jax.Array, jax.Array]: """Sorts valid source rows to the front of each row partition. + With ``packed`` (``moe_routing_index_math="count"``) the stable partition is one single-operand sort of + packed uint32 keys instead of a stable argsort; the permutation is the same. With ``unique_pair_argsort`` + (``moe_ragged_sort_unique_pair_argsort``) it is ``argsort_unique_pairs``, the same permutation again. + ``packed`` takes precedence. + Returns: sorted_by_validity: original row index of each slot after the stable sort, flattened across partitions and padded to ``row_chunk_size``. @@ -282,9 +308,22 @@ def _preprocess( row_partition_size = valid_rows_mask.shape[0] // num_row_partitions valid_rows_mask_2d = valid_rows_mask.reshape(num_row_partitions, -1) - # Stable sort of a boolean key is a stable partition: valid rows keep their - # relative order and move ahead of the invalid ones. - sorted_by_validity = jnp.argsort(~valid_rows_mask_2d, descending=False, stable=True, axis=-1) + if packed: + # Key (invalid << idx_bits) | index: valid rows sort first, each group in index order. The keys are all + # distinct, so an unstable sort has exactly one result and it is the stable partition below. It sorts ONE + # uint32 operand, where the stable boolean argsort sorts the key with an iota tie-break (XLA's stable-sort + # expansion adds a third operand in the forward pass). + idx_bits = max(1, (row_partition_size - 1).bit_length()) + key = (valid_rows_mask_2d == 0).astype(jnp.uint32) << jnp.uint32(idx_bits) + key = key | jnp.arange(row_partition_size, dtype=jnp.uint32)[None, :] + key = jax.lax.sort(key, dimension=1, is_stable=False) + sorted_by_validity = (key & jnp.uint32((1 << idx_bits) - 1)).astype(jnp.int32) + elif unique_pair_argsort: + sorted_by_validity = argsort_unique_pairs(~valid_rows_mask_2d, axis=-1) + else: + # Stable sort of a boolean key is a stable partition: valid rows keep their + # relative order and move ahead of the invalid ones. + sorted_by_validity = jnp.argsort(~valid_rows_mask_2d, descending=False, stable=True, axis=-1) sorted_by_validity += jnp.arange(num_row_partitions)[:, None] * row_partition_size pad_to = _align_to(row_partition_size, row_chunk_size) @@ -636,6 +675,8 @@ def col_loop(col_compute_offset): "flops_override", "bytes_accessed_override", "use_single_sparsecore", + "packed_preprocess", + "unique_pair_argsort", ), ) def ragged_gather_reduce( @@ -648,6 +689,8 @@ def ragged_gather_reduce( flops_override: int = -1, bytes_accessed_override: int = -1, use_single_sparsecore: bool = False, + packed_preprocess: bool = False, + unique_pair_argsort: bool = False, ) -> jax.Array: """Gathers ``x`` by ``indices``, weights and masks, then reduces by group. @@ -657,6 +700,11 @@ def ragged_gather_reduce( topk_weights: 1-D per-row weights, ``(input_size,)``. valid_rows_mask: 1-D bool mask of valid gathered rows, ``(input_size,)``. reduce_group_size: number of consecutive rows summed into one output row. + packed_preprocess: sort the valid rows to the front with one packed uint32 sort (``_preprocess(packed=True)``, + ``moe_routing_index_math="count"``). Same permutation, so the same output. Only the SparseCore path sorts; + the reference fallback (any non-TPU target) ignores it. + unique_pair_argsort: the same validity sort as a two-operand sort of unique (key, index) pairs + (``argsort_unique_pairs``, ``moe_ragged_sort_unique_pair_argsort``). Same permutation; SparseCore path only. Returns: Reduced output, ``(input_size // reduce_group_size, hidden_size)``. @@ -720,6 +768,8 @@ def ragged_gather_reduce( num_row_partitions, num_simd_lanes, row_chunk_size, + packed=packed_preprocess, + unique_pair_argsort=unique_pair_argsort, ) # Step 4: Launch the SparseCore kernel. diff --git a/src/maxtext/kernels/ragged/ragged_gather_tc.py b/src/maxtext/kernels/ragged/ragged_gather_tc.py new file mode 100644 index 0000000000..692b70eb71 --- /dev/null +++ b/src/maxtext/kernels/ragged/ragged_gather_tc.py @@ -0,0 +1,878 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TensorCore (Pallas, DMA-driven) ragged gather / gather-reduce kernels. + +Ported from Lineage's ops.py (ragged_gather_tc, +ragged_gather_reduce_tc and its metadata). Tokens are laid out as +(num_tokens, d0, d1) with d1 = 128 so that each token row is one contiguous DMA. +""" + +import dataclasses +import functools +import logging +import math + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + + +# --- Pallas Kernels for Ragged Gather on TensorCore --- + +_RAGGED_GATHER_TC_NUM_BUFFERS = 2 + + +def ragged_gather_tc( + x: jax.Array, + indices: jax.Array, + start: jax.Array | int, + num_tokens: jax.Array | int, + max_out_tokens: int, + block_size: int = 2048, +) -> jax.Array: + """Executes a ragged gather on TensorCore using a Pallas kernel. + + Gathers `out[i] = x[indices[start + i]]` for `i < num_tokens`. Output rows + `[num_tokens, max_out_tokens)` are left untouched (garbage from the + uninitialized output buffer). + + The caller must guarantee `0 <= num_tokens <= max_out_tokens` and + `0 <= start` with `start + num_tokens <= num_indices`; these are not checked + at runtime and violating them is undefined behavior. + + Args: + x: Input array of shape (num_input_tokens, d0, d1). + indices: Indices into `x` of shape (num_indices,). + start: Offset into `indices` of the first index to gather. May be a traced + scalar; the kernel does not recompile per value. + num_tokens: Number of indices to gather. May be a traced scalar; the kernel + does not recompile per value. + max_out_tokens: Length of the output buffer along axis 0. + block_size: Number of output rows gathered per grid step, i.e. the number of + concurrent row DMAs in flight per buffer. Clamped to `max_out_tokens`. + + Returns: + Gathered array in an output buffer of shape (max_out_tokens, d0, d1). + """ + if max_out_tokens <= 0: + raise ValueError(f"max_out_tokens must be positive, got {max_out_tokens}.") + if isinstance(num_tokens, int) and not 0 <= num_tokens <= max_out_tokens: + raise ValueError(f"num_tokens must be in [0, {max_out_tokens}], got {num_tokens}.") + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}.") + num_indices = indices.shape[0] + _, dim0, dim1 = x.shape + if num_indices == 0 or (isinstance(num_tokens, int) and num_tokens == 0): + return jnp.empty((max_out_tokens, dim0, dim1), dtype=x.dtype) + + block_size = min(block_size, max_out_tokens) + num_buffers = _RAGGED_GATHER_TC_NUM_BUFFERS + num_full_blocks = max_out_tokens // block_size + rem = max_out_tokens % block_size + num_blocks = num_full_blocks + (1 if rem > 0 else 0) + + # `indices` may exceed SMEM capacity, so each block stages only its window of + # indices in SMEM. The window starts at an offset aligned to the HBM tile + # (128 x i32) so the dynamic DMA slice needs no retiling, and has room for the + # misalignment. It is clamped to `max_win_start` to stay in bounds, which + # requires the (padded) indices length to be aligned and at least one window + # long. + idx_align = 128 + idx_window = pl.cdiv(block_size, idx_align) * idx_align + idx_align + idx_pad = max((-num_indices) % idx_align, idx_window - num_indices) + if idx_pad > 0: + indices = jnp.pad(indices, (0, idx_pad)) + max_win_start = indices.shape[0] - idx_window + + def _ragged_gather_kernel( + scalars_smem_ref, + idx_hbm_ref, + x_hbm_ref, + o_hbm_ref, + vmem_ref, + idx_smem_ref, + idx_sem, + data_recv_sem, + data_send_sem, + ): + block_idx = pl.program_id(0) + buf_idx = block_idx % num_buffers + start_idx = scalars_smem_ref[0] + num_valid = scalars_smem_ref[1] + + def _run_step(step_idx, b_idx, action: str): + base = step_idx * block_size + row0 = start_idx + base + win_start = jnp.minimum((row0 // idx_align) * idx_align, max_win_start) + idx_off = row0 - win_start + + def _fetch_indices(): + # Synchronous; the row gathers read the indices as scalars when issued, + # so a single SMEM buffer suffices. + copy = pltpu.make_async_copy( + idx_hbm_ref.at[pl.ds(pl.multiple_of(win_start, idx_align), idx_window)], + idx_smem_ref, + idx_sem, + ) + copy.start() + copy.wait() + + def _chunk(off, size: int, chunk_action: str): + """Issues `chunk_action` for block rows [off, off + size).""" + if chunk_action == "data_start": + step_unroll = math.gcd(size, 128) + + @pl.loop(0, size // step_unroll) + def _(chunk_i): + base_i = off + chunk_i * step_unroll + src_indices = [idx_smem_ref[idx_off + base_i + u] for u in range(step_unroll)] + for u in range(step_unroll): + pltpu.make_async_copy( + x_hbm_ref.at[pl.ds(src_indices[u], 1), :, :], + vmem_ref.at[b_idx, pl.ds(base_i + u, 1), :, :], + data_recv_sem.at[b_idx], + ).start() + + elif chunk_action == "recv_wait": + # Only the byte count matters for the wait; `o_hbm_ref` always has + # at least `block_size` rows, whereas `x` may have fewer. + pltpu.make_async_copy( + o_hbm_ref.at[pl.ds(0, size), :, :], + vmem_ref.at[b_idx, pl.ds(0, size), :, :], + data_recv_sem.at[b_idx], + ).wait() + else: + copy = pltpu.make_async_copy( + vmem_ref.at[b_idx, pl.ds(off, size), :, :], + o_hbm_ref.at[pl.ds(base + off, size), :, :], + data_send_sem.at[b_idx], + ) + if chunk_action == "send_start": + copy.start() + else: + copy.wait() + + def _do_len(length: int): + step_valid = jnp.minimum(jnp.maximum(num_valid - base, 0), length) + + def _over_valid_rows(chunk_action: str): + # Fast path: the whole block is valid, so use one static-size chunk. + @pl.when(step_valid == length) + def _(): + _chunk(0, length, chunk_action) + + # Partially valid block: DMA sizes must be static, so split the + # dynamic row count into power-of-two chunks, one per set bit. + @pl.when(step_valid < length) + def _(): + off = 0 + for k in reversed(range(length.bit_length())): + size = 1 << k + bit = step_valid & size + + def _chunk_if_bit_set(off=off, size=size): + _chunk(off, size, chunk_action) + + pl.when(bit != 0)(_chunk_if_bit_set) + + off = off + bit + + if action == "data_start": + + @pl.when(step_valid > 0) + def _(): + _fetch_indices() + + _over_valid_rows("data_start") + elif action == "copy_out_start": + # All recv waits must finish before any send starts, because the + # per-row gather DMAs share one semaphore and complete out of order. + _over_valid_rows("recv_wait") + _over_valid_rows("send_start") + elif action == "copy_out_wait": + _over_valid_rows("send_wait") + + if rem == 0: + _do_len(block_size) + elif isinstance(step_idx, int): + _do_len(rem if step_idx == num_blocks - 1 else block_size) + else: + jax.lax.cond( + step_idx == num_blocks - 1, + lambda: _do_len(rem), + lambda: _do_len(block_size), + ) + + @pl.when(block_idx == 0) + def _prologue(): + _run_step(0, 0, "data_start") + + _run_step(block_idx, buf_idx, "copy_out_start") + + # The grid covers only blocks overlapping `[0, num_tokens)`. + num_steps = pl.num_programs(0) + next_data_block = block_idx + 1 + + @pl.when(next_data_block < num_steps) + def _queue_next(): + next_data_buf = next_data_block % num_buffers + + @pl.when(block_idx >= 1) + def _wait_prev_out(): + _run_step(block_idx - 1, next_data_buf, "copy_out_wait") + + _run_step(next_data_block, next_data_buf, "data_start") + + @pl.when(block_idx == num_steps - 1) + def _epilogue(): + @pl.when(block_idx >= 1) + def _wait_prev_out(): + prev_block = block_idx - 1 + _run_step(prev_block, prev_block % num_buffers, "copy_out_wait") + + _run_step(block_idx, buf_idx, "copy_out_wait") + + manual_axis_type = getattr(jax.typeof(x), "manual_axis_type", None) + out_shape = jax.ShapeDtypeStruct( + (max_out_tokens, dim0, dim1), + x.dtype, + manual_axis_type=manual_axis_type, + ) + scalars = jnp.stack( + [ + jnp.asarray(start, dtype=jnp.int32), + jnp.asarray(num_tokens, dtype=jnp.int32), + ] + ) + # Only launch grid steps for blocks overlapping `[0, num_tokens)`. When + # `num_tokens` is traced, this is a dynamic grid size. + num_steps = pl.cdiv(num_tokens if isinstance(num_tokens, int) else scalars[1], block_size) + + return pl.pallas_call( + _ragged_gather_kernel, + out_shape=out_shape, + grid=(num_steps,), + in_specs=[ + pl.BlockSpec(memory_space=pltpu.SMEM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + ], + out_specs=pl.BlockSpec(memory_space=pltpu.HBM), + scratch_shapes=[ + pltpu.VMEM((num_buffers, block_size, dim0, dim1), x.dtype), + pltpu.SMEM((idx_window,), indices.dtype), + pltpu.SemaphoreType.DMA(()), + pltpu.SemaphoreType.DMA((num_buffers,)), + pltpu.SemaphoreType.DMA((num_buffers,)), + ], + compiler_params=pltpu.CompilerParams( + dimension_semantics=(pltpu.ARBITRARY,), + ), + name="ragged_gather_tc", + )(scalars, indices, x) + + +# --- Ragged Gather Reduce on TensorCore --- + + +# While block i is reduced, the row DMAs of block i + 2 are issued in the same +# loop and block i + 1's are in flight, and block i - 1 is being stored. +_RAGGED_GATHER_REDUCE_TC_NUM_BUFFERS = 4 +_RAGGED_GATHER_REDUCE_TC_LOOKAHEAD = _RAGGED_GATHER_REDUCE_TC_NUM_BUFFERS - 2 +# Static bodies per trip of the row DMA issue and reduce loops. +_RAGGED_GATHER_REDUCE_TC_ISSUE_UNROLL = 64 +_RAGGED_GATHER_REDUCE_TC_REDUCE_UNROLL = 8 +# Alignment of dynamic 1D int32 HBM slices. +_RAGGED_GATHER_REDUCE_TC_ALIGN = 128 +# Slots per chunk of the two-level search for block boundaries. +_RAGGED_GATHER_REDUCE_TC_SEARCH_CHUNK = 64 +# Block table entries per block: first slot, slot count, first token, token +# count. +_RAGGED_GATHER_REDUCE_TC_BLOCK_FIELDS = 4 +# VMEM of the block buffers, under the kernel's VMEM limit. +_RAGGED_GATHER_REDUCE_TC_VMEM_LIMIT = 64 * 1024 * 1024 +_RAGGED_GATHER_REDUCE_TC_VMEM_BUDGET = 56 * 1024 * 1024 + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class RaggedGatherReduceMetadata: + """Routing of `ragged_gather_reduce_tc`. + + Built by `ragged_gather_reduce_tc_metadata`. The valid slots of `indices` are + compacted so that each output token's valid rows are contiguous and in buffer + order, and output tokens are split into blocks of consecutive tokens whose + valid rows plus output rows fit in `block_size` VMEM rows. + + Attributes: + src_rows: Buffer row of each compacted slot, padded to an aligned length of + at least one SMEM slot window. + tokens: Output token of each compacted slot, padded like `src_rows`. + blocks: Number of blocks, followed by the first slot, slot count, first + token and token count of each block. + block_size: Maximum VMEM rows per block (valid input rows plus output + tokens), which also bounds the row DMAs in flight per block. + num_indices: Length of the `indices` the metadata was derived from + (`num_out_tokens * top_k`, or the buffer cap for windowed indices). + """ + + src_rows: jax.Array + tokens: jax.Array + blocks: jax.Array + block_size: int = jax.tree.static() + num_indices: int = jax.tree.static() + + +def _ragged_gather_reduce_tc_layout( + num_indices: int, + top_k: int, + block_size: int, + num_out_tokens: int | None = None, +) -> tuple[int, int, int]: + """Returns (max_blocks, window, slots_len) of the metadata layout. + + A token costs its valid rows plus its output row. Block `b` takes the tokens + whose exclusive cost prefix is in `[b * span, (b + 1) * span)`, where + `span = block_size - top_k`, so it costs at most `block_size` VMEM rows. Each + block stages a `window` of the compacted slots in SMEM, starting at an aligned + offset, so the slot arrays are padded to `slots_len`. + + Args: + num_indices: Length of the forward gather's indices. + top_k: Number of occurrences of each token id in the indices. + block_size: Maximum VMEM rows per block. + num_out_tokens: Number of output tokens; defaults to `num_indices // top_k`. + """ + if num_out_tokens is None: + num_out_tokens = num_indices // top_k + align = _RAGGED_GATHER_REDUCE_TC_ALIGN + max_blocks = pl.cdiv(num_indices + num_out_tokens, block_size - top_k) + window = pl.cdiv(block_size, align) * align + align + slots_len = pl.cdiv(max(num_indices, window), align) * align + return max_blocks, window, slots_len + + +def _check_ragged_gather_reduce_tc_args(top_k: int, block_size: int): + if top_k <= 0: + raise ValueError(f"top_k must be positive, got {top_k}.") + if block_size <= top_k: + raise ValueError( + f"block_size must exceed top_k ({top_k}) so that a block fits a token" f" and its rows, got {block_size}." + ) + + +@functools.cache +def _warn_packed_sort_fallback(n: int, max_key: int, pos_bits: int, key_bits: int) -> None: + """Logs, once per shape at trace time, that packed_stable_sort took the slower tuple sort.""" + logging.warning( + "packed_stable_sort: %d position bits (n=%d) + %d key bits (max_key=%d) > 32; " + "falling back to the (key, position) tuple sort.", + pos_bits, + n, + key_bits, + max_key, + ) + + +def packed_stable_sort(key: jax.Array, max_key: int) -> tuple[jax.Array, jax.Array]: + """Stable sort of `arange(len(key))` by non-negative int `key` <= max_key; returns (sorted_key, positions). + + Sorts the single uint32 array `key << pos_bits | position` when it fits in 32 bits (else the (key, position) + tuple sort). The positions are distinct, so every packed value is unique and the unstable single-array sort + gives the stable order. + """ + n = key.shape[0] + positions = jnp.arange(n, dtype=jnp.int32) + pos_bits = max(1, (n - 1).bit_length()) + key_bits = max(1, int(max_key).bit_length()) + if pos_bits + key_bits > 32: + _warn_packed_sort_fallback(n, int(max_key), pos_bits, key_bits) + sorted_key, sorted_positions = jax.lax.sort((key.astype(jnp.int32), positions), num_keys=1, is_stable=True) + return sorted_key, sorted_positions + shift = jnp.uint32(pos_bits) + packed_key = (key.astype(jnp.uint32) << shift) | positions.astype(jnp.uint32) + sorted_packed = jax.lax.sort(packed_key, is_stable=False) + sorted_positions = (sorted_packed & jnp.uint32((1 << pos_bits) - 1)).astype(jnp.int32) + return (sorted_packed >> shift).astype(jnp.int32), sorted_positions + + +def ragged_gather_reduce_tc_metadata( + indices: jax.Array, + start: jax.Array | int, + num_tokens: jax.Array | int, + *, + top_k: int, + block_size: int = 896, + num_out_tokens: int | None = None, + packed_sort: bool = False, +) -> RaggedGatherReduceMetadata: + """Derives the routing consumed by `ragged_gather_reduce_tc`. + + The result depends only on the routing, not on the token data, so it can be + computed once and shared by every `ragged_gather_reduce_tc` call with the + same `indices`, `start` and `num_tokens`. + + The valid slots, those in the `[start, start + num_tokens)` window, are + compacted by a stable sort on their token, so a token's valid rows are + contiguous and in buffer order. Output tokens are then split into blocks of + consecutive tokens that each fit `block_size` VMEM rows of input rows and + output tokens, so sparse windows get blocks of many tokens and dense windows + blocks of many rows. + + Each token id must appear exactly `top_k` times in `indices` (e.g. + `indices = sort_indices // top_k`). The caller must guarantee `0 <= start`, + `0 <= num_tokens` and `start + num_tokens <= num_indices`; these are not + checked at runtime and violating them is undefined behavior. + + Args: + indices: Indices used by the forward ragged gather, of shape (num_out_tokens + * top_k,), with values in [0, num_out_tokens). + start: Offset into `indices` of the index that formed the first buffer row. + May be a traced scalar. + num_tokens: Number of valid rows in the buffer. May be a traced scalar. + top_k: Number of occurrences of each token id in `indices`. + block_size: Maximum VMEM rows (valid input rows plus output tokens) per + block of the kernel. Must exceed `top_k`, and `ragged_gather_reduce_tc` + requires its block buffers to fit in VMEM for the token shape. + num_out_tokens: Number of output tokens. Defaults to `num_indices // top_k`; + set it when `indices` is already the `[start, start + cap)` window of the + full indices (then pass `start=0`). + packed_sort: Compact the slots with `packed_stable_sort` (same result, a + single-array sort) instead of the (key, position) tuple sort. + + Returns: + The compacted slots and block table. + """ + _check_ragged_gather_reduce_tc_args(top_k, block_size) + num_indices = indices.shape[0] + if num_out_tokens is None: + if num_indices % top_k: + raise ValueError(f"indices length ({num_indices}) must be divisible by top_k ({top_k}).") + num_out_tokens = num_indices // top_k + if isinstance(num_tokens, int) and num_tokens < 0: + raise ValueError(f"num_tokens must be non-negative, got {num_tokens}.") + if isinstance(start, int) and start < 0: + raise ValueError(f"start must be non-negative, got {start}.") + if isinstance(start, int) and isinstance(num_tokens, int) and start + num_tokens > num_indices: + raise ValueError(f"start + num_tokens ({start + num_tokens}) must not exceed indices" f" length ({num_indices}).") + max_blocks, _, slots_len = _ragged_gather_reduce_tc_layout( + num_indices, top_k, block_size, num_out_tokens=num_out_tokens + ) + if num_out_tokens == 0: + empty_slots = jnp.zeros((slots_len,), jnp.int32) + return RaggedGatherReduceMetadata( + src_rows=empty_slots, + tokens=empty_slots, + blocks=jnp.zeros((1,), jnp.int32), + block_size=block_size, + num_indices=num_indices, + ) + + start = jnp.asarray(start, dtype=jnp.int32) + num_tokens = jnp.asarray(num_tokens, dtype=jnp.int32) + positions = jnp.arange(num_indices, dtype=jnp.int32) + rel = positions - start + valid = (rel >= 0) & (rel < num_tokens) + # Valid slots first, by token and then buffer order; invalid slots last. + key = jnp.where(valid, indices, num_out_tokens).astype(jnp.int32) + if packed_sort: + sorted_key, sorted_positions = packed_stable_sort(key, num_out_tokens) + else: + sorted_key, sorted_positions = jax.lax.sort((key, positions), num_keys=1, is_stable=True) + pad = slots_len - num_indices + src_rows = jnp.pad(sorted_positions - start, (0, pad)) + tokens = jnp.pad(sorted_key, (0, pad)) + + # A token costs its valid rows plus its output row, so the tokens before + # slot i's token cost run_start[i] + sorted_key[i], where run_start[i] is + # the first slot of that token. This cost increases over the valid slots. + is_run_start = jnp.concatenate([jnp.ones((1,), bool), sorted_key[1:] != sorted_key[:-1]]) + run_start = jax.lax.cummax(jnp.where(is_run_start, positions, 0)) + int_max = jnp.iinfo(jnp.int32).max + slot_cost = jnp.where(sorted_key < num_out_tokens, run_start + sorted_key, int_max) + # Block b starts at the first token whose tokens before it cost at least + # threshold[b]. The slots of the tokens before it, slot_start[b], are the + # slots with a smaller cost. Searches over all slots are slow on TPU, so + # slot_start is counted in whole chunks of slots, then in the one chunk that + # straddles the threshold. + chunk = _RAGGED_GATHER_REDUCE_TC_SEARCH_CHUNK + num_chunks = pl.cdiv(num_indices, chunk) + cost_chunks = jnp.pad(slot_cost, (0, num_chunks * chunk - num_indices), constant_values=int_max).reshape( + num_chunks, chunk + ) + threshold = jnp.arange(max_blocks + 1, dtype=jnp.int32) * (block_size - top_k) + full_chunks = jnp.sum(cost_chunks[:, -1] < threshold[:, None], axis=1, dtype=jnp.int32) + row = jnp.minimum(full_chunks, num_chunks - 1) + in_chunk = jnp.sum(cost_chunks[row] < threshold[:, None], axis=1, dtype=jnp.int32) + slot_start = row * chunk + in_chunk + # The tokens between slot_start - 1's token and slot_start's token have no + # valid slots and cost 1 each, so block b starts at the first of them whose + # tokens before it reach the threshold, or else at slot_start's token. + prev_key = jnp.where(slot_start > 0, sorted_key[jnp.maximum(slot_start - 1, 0)], -1) + next_key = jnp.where( + slot_start < num_indices, + sorted_key[jnp.minimum(slot_start, num_indices - 1)], + num_out_tokens, + ) + token_start = jnp.clip(threshold - slot_start, prev_key + 1, next_key) + table = jnp.stack( + [ + slot_start[:-1], + jnp.diff(slot_start), + token_start[:-1], + jnp.diff(token_start), + ], + axis=1, + ) + num_blocks = jnp.sum(token_start[:-1] < num_out_tokens, dtype=jnp.int32) + blocks = jnp.concatenate([num_blocks[None], table.reshape(-1)]) + return RaggedGatherReduceMetadata( + src_rows=src_rows, + tokens=tokens, + blocks=blocks.astype(jnp.int32), + block_size=block_size, + num_indices=num_indices, + ) + + +def ragged_gather_reduce_tc( + x: jax.Array, + metadata: RaggedGatherReduceMetadata, + *, + num_out_tokens: int, + top_k: int, +) -> jax.Array: + """Transpose of `ragged_gather_tc`, reducing top-k duplicates by sum. + + `x` is a token buffer formed by `ragged_gather_tc(src, indices, start, + num_tokens, ...)`, i.e. `x[i] = src[indices[start + i]]` for + `i < num_tokens`. Each valid row is sent back to its source token, and rows + that land on the same token are summed: + `out[t] = sum_{i < num_tokens, indices[start + i] == t} x[i]`. Rows of `x` at + or past `num_tokens` are never read. Accumulation is done in float32 in + buffer order. An output token with no valid rows is zero. + + The routing is given by `metadata`, from + `ragged_gather_reduce_tc_metadata(indices, start, num_tokens, top_k=top_k, + ...)`. The caller must also guarantee `num_tokens <= max_input_tokens`; this + is not checked at runtime and violating it is undefined behavior. + + Args: + x: Ragged token buffer of shape (max_input_tokens, d0, d1). + metadata: Routing from `ragged_gather_reduce_tc_metadata`. + num_out_tokens: Number of output tokens the metadata was derived for. + top_k: Number of occurrences of each token id in `indices`. + + Returns: + Reduced tokens of shape (num_out_tokens, d0, d1). + + Raises: + ValueError: If the metadata was not derived for `num_out_tokens * top_k` + indices and `top_k`, or its block buffers do not fit in VMEM. + """ + max_input_tokens, dim0, dim1 = x.shape + block_size = metadata.block_size + _check_ragged_gather_reduce_tc_args(top_k, block_size) + if num_out_tokens < 0: + raise ValueError(f"num_out_tokens must be non-negative, got {num_out_tokens}.") + num_buffers = _RAGGED_GATHER_REDUCE_TC_NUM_BUFFERS + align = _RAGGED_GATHER_REDUCE_TC_ALIGN + fields = _RAGGED_GATHER_REDUCE_TC_BLOCK_FIELDS + max_blocks, window, slots_len = _ragged_gather_reduce_tc_layout( + metadata.num_indices, top_k, block_size, num_out_tokens=num_out_tokens + ) + expected_shapes = ((slots_len,), (slots_len,), (1 + fields * max_blocks,)) + shapes = ( + metadata.src_rows.shape, + metadata.tokens.shape, + metadata.blocks.shape, + ) + if shapes != expected_shapes: + raise ValueError( + "metadata (src_rows, tokens, blocks) must have shapes" + f" {expected_shapes}, got {shapes}; was it derived for" + f" num_out_tokens={num_out_tokens} with top_k={top_k}?" + ) + # A (d0, d1) token is padded to whole (sublanes, 128) tiles in VMEM. + itemsize = jnp.dtype(x.dtype).itemsize + sublanes = 8 * max(1, 4 // itemsize) + token_vmem_bytes = pl.cdiv(dim0, sublanes) * sublanes * pl.cdiv(dim1, 128) * 128 * itemsize + vmem_budget = _RAGGED_GATHER_REDUCE_TC_VMEM_BUDGET + if num_buffers * block_size * token_vmem_bytes > vmem_budget: + raise ValueError( + f"block_size {block_size} does not fit {num_buffers} VMEM buffers of" + f" ({dim0}, {dim1}) {x.dtype} tokens in {vmem_budget} bytes; use" + f" block_size <= {vmem_budget // (num_buffers * token_vmem_bytes)}." + ) + if num_out_tokens == 0 or max_input_tokens == 0: + return jnp.zeros((num_out_tokens, dim0, dim1), x.dtype) + + lookahead = _RAGGED_GATHER_REDUCE_TC_LOOKAHEAD + # A block's slot window is loaded two steps before its rows are issued and + # is live through its reduce, lookahead steps later, so windows cycle + # through a ring of lookahead + 2 SMEM slots. + ring = lookahead + 2 + + def _ragged_gather_reduce_kernel( + blocks_smem_ref, + src_hbm_ref, + tok_hbm_ref, + x_hbm_ref, + o_hbm_ref, + vmem_ref, + src_smem_ref, + tok_smem_ref, + recv_sem, + send_sem, + slots_sem, + ): + num_blocks = blocks_smem_ref[0] + + # Block blk reduces compacted slots [slot_start, slot_start + num_slots) + # into output tokens [token_start, token_start + num_tokens). In its VMEM + # buffer, rows [0, num_tokens) hold the output tokens and rows + # [num_tokens, num_tokens + num_slots) the gathered input rows. + def _block(blk): + entry = 1 + fields * blk + return tuple(blocks_smem_ref[entry + f] for f in range(fields)) + + def _window_start(slot_start): + return jnp.minimum((slot_start // align) * align, slots_len - window) + + def _slots_loads(blk): + r = blk % ring + window_start = pl.multiple_of(_window_start(_block(blk)[0]), align) + return [ + pltpu.make_async_copy( + hbm_ref.at[pl.ds(window_start, window)], + smem_ref.at[pl.ds(r * window, window)], + slots_sem.at[r], + ) + for hbm_ref, smem_ref in ( + (src_hbm_ref, src_smem_ref), + (tok_hbm_ref, tok_smem_ref), + ) + ] + + def _slot_base(blk): + """Returns the SMEM offset of block blk's first slot.""" + slot_start = _block(blk)[0] + return (blk % ring) * window + slot_start - _window_start(slot_start) + + def _pow2_chunks(n, fn): + """Calls fn(off, size) on static power-of-two sizes summing to n.""" + off = 0 + for k in reversed(range(block_size.bit_length())): + size = 1 << k + bit = n & size + + def _fn_if_bit_set(off=off, size=size): + fn(off, size) + + pl.when(bit != 0)(_fn_if_bit_set) + + off = off + bit + + def _wait_rows(b, n, sem): + def _wait(off, size): + del off + rows = vmem_ref.at[b, pl.ds(0, size), :, :] + pltpu.make_async_copy(rows, rows, sem.at[b]).wait() + + _pow2_chunks(n, _wait) + + def _dynamic_fori(lo, hi, unroll, body, carry): + """fori_loop over [lo, hi) with `unroll` static bodies per trip.""" + + def _chunk(c, carry): + for u in range(unroll): + carry = body(lo + c * unroll + u, carry) + return carry + + num_full = (hi - lo) // unroll + carry = jax.lax.fori_loop(0, num_full, _chunk, carry) + return jax.lax.fori_loop(lo + num_full * unroll, hi, body, carry) + + def _issuer(blk): + """Returns a loop body that starts block blk's i-th row DMA.""" + b = blk % num_buffers + num_tokens = _block(blk)[3] + base = _slot_base(blk) + + def _issue(i, carry): + pltpu.make_async_copy( + x_hbm_ref.at[pl.ds(src_smem_ref[base + i], 1), :, :], + vmem_ref.at[b, pl.ds(num_tokens + i, 1), :, :], + recv_sem.at[b], + ).start() + return carry + + return _issue + + def _in_start(blk): + _dynamic_fori( + 0, + _block(blk)[1], + _RAGGED_GATHER_REDUCE_TC_ISSUE_UNROLL, + _issuer(blk), + (), + ) + + def _in_wait(blk): + _wait_rows(blk % num_buffers, _block(blk)[1], recv_sem) + + def _reducer(blk): + """Zeroes block blk's tokens and returns its row reduce loop body.""" + b = blk % num_buffers + _, _, token_start, num_tokens = _block(blk) + + # Tokens without valid rows are zero. + @pl.loop(0, num_tokens) + def _zero_token(t): + vmem_ref[b, t, :, :] = jnp.zeros((dim0, dim1), x.dtype) + + # A token's valid rows are contiguous in the compacted order, so sum + # them in a running float32 total that restarts when the token changes. + # Every row stores the running total; the token's last row leaves the + # final sum. + base = _slot_base(blk) + + def _reduce(i, carry): + prev_t, total = carry + t = tok_smem_ref[base + i] - token_start + row = vmem_ref[b, num_tokens + i, :, :].astype(jnp.float32) + total = jnp.where(t == prev_t, total, 0.0) + row + vmem_ref[b, t, :, :] = total.astype(x.dtype) + return t, total + + return _reduce + + def _out_start(blk): + b = blk % num_buffers + _, _, token_start, num_tokens = _block(blk) + + def _store(off, size): + pltpu.make_async_copy( + vmem_ref.at[b, pl.ds(off, size), :, :], + o_hbm_ref.at[pl.ds(token_start + off, size), :, :], + send_sem.at[b], + ).start() + + _pow2_chunks(num_tokens, _store) + + def _out_wait(blk): + _wait_rows(blk % num_buffers, _block(blk)[3], send_sem) + + def _slots_wait(blk): + for load in _slots_loads(blk): + load.wait() + + for blk in range(ring): + + @pl.when(blk < num_blocks) + def _load_first_slots(blk=blk): + for load in _slots_loads(blk): + load.start() + + for blk in range(lookahead): + + @pl.when(blk < num_blocks) + def _start_first_rows(blk=blk): + _slots_wait(blk) + _in_start(blk) + + @pl.loop(0, num_blocks) + def _step(i): + _in_wait(i) + nxt = i + lookahead + has_nxt = nxt < num_blocks + + # Block nxt reuses the buffer of block nxt - num_buffers, whose store + # must be done. + @pl.when(nxt >= num_buffers) + def _free_buffer(): + _out_wait(nxt - num_buffers) + + @pl.when(has_nxt) + def _wait_next_slots(): + _slots_wait(nxt) + + # Block nxt's row DMAs are started in block i's reduce loop, so the + # scalar DMA issue shares bundles with the vector reduce. + nxt = jnp.minimum(nxt, num_blocks - 1) + num_reduce = _block(i)[1] + num_issue = jnp.where(has_nxt, _block(nxt)[1], 0) + num_both = jnp.minimum(num_reduce, num_issue) + issue = _issuer(nxt) + reduce = _reducer(i) + + def _issue_and_reduce(j, carry): + issue(j, ()) + return reduce(j, carry) + + unroll = _RAGGED_GATHER_REDUCE_TC_REDUCE_UNROLL + carry = (jnp.int32(-1), jnp.zeros((dim0, dim1), jnp.float32)) + carry = _dynamic_fori(0, num_both, unroll, _issue_and_reduce, carry) + _dynamic_fori(num_both, num_reduce, unroll, reduce, carry) + _dynamic_fori( + num_both, + num_issue, + _RAGGED_GATHER_REDUCE_TC_ISSUE_UNROLL, + issue, + (), + ) + _out_start(i) + + # Block i's slot window is free now that its reduce is done. + @pl.when(i + ring < num_blocks) + def _load_later_slots(): + for load in _slots_loads(i + ring): + load.start() + + # The last num_buffers - lookahead stores are still pending. + for k in range(num_buffers - lookahead): + + @pl.when(num_blocks - 1 - k >= 0) + def _wait_last_stores(k=k): + _out_wait(num_blocks - 1 - k) + + manual_axis_type = getattr(jax.typeof(x), "manual_axis_type", None) + out_shape = jax.ShapeDtypeStruct((num_out_tokens, dim0, dim1), x.dtype, manual_axis_type=manual_axis_type) + hbm_spec = pl.BlockSpec(memory_space=pltpu.HBM) + dma_sems = pltpu.SemaphoreType.DMA((num_buffers,)) + return pl.pallas_call( + _ragged_gather_reduce_kernel, + out_shape=out_shape, + in_specs=[ + pl.BlockSpec(memory_space=pltpu.SMEM), + hbm_spec, + hbm_spec, + hbm_spec, + ], + out_specs=hbm_spec, + scratch_shapes=[ + pltpu.VMEM((num_buffers, block_size, dim0, dim1), x.dtype), + pltpu.SMEM((ring * window,), jnp.int32), + pltpu.SMEM((ring * window,), jnp.int32), + dma_sems, + dma_sems, + pltpu.SemaphoreType.DMA((ring,)), + ], + compiler_params=pltpu.CompilerParams( + vmem_limit_bytes=_RAGGED_GATHER_REDUCE_TC_VMEM_LIMIT, + ), + name="ragged_gather_reduce_tc", + )(metadata.blocks, metadata.src_rows, metadata.tokens, x) diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py index dc563cef68..5c7469e7d2 100644 --- a/src/maxtext/kernels/ragged/ragged_sort.py +++ b/src/maxtext/kernels/ragged/ragged_sort.py @@ -15,10 +15,74 @@ """Ragged token sorting operations with custom VJP.""" import jax -from jax.ad_checkpoint import checkpoint_name import jax.numpy as jnp from maxtext.kernels.ragged.ragged_gather import ragged_gather +from maxtext.kernels.ragged.ragged_gather_reduce_v2 import argsort_unique_pairs from maxtext.kernels.ragged.ragged_gather_reduce_v2 import ragged_gather_reduce +from maxtext.kernels.ragged.ragged_gather_tc import packed_stable_sort + +_INDEX_MATH_MODES = ("sort", "count") + + +def _check_index_math(index_math): + if index_math not in _INDEX_MATH_MODES: + raise ValueError(f"index_math must be one of {_INDEX_MATH_MODES}, got {index_math!r}") + + +def _local_slot_positions(topk_indices_local, num_experts, shard_idx, ep_size, block=128): + """Position of every flat (token, k) slot routed to this shard, inside this shard's window of the sorted slots. + + Equals ``argsort(argsort(flat_ids)) - shard_output_start`` on the slots of this shard's experts (``argsort`` + stable, by expert id) and is -1 on every other slot. It needs no sort and no scatter: in the stable order by + expert id, slot (t, k) with local expert e comes after + - every slot of a lower local expert (an exclusive prefix sum of the local expert counts), + - every slot of expert e in an earlier token (an exclusive prefix count over tokens), and + - every earlier slot of token t that also picked e (0 unless a token repeats an expert). + The prefix count over tokens runs per block of ``block`` tokens as a strictly-lower-triangular matmul (its + inputs are counts <= topk, exact in bfloat16, and it accumulates in float32, exact below 2**24), plus an int32 + prefix sum over the block totals. + + Args: + topk_indices_local: ``[num_tokens, topk]`` expert ids. Ids outside this shard's experts (including the + ``num_experts`` padding id) are never local. + num_experts: total number of experts. + shard_idx: this shard's index on the expert-parallel axis (a traced scalar is fine). + ep_size: number of expert shards. + block: tokens per block of the prefix count. + + Returns: + ``[num_tokens * topk]`` int32 positions in ``[0, local slots)``, or -1. + """ + num_tokens, topk = topk_indices_local.shape + local_num_experts = num_experts // ep_size + pad = (-num_tokens) % block + ids = topk_indices_local.astype(jnp.int32) + if pad: + ids = jnp.pad(ids, ((0, pad), (0, 0)), constant_values=num_experts) # the padding id: never local + t = ids.shape[0] + le = ids - shard_idx * local_num_experts # local expert id, in [0, local_num_experts) iff the slot is local + is_local = (le >= 0) & (le < local_num_experts) + hit = le[:, :, None] == jnp.arange(local_num_experts, dtype=jnp.int32)[None, None, :] # [t, topk, E_local] + counts_per_token = jnp.sum(hit, axis=1, dtype=jnp.int32) # [t, E_local] + kk = jnp.arange(topk, dtype=jnp.int32) + dup_before = jnp.sum( + (le[:, :, None] == le[:, None, :]) & (kk[None, None, :] < kk[None, :, None]), axis=-1, dtype=jnp.int32 + ) # [t, topk]: earlier slots of the same token with the same expert + blocks = counts_per_token.reshape(t // block, block, local_num_experts) + strict_lower = jnp.tril(jnp.ones((block, block), jnp.bfloat16), k=-1) + # Slots of each local expert in the earlier tokens of the same block. + within = jnp.einsum( + "ij,bjk->bik", strict_lower, blocks.astype(jnp.bfloat16), preferred_element_type=jnp.float32 + ).astype(jnp.int32) + block_totals = jnp.sum(blocks, axis=1, dtype=jnp.int32) # [t // block, E_local] + before_block = jnp.cumsum(block_totals, axis=0) - block_totals + earlier = (within + before_block[:, None, :]).reshape(t, local_num_experts) # [t, E_local] + local_counts = jnp.sum(block_totals, axis=0) + local_offsets = jnp.cumsum(local_counts) - local_counts + # One select-reduce over the local experts picks each slot's own expert's terms (no gather). + pos = jnp.sum(jnp.where(hit, earlier[:, None, :] + local_offsets[None, None, :], 0), axis=-1) + dup_before + pos = jnp.where(is_local, pos, -1)[:num_tokens] + return pos.reshape(-1) def ring_ragged_sort( @@ -36,6 +100,9 @@ def ring_ragged_sort( gather_bytes_accessed_override=-1, gather_reduce_bytes_accessed_override=-1, use_single_sparsecore=False, + tag_routing_fn=None, + index_math="sort", + unique_pair_argsort=False, ): """Ragged-gather variant for AG-RS Expert Parallelism token routing. @@ -64,6 +131,19 @@ def ring_ragged_sort( ep_name: ``str`` identifying the expert parallel axis name. ep_size: scalar ``int`` representing the expert parallel mesh size. buffer_size: optional scalar ``int`` representing the size of the local buffer. + tag_routing_fn: optional function applied to each integer routing array (the sort permutation, its inverse, + the sorted token indices and the group sizes). + index_math: how the routing index maps are computed (``moe_routing_index_math``). + "sort" (default): the permutation and its inverse by two stable argsorts. + "count": the permutation by one single-operand packed uint32 sort, the sorted token indices by division, + and the inverse permutation by counting, on this shard's slots only (-1 on every other slot); the + gather-reduce of the backward packs its validity sort too. Every reader of the inverse masks it to this + shard's window, so the forward output and the gradients are the same. The unsort of these rows must then + get ``topk_argsort_indices`` and ``index_math="count"`` (a partial inverse cannot be argsorted back). + unique_pair_argsort: with ``index_math="sort"``, run both argsorts, and the validity sort of the backward's + gather-reduce, as unstable two-operand sorts of the unique (key, index) pairs (``argsort_unique_pairs``, + ``moe_ragged_sort_unique_pair_argsort``) instead of stable argsorts, which XLA lowers with a duplicated iota + operand. Same permutations, same values. Returns: A tuple containing: @@ -72,33 +152,60 @@ def ring_ragged_sort( only the tokens destined for local experts, padded with zeros elsewhere. - 1D tensor ``group_sizes_local`` tracking expert token counts. - 1D tensor ``topk_argsort_revert_indices`` for inverse routing. + - 1D tensor ``topk_argsort_indices``, the inverse of ``topk_argsort_revert_indices``. """ + _check_index_math(index_math) + # The routing index maps depend only on the (integer) expert ids, so they are computed outside the custom_vjp. + # Only the gather below is differentiated. + num_tokens_local = hidden_states_local.shape[0] + + topk_indices_flat = topk_indices_local.flatten() # num_tokens_local x topk + argsort = argsort_unique_pairs if unique_pair_argsort else jnp.argsort # the "sort" maps' argsort + if index_math == "count": + # The stable order by expert id as ONE single-operand sort of (expert << pos_bits | slot): the packed values + # are distinct, so the unstable sort has one result, the stable argsort's permutation. jnp.argsort sorts the + # (expert, iota) pair stably, which XLA lowers in the forward pass with a duplicated iota operand. + _, topk_argsort_indices = packed_stable_sort(topk_indices_flat.astype(jnp.int32), num_experts) + # Flat slot j belongs to token j // topk, so the sorted token indices need no gather of a repeated iota. + token_indices_sorted = (topk_argsort_indices // topk).astype(jnp.int32) # num_tokens_local x topk + else: + topk_argsort_indices = argsort(topk_indices_flat) # num_tokens_local x topk - @jax.custom_vjp - def _ring_ragged_sort(hidden_states_local, topk_indices_local): - """Sort and gather activations to different EP shards.""" - return _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local)[0] - - @jax.named_scope("ragged-sort-fwd") - def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): - """Sort and gather activations forward pass.""" + token_indices = jnp.arange(num_tokens_local, dtype=jnp.int32).repeat(topk) # num_tokens_local x topk + token_indices_sorted = token_indices[topk_argsort_indices] # num_tokens_local x topk - num_tokens_local = hidden_states_local.shape[0] + group_sizes_local = jax.nn.one_hot(topk_indices_flat, num_experts, dtype=jnp.int32).sum(axis=0) # GLOBAL_NUM_EXPERTS - topk_indices_flat = topk_indices_local.flatten() # num_tokens_local x topk - topk_argsort_indices = jnp.argsort(topk_indices_flat) # num_tokens_local x topk + if index_math == "count": + # Every reader of the inverse permutation masks it to [shard_output_start, shard_output_end), so it is + # computed on this shard's slots only, by counting instead of a second argsort, and is -1 elsewhere. + shard_idx = jax.lax.axis_index(ep_name) + local_pos = _local_slot_positions(topk_indices_local, num_experts, shard_idx, ep_size) + shard_start = jnp.sum(jnp.where(jnp.arange(num_experts) < shard_idx * (num_experts // ep_size), group_sizes_local, 0)) + topk_argsort_revert_indices = jnp.where(local_pos >= 0, shard_start + local_pos, -1).astype(jnp.int32) + else: + topk_argsort_revert_indices = argsort(topk_argsort_indices) # num_tokens_local x topk + packed_preprocess = index_math == "count" + unique_pair_preprocess = unique_pair_argsort and not packed_preprocess + + if tag_routing_fn is not None: + topk_argsort_indices, token_indices_sorted, group_sizes_local, topk_argsort_revert_indices = ( + tag_routing_fn(t) + for t in (topk_argsort_indices, token_indices_sorted, group_sizes_local, topk_argsort_revert_indices) + ) - token_indices = jnp.arange(num_tokens_local, dtype=jnp.int32).repeat(topk) # num_tokens_local x topk - token_indices_sorted = token_indices[topk_argsort_indices] # num_tokens_local x topk + @jax.custom_vjp + def _ring_ragged_sort(hidden_states_local, token_indices_sorted, group_sizes_local, topk_argsort_revert_indices): + """Gather activations to different EP shards.""" + return _ring_ragged_sort_fwd( + hidden_states_local, token_indices_sorted, group_sizes_local, topk_argsort_revert_indices + )[0] - group_sizes_local = jax.nn.one_hot(topk_indices_flat, num_experts, dtype=jnp.int32).sum(axis=0) # GLOBAL_NUM_EXPERTS + @jax.named_scope("ragged-sort-fwd") + def _ring_ragged_sort_fwd(hidden_states_local, token_indices_sorted, group_sizes_local, topk_argsort_revert_indices): + """Gather activations forward pass.""" - topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices) # num_tokens_local x topk - # Named for the `moe_route` remat key. Saving these integer results lets the backward replay - # the token gather, whose output the expert GMMs need again, without replaying both argsorts. - token_indices_sorted = checkpoint_name(token_indices_sorted, "moe_route") - group_sizes_local = checkpoint_name(group_sizes_local, "moe_route") - topk_argsort_revert_indices = checkpoint_name(topk_argsort_revert_indices, "moe_route") + num_tokens_local = hidden_states_local.shape[0] shard_idx = jax.lax.axis_index(ep_name) local_num_experts = num_experts // ep_size @@ -145,8 +252,6 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): use_single_sparsecore=use_single_sparsecore, ) - out = (x, group_sizes_local, topk_argsort_revert_indices) - res = ( topk_argsort_revert_indices, shard_output_start, @@ -155,7 +260,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): hidden_states_local.shape, ) - return out, res + return x, res @jax.named_scope("ragged-sort-bwd") def _ring_ragged_sort_bwd(res, g_out): @@ -173,7 +278,7 @@ def _ring_ragged_sort_bwd(res, g_out): local_buffer_size, _, ) = res - g_x, _, _ = g_out + g_x = g_out # Restrict to the [start, end) source range via a validity bitmask. The # ragged kernel packs valid rows to the front of each row-partition and # only iterates over the populated prefix, so we hand it the mask directly @@ -198,6 +303,8 @@ def _ring_ragged_sort_bwd(res, g_out): flops_override=gather_reduce_flops_override, bytes_accessed_override=gather_reduce_bytes_accessed_override, use_single_sparsecore=use_single_sparsecore, + packed_preprocess=packed_preprocess, + unique_pair_argsort=unique_pair_preprocess, ) else: # Buffering: g_x has size `local_buffer_size` (packed). @@ -224,12 +331,210 @@ def _ring_ragged_sort_bwd(res, g_out): flops_override=gather_reduce_flops_override, bytes_accessed_override=gather_reduce_bytes_accessed_override, use_single_sparsecore=use_single_sparsecore, + packed_preprocess=packed_preprocess, + unique_pair_argsort=unique_pair_preprocess, ) - return grad_hidden_states, None + return grad_hidden_states, None, None, None _ring_ragged_sort.defvjp(_ring_ragged_sort_fwd, _ring_ragged_sort_bwd) - return _ring_ragged_sort(hidden_states_local, topk_indices_local) + x = _ring_ragged_sort(hidden_states_local, token_indices_sorted, group_sizes_local, topk_argsort_revert_indices) + return x, group_sizes_local, topk_argsort_revert_indices, topk_argsort_indices + + +def _unsort_kwargs( + enforce_gather_fallback, + enforce_gather_reduce_fallback, + gather_flops_override, + gather_reduce_flops_override, + gather_bytes_accessed_override, + gather_reduce_bytes_accessed_override, + use_single_sparsecore, + index_math, + unique_pair_argsort, +): + """The keyword arguments of the unsort's ragged gather (backward) and ragged gather-reduce (forward).""" + _check_index_math(index_math) + gather_kwargs = { + "enforce_fallback": enforce_gather_fallback, + "flops_override": gather_flops_override, + "bytes_accessed_override": gather_bytes_accessed_override, + "use_single_sparsecore": use_single_sparsecore, + } + gather_reduce_kwargs = { + "enforce_fallback": enforce_gather_reduce_fallback, + "flops_override": gather_reduce_flops_override, + "bytes_accessed_override": gather_reduce_bytes_accessed_override, + "use_single_sparsecore": use_single_sparsecore, + "packed_preprocess": index_math == "count", + "unique_pair_argsort": unique_pair_argsort and index_math != "count", + } + return gather_kwargs, gather_reduce_kwargs + + +def _ring_unsort_window(group_sizes_local, local_num_experts, ep_name): + """This shard's ``[start, end)`` window in the globally expert-sorted slot order.""" + group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True) + + shard_idx = jax.lax.axis_index(ep_name) + experts_start = shard_idx * local_num_experts + experts_end = experts_start + local_num_experts + + return group_offsets[experts_start], group_offsets[experts_end] + + +def _ring_unsort_combine( + sorted_tokens_local, + topk_argsort_revert_indices, + topk_weights_flat, + shard_output_start, + shard_output_end, + topk, + gather_reduce_kwargs, +): + """The unsort's forward: each token's routing-weighted sum over its slots on this shard's experts.""" + buffer_size = sorted_tokens_local.shape[0] + num_tokens = topk_argsort_revert_indices.shape[0] + + # We support two buffering modes: + # 1. buffer_size >= num_tokens: sorted_tokens_local has size >= num_tokens, + # and local tokens are at their global positions [start, end). + # 2. buffer_size < num_tokens: sorted_tokens_local has size < num_tokens, + # and local tokens are packed at [0, local_num_tokens). + if buffer_size >= num_tokens: + # Express the scatter as a gather-reduce: each output row pulls + # from sorted_tokens_local at position `topk_argsort_revert_indices[i]` if + # that position is within this shard's [start, end) range, else zero. + # The routing weights are applied per-row before the topk reduction. + valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & ( + topk_argsort_revert_indices < shard_output_end + ) + return ragged_gather_reduce( + sorted_tokens_local, + topk_argsort_revert_indices, + topk_weights=topk_weights_flat, + valid_rows_mask=valid_rows_mask, + reduce_group_size=topk, + **gather_reduce_kwargs, + ) + # Shift indices so they map to the packed local buffer [0, local_num_tokens). + shifted_indices = topk_argsort_revert_indices - shard_output_start + local_num_tokens = shard_output_end - shard_output_start + limit = jnp.minimum(local_num_tokens, buffer_size) + valid_rows_mask = (shifted_indices >= 0) & (shifted_indices < limit) + safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0) + return ragged_gather_reduce( + sorted_tokens_local, + safe_indices, + topk_weights=topk_weights_flat, + valid_rows_mask=valid_rows_mask, + reduce_group_size=topk, + **gather_reduce_kwargs, + ) + + +def _ring_unsort_fan_out( + g_hidden_states_local, + topk_argsort_revert_indices, + topk_weights_flat, + shard_output_start, + shard_output_end, + buffer_size, + topk, + topk_argsort_indices, + gather_kwargs, + mask_rows, + index_math="sort", +): + """The first half of the unsort's backward: the output cotangent sent to the rows of the sorted buffer. + + Returns ``(gathered, weight_for_sorted, mask, to_slot_order)``: row ``j`` of ``gathered`` is the unweighted + cotangent of the token that sorted row ``j`` came from; ``weight_for_sorted[j]`` is that slot's routing + weight; ``mask`` marks the rows of this shard's experts; ``to_slot_order`` maps a per-row value that is + zero outside ``mask`` back to flat slot order, with zero for the slots of other shards and dropped slots. + + With ``mask_rows`` the rows of ``gathered`` outside ``mask`` are zeroed (:func:`ring_ragged_unsort`, whose + cotangent output is ``gathered`` scaled). Without it they are left as the ragged gather wrote them (unwritten by + the SparseCore kernel, other slots' rows on the reference fallback), which saves a pass over ``[rows, hidden]`` + for a caller whose consumers never read those rows (:func:`ring_ragged_unsort_down_projection`). + """ + n = topk_argsort_revert_indices.shape[0] + # Build the inverse permutation idx_inv such that idx_inv[j] = i + # where revert[i] = j. revert is itself argsort(topk_argsort_indices), so its inverse is + # topk_argsort_indices; reuse it when the caller has it rather than sorting again. + if topk_argsort_indices is not None: + idx_inv = topk_argsort_indices + elif index_math == "count": + # The "count" inverse holds only this shard's slots (-1 elsewhere), so it cannot be argsorted back. + raise ValueError("index_math='count' needs topk_argsort_indices (the forward permutation of ring_ragged_sort).") + else: + idx_inv = jnp.argsort(topk_argsort_revert_indices) + + # Handle the same two buffering modes for backward pass. + # ragged_gather does the fan-out, by indexing into the un-expanded + # g_hidden_states_local via idx_inv // topk. It gathers unweighted so the same rows + # feed both gradients: the activation one after scaling, the weight one after a dot. + if buffer_size >= n: + weight_for_sorted = topk_weights_flat[idx_inv] + gathered = ragged_gather( + g_hidden_states_local, + idx_inv // topk, + shard_output_start[None], + shard_output_end[None], + **gather_kwargs, + ) + # The gradients of the elements outside the valid shard output range are zero. + mask = (jnp.arange(n) >= shard_output_start) & (jnp.arange(n) < shard_output_end) + if mask_rows: + gathered = jnp.where(mask[:, None], gathered, 0.0) + + def from_buffer_rows(per_row): + # Row j of the sorted buffer is slot idx_inv[j], so slot i reads row revert[i]. + if index_math == "count": + # Other shards' slots hold -1, not a row outside the window, so mask them instead of reading row -1. + slot_valid = (topk_argsort_revert_indices >= shard_output_start) & ( + topk_argsort_revert_indices < shard_output_end + ) + return jnp.where(slot_valid, per_row[jnp.where(slot_valid, topk_argsort_revert_indices, 0)], 0.0) + return per_row[topk_argsort_revert_indices] + + return gathered, weight_for_sorted, mask, from_buffer_rows + + # Slice the inverse permutation and the per-slot routing weights to match the packed local buffer. + padded_idx_inv = jnp.pad(idx_inv, (0, buffer_size)) + sliced_idx_inv = jax.lax.dynamic_slice_in_dim(padded_idx_inv, shard_output_start, buffer_size, axis=0) + gather_end = jnp.minimum(shard_output_end - shard_output_start, buffer_size) + padded_weights = jnp.pad(topk_weights_flat[idx_inv], (0, buffer_size)) + sliced_weights = jax.lax.dynamic_slice_in_dim(padded_weights, shard_output_start, buffer_size, axis=0) + gathered = ragged_gather( + g_hidden_states_local, + sliced_idx_inv // topk, + jnp.int32(0)[None], + gather_end[None], + **gather_kwargs, + ) + # The gradients of the elements beyond the valid limit of the local buffer are zero. + limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size) + mask = jnp.arange(buffer_size) < limit + if mask_rows: + gathered = jnp.where(mask[:, None], gathered, 0.0) + + def from_packed_rows(per_row): + if index_math == "count": + # Route the per-row value back to flat slot order with a gather through the inverse permutation instead + # of a TensorCore scatter: each valid slot reads its own buffer row; other shards' slots (-1) and + # dropped slots read 0, as the scatter leaves them. The table is padded to n entries because XLA only + # offloads a 1-D gather to the SparseCore for a large table (v7x AOT: 225,280 entries -> TensorCore + # gather at ~7 ns/element; 655,360 -> SparseCore). + shifted = topk_argsort_revert_indices - shard_output_start + slot_valid = (shifted >= 0) & (shifted < limit) + dot_table = jnp.pad(per_row, (0, max(0, n - buffer_size))) + return jnp.where(slot_valid, dot_table[jnp.where(slot_valid, shifted, 0)], 0.0) + # Scatter the per-row value back to flat slot order; dropped slots stay zero. + slots = jnp.where(mask, sliced_idx_inv, n) + return jnp.zeros((n,), jnp.float32).at[slots].set(per_row, mode="drop") + + return gathered, sliced_weights, mask, from_packed_rows def ring_ragged_unsort( @@ -247,6 +552,9 @@ def ring_ragged_unsort( gather_bytes_accessed_override=-1, gather_reduce_bytes_accessed_override=-1, use_single_sparsecore=False, + topk_argsort_indices=None, + index_math="sort", + unique_pair_argsort=False, ): """Dual of :func:`ring_ragged_sort`. @@ -273,11 +581,32 @@ def ring_ragged_unsort( ep_name: ``str`` identifying the expert parallel axis name. topk_weights: ``[num_tokens_local * topk]`` tensor of per-slot routing weights. Differentiated: its gradient is what trains the router. + topk_argsort_indices: optional inverse of ``topk_argsort_revert_indices`` (as returned by + :func:`ring_ragged_sort`). The backward needs this permutation; when it is + None, the backward re-derives it with an argsort. Both give the same permutation. + Required with ``index_math="count"``, whose inverse covers only this shard's slots. + index_math: must match the :func:`ring_ragged_sort` call that produced the routing. With "count" the + gather-reduce packs its validity sort, and in the backward the routing-weight gradient goes back to slot + order by a masked gather through the inverse permutation instead of a scatter (each valid slot reads its + own row; other shards' and dropped slots read 0). Same values either way. + unique_pair_argsort: with ``index_math="sort"``, the forward gather-reduce's validity sort as a two-operand + sort of unique pairs (as in :func:`ring_ragged_sort`). Same values. Returns: A 2D ``[num_tokens_local, hidden]`` tensor with expert outputs scattered back to their original global sequence locations, with unpopulated positions zeroed out. """ + gather_kwargs, gather_reduce_kwargs = _unsort_kwargs( + enforce_gather_fallback, + enforce_gather_reduce_fallback, + gather_flops_override, + gather_reduce_flops_override, + gather_bytes_accessed_override, + gather_reduce_bytes_accessed_override, + use_single_sparsecore, + index_math, + unique_pair_argsort, + ) @jax.custom_vjp def _ring_ragged_unsort( @@ -285,6 +614,7 @@ def _ring_ragged_unsort( group_sizes_local, topk_argsort_revert_indices, topk_weights_flat, + topk_argsort_indices, ): """Unsort and scatter activations.""" return _ring_ragged_unsort_fwd( @@ -292,6 +622,7 @@ def _ring_ragged_unsort( group_sizes_local, topk_argsort_revert_indices, topk_weights_flat, + topk_argsort_indices, )[0] @jax.named_scope("ragged-unsort-fwd") @@ -300,73 +631,28 @@ def _ring_ragged_unsort_fwd( group_sizes_local, topk_argsort_revert_indices, topk_weights_flat, + topk_argsort_indices, ): """Executes unsorting sending tokens back.""" - group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True) - - shard_idx = jax.lax.axis_index(ep_name) - experts_start = shard_idx * local_num_experts - experts_end = experts_start + local_num_experts - - shard_output_start = group_offsets[experts_start] - shard_output_end = group_offsets[experts_end] - - buffer_size = sorted_tokens_local.shape[0] - num_tokens = topk_argsort_revert_indices.shape[0] - - # We support two buffering modes: - # 1. buffer_size >= num_tokens: sorted_tokens_local has size >= num_tokens, - # and local tokens are at their global positions [start, end). - # 2. buffer_size < num_tokens: sorted_tokens_local has size < num_tokens, - # and local tokens are packed at [0, local_num_tokens). - if buffer_size >= num_tokens: - # Express the scatter as a gather-reduce: each output row pulls - # from sorted_tokens_local at position `topk_argsort_revert_indices[i]` if - # that position is within this shard's [start, end) range, else zero. - # The routing weights are applied per-row before the topk reduction. - valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & ( - topk_argsort_revert_indices < shard_output_end - ) - out = ragged_gather_reduce( - sorted_tokens_local, - topk_argsort_revert_indices, - topk_weights=topk_weights_flat, - valid_rows_mask=valid_rows_mask, - reduce_group_size=topk, - enforce_fallback=enforce_gather_reduce_fallback, - flops_override=gather_reduce_flops_override, - bytes_accessed_override=gather_reduce_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) - else: - # Shift indices so they map to the packed local buffer [0, local_num_tokens). - shifted_indices = topk_argsort_revert_indices - shard_output_start - local_num_tokens = shard_output_end - shard_output_start - limit = jnp.minimum(local_num_tokens, buffer_size) - valid_rows_mask = (shifted_indices >= 0) & (shifted_indices < limit) - safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0) - - out = ragged_gather_reduce( - sorted_tokens_local, - safe_indices, - topk_weights=topk_weights_flat, - valid_rows_mask=valid_rows_mask, - reduce_group_size=topk, - enforce_fallback=enforce_gather_reduce_fallback, - flops_override=gather_reduce_flops_override, - bytes_accessed_override=gather_reduce_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) - + shard_output_start, shard_output_end = _ring_unsort_window(group_sizes_local, local_num_experts, ep_name) + out = _ring_unsort_combine( + sorted_tokens_local, + topk_argsort_revert_indices, + topk_weights_flat, + shard_output_start, + shard_output_end, + topk, + gather_reduce_kwargs, + ) res = ( sorted_tokens_local, topk_argsort_revert_indices, topk_weights_flat, shard_output_start, shard_output_end, - buffer_size, + sorted_tokens_local.shape[0], + topk_argsort_indices, ) - return out, res @jax.named_scope("ragged-unsort-bwd") @@ -392,66 +678,28 @@ def _ring_ragged_unsort_bwd(res, g_out): shard_output_start, shard_output_end, buffer_size, + topk_argsort_indices, ) = res - g_hidden_states_local = g_out - n = topk_argsort_revert_indices.shape[0] - # Build the inverse permutation idx_inv such that idx_inv[j] = i - # where revert[i] = j. - idx_inv = jnp.argsort(topk_argsort_revert_indices) - - # Handle the same two buffering modes for backward pass. - # ragged_gather does the fan-out, by indexing into the un-expanded - # g_hidden_states_local via idx_inv // topk. It gathers unweighted so the same rows - # feed both gradients: the activation one after scaling, the weight one after a dot. - if buffer_size >= n: - weight_for_sorted = topk_weights_flat[idx_inv] - gathered = ragged_gather( - g_hidden_states_local, - idx_inv // topk, - shard_output_start[None], - shard_output_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) - # Mask out gradients that correspond to elements outside the valid shard - # output range. - mask = (jnp.arange(n) >= shard_output_start) & (jnp.arange(n) < shard_output_end) - gathered = jnp.where(mask[:, None], gathered, 0.0) - grad_sorted_tokens = (gathered * weight_for_sorted[:, None]).astype(gathered.dtype) - # Row-wise dot in sorted order, then permuted back to flat slot order. - dot_sorted = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local[:n].astype(jnp.float32), axis=-1) - grad_topk_weights = dot_sorted[topk_argsort_revert_indices] - else: - # Slice the inverse permutation to match the packed local buffer. - padded_idx_inv = jnp.pad(idx_inv, (0, buffer_size)) - sliced_idx_inv = jax.lax.dynamic_slice_in_dim(padded_idx_inv, shard_output_start, buffer_size, axis=0) - gather_end = jnp.minimum(shard_output_end - shard_output_start, buffer_size) - # Slice the per-slot routing weights to match the packed local buffer. - padded_weights = jnp.pad(topk_weights_flat[idx_inv], (0, buffer_size)) - sliced_weights = jax.lax.dynamic_slice_in_dim(padded_weights, shard_output_start, buffer_size, axis=0) - gathered = ragged_gather( - g_hidden_states_local, - sliced_idx_inv // topk, - jnp.int32(0)[None], - gather_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) - # Mask out gradients for elements beyond the valid limit of the local buffer. - limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size) - mask = jnp.arange(buffer_size) < limit - gathered = jnp.where(mask[:, None], gathered, 0.0) - grad_sorted_tokens = (gathered * sliced_weights[:, None]).astype(gathered.dtype) - # Scatter the per-slot dot back to flat slot order; dropped slots stay zero. - dot_local = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local.astype(jnp.float32), axis=-1) - slots = jnp.where(mask, sliced_idx_inv, n) - grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop") - return grad_sorted_tokens, None, None, grad_topk_weights + gathered, weight_for_sorted, _, to_slot_order = _ring_unsort_fan_out( + g_out, + topk_argsort_revert_indices, + topk_weights_flat, + shard_output_start, + shard_output_end, + buffer_size, + topk, + topk_argsort_indices, + gather_kwargs, + mask_rows=True, + index_math=index_math, + ) + grad_sorted_tokens = (gathered * weight_for_sorted[:, None]).astype(gathered.dtype) + # Row-wise dot in sorted order, then back to flat slot order (a full buffer has its rows at [0, n)). + gathered_f32 = gathered.astype(jnp.float32) + sorted_rows = sorted_tokens_local[:n] if buffer_size >= n else sorted_tokens_local + grad_topk_weights = to_slot_order(jnp.sum(gathered_f32 * sorted_rows.astype(jnp.float32), axis=-1)) + return grad_sorted_tokens, None, None, grad_topk_weights, None _ring_ragged_unsort.defvjp(_ring_ragged_unsort_fwd, _ring_ragged_unsort_bwd) @@ -463,6 +711,163 @@ def _ring_ragged_unsort_bwd(res, g_out): group_sizes_local, topk_argsort_revert_indices, topk_weights_flat, + topk_argsort_indices, + ) + + +def ring_ragged_unsort_down_projection( + intermediate_local, + wo_local, + down_projection, + down_projection_operands, + group_sizes_local, + topk_argsort_revert_indices, + topk, + local_num_experts, + ep_name, + topk_weights, + topk_argsort_indices=None, + forward_tie=None, + enforce_gather_fallback=False, + enforce_gather_reduce_fallback=False, + gather_flops_override=-1, + gather_reduce_flops_override=-1, + gather_bytes_accessed_override=-1, + gather_reduce_bytes_accessed_override=-1, + use_single_sparsecore=False, + index_math="sort", + unique_pair_argsort=False, +): + """:func:`ring_ragged_unsort` of the down projection, with a backward that does not read its output. + + Forward: ``y = down_projection(intermediate_local, wo_local, *down_projection_operands)``, the + per-expert GEMM ``y_j = h_j @ wo[e(j)]`` over the sorted rows ``h_j``, then the forward of + ``ring_ragged_unsort(y, ...)``: the same helpers (``_ring_unsort_window``, ``_ring_unsort_combine``), so the + output is bit-identical. The backward's fan-out is ``ring_ragged_unsort``'s too (``_ring_unsort_fan_out``). + + Backward. Let ``g_j`` be the unweighted output cotangent of the token that sorted row ``j`` came + from and ``w_j`` that slot's routing weight. The two ops apart compute ``dy_j = w_j g_j``, + ``dh = dy @ wo^T`` (dlhs), ``dwo = h^T @ dy`` (drhs) and the routing-weight gradient + ``g_w_j = ``. That reads ``y``, so a rematerialized layer recomputes the down-projection + GEMM in the backward only for this dot, and the dot takes a pass over two ``[rows, hidden]`` + tensors that also writes ``dy``. Row ``j`` of the GEMM is linear in ``h_j`` and the weight is a + scalar per row, so:: + + = , (w_j g_j) wo^T = w_j (g_j wo^T), h^T (w * g) = (w * h)^T g. + + This backward therefore takes the down projection's own VJP at the input ``w * h`` with the + cotangent ``g``. Its two GEMMs give ``u = g wo^T`` and ``dwo = (w * h)^T g``, and then:: + + dh_j = w_j u_j, g_w_j = . + + The GEMMs keep their shapes and cost; ``y`` is no longer read, so XLA drops its recomputation; and + the weighting and the dot run over ``[rows, mlp]`` tensors instead of ``[rows, hidden]`` ones. + Each gradient still has one rounding to the activation dtype of a weighted operand or GEMM output, + but at a different point (``w * h`` instead of ``w * g``; ``dh`` after the product with ``w`` + instead of before the GEMM; ``g_w`` from ``g wo^T`` instead of ``y``). So the gradients agree with + the two ops apart to that rounding, not bit for bit. + + Rows of the sorted buffer outside this shard's experts' groups: the cotangent is left unmasked + there, which saves a pass over ``[rows, hidden]``. The down projection's backward GEMMs do not use + those rows for any row inside a group (the dlhs is row-wise and the drhs masks both operands of a + partial group tile), and every output of this backward is masked to the group rows. + + Args: + intermediate_local: ``[buffer_size, mlp]`` down-projection input, rows grouped by local expert. + wo_local: the down-projection kernel of the local experts, as ``down_projection`` takes it. + down_projection: ``(h, wo, *down_projection_operands) -> [buffer_size, hidden]``, a GEMM with its + own VJP. It must not close over traced arrays; pass them in ``down_projection_operands``. + down_projection_operands: tuple of non-differentiated operands of ``down_projection`` (group + sizes and the like). + group_sizes_local, topk_argsort_revert_indices, topk, local_num_experts, ep_name, topk_weights, + topk_argsort_indices: as for :func:`ring_ragged_unsort`. + forward_tie: optional array tied to ``y`` by an optimization barrier in the forward pass only, + so everything that reads either comes after both. It is returned as a second output and its + cotangent passes through unchanged. + index_math, unique_pair_argsort: as for :func:`ring_ragged_unsort` (they select how the combine's + gather-reduce sorts its valid rows, and with "count" how the weight gradient returns to slot order). + + Returns: + The unsorted ``[num_tokens_local, hidden]`` output, and ``forward_tie`` if one was given. + """ + num_slots = topk_argsort_revert_indices.shape[0] + if intermediate_local.shape[0] > num_slots: + # ring_ragged_sort's buffer has exactly num_slots rows (full mode) or fewer (packed mode). + raise ValueError( + f"ring_ragged_unsort_down_projection: {intermediate_local.shape[0]} buffer rows > {num_slots} slots." + ) + gather_kwargs, gather_reduce_kwargs = _unsort_kwargs( + enforce_gather_fallback, + enforce_gather_reduce_fallback, + gather_flops_override, + gather_reduce_flops_override, + gather_bytes_accessed_override, + gather_reduce_bytes_accessed_override, + use_single_sparsecore, + index_math, + unique_pair_argsort, + ) + + @jax.custom_vjp + def _unsort_down_projection(h, wo, operands, group_sizes, revert, weights, argsort_indices, tie): + """Down projection, then unsort and combine.""" + return _unsort_down_projection_fwd(h, wo, operands, group_sizes, revert, weights, argsort_indices, tie)[0] + + @jax.named_scope("ragged-unsort-down-projection-fwd") + def _unsort_down_projection_fwd(h, wo, operands, group_sizes, revert, weights, argsort_indices, tie): + y = down_projection(h, wo, *operands) + if tie is not None: + y, tie = jax.lax.optimization_barrier((y, tie)) + start, end = _ring_unsort_window(group_sizes, local_num_experts, ep_name) + out = _ring_unsort_combine(y, revert, weights, start, end, topk, gather_reduce_kwargs) + # y is not a residual: the backward below never reads it. + res = (h, wo, operands, revert, weights, start, end, argsort_indices) + return (out if tie is None else (out, tie)), res + + @jax.named_scope("ragged-unsort-down-projection-bwd") + def _unsort_down_projection_bwd(res, g_out): + h, wo, operands, revert, weights, start, end, argsort_indices = res + g_unsorted, g_tie = (g_out, None) if forward_tie is None else g_out + gathered, weight_for_sorted, mask, to_slot_order = _ring_unsort_fan_out( + g_unsorted, + revert, + weights, + start, + end, + h.shape[0], + topk, + argsort_indices, + gather_kwargs, + mask_rows=False, + index_math=index_math, + ) + weight = weight_for_sorted[:, None] + h32 = h.astype(jnp.float32) + # The routing weight goes on the narrow [rows, mlp] side: dwo = (w * h)^T g. + weighted_h = jnp.where(mask[:, None], h32 * weight, 0.0).astype(h.dtype) + # The down projection's own VJP at the input w * h: its input gradient g wo^T does not depend on the + # input, and its weight gradient is (w * h)^T g. Its forward GEMM has no reader, so XLA removes it. + _, project_vjp = jax.vjp(lambda h_, wo_: down_projection(h_, wo_, *operands), weighted_h, wo) + u, grad_wo = project_vjp(gathered) + u = u.astype(jnp.float32) + grad_h = jnp.where(mask[:, None], u * weight, 0.0).astype(h.dtype) + # The routing-weight gradient = , from the down projection's input. + dot = jnp.where(mask, jnp.sum(u * h32, axis=-1), 0.0) + grad_topk_weights = to_slot_order(dot) + grad_operands = jax.tree.map(lambda _: None, operands) + return grad_h, grad_wo, grad_operands, None, None, grad_topk_weights, None, g_tie + + _unsort_down_projection.defvjp(_unsort_down_projection_fwd, _unsort_down_projection_bwd) + + return _unsort_down_projection( + intermediate_local, + wo_local, + tuple(down_projection_operands), + group_sizes_local, + topk_argsort_revert_indices, + topk_weights.astype(jnp.float32), + topk_argsort_indices, + forward_tie, ) diff --git a/src/maxtext/kernels/ragged/ragged_sort_tc.py b/src/maxtext/kernels/ragged/ragged_sort_tc.py new file mode 100644 index 0000000000..6581919b2e --- /dev/null +++ b/src/maxtext/kernels/ragged/ragged_sort_tc.py @@ -0,0 +1,437 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Ring-of-experts ragged sort / unsort on TensorCore. + +Drop-in alternative to `ring_ragged_sort` / `ring_ragged_unsort` (truncated-buffer +case) that replaces the SparseCore ragged_gather / ragged_gather_reduce kernels +with the TensorCore DMA kernels in `ragged_gather_tc.py`. + +Layout: token rows are reshaped (N, D) -> (N, D // 128, 128) around the kernels +so each token is one contiguous row DMA. + +Routing (shared by sort and unsort of one MoE chunk): + token_ids_sorted[j]: source token of the j-th slot in global expert order. + slot_order[j]: flat (token * topk + k) slot of the j-th sorted slot. + start, count: this shard's window [start, start + count) of sorted slots, + count clamped to the buffer size. + reduce_meta: routing of the TC gather-reduce, derived once. + +Forward sort: buf[i] = x[token_ids_sorted[start + i]], i < count (else 0). +Forward unsort: y[t] = sum_{i < count, token_ids_sorted[start+i] == t} w_i * buf[i]. +Each op's transpose is the other op, implemented with the other TC kernel. + +keep_3d (moe_tc_ragged_3d_gmm): the sorted buffer stays in the kernels' 3D layout +(cap, D // 128, 128) on both sides (sort output / unsort input and their +cotangents), for GMM kernels that consume / produce that layout directly. As in +the 2D path, mask_padding zeros the buffer rows past `count` (the sort output and +the unsort-input cotangent). Without it they are left uninitialized: the 3D +gmm_v2 / tgmm_v2 consumers only use rows < sum(group_sizes) == count, but +whole-tensor statistics (e.g. absmax quantization of the gmm lhs or of the +backward gradient) would see the uninitialized rows. +""" + +import dataclasses +import functools + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +from maxtext.kernels.ragged import ragged_gather_tc as tc + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class TcRouting: + token_ids_sorted: jax.Array + slot_order: jax.Array + start: jax.Array + count: jax.Array + reduce_meta: tc.RaggedGatherReduceMetadata + num_out_tokens: int = jax.tree.static() + # slot_order / token_ids_sorted hold only this shard's [cap] window (start == 0), see `window_routing`. + windowed: bool = jax.tree.static() + + +def _to3d(x): + n, d = x.shape + if d % 128: + raise ValueError(f"TC ragged sort needs hidden % 128 == 0, got {d}.") + return x.reshape(n, d // 128, 128) + + +def _to2d(x): + return x.reshape(x.shape[0], x.shape[1] * x.shape[2]) + + +# --- Pallas relayout (N, D) <-> (N, D // 128, 128), after lineage ragged_flatten --- +# +# XLA lowers these reshapes to full-size layout copies. The kernels reshape in +# VMEM and only visit blocks overlapping the first `num_rows` rows; rows past +# that are left uninitialized unless `zero_fill`. + + +def _manual_axis_type(x): + return getattr(jax.typeof(x), "manual_axis_type", None) + + +def _flatten_kernel(n_ref, x_ref, o_ref, *, block, zero_fill): + """Relayouts one (block, D0, D1) block to (block, D0 * D1); with zero_fill, rows >= n are zeroed.""" + if not zero_fill: + o_ref[...] = x_ref[...].reshape(o_ref.shape) + return + start = pl.program_id(0) * block + n = n_ref[0] + + @pl.when(start + block <= n) + def _(): + o_ref[...] = x_ref[...].reshape(o_ref.shape) + + @pl.when(start >= n) + def _(): + o_ref[...] = jnp.zeros(o_ref.shape, o_ref.dtype) + + @pl.when(jnp.logical_and(start < n, start + block > n)) + def _(): + y = x_ref[...].reshape(o_ref.shape) + rows = start + jax.lax.broadcasted_iota(jnp.int32, o_ref.shape, 0) + o_ref[...] = jnp.where(rows < n, y, jnp.zeros((), y.dtype)) + + +def _tc_flatten(x, num_rows, block, zero_fill=False): + """(R, D0, D1) -> (R, D0 * D1) over the first num_rows rows.""" + rows, d0, d1 = x.shape + d = d0 * d1 + block = min(block, rows) + n = jnp.reshape(jnp.asarray(num_rows, jnp.int32), (1,)) + if zero_fill: + # Every block is written (zeros past num_rows); inputs past num_rows are + # clamped to the last valid block so they are not re-fetched. + grid = (pl.cdiv(rows, block),) + + def in_map(i, n_ref): + return (jnp.minimum(i, jnp.maximum(pl.cdiv(n_ref[0], block) - 1, 0)), 0, 0) + + else: + grid = (pl.cdiv(n[0], block),) + + def in_map(i, n_ref): + del n_ref + return (i, 0, 0) + + return pl.pallas_call( + functools.partial(_flatten_kernel, block=block, zero_fill=zero_fill), + out_shape=jax.ShapeDtypeStruct((rows, d), x.dtype, manual_axis_type=_manual_axis_type(x)), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=1, + grid=grid, + in_specs=[pl.BlockSpec((block, d0, d1), in_map)], + out_specs=pl.BlockSpec((block, d), lambda i, n_ref: (i, 0)), + ), + name="ragged_flatten_tc", + )(n, x) + + +def _unflatten_kernel(n_ref, x_ref, o_ref): + del n_ref + o_ref[...] = x_ref[...].reshape(o_ref.shape) + + +def _tc_unflatten(x, num_rows, block): + """(R, D) -> (R, D // 128, 128) over the first num_rows rows.""" + rows, d = x.shape + if d % 128: + raise ValueError(f"TC ragged sort needs hidden % 128 == 0, got {d}.") + d0, d1 = d // 128, 128 + block = min(block, rows) + n = jnp.reshape(jnp.asarray(num_rows, jnp.int32), (1,)) + return pl.pallas_call( + _unflatten_kernel, + out_shape=jax.ShapeDtypeStruct((rows, d0, d1), x.dtype, manual_axis_type=_manual_axis_type(x)), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=1, + grid=(pl.cdiv(n[0], block),), + in_specs=[pl.BlockSpec((block, d), lambda i, n_ref: (i, 0))], + out_specs=pl.BlockSpec((block, d0, d1), lambda i, n_ref: (i, 0, 0)), + ), + name="ragged_unflatten_tc", + )(n, x) + + +def _keep_3d(blocks): + return len(blocks) > 2 and bool(blocks[2]) + + +def _tokens_3d(blocks): + """Token-side tensors (sort input / unsort output) are already in the (N, D // 128, 128) layout.""" + return len(blocks) > 3 and bool(blocks[3]) + + +def _gather(x2d, routing, cap, blocks, mask_padding): + """TC ragged gather of the shard's sorted slots into a [cap, D] buffer ([cap, D // 128, 128] with keep_3d).""" + block_size, flat_block = blocks[:2] + if x2d.ndim == 3: + x3d = x2d + elif flat_block: + x3d = _tc_unflatten(x2d, x2d.shape[0], flat_block) + else: + x3d = _to3d(x2d) + out = tc.ragged_gather_tc( + x3d, + routing.token_ids_sorted, + start=routing.start, + num_tokens=routing.count, + max_out_tokens=cap, + block_size=block_size, + ) + if _keep_3d(blocks): + if mask_padding: + valid = jnp.arange(cap, dtype=jnp.int32) < routing.count + out = jnp.where(valid[:, None, None], out, jnp.zeros((), out.dtype)) + return out + if flat_block: + return _tc_flatten(out, routing.count, flat_block, zero_fill=mask_padding) + out = _to2d(out) + if mask_padding: + valid = jnp.arange(cap, dtype=jnp.int32) < routing.count + out = jnp.where(valid[:, None], out, jnp.zeros((), out.dtype)) + return out + + +def _gather_reduce(buf2d, routing, num_out_tokens, topk, blocks): + """TC ragged gather-reduce: sums each token's valid buffer rows into a [num_out_tokens, D] output.""" + flat_block = blocks[1] + if buf2d.ndim == 3: + buf3d = buf2d + elif flat_block: + buf3d = _tc_unflatten(buf2d, routing.count, flat_block) + else: + buf3d = _to3d(buf2d) + out = tc.ragged_gather_reduce_tc( + buf3d, + routing.reduce_meta, + num_out_tokens=num_out_tokens, + top_k=topk, + ) + if _tokens_3d(blocks): + return out + if flat_block: + return _tc_flatten(out, num_out_tokens, flat_block) + return _to2d(out) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(2, 3, 4, 5)) +def tc_sort_gather(x2d, routing, cap, topk, blocks, mask_padding): + """buf[i] = x[token_ids_sorted[start + i]] for i < count, zero-padded to cap.""" + return _gather(x2d, routing, cap, blocks, mask_padding) + + +def _tc_sort_gather_fwd(x2d, routing, cap, topk, blocks, mask_padding): + return _gather(x2d, routing, cap, blocks, mask_padding), (routing, x2d.shape[0]) + + +def _tc_sort_gather_bwd(cap, topk, blocks, mask_padding, res, g): + del cap, mask_padding + routing, n_tokens = res + return _gather_reduce(g, routing, n_tokens, topk, blocks), None + + +tc_sort_gather.defvjp(_tc_sort_gather_fwd, _tc_sort_gather_bwd) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(2, 3, 4, 5)) +def tc_unsort_reduce(buf2d, routing, num_out_tokens, topk, blocks, mask_padding): + """y[t] = sum of valid buffer rows whose source token is t.""" + return _gather_reduce(buf2d, routing, num_out_tokens, topk, blocks) + + +def _tc_unsort_reduce_fwd(buf2d, routing, num_out_tokens, topk, blocks, mask_padding): + return _gather_reduce(buf2d, routing, num_out_tokens, topk, blocks), (routing, buf2d.shape[0]) + + +def _tc_unsort_reduce_bwd(num_out_tokens, topk, blocks, mask_padding, res, g): + del num_out_tokens, topk + routing, cap = res + return _gather(g, routing, cap, blocks, mask_padding), None + + +tc_unsort_reduce.defvjp(_tc_unsort_reduce_fwd, _tc_unsort_reduce_bwd) + + +def ring_ragged_sort_tc( + hidden_states_local, + topk_indices_local, + num_experts, + topk, + ep_name, + ep_size, + buffer_size, + gather_block_size=1024, + reduce_block_size=896, + mask_padding=True, + flatten_block_size=0, + tag_routing_fn=None, + keep_3d=False, + window_routing=False, +): + """TC version of `ring_ragged_sort` for the truncated-buffer case. + + `tag_routing_fn` (as in `ring_ragged_sort`, e.g. a checkpoint_name tagger for moe_routing_maps) is applied to + each integer routing array (the TcRouting leaves, group sizes and the inverse permutation), so they can be + saved across the remat boundary instead of recomputed in the backward pass. + + With `window_routing` (moe_tc_window_routing), the expert sort is a packed uint32 sort, the routing keeps only + this shard's [start, start + cap) window of the sorted slots (so the reduce metadata sorts cap instead of + N * topk slots), and the inverse permutation is not computed: the third output is then the forward + permutation, which TC callers only use for its shape. + + Returns: + (sorted buffer [buffer_size, hidden], group_sizes [num_experts], + topk_argsort_revert_indices [N * topk], TcRouting). + """ + num_tokens_local = hidden_states_local.shape[0] + topk_indices_flat = topk_indices_local.flatten().astype(jnp.int32) + n = topk_indices_flat.shape[0] + if window_routing: + _, topk_argsort_indices = tc.packed_stable_sort(topk_indices_flat, num_experts - 1) + else: + topk_argsort_indices = jnp.argsort(topk_indices_flat, stable=True).astype(jnp.int32) + group_sizes_local = jax.nn.one_hot(topk_indices_flat, num_experts, dtype=jnp.int32).sum(axis=0) + + shard_idx = jax.lax.axis_index(ep_name) if ep_size > 1 else 0 + local_num_experts = num_experts // ep_size + group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True) + start = group_offsets[shard_idx * local_num_experts].astype(jnp.int32) + end = group_offsets[(shard_idx + 1) * local_num_experts].astype(jnp.int32) + cap = n if buffer_size is None else min(buffer_size, n) + count = jnp.minimum(end - start, cap).astype(jnp.int32) + + if window_routing: + topk_argsort_revert_indices = topk_argsort_indices + slot_order = jax.lax.dynamic_slice_in_dim(jnp.pad(topk_argsort_indices, (0, cap)), start, cap, axis=0) + token_ids_sorted = jax.lax.stop_gradient(slot_order // topk) + reduce_meta = tc.ragged_gather_reduce_tc_metadata( + token_ids_sorted, + 0, + count, + top_k=topk, + block_size=reduce_block_size, + num_out_tokens=num_tokens_local, + packed_sort=True, + ) + start = jnp.int32(0) + else: + topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices).astype(jnp.int32) + slot_order = topk_argsort_indices + token_ids_sorted = jax.lax.stop_gradient(topk_argsort_indices // topk) + reduce_meta = tc.ragged_gather_reduce_tc_metadata( + token_ids_sorted, start, count, top_k=topk, block_size=reduce_block_size + ) + routing = TcRouting( + token_ids_sorted=token_ids_sorted, + slot_order=slot_order, + start=start, + count=count, + reduce_meta=reduce_meta, + num_out_tokens=num_tokens_local, + windowed=window_routing, + ) + if tag_routing_fn is not None: + routing = jax.tree.map(tag_routing_fn, routing) + group_sizes_local = tag_routing_fn(group_sizes_local) + topk_argsort_revert_indices = tag_routing_fn(topk_argsort_revert_indices) + # A 3D (N, D // 128, 128) input (moe_tc_ragged_3d_dispatch) is gathered without a relayout, and its + # cotangent is produced in the same layout. + tokens_3d = hidden_states_local.ndim == 3 + x = tc_sort_gather( + hidden_states_local, + routing, + cap, + topk, + (gather_block_size, flatten_block_size, keep_3d, tokens_3d), + mask_padding, + ) + return x, group_sizes_local, topk_argsort_revert_indices, routing + + +def tc_buffer_row_weights(routing, topk_weights, cap): + """Routing weight of each buffer row: w[slot_order[start + i]] for i < count, else 0. + + Args: + routing: TcRouting from `ring_ragged_sort_tc`. + topk_weights: flat [N * topk] routing weights (differentiated). + cap: buffer rows. + + Returns: + [cap] float32 weights. + """ + w_flat = topk_weights.astype(jnp.float32) + iota = jnp.arange(cap, dtype=jnp.int32) + valid = iota < routing.count + if routing.windowed: + # Valid rows are distinct slots and padding rows get distinct out-of-range indices (filled with 0), so the + # gather is unique_indices and its transpose needs no scatter-sort. + rows = jnp.where(valid, routing.slot_order, w_flat.shape[0] + iota) + return w_flat.at[rows].get(mode="fill", fill_value=0.0, unique_indices=True) + padded_order = jnp.pad(routing.slot_order, (0, cap)) + rows = jax.lax.dynamic_slice_in_dim(padded_order, routing.start, cap, axis=0) + return jnp.where(valid, jnp.take(w_flat, jnp.where(valid, rows, 0), axis=0), 0.0) + + +def ring_ragged_unsort_tc( + sorted_tokens_local, + routing, + topk, + topk_weights, + gather_block_size=1024, + mask_padding=True, + flatten_block_size=0, + prescaled=False, + out_3d=False, +): + """TC version of `ring_ragged_unsort`: weighted top-k combine of this shard's slots. + + Args: + sorted_tokens_local: [buffer_size, hidden] expert outputs, or [buffer_size, hidden // 128, 128] + (moe_tc_ragged_3d_gmm), which is fed to the TC gather-reduce without a relayout. + routing: TcRouting from `ring_ragged_sort_tc`. + topk: routing top-k. + topk_weights: flat [N * topk] routing weights (differentiated). Unused if prescaled. + prescaled: rows are already multiplied by their routing weight (e.g. on the expert + activation, see `tc_buffer_row_weights`), so only the unweighted combine is done. + out_3d: return the combine in the TC kernels' (N, hidden // 128, 128) layout (no relayout); + the cotangent is then expected in that layout too. + + Returns: + [N, hidden] (or [N, hidden // 128, 128] if out_3d) partial combine for this shard's experts. + """ + cap = sorted_tokens_local.shape[0] + num_out_tokens = routing.num_out_tokens + if prescaled: + scaled = sorted_tokens_local + else: + w_rows = tc_buffer_row_weights(routing, topk_weights, cap) + w_rows = w_rows.reshape(-1, *([1] * (sorted_tokens_local.ndim - 1))) + scaled = (sorted_tokens_local.astype(jnp.float32) * w_rows).astype(sorted_tokens_local.dtype) + keep_3d = sorted_tokens_local.ndim == 3 + return tc_unsort_reduce( + scaled, + routing, + num_out_tokens, + topk, + (gather_block_size, flatten_block_size, keep_3d, out_3d), + mask_padding, + ) diff --git a/src/maxtext/kernels/tokamax_splash_attention/base.py b/src/maxtext/kernels/tokamax_splash_attention/base.py index df793fcec9..b0adce4265 100644 --- a/src/maxtext/kernels/tokamax_splash_attention/base.py +++ b/src/maxtext/kernels/tokamax_splash_attention/base.py @@ -135,7 +135,7 @@ def _attention_reference_custom_bwd( p = jnp.exp(logits - logsumexp[..., None]) p = jnp.where(mask, p, 0.0) - do = do.astype(jnp.float32) # pytype: disable=attribute-error + do = do.astype(jnp.float32) dv = jnp.einsum("pt,pd->td", p, do).astype(v.dtype) dp = jnp.einsum("pd,td->pt", do, v.astype(jnp.float32)) diff --git a/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py b/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py index 3db9e53a4e..769e85fff4 100644 --- a/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py +++ b/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py @@ -782,7 +782,7 @@ def manual_sharding_spec(self): def mask_info_spec(mask_info): if mask_info is None: return None - return MaskInfo( # pytype: disable=wrong-arg-types,bad-argument-type + return MaskInfo( mask_next=_resolve_spec(mask_info.mask_next), # pyrefly: ignore[bad-argument-type] active_rows=_resolve_spec(mask_info.active_rows), # pyrefly: ignore[bad-argument-type] active_cols=_resolve_spec(mask_info.active_cols), # pyrefly: ignore[bad-argument-type] diff --git a/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py b/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py index 4eec06d3a1..4100d0525c 100644 --- a/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py +++ b/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py @@ -236,7 +236,7 @@ def _apply_mask_and_soft_cap( # pyrefly: ignore[unsupported-operation] q_sequence_ref[...], (1, repeats), - ) # [bq, k_slice.size] # pyrefly: ignore[unsupported-operation] + ) # [bq, k_slice.size] else: assert q_sequence_ref.shape == (NUM_SUBLANES, bq) # pyrefly: ignore[missing-attribute] @@ -251,7 +251,7 @@ def _apply_mask_and_soft_cap( q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq)) assert q_sequence.shape == k_sequence.shape - computed_mask = mask_function(q_sequence, k_sequence) # pytype: disable=wrong-arg-count + computed_mask = mask_function(q_sequence, k_sequence) if computed_mask.dtype != jnp.dtype(jnp.bool_): raise ValueError("Mask function must return a boolean-valued array, but got:" f" {computed_mask.dtype}") masks.append(computed_mask) @@ -386,11 +386,11 @@ def init(): exp = jnp.exp2 if config.use_base2_exp else jnp.exp m_scratch_ref[...] = jnp.full_like(m_scratch_ref, max_logit_estimate) # pyrefly: ignore[bad-argument-type] l_scratch_ref[...] = exp( - # pyrefly: ignore[bad-argument-type, unsupported-operation] + # pyrefly: ignore[unsupported-operation] sink # pyrefly: ignore[bad-argument-type] - jnp.full_like(l_scratch_ref, max_logit_estimate) - ) # pyrefly: ignore[bad-argument-type, unsupported-operation] + ) def body(kv_compute_index, _, has_partial_mask=False): slice_k = pl.ds(kv_compute_index * bkv_compute, bkv_compute) @@ -430,7 +430,7 @@ def body(kv_compute_index, _, has_partial_mask=False): qk, mask = apply_mask_and_soft_cap() if max_logit_estimate is None: - m_curr = qk.max(axis=-1)[:, None] # pytype: disable=attribute-error + m_curr = qk.max(axis=-1)[:, None] assert m_curr.shape == (bq, 1) m_next = jnp.maximum(m_prev, m_curr) assert m_next.shape == (bq, NUM_LANES) @@ -798,7 +798,6 @@ def _fwd_cost_estimate( kv_segment_ids, mask_info.partial_mask_blocks, ] - # pyrefly: ignore[bad-argument-count] cost_estimate = config.fwd_cost_estimate or _fwd_cost_estimate( # pyrefly: ignore[bad-argument-type] *vmem_inputs, @@ -954,7 +953,7 @@ def _splash_attention_custom( # device. del dkv_mask_info - ret = _splash_attention_forward( # pytype: disable=wrong-arg-types + ret = _splash_attention_forward( fwd_mask_info, q, k, @@ -972,8 +971,8 @@ def _splash_attention_custom( if save_residuals: out, stats = ret if config.use_base2_exp: # for user, output values in natural base - stats = _base2_stats_to_natural_base(stats, mask_value) - return out, stats + stats = _base2_stats_to_natural_base(stats, mask_value) # pyrefly: ignore[bad-argument-type] + return out, stats # pyrefly: ignore[bad-return] else: return ret @@ -1000,7 +999,7 @@ def _splash_attention_fwd( # if save_residuals: # raise NotImplementedError("Higher-order AD not supported.") - out, stats = _splash_attention_forward( # pytype: disable=wrong-arg-types + out, stats = _splash_attention_forward( fwd_mask_info, q, k, @@ -1017,7 +1016,7 @@ def _splash_attention_fwd( ) logsumexp = stats["logsumexp"] # save in the config base for the bwd pass if config.use_base2_exp: # for user, output values in natural base - stats = _base2_stats_to_natural_base(stats, mask_value) + stats = _base2_stats_to_natural_base(stats, mask_value) # pyrefly: ignore[bad-argument-type] residuals = q, k, v, segment_ids, sinks, out, logsumexp, dkv_mask_info if save_residuals: return (out, stats), residuals # pyrefly: ignore[bad-return] @@ -1102,7 +1101,7 @@ def body(has_partial_mask: bool = False): q_segment_ids_ref, kv_segment_ids_ref, attn_logits_soft_cap=attn_logits_soft_cap, - k_slice=pl.ds(0, bkv), # pyrefly: ignore[bad-argument-type] + k_slice=pl.ds(0, bkv), k_offset=kv_index * bkv, bq=bq, mask_function=mask_function, @@ -1275,7 +1274,7 @@ def _load_kv(ref, layout): q_segment_ids_ref, kv_segment_ids_ref, attn_logits_soft_cap=attn_logits_soft_cap, - k_slice=slice_k, # pyrefly: ignore[bad-argument-type] + k_slice=slice_k, k_offset=kv_index * bkv + i * bkv_compute, bq=bq, k_in_lanes=False, @@ -1840,7 +1839,7 @@ def _splash_attention_bwd( q, k, v, segment_ids, sinks, o, logsumexp, dkv_mask_info = res # di: [num_heads, q_seq_len] - di = jnp.einsum("hsd,hsd->hs", o.astype(jnp.float32), do.astype(jnp.float32)) # pytype: disable=attribute-error + di = jnp.einsum("hsd,hsd->hs", o.astype(jnp.float32), do.astype(jnp.float32)) # pyrefly: ignore[missing-attribute] dq, dk, dv = _splash_attention_bwd_dkv( q, k, @@ -1971,7 +1970,7 @@ def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding): def mask_info_spec(mask_info): if mask_info is None: return None - return MaskInfo( # pytype: disable=wrong-arg-types + return MaskInfo( mask_next=_resolve_spec(mask_info.mask_next), # pyrefly: ignore[bad-argument-type] active_rows=_resolve_spec(mask_info.active_rows), # pyrefly: ignore[bad-argument-type] active_cols=_resolve_spec(mask_info.active_cols), # pyrefly: ignore[bad-argument-type] @@ -2132,7 +2131,7 @@ def process_mask_shard(mask): kernel = SplashAttentionKernel(fwd_mask_info, dkv_mask_info, **kwargs) return kernel - mask_info_specs = MaskInfo( # pytype: disable=wrong-arg-types + mask_info_specs = MaskInfo( mask_next=mask_spec, # pyrefly: ignore[bad-argument-type] active_rows=None, active_cols=None, diff --git a/src/maxtext/layers/attention_compressed.py b/src/maxtext/layers/attention_compressed.py index 2653fe1fbb..f7625af7dd 100644 --- a/src/maxtext/layers/attention_compressed.py +++ b/src/maxtext/layers/attention_compressed.py @@ -38,11 +38,10 @@ DEFAULT_MASK_VALUE, ) -from maxtext.layers import nnx_wrappers from maxtext.layers.attention_mla import indexer_losses from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import DeepSeekV4RotaryEmbedding -from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned +from maxtext.layers.initializers import nd_dense_init, NdInitializer from maxtext.layers.linears import DenseGeneral, DeepSeekV4GroupedLinear from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant @@ -1463,11 +1462,18 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No # Override the base rotary embedding with the correct theta for this layer. # CSA / HCA layers use compressed_rope_max_timescale (160000). # Sliding window prefix layers use rope_max_timescale (10000). - rope_theta = self.config.compressed_rope_max_timescale if self.compress_ratio > 0 else self.config.rope_max_timescale + is_compressed = self.compress_ratio > 0 + rope_theta = self.config.compressed_rope_max_timescale if is_compressed else self.config.rope_max_timescale self.rotary_embedding = DeepSeekV4RotaryEmbedding( head_dim=self.config.head_dim, partial_rotary_factor=self.config.qk_rope_head_dim / self.config.head_dim, rope_theta=rope_theta, + rope_type=self.config.rope_type if is_compressed else "default", + rope_factor=self.config.rope_factor, + beta_fast=self.config.beta_fast, + beta_slow=self.config.beta_slow, + original_max_position_embeddings=self.config.original_max_position_embeddings, + truncate=self.config.rope_truncate, fprop_dtype=self.dtype, ) @@ -1816,12 +1822,18 @@ def __call__( # Tokamax dynamic splash tile boundary alignment. Note: Tokamax kernel inside AttentionOp additionally # sets inner block size as min(block_kv, key_len) during kernel invocation. if self.attention_kernel == "flash": + if max_utils.is_eval(self.config): + block_q = self.config.eval_sa_block_q + block_kv = self.config.eval_sa_block_kv + else: + block_q = self.config.sa_block_q + block_kv = self.config.sa_block_kv comp_len = compressed_kv.shape[1] if compressed_kv is not None else 0 _, pad_kv_total = compute_hca_padding( q_len=inputs_q.shape[1], comp_len=comp_len, - block_q=self.config.sa_block_q, - block_kv=self.config.sa_block_kv, + block_q=block_q, + block_kv=block_kv, kv_len=kv.shape[1], ) @@ -2093,70 +2105,3 @@ def scan_body_heads(carry, xs): indexer_loss = (jnp.sum(kl_per_token) / num_valid_tokens) * scaling_factor return indexer_loss - - -def compressed_attention( - *, - config: Config, - num_query_heads: int, - num_kv_heads: int, - head_dim: int, - max_target_length: int, - mesh: Mesh, - attention_kernel: str, - inputs_q_shape: Tuple, - inputs_kv_shape: Tuple, - dtype: DType = jnp.float32, - weight_dtype: DType = jnp.float32, - max_prefill_predict_length: int = -1, - dropout_rate: float = 0.0, - kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "normal"), - float32_qk_product: bool = False, - float32_logits: bool = False, - quant: Optional[Quant] = None, - kv_quant: Optional[KVQuant] = None, - attention_type: AttentionType = AttentionType.COMPRESSED, - attn_logits_soft_cap: float | None = None, - sliding_window_size: int | None = None, - use_ragged_attention: bool = False, - ragged_block_size: int = 256, - use_qk_norm: bool = False, - query_pre_attn_scalar: float | None = None, - use_bias_in_projections: bool = False, - q_lora_rank: int = 1536, - name: str | None = None, -): - """Wrapper to create the CompressedAttention linen module.""" - return nnx_wrappers.to_linen( - CompressedAttention, - config=config, - num_query_heads=num_query_heads, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - max_target_length=max_target_length, - mesh=mesh, - attention_kernel=attention_kernel, - inputs_q_shape=inputs_q_shape, - inputs_kv_shape=inputs_kv_shape, - dtype=dtype, - weight_dtype=weight_dtype, - max_prefill_predict_length=max_prefill_predict_length, - dropout_rate=dropout_rate, - kernel_init=kernel_init, - float32_qk_product=float32_qk_product, - float32_logits=float32_logits, - quant=quant, - kv_quant=kv_quant, - attention_type=attention_type, - attn_logits_soft_cap=attn_logits_soft_cap, - sliding_window_size=sliding_window_size, - use_ragged_attention=use_ragged_attention, - ragged_block_size=ragged_block_size, - use_qk_norm=use_qk_norm, - query_pre_attn_scalar=query_pre_attn_scalar, - use_bias_in_projections=use_bias_in_projections, - q_lora_rank=q_lora_rank, - name=name, - metadata_fn=variable_to_logically_partitioned, - abstract_init=False, - ) diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 05ee1e2e38..1f6d40008c 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -530,6 +530,7 @@ def __init__( base_kv_cache = config.attention != "paged" and config.mla_naive_kvcache # Setting these before call to super because a field is used in super + self.use_absorbed_mqa = getattr(config, "use_mla_absorbed_mqa", False) self.q_lora_rank = q_lora_rank self.kv_lora_rank = kv_lora_rank self.qk_nope_head_dim = qk_nope_head_dim @@ -758,10 +759,55 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No self.out = self.init_out_w(output_dim=inputs_q_shape[-1]) + def _attention_op_num_kv_heads(self) -> int: + if self.use_absorbed_mqa: + return 1 + return self.num_kv_heads + @property def out_head_dim(self) -> int: return self.v_head_dim + def _absorbed_weights(self, dtype: DType) -> tuple[Array, Array]: + """Splits wkv_b kernel cleanly into W_UK and W_UV with logical sharding.""" + w_uk, w_uv = jnp.split(self.wkv_b.kernel.astype(dtype), [self.qk_nope_head_dim], axis=-1) + w_uk = self._maybe_shard_with_logical(w_uk, ("kv_lora", "q_heads", None)) + w_uv = self._maybe_shard_with_logical(w_uv, ("kv_lora", "q_heads", None)) + return w_uk, w_uv + + def absorb_query(self, q_nope: Array, q_pe: Array, w_uk: Optional[Array] = None) -> Array: + """Absorbs W_UK into Query: Q_absorbed = Q_nope @ W_UK -> [B, T, H, D_c + D_rope].""" + if w_uk is None: + w_uk, _ = self._absorbed_weights(q_nope.dtype) + q_absorbed = jnp.einsum("bthd, chd -> bthc", q_nope, w_uk, precision=self.config.matmul_precision) + q_absorbed = q_absorbed.astype(q_nope.dtype) + q_absorbed = self._maybe_shard_with_logical( + q_absorbed, ("activation_batch", "activation_length", "activation_heads", None) + ) + return jnp.concatenate([q_absorbed, q_pe], axis=-1) + + def get_kv_latent(self, low_rank_main: Array, key_rope: Array) -> tuple[Array, Array]: + """Constructs 1-head compressed latent KV: [B, S, 1, D_c + D_rope] and value [B, S, 1, D_c].""" + c_kv = low_rank_main[:, :, None, :] + key = jnp.concatenate([c_kv, key_rope], axis=-1) + return key, c_kv + + def absorb_output(self, out_latent: Array, w_uv: Optional[Array] = None) -> Array: + """Absorbs W_UV into Output: O = O_latent @ W_UV -> [B, T, H, D_v].""" + if w_uv is None: + _, w_uv = self._absorbed_weights(out_latent.dtype) + out = jnp.einsum("bthc, chv -> bthv", out_latent, w_uv, precision=self.config.matmul_precision) + out = out.astype(out_latent.dtype) + return self._maybe_shard_with_logical( + out, ("activation_batch", "activation_length", "activation_heads", "activation_kv") + ) + + def _build_key_value(self, low_rank_main: Array, key_rope: Array, model_mode: str) -> tuple[Array, Array]: + """Dispatches key and value construction based on absorption mode.""" + if self.use_absorbed_mqa: + return self.get_kv_latent(low_rank_main, key_rope) + return self.mla_get_key_value(low_rank_main, key_rope, model_mode) + def mla_query_projection( self, inputs_q: Array, inputs_positions: Array, model_mode ) -> tuple[jax.Array, Optional[jax.Array]]: @@ -816,7 +862,11 @@ def mla_query_projection( q_pe = self._maybe_shard_with_logical(q_pe, query_logical_name) # Query projection is scaled by self.softmax_scale to be consistent MaxText implementation. # DeepSeek v3 was doing it in attention score computation. - query = jnp.concatenate([q_nope, q_pe], axis=-1) * self.softmax_scale + if self.use_absorbed_mqa: + query = self.absorb_query(q_nope, q_pe) * self.softmax_scale + query_logical_name = query_logical_name[:-1] + (None,) + else: + query = jnp.concatenate([q_nope, q_pe], axis=-1) * self.softmax_scale if self.config.experimental_sa_quant_q_fp8: query = query.astype(jnp.float8_e4m3fn) @@ -937,14 +987,14 @@ def update_mla_kv_caches(self, low_rank_main, key_rope, decoder_segment_ids, mod if prefill_mla_cache: low_rank_main, key_rope, decoder_segment_ids = prefill_mla_cache - key, value = self.mla_get_key_value(low_rank_main, key_rope, model_mode) + key, value = self._build_key_value(low_rank_main, key_rope, model_mode) prefill_kv_cache = key, value, decoder_segment_ids else: prefill_kv_cache = None if ar_mla_cache: low_rank_main, key_rope, decoder_segment_ids, lengths = ar_mla_cache - key, value = self.mla_get_key_value(low_rank_main, key_rope, model_mode) + key, value = self._build_key_value(low_rank_main, key_rope, model_mode) ar_kv_cache = key, value, decoder_segment_ids, lengths else: ar_kv_cache = None @@ -966,7 +1016,7 @@ def mla_kv_projection(self, inputs: Array, inputs_positions: Array, decoder_segm key_rope = jnp.expand_dims(low_rank_rope, axis=2) key_rope = self.apply_rotary_embedding(key_rope, inputs_positions=inputs_positions) - key, value = self.mla_get_key_value(low_rank_main, key_rope, model_mode) + key, value = self._build_key_value(low_rank_main, key_rope, model_mode) cached_values = [None, None] if self.config.attention != "paged" and model_mode != MODEL_MODE_TRAIN: if self.config.mla_naive_kvcache: @@ -1028,6 +1078,12 @@ def calculate_indexer_loss( indexer_probs = jax.nn.softmax(indexer_score.astype(jnp.float32), axis=-1) batch, q_len, heads, dim = query.shape + # is_1_kv_head is True when: + # 1. use_mla_absorbed_mqa is active: the key is the 1-head compressed latent [B, S, 1, D_c + D_rope]. + # 2. Any single-KV-head configuration (e.g. num_query_heads == 1). + # In these cases, contracting query against squeezed key [B, S, D] directly ("bthd, bsd -> bhts") + # avoids materializing a broadcasted [B, S, H, D] key tensor in HBM. + is_1_kv_head = key.shape[2] == 1 # Chunk across the 'heads' dimension manually using jax.lax.scan # Control the HBM footprint of QK tensor: [batch, q_len, s_len, heads] @@ -1040,19 +1096,32 @@ def calculate_indexer_loss( # Transpose and reshape to put chunk dimension first for jax.lax.scan # query: [b, t, h, d] -> [h, b, t, d] -> [num_chunks, head_chunk_size, b, t, d] q_h = query.transpose(2, 0, 1, 3).reshape(num_chunks, head_chunk_size, batch, q_len, dim) - k_h = key.transpose(2, 0, 1, 3).reshape(num_chunks, head_chunk_size, batch, key.shape[1], dim) + + if is_1_kv_head: + k_squeezed = key.squeeze(2) # [b, s, d] + scan_inputs = {"q": q_h} + else: + k_h = key.transpose(2, 0, 1, 3).reshape(num_chunks, head_chunk_size, batch, key.shape[1], dim) + scan_inputs = {"q": q_h, "k": k_h} def scan_body_heads(carry, xs): q_c = xs["q"] # [h_chunk, b, t, d] - k_c = xs["k"] # [h_chunk, b, s, d] - - # Directly use the chunked shapes in einsum to avoid transposes inside the loop - attn_chunk = jnp.einsum( - "hbtd, hbsd -> bhts", - q_c, - k_c, - precision=self.config.matmul_precision, - ) + if is_1_kv_head: + attn_chunk = jnp.einsum( + "hbtd, bsd -> bhts", + q_c, + k_squeezed, + precision=self.config.matmul_precision, + ) + else: + k_c = xs["k"] # [h_chunk, b, s, d] + # Directly use the chunked shapes in einsum to avoid transposes inside the loop + attn_chunk = jnp.einsum( + "hbtd, hbsd -> bhts", + q_c, + k_c, + precision=self.config.matmul_precision, + ) if sparse_loss: attn_chunk = attn_chunk + indexer_mask[:, None, :, :] @@ -1065,16 +1134,24 @@ def scan_body_heads(carry, xs): return carry + probs_chunk_sum, None init_probs = jnp.zeros((batch, q_len, key.shape[1]), dtype=jnp.float32) - attention_probs, _ = jax.lax.scan(scan_body_heads, init_probs, {"q": q_h, "k": k_h}) + attention_probs, _ = jax.lax.scan(scan_body_heads, init_probs, scan_inputs) else: # Native implementation (default) if chunking is disabled - attention_scores = jnp.einsum( - "bthd, bshd -> bhts", - query, - key, - precision=self.config.matmul_precision, - ) + if is_1_kv_head: + attention_scores = jnp.einsum( + "bthd, bsd -> bhts", + query, + key.squeeze(2), + precision=self.config.matmul_precision, + ) + else: + attention_scores = jnp.einsum( + "bthd, bshd -> bhts", + query, + key, + precision=self.config.matmul_precision, + ) if sparse_loss: attention_scores = attention_scores + indexer_mask[:, None, :, :] elif attention_mask is not None: @@ -1205,6 +1282,9 @@ def __call__( record_max_logits=use_qk_clip, ) + if self.use_absorbed_mqa: + out = self.absorb_output(out) + out = self._maybe_shard_with_logical(out, self.out_axis_names) out = jax.ad_checkpoint.checkpoint_name(out, "attention_out") diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 0c6d86cd9f..3b9f319aa3 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -22,7 +22,6 @@ from flax import linen as nn from flax import nnx -from flax.linen import partitioning import jax from jax import lax from jax.ad_checkpoint import checkpoint_name @@ -73,7 +72,14 @@ from maxtext.layers import nnx_wrappers from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils -from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec, get_logical_axis_rules, mesh_axes_for_dim +from maxtext.utils.sharding import ( + get_logical_axis_rules, + logical_to_mesh_axes, + logical_to_spec, + maybe_shard_with_pspec, + mesh_axes_for_dim, + with_logical_constraint, +) import numpy as np from tokamax._src.ops.attention import base as tokamax_attention_base from tokamax._src.ops.attention import pallas_triton as tokamax_pallas_triton @@ -623,6 +629,8 @@ def __init__( self.max_target_length = max_target_length self.num_query_heads = num_query_heads self.num_kv_heads = num_kv_heads + self.attention_type = _resolve_attention_type(self.config, attention_type) + self.is_absorbed_mqa = getattr(config, "use_mla_absorbed_mqa", False) and self.attention_type == AttentionType.MLA self.float32_qk_product = float32_qk_product self.max_prefill_predict_length = max_prefill_predict_length self.float32_logits = float32_logits @@ -641,7 +649,6 @@ def __init__( self.dtype = dtype self.quant = quant self.kv_quant = kv_quant - self.attention_type = _resolve_attention_type(self.config, attention_type) self.causal_block_size = getattr(self.config, "causal_block_size", None) if self.attention_type == AttentionType.BLOCK_DIFFUSION: if self.causal_block_size is None or self.causal_block_size <= 0: @@ -662,6 +669,9 @@ def __init__( self.block_q = self.config.local_sa_block_q self.block_kv = self.config.local_sa_block_kv self.block_kv_compute = self.config.local_sa_block_kv_compute + self.eval_block_q = self.config.eval_local_sa_block_q + self.eval_block_kv = self.config.eval_local_sa_block_kv + self.eval_block_kv_compute = self.config.eval_local_sa_block_kv_compute self.block_q_dkv = self.config.local_sa_block_q_dkv self.block_kv_dkv = self.config.local_sa_block_kv_dkv self.block_kv_dkv_compute = self.config.local_sa_block_kv_dkv_compute @@ -671,13 +681,22 @@ def __init__( self.q_layout = self.config.local_sa_q_layout self.k_layout = self.config.local_sa_k_layout self.v_layout = self.config.local_sa_v_layout + self.eval_q_layout = self.config.eval_local_sa_q_layout + self.eval_k_layout = self.config.eval_local_sa_k_layout + self.eval_v_layout = self.config.eval_local_sa_v_layout self.use_splash_scheduler = self.config.local_use_splash_scheduler self.fuse_reciprocal = self.config.local_sa_fuse_reciprocal self.use_base2_exp = self.config.local_sa_use_base2_exp + # Diagonal sub-tile skipping requires a pure causal mask. + self.qk_diag_skip = False + self.sv_diag_skip = False else: self.block_q = self.config.sa_block_q self.block_kv = self.config.sa_block_kv self.block_kv_compute = self.config.sa_block_kv_compute + self.eval_block_q = self.config.eval_sa_block_q + self.eval_block_kv = self.config.eval_sa_block_kv + self.eval_block_kv_compute = self.config.eval_sa_block_kv_compute self.block_q_dkv = self.config.sa_block_q_dkv self.block_kv_dkv = self.config.sa_block_kv_dkv self.block_kv_dkv_compute = self.config.sa_block_kv_dkv_compute @@ -687,9 +706,14 @@ def __init__( self.q_layout = self.config.sa_q_layout self.k_layout = self.config.sa_k_layout self.v_layout = self.config.sa_v_layout + self.eval_q_layout = self.config.eval_sa_q_layout + self.eval_k_layout = self.config.eval_sa_k_layout + self.eval_v_layout = self.config.eval_sa_v_layout self.use_splash_scheduler = self.config.use_splash_scheduler self.fuse_reciprocal = self.config.sa_fuse_reciprocal self.use_base2_exp = self.config.sa_use_base2_exp + self.qk_diag_skip = self.config.sa_qk_diag_skip + self.sv_diag_skip = self.config.sa_sv_diag_skip self.attn_logits_soft_cap = attn_logits_soft_cap self.sliding_window_size = sliding_window_size self.chunk_attn_window_size = chunk_attn_window_size @@ -1645,6 +1669,14 @@ def apply_attention( wv_product_einsum: Callable[..., Array], ): """Apply attention""" + if self.attention_kernel == "cudnn_flash_jax": + validate_gpu_flash_attention(sinks, record_max_logits) + if isinstance(key, KVTensor): + key = key.dequant() + if isinstance(value, KVTensor): + value = value.dequant() + query, key, value = self._align_qkv_for_cudnn_flash(query, key, value) + self.check_attention_inputs(query, key, value) length = query.shape[-3] target_hardware = self.mesh.devices[(0,) * self.mesh.devices.ndim].platform @@ -1809,11 +1841,6 @@ def apply_attention( None, ) elif self.attention_kernel == "cudnn_flash_jax": - validate_gpu_flash_attention(sinks, record_max_logits) - if isinstance(key, KVTensor): - key = key.dequant() - if isinstance(value, KVTensor): - value = value.dequant() return ( *self.cudnn_jax_flash_attention(query, key, value, decoder_segment_ids, model_mode), None, @@ -1985,16 +2012,41 @@ def tpu_flash_attention( pad_kv = 0 decoder_segment_ids_kv_in = decoder_segment_ids_kv if decoder_segment_ids_kv is not None else decoder_segment_ids + if max_utils.is_eval(self.config): + block_q = self.eval_block_q + block_kv = self.eval_block_kv + block_kv_compute = self.eval_block_kv_compute + block_q_dkv = block_q + block_kv_dkv = block_kv + block_kv_dkv_compute = block_kv_compute + block_q_dq = block_q + block_kv_dq = block_kv + q_layout = self.eval_q_layout + k_layout = self.eval_k_layout + v_layout = self.eval_v_layout + else: + block_q = self.block_q + block_kv = self.block_kv + block_kv_compute = self.block_kv_compute + block_q_dkv = self.block_q_dkv + block_kv_dkv = self.block_kv_dkv + block_kv_dkv_compute = self.block_kv_dkv_compute + block_q_dq = self.block_q_dq + block_kv_dq = self.block_kv_dq + q_layout = self.q_layout + k_layout = self.k_layout + v_layout = self.v_layout + # Pad sequences to block-sized boundaries upfront for AttentionType.COMPRESSED if self.attention_type == AttentionType.COMPRESSED: - if query.shape[2] % self.block_q != 0: - pad_q = (self.block_q - (query.shape[2] % self.block_q)) % self.block_q + if query.shape[2] % block_q != 0: + pad_q = (block_q - (query.shape[2] % block_q)) % block_q query = jnp.pad(query, ((0, 0), (0, 0), (0, pad_q), (0, 0))) if decoder_segment_ids is not None: decoder_segment_ids = jnp.pad(decoder_segment_ids, ((0, 0), (0, pad_q)), constant_values=-1) - if key.shape[2] % self.block_kv != 0: - pad_kv = (self.block_kv - (key.shape[2] % self.block_kv)) % self.block_kv + if key.shape[2] % block_kv != 0: + pad_kv = (block_kv - (key.shape[2] % block_kv)) % block_kv key = jnp.pad(key, ((0, 0), (0, 0), (0, pad_kv), (0, 0))) value = jnp.pad(value, ((0, 0), (0, 0), (0, pad_kv), (0, 0))) if decoder_segment_ids_kv_in is not None: @@ -2065,17 +2117,32 @@ def tpu_flash_attention( # create_splash_attention config def create_sa_config(config, query, key, attn_logits_soft_cap): if config.use_tokamax_splash: + # Tokamax only allows diagonal sub-tile skipping with a pure `CausalMask`, which + # GLOBAL/MLA get unless load-balanced context parallelism permutes the sequence. + use_diag_skip = ( + self.attention_type in (AttentionType.GLOBAL, AttentionType.MLA) + and query.shape[2] == key.shape[2] + and not (cp_size > 1 and load_balanced_context_parallel) + ) + diag_skip_kwargs = {} + if use_diag_skip and (self.qk_diag_skip or self.sv_diag_skip): + # Only passed when enabled so older Tokamax releases (< 0.0.15) keep working by default. + diag_skip_kwargs = { + "qk_diag_skip": self.qk_diag_skip, + "sv_diag_skip": self.sv_diag_skip, + "qk_diag_grid": config.sa_qk_diag_grid, + } sa_config = tokamax_splash_kernel.SplashConfig( - block_q=min(self.block_q, query.shape[2]), - block_kv=min(self.block_kv, key.shape[2]), - block_kv_compute=min(self.block_kv_compute, key.shape[2]), - block_q_dkv=min(self.block_q_dkv, query.shape[2]), - block_kv_dkv=min(self.block_kv_dkv, key.shape[2]), - block_kv_dkv_compute=min(self.block_kv_dkv_compute, key.shape[2]), + block_q=min(block_q, query.shape[2]), + block_kv=min(block_kv, key.shape[2]), + block_kv_compute=min(block_kv_compute, key.shape[2]), + block_q_dkv=min(block_q_dkv, query.shape[2]), + block_kv_dkv=min(block_kv_dkv, key.shape[2]), + block_kv_dkv_compute=min(block_kv_dkv_compute, key.shape[2]), use_fused_bwd_kernel=True, # tokamax only supports fused bwd kernel - q_layout=tokamax_splash_kernel.QKVLayout[self.q_layout], - k_layout=tokamax_splash_kernel.QKVLayout[self.k_layout], - v_layout=tokamax_splash_kernel.QKVLayout[self.v_layout], + q_layout=tokamax_splash_kernel.QKVLayout[q_layout], + k_layout=tokamax_splash_kernel.QKVLayout[k_layout], + v_layout=tokamax_splash_kernel.QKVLayout[v_layout], attn_logits_soft_cap=attn_logits_soft_cap, fuse_reciprocal=self.fuse_reciprocal, use_base2_exp=self.use_base2_exp, @@ -2103,21 +2170,22 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): else (_COMPRESSED_DQ_REDUCTION_STEPS if self.attention_type == AttentionType.COMPRESSED else None) ), use_experimental_scheduler=self.use_splash_scheduler, + **diag_skip_kwargs, ) else: sa_config = splash_attention_kernel.BlockSizes( - block_q=min(self.block_q, query.shape[2]), - block_kv=min(self.block_kv, key.shape[2]), - block_kv_compute=min(self.block_kv_compute, key.shape[2]), - block_q_dkv=min(self.block_q_dkv, query.shape[2]), - block_kv_dkv=min(self.block_kv_dkv, key.shape[2]), - block_kv_dkv_compute=min(self.block_kv_dkv_compute, key.shape[2]), - block_q_dq=None if self.use_fused_bwd_kernel else min(self.block_q_dq, query.shape[2]), - block_kv_dq=None if self.use_fused_bwd_kernel else min(self.block_kv_dq, query.shape[2]), + block_q=min(block_q, query.shape[2]), + block_kv=min(block_kv, key.shape[2]), + block_kv_compute=min(block_kv_compute, key.shape[2]), + block_q_dkv=min(block_q_dkv, query.shape[2]), + block_kv_dkv=min(block_kv_dkv, key.shape[2]), + block_kv_dkv_compute=min(block_kv_dkv_compute, key.shape[2]), + block_q_dq=None if self.use_fused_bwd_kernel else min(block_q_dq, query.shape[2]), + block_kv_dq=None if self.use_fused_bwd_kernel else min(block_kv_dq, key.shape[2]), use_fused_bwd_kernel=self.use_fused_bwd_kernel, - q_layout=splash_attention_kernel.QKVLayout[self.q_layout], - k_layout=splash_attention_kernel.QKVLayout[self.k_layout], - v_layout=splash_attention_kernel.QKVLayout[self.v_layout], + q_layout=splash_attention_kernel.QKVLayout[q_layout], + k_layout=splash_attention_kernel.QKVLayout[k_layout], + v_layout=splash_attention_kernel.QKVLayout[v_layout], ) return sa_config @@ -2287,7 +2355,7 @@ def wrap_tokamax_splash_kernel(single_head_mask): elif not use_tokamax_ring and not use_ulysses and not use_usp and self.config.use_jax_splash: if self.config.use_max_logit_estimate > 0: sa_config = dataclasses.replace(sa_config, max_logit_const=self.config.use_max_logit_estimate) - segment_axis_names_splash_kernel = nn.logical_to_mesh_axes((Q_LENGTH,)) + segment_axis_names_splash_kernel = logical_to_spec((Q_LENGTH,)) elif not use_tokamax_ring and not use_ulysses and not use_usp: # Create multi-head mask multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1]) @@ -2835,6 +2903,25 @@ def _sequence_descriptor(segment_ids): _inject_te_softmax_offset(dpa_layer, _sinks_to_te_softmax_offset(sinks, self.num_query_heads)) return dpa_layer(query, key, value, sequence_descriptor=attn_mask) + def _align_qkv_for_cudnn_flash( + self, + query: Array, + key: Array, + value: Array, + ) -> tuple[Array, Array, Array]: + """Broadcasts key/value batch dim to match query batch dim for cuDNN flash attention.""" + if query.shape[0] != key.shape[0]: + if key.shape[0] == 1 and query.shape[0] > 1: + key = jnp.broadcast_to(key, (query.shape[0], *key.shape[1:])) + value = jnp.broadcast_to(value, (query.shape[0], *value.shape[1:])) + else: + raise ValueError( + "Query and key/value batch sizes must match for cuDNN flash attention: " + f"{query.shape=}, {key.shape=}, {value.shape=}" + ) + + return query, key, value + def cudnn_jax_flash_attention( self, query: Array, @@ -2855,6 +2942,8 @@ def cudnn_jax_flash_attention( if model_mode == MODEL_MODE_AUTOREGRESSIVE: lengths = jnp.sum(decoder_segment_ids, axis=-1) + if lengths.shape[0] == 1 and query.shape[0] > 1: + lengths = jnp.broadcast_to(lengths, (query.shape[0],)) output, lse = dot_product_attention( query, @@ -2933,9 +3022,9 @@ def compute_local_attention( local_out = self.wv_product(local_exps, value, model_mode, wv_product_einsum) if model_mode == MODEL_MODE_AUTOREGRESSIVE and self.is_partition_in_decode(q_seq_len): - local_out = partitioning.with_sharding_constraint(local_out, (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV)) + local_out = with_logical_constraint(local_out, (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV)) elif model_mode == MODEL_MODE_PREFILL: - local_out = partitioning.with_sharding_constraint(local_out, (BATCH_ATTN, KV_LENGTH, HEAD, D_KV)) + local_out = with_logical_constraint(local_out, (BATCH_ATTN, KV_LENGTH, HEAD, D_KV)) if self.reshape_q and q_seq_len == 1: local_max = local_max[:, 0:1, :, :] @@ -2943,9 +3032,9 @@ def compute_local_attention( local_out = local_out[:, 0:1, :, :] if model_mode == MODEL_MODE_AUTOREGRESSIVE and self.is_partition_in_decode(q_seq_len): - local_max = partitioning.with_sharding_constraint(local_max, (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV)) - local_sum = partitioning.with_sharding_constraint(local_sum, (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV)) - local_out = partitioning.with_sharding_constraint(local_out, (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV)) + local_max = with_logical_constraint(local_max, (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV)) + local_sum = with_logical_constraint(local_sum, (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV)) + local_out = with_logical_constraint(local_out, (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV)) return local_out, local_max, local_sum @@ -2983,34 +3072,45 @@ def apply_attention_dot( # special sharding for decode q_seq_len = query.shape[1] - prefill_qkv_sharding = (BATCH_ATTN, PREFILL_LENGTH, HEAD, D_KV) - decode_qkv_sharding = (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV) + if self.is_absorbed_mqa: + prefill_q_sharding = (BATCH_ATTN, PREFILL_LENGTH, HEAD, None) + decode_q_sharding = (DECODE_BATCH, DECODE_LENGTH, HEAD, None) + prefill_kv_sharding = (BATCH_ATTN, PREFILL_LENGTH, None, None) + decode_kv_sharding = (DECODE_BATCH, DECODE_LENGTH, None, None) + weights_decode_shd = (KV_LENGTH, None, HEAD, None, None) + weights_prefill_shd = (BATCH_ATTN, None, HEAD, PREFILL_LENGTH, KV_LENGTH) + else: + prefill_q_sharding = (BATCH_ATTN, PREFILL_LENGTH, HEAD, D_KV) + decode_q_sharding = (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV) + prefill_kv_sharding = prefill_q_sharding + decode_kv_sharding = decode_q_sharding + weights_decode_shd = (KV_LENGTH, HEAD, None, None, None) + weights_prefill_shd = (BATCH_ATTN, HEAD, None, PREFILL_LENGTH, KV_LENGTH) + if self.is_partition_in_decode(q_seq_len): - query = partitioning.with_sharding_constraint(query, decode_qkv_sharding) + query = with_logical_constraint(query, decode_q_sharding) # avoid sharding scale tensor when using kv cache quantization if self.kv_quant and isinstance(key, KVTensor) and isinstance(value, KVTensor): - key.qvalue = partitioning.with_sharding_constraint(key.qvalue, decode_qkv_sharding) - value.qvalue = partitioning.with_sharding_constraint(value.qvalue, decode_qkv_sharding) + key.qvalue = with_logical_constraint(key.qvalue, decode_kv_sharding) + value.qvalue = with_logical_constraint(value.qvalue, decode_kv_sharding) else: - key = partitioning.with_sharding_constraint(key, decode_qkv_sharding) - value = partitioning.with_sharding_constraint(value, decode_qkv_sharding) + key = with_logical_constraint(key, decode_kv_sharding) + value = with_logical_constraint(value, decode_kv_sharding) elif model_mode == MODEL_MODE_PREFILL: - query = partitioning.with_sharding_constraint(query, prefill_qkv_sharding) + query = with_logical_constraint(query, prefill_q_sharding) # avoid sharding scale tensor when using kv cache quantization if self.kv_quant and isinstance(key, KVTensor) and isinstance(value, KVTensor): - key.qvalue = partitioning.with_sharding_constraint(key.qvalue, prefill_qkv_sharding) - value.qvalue = partitioning.with_sharding_constraint(value.qvalue, prefill_qkv_sharding) + key.qvalue = with_logical_constraint(key.qvalue, prefill_kv_sharding) + value.qvalue = with_logical_constraint(value.qvalue, prefill_kv_sharding) else: - key = partitioning.with_sharding_constraint(key, prefill_qkv_sharding) - value = partitioning.with_sharding_constraint(value, prefill_qkv_sharding) + key = with_logical_constraint(key, prefill_kv_sharding) + value = with_logical_constraint(value, prefill_kv_sharding) attn_weights = self.qk_product(query, key, q_seq_len, model_mode, qk_product_einsum) if self.is_partition_in_decode(q_seq_len): - attn_weights = partitioning.with_sharding_constraint(attn_weights, (KV_LENGTH, HEAD, None, None, None)) + attn_weights = with_logical_constraint(attn_weights, weights_decode_shd) elif model_mode == MODEL_MODE_PREFILL: - attn_weights = partitioning.with_sharding_constraint( - attn_weights, (BATCH_ATTN, HEAD, None, PREFILL_LENGTH, KV_LENGTH) - ) + attn_weights = with_logical_constraint(attn_weights, weights_prefill_shd) if self.attn_logits_soft_cap: attn_weights = jnp.tanh(attn_weights / self.attn_logits_soft_cap) @@ -3059,9 +3159,9 @@ def apply_attention_dot( attn_weights = apply_mask_to_logits(attn_weights, indexer_mask) if self.is_partition_in_decode(q_seq_len): - attn_mask = partitioning.with_sharding_constraint(attn_mask, (KV_LENGTH, HEAD, None, None, None)) + attn_mask = with_logical_constraint(attn_mask, (KV_LENGTH, HEAD, None, None, None)) elif model_mode == MODEL_MODE_PREFILL: - attn_mask = partitioning.with_sharding_constraint(attn_mask, (BATCH_ATTN, HEAD, None, PREFILL_LENGTH, KV_LENGTH)) + attn_mask = with_logical_constraint(attn_mask, (BATCH_ATTN, HEAD, None, PREFILL_LENGTH, KV_LENGTH)) if attn_mask is not None: attn_weights = apply_mask_to_logits(attn_weights, attn_mask) @@ -3106,7 +3206,7 @@ def qk_product( """ b, t, n, d = query.shape n_kv = key.shape[-2] - assert n_kv == self.num_kv_heads + assert n_kv == self.num_kv_heads and n % n_kv == 0 precision_kwargs = {"precision": self.config.matmul_precision} if einsum is jnp.einsum else {} if model_mode == MODEL_MODE_TRAIN or self.compute_axis_order == ( 0, diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index 9945729fa8..bc273ca1a9 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -51,7 +51,6 @@ MODEL_MODE_PREFILL, AttentionType, ) -from maxtext.layers import nnx_wrappers from maxtext.layers.attention_op import AttentionOp, _resolve_attention_type from maxtext.layers.embeddings import ( LLaMARotaryEmbedding, @@ -63,7 +62,7 @@ PartialRotaryEmbedding, Gemma4PartialRotaryEmbedding, ) -from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned, default_bias_init +from maxtext.layers.initializers import nd_dense_init, NdInitializer, default_bias_init from maxtext.layers.linears import DenseGeneral, canonicalize_tuple, normalize_axes from maxtext.layers.normalizations import RMSNorm, Qwen3NextRMSNorm, GlobalRMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant @@ -91,131 +90,6 @@ def __call__(self, x): return x * jax.lax.rsqrt(jnp.mean(x**2, axis=-1, keepdims=True) + self.eps) -def attention_as_linen( - *, - config: Config, - num_query_heads: int, - num_kv_heads: int, - head_dim: int, - max_target_length: int, - mesh: Mesh, - attention_kernel: str, - inputs_q_shape: Tuple, - inputs_kv_shape: Tuple, - dtype: DType = jnp.float32, - weight_dtype: DType = jnp.float32, - max_prefill_predict_length: int = -1, - dropout_rate: float = 0.0, - kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "normal"), - float32_qk_product: bool = False, # computes logits in float32 for stability. - float32_logits: bool = False, # cast logits in float32 for stability. - quant: Optional[Quant] = None, - kv_quant: Optional[KVQuant] = None, - attention_type: AttentionType = AttentionType.GLOBAL, - attn_logits_soft_cap: float | None = None, - sliding_window_size: int | None = None, - use_ragged_attention: bool = False, - ragged_block_size: int = 256, - use_qk_norm: bool = False, - query_pre_attn_scalar: float | None = None, - use_bias_in_projections: bool = False, # Set to True will enable bias in q, k, v, o projections - share_kv_projections: bool = False, # If true, Key and Value use the same projection - # Temperature tuning parameters used for Llama4 - temperature_tuning: bool = False, - temperature_tuning_scale: float = 0.1, - temperature_tuning_floor_scale: float = 8192.0, - # Shard the query activation as the same as the key and value. - # TODO: Find a better sharding axis name. - # TODO: Further break down the Training and Inference axes for the q, k, v. - prefill_query_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), - prefill_key_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), - prefill_value_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), - query_axis_names: AxisNames = (KV_BATCH, ATTN_LENGTH, KV_HEAD, KV_HEAD_DIM), - key_axis_names: AxisNames = (KV_BATCH, ATTN_LENGTH, KV_HEAD, KV_HEAD_DIM), - value_axis_names: AxisNames = (KV_BATCH, ATTN_LENGTH, KV_HEAD, KV_HEAD_DIM), - input_axis_names: AxisNames = (BATCH_ATTN, ATTN_INPUT_LENGTH, ATTN_EMBED), - out_axis_names: AxisNames = (BATCH_ATTN, ATTN_LENGTH, HEAD, D_KV), - prefill_input_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, ATTN_EMBED), - decode_input_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, ATTN_EMBED), - prefill_out_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, HEAD, D_KV), - decode_out_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV), - prefill_cache_axis_order: AxisIdxes = (1, 2, 0, 3), - ar_cache_axis_order: AxisIdxes = (1, 2, 0, 3), - compute_axis_order: AxisIdxes = (0, 1, 2, 3), - reshape_q: bool = False, - is_nope_layer: bool = False, - is_vision: bool = False, - model_mode: str = MODEL_MODE_TRAIN, - use_mrope: bool = False, - mrope_section: tuple[int, int, int] | None = None, - name: str | None = None, - rope_type: str | None = None, -): - """A factory function to create an Attention as a Linen module. - - This function serves as a bridge to use the NNX-based `Attention` within a - Linen model. - """ - return nnx_wrappers.to_linen( - Attention, - config=config, - num_query_heads=num_query_heads, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - max_target_length=max_target_length, - mesh=mesh, - attention_kernel=attention_kernel, - inputs_q_shape=inputs_q_shape, - inputs_kv_shape=inputs_kv_shape, - dtype=dtype, - weight_dtype=weight_dtype, - max_prefill_predict_length=max_prefill_predict_length, - dropout_rate=dropout_rate, - kernel_init=kernel_init, - float32_qk_product=float32_qk_product, - float32_logits=float32_logits, - quant=quant, - kv_quant=kv_quant, - attention_type=attention_type, - attn_logits_soft_cap=attn_logits_soft_cap, - sliding_window_size=sliding_window_size, - use_ragged_attention=use_ragged_attention, - ragged_block_size=ragged_block_size, - use_qk_norm=use_qk_norm, - query_pre_attn_scalar=query_pre_attn_scalar, - use_bias_in_projections=use_bias_in_projections, - share_kv_projections=share_kv_projections, - temperature_tuning=temperature_tuning, - temperature_tuning_scale=temperature_tuning_scale, - temperature_tuning_floor_scale=temperature_tuning_floor_scale, - prefill_query_axis_names=prefill_query_axis_names, - prefill_key_axis_names=prefill_key_axis_names, - prefill_value_axis_names=prefill_value_axis_names, - query_axis_names=query_axis_names, - key_axis_names=key_axis_names, - value_axis_names=value_axis_names, - input_axis_names=input_axis_names, - out_axis_names=out_axis_names, - prefill_input_axis_names=prefill_input_axis_names, - decode_input_axis_names=decode_input_axis_names, - prefill_out_axis_names=prefill_out_axis_names, - decode_out_axis_names=decode_out_axis_names, - prefill_cache_axis_order=prefill_cache_axis_order, - ar_cache_axis_order=ar_cache_axis_order, - compute_axis_order=compute_axis_order, - reshape_q=reshape_q, - is_nope_layer=is_nope_layer, - is_vision=is_vision, - model_mode=model_mode, - use_mrope=use_mrope, - mrope_section=mrope_section, - name=name, - rope_type=rope_type, - metadata_fn=variable_to_logically_partitioned, - abstract_init=False, - ) - - class Attention(nnx.Module): """Attention Module. @@ -447,7 +321,7 @@ def __init__( quant=self.quant, kv_quant=self.kv_quant, num_query_heads=self.num_query_heads, - num_kv_heads=self.num_kv_heads, + num_kv_heads=self._attention_op_num_kv_heads(), dropout_rate=self.dropout_rate, dtype=self.dtype, compute_axis_order=self.compute_axis_order, @@ -557,6 +431,9 @@ def _logical_to_mesh_axes(self, logical_name): logical_rules = None if self.config.using_pipeline_parallelism else self.config.logical_axis_rules return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=logical_rules) + def _attention_op_num_kv_heads(self) -> int: + return self.num_kv_heads + def _validate_kv_heads(self) -> None: """Validates the number of key/value heads.""" if self.num_kv_heads == -1: @@ -892,7 +769,7 @@ def init_rotary_embedding(self): rope_type = self.rope_type rope_use_scale = self.config.rope_use_scale if self.is_vision: - if self.config.model_name.startswith("qwen3") or self.config.model_name.startswith("cosmos3"): + if self.config.model_name.startswith("qwen3") or self.config.model_name.startswith("weaver"): rotary_embedding = Qwen3OmniMoeVisionRotaryEmbedding( hidden_size=self.config.hidden_size_for_vit, num_attention_heads=self.config.num_attention_heads_for_vit, @@ -955,6 +832,8 @@ def init_rotary_embedding(self): truncate=self.config.rope_truncate, attention_scaling=self.config.rope_attention_scaling, pairwise=self.config.rope_pairwise, + direct_position_freqs=self.config.rope_direct_position_freqs, + dot_pairwise=self.config.rope_dot_pairwise, shard_mode=self.config.shard_mode, rngs=self.rngs, ) diff --git a/src/maxtext/layers/embeddings.py b/src/maxtext/layers/embeddings.py index 794531bc50..5ffd7020b5 100644 --- a/src/maxtext/layers/embeddings.py +++ b/src/maxtext/layers/embeddings.py @@ -138,6 +138,17 @@ def __call__(self, inputs: Array, model_mode: str = MODEL_MODE_TRAIN) -> Array: return output + def embed_single_token(self, token_id: int) -> Array: + """Returns the [num_features] embedding of one token id, bitwise equal to the row `__call__` returns.""" + embedding = jnp.asarray( + _maybe_move_embedding_to_device(self.embedding.get_value(), self.config), + self.dtype, + ) + row = embedding[token_id] + if self.config.use_iota_embed: + row = jnp.where(row == 0, jnp.zeros_like(row), row) + return row + def attend(self, query: Array, out_sharding: NamedSharding | None = None) -> Array: """Attend over the embedding using a query array. @@ -585,13 +596,13 @@ class YarnRotaryEmbedding(nnx.Module): when `rope_factor > 1`. This scaling can be applied within this layer (if `attention_scaling=True`) or externally. - RoPE Implementation Details (General): - - Arithmetic: Uses complex number arithmetic. Real number arithmetic is not implemented here, - though the resulting embeddings would be equivalent. + - Arithmetic: Uses complex number arithmetic by default. `pairwise=True` uses real number arithmetic + on (real, imag) pairs instead; the resulting embeddings are equivalent. - Input Layout: Supports both interleaved (`interleave=True`, e.g., [real1, img1, real2, img2]) and concatenated (`interleave=False`, e.g., [real1, real2, img1, img2]) formats. - - Output Layout: Always returns concatenated format ([real, imag]). Interleaved output is not - implemented: While the embedding is different, attention scores are invariant, as long as we apply - the same output layout for Q and K. + - Output Layout: Returns concatenated format ([real, imag]) by default, or interleaved format when + `pairwise=True`. While the embedding is different, attention scores are invariant, as long as we + apply the same output layout for Q and K. Attributes: embedding_dims: Dimension of the embedding to be generated. @@ -606,6 +617,14 @@ class YarnRotaryEmbedding(nnx.Module): rope_interleave: Whether complex representation is interleaved or concatenated. rope_truncate: Whether or not to floor lower bound and ceil upper bound for correction range. rope_attention_scaling: Whether or not to scale the rotary embedding output. + pairwise: Whether to rotate interleaved inputs as real (real, imag) pairs and return interleaved output. + Requires `interleave=True`. + direct_position_freqs: Whether to compute angles directly as `position * inv_freqs` instead of + gathering from a precomputed [max_position_embeddings, half_dim] table. Outputs match for positions + < max_position_embeddings. Beyond that, the table gather clamps every position to the last row, while + the direct computation keeps using the true angle. + dot_pairwise: Whether to keep the pairwise rotation in rank 4 by applying the pair swap as a block-skew + matmul instead of reshaping to rank 5. Requires `pairwise=True`. rngs: rng keys passed in by nnx.bridge.to_linen. """ @@ -626,6 +645,8 @@ def __init__( truncate=True, attention_scaling=False, pairwise=False, + direct_position_freqs=False, + dot_pairwise=False, # Not used in YarnRotaryEmbedding but passed in by nnx.bridge.to_linen. # TODO: Remove when bridge no longer needed rngs: nnx.Rngs = None, @@ -646,9 +667,13 @@ def __init__( self.shard_mode = shard_mode self.attention_scaling = attention_scaling self.pairwise = pairwise + self.direct_position_freqs = direct_position_freqs + self.dot_pairwise = dot_pairwise if self.pairwise and not self.interleave: raise ValueError("rope_pairwise=True requires rope_interleave=True.") + if self.dot_pairwise and not self.pairwise: + raise ValueError("rope_dot_pairwise=True requires rope_pairwise=True.") self.freqs_sharding = ( create_sharding(mesh, ("activation_batch", "activation_length", "q_heads")) @@ -660,8 +685,8 @@ def __init__( raise ValueError("Embedding dim for rotary position embedding must be a multiple of 2.") @property - def freqs_cis(self): - """Frequencies for rotary embedding.""" + def inv_freqs(self) -> Array: + """1D base inverse frequencies of shape [half_dim] in float32.""" half_dim = self.embedding_dims // 2 # Compute base frequencies for each (even-indexed) dimension. # (Note: We use jnp.arange with float32 for precision.) @@ -677,7 +702,12 @@ def freqs_cis(self): ) smooth = 1 - self._linear_ramp_factor(low, high, half_dim) # The corrected frequency is a weighted mix of the scaled and base values. - freqs = freqs / self.rope_factor * (1 - smooth) + freqs * smooth + return freqs / self.rope_factor * (1 - smooth) + freqs * smooth + + @property + def freqs_cis(self): + """Frequencies for rotary embedding.""" + freqs = self.inv_freqs # Precompute frequencies for all positions by taking the outer product. t = jnp.arange(self.max_position_embeddings, dtype=jnp.float32) # shape [max_position_embeddings] @@ -687,12 +717,13 @@ def freqs_cis(self): # Compute the complex “cis” values: exp(i * theta). return jnp.exp(1j * freqs) # shape [max_position_embeddings, half_dim] - def _find_correction_dim(self, num_rotations: float, dim: int, base: float, max_position_embeddings: int) -> float: + @staticmethod + def _find_correction_dim(num_rotations: float, dim: int, base: float, max_position_embeddings: int) -> float: """Compute the correction dimension for a given number of rotations.""" return dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) / (2 * math.log(base)) + @staticmethod def _find_correction_range( - self, low_rot: float, high_rot: float, dim: int, @@ -713,8 +744,8 @@ def _find_correction_range( Returns: tuple[int, int]: The range of correction dimensions (low, high), clamped to valid indices. """ - low = self._find_correction_dim(low_rot, dim, base, max_position_embeddings) - high = self._find_correction_dim(high_rot, dim, base, max_position_embeddings) + low = YarnRotaryEmbedding._find_correction_dim(low_rot, dim, base, max_position_embeddings) + high = YarnRotaryEmbedding._find_correction_dim(high_rot, dim, base, max_position_embeddings) if truncate: low = math.floor(low) high = math.ceil(high) @@ -722,7 +753,8 @@ def _find_correction_range( high = min(high, dim - 1) return low, high - def _linear_ramp_factor(self, min_val: float, max_val: float, dim: int) -> Array: + @staticmethod + def _linear_ramp_factor(min_val: float, max_val: float, dim: int) -> Array: """Computes a linear ramp over the dimension. Returns a jax.Array of shape (dim,) with values between 0 and 1. @@ -732,6 +764,34 @@ def _linear_ramp_factor(self, min_val: float, max_val: float, dim: int) -> Array linear_func = (jnp.arange(dim, dtype=jnp.float32) - min_val) / (max_val - min_val) return jnp.clip(linear_func, 0, 1) + def _position_angles(self, position: Array, inv_freqs: Array) -> Array: + """Computes float32 RoPE angles position * inv_freqs of shape [B, S, 1, len(inv_freqs)].""" + return ( + position.astype(jnp.float32)[:, :, jnp.newaxis, jnp.newaxis] * inv_freqs[jnp.newaxis, jnp.newaxis, jnp.newaxis, :] + ) + + def _gather_freqs_cis(self, position: Array) -> Array: + """Looks up precomputed complex frequencies, returning shape [B, S, 1, half_dim].""" + freqs = self.freqs_cis.at[position].get(out_sharding=self.freqs_sharding) # shape: [B, S, half_dim] + return freqs[:, :, jnp.newaxis, :] + + def _cos_sin(self, position: Array, repeat_interleave: bool = False) -> tuple[Array, Array]: + """Returns float32 (cos, sin) of the RoPE angles, each of shape [B, S, 1, half_dim]. + + With repeat_interleave=True, each frequency is repeated twice ([f_0, f_0, f_1, f_1, ...]) to match the + interleaved layout, giving shape [B, S, 1, dim]. + """ + if self.direct_position_freqs: + inv_freqs = jnp.repeat(self.inv_freqs, 2) if repeat_interleave else self.inv_freqs + angles = self._position_angles(position, inv_freqs) + return jnp.cos(angles), jnp.sin(angles) + freqs = self._gather_freqs_cis(position) + cos, sin = jnp.real(freqs), jnp.imag(freqs) + if repeat_interleave: + cos = jnp.repeat(cos, 2, axis=-1) + sin = jnp.repeat(sin, 2, axis=-1) + return cos, sin + def __call__(self, inputs: Array, position: None | Array = None) -> Array: """Applies the rotary positional embedding using the precomputed complex frequencies. @@ -756,29 +816,47 @@ def __call__(self, inputs: Array, position: None | Array = None) -> Array: else: position = position.astype(jnp.int32) - # Lookup the precomputed frequencies using the position indices. - # self.freqs_cis has shape [max_position_embeddings, half_dim] so we use jnp.take along axis 0. - # After indexing, shape becomes [B, S, half_dim]; we then add an axis for the heads. - freqs = self.freqs_cis.at[position].get(out_sharding=self.freqs_sharding) # shape: [B, S, half_dim] - freqs = freqs[:, :, jnp.newaxis, :] # shape: [B, S, 1, half_dim] - if self.interleave and self.pairwise: with jax.named_scope("rope_pairwise"): b, s, n, h = inputs.shape half_dim = h // 2 - pairs = inputs.reshape(b, s, n, half_dim, 2) - pairs = pairs.astype(jnp.float32) - cos = jnp.real(freqs)[..., jnp.newaxis] - sin = jnp.imag(freqs)[..., jnp.newaxis] - if self.shard_mode == ShardMode.EXPLICIT: - rotated_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", None, None, None)) - cos = jnp.broadcast_to(cos, pairs.shape, out_sharding=rotated_sharding) - sin = jnp.broadcast_to(sin, pairs.shape, out_sharding=rotated_sharding) - swapped = jnp.flip(pairs, axis=-1) - sign = jnp.asarray([-1.0, 1.0], dtype=jnp.float32) - rotated_pairs = pairs * cos + swapped * sin * sign - output = rotated_pairs.reshape(b, s, n, h) + if self.dot_pairwise: + cos, sin = self._cos_sin(position, repeat_interleave=True) # shape: [B, S, 1, dim] + if self.shard_mode == ShardMode.EXPLICIT: + rotated_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", None, None)) + cos = jnp.broadcast_to(cos, inputs.shape, out_sharding=rotated_sharding) + sin = jnp.broadcast_to(sin, inputs.shape, out_sharding=rotated_sharding) + # Swap-and-sign [x_0, x_1, ...] -> [-x_1, x_0, ...] as a matmul against the block-skew + # matrix I_{half_dim} (x) [[0, 1], [-1, 0]], avoiding the rank-5 reshape relayout copies. + p_swap = jnp.kron(jnp.eye(half_dim, dtype=inputs.dtype), jnp.asarray([[0, 1], [-1, 0]], dtype=inputs.dtype)) + swapped_signed = jax.lax.dot_general( + inputs, + p_swap, + (((3,), (0,)), ((), ())), + # The default precision rounds non-bf16 inputs to bf16 on the TPU MXU; p_swap is exact in any dtype. + precision=None if inputs.dtype == jnp.bfloat16 else jax.lax.Precision.HIGHEST, + preferred_element_type=jnp.float32, + ) + output = inputs.astype(jnp.float32) * cos + swapped_signed * sin + else: + pairs = inputs.reshape(b, s, n, half_dim, 2) + pairs = pairs.astype(jnp.float32) + cos, sin = self._cos_sin(position) # shape: [B, S, 1, half_dim] + cos = cos[..., jnp.newaxis] + sin = sin[..., jnp.newaxis] + if self.shard_mode == ShardMode.EXPLICIT: + rotated_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", None, None, None)) + cos = jnp.broadcast_to(cos, pairs.shape, out_sharding=rotated_sharding) + sin = jnp.broadcast_to(sin, pairs.shape, out_sharding=rotated_sharding) + swapped = jnp.flip(pairs, axis=-1) + sign = jnp.asarray([-1.0, 1.0], dtype=jnp.float32) + rotated_pairs = pairs * cos + swapped * sin * sign + output = rotated_pairs.reshape(b, s, n, h) else: + if self.direct_position_freqs: + freqs = jnp.exp(1j * self._position_angles(position, self.inv_freqs)) # shape: [B, S, 1, half_dim] + else: + freqs = self._gather_freqs_cis(position) # shape: [B, S, 1, half_dim] if self.interleave: # Inputs with interleaved format [real1, img1, real2, img2, ...] at last dimension # Convert the last dimension into a complex representation. @@ -1661,6 +1739,12 @@ def __init__( min_timescale: int = 10000, max_timescale: int = 10000, mesh: Any = None, + rope_type: str = "default", + rope_factor: float = 16.0, + beta_fast: float = 32.0, + beta_slow: float = 1.0, + original_max_position_embeddings: int = 65536, + truncate: bool = True, **kwargs, ): super().__init__( @@ -1674,16 +1758,37 @@ def __init__( self.partial_rotary_factor = partial_rotary_factor self.rope_theta = rope_theta self.fprop_dtype = fprop_dtype + self.rope_type = rope_type + self.rope_factor = rope_factor + self.beta_fast = beta_fast + self.beta_slow = beta_slow + self.original_max_position_embeddings = original_max_position_embeddings + self.truncate = truncate # Compute the partial rotary dimension (rope_head_dim) self.dim = int(head_dim * partial_rotary_factor) @property def inv_freq(self): + """Computes inverse frequencies with optional YaRN scaling.""" # Compute base inverse frequencies for half of self.dim half_dim = self.dim // 2 fraction = 2 * jnp.arange(0, half_dim, dtype=jnp.float32) / self.dim - return 1.0 / (self.rope_theta**fraction) + freqs = 1.0 / (self.rope_theta**fraction) + if self.rope_type != "yarn": + return freqs + + # pylint: disable=protected-access + low, high = YarnRotaryEmbedding._find_correction_range( + self.beta_fast, + self.beta_slow, + self.dim, + self.rope_theta, + self.original_max_position_embeddings, + self.truncate, + ) + smooth = 1 - YarnRotaryEmbedding._linear_ramp_factor(low, high, half_dim) + return freqs / self.rope_factor * (1 - smooth) + freqs * smooth def get_freqs(self, position_ids: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]: # position_ids: [B, S] diff --git a/src/maxtext/layers/initializers.py b/src/maxtext/layers/initializers.py index 742e00dab5..83674225d2 100644 --- a/src/maxtext/layers/initializers.py +++ b/src/maxtext/layers/initializers.py @@ -33,7 +33,7 @@ def _default_embed_init(key, shape, dtype=jnp.float32): target_dtype = dtype sample_dtype = jnp.float32 if is_fp8_dtype(dtype) else dtype - fn = nn.initializers.variance_scaling(1.0, "fan_in", "normal", out_axis=0) + fn = jax.nn.initializers.variance_scaling(1.0, "fan_in", "normal", out_axis=0) return fn(key, shape, sample_dtype).astype(target_dtype) diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index 80f14b62dc..5f65e321f6 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -32,10 +32,12 @@ from maxtext.common.common_types import MODEL_MODE_PREFILL from maxtext.layers import nnx_wrappers, quantizations from maxtext.layers import normalizations -from maxtext.layers.initializers import NdInitializer, nd_dense_init, default_bias_init, variable_to_logically_partitioned, Initializer +from maxtext.layers.initializers import NdInitializer, nd_dense_init, default_bias_init, Initializer from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_logging from maxtext.utils import max_utils +from maxtext.utils.sharding import get_logical_axis_rules +from maxtext.utils.sharding import logical_to_mesh_axes from maxtext.utils.sharding import maybe_shard_with_logical from maxtext.utils.sharding import maybe_shard_with_name from maxtext.utils.sharding import get_physical_spec_without_axes @@ -78,54 +80,34 @@ def canonicalize_tuple(x): WeightQuantConfig = quantizations.WeightQuantConfig -def _compute_dot_general( - inputs, - kernel, - kernel_axes, - axis, - contract_ind, - matmul_precision, - quant, - kernel_scale: Array | None = None, - compute_dtype: DType | None = None, -): - """Computes a dot_general operation that may be quantized.""" - dot_general = lax.dot_general - matmul_precision = lax.Precision(matmul_precision) - if quant: - dot_general_cls = quant.dot_general_cls(mesh_axes=kernel_axes) - dot_general = dot_general_cls() - return dot_general(inputs, kernel, ((axis, contract_ind), ((), ())), precision=None) - - if kernel_scale is not None or is_fp8_dtype(getattr(kernel, "dtype", None)): - if compute_dtype is None: - compute_dtype = inputs.dtype - kernel = dequantize_weight(kernel, kernel_scale, compute_dtype=compute_dtype) - - return dot_general(inputs, kernel, ((axis, contract_ind), ((), ())), precision=matmul_precision) - - def _compute_dot_general_nnx( inputs, kernel, axis, contract_ind, matmul_precision, - quant_dot_general: nnx_wrappers.ToNNX | None, + quant_dot_general: nnx.Module | None, initializing: bool, out_sharding: NamedSharding | None = None, kernel_scale: Array | None = None, compute_dtype: DType | None = None, native_fp8_compute: bool = False, scale_block_size: int | tuple[int, ...] | None = None, + act_calibration_method: str = "absmax", ): """Computes a dot_general operation that may be quantized.""" dot_general = lax.dot_general matmul_precision = lax.Precision(matmul_precision) if quant_dot_general is not None: + dimension_numbers = ((axis, contract_ind), ((), ())) + if not isinstance(quant_dot_general, nnx_wrappers.ToNNX): + return quant_dot_general(inputs, kernel, dimension_numbers, precision=None) if initializing: - quant_dot_general.lazy_init(inputs, kernel, ((axis, contract_ind), ((), ())), precision=None) - return quant_dot_general(inputs, kernel, ((axis, contract_ind), ((), ())), precision=None, mutable=["aqt"]) + quant_dot_general.lazy_init(inputs, kernel, dimension_numbers, precision=None) + return quant_dot_general(inputs, kernel, dimension_numbers, precision=None, mutable=["aqt"]) + + if compute_dtype is None: + compute_dtype = inputs.dtype if compute_dtype is None: compute_dtype = inputs.dtype @@ -152,6 +134,7 @@ def _compute_dot_general_nnx( compute_dtype=compute_dtype, precision=matmul_precision, out_sharding=out_sharding, + act_calibration_method=act_calibration_method, ) if kernel_scale is not None or is_fp8_dtype(getattr(kernel, "dtype", None)): @@ -162,6 +145,67 @@ def _compute_dot_general_nnx( ) +# Logical axis a custom rule set maps to every mesh axis the tokens are sharded over when it keeps the dense weights +# off some of them (cp-as-ep-dense-no-cp, cp-as-ep-dense-fsdp); see `route_dense_weight_grad`. +DENSE_GRAD_AXIS = "embed_grad" +_DENSE_EMBED_AXES = ("embed", "embed_attn") + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2)) +def _reshard_cotangent(x, via, to): + """Identity whose cotangent is constrained to the PartitionSpec `via` and then to `to`.""" + del via, to + return x + + +def _reshard_cotangent_fwd(x, via, to): + del via, to + return x, None + + +def _reshard_cotangent_bwd(via, to, _, cotangent): + # Without the barrier, XLA:TPU folds the reduce-scatter into `via` and the all-gather into `to` back into one + # full-size all-reduce. + cotangent = jax.lax.optimization_barrier(jax.lax.with_sharding_constraint(cotangent, via)) + return (jax.lax.with_sharding_constraint(cotangent, to),) + + +_reshard_cotangent.defvjp(_reshard_cotangent_fwd, _reshard_cotangent_bwd) + + +def route_dense_weight_grad(kernel: Array, kernel_axes: tuple[None | str, ...]) -> Array: + """Reduce-scatters a dense weight's gradient over every token axis before resharding it to the weight's sharding. + + A weight's gradient is a partial sum over every mesh axis the tokens are sharded on. When the weight is sharded + over only some of them (cp-as-ep-dense-no-cp and cp-as-ep-dense-fsdp replicate the dense weights over context), + GSPMD reduces that partial sum with a full-size all-reduce over every device and keeps a slice of it. Under a rule + set that maps `embed_grad` to the token axes, the backward of every kernel with an `embed` or `embed_attn` + dimension instead constrains the gradient to the kernel's axes with that dimension on `embed_grad` (one + reduce-scatter over all of the token axes' devices, as under cp-as-ep), and then, across an optimization barrier, + to the kernel's own sharding (an all-gather of that shard over the axes the weight is replicated on). The forward + is unchanged, and under every rule set that does not map `embed_grad` the kernel is returned as is. + + Args: + kernel: The kernel, as the dot product consumes it. + kernel_axes: Its logical axes. + + Returns: + The kernel. + """ + rules = get_logical_axis_rules() + if not rules or not any(rule[0] == DENSE_GRAD_AXIS for rule in rules): + return kernel + if not any(isinstance(axis, str) and axis in _DENSE_EMBED_AXES for axis in kernel_axes): + return kernel + mesh = jax.sharding.get_abstract_mesh() + if mesh.empty: + raise ValueError(f"The logical axis rules map {DENSE_GRAD_AXIS!r}, which needs a mesh set with jax.set_mesh.") + via_axes = tuple(DENSE_GRAD_AXIS if isinstance(a, str) and a in _DENSE_EMBED_AXES else a for a in kernel_axes) + via = logical_to_mesh_axes(via_axes, mesh, rules=rules) + to = logical_to_mesh_axes(tuple(kernel_axes), mesh, rules=rules) + return _reshard_cotangent(kernel, via, to) + + class DenseGeneral(nnx.Module): """A linear transformation with flexible axes.""" @@ -336,22 +380,21 @@ def __init__( if quant and not isinstance(quant, quantizations.ServeFp8WeightQuantization): # Native FP8 compute runs in _compute_dot_general_nnx; no Linen quantizer needed. - dot_general_cls = quant.dot_general_cls(mesh_axes=kernel_axes) - dot_general_linen = dot_general_cls() - quant_dot_general = nnx_wrappers.ToNNX(dot_general_linen, rngs=rngs) - self._quant_dot_general_name = f"{type(dot_general_linen).__name__}_0" + self._quant_dot_general_name, quant_dot_general = quantizations.make_dot_general(quant, kernel_axes, rngs) setattr(self, self._quant_dot_general_name, quant_dot_general) - block_size = getattr(quant, "get_block_size", lambda: 1)() # needed for TE MXFP8 - dummy_inputs = jnp.zeros((block_size, *self.in_features_shape), dtype=self.dtype) - self(dummy_inputs, _initializing=True) - # Backends that never draw at apply time leave dead RNG state in the model. - if not quant.needs_apply_rngs: - quant_dot_general.release_rngs() + if isinstance(quant_dot_general, nnx_wrappers.ToNNX): + # A bridged Linen quantizer creates its state on the first call. + block_size = getattr(quant, "get_block_size", lambda: 1)() # needed for TE MXFP8 + dummy_inputs = jnp.zeros((block_size, *self.in_features_shape), dtype=self.dtype) + self(dummy_inputs, _initializing=True) + # Backends that never draw at apply time leave dead RNG state in the model. + if not quant.needs_apply_rngs: + quant_dot_general.release_rngs() else: self._quant_dot_general_name = None @property - def quant_dot_general(self) -> nnx_wrappers.ToNNX | None: + def quant_dot_general(self) -> nnx.Module | None: if self._quant_dot_general_name is None: return None return getattr(self, self._quant_dot_general_name) @@ -464,12 +507,14 @@ def __call__( kernel = kernel[..., begin:end] kernel = self._maybe_two_stage_all_gather(kernel) + kernel = route_dense_weight_grad(kernel, self.kernel_axes) # out_sharding should be None for auto mesh axis if self.shard_mode != ShardMode.EXPLICIT: out_sharding = None contract_ind = tuple(range(0, len(self.axis))) + native_fp8_compute = isinstance(self.quant, quantizations.ServeFp8WeightQuantization) output = _compute_dot_general_nnx( inputs, kernel, @@ -481,8 +526,9 @@ def __call__( out_sharding, kernel_scale=kernel_scale, compute_dtype=self.dtype, - native_fp8_compute=isinstance(self.quant, quantizations.ServeFp8WeightQuantization), + native_fp8_compute=native_fp8_compute, scale_block_size=self.block_size, + act_calibration_method=self.quant.act_calibration_method if native_fp8_compute else "absmax", ) if self.bias is not None: @@ -494,92 +540,6 @@ def __call__( return output -def dense_general( - *, - inputs_shape: tuple[int, ...] | None = None, - in_features_shape: tuple[int, ...] | int | None = None, - out_features_shape: Iterable[int] | int, - axis: Iterable[int] | int = -1, - weight_dtype: DType = jnp.float32, - dtype: DType = jnp.float32, - kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "truncated_normal"), - kernel_axes: tuple[None | str, ...] = (), - quant: None | Quant = None, - use_bias: bool = False, - shard_mode: ShardMode = ShardMode.AUTO, - matmul_precision: str = "default", - parameter_memory_host_offload: bool = False, - has_scale: bool | None = None, - kernel_scale_init: Initializer | None = None, - scale_shape: Shape | None = None, - scale_axes: tuple[None | str, ...] | None = None, - scale_dtype: DType = jnp.float32, - block_size: int | tuple[int, ...] | None = None, - weight_quant: quantizations.WeightQuantConfig | None = None, - name: None | str = None, -): - """Creates a DenseGeneral Linen module using nnx.bridge.to_linen. - - Args: - inputs_shape: tuple with the shape of the inputs - in_features_shape: tuple with numbers of input features for axes specified in - 'axis'. - out_features_shape: tuple with numbers of output features. - axis: tuple with axes to apply the transformation on. - weight_dtype: the dtype of the weights (default: float32). - dtype: the dtype of the computation (default: float32). - kernel_init: initializer function for the weight matrix. - kernel_axes: logical axes for partitioning the kernel. - quant: quantization config, defaults to None implying no quantization. - use_bias: whether to add bias in linear transformation. - shard_mode: indicating the shard mode - matmul_precision: Precision for matrix multiplication. - parameter_memory_host_offload: Determines whether to offload params to host - has_scale: whether to initialize a separate scale parameter (kernel_scale). - kernel_scale_init: initializer function for kernel_scale. - scale_shape: explicit shape for kernel_scale. - scale_axes: logical axes for partitioning kernel_scale. - scale_dtype: dtype of kernel_scale (default: float32). - block_size: block size for block-wise quantization scales. - weight_quant: optional WeightQuantConfig decoupling weight quantization settings. - name: name passed to the ToLinen Module - """ - if not (inputs_shape is not None) ^ (in_features_shape is not None): - raise ValueError("Exactly one of inputs_shape or in_features must be specified.") - - if inputs_shape is not None: - axis = canonicalize_tuple(axis) - in_features_shape = tuple(inputs_shape[ax] for ax in normalize_axes(axis, len(inputs_shape))) - else: - assert in_features_shape is not None - module = nnx_wrappers.to_linen( - DenseGeneral, - in_features_shape=in_features_shape, - out_features_shape=out_features_shape, - axis=axis, - weight_dtype=weight_dtype, - dtype=dtype, - kernel_init=kernel_init, - kernel_axes=kernel_axes, - quant=quant, - use_bias=use_bias, - shard_mode=shard_mode, - matmul_precision=matmul_precision, - parameter_memory_host_offload=parameter_memory_host_offload, - has_scale=has_scale, - kernel_scale_init=kernel_scale_init, - scale_shape=scale_shape, - scale_axes=scale_axes, - scale_dtype=scale_dtype, - block_size=block_size, - weight_quant=weight_quant, - name=name, - metadata_fn=variable_to_logically_partitioned, - abstract_init=False, - ) - return module - - class Dropout(nnx.Dropout): """Forked nnx.Dropout that is easier to use with bridge""" @@ -826,46 +786,6 @@ def __call__( return output -def mlp_block( - *, - config: Config, - mesh: Mesh, - in_features: int, - intermediate_dim: int = 2048, - activations: Sequence[str | Callable[..., Any]] = ("relu",), - kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "truncated_normal"), - intermediate_dropout_rate: float = 0.1, - dtype: Any = jnp.float32, - weight_dtype: Any = jnp.float32, - use_bias: bool = False, - use_pre_norm: bool = False, - quant: None | Quant = None, - model_mode: None | str = None, - name: None | str = None, -): - """Creates a MlpBlock Linen module using nnx.bridge.to_linen.""" - module = nnx_wrappers.to_linen( - MlpBlock, - config=config, - mesh=mesh, - in_features=in_features, - intermediate_dim=intermediate_dim, - activations=activations, - kernel_init=kernel_init, - intermediate_dropout_rate=intermediate_dropout_rate, - dtype=dtype, - weight_dtype=weight_dtype, - use_bias=use_bias, - use_pre_norm=use_pre_norm, - quant=quant, - model_mode=model_mode, - name=name, - metadata_fn=variable_to_logically_partitioned, - abstract_init=False, - ) - return module - - class DeepSeekV4GroupedLinear(nnx.Module): """Block-diagonal grouped linear used by the grouped output projection in DeepSeek-V4. @@ -972,35 +892,3 @@ def __call__(self, inputs: Array) -> Array: output = jnp.einsum("...gi,gio->...go", inputs, kernel, precision=lax.Precision(self.matmul_precision)) return output - - -def deepseek_v4_grouped_linear( - *, - in_features_per_group: int, - out_features: int, - n_groups: int, - weight_dtype: DType = jnp.float32, - dtype: DType = jnp.float32, - kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "truncated_normal"), - kernel_axes: tuple[None | str, ...] = ("groups", "embed", "mlp"), - matmul_precision: str = "default", - parameter_memory_host_offload: bool = False, - name: None | str = None, -): - """Creates a DeepSeekV4GroupedLinear Linen module using nnx.bridge.to_linen.""" - module = nnx_wrappers.to_linen( - DeepSeekV4GroupedLinear, - in_features_per_group=in_features_per_group, - out_features=out_features, - n_groups=n_groups, - weight_dtype=weight_dtype, - dtype=dtype, - kernel_init=kernel_init, - kernel_axes=kernel_axes, - matmul_precision=matmul_precision, - parameter_memory_host_offload=parameter_memory_host_offload, - name=name, - metadata_fn=variable_to_logically_partitioned, - abstract_init=False, - ) - return module diff --git a/src/maxtext/layers/mhc.py b/src/maxtext/layers/mhc.py index 81f93b58fa..450f03efca 100644 --- a/src/maxtext/layers/mhc.py +++ b/src/maxtext/layers/mhc.py @@ -26,8 +26,7 @@ from maxtext.common.common_types import Array, Config from maxtext.common.common_types import HyperConnectionType from maxtext.kernels.mhc import api as mhc_kernel -from maxtext.layers import nnx_wrappers -from maxtext.layers.initializers import default_bias_init, default_scalar_init, nd_dense_init, variable_to_logically_partitioned +from maxtext.layers.initializers import default_bias_init, default_scalar_init, nd_dense_init from maxtext.layers.normalizations import RMSNorm @@ -437,9 +436,3 @@ def __call__(self, x: Array) -> Array: x_f32 = x.astype(jnp.float32) out = jnp.sum(pre[:, :, :, None] * x_f32, axis=2) return out.astype(self.dtype) - - -DeepSeek4HyperHeadToLinen = nnx_wrappers.to_linen_class( - DeepSeek4HyperHead, - base_metadata_fn=variable_to_logically_partitioned, -) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index dda2f599eb..1b4237b02c 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -18,16 +18,22 @@ import dataclasses import enum import functools +import inspect import math import random -from typing import Iterable, Optional, Tuple, Union +from typing import Any, Iterable, Optional, Tuple, Union from aqt.jax.v2 import aqt_tensor as aqt from flax import nnx from flax import struct import jax from jax import ad_checkpoint as adc -from jax.experimental.compute_on import compute_on +try: + # jax 0.11.0 (the RL trainer image) has the keyword form (compute_type, out_memory_spaces, compiler_options) only as + # compute_on2; newer jax renamed it compute_on. The SparseCore pins below use the keyword form. + from jax.experimental.compute_on import compute_on2 as compute_on +except ImportError: + from jax.experimental.compute_on import compute_on from jax.experimental import xla_metadata from jax.experimental.layout import Layout, with_layout_constraint import jax.numpy as jnp @@ -41,17 +47,23 @@ from maxtext.kernels.ragged.ragged_sort import a2a_ragged_unsort from maxtext.kernels.ragged.ragged_sort import ring_ragged_sort from maxtext.kernels.ragged.ragged_sort import ring_ragged_unsort +from maxtext.kernels.ragged.ragged_sort import ring_ragged_unsort_down_projection +from maxtext.kernels.ragged.ragged_sort_tc import ring_ragged_sort_tc +from maxtext.kernels.ragged.ragged_sort_tc import ring_ragged_unsort_tc +from maxtext.kernels.ragged.ragged_sort_tc import tc_buffer_row_weights from maxtext.layers import attentions, linears, nnx_wrappers, quantizations -from maxtext.layers.initializers import NdInitializer, default_bias_init, nd_dense_init, variable_to_logically_partitioned +from maxtext.layers.initializers import NdInitializer, default_bias_init, nd_dense_init from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils.sharding import ( + CompoundLogicalAxis, create_sharding, get_logical_axis_rules, logical_to_mesh_axes, maybe_shard_with_logical, maybe_shard_with_pspec, + mesh_axes_for_dim, remove_expert_from_partition_spec, remove_incompatible_mesh_axes_from_partition_spec, remove_mesh_axes_from_partition_spec, @@ -106,6 +118,18 @@ class RouteOutput: # Shape [local experts], tracks number of local tokens routed to every local expert. local_group_sizes: Optional[jax.Array] = None has_overflow: bool | jax.Array = False + # log_required_ragged_buffer_factor probe: minimum ragged_buffer_factor that would have avoided drops, max over + # the mesh (float32 scalar). None when the probe is off. + required_rbf: Optional[jax.Array] = None + # Scalar max over expert shards (and, after the mesh reduction, over EP groups) of the + # tokens routed to a shard divided by the perfectly balanced load. Only set when + # moe_log_max_load_ratio=True on the ring-of-experts ragged path. + max_load_ratio: Optional[jax.Array] = None + # Inverse of sorted_selected_experts on the ring-of-experts ragged path, so the unsort backward does not re-sort to + # get it. Set when moe_routing_maps is not "remat", and always with moe_routing_index_math="count". + topk_argsort_indices: Optional[jax.Array] = None + # moe_tc_ragged_sort: TcRouting shared by the TensorCore ragged sort and unsort of this chunk. + tc_routing: Optional[Any] = None # log_moe_drops=True on the ring-of-experts ragged path only (None otherwise): the drop audit # `(dropped_assignments, max_load_ratio)`, see `RoutedMoE.permute`. Logging only. moe_audit: Optional[tuple[jax.Array, jax.Array]] = None @@ -125,23 +149,6 @@ def combine_moe_audits(audits): MOE_AUDIT_NOT_COMPUTED = -1 -def max_expert_shard_load_ratio(group_sizes, num_expert_shards, balanced_load): - """Largest per-expert-shard load divided by the perfectly balanced load (probe port of MaxText main). - - Args: - group_sizes: Untruncated token assignments per expert, shape (num_experts,), experts ordered by shard. - num_expert_shards: Number of expert-parallel shards. - balanced_load: Token assignments per shard under perfectly balanced routing, i.e. the ragged - buffer size at ragged_buffer_factor=1.0. - - Returns: - float32 scalar. The ring-of-experts ragged buffer overflows iff max shard load > int(balanced_load * rbf), - which for integer loads is max shard load > balanced_load * rbf, i.e. this ratio > rbf (up to float rounding). - """ - shard_loads = jnp.sum(group_sizes.reshape(num_expert_shards, -1), axis=-1) - return jnp.max(shard_loads).astype(jnp.float32) / jnp.float32(balanced_load) - - def _truncate_matrix(all_shards_group_sizes: jax.Array, buffer_size: int) -> jax.Array: """Truncates the traffic matrix to fit in buffer_size on receiver side. @@ -275,6 +282,39 @@ def barrier_tie(x, dep, backward: bool = False): return barrier((x, dep)) +def _prefused_ffn_activation(act_fn, h01): + """`act_fn(h01[..., :n], h01[..., n:])` on the prefused up-projection output `h01 = [gate | up]`. + + moe_prefused_swiglu_vjp. The forward is exactly the slice + `act_fn` the split path runs. The backward calls + `jax.vjp` of the same `act_fn` at the same primal point, so d(gate) and d(up) come from the ops autodiff emits for + it, and returns them as ONE `[rows, 2n]` array, so the concatenate fuses into the elementwise backward. Split, the + two `[rows, n]` cotangents leave the SwiGLU-backward fusion separately and XLA joins them in an extra pass over + the whole expert buffer (the transpose of the two slices, `add_any(pad, pad)`). The values are the same; the bits + can differ only in the sign of an exact zero, since the add of zero pads gives -0.0 + 0.0 = +0.0 where the + concatenate keeps -0.0 (measured on XLA:CPU). The residual is `h01` itself, from which the backward recomputes the + slices, so no separate gate slice is kept for the backward. + """ + n = h01.shape[-1] // 2 + + def split(h): + return h[..., :n], h[..., n:] + + @jax.custom_vjp + def fused_act(h): + return act_fn(*split(h)) + + def fused_act_fwd(h): + return act_fn(*split(h)), h + + def fused_act_bwd(h, cotangent): + _, act_vjp = jax.vjp(act_fn, *split(h)) + d_gate, d_up = act_vjp(cotangent) + return (jnp.concatenate([d_gate, d_up], axis=-1),) + + fused_act.defvjp(fused_act_fwd, fused_act_bwd) + return fused_act(h01) + + @functools.partial(jax.custom_vjp, nondiff_argnums=(1,)) def _forward_layout_constraint(x, major_to_minor): """`with_layout_constraint` on the forward value and the identity on its cotangent.""" @@ -299,6 +339,35 @@ def forward_row_major(x: jax.Array) -> jax.Array: Only the forward value is constrained; the cotangent keeps whatever layout XLA picks. """ return _forward_layout_constraint(x, tuple(range(x.ndim))) +@functools.partial(jax.custom_vjp, nondiff_argnums=(2,)) +def take_along_last_axis_dense_vjp(x: jax.Array, indices: jax.Array, num_classes: int) -> jax.Array: + """`jnp.take_along_axis(x, indices, axis=-1)` whose backward is dense rather than a scatter-add. + + The transpose of `take_along_axis` is a scatter-add of the `[..., k]` cotangent into `[..., num_classes]`. On the + TPU, XLA sorts the flattened indices and then runs an element-serial scatter; for DeepSeek routing + ([tokens, 8] into [tokens, 256]) that is the largest VPU-lane op in the step. Here the backward is + `d x[..., e] = sum_k g[..., k] * (indices[..., k] == e)`, one fused compare-select-reduce. + + Requires `0 <= indices < num_classes` (true for `lax.top_k` output). When the indices along the last axis are + distinct (top-k), each output element receives at most one nonzero term, so the gradient equals the scatter-add + result exactly; with repeated indices it is the same sum in a different order. + """ + del num_classes + return jnp.take_along_axis(x, indices, axis=-1) + + +def _take_along_last_axis_dense_vjp_fwd(x, indices, num_classes): + del num_classes + return jnp.take_along_axis(x, indices, axis=-1), indices + + +def _take_along_last_axis_dense_vjp_bwd(num_classes, indices, g): + iota = jax.lax.broadcasted_iota(indices.dtype, (1,) * indices.ndim + (num_classes,), indices.ndim) + hit = indices[..., None] == iota # [..., k, num_classes] + return jnp.sum(jnp.where(hit, g[..., None], jnp.zeros((), g.dtype)), axis=-2), None + + +take_along_last_axis_dense_vjp.defvjp(_take_along_last_axis_dense_vjp_fwd, _take_along_last_axis_dense_vjp_bwd) def get_batchsplit_init_kernel_axes(): @@ -429,6 +498,25 @@ def _top_2_in_group_sum(scores_grouped: jax.Array) -> jax.Array: return max_1.astype(jnp.float32) + max_2.astype(jnp.float32) +def topk_indices_by_argmax(scores: jax.Array, k: int) -> jax.Array: + """`jax.lax.top_k(scores, k)[1]` without a sort (moe_routing_index_math="count"). + + k passes of argmax over the last axis, each masking its winner to -inf. argmax returns the first maximal index, + so ties go to the lower index, in the order top_k returns them. With k small against the expert count (10 of 512 + for Qwen3.5-397B) the k reductions are expected to cost less than the sort top_k lowers to; this was timed only as + part of moe_routing_index_math="count", not on its own. Needs at least k entries above -inf per row, true of the + router probabilities it is used on. + """ + lane = jax.lax.broadcasted_iota(jnp.int32, scores.shape, scores.ndim - 1) + masked = scores + picks = [] + for _ in range(k): + pick = jnp.argmax(masked, axis=-1).astype(jnp.int32) + picks.append(pick) + masked = jnp.where(lane == pick[..., None], jnp.array(-jnp.inf, scores.dtype), masked) + return jnp.stack(picks, axis=-1) + + def calculate_load_balance_updates(top_k_indices, num_experts, rate, axis_names=None): """Computes a bias adjustment update based on expert load. @@ -447,19 +535,165 @@ def calculate_load_balance_updates(top_k_indices, num_experts, rate, axis_names= Returns: update: The value to add to the expert bias. Shape (num_experts,). """ + expert_counts = calculate_expert_counts(top_k_indices, num_experts, axis_names=axis_names) + return load_balance_updates_from_counts(expert_counts, num_experts, rate) + + +def calculate_expert_counts(top_k_indices, num_experts, axis_names=None): + """Counts the token assignments routed to each expert. + + Args: + top_k_indices: Integer expert ids of any shape, e.g. (batch, sequence, top_k). + num_experts: Total number of experts. + axis_names: Optional mesh axis names to psum the counts across. + + Returns: + int32 array of shape (num_experts,). + """ flat_indices = top_k_indices.ravel() # one_hot rather than bincount: bincount clips out-of-range values, so the -1 # padding that forced routing uses would all be counted as expert 0. expert_counts = jnp.sum(jax.nn.one_hot(flat_indices, num_experts, dtype=jnp.int32), axis=0) if axis_names: expert_counts = jax.lax.psum(expert_counts, axis_names) - total_tokens = jnp.sum(expert_counts) + return expert_counts + + +def load_balance_updates_from_counts(expert_counts, num_experts, rate): + """Loss-free-balancing bias update, rate * sign(mean_load - load), from expert counts. + + Matches Megatron-LM `get_updated_expert_bias` when `expert_counts` covers the full global batch. + Supports shape (..., num_experts), e.g. (num_experts,) or scanned (num_layers, num_experts). + """ + total_tokens = jnp.sum(expert_counts, axis=-1, keepdims=True) average_load = total_tokens / num_experts direction = jnp.sign(average_load - expert_counts) output = direction * rate return output +def finalize_deferred_bias_signal(partial_counts, config): + """Reduces the stacked partial expert counts of defer_small_all_reduces into the routed-bias signal. + + Under defer_small_all_reduces each ring-of-experts MoE layer emits its local (per-shard) expert counts, per token + chunk, instead of all-reducing them inside the layer. The layer scan stacks them, so the value arrives here with + shape (..., num_parts, num_chunks, num_experts), where num_parts is the product of the mesh axes the counts were + summed over (the batch axes except EP) and the leading axes are the stacked layers, if any. The sum over num_parts + is the one all-reduce that replaces the per-layer ones. The rest repeats what the layer does without deferral: the + counts are summed over chunks, then either one rate * sign(mean - counts) per layer is taken + (load_balance_updates_from_counts), or, under gradient accumulation, the summed integer counts are returned so + gradient_accumulation_loss_and_grad sums them over microbatches before its single sign. Integer counts make the + result identical to the non-deferred path. + + Returns: + (..., num_experts): the float bias update, or the int32 counts when gradient_accumulation_steps > 1, as the + layer emits it without deferral. + """ + counts = jnp.sum(partial_counts, axis=(-3, -2)) # (..., num_experts) + if getattr(config, "gradient_accumulation_steps", 1) > 1: + return counts + return load_balance_updates_from_counts(counts, counts.shape[-1], config.routed_bias_update_rate) + + +def finalize_deferred_intermediates(intermediate_outputs, config): + """Applies finalize_deferred_bias_signal to every "moe_bias_updates" leaf when defer_small_all_reduces is on. + + The "moe_has_overflow" leaves need no rewrite: the layers emit them unreduced (per-device local flags) with or + without the flag, and loss_fn reduces them once with jnp.any. + """ + if not getattr(config, "defer_small_all_reduces", False): + return intermediate_outputs + + def _fix(path, leaf): + keys = [getattr(k, "key", getattr(k, "name", None)) for k in path] + if "moe_bias_updates" in keys and leaf is not None and getattr(leaf, "ndim", 0) >= 3: + return finalize_deferred_bias_signal(leaf, config) + return leaf + + return jax.tree_util.tree_map_with_path(_fix, intermediate_outputs) + + +def pack_routing_for_all_gather(weights: jax.Array, ids: jax.Array) -> jax.Array: + """`[..., k]` routing weights and expert ids as ONE float32 `[..., 2k]` array (moe_roe_packed_routing_gather). + + Exact: expert ids are integers in [-1, num_experts), and float32 holds every integer up to 2**24; float32 weights + are kept as they are and narrower ones convert exactly. `unpack_routing` inverts it. + """ + return jnp.concatenate([weights.astype(jnp.float32), ids.astype(jnp.float32)], axis=-1) + + +def unpack_routing(packed: jax.Array, weights_dtype, ids_dtype) -> tuple[jax.Array, jax.Array]: + """The routing weights and expert ids of a `pack_routing_for_all_gather` array, in their own dtypes.""" + k = packed.shape[-1] // 2 + return packed[..., :k].astype(weights_dtype), packed[..., k:].astype(ids_dtype) + + +def routing_tag_fn(config): + """Returns the checkpoint_name tagger for the routing arrays, or None under moe_routing_maps=remat. + + With moe_routing_maps=device (or offload) and remat_policy=custom, the tagged top-k expert ids and the ragged-sort + permutations (and, on the TC path with moe_tc_window_routing, the group sizes) are saved across the remat boundary, + so the backward re-runs neither top-k nor the sorts. Saving them instead of recomputing them does not change any value. + """ + if getattr(config, "moe_routing_maps", "remat") == "remat": + return None + return lambda t: adc.checkpoint_name(t, "moe_routing_maps") + + +def uses_megatron_seq_aux_loss(config) -> bool: + """Whether the MoE aux loss is Megatron-LM's sigmoid-router seq_aux_loss. + + When moe_use_megatron_seq_aux_loss=True, sigmoid routers (DeepSeek-V3 family) use it: + fp32 per-token normalized sigmoid scores, token fractions from an unbiased, ungrouped top-k, + one loss per full sequence, summed over MoE layers. By default (moe_use_megatron_seq_aux_loss=False), + and for other routers or te_moe_block (whose aux loss TransformerEngine computes), MaxText uses + the per-layer Switch-style loss averaged over MoE layers. + """ + return ( + getattr(config, "moe_use_megatron_seq_aux_loss", False) + and getattr(config, "routed_score_func", "") == "sigmoid" + and not getattr(config, "te_moe_block", False) + ) + + +def megatron_aux_loss_inputs(scores, num_experts_per_tok): + """Aux-loss probabilities and token assignments of the Megatron-LM sigmoid router. + + Mirrors Megatron-LM `compute_routing_scores_for_aux_loss(score_function="sigmoid")`: the + probabilities are the fp32 sigmoid scores normalized per token, and the token assignments + come from a plain top-k over the scores, i.e. without the expert bias and without the + group-limited mask used for dispatch. + + Args: + scores: Sigmoid router scores before the expert bias, shape (batch, sequence, num_experts). + num_experts_per_tok: Router top-k. + + Returns: + (aux_probs, aux_top_k_indices) with shapes (batch, sequence, num_experts) and + (batch, sequence, num_experts_per_tok). + """ + scores_f32 = scores.astype(jnp.float32) + aux_probs = scores_f32 / (jnp.sum(scores_f32, axis=-1, keepdims=True) + 1e-20) + _, aux_top_k_indices = jax.lax.top_k(scores_f32, k=num_experts_per_tok) + return aux_probs, aux_top_k_indices + + +def max_expert_shard_load_ratio(group_sizes, num_expert_shards, balanced_load): + """Largest per-expert-shard load divided by the perfectly balanced load. + + Args: + group_sizes: Token assignments per expert, shape (num_experts,), experts ordered by shard. + num_expert_shards: Number of expert-parallel shards. + balanced_load: Token assignments per shard under a perfectly balanced routing. + + Returns: + float32 scalar. The ring-of-experts ragged buffer overflows iff this exceeds + ragged_buffer_factor (up to the int() rounding of the buffer size). + """ + shard_loads = jnp.sum(group_sizes.reshape(num_expert_shards, -1), axis=-1) + return jnp.max(shard_loads).astype(jnp.float32) / balanced_load + + def _batch_axis_names(pspec) -> tuple[str, ...] | None: """Returns the mesh axes the (batch, sequence) dims are partitioned over.""" if not pspec: @@ -505,6 +739,27 @@ def _filter_axis_names( return filtered or None +def fused_moe_lhs_scale_kwargs(quant, fused_moe_fn) -> dict[str, jax.Array]: + """Static activation-scale kwargs for tpu-inference's `fused_moe_func`. + + With `serve_fp8_weight` and a fixed `act_quantization_calibration_method`, both expert + grouped matmuls quantize their input with the same static per-tensor scale as the dense + layers, instead of the kernel's dynamic per-row, per-K-block scale. Returns {} (dynamic + scaling, the kernel default) otherwise. + """ + if not isinstance(quant, quantizations.ServeFp8WeightQuantization): + return {} + scale = quant.static_act_scale() + if scale is None: + return {} + if "w1_lhs_scale" not in inspect.signature(fused_moe_fn).parameters: + raise ValueError( + f"act_quantization_calibration_method={quant.act_calibration_method!r} requires a tpu-inference whose" + " fused_moe_func accepts w1_lhs_scale / w2_lhs_scale; upgrade tpu-inference or use 'absmax'." + ) + return {"w1_lhs_scale": scale, "w2_lhs_scale": scale} + + class Tid2EidVar(nnx.Variable): """Custom variable to hold tid2eid without trainable param overhead.""" @@ -603,22 +858,20 @@ def __init__( self.bias = None if quant and not isinstance(quant, quantizations.ServeFp8WeightQuantization): - dot_general_cls = quant.dot_general_cls(mesh_axes=kernel_axes) - dot_general_linen = dot_general_cls() - quant_dot_general = nnx_wrappers.ToNNX(dot_general_linen, rngs=rngs) - self._quant_dot_general_name = f"{type(dot_general_linen).__name__}_0" + self._quant_dot_general_name, quant_dot_general = quantizations.make_dot_general(quant, kernel_axes, rngs) setattr(self, self._quant_dot_general_name, quant_dot_general) - block_size = getattr(quant, "get_block_size", lambda: 1)() # needed for TE block scaling - dummy_inputs = jnp.zeros((block_size, *self.in_features_shape), dtype=self.dtype) - self(dummy_inputs, _initializing=True) - # See the matching comment in linears.py. - if not quant.needs_apply_rngs: - quant_dot_general.release_rngs() + # See the matching comments in linears.py. + if isinstance(quant_dot_general, nnx_wrappers.ToNNX): + block_size = getattr(quant, "get_block_size", lambda: 1)() # needed for TE block scaling + dummy_inputs = jnp.zeros((block_size, *self.in_features_shape), dtype=self.dtype) + self(dummy_inputs, _initializing=True) + if not quant.needs_apply_rngs: + quant_dot_general.release_rngs() else: self._quant_dot_general_name = None @property - def quant_dot_general(self) -> nnx_wrappers.ToNNX | None: + def quant_dot_general(self) -> nnx.Module | None: if self._quant_dot_general_name is None: return None return getattr(self, self._quant_dot_general_name) @@ -635,6 +888,7 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax. else: kernel = self.kernel[...] kernel = jnp.asarray(kernel, self.dtype) + kernel = linears.route_dense_weight_grad(kernel, self.kernel_axes) contract_ind = tuple(range(0, len(norm_axis))) output_sharding = ( @@ -667,9 +921,239 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax. return output, pre_bias_logits +def _all_gather_quantized_payload( + x: jax.Array, + axis_name: str | tuple[str, ...], + *, + axis: int, + tiled: bool, + method: str, + qtype: jnp.dtype = jnp.float8_e5m2, + num_feature_axes: int = 1, +) -> jax.Array: + """jax.lax.all_gather with a quantized wire payload; returns dequantized x.dtype. + + Args: + x: Local shard to gather. + axis_name: Mesh axis name(s) to gather over. + axis: Axis of `x` to gather along. + tiled: Same as `jax.lax.all_gather`. + method: One of + 'fixed,': static per-tensor scale b/qmax. The scale is identical on all + shards, so only qvalues are gathered. For E5M2, |x| > ~b saturates and + |x| < ~1.3e-10*b flushes to zero. + 'rowwise': dynamic per-token absmax scale over the token axes (all but + the trailing `num_feature_axes`, e.g. (batch, seq) for 3D, (tokens,) + for 2D), matching the per-row gradient quantization in megablox + `_bwd_quantize_gradient`. Scales differ per shard, so they are + gathered alongside qvalues. + qtype: Wire dtype of the quantized payload. + num_feature_axes: Number of trailing axes that form one token's features. + 'rowwise' computes one scale per token over these axes, e.g. 2 for the + (batch, seq, emb // 128, 128) layout of moe_tc_ragged_3d_dispatch. + + Returns: + The gathered, dequantized array in `x.dtype`. + """ + static = method.startswith("fixed") + token_axes = tuple(range(x.ndim - num_feature_axes)) + if not static and method != "rowwise": + raise ValueError(f"Unsupported method: {method!r}. Supported: 'rowwise', 'fixed,'.") + if not static and axis % x.ndim not in token_axes: + raise ValueError(f"'rowwise' requires gathering along a token axis, got axis={axis}.") + + x_f32 = x.astype(jnp.float32) + if static: + x_q = qpl.quantize( + x_f32, + qtype=qtype, + channelwise_axes=(), + calibration_method=method, + ) + scale = x_q.scale + else: # dynamic rowwise with per-token scale + x_q = qpl.quantize( + x_f32, + qtype=qtype, + channelwise_axes=token_axes, + calibration_method="absmax", + ) + scale = jax.lax.all_gather(x_q.scale, axis_name=axis_name, axis=axis, tiled=tiled) + gathered_qvals = jax.lax.all_gather(x_q.qvalue, axis_name=axis_name, axis=axis, tiled=tiled) + return (gathered_qvals.astype(jnp.float32) * scale).astype(x.dtype) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3, 4)) +def _moe_combine_psum_scatter( + x: jax.Array, + axis_name: str | tuple[str, ...], + scatter_dimension: int = 0, + tiled: bool = True, + bwd_method: str = "", +) -> jax.Array: + """Evaluates psum_scatter in forward, and quantized All-Gather in backward. + + Shapes: + Forward: + x: [global_batch, sequence_length, hidden_dim], where global_batch is + the full batch across all expert parallelism (EP) shards. + Output: [local_batch, sequence_length, hidden_dim], scattered along + scatter_dimension=0 (local_batch = global_batch // num_ep). + Backward: + grads: [local_batch, sequence_length, hidden_dim], incoming cotangent. + Output: [global_batch, sequence_length, hidden_dim], gathered along + scatter_dimension=0. + With moe_tc_ragged_3d_dispatch, hidden_dim is split as (hidden_dim // 128, 128). + """ + return jax.lax.psum_scatter(x, axis_name=axis_name, scatter_dimension=scatter_dimension, tiled=tiled) + + +def _moe_combine_psum_scatter_fwd( + x: jax.Array, + axis_name: str | tuple[str, ...], + scatter_dimension: int = 0, + tiled: bool = True, + bwd_method: str = "", +) -> tuple[jax.Array, None]: + """Custom VJP forward pass for _moe_combine_psum_scatter.""" + return _moe_combine_psum_scatter(x, axis_name, scatter_dimension, tiled, bwd_method), None + + +def _moe_combine_psum_scatter_bwd( + axis_name: str | tuple[str, ...], + scatter_dimension: int, + tiled: bool, + bwd_method: str, + res: Any, + grads: jax.Array, +) -> tuple[jax.Array]: + """Custom VJP backward pass for _moe_combine_psum_scatter.""" + if not bwd_method: + gathered_grads = jax.lax.all_gather(grads, axis_name=axis_name, tiled=tiled, axis=scatter_dimension) + return (gathered_grads,) + + # The combine tensor is (batch, seq, *features): features are (emb,) or, with moe_tc_ragged_3d_dispatch, + # (emb // 128, 128). 'rowwise' keeps one scale per token in both layouts. + gathered_grads = _all_gather_quantized_payload( + grads, axis_name, axis=scatter_dimension, tiled=tiled, method=bwd_method, num_feature_axes=grads.ndim - 2 + ) + return (gathered_grads,) + + +_moe_combine_psum_scatter.defvjp(_moe_combine_psum_scatter_fwd, _moe_combine_psum_scatter_bwd) + + +def _param_path(path) -> tuple: + """A params pytree path as the module path keys `nnx.iter_modules` yields, plus the param name. + + Drops the trailing `GetAttrKey("value")` of each Variable leaf. + """ + keys = [] + for k in path: + if isinstance(k, jax.tree_util.DictKey): + keys.append(k.key) + elif isinstance(k, jax.tree_util.SequenceKey): + keys.append(k.idx) + return tuple(keys) + + +def _gathered_kernels(layers: nnx.Module) -> dict: + """{kernel path within a scanned layer stack: its module}, for the kernels gathered once per stack.""" + kernels = {} + for path, module in nnx.iter_modules(layers): + if isinstance(module, GateLogit): + kernels[path + ("kernel",)] = module + elif isinstance(module, attentions.Attention) and hasattr(module, "query"): + # The backward pass otherwise gathers the query kernel twice per layer: for the rematerialized + # projection and for the input gradient. Under fused_qkv there is no `query` (it is `qkv_proj`). + kernels[path + ("query", "kernel")] = module.query + return kernels + + +def gather_scanned_kernels_once(layers: nnx.Module, params, mesh: Mesh, shard_mode): + """Gathers small FSDP-sharded kernels of a scanned layer stack once, before the scan, not in every layer. + + The MoE router kernel (embed x num_experts, f32 under float32_gate_logits) is small but FSDP-sharded, + so every layer all-gathers it on the critical path, synchronously on TPUs that run all-gathers on the + TensorCore, and the backward pass gathers it again. Replicating the stacked kernels up front is one + all-gather for all layers. `params` is the stack's Param state with the scan axis leading; pair with + `reduce_scatter_gathered_grads` inside the scan so each layer's gradient still lands sharded. + """ + gathered = _gathered_kernels(layers) + replicated = NamedSharding(mesh, P()) + + def replicate(path, x): + if _param_path(path) not in gathered: + return x + if shard_mode == ShardMode.EXPLICIT: + return jax.sharding.reshard(x, replicated) + return jax.lax.with_sharding_constraint(x, replicated) + + return jax.tree_util.tree_map_with_path(replicate, params) + + +def gathered_kernel_grad_shardings(layers: nnx.Module, mesh: Mesh, rules) -> dict: + """{kernel path within one layer: the kernel's own sharding} for the kernels `gather_scanned_kernels_once` gathers.""" + return { + path: NamedSharding(mesh, logical_to_mesh_axes(tuple(module.kernel_axes), mesh, rules)) + for path, module in _gathered_kernels(layers).items() + } + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1,)) +def _identity_with_grad_sharding(x, sharding): + """The identity on `x`; its cotangent is constrained to `sharding`.""" + return x + + +def _identity_with_grad_sharding_fwd(x, sharding): + del sharding + return x, None + + +def _identity_with_grad_sharding_bwd(sharding, _, cotangent): + return (jax.lax.with_sharding_constraint(cotangent, sharding),) + + +_identity_with_grad_sharding.defvjp(_identity_with_grad_sharding_fwd, _identity_with_grad_sharding_bwd) + + +def reduce_scatter_gathered_grads(params, shardings: dict): + """Identity on one layer's params whose backward puts the gathered kernels' gradients back on their sharding. + + A layer handed a slice of a replicated kernel otherwise gets a replicated gradient: an all-reduce per + layer, stacked at full size across the scan. Constraining it to the kernel's own sharding makes it a + reduce-scatter into a stacked shard, as when the layer gathered the kernel itself. + """ + + def apply(path, x): + keys = _param_path(path) + return _identity_with_grad_sharding(x, shardings[keys]) if keys in shardings else x + + return jax.tree_util.tree_map_with_path(apply, params) + + class RoutedMoE(nnx.Module): """Implements a routed MoE block.""" + def required_ragged_buffer_factor(self, group_sizes, bsz_times_seq_len, num_ep, expert_shard_id): + """log_required_ragged_buffer_factor probe: the smallest ragged_buffer_factor that would not have dropped tokens. + + group_sizes are the global per-expert token counts of the (EP-gathered) batch, as ring_ragged_sort computes them + for the overflow check. The per-shard buffer is int(balanced_size * factor) rows, so a shard drops tokens iff + sum(local group sizes) > balanced_size * factor. The result is max over all devices of + sum(local group sizes) / balanced_size: one float32 scalar max all-reduce over the mesh axes (no other change). + """ + local_num_experts = self.config.num_experts // num_ep + local_group_sizes = jax.lax.dynamic_slice_in_dim( + group_sizes, expert_shard_id * local_num_experts, local_num_experts, axis=0 + ) + balanced_size = (bsz_times_seq_len // num_ep) * self.num_experts_per_tok + required = jnp.sum(local_group_sizes).astype(jnp.float32) / jnp.float32(balanced_size) + if self.mesh is not None and self.mesh.axis_names: + required = jax.lax.pmax(required, tuple(self.mesh.axis_names)) + return required + def __init__( self, config: ctypes.Config, @@ -763,6 +1247,13 @@ def __init__( else: self.wi_kernel_axes = ("exp", "embed_moe", "mlp_moe") self.wo_kernel_axes = ("exp", "mlp_moe", "embed_moe") + if self.config.moe_flat_fsdp_weights: + # Flat layout: wi [E * embed, mlp], wo [E * mlp, embed]; dim 0 is sharded over (expert, fsdp) via the + # compound logical axis ("exp", "embed_moe") so every FSDP shard is a contiguous tile-aligned row block + # and the FSDP-gathered weight is a free bitcast of the [E_local, rows, cols] operand the gmm kernels pin. + # wo's rows are mlp but reuse the "embed_moe" logical name to get the same (expert, fsdp) dim-0 sharding. + self.wi_kernel_axes = (CompoundLogicalAxis(("exp", "embed_moe")), "mlp_moe") + self.wo_kernel_axes = (CompoundLogicalAxis(("exp", "embed_moe")), None) if self.config.attention in ("vllm_rpa", "vllm_batched_rpa", "vllm_batched_rpa_long_ctx"): # vLLM uses 'model' as the tensor parallelism axis name @@ -775,8 +1266,12 @@ def __init__( if self.mesh is not None and "dcp" in self.mesh.shape and self.mesh.shape["dcp"] > 1: expert_axes.append("dcp") self._expert_parallelism_name = tuple(expert_axes) - elif self.config.custom_mesh_and_rule == ctypes.CustomRule.CP_AS_EP: - # when custom mesh and rule is cp-as-ep, context axis is same with expert in MoE component + elif self.config.custom_mesh_and_rule in ( + ctypes.CustomRule.CP_AS_EP, + ctypes.CustomRule.CP_AS_EP_DENSE_NO_CP, + ctypes.CustomRule.CP_AS_EP_DENSE_FSDP, + ): + # when custom mesh and rule is cp-as-ep (or a variant of it), context axis is same with expert in MoE component self._expert_parallelism_name = ("context", "expert") else: self._expert_parallelism_name = "expert" @@ -884,38 +1379,25 @@ def __init__( out_sharding=self.wo_kernel_axes, ) else: + + def _init_expert_kernel(shape, dtype): + # Always init in 3D [E, rows, cols] so the fan-in/out (and RNG draw) match the 3D layout, then view it + # as the flat [E * rows, cols] matrix under moe_flat_fsdp_weights (see wi_kernel_axes above). + kernel = self.kernel_init(self.rngs.params(), shape, dtype, kernel_in_axis, kernel_out_axis) + if self.config.moe_flat_fsdp_weights: + kernel = kernel.reshape(shape[0] * shape[1], shape[2]) + return kernel + self.wi_0 = nnx.Param( - self.kernel_init( - self.rngs.params(), - (num_experts, self.moe_expert_input_dim, moe_intermediate_dim), - weight_dtype, - kernel_in_axis, - kernel_out_axis, - ), + _init_expert_kernel((num_experts, self.moe_expert_input_dim, moe_intermediate_dim), weight_dtype), out_sharding=self.wi_kernel_axes, ) self.wi_1 = nnx.Param( - self.kernel_init( - self.rngs.params(), - (num_experts, self.moe_expert_input_dim, moe_intermediate_dim), - weight_dtype, - kernel_in_axis, - kernel_out_axis, - ), + _init_expert_kernel((num_experts, self.moe_expert_input_dim, moe_intermediate_dim), weight_dtype), out_sharding=self.wi_kernel_axes, ) self.wo = nnx.Param( - self.kernel_init( - self.rngs.params(), - ( - self.num_experts, - moe_intermediate_dim, - self.moe_expert_input_dim, - ), - self.weight_dtype, - kernel_in_axis, - kernel_out_axis, - ), + _init_expert_kernel((self.num_experts, moe_intermediate_dim, self.moe_expert_input_dim), self.weight_dtype), out_sharding=self.wo_kernel_axes, ) @@ -1136,7 +1618,12 @@ def get_topk( valid_token_mask = None if forced_routed_experts is not None: router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1) - _, auto_indices = jax.lax.top_k(router_probs, self.num_experts_per_tok) + if self.config.moe_routing_index_math == "count": + # The router's own choice for tokens without (valid) forced ids: the same indices as top_k, from k argmax + # passes instead of a sort. Under moe_exclude_padding_from_routing every token takes this branch. + auto_indices = topk_indices_by_argmax(router_probs, self.num_experts_per_tok) + else: + _, auto_indices = jax.lax.top_k(router_probs, self.num_experts_per_tok) slot_valid = valid_expert_mask(forced_routed_experts, self.num_experts) # A token has real forced routing only if its slots are within [0, num_experts). # Tokens with UNSET_ROUTED_EXPERT (-1) slots must fall back to auto_indices @@ -1301,7 +1788,14 @@ def deepseek_routing(self, gate_logits: jax.Array, pre_bias_logits: jax.Array) - jnp.where(expert_mask > 0, gate_logits, -jnp.inf), k=self.num_experts_per_tok, ) - top_k_weights = jnp.take_along_axis(pre_bias_logits, top_k_indices, axis=-1) + # Tag before the weight gather, so the gather's backward reads the saved ids instead of re-running top-k. + tag_routing_fn = routing_tag_fn(self.config) + if tag_routing_fn is not None: + top_k_indices = tag_routing_fn(top_k_indices) + if getattr(self.config, "router_topk_matmul_vjp", False): + top_k_weights = take_along_last_axis_dense_vjp(pre_bias_logits, top_k_indices, pre_bias_logits.shape[-1]) + else: + top_k_weights = jnp.take_along_axis(pre_bias_logits, top_k_indices, axis=-1) return top_k_weights, top_k_indices def apply_ffn_activation(self, layer_w0, layer_w1): @@ -1335,6 +1829,43 @@ def apply_ffn_activation(self, layer_w0, layer_w1): intermediate_layer = jnp.multiply(layer_act, layer_w1) return intermediate_layer.astype(self.dtype) + def apply_ffn_activation_maybe_fused(self, layer_w0, layer_w1): + """`apply_ffn_activation`; `layer_w1=None` means `layer_w0` is the whole [gate | up] up-projection output. + + gmm_up returns that layout under moe_prefused_swiglu_vjp, and the activation then runs as one custom-VJP op + (`_prefused_ffn_activation`) whose backward writes d[gate | up] in one pass. + """ + if layer_w1 is None: + return _prefused_ffn_activation(self.apply_ffn_activation, layer_w0) + return self.apply_ffn_activation(layer_w0, layer_w1) + + def _keep_prefused_wi_whole(self): + """moe_unsplit_prefused_wi: whether the prefused `wi = [w0 | w1]` stays one tensor up to the up-projection GEMM. + + Then `wi` is fake-quantized, FSDP-all-gathered and carried through the expert-weight prefetch barrier as one + tensor, and the GEMM reads it directly (`w1_kernel=None` marks the layout), instead of being sliced into w0/w1 + here and concatenated back after the gather. The config validator admits the key only where every reader of + w0/w1 handles that layout, including an expert MLP dimension that the default logical axis rules leave + unsharded; this checks what only the built layer knows (quantization objects, sparsity, the partition spec under + the actual axis rules), and raises rather than fall back. + """ + if not self.config.moe_unsplit_prefused_wi: + return False + blockers = [] + if self.quant is not None: + blockers.append("weight quantization") + if self.wi_0_sparsity_module is not None: + blockers.append("weight sparsity") + # The whole [w0 | w1] last dim must be on every device of the expert shard_map: sharded, a device would hold a + # slice of [w0 | w1] instead of matching slices of w0 and w1. Same pspec as sparse_matmul's default branch. + wi_pspec = remove_mesh_axes_from_partition_spec(self._logical_to_mesh_axes(("exp", None, "mlp_no_fsdp")), ("fsdp",)) + f_shards = math.prod(self.mesh.shape[a] for a in mesh_axes_for_dim(wi_pspec[2]) if a in self.mesh.shape) + if f_shards != 1: + blockers.append(f"the MLP dimension sharded {f_shards}-way in the expert shard_map") + if blockers: + raise ValueError(f"moe_unsplit_prefused_wi=True cannot apply: {', '.join(blockers)}.") + return True + def permute( self, inputs, @@ -1348,34 +1879,69 @@ def permute( force_dropless=False, mesh_axis_names=None, topk=None, + compute_lb_loss=True, + return_expert_counts=False, + defer_reductions=False, + precomputed_topk=None, + precomputed_density_prob=None, ): """Permute tokens to group by expert to fit gmm call. - `topk`, if given, is `(weights, selected_experts, lb_loss)` already computed for `inputs`. Then - `gate_logits`, `pre_bias_logits`, `rngs` and `input_ids` are not used, and of - `forced_routed_experts` only whether it is None matters: it says padded slots may be present. + `topk`, if given, is `(weights, selected_experts, lb_loss)` already computed for `inputs` + (ring_of_experts_local_routing). Then `gate_logits`, `pre_bias_logits`, `rngs` and `input_ids` are not + used, and of `forced_routed_experts` only whether it is None matters: it says padded slots may be present. + + Args: + compute_lb_loss: If False, skip the load-balance loss here (returned as None) because the + caller computes the sigmoid-router loss on the full sequences. + return_expert_counts: If True, the bias-update slot of the return value holds the raw expert + counts (psum'd over `mesh_axis_names`) instead of the per-call sign update, so the caller + can sum them over token chunks and take a single sign over the full batch. + defer_reductions: defer_small_all_reduces. The routed-bias signal is this shard's local expert + counts, not all-reduced here (see finalize_deferred_bias_signal). Takes precedence over + return_expert_counts. + precomputed_topk: Optional (weights, selected_experts) already computed by the caller; skips + get_topk. gate_logits / pre_bias_logits are then unused. + precomputed_density_prob: Optional per-sequence mean router probabilities (batch, num_experts) + used for the load-balance loss instead of recomputing them from the logits. """ is_qarray = isinstance(inputs, qpl.QArray) raw_inputs = inputs.qvalue if is_qarray else inputs # reshape inputs (batch, sequence, emb) to (batch * sequence, emb) inputs_shape = raw_inputs.shape bsz_times_seq_len = inputs_shape[0] * inputs_shape[1] - inputs_2d = jnp.reshape(raw_inputs, (bsz_times_seq_len, inputs_shape[2])) - if topk is None: - weights, selected_experts = self.get_topk(gate_logits, pre_bias_logits, rngs, input_ids, forced_routed_experts) - lb_loss = None - if self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing: - lb_loss = self.router_load_balance_loss(selected_experts, gate_logits, pre_bias_logits) - else: + # moe_tc_ragged_3d_dispatch: inputs arrive as (batch, seq, emb // 128, 128) and stay 3D for the TC sort. + inputs_2d = jnp.reshape(raw_inputs, (bsz_times_seq_len, *inputs_shape[2:])) + if topk is not None: weights, selected_experts, lb_loss = topk + else: + if precomputed_topk is not None: + weights, selected_experts = precomputed_topk + else: + weights, selected_experts = self.get_topk(gate_logits, pre_bias_logits, rngs, input_ids, forced_routed_experts) + lb_loss = None + if compute_lb_loss and self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing: + if precomputed_density_prob is not None: + lb_loss = self.load_balance_loss(selected_experts, None, density_prob=precomputed_density_prob) + else: + lb_loss = self.router_load_balance_loss(selected_experts, gate_logits, pre_bias_logits) if self.should_update_load_balance(): - bias_updates = calculate_load_balance_updates( - selected_experts, - self.config.num_experts, - self.config.routed_bias_update_rate, - axis_names=mesh_axis_names, - ) + if defer_reductions: + bias_updates = calculate_expert_counts(selected_experts, self.config.num_experts) + elif return_expert_counts: + bias_updates = calculate_expert_counts( + selected_experts, + self.config.num_experts, + axis_names=mesh_axis_names, + ) + else: + bias_updates = calculate_load_balance_updates( + selected_experts, + self.config.num_experts, + self.config.routed_bias_update_rate, + axis_names=mesh_axis_names, + ) else: bias_updates = None @@ -1395,6 +1961,9 @@ def permute( # local prefix of valid rows. use_ragged_in_permute = self.config.use_ragged_sort and self.config.use_ring_of_experts buffer_size = None + topk_argsort_indices = None + use_tc_sort = False + tc_routing = None if use_ragged_in_permute: topk_indices_2d = jnp.reshape(selected_experts, (bsz_times_seq_len, selected_experts.shape[2])) if forced_routed_experts is not None: @@ -1407,20 +1976,54 @@ def permute( self.config.num_experts, ) # roll_to_expert_id is not directly used in the kernel, ep axis id is directly called - rbf = self.get_ragged_buffer_factor() - if not force_dropless and rbf > 0.0: + ragged_buffer_factor = self.get_ragged_buffer_factor() + if not force_dropless and ragged_buffer_factor > 0.0: balanced_size = (bsz_times_seq_len // num_expert_parallelism) * self.num_experts_per_tok buffer_size = self.get_ragged_buffer_size( balanced_size, num_expert_parallelism, self.config.num_experts, self.num_experts_per_tok, - rbf, + ragged_buffer_factor, ) else: buffer_size = None - sorted_inputs, group_size, sorted_selected_experts = ring_ragged_sort( + use_tc_sort = ( + getattr(self.config, "moe_tc_ragged_sort", False) + and buffer_size is not None + and buffer_size < bsz_times_seq_len * self.num_experts_per_tok + ) + if getattr(self.config, "moe_tc_ragged_sort", False) and not use_tc_sort: + # E.g. the dropless program (force_dropless / moe_dropless_fallback) or a buffer covering every slot. + max_logging.warning( + "moe_tc_ragged_sort=True but the TensorCore ragged sort needs a truncated buffer " + f"(buffer_size={buffer_size}, total slots={bsz_times_seq_len * self.num_experts_per_tok}); " + "using the SparseCore ragged sort for this call." + ) + if inputs_2d.ndim == 3 and not use_tc_sort: + # 3D dispatch layout but this call does not use the TC sort (e.g. dropless fallback): back to 2D. + inputs_2d = jnp.reshape(inputs_2d, (bsz_times_seq_len, -1)) + if use_ragged_in_permute and use_tc_sort: + sorted_inputs, group_size, sorted_selected_experts, tc_routing = ring_ragged_sort_tc( + inputs_2d, + topk_indices_2d, + self.config.num_experts, + self.num_experts_per_tok, + self._expert_parallelism_name, + num_expert_parallelism, + buffer_size, + gather_block_size=self.config.moe_tc_ragged_gather_block_size, + reduce_block_size=self.config.moe_tc_ragged_reduce_block_size, + mask_padding=self.config.moe_tc_ragged_mask_padding, + flatten_block_size=self.config.moe_tc_ragged_flatten_block_size, + tag_routing_fn=routing_tag_fn(self.config), + keep_3d=self.config.moe_tc_ragged_3d_gmm, + window_routing=self.config.moe_tc_window_routing, + ) + elif use_ragged_in_permute: + tag_routing_fn = routing_tag_fn(self.config) + sorted_inputs, group_size, sorted_selected_experts, topk_argsort_indices = ring_ragged_sort( inputs_2d, topk_indices_2d, self.config.num_experts, @@ -1435,7 +2038,15 @@ def permute( gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, + tag_routing_fn=tag_routing_fn, + index_math=self.config.moe_routing_index_math, + unique_pair_argsort=self.config.moe_ragged_sort_unique_pair_argsort, ) + if tag_routing_fn is None and self.config.moe_routing_index_math == "sort": + # moe_routing_maps=remat: unpermute re-derives the permutation in its backward (existing behavior). + # moe_routing_index_math="count" always passes it: its inverse covers only this shard's slots, so it + # cannot be argsorted back, and under remat the replay recomputes it anyway (no memory is saved). + topk_argsort_indices = None else: flatten_selected_experts = jnp.ravel(selected_experts) @@ -1472,14 +2083,24 @@ def permute( num_tokens = bsz_times_seq_len * self.num_experts_per_tok use_truncated_buffer = use_ragged_in_permute and buffer_size is not None and buffer_size < num_tokens has_overflow = False + max_load_ratio = None + if use_ragged_in_permute and getattr(self.config, "moe_log_max_load_ratio", False): + # group_size is the untruncated per-expert count of this EP group's (all-gathered) tokens, + # so this is the load every shard of the group would receive with an unbounded buffer. + balanced_load = (bsz_times_seq_len / num_expert_parallelism) * self.num_experts_per_tok + max_load_ratio = max_expert_shard_load_ratio(group_size, num_expert_parallelism, balanced_load) + if self.mesh is not None and self.mesh.axis_names: + max_load_ratio = jax.lax.pmax(max_load_ratio, tuple(self.mesh.axis_names)) moe_audit = None if use_ragged_in_permute and getattr(self.config, "log_moe_drops", False): # Drop audit (log_moe_drops), logging only. group_size is the untruncated per-expert count of this EP # group's (all-gathered) tokens, so a shard's sum is what it would receive with an unbounded buffer. # balanced_load uses the same integer expression as the buffer sizing above (buffer_size = # int(balanced_size * rbf)), so overflow at a factor rbf is exactly max_shard_load > balanced_load * rbf. - balanced_load = (bsz_times_seq_len // num_expert_parallelism) * self.num_experts_per_tok - max_load_ratio = max_expert_shard_load_ratio(group_size, num_expert_parallelism, balanced_load) + audit_balanced_load = (bsz_times_seq_len // num_expert_parallelism) * self.num_experts_per_tok + audit_ratio = max_expert_shard_load_ratio( + group_size, num_expert_parallelism, jnp.float32(audit_balanced_load) + ) if use_truncated_buffer: # This device's own shard: the assignments past the buffer `_truncate_matrix` cuts below (the same # buffer_size), i.e. max(0, shard_load - buffer_size). Summed over the mesh, every EP shard of every @@ -1487,13 +2108,13 @@ def permute( audit_shard = jax.lax.axis_index(self._expert_parallelism_name) if num_expert_parallelism > 1 else 0 audit_local = self.config.num_experts // num_expert_parallelism shard_load = jnp.sum(jax.lax.dynamic_slice_in_dim(group_size, audit_shard * audit_local, audit_local, axis=0)) - dropped = jnp.maximum(shard_load.astype(jnp.int32) - jnp.int32(buffer_size), 0) + audit_dropped = jnp.maximum(shard_load.astype(jnp.int32) - jnp.int32(buffer_size), 0) else: - dropped = jnp.int32(0) # dropless: buffer sized for the worst case + audit_dropped = jnp.int32(0) # dropless: buffer sized for the worst case if self.mesh is not None and self.mesh.axis_names: - max_load_ratio = jax.lax.pmax(max_load_ratio, tuple(self.mesh.axis_names)) - dropped = jax.lax.psum(dropped, tuple(self.mesh.axis_names)) - moe_audit = (dropped, max_load_ratio) + audit_ratio = jax.lax.pmax(audit_ratio, tuple(self.mesh.axis_names)) + audit_dropped = jax.lax.psum(audit_dropped, tuple(self.mesh.axis_names)) + moe_audit = (audit_dropped, audit_ratio) if use_truncated_buffer: local_num_experts = self.config.num_experts // num_expert_parallelism shard_idx = jax.lax.axis_index(self._expert_parallelism_name) if num_expert_parallelism > 1 else 0 @@ -1504,14 +2125,12 @@ def permute( local_num_experts, axis=0, ) - local_overflow = (jnp.sum(local_group_size) > buffer_size).astype(jnp.int32) + has_overflow = jnp.sum(local_group_size) > buffer_size # Clamp local_group_size to buffer_size to ensure we don't exceed buffer # capacity by leveraging the helper _truncate_matrix. local_group_size = _truncate_matrix(local_group_size[:, None], buffer_size)[:, 0] - if self.mesh is not None and self.mesh.axis_names: - has_overflow = jax.lax.psum(local_overflow, tuple(self.mesh.axis_names)) > 0 - else: - has_overflow = local_overflow > 0 + if use_tc_sort and self.config.moe_tc_window_routing and (tag_fn := routing_tag_fn(self.config)): + local_group_size = tag_fn(local_group_size) expert_indices = jnp.arange(local_num_experts) sorted_experts = jnp.repeat( expert_indices, @@ -1541,6 +2160,9 @@ def permute( bias_updates, local_group_size, has_overflow, + max_load_ratio, + topk_argsort_indices, + tc_routing, moe_audit, ) @@ -1553,10 +2175,33 @@ def unpermute( sequence_length, use_custom_sort_vjp=True, group_sizes=None, + topk_argsort_indices=None, + tc_routing=None, + tc_prescaled=False, + tc_out_3d=False, ): - """Unpermute tokens to original order and combine weights.""" + """Unpermute tokens to original order and combine weights. - if self.config.use_ragged_sort and self.config.use_ring_of_experts: + With tc_out_3d (moe_tc_ragged_3d_dispatch), the TC combine stays in the (emb // 128, 128) layout and + the result is (batch, seq, emb // 128, 128). + """ + + if tc_routing is not None: + output = ring_ragged_unsort_tc( + intermediate, + tc_routing, + self.num_experts_per_tok, + jnp.ravel(weights).astype(jnp.float32), + gather_block_size=self.config.moe_tc_ragged_gather_block_size, + mask_padding=self.config.moe_tc_ragged_mask_padding, + flatten_block_size=self.config.moe_tc_ragged_flatten_block_size, + prescaled=tc_prescaled, + out_3d=tc_out_3d, + ) + if tc_out_3d: + # Rows cover the EP-gathered batch (batch_size is the local batch), so the leading dim is inferred. + return output.reshape(-1, sequence_length, *output.shape[1:]).astype(self.dtype) + elif self.config.use_ragged_sort and self.config.use_ring_of_experts: local_num_experts = self.config.num_experts // self.get_expert_parallelism_size() # Build the flat routing weights in the same layout as # topk_argsort_revert_indices (i.e. the flat token×topk order before @@ -1577,6 +2222,9 @@ def unpermute( gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, + topk_argsort_indices=topk_argsort_indices, + index_math=self.config.moe_routing_index_math, + unique_pair_argsort=self.config.moe_ragged_sort_unique_pair_argsort, ) else: unsort_intermediate = _sort_activations( @@ -1907,15 +2555,18 @@ def transform_bias(self, experts_index, *biases): """Selects bias values for a variable number of bias tensors based on chosen experts.""" return tuple(bias[experts_index] for bias in biases) - def get_ragged_buffer_factor(self) -> float: - """Returns the ragged buffer factor to use for this RoutedMoE instance. + def get_ragged_buffer_factor(self): + """Reads ragged_buffer_factor for this module. - If `ragged_buffer_factor_override` is set on this instance (e.g. for a first-phase - graphdef), that is used. Otherwise `self.config.ragged_buffer_factor` is used. + A per-graphdef `ragged_buffer_factor_override` (set post-construction on the first-phase graphdef via + first_phase_ragged_buffer_factor, or on the eval graphdef via eval_ragged_buffer_factor) wins when it is > 0. + Otherwise `eval_ragged_buffer_factor` is used under eval axis rules, and `ragged_buffer_factor` elsewhere. """ override = getattr(self, "ragged_buffer_factor_override", None) if override is not None and override > 0: return override + if max_utils.is_eval(self.config): + return self.config.eval_ragged_buffer_factor return self.config.ragged_buffer_factor @staticmethod @@ -1954,6 +2605,41 @@ def get_ragged_buffer_size(local_batch, ep_degree, global_experts, top_k, ragged worst_case_factor = min(ep_degree, global_experts / top_k) return int(balanced_size * worst_case_factor) + def get_dropless_fallback_num_chunks(self, seq_len, n_chunks): + """Chunk count for the in-layer dropless branch, or None if the ragged buffer can never overflow. + + A dropless chunk allocates EP * balanced rows versus factor * balanced for the capped path, so + n_chunks * EP / factor chunks (rounded up to a divisor of seq_len) keep the fallback's per-chunk + buffer, and hence peak HBM, no larger than the capped path's. + """ + ep = self.get_expert_parallelism_size() + factor = self.get_ragged_buffer_factor() + if factor <= 0.0 or factor >= min(ep, self.config.num_experts / self.num_experts_per_tok): + return None + target = min(math.ceil(n_chunks * ep / factor), seq_len) + return next(d for d in range(target, seq_len + 1) if seq_len % d == 0) + + def _check_ragged_overflow(self, logits, pre_bias_logits, rngs, input_ids, forced_routed_experts, n_chunks): + """Returns True if permute()'s ragged buffer would overflow for any chunk on any EP shard. + + Counts this device's (token, k) pairs per chunk per destination EP shard, all-reduces the counts + over EP (what each shard receives after the all-gather) and compares them with the buffer size. + """ + num_ep = self.get_expert_parallelism_size() + _, selected_experts = self.get_topk(logits, pre_bias_logits, rngs, input_ids, forced_routed_experts) + batch, seq_len, top_k = selected_experts.shape + # Forced-routing padding (-1) maps to no shard, so one_hot gives it an all-zero row. + target_shard = selected_experts // (self.config.num_experts // num_ep) + counts = jax.nn.one_hot(target_shard, num_ep, dtype=jnp.int32) + counts = counts.reshape(batch, n_chunks, -1, num_ep).sum(axis=(0, 2)) # [n_chunks, EP] + if num_ep > 1: + counts = jax.lax.psum(counts, self._expert_parallelism_name) + buffer_size = self.get_ragged_buffer_size( + batch * (seq_len // n_chunks) * top_k, num_ep, self.config.num_experts, top_k, self.get_ragged_buffer_factor() + ) + overflow = jnp.any(counts > buffer_size).astype(jnp.int32) + return jax.lax.psum(overflow, tuple(self.mesh.axis_names)) > 0 + def sparse_matmul( self, inputs, @@ -1968,10 +2654,27 @@ def sparse_matmul( input_ids=None, forced_routed_experts=None, ): - """Perform sparse matrix multiplication of inputs and Experts.""" + """Perform sparse matrix multiplication of inputs and Experts. + + `w1_kernel=None` (moe_unsplit_prefused_wi) means `w0_kernel` is the whole prefused [w0 | w1] weight. + """ + if w1_kernel is None and not self.config.prefuse_moe_weights: + raise ValueError("sparse_matmul got w1_kernel=None (the whole [w0 | w1] layout), which needs prefuse_moe_weights.") # log_moe_drops only: the shard_map gets a 5th (moe_audit) output only when requested, so the # compiled program is unchanged with the flag off. log_moe_drops = getattr(self.config, "log_moe_drops", False) + # moe_router_grad_from_wo_input: the ring-of-experts down projection and unsort run as one op + # (_ring_down_project_and_unpermute). The config validator checks the rest; the mesh decides this. + router_grad_from_wo_input = self.config.moe_router_grad_from_wo_input + if router_grad_from_wo_input and self.get_tensor_parallelism_size() > 1: + raise ValueError( + "moe_router_grad_from_wo_input=True does not support tensor parallelism in the MoE layer: the down" + " projection must be the bare GEMM of the sorted rows, without its reduce-scatter." + ) + + def fwd_tile(name): + """Reads a forward GMM tile size, preferring `eval_*` under eval axis rules.""" + return getattr(self.config, f"eval_{name}" if max_utils.is_eval(self.config) else name) def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments, padding_amount): """Execute jax.lax.ragged_dot, with potential quantization""" @@ -2049,6 +2752,11 @@ def gmm( weight_gather_axes, group_offset, partial_sum=None, + return_lhs=False, + return_rhs=False, + out_is_3d=False, + tc_routing=None, + tc_unsort_cfg=None, ): def extract_vma(tensor): # Extract underlying array from QArray to inspect sharding annotation string. @@ -2079,12 +2787,12 @@ def extract_vma(tensor): orig_inputs_shape = inputs.shape # save shape of inputs before potentially padding. # Pad only the underlying qvalue buffer to tile size boundaries, leaving scale intact. if isinstance(inputs, qpl.QArray): - padded_qval, padding_amount = max_utils.maybe_pad(inputs.qvalue, self.config.wi_tile_fwd_batch_seq) + padded_qval, padding_amount = max_utils.maybe_pad(inputs.qvalue, fwd_tile("wi_tile_fwd_batch_seq")) inputs = dataclasses.replace(inputs, qvalue=padded_qval) else: - inputs, padding_amount = max_utils.maybe_pad(inputs, self.config.wi_tile_fwd_batch_seq) + inputs, padding_amount = max_utils.maybe_pad(inputs, fwd_tile("wi_tile_fwd_batch_seq")) if padding_amount > 0 and partial_sum is not None: - partial_sum = jnp.pad(partial_sum, ((0, padding_amount), (0, 0))) + partial_sum = jnp.pad(partial_sum, ((0, padding_amount),) + ((0, 0),) * (partial_sum.ndim - 1)) if not isinstance(inputs, qpl.QArray): inputs = inputs.astype(self.dtype) if not is_native_kernel: @@ -2123,6 +2831,26 @@ def extract_vma(tensor): ) elif use_custom_vjp_gmm: # tokamax gmm v1 (quantized), tokamax gmm v2 (quantized, unquantized), older forked megablox + flat_fsdp_shard = None + if self.config.moe_flat_fsdp_weights and not is_native_kernel and kernel.ndim == 2: + # `kernel` is the local [rows_local, cols] block of the flat (expert, fsdp)-sharded weight. After the + # dim-0 FSDP all-gather it holds this shard's E_local experts back to back: a free bitcast to + # [E_local, rows, cols]. With no FSDP gather axis the local block already is the whole expert group. + num_local_experts = self.num_experts // self.get_expert_parallelism_size() + gather_factor = 1 + for axis_name, axis_idx in weight_gather_axes: + if axis_idx == 0: + gather_factor *= self.mesh.shape[axis_name] + gathered_rows = kernel.shape[0] * gather_factor + if gathered_rows % num_local_experts != 0: + raise ValueError( + f"moe_flat_fsdp_weights: {gathered_rows} gathered rows not divisible by {num_local_experts} experts" + ) + gathered_shape = (num_local_experts, gathered_rows // num_local_experts, kernel.shape[1]) + if weight_gather_axes: + flat_fsdp_shard = gathered_shape + else: + kernel = kernel.reshape(gathered_shape) output = mblx.gmm( lhs=inputs, rhs=kernel, @@ -2139,8 +2867,16 @@ def extract_vma(tensor): rhs_vma_axes=rhs_vma_axes, use_gmm_v2=self.config.use_gmm_v2, use_gmm_v2_heuristic_tiling=self.config.use_gmm_v2_heuristic_tiling, + use_dlhs_transpose_rhs=self.config.moe_gmm_v2_dlhs_transpose_rhs, partial_sum=partial_sum, interpret=megablox_interpret, + return_lhs=return_lhs, + fuse_dlhs_scale=self.config.moe_accumulate_wi_dlhs, + return_rhs=return_rhs, + flat_fsdp_shard=flat_fsdp_shard, + fuse_3d_dlhs_ps=self.config.moe_3d_dlhs_partial_sum, + **({"out_is_3d": True} if out_is_3d else {}), + **({"tc_routing": tc_routing, "tc_unsort_cfg": tc_unsort_cfg} if tc_unsort_cfg is not None else {}), ) else: # jax.lax.ragged_dot @@ -2153,9 +2889,27 @@ def extract_vma(tensor): padding_amount, ) - if padding_amount > 0: - output = output[: orig_inputs_shape[0]] - return output + def unpad(t): + if padding_amount == 0: + return t + # Mirror the padding above: only the qvalue of a QArray was padded. + if isinstance(t, qpl.QArray): + return dataclasses.replace(t, qvalue=t.qvalue[: orig_inputs_shape[0]]) + return t[: orig_inputs_shape[0]] + + # With tc_unsort_cfg the output is already reduced to token order, so it has no padding rows. + unpad_out = (lambda t: t) if tc_unsort_cfg is not None else unpad + # The returned lhs is padded like the output; the returned rhs (weights) is not. + if return_lhs and return_rhs: + output, lhs_out, rhs_out = output + return unpad_out(output), unpad(lhs_out), rhs_out + if return_lhs: + output, lhs_out = output + return unpad_out(output), unpad(lhs_out) + if return_rhs: + output, rhs_out = output + return unpad_out(output), rhs_out + return unpad_out(output) def is_batch_sharded_by_ep(input_activation): # The batch is sharded by expert, except during inference decoding (where batch size == 1). @@ -2236,9 +2990,15 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): wo_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) elif self.config.shard_embed_moe_on_fsdp and explicitly_weight_ag(): # Keep embed_moe sharded so we can manually QAG it over FSDP - w0_pspec = self._logical_to_mesh_axes(("exp", "embed_moe", "mlp_no_fsdp")) - w1_pspec = self._logical_to_mesh_axes(("exp", "embed_moe", "mlp_no_fsdp")) - wo_pspec = self._logical_to_mesh_axes(("exp", "mlp_no_fsdp", "embed_moe")) + if self.config.moe_flat_fsdp_weights: + # Flat [E * rows, cols] weights: keep the (expert, fsdp) row sharding; gmm QAGs dim 0 over FSDP. + w0_pspec = self._logical_to_mesh_axes((CompoundLogicalAxis(("exp", "embed_moe")), "mlp_no_fsdp")) + w1_pspec = self._logical_to_mesh_axes((CompoundLogicalAxis(("exp", "embed_moe")), "mlp_no_fsdp")) + wo_pspec = self._logical_to_mesh_axes((CompoundLogicalAxis(("exp", "embed_moe")), None)) + else: + w0_pspec = self._logical_to_mesh_axes(("exp", "embed_moe", "mlp_no_fsdp")) + w1_pspec = self._logical_to_mesh_axes(("exp", "embed_moe", "mlp_no_fsdp")) + wo_pspec = self._logical_to_mesh_axes(("exp", "mlp_no_fsdp", "embed_moe")) else: # Tell XLA to automatically gather the D dimension (e.g. for dynamic absmax scaling) w0_pspec = self._logical_to_mesh_axes(("exp", None, "mlp_no_fsdp")) @@ -2264,6 +3024,11 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): is_batch_sharded_by_expert = is_batch_sharded_by_ep(inputs) weight_gather = explicitly_weight_ag() + if self.config.moe_flat_fsdp_weights and not weight_gather: + raise NotImplementedError( + "moe_flat_fsdp_weights needs the explicit fp8 weight all-gather" + " (shard_embed_moe_on_fsdp + fixed weight calibration)." + ) ( batch_logical_axis, input_partition_pspec, @@ -2296,6 +3061,55 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): # Exclude the EP axis so psum doesn't redundantly reduce and scale counts by num_ep # for calculate_load_balance_updates(). roe_bias_axis_names = _filter_axis_names(bias_update_axis_names, self._expert_parallelism_name) + # Sigmoid routers use Megatron-LM's seq_aux_loss, computed once on the full sequences after the + # shard_map (below), so the per-chunk loss inside the ring-of-experts loop is skipped. + use_megatron_seq_aux_loss = ( + self.config.load_balance_loss_weight > 0.0 + and not self.is_hash_routing + and uses_megatron_seq_aux_loss(self.config) + ) + # Debug only: the shard_map gets a 5th (max_load_ratio) output only when requested, so the + # default shard_map signature is unchanged. Stub configs in tests may lack the field. + log_max_load_ratio = getattr(self.config, "moe_log_max_load_ratio", False) + + def finalize_bias_updates(summed_expert_counts): + """Loss-free-balancing update from expert counts summed over all token chunks (the full batch).""" + if summed_expert_counts is None: + return None + if getattr(self.config, "gradient_accumulation_steps", 1) > 1: + # Under gradient accumulation, emit raw per-microbatch expert counts so + # gradient_accumulation_loss_and_grad sums counts across microbatches + # before taking a single sign() over the global batch. + return summed_expert_counts + return load_balance_updates_from_counts( + summed_expert_counts, self.config.num_experts, self.config.routed_bias_update_rate + ) + + # defer_small_all_reduces: the ring-of-experts layer emits local partial counts (leading axis sharded over the + # axes the per-layer psum would have used) and loss_fn reduces the stacked arrays once after the layer loop. + defer_small_ars = bool(getattr(self.config, "defer_small_all_reduces", False)) and self.config.use_ring_of_experts + + def use_tc_3d_dispatch(): + """moe_tc_ragged_3d_dispatch: tokens are all-gathered, TC-sorted and combined in (emb // 128, 128) layout. + + Raises instead of silently falling back to the 2D layout when the flag is set but cannot apply. + """ + if not self.config.moe_tc_ragged_3d_dispatch: + return False + requirements = { + "moe_tc_ragged_sort=True": self.config.moe_tc_ragged_sort, + "moe_tc_ragged_3d_gmm=True": self.config.moe_tc_ragged_3d_gmm, + "use_ring_of_experts=True": self.config.use_ring_of_experts, + "use_ragged_sort=True": self.config.use_ragged_sort, + "num_moe_emb_chunks <= 0": self.config.num_moe_emb_chunks <= 0, + "no tensor parallelism": self.get_tensor_parallelism_size() == 1, + "mlp_bias=False": not self.config.mlp_bias, + f"MoE expert input dim ({self.moe_expert_input_dim}) % 128 == 0": self.moe_expert_input_dim % 128 == 0, + } + unmet = [name for name, ok in requirements.items() if not ok] + if unmet: + raise ValueError(f"moe_tc_ragged_3d_dispatch=True requires: {', '.join(unmet)}.") + return True def quantize_and_all_gather_tokens( x: jax.Array, @@ -2314,6 +3128,13 @@ def quantize_and_all_gather_tokens( quantization_rule = qpl.get_current_rule("gmm") if quantization_rule is None: raise ValueError("Quantization rule is None, cannot quantize activation.") + # The gmm backward (ops._lhs_rhs_share_fixed_scale) assumes this QArray's fixed scale is the rule's act + # scale, so it must be quantized with the rule's act calibration. + if calibration_method != quantization_rule.act_calibration_method: + raise ValueError( + f"moe_quantize_token_all_gather calibration {calibration_method!r} must match the gmm rule's" + f" act_calibration_method {quantization_rule.act_calibration_method!r}." + ) act_qtype = quantization_rule.act_qtype # Quantize local token activations using the GMM activation quantization rule. @@ -2327,6 +3148,25 @@ def quantize_and_all_gather_tokens( x_gathered_qval = all_gather_fn(x_input.qvalue) return dataclasses.replace(x_input, qvalue=x_gathered_qval) + # The ring of experts' all-gather over the expert-parallel axis (tokens and routing), on a SparseCore with + # moe_pin_sparse_core_all_gathers / moe_pin_sparse_core_ep_all_gathers. roe_ag_and_route and the packed routing + # gather of the pipelined chunk loop (moe_roe_packed_routing_gather) both use it. + if self.config.moe_pin_sparse_core_all_gathers or getattr(self.config, "moe_pin_sparse_core_ep_all_gathers", False): + + @functools.partial( + compute_on, + compute_type="tpu_sparsecore", + out_memory_spaces=jax.memory.Space.Device, + compiler_options={"sparse_core_config": {"core_ids": [self.config.moe_ep_all_gather_sparse_core_id]}}, + ) + def _ep_all_gather(z): + return jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True) + + else: + + def _ep_all_gather(z): + return jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True) + def roe_ag_and_route( x, logits, @@ -2337,22 +3177,22 @@ def roe_ag_and_route( input_ids=None, forced_routed_experts=None, force_dropless=False, + gathered_topk=None, ): - if self.config.moe_pin_sparse_core_all_gathers: - - @functools.partial( - compute_on, - compute_type="tpu_sparsecore", - out_memory_spaces=jax.memory.Space.Device, - compiler_options={"sparse_core_config": {"core_ids": [self.config.moe_ep_all_gather_sparse_core_id]}}, - ) - def _ep_all_gather(z): - return jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True) + """Ring of experts: all-gathers this chunk's tokens over the expert-parallel axis and routes them. + + `gathered_topk` (moe_roe_packed_routing_gather): `(weights, expert ids)` of these tokens, already routed and + all-gathered over the expert-parallel axis, plus this shard's own (local, not gathered) expert ids for the + load-balance loss. The local top-k and the two routing all-gathers of this chunk are then skipped. + """ + if use_tc_3d_dispatch(): + # moe_tc_ragged_3d_dispatch: all-gather the tokens in the TC kernels' (emb // 128, 128) layout; the + # reshape is on the local tokens, so the gathered tensor needs no relayout before the TC sort. + def token_all_gather(z): + return _ep_all_gather(z.reshape(*z.shape[:-1], z.shape[-1] // 128, 128)) else: - - def _ep_all_gather(z): - return jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True) + token_all_gather = _ep_all_gather # Duplicate token inputs across all expert shards. if self.config.moe_quantize_token_all_gather: @@ -2362,17 +3202,41 @@ def _ep_all_gather(z): x, axis_name=self._expert_parallelism_name, calibration_method=self.config.act_quantization_calibration_method, - all_gather_fn=_ep_all_gather, + all_gather_fn=token_all_gather, ) else: - x = _ep_all_gather(x) + x = token_all_gather(x) + local_tokens = logits.shape[0] * logits.shape[1] topk = None - if self.config.ring_of_experts_local_routing: - # Routing is per token, so each shard routes only its own tokens and all-gathers the top-k + precomputed_topk = None + precomputed_density_prob = None + topk_before_all_gather = ( + self.config.moe_topk_before_ep_all_gather + and forced_routed_experts is None + and not self.is_hash_routing + and not self.config.use_random_routing + ) + if gathered_topk is not None: + # moe_roe_packed_routing_gather: the chunk pipeline routed this shard's tokens of every chunk at once and + # all-gathered weights and ids as one packed array; this chunk's slice arrives here. Same values as the + # local-routing branch below (routing is per token), including its per-chunk load-balance loss, which is + # computed from this shard's own ids of these tokens exactly as there. + weights, selected_experts, local_selected_experts = gathered_topk + lb_loss = None + if self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing and not use_megatron_seq_aux_loss: + lb_loss = jax.lax.pmean( + self.router_load_balance_loss(local_selected_experts, logits, pre_bias_logits), + self._expert_parallelism_name, + ) + topk = (weights, selected_experts, lb_loss) + logits, pre_bias_logits = None, None + elif self.config.ring_of_experts_local_routing: + # (68792f5 pin) Routing is per token, so each shard routes only its own tokens and all-gathers the top-k # ids and weights the way the tokens were gathered. The [tokens, num_experts] logits never # cross shards, so neither does their gradient, and the gradient of the top-k gather is a - # scatter into this shard's logits rather than into the whole gathered batch's. + # scatter into this shard's logits rather than into the whole gathered batch's. Unlike + # moe_topk_before_ep_all_gather below, this also covers router replay (forced_routed_experts). weights, selected_experts = self.get_topk( logits, pre_bias_logits, @@ -2383,7 +3247,7 @@ def _ep_all_gather(z): token_shard_index=expert_shard_id, ) lb_loss = None - if self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing: + if self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing and not use_megatron_seq_aux_loss: # The loss averages per-sequence terms over the batch, and every shard holds the same # number of sequences, so the mean over shards is the loss over the gathered batch. lb_loss = jax.lax.pmean( @@ -2392,6 +3256,20 @@ def _ep_all_gather(z): ) topk = (_ep_all_gather(weights), _ep_all_gather(selected_experts), lb_loss) logits, pre_bias_logits = None, None + elif topk_before_all_gather: + # Top-k is per token, so run it on the local tokens and all-gather only the (batch, seq, k) + # weights and expert ids instead of the (batch, seq, num_experts) logits: every shard + # otherwise redoes top-k on the full gathered batch, and the logit gather (and its + # reduce-scatter in the backward) moves num_experts / k times more bytes. Only the index map + # is computed here; permute() below still sorts the gathered tokens. + weights, selected_experts = self.get_topk(logits, pre_bias_logits, rngs, input_ids) + precomputed_topk = (_ep_all_gather(weights), _ep_all_gather(selected_experts)) + if self.config.load_balance_loss_weight > 0.0 and not use_megatron_seq_aux_loss: + # Per-sequence statistic, so gathering it over EP gives the same loss as the full batch. + probs_logits = pre_bias_logits if pre_bias_logits is not None else logits + softmax_probs = jax.nn.softmax(probs_logits.astype(jnp.float32), axis=-1).astype(self.dtype) + precomputed_density_prob = _ep_all_gather(jnp.mean(softmax_probs, axis=1)) + logits, pre_bias_logits = None, None else: # Duplicate routing inputs across all expert shards. Routing evaluates # the full gathered batch; logits, pre_bias_logits, and @@ -2413,6 +3291,9 @@ def _ep_all_gather(z): bias_updates, local_group_sizes, has_overflow, + max_load_ratio, + topk_argsort_indices, + tc_routing, moe_audit, ) = self.permute( x, @@ -2426,7 +3307,15 @@ def _ep_all_gather(z): force_dropless=force_dropless, mesh_axis_names=roe_bias_axis_names, topk=topk, + compute_lb_loss=not use_megatron_seq_aux_loss, + return_expert_counts=True, + defer_reductions=defer_small_ars, + precomputed_topk=precomputed_topk, + precomputed_density_prob=precomputed_density_prob, ) + required_rbf = None + if getattr(self.config, "log_required_ragged_buffer_factor", False): + required_rbf = self.required_ragged_buffer_factor(group_sizes, local_tokens * num_ep, num_ep, expert_shard_id) return ( x, RouteOutput( @@ -2438,6 +3327,10 @@ def _ep_all_gather(z): bias_updates=bias_updates, local_group_sizes=local_group_sizes, has_overflow=has_overflow, + required_rbf=required_rbf, + max_load_ratio=max_load_ratio, + topk_argsort_indices=topk_argsort_indices, + tc_routing=tc_routing, moe_audit=moe_audit, ), RouteMetadata( @@ -2472,6 +3365,9 @@ def ra2a_and_route( local_group_sizes, _, _, + _, + _, + _, ) = self.permute( x, logits, @@ -2481,6 +3377,8 @@ def ra2a_and_route( input_ids=input_ids, forced_routed_experts=forced_routed_experts, mesh_axis_names=bias_update_axis_names, + compute_lb_loss=not use_megatron_seq_aux_loss, + return_expert_counts=True, ) if num_ep > 1: @@ -2490,7 +3388,6 @@ def ra2a_and_route( reshaped_group_sizes = jnp.sum(group_sizes.reshape(-1, local_expert_size), axis=1) global_group_sizes = group_sizes - rbf = self.get_ragged_buffer_factor() if is_batch_sharded_by_expert: all_shards_group_sizes = jax.lax.all_gather(reshaped_group_sizes, axis_name=batch_axis) buffer_size = self.get_ragged_buffer_size( @@ -2498,13 +3395,13 @@ def ra2a_and_route( num_ep, self.config.num_experts, self.config.num_experts_per_tok, - rbf, + self.get_ragged_buffer_factor(), ) input_offsets, send_sizes, output_offsets, recv_sizes = RoutedMoE.get_all_to_all_params( all_shards_group_sizes, expert_shard_id, num_ep, - ragged_buffer_factor=rbf, + ragged_buffer_factor=self.get_ragged_buffer_factor(), buffer_size=buffer_size, ) @@ -2527,7 +3424,7 @@ def ra2a_and_route( shard_index=expert_shard_id, use_custom_sort_vjp=self.config.use_custom_sort_vjp, use_ragged_sort=self.config.use_ragged_sort, - ragged_buffer_factor=rbf, + ragged_buffer_factor=self.get_ragged_buffer_factor(), use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) else: @@ -2540,7 +3437,7 @@ def ra2a_and_route( global_sorted_experts=selected_experts, use_custom_sort_vjp=self.config.use_custom_sort_vjp, use_ragged_sort=self.config.use_ragged_sort, - ragged_buffer_factor=rbf, + ragged_buffer_factor=self.get_ragged_buffer_factor(), use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) @@ -2571,6 +3468,7 @@ def route( input_ids=None, forced_routed_experts=None, force_dropless=False, + gathered_topk=None, ): """Performs both across device and within device token routing/sorting""" num_ep = self.get_expert_parallelism_size() @@ -2587,6 +3485,7 @@ def route( input_ids=input_ids, forced_routed_experts=forced_routed_experts, force_dropless=force_dropless, + gathered_topk=gathered_topk, ) else: return ra2a_and_route( @@ -2600,6 +3499,15 @@ def route( forced_routed_experts=forced_routed_experts, ) + def non_expert_axes(pspec_dim_axes): + """Drops the expert-parallel mesh axis name(s) from one PartitionSpec entry (flat dim 0 = (expert, fsdp)).""" + if pspec_dim_axes is None: + return None + ep = self._expert_parallelism_name + ep_names = set(ep) if isinstance(ep, tuple) else {ep} + axes = (pspec_dim_axes,) if isinstance(pspec_dim_axes, str) else tuple(pspec_dim_axes) + return tuple(ax for ax in axes if ax not in ep_names) + def get_active_sharding_axes(pspec_dim_axes, tensor_dim_index): if pspec_dim_axes is None: return [] @@ -2614,18 +3522,23 @@ def get_wi_gmm_params(): wi_gather_axes = [] if weight_gather: # weight_gather implies either exp or embed_moe is sharded. - if self.config.shard_exp_on_fsdp: - # wi [Experts, In, Hidden] -> Gather Exp(0) - wi_gather_axes.extend(get_active_sharding_axes(w0_pspec[0], 0)) - else: - # Gather In(1) where embed_moe is sharded. + if self.config.moe_flat_fsdp_weights: + # Flat wi [E * In, Hidden]: gather the non-expert (FSDP) axes of dim 0, then Hidden(1). + wi_gather_axes.extend(get_active_sharding_axes(non_expert_axes(w0_pspec[0]), 0)) wi_gather_axes.extend(get_active_sharding_axes(w0_pspec[1], 1)) - # Gather Hidden(2) - wi_gather_axes.extend(get_active_sharding_axes(w0_pspec[2], 2)) + else: + if self.config.shard_exp_on_fsdp: + # wi [Experts, In, Hidden] -> Gather Exp(0) + wi_gather_axes.extend(get_active_sharding_axes(w0_pspec[0], 0)) + else: + # Gather In(1) where embed_moe is sharded. + wi_gather_axes.extend(get_active_sharding_axes(w0_pspec[1], 1)) + # Gather Hidden(2) + wi_gather_axes.extend(get_active_sharding_axes(w0_pspec[2], 2)) wi_tile_size = ( - self.config.wi_tile_fwd_batch_seq, # m (LHS batch) - self.config.wi_tile_fwd_embed_dim, # k (contracting) - self.config.wi_tile_fwd_mlp_dim, # n (RHS batch) + fwd_tile("wi_tile_fwd_batch_seq"), # m (LHS batch) + fwd_tile("wi_tile_fwd_embed_dim"), # k (contracting) + fwd_tile("wi_tile_fwd_mlp_dim"), # n (RHS batch) self.config.wi_tile_dlhs_batch_seq, # m (LHS batch) self.config.wi_tile_dlhs_mlp_dim, # k (contracting) self.config.wi_tile_dlhs_embed_dim, # n (RHS batch) @@ -2639,18 +3552,23 @@ def get_wo_gmm_params(): wo_gather_axes = [] if weight_gather: # weight_gather implies either exp or embed_moe is sharded. - if self.config.shard_exp_on_fsdp: - # wo [Experts, Hidden, Out] -> Gather Exp(0) - wo_gather_axes.extend(get_active_sharding_axes(wo_pspec[0], 0)) + if self.config.moe_flat_fsdp_weights: + # Flat wo [E * Hidden, Out]: gather the non-expert (FSDP) axes of dim 0, then Out(1). + wo_gather_axes.extend(get_active_sharding_axes(non_expert_axes(wo_pspec[0]), 0)) + wo_gather_axes.extend(get_active_sharding_axes(wo_pspec[1], 1)) else: - # Gather Out(2) where embed_moe is sharded. - wo_gather_axes.extend(get_active_sharding_axes(wo_pspec[2], 2)) - # Gather Hidden(1) - wo_gather_axes.extend(get_active_sharding_axes(wo_pspec[1], 1)) + if self.config.shard_exp_on_fsdp: + # wo [Experts, Hidden, Out] -> Gather Exp(0) + wo_gather_axes.extend(get_active_sharding_axes(wo_pspec[0], 0)) + else: + # Gather Out(2) where embed_moe is sharded. + wo_gather_axes.extend(get_active_sharding_axes(wo_pspec[2], 2)) + # Gather Hidden(1) + wo_gather_axes.extend(get_active_sharding_axes(wo_pspec[1], 1)) wo_tile_size = ( - self.config.wo_tile_fwd_batch_seq, # m (LHS batch) - self.config.wo_tile_fwd_mlp_dim, # k (contracting) - self.config.wo_tile_fwd_embed_dim, # n (RHS batch) + fwd_tile("wo_tile_fwd_batch_seq"), # m (LHS batch) + fwd_tile("wo_tile_fwd_mlp_dim"), # k (contracting) + fwd_tile("wo_tile_fwd_embed_dim"), # n (RHS batch) self.config.wo_tile_dlhs_batch_seq, # m (LHS batch) self.config.wo_tile_dlhs_embed_dim, # k (contracting) self.config.wo_tile_dlhs_mlp_dim, # n (RHS) @@ -2671,13 +3589,27 @@ def gmm_up( partial_accum0=None, partial_accum1=None, mask=None, + gather_weights=True, + return_weights=False, ): """Run the two up-projections (gate + up) and apply the FFN activation.""" wi_gather_axes, wi_tile_size = get_wi_gmm_params() + if not gather_weights: + wi_gather_axes = [] if self.config.prefuse_moe_weights: - # Weights are stored as (G,K,2N); w0/w1 are adjacent slices so XLA elides this concat. - w_fused = jnp.concatenate([w0, w1], axis=-1) + if w1 is None: + # moe_unsplit_prefused_wi: w0 is already the gathered (and fake-quantized) whole [w0 | w1] weight. + w_fused = w0 + else: + # Weights are stored as (G,K,2N); w0/w1 are adjacent slices so XLA elides this concat. + w_fused = jnp.concatenate([w0, w1], axis=-1) out = gmm_fn(x, w_fused, tiling=wi_tile_size, weight_gather_axes=wi_gather_axes) + if partial_accum0 is None and partial_accum1 is None and self.config.moe_prefused_swiglu_vjp: + # moe_prefused_swiglu_vjp: the whole [gate | up] output goes to the activation as one array + # (apply_ffn_activation_maybe_fused, which layer_w1=None selects), whose backward then writes d[gate | up] + # in one pass. There are no per-half outputs to tag; the config validator rejects a remat policy that + # would save them, since it would silently stop saving the MoE tensor. + return out, None n = out.shape[-1] // 2 layer_w0, layer_w1 = out[:, :n], out[:, n:] if self.config.mlp_bias and w0_bias is not None and w1_bias is not None: @@ -2688,14 +3620,28 @@ def gmm_up( layer_w1 = jnp.where(mask[:, None], layer_w1, 0) layer_w0 = adc.checkpoint_name(adc.checkpoint_name(layer_w0, "mlpwi_0"), "moe_mlpwi_0") layer_w1 = adc.checkpoint_name(layer_w1, "moe_mlpwi_1") + return layer_w0, layer_w1 else: + # With moe_accumulate_wi_dlhs, the wi_0 GMM returns its input and the wi_1 GMM consumes that, so in the + # backward pass the wi_1 DLHS arrives as the cotangent of the returned input and the wi_0 DLHS GMM + # accumulates onto it in place (gmm_v2 partial_sum) instead of a separate add of the two DLHS. + accum_wi_dlhs = self.config.moe_accumulate_wi_dlhs layer_w0 = gmm_fn( x, w0, tiling=wi_tile_size, weight_gather_axes=wi_gather_axes, partial_sum=partial_accum0, + return_lhs=accum_wi_dlhs, + return_rhs=return_weights, ) + x_for_w1 = x + if accum_wi_dlhs and return_weights: + layer_w0, x_for_w1, w0 = layer_w0 + elif accum_wi_dlhs: + layer_w0, x_for_w1 = layer_w0 + elif return_weights: + layer_w0, w0 = layer_w0 if self.config.mlp_bias and w0_bias is not None: layer_w0 = layer_w0 + w0_bias if mask is not None: @@ -2703,18 +3649,23 @@ def gmm_up( layer_w0 = adc.checkpoint_name(adc.checkpoint_name(layer_w0, "mlpwi_0"), "moe_mlpwi_0") layer_w1 = gmm_fn( - x, + x_for_w1, w1, tiling=wi_tile_size, weight_gather_axes=wi_gather_axes, partial_sum=partial_accum1, + return_rhs=return_weights, ) + if return_weights: + layer_w1, w1 = layer_w1 if self.config.mlp_bias and w1_bias is not None: layer_w1 = layer_w1 + w1_bias if mask is not None: layer_w1 = jnp.where(mask[:, None], layer_w1, 0) layer_w1 = adc.checkpoint_name(layer_w1, "moe_mlpwi_1") - return layer_w0, layer_w1 + if return_weights: + return layer_w0, layer_w1, w0, w1 + return layer_w0, layer_w1 def valid_token_count(x, routing, route_metadata): """Rows of `x` the gmm actually computes; the rest is ragged-buffer padding.""" @@ -2727,28 +3678,105 @@ def valid_token_count(x, routing, route_metadata): return jnp.sum(jax.lax.dynamic_slice_in_dim(routing.group_sizes, experts_start, num_experts_per_shard)) return jnp.sum(routing.group_sizes) - def get_gmm_for_local_experts(x, routing, route_metadata): - """Return a partial GMM function with preconfigured routing params.""" + # moe_wi_gmm_fused_swiglu: apply the SwiGLU in the wi GMM's epilogue (megablox ops.gmm_silu). The unfused path + # writes the [buffer, 2 * mlp] GMM output and reads it back in an elementwise pass over every buffer row, padding + # included; the fused GMM writes the [buffer, mlp] activation of the rows it computes. When differentiated it also + # writes gate and up, and the backward is the unfused one. Works with split (w0, w1) and with the whole prefused + # wi of moe_unsplit_prefused_wi (w0 = [w0 | w1], w1 = None). Off, gmm_up_or_silu is gmm_up and nothing below runs. + fuse_silu_in_wi_gmm = self.config.moe_wi_gmm_fused_swiglu + if fuse_silu_in_wi_gmm: + # The config validator checks the config-level requirements; these two only the built layer knows. + unmet = [ + name + for name, hit in ( + ("weight quantization", self.quant is not None), + ("an explicit expert-weight all-gather", weight_gather), + ) + if hit + ] + if unmet: + raise ValueError(f"moe_wi_gmm_fused_swiglu=True does not support {' or '.join(unmet)}.") + + def gmm_up_or_silu( + x, + w0, + w1, + w0_bias, + w1_bias, + gmm_fn, + gmm_operands, + _weight_gather, + mask=None, + gather_weights=True, + return_weights=False, + ): + """gmm_up; with moe_wi_gmm_fused_swiglu, (silu(x @ w0) * (x @ w1), None) from one fused GMM. + + `gmm_operands` are the routing operands that `gmm_fn` is bound to (get_gmm_operands_for_local_experts); the + fused GMM takes them explicitly. An output1 of None tells the caller the activation is already applied. + Biases, the mask and gather_weights only matter in configurations the key does not support (mlp_bias, an + explicit weight all-gather). + """ + if not fuse_silu_in_wi_gmm: + return gmm_up( + x, + w0, + w1, + w0_bias, + w1_bias, + gmm_fn, + _weight_gather, + mask=mask, + gather_weights=gather_weights, + return_weights=return_weights, + ) + if return_weights: + raise NotImplementedError("moe_wi_gmm_fused_swiglu does not support moe_accumulate_chunk_wgrad.") + del w0_bias, w1_bias, mask, gather_weights + _, wi_tile_size = get_wi_gmm_params() + # As gmm(): pad the rows to the forward tile and cast to the compute dtype. + inputs, padding_amount = max_utils.maybe_pad(x, fwd_tile("wi_tile_fwd_batch_seq")) + # Weights are stored as (G,K,2N): w0/w1 are adjacent slices (the same concatenate as gmm_up), or w0 is already + # the whole [w0 | w1] (moe_unsplit_prefused_wi). + w_fused = w0 if w1 is None else jnp.concatenate([w0, w1], axis=-1) + act = mblx.gmm_silu( + inputs.astype(self.dtype), + w_fused.astype(self.dtype), + gmm_operands["group_sizes"], + preferred_element_type=self.dtype, + tiling=wi_tile_size, + group_offset=gmm_operands["group_offset"], + use_gmm_v2_heuristic_tiling=self.config.use_gmm_v2_heuristic_tiling, + use_dlhs_transpose_rhs=self.config.moe_gmm_v2_dlhs_transpose_rhs, + fuse_dlhs_scale=self.config.moe_accumulate_wi_dlhs, + fuse_3d_dlhs_ps=self.config.moe_3d_dlhs_partial_sum, + ) + return (act[: x.shape[0]] if padding_amount else act), None + + def get_gmm_operands_for_local_experts(x, routing, route_metadata): + """The routing operands of the local experts' GMMs: group sizes, expert assignments, group offset.""" num_ep = self.get_expert_parallelism_size() num_experts_per_shard = self.config.num_experts // num_ep if self.config.use_ring_of_experts and x.shape[0] < routing.sorted_selected_experts.shape[0]: local_group_sizes = routing.local_group_sizes - return functools.partial( - gmm, - group_sizes=local_group_sizes, - expert_assignments=routing.selected_experts, - group_offset=0, - ) + return { + "group_sizes": local_group_sizes, + "expert_assignments": routing.selected_experts, + "group_offset": 0, + } if self.config.use_ragged_sort and self.config.use_ring_of_experts: experts_start = route_metadata.expert_shard_id * num_experts_per_shard else: experts_start = 0 - return functools.partial( - gmm, - group_sizes=routing.group_sizes, - expert_assignments=routing.selected_experts, - group_offset=experts_start, - ) + return { + "group_sizes": routing.group_sizes, + "expert_assignments": routing.selected_experts, + "group_offset": experts_start, + } + + def get_gmm_for_local_experts(x, routing, route_metadata): + """Return a partial GMM function with preconfigured routing params.""" + return functools.partial(gmm, **get_gmm_operands_for_local_experts(x, routing, route_metadata)) def unsort_output_and_ra2a( intermediate_output, @@ -2779,12 +3807,11 @@ def unsort_output_and_ra2a( ) buffer_size = intermediate_output.shape[0] - rbf = self.get_ragged_buffer_factor() input_offsets, send_sizes, output_offsets, recv_sizes = RoutedMoE.get_all_to_all_params( route_metadata.all_shards_group_sizes, route_metadata.expert_shard_id, self.get_expert_parallelism_size(), - ragged_buffer_factor=rbf, + ragged_buffer_factor=self.get_ragged_buffer_factor(), buffer_size=buffer_size, is_dispatch=False, ) @@ -2918,27 +3945,85 @@ def scan_fn(carry, _): ) return output0, output1, gmm_fn, routing, route_metadata, wo_bias - def _up_projection(x, routing, route_metadata, w0, w1, w0_bias, w1_bias, wo_bias): + # Helpers for the pipelined ring-of-experts chunk loop (moe_chunk_pipeline, from the 68792f5 pin). They + # mirror the inline ring-of-experts path of _moe_body below (upstream main 2ee266287), including the + # moe_x_sorted tags, the TC ragged sort (routing weights on the activation, 3D gmm) and + # moe_quantize_combine_bwd_method, so the pipelined and plain chunk loops compute the same thing. + def _up_projection( + x, routing, route_metadata, w0, w1, w0_bias, w1_bias, wo_bias, gather_weights=True, return_weights=False + ): """The gate and up projections of the routed tokens `x`.""" + + def _cn(t): + return adc.checkpoint_name(t, "moe_x_sorted") if isinstance(t, jax.Array) else t + + x = jax.tree.map(_cn, x) + routing = jax.tree.map(_cn, routing) + route_metadata = jax.tree.map(_cn, route_metadata) mask = jnp.arange(x.shape[0]) < valid_token_count(x, routing, route_metadata) if self.config.mlp_bias: w0_bias, w1_bias, wo_bias = self.transform_bias(routing.selected_experts, w0_bias, w1_bias, wo_bias) - gmm_fn = get_gmm_for_local_experts(x, routing, route_metadata) - output0, output1 = gmm_up(x, w0, w1, w0_bias, w1_bias, gmm_fn, weight_gather, mask=mask) - return output0, output1, gmm_fn, wo_bias + gmm_operands = get_gmm_operands_for_local_experts(x, routing, route_metadata) + gmm_fn = functools.partial(gmm, **gmm_operands) + up_res = gmm_up_or_silu( + x, + w0, + w1, + w0_bias, + w1_bias, + gmm_fn, + gmm_operands, + weight_gather, + mask=mask, + gather_weights=gather_weights, + return_weights=return_weights, + ) + if return_weights: + output0, output1, w0, w1 = up_res # pylint: disable=unbalanced-tuple-unpacking + else: + output0, output1 = up_res # pylint: disable=unbalanced-tuple-unpacking + return x, routing, route_metadata, output0, output1, gmm_fn, wo_bias, (w0, w1) - def _down_projection(output0, output1, gmm_fn, wo, wo_bias, routing, route_metadata): - """The activation and the down projection, per routed token.""" - intermediate_layer = self.apply_ffn_activation(output0, output1) + def _down_projection( + x, output0, output1, gmm_fn, wo, wo_bias, routing, route_metadata, gather_weights=True, return_weights=False + ): + """The activation and the down projection, per routed token. + + Returns (intermediate_output, wo, tc_prescaled, x_is_3d). + """ + tc_prescaled = ( + self.config.use_ring_of_experts + and getattr(routing, "tc_routing", None) is not None + and self.config.moe_tc_ragged_weights_on_activation + ) + if tc_prescaled: + w_rows = tc_buffer_row_weights(routing.tc_routing, jnp.ravel(routing.weights), output0.shape[0]) + if self.config.moe_tc_routing_checkpoint != "remat": + w_rows = adc.checkpoint_name(w_rows, "moe_tc_routing_checkpoint") + # moe_wi_gmm_fused_swiglu: the wi GMM already applied the activation (output1 is None). + intermediate_layer = output0 if fuse_silu_in_wi_gmm else self.apply_ffn_activation_maybe_fused(output0, output1) + if tc_prescaled: + intermediate_layer = (intermediate_layer * w_rows[:, None]).astype(intermediate_layer.dtype) wo_gather_axes, wo_tile_size = get_wo_gmm_params() + if not gather_weights: + wo_gather_axes = [] + x_is_3d = self.config.num_moe_emb_chunks <= 0 and (x.qvalue if isinstance(x, qpl.QArray) else x).ndim == 3 + if x_is_3d and (self.get_tensor_parallelism_size() > 1 or self.config.mlp_bias): + raise NotImplementedError("moe_tc_ragged_3d_gmm does not support tensor parallelism or mlp_bias.") + if self.config.moe_bwd_prequant_before_unsort: + raise NotImplementedError("moe_chunk_pipeline does not support moe_bwd_prequant_before_unsort.") intermediate_output = gmm_fn( intermediate_layer, wo, tiling=wo_tile_size, weight_gather_axes=wo_gather_axes, + return_rhs=return_weights, + **({"out_is_3d": True} if x_is_3d else {}), ) + if return_weights: + intermediate_output, wo = intermediate_output if self.get_tensor_parallelism_size() > 1: intermediate_output = jax.lax.psum_scatter( intermediate_output, @@ -2948,13 +4033,15 @@ def _down_projection(output0, output1, gmm_fn, wo, wo_bias, routing, route_metad ) if self.config.mlp_bias: mask = jnp.arange(intermediate_output.shape[0]) < valid_token_count(intermediate_output, routing, route_metadata) + if tc_prescaled: + wo_bias = (wo_bias * w_rows[:, None]).astype(wo_bias.dtype) intermediate_output = intermediate_output + wo_bias intermediate_output = jnp.where(mask[:, None], intermediate_output, 0) - return adc.checkpoint_name(adc.checkpoint_name(intermediate_output, "mlpwo"), "moe_mlpwo") + intermediate_output = adc.checkpoint_name(adc.checkpoint_name(intermediate_output, "mlpwo"), "moe_mlpwo") + return intermediate_output, wo, tc_prescaled, x_is_3d - def _ring_unpermute(intermediate_output, routing, batch_size, sequence_length): - """Ring of experts: unsort, weight and sum the outputs locally, as [expert shards x batch, seq, embed].""" - # Unsort and deduplicate the outputs locally. + def _ring_unpermute(intermediate_output, routing, batch_size, sequence_length, tc_prescaled=False, x_is_3d=False): + """Ring of experts: unsort, weight and sum the outputs locally, as [expert shards x batch, seq, ...].""" output = self.unpermute( intermediate_output, routing.sorted_selected_experts, @@ -2963,7 +4050,14 @@ def _ring_unpermute(intermediate_output, routing, batch_size, sequence_length): sequence_length=sequence_length, use_custom_sort_vjp=self.config.use_custom_sort_vjp, group_sizes=routing.group_sizes, + topk_argsort_indices=routing.topk_argsort_indices, + tc_routing=routing.tc_routing, + tc_prescaled=tc_prescaled, + tc_out_3d=x_is_3d and use_tc_3d_dispatch(), ) + if output.ndim == 4: + # moe_tc_ragged_3d_dispatch: reduce-scatter in the (emb // 128, 128) layout. + return jnp.reshape(output, (-1, sequence_length, *output.shape[2:])) return jnp.reshape( output, ( @@ -2973,6 +4067,64 @@ def _ring_unpermute(intermediate_output, routing, batch_size, sequence_length): ), ) + def _ring_down_project_and_unpermute( + output0, output1, wo, routing, route_metadata, batch_size, sequence_length, forward_tie=None + ): + """moe_router_grad_from_wo_input: `_ring_unpermute(_down_projection(...))` as one op. + + The fused op's backward does not read the down projection's output: the routing-weight gradient + is computed as from the down projection's input, and the routing weight is applied + on the [rows, mlp] side, so a rematerialized layer no longer recomputes the down-projection GEMM in the + backward (ring_ragged_unsort_down_projection). The config validator admits the key only where the down + projection is the bare GEMM of the sorted rows (no bias, TC sort, tensor parallelism or quantization). + `forward_tie` is the pipelined loop's tie 2, applied inside the op to the down-projection output. + Returns the output as [expert shards x batch, seq, embed], as `_ring_unpermute`, and the tie's output. + """ + # moe_wi_gmm_fused_swiglu: the wi GMM already applied the activation (output1 is None). + intermediate_layer = output0 if fuse_silu_in_wi_gmm else self.apply_ffn_activation_maybe_fused(output0, output1) + wo_gather_axes, wo_tile_size = get_wo_gmm_params() + + def down_projection(h, kernel, group_sizes, expert_assignments, group_offset): + return gmm( + h, + kernel, + tiling=wo_tile_size, + group_sizes=group_sizes, + expert_assignments=expert_assignments, + weight_gather_axes=wo_gather_axes, + group_offset=group_offset, + ) + + operands = get_gmm_operands_for_local_experts(intermediate_layer, routing, route_metadata) + output = ring_ragged_unsort_down_projection( + intermediate_layer, + wo, + down_projection, + (operands["group_sizes"], operands["expert_assignments"], operands["group_offset"]), + routing.group_sizes, + routing.sorted_selected_experts, + self.num_experts_per_tok, + self.config.num_experts // self.get_expert_parallelism_size(), + self._expert_parallelism_name, + topk_weights=jnp.ravel(routing.weights).astype(jnp.float32), + topk_argsort_indices=routing.topk_argsort_indices, + forward_tie=forward_tie, + enforce_gather_fallback=self.config.ragged_gather_fallback, + enforce_gather_reduce_fallback=self.config.ragged_gather_reduce_fallback, + gather_flops_override=self.config.ragged_gather_cost_estimate_flops, + gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, + gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, + gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, + use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, + index_math=self.config.moe_routing_index_math, + unique_pair_argsort=self.config.moe_ragged_sort_unique_pair_argsort, + ) + if forward_tie is not None: + output, forward_tie = output + # The reshapes of unpermute and _ring_unpermute (the embed dim is whole: no tensor parallelism here). + output = output.reshape(batch_size, sequence_length, -1).astype(self.dtype) + return jnp.reshape(output, (-1, sequence_length, self.moe_expert_input_dim)), forward_tie + def _ring_reduce_scatter(output): """Ring of experts: sum the partial outputs across the expert shards.""" if self.config.ring_of_experts_row_major_reduce_scatter: @@ -2981,11 +4133,40 @@ def _ring_reduce_scatter(output): # layout and the reduce-scatter's result back to row-major. Row-major makes the reduce-scatter read the # combine output where it is. output = forward_row_major(output) - return jax.lax.psum_scatter( + combine_bwd_method = self.config.moe_quantize_combine_bwd_method + if combine_bwd_method: + output = _moe_combine_psum_scatter( + output, + self._expert_parallelism_name, + scatter_dimension=0, + tiled=True, + bwd_method=combine_bwd_method, + ) + else: + output = jax.lax.psum_scatter( + output, + self._expert_parallelism_name, + scatter_dimension=0, + tiled=True, + ) + if output.ndim == 4: + output = jnp.reshape(output, (*output.shape[:2], -1)) + return output + + def _moe_body_outputs(output, routing, weights): + """What _moe_body returns: the layer output, the routing's side outputs and the weights `(w0, w1, wo)`. + + The weights are as the GMMs return them under moe_accumulate_chunk_wgrad. + """ + return ( output, - self._expert_parallelism_name, - scatter_dimension=0, - tiled=True, + routing.lb_loss, + routing.bias_updates, + routing.has_overflow, + routing.required_rbf, + routing.max_load_ratio, + routing.moe_audit, + weights, ) def _moe_body( @@ -3002,7 +4183,12 @@ def _moe_body( rngs, forced_routed_experts=None, force_dropless=False, + gather_weights=True, + return_weights=False, ): + if router_grad_from_wo_input and (return_weights or not gather_weights): + # The fused op does not return the weights (the config validator rejects the two keys together). + raise NotImplementedError("moe_router_grad_from_wo_input does not support moe_accumulate_chunk_wgrad.") batch_size, sequence_length, embed_dim = x.shape if self.config.num_moe_emb_chunks > 0: output0, output1, gmm_fn, routing, route_metadata, wo_bias = moe_emb_chunking( @@ -3029,13 +4215,168 @@ def _moe_body( forced_routed_experts=forced_routed_experts, force_dropless=force_dropless, ) - output0, output1, gmm_fn, wo_bias = _up_projection(x, routing, route_metadata, w0, w1, w0_bias, w1_bias, wo_bias) - intermediate_output = _down_projection(output0, output1, gmm_fn, wo, wo_bias, routing, route_metadata) + # moe_x_sorted: tag the routed (expert-sorted) MoE input and its small routing/metadata + # bundle for the custom remat policy. With moe_x_sorted=device the backward loads these + # instead of re-running route(), which removes the rematted dispatch token all-gather (EP + # axis) and the SparseCore ragged sort from the backward pass; the expert GMMs re-run from + # the saved tensor. The routing/metadata leaves (indices, group sizes, weights) are tiny but + # must be saved too, or the sort re-runs just to reproduce them. Tags are inert under the + # default moe_x_sorted=remat. tree.map so a QArray input (moe_quantize_token_all_gather) + # gets its qvalue/scale leaves tagged. + def _cn(t): + return adc.checkpoint_name(t, "moe_x_sorted") if isinstance(t, jax.Array) else t + + x = jax.tree.map(_cn, x) + routing = jax.tree.map(_cn, routing) + route_metadata = jax.tree.map(_cn, route_metadata) + mask = jnp.arange(x.shape[0]) < valid_token_count(x, routing, route_metadata) + + if self.config.mlp_bias: + w0_bias, w1_bias, wo_bias = self.transform_bias(routing.selected_experts, w0_bias, w1_bias, wo_bias) + + gmm_operands = get_gmm_operands_for_local_experts(x, routing, route_metadata) + gmm_fn = functools.partial(gmm, **gmm_operands) + up_res = gmm_up_or_silu( + x, + w0, + w1, + w0_bias, + w1_bias, + gmm_fn, + gmm_operands, + weight_gather, + mask=mask, + gather_weights=gather_weights, + return_weights=return_weights, + ) + if return_weights: + output0, output1, w0, w1 = up_res # pylint: disable=unbalanced-tuple-unpacking + else: + output0, output1 = up_res # pylint: disable=unbalanced-tuple-unpacking + + if router_grad_from_wo_input: + # moe_router_grad_from_wo_input: the activation, the down projection and the unsort as one op whose + # backward does not read the down projection's output; then the same reduce-scatter as below. + output, _ = _ring_down_project_and_unpermute( + output0, output1, wo, routing, route_metadata, batch_size, sequence_length + ) + return _moe_body_outputs(_ring_reduce_scatter(output), routing, (w0, w1, wo)) + + tc_prescaled = ( + self.config.use_ring_of_experts + and getattr(routing, "tc_routing", None) is not None + and self.config.moe_tc_ragged_weights_on_activation + ) + if tc_prescaled: + # Routing weights applied on the (buffer, mlp) activation (fuses into silu * mul), as in + # lineage ragged_silu_mul, instead of on the (buffer, emb) expert output before the unsort. + w_rows = tc_buffer_row_weights(routing.tc_routing, jnp.ravel(routing.weights), output0.shape[0]) + if self.config.moe_tc_routing_checkpoint != "remat": + w_rows = adc.checkpoint_name(w_rows, "moe_tc_routing_checkpoint") + # moe_wi_gmm_fused_swiglu: the wi GMM already applied the activation (output1 is None). + intermediate_layer = output0 if fuse_silu_in_wi_gmm else self.apply_ffn_activation_maybe_fused(output0, output1) + if tc_prescaled: + intermediate_layer = (intermediate_layer * w_rows[:, None]).astype(intermediate_layer.dtype) + wo_gather_axes, wo_tile_size = get_wo_gmm_params() + if not gather_weights: + wo_gather_axes = [] + # moe_tc_ragged_3d_gmm: the TC sort left the token buffer in the (buffer, emb // 128, 128) layout; + # the wo gmm then also writes its output in that layout for the TC unsort. + x_is_3d = self.config.num_moe_emb_chunks <= 0 and (x.qvalue if isinstance(x, qpl.QArray) else x).ndim == 3 + if x_is_3d and (self.get_tensor_parallelism_size() > 1 or self.config.mlp_bias): + raise NotImplementedError("moe_tc_ragged_3d_gmm does not support tensor parallelism or mlp_bias.") + # moe_bwd_prequant_before_unsort: fuse the TC unsort (and with moe_combine_bwd_direct_qarray, the EP + # combine) into the wo gmm so its backward quantizes the token-order cotangent before the TC gather. + # The config validation guarantees the 3D TC dispatch with routing weights on the activation. + use_fused_tc_unsort = self.config.moe_bwd_prequant_before_unsort and tc_prescaled and x_is_3d + use_fused_combine_bwd_qarray = use_fused_tc_unsort and self.config.moe_combine_bwd_direct_qarray + tc_unsort_kwargs = {} + if use_fused_tc_unsort: + tc_unsort_kwargs = { + "tc_routing": routing.tc_routing, + "tc_unsort_cfg": mblx.TcUnsortCfg( + unpadded_cap=intermediate_layer.shape[0], + num_out_tokens=routing.tc_routing.num_out_tokens, + topk=self.num_experts_per_tok, + blocks=( + self.config.moe_tc_ragged_gather_block_size, + self.config.moe_tc_ragged_flatten_block_size, + True, + True, + ), + mask_padding=self.config.moe_tc_ragged_mask_padding, + combine_rs_cfg=( + (self._expert_parallelism_name, sequence_length) if use_fused_combine_bwd_qarray else None + ), + ), + } + intermediate_output = gmm_fn( + intermediate_layer, + wo, + tiling=wo_tile_size, + weight_gather_axes=wo_gather_axes, + return_rhs=return_weights, + **({"out_is_3d": True} if x_is_3d else {}), + **tc_unsort_kwargs, + ) + if return_weights: + intermediate_output, wo = intermediate_output + if self.get_tensor_parallelism_size() > 1: + intermediate_output = jax.lax.psum_scatter( + intermediate_output, + self._tensor_parallelism_name, + scatter_dimension=1, + tiled=True, + ) + if self.config.mlp_bias: + mask = jnp.arange(intermediate_output.shape[0]) < valid_token_count(intermediate_output, routing, route_metadata) + if tc_prescaled: + # The unsort skips the routing weights, so weight the bias too: w * (h @ wo + b) = (w * h) @ wo + w * b. + wo_bias = (wo_bias * w_rows[:, None]).astype(wo_bias.dtype) + intermediate_output = intermediate_output + wo_bias + intermediate_output = jnp.where(mask[:, None], intermediate_output, 0) + intermediate_output = adc.checkpoint_name(adc.checkpoint_name(intermediate_output, "mlpwo"), "moe_mlpwo") if self.config.use_ring_of_experts: - output = _ring_reduce_scatter(_ring_unpermute(intermediate_output, routing, batch_size, sequence_length)) - return output, routing.lb_loss, routing.bias_updates, routing.has_overflow, routing.moe_audit + if use_fused_combine_bwd_qarray: + output = intermediate_output + else: + if use_fused_tc_unsort: + output = intermediate_output.reshape(-1, sequence_length, *intermediate_output.shape[1:]).astype(self.dtype) + else: + # Unsort and deduplicate the outputs locally. + output = self.unpermute( + intermediate_output, + routing.sorted_selected_experts, + routing.weights, + batch_size=batch_size, + sequence_length=sequence_length, + use_custom_sort_vjp=self.config.use_custom_sort_vjp, + group_sizes=routing.group_sizes, + topk_argsort_indices=routing.topk_argsort_indices, + tc_routing=routing.tc_routing, + tc_prescaled=tc_prescaled, + tc_out_3d=x_is_3d and use_tc_3d_dispatch(), + ) + + # Sum up the partial outputs across the expert shards. + if output.ndim == 4: + # moe_tc_ragged_3d_dispatch: reduce-scatter in the (emb // 128, 128) layout; only the local result is + # reshaped back below. + output = jnp.reshape(output, (-1, sequence_length, *output.shape[2:])) + else: + output = jnp.reshape( + output, + ( + -1, + sequence_length, + self.moe_expert_input_dim // self.get_tensor_parallelism_size(), + ), + ) + # Same reduce-scatter as the pipelined chunk loop (adds ring_of_experts_row_major_reduce_scatter). + output = _ring_reduce_scatter(output) + return _moe_body_outputs(output, routing, (w0, w1, wo)) if self.get_expert_parallelism_size() > 1: original_inputs_first_dim = batch_size * sequence_length * self.config.num_experts_per_tok @@ -3067,34 +4408,45 @@ def _moe_body( group_sizes=routing.group_sizes, ) - return output, routing.lb_loss, routing.bias_updates, routing.has_overflow, routing.moe_audit + return _moe_body_outputs(output, routing, (w0, w1, wo)) + + in_specs = ( + input_partition_pspec, + gate_logits_pspec, + pre_bias_logits_pspec, + w0_pspec, + w1_pspec, + wo_pspec, + w0_bias_pspec, + w1_bias_pspec, + wo_bias_pspec, + decoder_tokens_pspec, + P(), # Replicate the input key + # Reuse gate_logits_pspec (which forced_routed_experts is sharded + # with below); recomputing it would skip + # maybe_replicate_incompatible_batch. + gate_logits_pspec if forced_routed_experts is not None else None, + ) + # Not every mesh axis: shard_map rejects axes an enclosing vmap mapped away ("stage", "diloco"). + input_mesh_axes = {axis for spec in in_specs if spec is not None for dim in spec for axis in mesh_axes_for_dim(dim)} + overflow_mesh_axes = tuple(axis for axis in self.mesh.axis_names if axis in input_mesh_axes) @functools.partial( jax.shard_map, mesh=self.mesh, - in_specs=( - input_partition_pspec, - gate_logits_pspec, - pre_bias_logits_pspec, - w0_pspec, - w1_pspec, - wo_pspec, - w0_bias_pspec, - w1_bias_pspec, - wo_bias_pspec, - decoder_tokens_pspec, - P(), # Replicate the input key - # Reuse gate_logits_pspec (which forced_routed_experts is sharded - # with below); recomputing it would skip - # maybe_replicate_incompatible_batch. - gate_logits_pspec if forced_routed_experts is not None else None, - ), + in_specs=in_specs, out_specs=( output_pspec, P(), # Handle None or replicate the output - P(), # Handle None or replicate the output - P(), # has_overflow: replicated scalar, all-reduced across entire mesh + # bias_updates: replicated, or under defer_small_all_reduces the local counts (1, n_chunks, E) per shard of + # the axes the per-layer psum would have used. + P(roe_bias_axis_names, None, None) if defer_small_ars else P(), + P(overflow_mesh_axes), # has_overflow: unreduced, reduced once per step in loss_fn + P(), # took_fallback: replicated scalar, True if the in-layer dropless branch ran + P(), # required_rbf (probe): replicated scalar max-reduced across entire mesh, or None ) + # max_load_ratio (moe_log_max_load_ratio only): replicated scalar, max-reduced across the mesh. + + ((P(),) if log_max_load_ratio else ()) # moe_audit (log_moe_drops only): (dropped_assignments psum-reduced, max_load_ratio pmax-reduced) scalars. + (((P(), P()),) if log_moe_drops else ()), check_vma=self.config.check_vma, @@ -3132,6 +4484,10 @@ def sparse_matmul_route_and_compute( "moe_chunk_pipeline=True is skipped: --xla_tpu_aggressive_opt_barrier_removal=true (or ENABLED)" " removes optimization barriers before scheduling. Set it to false to use the pipeline." ) + if self.config.moe_roe_packed_routing_gather: + max_logging.warning( + "moe_roe_packed_routing_gather=True is skipped too: it only runs in the pipelined chunk loop." + ) pipeline_enabled = False def _pipelined_route_and_compute(force_dropless): @@ -3169,6 +4525,33 @@ def _pipelined_route_and_compute(force_dropless): seq_len = x.shape[1] chunk = seq_len // n_chunks + packed_topk = None + if self.config.moe_roe_packed_routing_gather: + # moe_roe_packed_routing_gather: route ALL of this shard's tokens once (every chunk) and all-gather the + # top-k weights and expert ids as ONE float32 [batch, seq, 2k] array: one all-gather (one SparseCore + # offload) per layer instead of two per chunk, and in the backward one reduce-scatter of the weight + # cotangent instead of one per chunk. Routing is per token, so each chunk's slice holds the values the + # per-chunk local routing would compute (pack_routing_for_all_gather: exact). + num_expert_shards = self.get_expert_parallelism_size() + expert_shard = jax.lax.axis_index(self._expert_parallelism_name) if num_expert_shards > 1 else 0 + local_weights, local_ids = self.get_topk( + logits, + pre_bias_logits, + rngs, + sharded_input_ids, + forced_routed_experts, + num_token_shards=num_expert_shards, + token_shard_index=expert_shard, + ) + packed = _ep_all_gather(pack_routing_for_all_gather(local_weights, local_ids)) + # A routing map: with moe_routing_maps=device (or offload) the backward reuses the gathered array and the + # rematerialized forward skips the gather. Untagged under moe_routing_maps=remat. + tag_fn = routing_tag_fn(self.config) + if tag_fn is not None: + packed = tag_fn(packed) + gathered_weights, gathered_ids = unpack_routing(packed, local_weights.dtype, local_ids.dtype) + packed_topk = (gathered_weights, gathered_ids, local_ids) + def _dispatch(c, x_c): sl = slice(c * chunk, (c + 1) * chunk) return route( @@ -3179,34 +4562,70 @@ def _dispatch(c, x_c): input_ids=None if sharded_input_ids is None else sharded_input_ids[:, sl], forced_routed_experts=None if forced_routed_experts is None else forced_routed_experts[:, sl, :], force_dropless=force_dropless, + gathered_topk=None if packed_topk is None else tuple(t[:, sl, :] for t in packed_topk), ) def _with_routed(dispatched_c, routed): return (routed,) + tuple(dispatched_c[1:]) dispatched = [_dispatch(0, x[:, 0:chunk, :])] + [None] * (n_chunks - 1) - outs, lb_losses, bias_updates_list, has_overflows, moe_audits = [], [], [], [], [] + outs, lb_losses, bias_updates_list, has_overflows, required_rbfs, max_load_ratios = [], [], [], [], [], [] + moe_audits = [] unpermuted = None # the previous chunk's combine input, not yet reduce-scattered + # moe_accumulate_chunk_wgrad (main): chunk 0 gathers the expert weights and every later chunk reuses them. + accum_chunk_wgrad = self.config.moe_accumulate_chunk_wgrad + cur_w0, cur_w1, cur_wo = w0, w1, wo for c in range(n_chunks): + gather_c = not accum_chunk_wgrad or c == 0 + return_c = accum_chunk_wgrad and c < n_chunks - 1 + if router_grad_from_wo_input and (return_c or not gather_c): + # The fused op does not return the weights (the config validator rejects the two keys together). + raise NotImplementedError("moe_router_grad_from_wo_input does not support moe_accumulate_chunk_wgrad.") if c + 1 < n_chunks: # 1. The next chunk's token all-gather waits for this chunk's routed tokens. x_next, routed = barrier_tie(x[:, (c + 1) * chunk : (c + 2) * chunk, :], dispatched[c][0]) dispatched[c] = _with_routed(dispatched[c], routed) dispatched[c + 1] = _dispatch(c + 1, x_next) x_c, routing, route_metadata = dispatched[c] - output0, output1, gmm_fn, wo_bias_c = _up_projection( - x_c, routing, route_metadata, w0, w1, w0_bias, w1_bias, wo_bias + x_c, routing, route_metadata, output0, output1, gmm_fn, wo_bias_c, (w0_c, w1_c) = _up_projection( + x_c, + routing, + route_metadata, + cur_w0, + cur_w1, + w0_bias, + w1_bias, + wo_bias, + gather_weights=gather_c, + return_weights=return_c, ) if unpermuted is not None: # 4. The previous chunk's reduce-scatter waits for this chunk's first expert GEMM, # so the TensorCore runs that GEMM while the gather-reduce is on the SparseCore. unpermuted, output0 = barrier_tie(unpermuted, output0) outs.append(_ring_reduce_scatter(unpermuted)) - intermediate_output = _down_projection(output0, output1, gmm_fn, wo, wo_bias_c, routing, route_metadata) - if c + 1 < n_chunks: - # 2. This chunk's gather-reduce waits for the next chunk's routed tokens. - intermediate_output, routed = barrier_tie(intermediate_output, dispatched[c + 1][0]) - dispatched[c + 1] = _with_routed(dispatched[c + 1], routed) + if router_grad_from_wo_input: + # moe_router_grad_from_wo_input: the fused op below runs the down projection, after tie 3. + intermediate_output, tc_prescaled, x_is_3d = None, False, False + else: + intermediate_output, wo_c, tc_prescaled, x_is_3d = _down_projection( + x_c, + output0, + output1, + gmm_fn, + cur_wo, + wo_bias_c, + routing, + route_metadata, + gather_weights=gather_c, + return_weights=return_c, + ) + if return_c: + cur_w0, cur_w1, cur_wo = w0_c, w1_c, wo_c + if c + 1 < n_chunks: + # 2. This chunk's gather-reduce waits for the next chunk's routed tokens. + intermediate_output, routed = barrier_tie(intermediate_output, dispatched[c + 1][0]) + dispatched[c + 1] = _with_routed(dispatched[c + 1], routed) if outs: # 3. This chunk's gather-reduce waits for the previous chunk's reduce-scatter. The tie holds the # gather-reduce's routing weights, a small input: holding the down-projection output instead @@ -3217,17 +4636,46 @@ def _with_routed(dispatched_c, routed): # run before the previous chunk's reduce-scatter. weights, outs[-1] = barrier_tie(routing.weights, outs[-1]) routing = routing.replace(weights=weights) - unpermuted = _ring_unpermute(intermediate_output, routing, x.shape[0], chunk) + if router_grad_from_wo_input: + # moe_router_grad_from_wo_input: the down projection and the unsort as one op. Tie 2 moves inside it + # (forward only, as barrier_tie's backward=False): this chunk's down-projection output and the next + # chunk's routed tokens. Tie 3 above already holds the routing weights the op's gather-reduce reads. + unpermuted, routed = _ring_down_project_and_unpermute( + output0, + output1, + cur_wo, + routing, + route_metadata, + x.shape[0], + chunk, + forward_tie=dispatched[c + 1][0] if c + 1 < n_chunks else None, + ) + if c + 1 < n_chunks: + dispatched[c + 1] = _with_routed(dispatched[c + 1], routed) + else: + unpermuted = _ring_unpermute( + intermediate_output, routing, x.shape[0], chunk, tc_prescaled=tc_prescaled, x_is_3d=x_is_3d + ) lb_losses.append(routing.lb_loss) bias_updates_list.append(routing.bias_updates) has_overflows.append(routing.has_overflow) + required_rbfs.append(routing.required_rbf) + max_load_ratios.append(routing.max_load_ratio) moe_audits.append(routing.moe_audit) outs.append(_ring_reduce_scatter(unpermuted)) output = jnp.concatenate(outs, axis=1) lb_loss = None if lb_losses[0] is None else sum(lb_losses) / n_chunks - bias_updates = None if bias_updates_list[0] is None else sum(bias_updates_list) / n_chunks + # Same bias-signal convention as the plain chunk loop below (permute returns expert counts). + if bias_updates_list[0] is None: + bias_updates = None + elif defer_small_ars: + bias_updates = jnp.stack(bias_updates_list) + else: + bias_updates = finalize_bias_updates(sum(bias_updates_list)) has_overflow = jnp.any(jnp.stack(has_overflows)) - return output, lb_loss, bias_updates, has_overflow, combine_moe_audits(moe_audits) + required_rbf = None if required_rbfs[0] is None else jnp.max(jnp.stack(required_rbfs)) + max_load_ratio = None if max_load_ratios[0] is None else jnp.max(jnp.stack(max_load_ratios)) + return output, lb_loss, bias_updates, has_overflow, required_rbf, max_load_ratio, combine_moe_audits(moe_audits) def _scanned_dropless_chunks(): """The dropless chunks as a `lax.scan`, each chunk rematerialized on its own (the dropless replay program). @@ -3248,7 +4696,7 @@ def split(a): # [batch, seq, ...] -> [n_chunks, batch, chunk, ...] @functools.partial(jax.checkpoint, policy=jax.checkpoint_policies.nothing_saveable) def body(carry, xs_c): x_c, logits_c, pre_bias_logits_c, input_ids_c, forced_c = xs_c - out_c, lb_c, bu_c, ov_c, audit_c = _moe_body( + out_c, lb_c, bu_c, ov_c, req_c, ratio_c, audit_c, _ = _moe_body( x_c, logits_c, pre_bias_logits_c, @@ -3263,25 +4711,33 @@ def body(carry, xs_c): forced_c, force_dropless=True, ) - return carry, (out_c, lb_c, bu_c, jnp.asarray(ov_c, dtype=jnp.bool_), audit_c) + return carry, (out_c, lb_c, bu_c, jnp.asarray(ov_c, dtype=jnp.bool_), req_c, ratio_c, audit_c) xs = tuple(split(a) for a in (x, logits, pre_bias_logits, sharded_input_ids, forced_routed_experts)) - _, (outs, lb, bias_upd, ov, audit) = jax.lax.scan(body, None, xs) + _, (outs, lb, bias_upd, ov, req, ratio, audit) = jax.lax.scan(body, None, xs) output = jnp.moveaxis(outs, 0, 1).reshape(outs.shape[1], seq_len, *outs.shape[3:]) lb_loss = None if lb is None else jnp.mean(lb, axis=0) - bias_updates = None if bias_upd is None else jnp.mean(bias_upd, axis=0) + # Same bias-signal convention as the unrolled chunk loops (permute returns expert counts). + if bias_upd is None: + bias_updates = None + elif defer_small_ars: + bias_updates = bias_upd + else: + bias_updates = finalize_bias_updates(jnp.sum(bias_upd, axis=0)) + required_rbf = None if req is None else jnp.max(req) + max_load_ratio = None if ratio is None else jnp.max(ratio) # log_moe_drops: per-chunk audits stacked on the scan axis, combined as the unrolled loops combine them. audit = None if audit is None else (jnp.sum(audit[0]), jnp.max(audit[1])) - return output, lb_loss, bias_updates, jnp.any(ov), audit + return output, lb_loss, bias_updates, jnp.any(ov), required_rbf, max_load_ratio, audit - def _route_and_compute(force_dropless): + def _route_and_compute(force_dropless, n_chunks=n_chunks, barrier_enabled=barrier_enabled): """Runs route+compute once; force_dropless=True redoes all n_chunks, not just the overflowing one(s).""" if force_dropless and n_chunks > 1 and self.config.use_ring_of_experts and self.dropless_scan_chunks: return _scanned_dropless_chunks() if pipeline_enabled: return _pipelined_route_and_compute(force_dropless) if n_chunks <= 1 or not self.config.use_ring_of_experts: - return _moe_body( + out, lb_loss, bias_updates, has_overflow, required_rbf, max_load_ratio, moe_audit, _ = _moe_body( x, logits, pre_bias_logits, @@ -3296,18 +4752,30 @@ def _route_and_compute(force_dropless): forced_routed_experts, force_dropless=force_dropless, ) + if not defer_small_ars: + bias_updates = finalize_bias_updates(bias_updates) + return out, lb_loss, bias_updates, has_overflow, required_rbf, max_load_ratio, moe_audit # Chunked ring-of-experts pipeline: split the per-shard tokens along the # sequence dim into `n_chunks` data-independent chunks. Each chunk runs the # full route -> GMM -> combine path; with no barrier between them XLA is # free to overlap chunk (c+1)'s EP all-gather and chunk (c-1)'s # reduce-scatter with chunk c's GMM compute. Token routing is per-token, so - # the main (lm) output is identical to n_chunks=1; only the aggregate - # load-balance loss / bias updates are averaged across chunks. + # the main (lm) output is identical to n_chunks=1. The bias update (one sign + # of the expert counts summed over chunks) and the sigmoid-router aux loss + # (computed on full sequences after the shard_map) do not depend on n_chunks; + # the Switch-style loss of other routers is averaged across chunks. seq_len = x.shape[1] chunk = seq_len // n_chunks - outs, lb_losses, bias_updates_list, has_overflows, moe_audits = [], [], [], [], [] + outs, lb_losses, bias_updates_list, has_overflows, required_rbfs, max_load_ratios = [], [], [], [], [], [] + moe_audits = [] _prev = None + # With moe_accumulate_chunk_wgrad, each chunk's GMMs return the (gathered, possibly quantized) + # weights they consumed and the next chunk reuses them. The weights are then gathered and + # quantized once (chunk 0), and in bwd each chunk's tgmm accumulates its weight gradient in place + # onto the later chunks' gradient instead of materializing and summing one per chunk. + accum_chunk_wgrad = self.config.moe_accumulate_chunk_wgrad + cur_w0, cur_w1, cur_wo = w0, w1, wo for c in range(n_chunks): sl = slice(c * chunk, (c + 1) * chunk) x_c = x[:, sl, :] @@ -3317,13 +4785,13 @@ def _route_and_compute(force_dropless): # loss stays bit-exact. if barrier_enabled and _prev is not None: x_c, _prev = jax.lax.optimization_barrier((x_c, _prev)) - out_c, lb_c, bu_c, ov_c, audit_c = _moe_body( + out_c, lb_c, bu_c, ov_c, req_c, ratio_c, audit_c, chunk_weights = _moe_body( x_c, logits[:, sl, :], None if pre_bias_logits is None else pre_bias_logits[:, sl, :], - w0, - w1, - wo, + cur_w0, + cur_w1, + cur_wo, w0_bias, w1_bias, wo_bias, @@ -3331,28 +4799,106 @@ def _route_and_compute(force_dropless): rngs, None if forced_routed_experts is None else forced_routed_experts[:, sl, :], force_dropless=force_dropless, + gather_weights=not accum_chunk_wgrad or c == 0, + return_weights=accum_chunk_wgrad and c < n_chunks - 1, ) + if accum_chunk_wgrad: + cur_w0, cur_w1, cur_wo = chunk_weights if barrier_enabled: _prev = out_c outs.append(out_c) lb_losses.append(lb_c) bias_updates_list.append(bu_c) has_overflows.append(ov_c) + required_rbfs.append(req_c) + max_load_ratios.append(ratio_c) moe_audits.append(audit_c) output = jnp.concatenate(outs, axis=1) lb_loss = None if lb_losses[0] is None else sum(lb_losses) / n_chunks - bias_updates = None if bias_updates_list[0] is None else sum(bias_updates_list) / n_chunks + if bias_updates_list[0] is None: + bias_updates = None + elif defer_small_ars: + # Local per-chunk counts, reduced over the mesh and combined over chunks in finalize_deferred_bias_signal. + bias_updates = jnp.stack(bias_updates_list) + else: + bias_updates = finalize_bias_updates(sum(bias_updates_list)) has_overflow = jnp.any(jnp.stack(has_overflows)) - return output, lb_loss, bias_updates, has_overflow, combine_moe_audits(moe_audits) + # Each chunk has its own buffer, so the step needs the largest per-chunk factor. + required_rbf = None if required_rbfs[0] is None else jnp.max(jnp.stack(required_rbfs)) + max_load_ratio = None if max_load_ratios[0] is None else jnp.max(jnp.stack(max_load_ratios)) + return output, lb_loss, bias_updates, has_overflow, required_rbf, max_load_ratio, combine_moe_audits(moe_audits) force_dropless = getattr(self, "force_dropless", False) - out, lb_loss, bias_updates, has_overflow, moe_audit = _route_and_compute(force_dropless=force_dropless) - if log_moe_drops: - if moe_audit is None: - # Flag on but not on the ring-of-experts ragged path: -1 marks "not computed" (never reads as 0). - moe_audit = (jnp.int32(MOE_AUDIT_NOT_COMPUTED), jnp.float32(MOE_AUDIT_NOT_COMPUTED)) - return out, lb_loss, bias_updates, has_overflow, moe_audit - return out, lb_loss, bias_updates, has_overflow + + def _per_device_flag(has_overflow): + return jnp.reshape(jnp.asarray(has_overflow, dtype=jnp.bool_), (1,)) + + fallback_chunks = None + if self.config.moe_dropless_fallback == "layer": + fallback_chunks = self.get_dropless_fallback_num_chunks(x.shape[1], n_chunks) + + def _extra_outputs(max_load_ratio, moe_audit): + """The optional trailing shard_map outputs: max_load_ratio, then the log_moe_drops audit.""" + extra = () + if log_max_load_ratio: + extra += (max_load_ratio,) + if log_moe_drops: + if moe_audit is None: + # Flag on but not on the ring-of-experts ragged path: -1 marks "not computed" (never reads as 0). + moe_audit = (jnp.int32(MOE_AUDIT_NOT_COMPUTED), jnp.float32(MOE_AUDIT_NOT_COMPUTED)) + extra += (moe_audit,) + return extra + + if fallback_chunks is None: + out, lb_loss, bias_updates, has_overflow, required_rbf, max_load_ratio, moe_audit = _route_and_compute( + force_dropless + ) + has_overflow = _per_device_flag(has_overflow) + if defer_small_ars and bias_updates is not None: + if bias_updates.ndim == 1: # unchunked path: one chunk + bias_updates = bias_updates[None, :] + bias_updates = bias_updates[None] # (1, n_chunks, E): this shard's partial counts + return (out, lb_loss, bias_updates, has_overflow, jnp.bool_(False), required_rbf) + _extra_outputs( + max_load_ratio, moe_audit + ) + + # In-layer dropless fallback: calculate the overflow, then run either the + # normal capped path or a chunked dropless path for this layer only (no step replay). + took_fallback = self._check_ragged_overflow( + logits, pre_bias_logits, rngs, sharded_input_ids, forced_routed_experts, n_chunks + ) + + # Both branches return required_rbf with the same structure: None when the + # log_required_ragged_buffer_factor probe is off, a float32 scalar otherwise. The dropless + # branch returns the probe value of its own (fallback_chunks) chunking. + def _capped_branch(): + out, lb_loss, bias_updates, has_overflow, required_rbf, max_load_ratio, moe_audit = _route_and_compute(False) + return ( + out, + lb_loss, + bias_updates, + jnp.asarray(has_overflow, dtype=jnp.bool_), + required_rbf, + max_load_ratio, + moe_audit, + ) + + # Save nothing from the rare dropless branch: lax.cond gives both branches the union of + # their residuals, so the (larger) dropless residuals would otherwise be held every step. + @functools.partial(jax.checkpoint, policy=jax.checkpoint_policies.nothing_saveable) + def _dropless_branch(): + out, lb_loss, bias_updates, _, required_rbf, max_load_ratio, moe_audit = _route_and_compute( + True, n_chunks=fallback_chunks, barrier_enabled=True + ) + return out, lb_loss, bias_updates, jnp.bool_(False), required_rbf, max_load_ratio, moe_audit + + out, lb_loss, bias_updates, has_overflow, required_rbf, max_load_ratio, moe_audit = jax.lax.cond( + took_fallback, _dropless_branch, _capped_branch + ) + has_overflow = _per_device_flag(has_overflow) + return (out, lb_loss, bias_updates, has_overflow, took_fallback, required_rbf) + _extra_outputs( + max_load_ratio, moe_audit + ) # Note: SparseCore pinning (`moe_pin_sparse_core_all_gathers`) is currently not supported # with `moe_fsdp_use_two_stage_all_gather` (enforced via validation in types.py). @@ -3400,7 +4946,8 @@ def _route_and_compute(force_dropless): logical_axes=gate_logits_logical_axes, ) - if self.config.moe_pin_sparse_core_all_gathers: + if self.config.moe_pin_sparse_core_all_gathers or getattr(self.config, "moe_pin_sparse_core_fsdp_all_gathers", False): + fwd_only = getattr(self.config, "moe_pin_sparse_core_fsdp_all_gathers_fwd_only", False) def _fsdp_all_gather(w, pspec): if w is None or pspec is None: @@ -3415,7 +4962,24 @@ def _fsdp_all_gather(w, pspec): def _reshard_fn(x): return self._maybe_shard_with_pspec(x, pspec) - return _reshard_fn(w) + if not fwd_only: + return _reshard_fn(w) + + # Forward: the pinned SparseCore all-gather. Backward: the transpose of the unpinned reshard (a sharding + # constraint on the cotangent), so the weight-gradient reduce-scatter and its cross-slice all-reduce are + # not wrapped in a SparseCore compute_on region. + @jax.custom_vjp + def _pinned_fwd_reshard(x): + return _reshard_fn(x) + + def _pinned_fwd_reshard_fwd(x): + return _reshard_fn(x), None + + def _pinned_fwd_reshard_bwd(_, g): + return (self._maybe_shard_with_pspec(g, pspec),) + + _pinned_fwd_reshard.defvjp(_pinned_fwd_reshard_fwd, _pinned_fwd_reshard_bwd) + return _pinned_fwd_reshard(w) else: @@ -3479,13 +5043,28 @@ def _fsdp_all_gather(w, pspec): self.rngs, forced_routed_experts, ) - output, lb_loss, bias_updates, has_overflow = route_and_compute_outputs[:4] - self.sow(nnx.Intermediate, "moe_has_overflow", has_overflow) - if log_moe_drops: + output, lb_loss, bias_updates, has_overflow, took_fallback, required_rbf = route_and_compute_outputs[:6] + max_load_ratio = route_and_compute_outputs[6] if log_max_load_ratio else None + moe_audit = route_and_compute_outputs[-1] if log_moe_drops else None + if use_megatron_seq_aux_loss: + # Global (batch, sequence) arrays: one loss per full sequence (not per token chunk), averaged + # over the global batch. + lb_loss = self.megatron_seq_aux_loss(pre_bias_logits if pre_bias_logits is not None else gate_logits) + if getattr(self, "force_dropless", False): + self.sow(nnx.Intermediate, "moe_has_overflow", jnp.zeros((), jnp.bool_)) + else: + self.sow(nnx.Intermediate, "moe_has_overflow", has_overflow) + if self.config.moe_dropless_fallback == "layer": + self.sow(nnx.Intermediate, "moe_dropless_fallback", took_fallback) + if required_rbf is not None: + self.sow(nnx.Intermediate, "moe_required_rbf", required_rbf) + if max_load_ratio is not None: + self.sow(nnx.Intermediate, "moe_max_load_ratio", max_load_ratio) + if moe_audit is not None: # One entry per MoE layer call (token chunks combined): assignments the ragged buffer dropped, summed # over every EP shard of the mesh, and the max expert-shard load / balanced load. - self.sow(nnx.Intermediate, "moe_dropped_assignments", route_and_compute_outputs[4][0]) - self.sow(nnx.Intermediate, "moe_max_load_ratio", route_and_compute_outputs[4][1]) + self.sow(nnx.Intermediate, "moe_dropped_assignments", moe_audit[0]) + self.sow(nnx.Intermediate, "moe_max_load_ratio", moe_audit[1]) return output, lb_loss, bias_updates def reshape_and_update_weights(self, weights, indices, safe_updates=False): @@ -3711,7 +5290,7 @@ def router_load_balance_loss(self, top_k_indices, gate_logits, pre_bias_logits) return self.load_balance_loss(top_k_indices, softmax_probs) # See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. - def load_balance_loss(self, top_k_indices, logits) -> jax.Array: + def load_balance_loss(self, top_k_indices, logits, density_prob=None) -> jax.Array: """Compute the sequence-wise load balance loss. For DeepSeek V4 like models, standard load balancing across an entire batch can @@ -3724,6 +5303,8 @@ def load_balance_loss(self, top_k_indices, logits) -> jax.Array: into the total training loss. By minimizing this scaled auxiliary loss, the optimizer updates the routing parameters to actively enforce an even distribution of tokens to experts within each individual sequence. + + `density_prob` optionally passes the precomputed per-sequence mean of `logits` (batch, num_experts). """ expert_mask = jax.nn.one_hot(top_k_indices, num_classes=self.num_experts, dtype=jnp.int32) summed_expert_mask = jnp.sum(expert_mask, axis=2) @@ -3733,12 +5314,28 @@ def load_balance_loss(self, top_k_indices, logits) -> jax.Array: density = jnp.mean(summed_expert_mask, axis=1) / self.num_experts_per_tok # get fraction of probability allocated to each expert # jnp.mean over axis=1 isolates the routing probability per sequence. - density_prob = jnp.mean(logits, axis=1) + if density_prob is None: + density_prob = jnp.mean(logits, axis=1) # The sequence-wise densities and probabilities are multiplied and then averaged # over the batch dimension, scaled by the required constant. loss = jnp.mean(density * density_prob) * (self.num_experts**2) * self.config.load_balance_loss_weight return loss + def megatron_seq_aux_loss(self, scores) -> jax.Array: + """Sequence-wise aux loss of the Megatron-LM sigmoid router (seq_aux_loss). + + Used for sigmoid routers (see `uses_megatron_seq_aux_loss`). Pass whole sequences: the token + fractions and probabilities are per-sequence statistics. + + Args: + scores: Sigmoid router scores before the expert bias, shape (batch, sequence, num_experts). + + Returns: + Scalar loss averaged over the batch, already scaled by load_balance_loss_weight. + """ + aux_probs, aux_top_k_indices = megatron_aux_loss_inputs(scores, self.num_experts_per_tok) + return self.load_balance_loss(aux_top_k_indices, aux_probs) + def get_einsum( self, rhs_mesh_axes: Tuple[Optional[str], ...] = (), @@ -3767,7 +5364,7 @@ def quant_einsum(*args, **kwargs): # pylint: disable=unused-argument f"Einsum name '{op_id}' is not registered in quant_einsums. " f"Available names: {list(self.quant_einsums.keys())}" ) - return self.quant_einsums[op_id](*args, mutable=["_overwrite_with_gradient"]) + return self.quant_einsums[op_id](*args) einsum = self.quant.einsum(mesh_axes=rhs_mesh_axes) # pytype: disable=attribute-error return quantizations.apply_einsum_in_nnx(self, op_id, einsum, ["aqt"], *args) @@ -3840,10 +5437,14 @@ def dense_matmul( # DeepSeek V4 uses Hash Routing for the first `first_num_hash_layers` layers. # These layers route deterministically based on token IDs and do not generate auxiliary loss. if self.config.model_call_mode != "inference" and not self.is_hash_routing: - softmax_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1).astype(self.dtype) - lb_loss = ( - self.load_balance_loss(top_k_indices, softmax_probs) if self.config.load_balance_loss_weight > 0.0 else None - ) + if self.config.load_balance_loss_weight > 0.0 and uses_megatron_seq_aux_loss(self.config): + # Same loss as sparse_matmul: Megatron-LM's seq_aux_loss on the pre-bias sigmoid scores. + lb_loss = self.megatron_seq_aux_loss(pre_bias_logits if pre_bias_logits is not None else gate_logits) + else: + softmax_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1).astype(self.dtype) + lb_loss = ( + self.load_balance_loss(top_k_indices, softmax_probs) if self.config.load_balance_loss_weight > 0.0 else None + ) # TODO(dipakg-lang, b/521990776): Add sequence-wise balance loss * 0.0001 else: lb_loss = None @@ -3853,12 +5454,19 @@ def dense_matmul( if self.should_update_load_balance(): # Under plain JIT (dense_matmul), GSPMD automatically inserts the all-reduce # across partitioned batch/sequence axes on global arrays, so axis_names=None. - bias_updates = calculate_load_balance_updates( - top_k_indices, - self.config.num_experts, - self.config.routed_bias_update_rate, - axis_names=None, - ) + if getattr(self.config, "gradient_accumulation_steps", 1) > 1: + bias_updates = calculate_expert_counts( + top_k_indices, + self.config.num_experts, + axis_names=None, + ) + else: + bias_updates = calculate_load_balance_updates( + top_k_indices, + self.config.num_experts, + self.config.routed_bias_update_rate, + axis_names=None, + ) else: bias_updates = None @@ -4193,6 +5801,9 @@ def fused_moe_matmul( rule = quantizations.get_fused_moe_rule() quantized_w1, w1_scale = quantizations.quantize_weight_for_fused_moe(fused_kernel, rule) quantized_w2, w2_scale = quantizations.quantize_weight_for_fused_moe(wo_kernel, rule) + # serve_fp8_weight with a fixed act_quantization_calibration_method: the experts use the + # same static activation scale as the dense layers ({} keeps the kernel's dynamic scale). + lhs_scale_kwargs = fused_moe_lhs_scale_kwargs(getattr(self, "quant", None), fused_moe_func) fused_moe = quantizations.without_qwix_interception(fused_moe_func) output_2d = fused_moe( @@ -4210,6 +5821,7 @@ def fused_moe_matmul( use_ep=use_ep, activation=activation, scoring_fn=scoring_fn, + **lhs_scale_kwargs, # Forward the same environment-backed kernel knobs that tpu-inference passes on its # own serving path (tpu_inference/layers/common/moe.py). Without these, env vars such # as ONEHOT_MOE_PERMUTE_THRESHOLD and VLLM_MOE_CHUNK_SIZE are silently ignored when a @@ -4466,9 +6078,15 @@ def _maybe_native_fused_weight(kernel, scale): elif cfg.prefuse_moe_weights: wi_scale = self.wi_scale[...] if self.wi_scale is not None else None wi = quantizations.dequantize_weight(self.wi[...], wi_scale, self.dtype) - n = wi.shape[-1] // 2 - w0_kernel = wi[..., :n] - w1_kernel = wi[..., n:] + if self._keep_prefused_wi_whole(): + # moe_unsplit_prefused_wi: w0_kernel carries the whole [w0 | w1] and w1_kernel=None marks the layout, so + # the fake-quant, the FSDP all-gather and the prefetch barrier see one tensor and the up-projection reads it + # without the concatenate that would otherwise put the gathered halves back together. + w0_kernel, w1_kernel = wi, None + else: + n = wi.shape[-1] // 2 + w0_kernel = wi[..., :n] + w1_kernel = wi[..., n:] else: wi_0_scale = self.wi_0_scale[...] if self.wi_0_scale is not None else None wi_1_scale = self.wi_1_scale[...] if self.wi_1_scale is not None else None @@ -4485,7 +6103,8 @@ def _maybe_native_fused_weight(kernel, scale): # Axis 2 is the output dim of wi_0/wi_1 (mlp) and wo (embed). An fp8_moe rollout quantizes the # fused wi = [wi_0 | wi_1] per output channel, so rounding the halves separately is identical. w0_kernel = quantizations.fake_quantize_weight_per_channel(w0_kernel) - w1_kernel = quantizations.fake_quantize_weight_per_channel(w1_kernel) + if w1_kernel is not None: # None: moe_unsplit_prefused_wi, and the line above rounded both halves at once + w1_kernel = quantizations.fake_quantize_weight_per_channel(w1_kernel) wo_kernel = quantizations.fake_quantize_weight_per_channel(wo_kernel) # For fused MoE path (inference only), if we have not fused expert @@ -4612,6 +6231,8 @@ def __init__( self.config.emb_dim if self.config.moe_expert_input_dim <= 0 else self.config.moe_expert_input_dim ) + # The router ('embed_router') has its own logical axis so a rule set can + # replicate or shard it independently from the routed experts ('embed_moe'). # NOTE: the name MoeBlock_0 is to ensure reverse compatibility with # existing checkpoints for routed experts. self.MoeBlock_0 = RoutedMoE( @@ -4620,7 +6241,7 @@ def __init__( num_experts_per_tok=self.config.num_experts_per_tok, mesh=self.mesh, kernel_init=self.kernel_init, - kernel_axes=("embed_moe", None), + kernel_axes=("embed_router", None), intermediate_dim=self.config.moe_mlp_dim, dtype=self.config.dtype, weight_dtype=self.config.weight_dtype, @@ -4686,110 +6307,3 @@ def __call__( out_sharding=out_sharding, ) return routed_experts + shared_experts, load_balance_loss, moe_bias_updates - - -def get_gate_logit( - inputs_shape: tuple[int, ...], - out_features_shape: Union[Iterable[int], int], - model_name: str, - axis: Union[Iterable[int], int] = -1, - weight_dtype: ctypes.DType = jnp.float32, - dtype: ctypes.DType = jnp.float32, - kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "truncated_normal"), - kernel_axes: Tuple[Optional[str], ...] = (), - use_bias: bool = False, - score_func: str = "", - quant: Optional[quantizations.AqtQuantization] = None, - matmul_precision: str = "default", - name: Optional[str] = None, -): - """Creates a GateLogit Linen module.""" - - axis = linears.canonicalize_tuple(axis) - in_features_shape = tuple(inputs_shape[ax] for ax in linears.normalize_axes(axis, len(inputs_shape))) - - module = nnx_wrappers.to_linen( - GateLogit, - in_features_shape=in_features_shape, - out_features_shape=out_features_shape, - model_name=model_name, - axis=axis, - weight_dtype=weight_dtype, - dtype=dtype, - kernel_init=kernel_init, - kernel_axes=kernel_axes, - use_bias=use_bias, - score_func=score_func, - quant=quant, - matmul_precision=matmul_precision, - name=name, - metadata_fn=variable_to_logically_partitioned, - abstract_init=False, - ) - return module - - -def get_routed_moe( - config: ctypes.Config, - num_experts: int, - num_experts_per_tok: int, - mesh: jax.sharding.Mesh, - kernel_init: NdInitializer, - kernel_axes: Tuple[Optional[str], ...], - intermediate_dim: int = 2048, - weight_dtype: ctypes.DType = jnp.float32, - dtype: ctypes.DType = jnp.float32, - quant: Optional[quantizations.AqtQuantization] = None, - name: Optional[str] = None, - force_dropless: bool = False, -): - """Creates a RoutedMoE Linen module.""" - - module = nnx_wrappers.to_linen( - RoutedMoE, - config=config, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - mesh=mesh, - kernel_init=kernel_init, - kernel_axes=kernel_axes, - intermediate_dim=intermediate_dim, - weight_dtype=weight_dtype, - dtype=dtype, - quant=quant, - name=name, - force_dropless=force_dropless, - metadata_fn=variable_to_logically_partitioned, - abstract_init=False, - ) - return module - - -def get_routed_and_shared_moe( - config: ctypes.Config, - mesh: jax.sharding.Mesh, - kernel_init: NdInitializer, - kernel_axes: Tuple[Optional[str], ...], - weight_dtype: ctypes.DType = jnp.float32, - dtype: ctypes.DType = jnp.float32, - quant: Optional[quantizations.AqtQuantization] = None, - name: Optional[str] = None, - is_hash_routing: bool = False, -): - """Creates a RoutedAndSharedMoE Linen module.""" - - module = nnx_wrappers.to_linen( - RoutedAndSharedMoE, - config=config, - mesh=mesh, - kernel_init=kernel_init, - kernel_axes=kernel_axes, - weight_dtype=weight_dtype, - dtype=dtype, - quant=quant, - name=name, - is_hash_routing=is_hash_routing, - metadata_fn=variable_to_logically_partitioned, - abstract_init=False, - ) - return module diff --git a/src/maxtext/layers/multi_token_prediction.py b/src/maxtext/layers/multi_token_prediction.py index d751f9fb96..9d1451641f 100644 --- a/src/maxtext/layers/multi_token_prediction.py +++ b/src/maxtext/layers/multi_token_prediction.py @@ -22,10 +22,12 @@ import jax.numpy as jnp from jax.sharding import Mesh from maxtext.common.common_types import Config, DecoderBlockType, MODEL_MODE_TRAIN, ShardMode +from maxtext.layers import mhc from maxtext.layers import moe from maxtext.layers.linears import DenseGeneral from maxtext.layers.nnx_decoders import NNXDecoderLayer from maxtext.layers.normalizations import RMSNorm +from maxtext.models import deepseek4 from maxtext.models import deepseek_batchsplit from maxtext.utils import max_utils from maxtext.utils import maxtext_utils @@ -33,6 +35,15 @@ from maxtext.utils.globals import EPS +def _is_deepseek4(config: Config) -> bool: + """Returns True when the model uses the DeepSeek-V4 decoder block. + + `decoder_block` may hold either the `DecoderBlockType` enum or its raw string value + depending on how the config was constructed, so both are accepted. + """ + return config.decoder_block in (DecoderBlockType.DEEPSEEK4, DecoderBlockType.DEEPSEEK4.value) + + # Custom Variable types for MTP intermediate outputs # These will be automatically converted to Linen mutable collections by ToLinen wrapper # The class names become collection names directly (no case conversion) @@ -65,6 +76,14 @@ def _shift_left_one_cp_aware(x: jnp.ndarray, axis_name: str = "context") -> jnp. Returns: Array shaped like x, left-shifted by 1 across CP boundaries. """ + # jnp.roll cannot slice an explicitly sharded sequence axis, so shift a copy + # gathered along it and restore the original sharding. + x_sharding = jax.typeof(x).sharding + spec = x_sharding.spec.partitions + if len(spec) > 1 and spec[1] is not None: + gathered = jax.sharding.NamedSharding(x_sharding.mesh, jax.sharding.PartitionSpec(spec[0], None, *spec[2:])) + return jax.reshard(_shift_left_one_cp_aware(jax.reshard(x, gathered), axis_name), x_sharding) + local_rolled = jnp.roll(x, -1, axis=1) # Mask for the last position along axis=1 (avoids .at[...].set() scatter). @@ -198,12 +217,42 @@ def __init__( rngs=rngs, ) # Use MODEL_MODE_TRAIN for initialization; runtime model_mode is passed dynamically. - self.transformer_layer = transformer_layer_module( - config=cfg, - mesh=mesh, - model_mode=MODEL_MODE_TRAIN, - rngs=rngs, - ) + self.is_mhc = _is_deepseek4(cfg) and cfg.mhc_expansion_rate > 1 + if _is_deepseek4(cfg): + # `Decoder.get_decoder_layers()` returns `DeepSeek4ScannableBlock` when + # `scan_layers=True`, which bundles one HCA and one CSA layer. An MTP module is a + # single standard layer, so build `DeepSeek4DecoderLayer` directly rather than + # using the passed-in blueprint. `compress_ratio=0` selects sliding-window + # attention and `is_hash_routing=False` keeps it out of the hash-routed prefix, + # matching the reference MTP block. `layer_idx` places this layer after the main + # trunk, as the reference does with `MTPBlock(args.n_layers + layer_id)`. + self.transformer_layer = deepseek4.DeepSeek4DecoderLayer( + config=cfg, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=rngs, + layer_idx=cfg.base_num_decoder_layers + layer_number - 1, + compress_ratio=0, + is_hash_routing=False, + ) + else: + self.transformer_layer = transformer_layer_module( + config=cfg, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=rngs, + ) + + if self.is_mhc: + # Collapses this module's `[batch, seq, mhc_expansion_rate, emb]` output down to + # `[batch, seq, emb]` for the logits head. The main decoder's `hc_head` is not + # reused: each MTP depth is a separate sublayer with its own learned collapse, + # mirroring the reference implementation. + self.hc_head = mhc.DeepSeek4HyperHead( + config=cfg, + mesh=mesh, + rngs=rngs, + ) self.final_norm = RMSNorm( num_features=cfg.emb_dim, @@ -288,6 +337,14 @@ def __call__( embedding_norm = self.embedding_norm(target_token_embedding) hidden_state_norm = self.hidden_state_norm(prev_hidden_state) + + if self.is_mhc: + # `prev_hidden_state` arrives as `[batch, seq, mhc_expansion_rate, emb]` because the + # decoder no longer collapses its mHC streams. The target embedding is `[batch, seq, + # emb]`, so broadcast it across the stream axis to make the two concatenable. + mhc_expand, _ = mhc.get_functions(self.config.mhc_expansion_rate) + embedding_norm = mhc_expand(embedding_norm) + if self.config.shard_mode == ShardMode.EXPLICIT: # Ensure that the embedding norm has the same sharding as the hidden state # norm. @@ -432,7 +489,22 @@ def __call__( decoder_segment_ids, model_mode, deterministic, + main_token_embeddings=None, ) -> dict: + """Runs the MTP layers and sows their losses. + + Args: + shared_embedding: Shared token embedding, used for the MTP input embeddings and the output head. + main_hidden_state: Final (pre-norm) hidden state of the main decoder. + input_ids: Main decoder input tokens. + target_ids: Main decoder target tokens. + target_mask: Main decoder target mask or packed segment ids. + position_ids: Main decoder positions. + decoder_segment_ids: Main decoder segment ids. + model_mode: Operational mode. + deterministic: Whether dropout is disabled. + main_token_embeddings: Token embeddings of ``input_ids`` from the main decoder (mtp_reuse_input_embedding). + """ cfg = self.config cp_size = self.mesh.shape.get(cfg.context_sharding, 1) if self.mesh is not None else 1 @@ -492,7 +564,22 @@ def __call__( mtp_preds_list = [] mtp_masks_list = [] + # Token embeddings of rolled_input_ids, when shifted from the main decoder's embeddings. + rolled_token_embeddings = main_token_embeddings + if main_token_embeddings is not None: + # Rolled-in positions hold token id 0, so they take the embedding of token 0. + token_zero_embedding = shared_embedding.embed_single_token(0) + not_filled = jnp.ones(input_ids.shape, dtype=jnp.int32) + for k in range(1, cfg.mtp_num_layers + 1): + if main_token_embeddings is not None: + # roll_and_mask_by_segment zero-fills the positions whose token id it replaces with 0. + is_filled = roll_and_mask_by_segment(not_filled, rolled_segment_ids) == 0 + rolled_token_embeddings = jnp.where( + is_filled[..., None], + token_zero_embedding, + roll_and_mask_by_segment(rolled_token_embeddings, rolled_segment_ids), + ) rolled_input_ids = roll_and_mask_by_segment(rolled_input_ids, rolled_segment_ids) rolled_target_ids = roll_and_mask_by_segment(rolled_target_ids, rolled_segment_ids) rolled_target_mask = roll_and_mask_by_segment(rolled_target_mask, rolled_segment_ids) @@ -501,12 +588,14 @@ def __call__( if rolled_segment_ids is not None: rolled_segment_ids = roll_and_mask(rolled_segment_ids) + embedding_kwargs = {} if main_token_embeddings is None else {"decoder_input_embeddings": rolled_token_embeddings} target_token_embedding = self.decoder._apply_embedding( shared_embedding, rolled_input_ids, rolled_position_id, deterministic, model_mode=self.decoder.model_mode, + **embedding_kwargs, ) mtp_layer = getattr(self, f"mtp_layer_{k}") @@ -519,8 +608,15 @@ def __call__( model_mode=self.decoder.model_mode, ) + # Collapse the parallel mHC streams before the head. `mtp_hidden_state` itself stays + # in its expanded, unnormalized form so the next MTP depth receives the full state, + # exactly as the main decoder hands it over. + mtp_hidden_state_to_head = mtp_hidden_state + if mtp_layer.is_mhc: + mtp_hidden_state_to_head = mtp_layer.hc_head(mtp_hidden_state) + # Apply separate normalization to the MTP hidden state before projecting to logits. - normed_mtp_hidden_state = mtp_layer.final_norm(mtp_hidden_state) + normed_mtp_hidden_state = mtp_layer.final_norm(mtp_hidden_state_to_head) normed_mtp_hidden_state = sharding.maybe_shard_with_logical( normed_mtp_hidden_state, ("activation_batch", "activation_length", "activation_embed"), @@ -529,8 +625,15 @@ def __call__( sharding.get_logical_axis_rules(), ) + # `reduce_mhc=False` because the collapse already happened above with this layer's own + # `hc_head`; `normalize_y=False` because `final_norm` already ran. mtp_logits = self.decoder.apply_output_head( - shared_embedding, normed_mtp_hidden_state, deterministic, model_mode, normalize_y=False + shared_embedding, + normed_mtp_hidden_state, + deterministic, + model_mode, + normalize_y=False, + reduce_mhc=False, ) logits_logical_axes = ( diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 890bd01179..448c59eb54 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -22,8 +22,8 @@ from typing import Any import warnings -from flax import linen as nn from flax import nnx +from flax.core.spmd import logical_axis_rules import jax from jax.ad_checkpoint import checkpoint_name import jax.numpy as jnp @@ -40,7 +40,7 @@ ) from maxtext.configs.types import check_forced_routing_support from maxtext.layers import linears, mhc, moe, normalizations, quantizations -from maxtext.layers import nnx_scan, nnx_wrappers +from maxtext.layers import nnx_scan from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding from maxtext.layers.normalizations import RMSNorm @@ -69,7 +69,9 @@ qwen3_5, qwen3_custom, simple_layer, + weaver, ) + from maxtext.multimodal import utils as mm_utils from maxtext.utils import max_logging, max_utils, maxtext_utils, maxtext_utils_nnx, sharding from maxtext.utils.sharding import create_sharding @@ -448,7 +450,7 @@ def __init__( num_embeddings=config.trainable_position_size, num_features=config.emb_dim, dtype=config.dtype, - embedding_init=nn.initializers.normal(stddev=1.0), + embedding_init=jax.nn.initializers.normal(stddev=1.0), config=config, mesh=self.mesh, rngs=rngs, @@ -502,7 +504,7 @@ def _init_decoder_layers(self, decoder_block_classes, rngs, mesh): if self.is_gemma4_small: # Gemma4 E2B/E4B: per-layer-index KV-share donor threading and a distinct attention_type - # per layer are not expressible inside nn.scan; pipeline parallelism is also unsupported. + # per layer are not expressible inside jax.lax.scan; pipeline parallelism is also unsupported. if getattr(config, "using_pipeline_parallelism", False) or getattr(config, "scan_layers", False): raise ValueError("gemma4_small (Gemma4 E2B/E4B) does not support pipeline parallelism or scan_layers.") self._init_gemma4_small_layers(rngs) @@ -814,7 +816,14 @@ def _init_scanned_generic(self, decoder_block_classes, rngs): """Initializes scanned generic decoder layers.""" config = self.config layer_cls = decoder_block_classes[0] - num_layers = int(config.num_decoder_layers / config.inhomogeneous_layer_cycle_interval) + cycle_interval = config.inhomogeneous_layer_cycle_interval + if config.num_decoder_layers % cycle_interval: + raise ValueError( + f"num_decoder_layers ({config.num_decoder_layers}) must be a multiple of" + f" inhomogeneous_layer_cycle_interval ({cycle_interval}) to scan {config.decoder_block.value} layers:" + " each scan step builds one whole cycle, so the remaining layers would be dropped." + ) + num_layers = config.num_decoder_layers // cycle_interval layer_kwargs = {} if config.decoder_block == DecoderBlockType.LLAMA4: layer_kwargs = { @@ -949,25 +958,15 @@ def _create_and_register_layer(self, layer_cls, rngs, base_name, i, **layer_kwar setattr(self, attr_name, layer) def _create_single_layer(self, decoder_layer_class, rngs, **kwargs): - """Helper to create a single layer (Linen or NNX).""" - if issubclass(decoder_layer_class, nnx.Module): - return decoder_layer_class( - config=self.config, - mesh=self.mesh, - quant=self.quant, - model_mode=self.model_mode, - rngs=rngs, - **kwargs, - ) - else: - layer_linen = decoder_layer_class( - config=self.config, - mesh=self.mesh, - quant=self.quant, - model_mode=self.model_mode, - **kwargs, - ) - return nnx_wrappers.ToNNX(layer_linen, rngs=rngs) + """Helper to create a single layer.""" + return decoder_layer_class( + config=self.config, + mesh=self.mesh, + quant=self.quant, + model_mode=self.model_mode, + rngs=rngs, + **kwargs, + ) def _create_scanned_layers( self, @@ -1071,6 +1070,14 @@ def _apply_layers_sequentially( scan_axis = self.config.param_scan_axis if scan_axis != 0: params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0) if x.ndim > scan_axis else x, params) + gathered_grad_shardings = {} + if ( + self.config.gather_scanned_kernels_once + and kv_caches_stacked is None + and not self.config.parameter_memory_host_offload + ): + params = moe.gather_scanned_kernels_once(layers, params, self.mesh, self.config.shard_mode) + gathered_grad_shardings = moe.gathered_kernel_grad_shardings(layers, self.mesh, self.config.logical_axis_rules) layer_cls = layers.__class__ sig = inspect.signature(layer_cls.__call__) @@ -1105,6 +1112,9 @@ def layer_fn(carry, scanned_vars): kv_cache_layer = None forced_routed_experts_layer = None + if gathered_grad_shardings: + current_params = moe.reduce_scatter_gathered_grads(current_params, gathered_grad_shardings) + if self.config.parameter_memory_host_offload: current_params = jax.tree.map( lambda x: jax.device_put(x, max_utils.device_space()), @@ -1242,6 +1252,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: @@ -1249,6 +1260,8 @@ def get_deepseek(): return layer_map[cfg.decoder_block] + # configs/types.py _REMAT_PRESET_MLP_NAMES lists the MLP checkpoint names that minimal_policy and the presets of + # get_remat_policy save (the MoE key validators read it): keep the two in sync. def minimal_policy(self, with_context=False, with_quantization=False): """Helper for creating minimal checkpoint policies.""" names = [ @@ -1453,7 +1466,7 @@ def _apply_embedding( y = ( decoder_input_embeddings if decoder_input_embeddings is not None - else shared_embedding(decoder_input_tokens.astype("int32"), model_mode=model_mode) + else self.embed_tokens(shared_embedding, decoder_input_tokens, model_mode) ) # Precomputed embeddings are complete (including any multimodal replacements), @@ -1489,8 +1502,8 @@ def _apply_embedding( "qwen3.5-397b-a17b", "qwen3.5-397b-a17b-fp8", "maxtext-omni-gemma3-qwen3", - "cosmos3-nano-reasoner", - "cosmos3-super-reasoner", + "weaver-mini", + "weaver-max", }: y = mm_utils.merge_mm_embeddings( text_embeddings=y, @@ -1512,8 +1525,8 @@ def _apply_embedding( "qwen3.5-35b-fp8", "qwen3.5-397b-a17b", "qwen3.5-397b-a17b-fp8", - "cosmos3-nano-reasoner", - "cosmos3-super-reasoner", + "weaver-mini", + "weaver-max", }: y = mm_utils.merge_mm_embeddings( text_embeddings=y, @@ -1546,7 +1559,34 @@ def _apply_embedding( return y - def apply_output_head(self, shared_embedding, y, deterministic, model_mode, normalize_y=True): + def embed_tokens(self, shared_embedding, decoder_input_tokens, model_mode): + """Looks up the token embeddings, before dropout and positional embeddings.""" + return shared_embedding(decoder_input_tokens.astype("int32"), model_mode=model_mode) + + def _reduce_mhc_streams(self, y): + """Collapses a 4D mHC hidden state down to 3D. + + `[batch, length, mhc_expansion_rate, emb_dim]` --> `[batch, length, emb_dim]`. + + This is a no-op for non-mHC models and for states that are already collapsed, so it + is safe to call on any path that requires a 3D state. + + Args: + y: Hidden state tensor, either already 3D or an expanded 4D mHC state. + + Returns: + The hidden state with the parallel-stream axis removed. + """ + cfg = self.config + if cfg.mhc_expansion_rate <= 1 or y.ndim < 4: + return y + if self.is_deepseek4: + return self.hc_head(y) + # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim) + _, mhc_reduce_fn = mhc.get_functions(cfg.mhc_expansion_rate) + return mhc_reduce_fn(y) + + def apply_output_head(self, shared_embedding, y, deterministic, model_mode, normalize_y=True, reduce_mhc=True): """Applies final normalization and projects hidden states to logits. Args: @@ -1558,9 +1598,16 @@ def apply_output_head(self, shared_embedding, y, deterministic, model_mode, norm (decoder_norm) before projecting to logits. Set to False when called from Multi-Token Prediction (MTP), which applies its own dedicated final norm (mtp_k_final_norm) to avoid double normalization. + reduce_mhc: If True (default), collapses the parallel hyper-connection (mHC) + streams of a 4D `[batch, seq, mhc_expansion_rate, emb]` hidden state down to + 3D before normalizing. Set to False when the caller has already reduced the + state itself (MTP applies its own per-depth `hc_head`). """ cfg = self.config + if reduce_mhc: + y = self._reduce_mhc_streams(y) + if normalize_y: if cfg.shard_mode == ShardMode.EXPLICIT: norm_out_sharding = create_sharding( @@ -1588,12 +1635,7 @@ def apply_output_head(self, shared_embedding, y, deterministic, model_mode, norm # [batch, length, emb_dim] -> [batch, length, vocab_size] if cfg.logits_via_embedding: # Use the transpose of embedding matrix for logit transform. - if isinstance(shared_embedding, nnx.Module): - embedding_table = shared_embedding.embedding[...] - else: - embedding_table = shared_embedding.variables["params"]["embedding"] - if isinstance(embedding_table, nn.spmd.LogicallyPartitioned): - embedding_table = embedding_table.unbox() + embedding_table = shared_embedding.embedding[...] attend_dtype = jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype logits = attend_on_embedding(y, embedding_table, attend_dtype, self.config, out_sharding) @@ -1633,6 +1675,7 @@ def _build_linen_params(self, moe_stack: nnx.Module) -> dict: "post_self_attention_layer_norm": state_dict["post_self_attention_layer_norm"], "self_attention": state_dict["self_attention"], "DeepSeekMoeBlock_0": moe_block, + "mlp": state_dict.get("mlp"), } def _find_next_boundary(self, current_idx, end_idx, engram_indices): @@ -1794,9 +1837,8 @@ def __call__( decoder_input_embeddings=decoder_input_embeddings, ) - mhc_reduce = None if hasattr(cfg, "mhc_expansion_rate"): - mhc_expand, mhc_reduce = mhc.get_functions(cfg.mhc_expansion_rate) + mhc_expand, _ = mhc.get_functions(cfg.mhc_expansion_rate) if cfg.mhc_expansion_rate > 1: # (batch, length, emb_dim) --> (batch, length, mhc_expansion_rate, emb_dim) y = mhc_expand(y) @@ -1838,7 +1880,7 @@ def __call__( if self.is_deepseek: # Pre-pipeline: dense layers + outside-pipeline MoE layers under PP-as-DP axis rules. logical_axis_rules_pp_as_dp = sharding.logical_axis_rules_pp_act_as_dp(cfg.logical_axis_rules) - with self.mesh, nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp): + with self.mesh, logical_axis_rules(logical_axis_rules_pp_as_dp): if cfg.scan_layers: if getattr(self, "dense_layers", None) is not None and cfg.first_num_dense_layers > 0: y, self.dense_layers, _ = self._apply_layers_sequentially( @@ -1913,7 +1955,7 @@ def __call__( logical_axis_rules_pp_as_dp = sharding.logical_axis_rules_pp_act_as_dp(cfg.logical_axis_rules) with ( self.mesh, - nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp), + logical_axis_rules(logical_axis_rules_pp_as_dp), ): if cfg.scan_layers and hasattr(self, "layers_outside_pipeline"): remaining = cfg.num_decoder_layers - cfg.pipeline_parallel_layers @@ -2264,17 +2306,21 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in, valid_kwargs): assert isinstance(y, jax.Array) # After the final transformer layer, `y` holds the raw, un-normalized hidden state. - if getattr(cfg, "mhc_expansion_rate", 1) > 1: - if cfg.decoder_block == DecoderBlockType.DEEPSEEK4: - hidden_state = self.hc_head(y) - else: - # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim) - hidden_state = mhc_reduce(y) - else: - hidden_state = y + # For DeepSeek-V4 it is deliberately left in its expanded 4D + # `[batch, length, mhc_expansion_rate, emb_dim]` form: the DeepSeek-V4 MTP block consumes this + # return value directly and needs the full parallel-stream state to stay faithful to the + # reference architecture. Every consumer that instead requires a collapsed 3D state reduces it + # explicitly: `apply_output_head` does so for the logit paths, and the two branches below + # do so for the consumers that bypass `apply_output_head`. Other mHC models have a generic MTP + # block that expects a 3D state, so their streams are collapsed here (a no-op for non-mHC models). + hidden_state = y if self.is_deepseek4 else self._reduce_mhc_streams(y) # When invoking from vLLM with RPA attention, logit computation is deferred to a later stage. if cfg.attention in ("vllm_rpa", "vllm_batched_rpa", "vllm_batched_rpa_long_ctx"): + # The vLLM adapter immediately flattens this to `[tokens, emb]` and only later feeds it + # back through `apply_output_head`, by which point the stream axis is indistinguishable + # from the token axis. So the collapse has to happen here, before we hand the state out. + hidden_state = self._reduce_mhc_streams(hidden_state) logits = None # When in the Indexer Dense Warm-up stage, skip the expensive output head projection @@ -2290,7 +2336,11 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in, valid_kwargs): # Instead, we keep track on the hidden states, which has smaller size compared to full logits elif cfg.num_vocab_tiling > 1 and model_mode == MODEL_MODE_TRAIN: logits = None - self.sow(nnx.Intermediate, "hidden_states", hidden_state) + # The tiled loss reshapes on exactly three axes and differentiates only the output-head + # params, so the collapse belongs out here rather than inside the tiled region. Note we + # sow a collapsed copy and leave `hidden_state` itself 4D, so MTP still receives the + # full stream state when both features are enabled. + self.sow(nnx.Intermediate, "hidden_states", self._reduce_mhc_streams(hidden_state)) elif skip_lm_head: logits = None diff --git a/src/maxtext/layers/normalizations.py b/src/maxtext/layers/normalizations.py index a408f8ff35..7e8ec3fac6 100644 --- a/src/maxtext/layers/normalizations.py +++ b/src/maxtext/layers/normalizations.py @@ -22,8 +22,7 @@ import jax.numpy as jnp from jax.sharding import NamedSharding from maxtext.common.common_types import Array, DType, ShardMode, is_fp8_dtype -from maxtext.layers import nnx_wrappers -from maxtext.layers.initializers import Initializer, variable_to_logically_partitioned +from maxtext.layers.initializers import Initializer from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils.sharding import truncate_out_sharding @@ -250,36 +249,6 @@ def __call__( return gated_states.astype(self.dtype) -def rms_norm( - num_features: int, - epsilon: float = 1e-6, - dtype: Any = jnp.float32, - weight_dtype: Any = jnp.float32, - shard_mode: ShardMode = ShardMode.AUTO, - kernel_axes: tuple[None | str, ...] = (), - scale_init: Initializer = jax.nn.initializers.ones, - name: None | str = None, - parameter_memory_host_offload: bool = False, - with_scale: bool = True, -): - """Creates a RMSNorm module.""" - module = nnx_wrappers.to_linen( - RMSNorm, - num_features=num_features, - epsilon=epsilon, - dtype=dtype, - weight_dtype=weight_dtype, - shard_mode=shard_mode, - kernel_axes=kernel_axes, - scale_init=scale_init, - parameter_memory_host_offload=parameter_memory_host_offload, - with_scale=with_scale, - name=name, - metadata_fn=variable_to_logically_partitioned, - ) - return module - - def l2norm(x: Array, dim: int = -1, eps: float = 1e-6) -> Array: """L2 normalization function. Normalizes a vector to have a length of 1. @@ -298,11 +267,3 @@ def l2norm(x: Array, dim: int = -1, eps: float = 1e-6) -> Array: safe_norm_sq = jnp.where(is_zero, 1.0, norm_sq) inv_norm = jnp.where(is_zero, 0.0, jax.lax.rsqrt(safe_norm_sq + jnp.float32(eps))) return jnp.asarray(x_f32 * inv_norm, x.dtype) - - -Qwen3NextRMSNormLinen = nnx_wrappers.to_linen_class( - RMSNorm, - base_metadata_fn=variable_to_logically_partitioned, - scale_init=jax.nn.initializers.zeros, - scale_offset=1.0, -) diff --git a/src/maxtext/layers/pipeline.py b/src/maxtext/layers/pipeline.py index bf66fbcce8..2fe2ae1747 100644 --- a/src/maxtext/layers/pipeline.py +++ b/src/maxtext/layers/pipeline.py @@ -28,8 +28,7 @@ from flax.core import lift as flax_lift from flax.core import scope as flax_scope from flax import nnx -from maxtext.layers import initializers -from maxtext.layers.nnx_wrappers import is_linen_initializing, to_linen_class +from maxtext.layers.nnx_wrappers import is_linen_initializing from maxtext.common.common_types import Config, MODEL_MODE_TRAIN, ShardMode from maxtext.utils.sharding import ( @@ -1821,33 +1820,3 @@ def create_nnx_pipeline( config=config, stage_factory=stage_factory, mesh=mesh, remat_policy=remat_policy, rngs=rngs ) return NNXPipeline(config=config, stage_factory=stage_factory, mesh=mesh, remat_policy=remat_policy, rngs=rngs) - - -Pipeline = to_linen_class( - NNXPipeline, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) -CircularPipeline = to_linen_class( - NNXCircularPipeline, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - -def create_pipeline( - config: Config, - layers=None, - mesh: Mesh = None, - remat_policy: Any = None, -) -> nn.Module: - """Returns the ToLinen-wrapped NNX pipeline appropriate for the config. - - For raw NNX pipeline classes (no Linen wrapping), use create_nnx_pipeline() instead. - - Args: - config: Model configuration. - layers: Callable[[nnx.Rngs], nnx.Module] constructing one pipeline stage. - mesh: JAX device mesh for sharding. - remat_policy: Optional rematerialization policy. - """ - cls = CircularPipeline if config.pipeline_fsdp_ag_per_repeat else Pipeline - return cls(config=config, stage_factory=layers, mesh=mesh, remat_policy=remat_policy) diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index b78d6b9b5f..51c4ada7b6 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -40,9 +40,9 @@ from jax.sharding import NamedSharding from flax.linen import fp8_ops -from flax.linen import initializers as flax_initializers import flax.linen as nn from flax import nnx +from flax.nnx import variablelib # Support different packaging structures across environments even within # the same Qwix version identifier (imports from _src.utils vs _src). try: @@ -223,6 +223,7 @@ def native_fp8_dot_general( compute_dtype: DType = jnp.bfloat16, precision: jax.lax.PrecisionLike = None, out_sharding: NamedSharding | None = None, + act_calibration_method: str = "absmax", ) -> jnp.ndarray: """Computes dot_general with pre-quantized FP8 weights without dequantization.""" if len(axis) != 1 or len(contract_ind) != 1: @@ -262,12 +263,15 @@ def native_fp8_dot_general( rhs_block_size = scale_block_size[rhs_axis] if isinstance(scale_block_size, (list, tuple)) else scale_block_size channelwise_axes = [d for d in range(inputs.ndim) if d != lhs_axis] tiled_axes = {lhs_axis: rhs_block_size} + if is_static_calibration(act_calibration_method): + channelwise_axes = [] + tiled_axes = {} lhs = qwix.quantize( jnp.asarray(inputs, compute_dtype), jnp.float8_e4m3fn, channelwise_axes=channelwise_axes, tiled_axes=tiled_axes, - calibration_method="absmax", + calibration_method=act_calibration_method, ) dimension_numbers = (((lhs_axis,), (rhs_axis,)), ((), ())) out = qwix.dot_general(lhs, rhs, dimension_numbers, precision=precision, preferred_element_type=jnp.float32) @@ -279,9 +283,40 @@ def native_fp8_dot_general( @dataclass class ServeFp8WeightQuantization(Quantization): - """Marks a layer as using native FP8 compute without weight dequantization.""" + """Marks a layer as using native FP8 compute without weight dequantization. + + Weights are FP8 with static scales loaded from the checkpoint. Activations (matmul + inputs) are quantized to FP8 according to `act_calibration_method`: + + - "absmax" (default): dynamic scale computed from the activation at runtime. + - "fixed,-A,A" or "fixed,A": static per-tensor scale A / 448, no runtime reduction. + + The setting applies to the dense layers (`native_fp8_dot_general`) and to the routed + experts on the tpu-inference fused MoE path (`static_act_scale`). + """ quant_mode = None + act_calibration_method: str = "absmax" + + def __post_init__(self): + if is_static_calibration(self.act_calibration_method): + # Fail at config time on e.g. an asymmetric range instead of at trace time. + get_static_scale(jnp.float8_e4m3fn, self.act_calibration_method) + + def static_act_scale(self) -> jax.Array | None: + """Returns the static per-tensor activation scale as a float32 [1, 1] array. + + Returns None when activations use dynamic scaling. + """ + if not is_static_calibration(self.act_calibration_method): + return None + scale = get_static_scale(jnp.float8_e4m3fn, self.act_calibration_method) + return jnp.full((1, 1), scale, dtype=jnp.float32) + + +def is_static_calibration(calibration_method: str | None) -> bool: + """Whether `calibration_method` is a fixed-range (static scale) calibration.""" + return bool(calibration_method) and calibration_method.lower().startswith("fixed") def prepare_fused_gmm_scale(scale: jnp.ndarray, kernel_shape: Tuple[int, ...]) -> jnp.ndarray: @@ -477,12 +512,12 @@ def einsum(self, mesh_axes: Tuple[str, ...] = ()): return QwixEinsum(config=self._get_fp8_full_qwix_config()) -class QwixDotGeneral(nn.Module): +class QwixDotGeneral(nnx.Module): """A callable class for Qwix dot_general.""" - config: dot_general_qt.DotGeneralQtConfig + def __init__(self, config: dot_general_qt.DotGeneralQtConfig): + self.config = config - @nn.compact def __call__( self, lhs: jax.Array, @@ -496,12 +531,12 @@ def __call__( return dot_general_qt.dot_general_qt(lhs, rhs, dimension_numbers, self.config) -class QwixEinsum(nn.Module): +class QwixEinsum(nnx.Module): """A callable class for Qwix einsum.""" - config: dot_general_qt.DotGeneralQtConfig + def __init__(self, config: dot_general_qt.DotGeneralQtConfig): + self.config = config - @nn.compact def __call__( self, einsum_str: str, @@ -539,92 +574,138 @@ class Fp8Quantization(Quantization): def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): """Returns dot_general configured with aqt params.""" - return nn.Fp8DirectDotGeneralOp + return Fp8DirectDotGeneralOp def einsum(self, dtype: DType = jnp.float32): return _Fp8EinsumWrapper(dtype=dtype) -class _Fp8EinsumWrapper(nn.Module): - """Wrapper for nn.Fp8Einsum to handle computation dtype.""" +_OVERWRITE_WITH_GRADIENT = variablelib.variable_type_from_name("_overwrite_with_gradient", allow_register=True) - dtype: DType - @nn.compact - def __call__(self, eqn, lhs, rhs, **kwargs): - # nn.Fp8Einsum determines compute dtype from rhs. - # We cast rhs to the desired computation dtype. - # nn.Fp8Einsum will then cast lhs to the same dtype. - rhs = rhs.astype(self.dtype) - return nn.Fp8Einsum(name="fp8_einsum")(eqn, lhs, rhs, **kwargs) +def _fp8_dot_inputs(*args, **kwargs): + """Splits dot_general arguments the way Flax's fp8 ops do: the kernel's dtype is the computation dtype.""" + del kwargs + assert len(args) == 3 + x, k, dimension_numbers = args + return jnp.asarray(x, k.dtype), k, dimension_numbers, k.dtype -class Fp8Einsum(nn.Module): - """An fp8 einsum op.""" +class _Fp8Stats(nnx.Module): + """The scales and amax histories of an fp8 op, kept the way Flax's `Fp8DotGeneralBase` keeps them. - #: size of the amax history. - amax_history_length: int = 1024 - #: e4m3 variants, e.g., e4m3fn, e4m3fnuz. - e4m3_dtype: DType = jnp.float8_e4m3fn - #: e5m2 variants, e.g., e5m2, e5m2fnuz. - e5m2_dtype: DType = jnp.float8_e5m2 - #: computation dtype. - dtype: DType = jnp.float32 - - def setup(self) -> None: - """init with input_amax_history, kernel_amax_history, output_grad_amax_history, - input_scale, kernel_scale, output_grad_scale""" - scale_args = ( - flax_initializers.ones_init(), - jax.random.PRNGKey(0), - (1,), - jnp.float32, - ) - amax_history_args = ( - flax_initializers.zeros_init(), - jax.random.PRNGKey(0), - (self.amax_history_length,), - jnp.float32, + They are `_overwrite_with_gradient` variables: the fp8 ops return their new values as gradients, + and the train step writes those back instead of applying the optimizer. + """ + + def __init__(self, amax_history_length: int = 1024): + self.input_amax_history = _OVERWRITE_WITH_GRADIENT(jnp.zeros((amax_history_length,), jnp.float32)) + self.kernel_amax_history = _OVERWRITE_WITH_GRADIENT(jnp.zeros((amax_history_length,), jnp.float32)) + self.output_grad_amax_history = _OVERWRITE_WITH_GRADIENT(jnp.zeros((amax_history_length,), jnp.float32)) + self.input_scale = _OVERWRITE_WITH_GRADIENT(jnp.ones((1,), jnp.float32)) + self.kernel_scale = _OVERWRITE_WITH_GRADIENT(jnp.ones((1,), jnp.float32)) + self.output_grad_scale = _OVERWRITE_WITH_GRADIENT(jnp.ones((1,), jnp.float32)) + + def scaled_dot_general(self, compute_dtype: DType): + """Flax's `fp8_scaled_dot_general` bound to these scales and amax histories.""" + return functools.partial( + fp8_ops.fp8_scaled_dot_general, + lhs_scale=self.input_scale[...], + rhs_scale=self.kernel_scale[...], + grad_scale=self.output_grad_scale[...], + lhs_amax_history=self.input_amax_history[...], + rhs_amax_history=self.kernel_amax_history[...], + grad_amax_history=self.output_grad_amax_history[...], + quantize_compute_type=compute_dtype, ) - OVERWRITE_WITH_GRADIENT = "_overwrite_with_gradient" - self.input_amax_history = self.variable(OVERWRITE_WITH_GRADIENT, "input_amax_history", *amax_history_args) - self.kernel_amax_history = self.variable(OVERWRITE_WITH_GRADIENT, "kernel_amax_history", *amax_history_args) - self.output_grad_amax_history = self.variable(OVERWRITE_WITH_GRADIENT, "output_grad_amax_history", *amax_history_args) + def quantize_inputs(self, compute_dtype: DType, e4m3_dtype: DType, x, k): + x_qdq = fp8_ops.in_qdq(compute_dtype, e4m3_dtype, x, self.input_scale[...], self.input_amax_history[...]) + k_qdq = fp8_ops.in_qdq(compute_dtype, e4m3_dtype, k, self.kernel_scale[...], self.kernel_amax_history[...]) + return x_qdq, k_qdq - self.input_scale = self.variable(OVERWRITE_WITH_GRADIENT, "input_scale", *scale_args) - self.kernel_scale = self.variable(OVERWRITE_WITH_GRADIENT, "kernel_scale", *scale_args) - self.output_grad_scale = self.variable(OVERWRITE_WITH_GRADIENT, "output_grad_scale", *scale_args) + def quantize_output_grad(self, compute_dtype: DType, e5m2_dtype: DType, y): + return fp8_ops.out_qdq(compute_dtype, e5m2_dtype, y, self.output_grad_scale[...], self.output_grad_amax_history[...]) - def __call__(self, eqn, *args, **kwargs): - assert len(args) == 2 - x = args[0] - k = args[1] - comp_dtype = self.dtype - k = jnp.asarray(k, comp_dtype) - x = jnp.asarray(x, comp_dtype) +class Fp8DirectDotGeneralOp(_Fp8Stats): + """Flax's `Fp8DirectDotGeneralOp` as an NNX module. + + It keeps Flax's class name: layers store it as "_0", and checkpoints use that path. + """ - x_qdq = fp8_ops.in_qdq(comp_dtype, self.e4m3_dtype, x, self.input_scale.value, self.input_amax_history.value) - k_qdq = fp8_ops.in_qdq(comp_dtype, self.e4m3_dtype, k, self.kernel_scale.value, self.kernel_amax_history.value) + def __call__(self, *args, **kwargs): + x, k, dimension_numbers, compute_dtype = _fp8_dot_inputs(*args, **kwargs) + return self.scaled_dot_general(compute_dtype)(x, k, dimension_numbers, precision=None, preferred_element_type=x.dtype) + + +class NANOOFp8DotGeneralOp(_Fp8Stats): + """Flax's `NANOOFp8DotGeneralOp` as an NNX module, using the fnuz fp8 formats of AMD MI300/MI325.""" + + def __call__(self, *args, **kwargs): + x, k, dimension_numbers, compute_dtype = _fp8_dot_inputs(*args, **kwargs) + x_qdq, k_qdq = self.quantize_inputs(compute_dtype, jnp.float8_e4m3fnuz, x, k) + y_qdq = fp8_ops.dot_general_with_precision(x_qdq, k_qdq, dimension_numbers) + return self.quantize_output_grad(compute_dtype, jnp.float8_e5m2fnuz, y_qdq) - y_qdq = jnp.einsum(eqn, x_qdq, k_qdq, _dot_general=fp8_ops.dot_general_with_precision) - y = fp8_ops.out_qdq( - comp_dtype, - self.e5m2_dtype, - y_qdq, - self.output_grad_scale.value, - self.output_grad_amax_history.value, +class Fp8DirectEinsum(_Fp8Stats): + """Flax's `Fp8Einsum` as an NNX module. The rhs is the weight and sets the computation dtype.""" + + def __call__(self, eqn, lhs, rhs, precision=None, preferred_element_type=None): + compute_dtype = rhs.dtype + return jnp.einsum( + eqn, + lhs.astype(compute_dtype), + rhs, + precision=precision, + preferred_element_type=preferred_element_type, + _dot_general=self.scaled_dot_general(compute_dtype), ) - return y + + +class _Fp8EinsumWrapper(nnx.Module): + """Runs `Fp8DirectEinsum` in a given computation dtype.""" + + def __init__(self, dtype: DType): + self.dtype = dtype + self.fp8_einsum = Fp8DirectEinsum() + + def __call__(self, eqn, lhs, rhs, **kwargs): + # The fp8 einsum takes its computation dtype from rhs and casts lhs to it. + return self.fp8_einsum(eqn, lhs, rhs.astype(self.dtype), **kwargs) + + +class Fp8Einsum(_Fp8Stats): + """An fp8 einsum op.""" + + def __init__( + self, + amax_history_length: int = 1024, + e4m3_dtype: DType = jnp.float8_e4m3fn, + e5m2_dtype: DType = jnp.float8_e5m2, + dtype: DType = jnp.float32, + ): + super().__init__(amax_history_length) + self.e4m3_dtype = e4m3_dtype + self.e5m2_dtype = e5m2_dtype + self.dtype = dtype + + def __call__(self, eqn, *args, **kwargs): + assert len(args) == 2 + comp_dtype = self.dtype + x = jnp.asarray(args[0], comp_dtype) + k = jnp.asarray(args[1], comp_dtype) + x_qdq, k_qdq = self.quantize_inputs(comp_dtype, self.e4m3_dtype, x, k) + y_qdq = jnp.einsum(eqn, x_qdq, k_qdq, _dot_general=fp8_ops.dot_general_with_precision) + return self.quantize_output_grad(comp_dtype, self.e5m2_dtype, y_qdq) @dataclass class NANOOFp8Quantization(Quantization): """Configures NANOO Fp8 quantization for AMD MI300/MI325 GPUs""" - # Same as Fp8Quantization: nn.NANOOFp8DotGeneralOp never draws at apply time. + # Same as Fp8Quantization: the NANOO fp8 op never draws at apply time. needs_apply_rngs: ClassVar[bool] = False quant_mode = "train" @@ -632,7 +713,7 @@ class NANOOFp8Quantization(Quantization): def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): """Returns dot_general configured with aqt params.""" - return nn.NANOOFp8DotGeneralOp + return NANOOFp8DotGeneralOp def einsum(self, dtype: DType = jnp.float32): """Returns an einsum using the NANOO (fnuz) fp8 formats of AMD MI300/MI325.""" @@ -902,7 +983,7 @@ def get_quant_mode(quant_mode_str: str = "train"): def configure_quantization(config: Config, quant_mode_str: str = "train"): """Configure quantization based on user config and quant mode.""" if config.quantization == "serve_fp8_weight": - return ServeFp8WeightQuantization() + return ServeFp8WeightQuantization(act_calibration_method=config.act_quantization_calibration_method) if getattr(config, "use_batch_split_schedule", False) and config.quantization: # The older version of batch-split that fully uses qwix quantization. @@ -982,52 +1063,56 @@ def configure_kv_quant(config): return None if not config.quantize_kvcache else KVQuant(config) -def _apply_linen_module_in_nnx(linen_module_cls, op_id, *args, **kwargs): - """Applies a Linen module within an NNX context.""" +def _apply_fp8_op_in_nnx(op_cls, op_id, *args, **kwargs): + """Applies an fp8 op for a Qwix provider, keeping its state on the NNX module Qwix intercepts.""" try: parent = flax_util.get_current_module() - is_nnx = isinstance(parent, nnx.Module) except ValueError: - is_nnx = False - - if is_nnx: - attr_name = f"_qwix_fp8_gpu_{op_id}" - if not hasattr(parent, attr_name): # pyrefly: ignore[unbound-name] - rngs = getattr(parent, "qwix_rngs", None) - if rngs is None: - parent_rngs = getattr(parent, "rngs", None) - if parent_rngs is not None and hasattr(parent_rngs, "fork"): - rngs = parent_rngs.fork() - else: - rngs = nnx.Rngs(0) - wrapper = nnx_wrappers.ToNNX(linen_module_cls(name=op_id), rngs=rngs) - wrapper.lazy_init(*args, **kwargs) - setattr(parent, attr_name, wrapper) - return getattr(parent, attr_name)(*args, mutable=["_overwrite_with_gradient"], **kwargs) - else: - return linen_module_cls(name=op_id)(*args, **kwargs) - - -def create_fp8_einsum(quant: Quantization, dtype: DType, rngs: nnx.Rngs) -> nnx_wrappers.ToNNX: - """Creates an fp8 einsum that an `nnx.Module` can call. - - An fp8 einsum holds its scaling factors and amax histories in Linen variables, which can - only be created while the module is bound to a Linen scope. The bridge into NNX therefore - has to happen while the parent module is being built; creating the state on the first call - instead would grow the module graph inside the scanned layer loop, which NNX rejects. - - The state has a fixed shape, so a canonical pair of operands is enough to materialize it. - The returned einsum still accepts operands of any shape. + parent = None + if not isinstance(parent, nnx.Module): + raise ValueError(f"The fp8 op {op_id} needs an NNX module to hold its scales; got {type(parent).__name__}.") + attr_name = f"_qwix_fp8_gpu_{op_id}" + if not hasattr(parent, attr_name): + # Fork the parent's rngs as the bridge did, so its later draws (e.g. dropout) stay the same. + rngs = getattr(parent, "qwix_rngs", None) + if rngs is None: + rngs = getattr(parent, "rngs", None) + if hasattr(rngs, "fork"): + rngs.fork() + setattr(parent, attr_name, op_cls()) + return getattr(parent, attr_name)(*args, **kwargs) + + +def create_fp8_einsum(quant: Quantization, dtype: DType, rngs: nnx.Rngs) -> nnx.Module: + """Creates the fp8 einsum an MoE layer registers in `quant_einsums`. + + Its scales and amax histories live in the einsum, so it is built with the layer rather than + on the first call, which would grow the module graph inside the scanned layer loop. Args: quant: The fp8 quantization providing the einsum. dtype: The computation dtype of the einsum. - rngs: The `nnx.Rngs` of the parent module. + rngs: The `nnx.Rngs` of the parent module, forked as in `make_dot_general`. """ - wrapper = nnx_wrappers.ToNNX(quant.einsum(dtype=dtype), rngs=rngs) # pytype: disable=attribute-error - dummy_operand = jnp.zeros((1, 1), dtype=dtype) - wrapper.lazy_init("ab,bc->ac", dummy_operand, dummy_operand) - return wrapper + rngs.fork() + return quant.einsum(dtype=dtype) # pytype: disable=attribute-error + + +def make_dot_general(quant, kernel_axes: Tuple[str, ...], rngs: nnx.Rngs) -> tuple[str, nnx.Module]: + """Builds the quantized dot_general of an NNX layer and the attribute name the layer stores it under. + + The name follows the quantizer's class, e.g. "Fp8DirectDotGeneralOp_0", and checkpoints use it. + Linen quantizers (AQT, Transformer Engine) are bridged into NNX and create their state on the + first call; the others are NNX modules and are used as they are. + """ + dot_general = quant.dot_general_cls(mesh_axes=kernel_axes)() + name = f"{type(dot_general).__name__}_0" + if isinstance(dot_general, nn.Module): + dot_general = nnx_wrappers.ToNNX(dot_general, rngs=rngs) + else: + # The bridge forks rngs; do the same so the weights created after this one draw the same values. + rngs.fork() + return name, dot_general def apply_einsum_in_nnx(parent: nnx.Module, op_id: str, einsum, mutable: Sequence[str], *args): @@ -1062,20 +1147,20 @@ def apply_einsum_in_nnx(parent: nnx.Module, op_id: str, einsum, mutable: Sequenc class NvidaFp8Provider(qwix.QtProvider): - """Wraps nn.Fp8DirectDotGeneralOp with Qwix's provider interface.""" + """Runs the fp8 ops of NVIDIA GPUs through Qwix's provider interface.""" def dot_general(self, *args, **kwargs): # Here we only check if the rule is None or not. rule, op_id = self._get_current_rule_and_op_id("dot_general") if rule is None: return jax.lax.dot_general(*args, **kwargs) - return _apply_linen_module_in_nnx(nn.Fp8DirectDotGeneralOp, op_id, *args, **kwargs) + return _apply_fp8_op_in_nnx(Fp8DirectDotGeneralOp, op_id, *args, **kwargs) def einsum(self, *args, **kwargs): rule, op_id = self._get_current_rule_and_op_id("einsum") if rule is None: return jnp.einsum(*args, **kwargs) - return _apply_linen_module_in_nnx(nn.Fp8Einsum, op_id, *args, **kwargs) + return _apply_fp8_op_in_nnx(Fp8DirectEinsum, op_id, *args, **kwargs) class NANOOFp8Provider(qwix.QtProvider): @@ -1085,7 +1170,7 @@ def dot_general(self, *args, **kwargs): rule, op_id = self._get_current_rule_and_op_id("dot_general") if rule is None: return jax.lax.dot_general(*args, **kwargs) - return _apply_linen_module_in_nnx(nn.NANOOFp8DotGeneralOp, op_id, *args, **kwargs) + return _apply_fp8_op_in_nnx(NANOOFp8DotGeneralOp, op_id, *args, **kwargs) def _get_router_proj_unquantized_rule() -> qwix.QtRule: @@ -1106,6 +1191,12 @@ def _get_router_proj_unquantized_rule() -> qwix.QtRule: ) +def _drhs_grad_calibration_override(config: Config) -> dict: + """Qwix DotGeneralQtConfig override for the weight-gradient arm's cotangent calibration (see types.py).""" + m = getattr(config, "drhs_grad_quantization_calibration_method", None) + return {"drhs_grad_calibration_method": m} if m else {} + + def get_fp8_full_qwix_rule_w_sparsity(config: Config): """Returns Qwix quantization rules for fp8_full with optional weight sparsity.""" sparsity_rule = None @@ -1117,15 +1208,54 @@ def get_fp8_full_qwix_rule_w_sparsity(config: Config): weight_sparsity_start_step=config.weight_sparsity_start_step, ) - if config.quantize_mtp: - module_path = "(decoder/.*layers.*|mtp_block/.*)" - else: - module_path = "decoder/.*layers.*" - rules = [] if not config.quantize_router_proj: rules.append(_get_router_proj_unquantized_rule()) + if config.quantize_logits_proj: + logits_calib = config.logits_proj_quant_calibration_method or None + logits_bwd_calib = config.logits_proj_bwd_quant_calibration_method or None + # A logits-specific bwd calibration applies to both gradient arms and takes precedence over the + # global drhs_grad_quantization_calibration_method. + if logits_bwd_calib: + logits_drhs_override = {"drhs_grad_calibration_method": logits_bwd_calib} + else: + logits_drhs_override = _drhs_grad_calibration_override(config) + rules.append( + qwix.QtRule( + module_path="decoder/logits_dense.*", + weight_qtype=jnp.float8_e4m3fn, + act_qtype=jnp.float8_e4m3fn, + bwd_qtype=jnp.float8_e5m2, + weight_calibration_method=logits_calib or config.weight_quantization_calibration_method, + act_calibration_method=logits_calib or config.act_quantization_calibration_method, + bwd_calibration_method=logits_bwd_calib or config.bwd_quantization_calibration_method, + additional_qt_config=logits_drhs_override or None, + op_names=("dot_general",), + ) + ) + + paths = ["decoder/.*layers.*"] # Main transformer layers + if config.quantize_mtp: + paths.append("mtp_block/.*") # Multi-token prediction block + # Disjunct regex paths, e.g. "(path1|path2|...)" + module_path = f"({'|'.join(paths)})" if len(paths) > 1 else paths[0] + + if config.sa_qproj_bwd_bf16: + rules.append( + qwix.QtRule( + module_path=(rf"{module_path}/(self_)?attention\w*/(query|wq_a|wq_b)"), + weight_qtype=jnp.float8_e4m3fn, + act_qtype=jnp.float8_e4m3fn, + bwd_qtype=None, + weight_calibration_method=config.weight_quantization_calibration_method, + act_calibration_method=config.act_quantization_calibration_method, + bwd_calibration_method=config.bwd_quantization_calibration_method, + additional_qt_config={"sparsity_rule": sparsity_rule}, + op_names=("dot_general",), + ) + ) + rules.append( qwix.QtRule( module_path=module_path, @@ -1135,7 +1265,7 @@ def get_fp8_full_qwix_rule_w_sparsity(config: Config): weight_calibration_method=config.weight_quantization_calibration_method, act_calibration_method=config.act_quantization_calibration_method, bwd_calibration_method=config.bwd_quantization_calibration_method, - additional_qt_config={"sparsity_rule": sparsity_rule}, + additional_qt_config={"sparsity_rule": sparsity_rule, **_drhs_grad_calibration_override(config)}, op_names=("dot_general", "gmm", "ragged_dot"), ) ) @@ -1155,7 +1285,9 @@ def make_qt_rule(dtype) -> list[qwix.QtRule]: weight_qtype=dtype, act_qtype=dtype, bwd_qtype=dtype, - bwd_weight_grad_tile_size=1 / config.quantization_local_shard_count, + bwd_weight_grad_tile_size=( + 1 / config.qwix_bwd_weight_grad_tile_count if config.qwix_bwd_weight_grad_tile_count > 1 else None + ), disable_channelwise_axes=False, op_names=("dot_general",), ) diff --git a/src/maxtext/m3/README.md b/src/maxtext/m3/README.md new file mode 100644 index 0000000000..05e30a2030 --- /dev/null +++ b/src/maxtext/m3/README.md @@ -0,0 +1,79 @@ + + +# m3 — Modern MaxText Models + +m3 organizes models as a model garden: one self-contained directory per model family. A family +defines its own attention, MLP and decoder blocks instead of routing through a shared attention +class and a central `Decoder`, so a model can be read top to bottom in one place. Scalability +concerns are applied as post-construction transforms rather than branched on inside model code. + +> **Status: scaffolding.** Opt in with `use_m3_model=true`. `qwen3-0.6b` has been tested +> for correctness here, but performance features are not yet available. +> Legacy `src/maxtext/layers/` remains the production path +> until the migration milestones complete. + +## Layout + +```text +m3/ + core/ # RoPE, sharding, interfaces, checkpoint, data — only what NNX lacks + infra/ # remat, scan, quantization, offload — applied after construction + configs/ # slim Pydantic config types + base.yml + models/ # one directory per model family + create_model dispatch + train/ # m3's own training loop (one of two supported drivers) +``` + +There is no `layers/` directory: NNX *is* the layer library. + +## Architecture rules + +1. **Use NNX built-ins directly.** Compose `nnx.Linear`, `nnx.Einsum`, `nnx.Embed`, `nnx.RMSNorm`, + `nnx.Dropout` and `nnx.List` as-is — no wrapper layer classes, no Linen, no bridges. +2. **Models own their architecture.** Each model file defines its own attention, MLP and decoder + block. No shared attention class parameterized for every variant, and no model-name branching. + Some duplication across model files is cheaper than a 60-parameter shared layer. +3. **Sharding is an annotation.** Models receive a `Sharding` instance at construction and refer to + logical axis names only (`"embed"`, `"heads"`, `"mlp"`). Model code never names physical mesh + axes or devices. +4. **Infrastructure stays out of model code.** No `maybe_remat`, no `dot_general` injection, no + `model_mode` conditionals. Only zero-cost annotations belong in `__call__`: sharding specs and + `checkpoint_name` tags. +5. **Config-driven.** A typed Pydantic config selects architecture and hyperparameters; per-model + YAML files override a base. +6. **Readability is a feature.** A model file should read top to bottom without jumping through a + shared decoder. + +## Invariants for the training-engine and RL path + +m3 models must be drivable both by `m3/train/loop.py` and by the existing +`maxtext.training_engine.maxtext_engine` (which the RL stack targets). These invariants are load +bearing for that second driver — breaking them tends to fail silently: + +| Invariant | Why | +| :---- | :---- | +| `nnx.Param` is the only trainable variable type | The engine differentiates `nnx.Param` only; other variable subclasses silently receive zero gradient | +| Do not `sow` intermediates on the engine path | Widening non-parameter state changes the state treedef across the forward pass and permanently drops the engine onto its slow path | +| Remat and scan wrappers must be path-transparent | Parameter paths feed checkpointing, HF mappings, weight sync, optimizer masks and unscan; an extra path segment breaks them, in several cases without an error | +| Auxiliary losses fold into the primary loss | Only the primary loss is differentiated; a router loss reported alongside it contributes no gradient | +| One `Mesh` object shared with the driver | Sharding whose mesh differs from the engine's is silently replaced by full replication | +| No parameter sharding constraints after init | Parameters are sharded once by the initializer; re-asserting collides with the engine's deferred all-reduce and Zero-1 resharding | + +## Tests + +Unit tests live in `tests/unit/m3/`, integration tests in `tests/integration/m3/`. They are picked +up by the standard MaxText pytest markers and run on every PR — no workflow changes are needed to +add a test file. diff --git a/src/maxtext/m3/__init__.py b/src/maxtext/m3/__init__.py new file mode 100644 index 0000000000..e27184e773 --- /dev/null +++ b/src/maxtext/m3/__init__.py @@ -0,0 +1,42 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""m3: Modern MaxText Models. + +m3 is the pure-NNX, single-file-per-model rewrite of the MaxText modeling stack. +Models own their own attention, MLP and decoder blocks by composing NNX +primitives; scalability concerns (remat, scan, quantization, host offload) are +applied post-construction as transforms rather than being branched on inside +model code. + +Subpackages: + * `core` — primitives m3 must write itself (RoPE, sharding, checkpoint, data). + * `infra` — post-construction transforms (remat, scan, quantization, offload). + * `configs` — slim Pydantic config types plus the YAML defaults. + * `models` — one self-contained file per model family, plus its sharding rules. + * `train` — m3's own training loop (one of two supported drivers; the other is + `maxtext.training_engine.maxtext_engine`). + +This package is under active development and is not yet wired into any MaxText +entry point. See `README.md` in this directory for the architecture rules that +contributions must follow. +""" + +__all__ = [ + "configs", + "core", + "infra", + "models", + "train", +] diff --git a/src/maxtext/m3/configs/__init__.py b/src/maxtext/m3/configs/__init__.py new file mode 100644 index 0000000000..68061e5eb1 --- /dev/null +++ b/src/maxtext/m3/configs/__init__.py @@ -0,0 +1,26 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""m3 configuration: slim Pydantic types plus YAML defaults. + +Planned contents: + * `config.py` — `ModelConfig`, `TransformConfig`, `RootConfig`, `load_config`. + * `base.yml` — base defaults that model YAMLs override. + +m3 configs are intentionally small and additive: fields are added as +architectures require them, not preemptively. Interop with the legacy +`maxtext.configs.pyconfig.HyperParameters` surface (required by the training +engine and the RL path) is handled by a separate bridge, not by widening these +types. +""" diff --git a/src/maxtext/m3/configs/config.py b/src/maxtext/m3/configs/config.py new file mode 100644 index 0000000000..000db1a0b5 --- /dev/null +++ b/src/maxtext/m3/configs/config.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Unified Modern MaxText Models (M3) configuration and initialization.""" diff --git a/src/maxtext/m3/configs/diffusers_config.py b/src/maxtext/m3/configs/diffusers_config.py new file mode 100644 index 0000000000..753406aff3 --- /dev/null +++ b/src/maxtext/m3/configs/diffusers_config.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Diffusion model configuration schema and hyperparameter definitions.""" diff --git a/src/maxtext/m3/core/__init__.py b/src/maxtext/m3/core/__init__.py new file mode 100644 index 0000000000..a05761c55c --- /dev/null +++ b/src/maxtext/m3/core/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Core m3 primitives: the pieces NNX does not already provide. + +Modules: + * `rope.py` — rotary positional embedding (base implementation). + * `sharding.py` — `Sharding` ABC, logical-axis specs, mesh creation. + +Planned modules: + * `interfaces.py` — shared ABCs and `CausalLMOutput`. + * `checkpoint.py` — Orbax wrapper. + * `data.py` — data loading helpers on top of the MaxText input pipeline. + +Leaf layers (`nnx.Linear`, `nnx.Einsum`, `nnx.Embed`, `nnx.RMSNorm`, ...) are +used directly from NNX and are deliberately not wrapped here. +""" diff --git a/src/maxtext/m3/core/rope.py b/src/maxtext/m3/core/rope.py new file mode 100644 index 0000000000..30a5788b6e --- /dev/null +++ b/src/maxtext/m3/core/rope.py @@ -0,0 +1,106 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Rotary positional embedding (RoPE). + +This is one "layer" m3 writes itself, because NNX has no RoPE primitive. It +is a pure function rather than a module: it holds no state, allocates no +parameters, and so has nothing to place on a mesh. + +Only the base ("default") variant lives here. Models needing a specialized +variant -- YaRN, LLaMA 3.1 extended context, partial-rotary, mRoPE -- implement +it in their own model directory rather than growing a parameter on this +function. + +Convention +---------- +This is the half-split ("rotate half") formulation used by Hugging Face and by +the legacy MaxText `RotaryEmbedding`, *not* the interleaved (GPT-J) pairing. +The first and second halves of `head_dim` form the rotation pairs, so dimension +`i` is paired with dimension `i + head_dim // 2`. + +Numerics +-------- +Angles and their sinusoids are computed in float32 regardless of the input +dtype, then cast down to the input dtype before the rotation itself. This +mirrors legacy MaxText exactly, which keeps bf16 training runs bit-comparable +across the migration. Computing the rotation in float32 and casting the result +instead would be marginally more accurate, but would silently diverge from the +numbers the legacy stack produces. +""" + +import jax +import jax.numpy as jnp + + +def apply_rope( + x: jax.Array, + positions: jax.Array, + *, + max_timescale: float = 10_000.0, + min_timescale: float = 1.0, +) -> jax.Array: + """Applies rotary positional embedding to a projected query or key tensor. + + Args: + x: Projected queries or keys, shape `(batch, length, heads, head_dim)`. + `head_dim` must be even. Applied after the q/k projection and before + attention. + positions: Integer token positions, shape `(batch, length)`. Supplying + explicit positions (rather than assuming `arange(length)`) is what makes + packed sequences work: each packed segment restarts at zero. + max_timescale: Longest rotation period, commonly called `rope_theta`. + 10_000 is the original RoPE value; long-context models raise it (Qwen3 + uses 1_000_000). + min_timescale: Shortest rotation period. Effectively always 1. + + Returns: + An array of the same shape and dtype as `x`, rotated by position. + + Raises: + ValueError: If `x` is not rank 4, if `head_dim` is odd, or if `positions` + does not match the leading `(batch, length)` dimensions of `x`. + """ + if x.ndim != 4: + raise ValueError(f"apply_rope expects x of shape (batch, length, heads, head_dim); got shape {x.shape}.") + + head_dim = x.shape[-1] + if head_dim % 2: + raise ValueError(f"apply_rope requires an even head_dim to form rotation pairs; got {head_dim}.") + + if positions.shape != x.shape[:2]: + raise ValueError( + f"apply_rope expects positions of shape (batch, length) = {x.shape[:2]}; got shape {positions.shape}." + ) + + half_dim = head_dim // 2 + + # Geometric series of rotation periods, one per dimension pair. The integer + # `arange` is deliberate: it matches the legacy `RotaryEmbedding.timescale` + # expression, which is what keeps the two bit-exact. + fraction = 2 * jnp.arange(half_dim) / head_dim + timescale = min_timescale * (max_timescale / min_timescale) ** fraction + + # (batch, length) -> (batch, length, 1, half_dim), broadcasting over heads. + angles = positions[:, :, jnp.newaxis, jnp.newaxis].astype(jnp.float32) / timescale + sin = jnp.sin(angles).astype(x.dtype) + cos = jnp.cos(angles).astype(x.dtype) + + # Rotate each (x1[i], x2[i]) pair by its angle, pairing dimension i with + # dimension i + half_dim. Equivalent to the legacy + # `x * cos + rotate_half(x) * sin`, without materializing the doubled + # sin/cos tables. + x1 = x[..., :half_dim] + x2 = x[..., half_dim:] + return jnp.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1) diff --git a/src/maxtext/m3/core/sharding.py b/src/maxtext/m3/core/sharding.py new file mode 100644 index 0000000000..20b2c38cb3 --- /dev/null +++ b/src/maxtext/m3/core/sharding.py @@ -0,0 +1,367 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sharding rules and mesh creation. + +Model code never names a physical mesh axis. It refers to logical axes -- +`"embed"`, `"heads"`, `"mlp"` -- and a `Sharding` instance supplied at +construction translates those into `PartitionSpec`s. Swapping a parallelism +strategy is then a change to one `Sharding` subclass, not a change to a model. + +Two call sites, and only two: + + * At construction, for weight metadata:: + + nnx.Linear(..., kernel_metadata={"sharding": s("gate", ("embed", "mlp"))}) + + * Inside `__call__`, for activation constraints:: + + x = s.constrain(x, "mlp_mid", ("batch", "length", "mlp")) + +Parameters are sharded once by the initializer. Re-asserting a constraint on a +parameter afterwards collides with the training engine's deferred all-reduce +and Zero-1 resharding, so `constrain` is hardwired to `TensorType.ACTIVATION` +and cannot be used on weights. +""" + +import abc +import enum +import math +from collections.abc import Mapping, Sequence + +import jax +from jax.experimental import mesh_utils +from jax.sharding import Mesh +from jax.sharding import PartitionSpec as P + +# Data, fully-sharded data, tensor and expert parallelism, in mesh order. +# +# Order is not cosmetic: it is the order handed to `mesh_utils`, which places +# axes against the physical topology. Inserting an axis in the middle changes +# the placement of everything after it. +# +# TODO: This is a basic starting point, this list will be expanded soon. +# Legacy carries twelve axes, but that list is a union -- its seven strategy +# configs under `configs/custom_mesh_and_rule/` use between one and six. The +# ones m3 is likely to need next are `stage` (pipeline), `context` and +# `context_usp_ulysses` (context parallelism), and `tensor_sequence`. +MESH_AXIS_NAMES = ("dp", "fsdp", "tp", "ep") + +# A logical axis maps to one mesh axis, several (their product shards the +# dimension), or None to leave the dimension replicated. +MeshAxes = str | tuple[str, ...] | None + + +class TensorType(enum.Enum): + """Whether an axis is being mapped for a parameter or for an activation. + + The same logical axis often wants different treatment in each case. Under + FSDP, `"embed"` shards weights across `fsdp` but is usually left replicated + for activations. + """ + + WEIGHT = "weight" + ACTIVATION = "activation" + + +def _ambient_mesh_is_explicit() -> bool: + """Reports whether the mesh in scope uses `Explicit` axis types. + + The two sharding modes need different APIs and each rejects the other's: + `with_sharding_constraint` asserts rather than shards under an explicit mesh, + and `reshard` refuses a `PartitionSpec` naming auto axes. So `constrain` has + to know which mode it is in. + + The mode is read off the mesh in scope rather than taken as a constructor + argument or config flag, because the mesh is what the two APIs actually + check. A flag can disagree with the mesh; the mesh cannot disagree with + itself. Legacy does the same in `optimizers/reshape_utils.py`, though + `utils/sharding.py` threads a `ShardMode` through instead. + + Raises: + ValueError: If the mesh mixes explicit and non-explicit axes, where no + single API is correct for the whole spec. + """ + axis_types = tuple(jax.sharding.get_abstract_mesh().axis_types) + explicit = sum(1 for t in axis_types if t == jax.sharding.AxisType.Explicit) + if 0 < explicit < len(axis_types): + raise ValueError(f"Mesh mixes explicit and non-explicit axis types ({axis_types}), which m3 does not support.") + return explicit == len(axis_types) and explicit > 0 + + +def _validate_mesh_axes(mapped: MeshAxes, axis: str, tensor_name: str) -> None: + """Rejects mesh axis names that are not on the mesh. + + Without this a typo such as `"tensor"` for `"tp"` surfaces much later as an + opaque sharding error, far from the rule that produced it. + """ + if mapped is None: + return + names = (mapped,) if isinstance(mapped, str) else tuple(mapped) + unknown = [name for name in names if name not in MESH_AXIS_NAMES] + if unknown: + raise ValueError( + f"Sharding rule for logical axis {axis!r} of tensor {tensor_name!r} returned unknown mesh " + f"axes {unknown}; expected names from {MESH_AXIS_NAMES}." + ) + + +class Sharding(abc.ABC): + """Maps a model's logical axis names onto mesh axes. + + A subclass encodes one parallelism strategy. Strategies are meant to be + reused across models, which is why the common logical vocabulary -- + `"embed"`, `"mlp"`, `"heads"`, `"kv_heads"`, `"norm"`, `"layers"` -- should + mean the same thing everywhere. A model with a genuinely novel axis extends a + strategy and delegates the rest:: + + class PureFsdpSharding(Sharding): + def map_axis(self, axis, tensor_name, tensor_type): + match axis, tensor_type: + case "batch", TensorType.ACTIVATION: return ("dp", "fsdp") + case "embed", TensorType.WEIGHT: return "fsdp" + case "embed", TensorType.ACTIVATION: return None + case ("heads" | "kv_heads"), _: return "tp" + case "mlp", _: return "tp" + case _: raise KeyError(axis) + + class Qwen3NextSharding(PureFsdpSharding): + def map_axis(self, axis, tensor_name, tensor_type): + if axis == "gdn_head": + return "tp" + return super().map_axis(axis, tensor_name, tensor_type) + + Keeping rules per-strategy rather than per-model matters: legacy's rule table + is model-agnostic, so N models and S strategies cost S tables. A rule set + written per model costs N x S instead. + + Prefer raising on an unrecognised axis over a catch-all that returns None. A + catch-all turns a misspelled logical axis into silent full replication, which + costs performance without ever failing a test. Raising surfaces it the first + time the model is constructed. + """ + + @abc.abstractmethod + def map_axis(self, axis: str, tensor_name: str, tensor_type: TensorType) -> MeshAxes: + """Maps one logical axis to mesh axis name(s), or None to leave it replicated. + + Args: + axis: Logical axis name, e.g. `"embed"`. + tensor_name: Name of the tensor being sharded, for rules that need to + distinguish otherwise identical axes. + tensor_type: Whether this is a parameter or an activation. + """ + + def __call__( + self, + tensor_name: str, + axes: Sequence[str | None], + tensor_type: TensorType = TensorType.WEIGHT, + ) -> P: + """Resolves a tuple of logical axes into a `PartitionSpec`. + + A `None` entry in `axes` marks a dimension with no logical name, which is + left replicated without consulting `map_axis`. + + Raises if two dimensions claim the same mesh axis. Legacy resolves that + case through Flax's fallback rules, which skip a rule whose mesh axis is + already taken and fall through to a shorter one -- hence the duplicate keys + in `logical_axis_rules`. m3 has no fallback, so the conflict is reported + here rather than surfacing as an opaque XLA error later. + """ + resolved = [] + claimed: dict[str, str] = {} + for axis in axes: + if axis is None: + resolved.append(None) + continue + mapped = self.map_axis(axis, tensor_name, tensor_type) + _validate_mesh_axes(mapped, axis, tensor_name) + for mesh_axis in (mapped,) if isinstance(mapped, str) else tuple(mapped or ()): + if mesh_axis in claimed: + raise ValueError( + f"Tensor {tensor_name!r} maps both {claimed[mesh_axis]!r} and {axis!r} onto mesh axis " + f"{mesh_axis!r}; a mesh axis can shard only one dimension of a tensor." + ) + claimed[mesh_axis] = axis + resolved.append(mapped) + return P(*resolved) + + def constrain(self, x: jax.Array, tensor_name: str, axes: Sequence[str | None]) -> jax.Array: + """Applies an activation sharding constraint to `x`. + + Activations only, by construction -- see the module docstring on why + re-constraining parameters after init breaks the training engine. + + Dispatches on the sharding mode of the mesh in scope. Under an explicit + mesh `with_sharding_constraint` is an assertion, not a constraint, and + fails on anything not already sharded as asked; `reshard` is the operation + that moves data there. + """ + spec = self(tensor_name, axes, TensorType.ACTIVATION) + if _ambient_mesh_is_explicit(): + return jax.sharding.reshard(x, spec) + return jax.lax.with_sharding_constraint(x, spec) + + +def _resolve_axis_sizes(sizes: Sequence[int], target_product: int, kind: str, target_desc: str) -> tuple[int, ...]: + """Fills in a single `-1` axis and checks the result tiles `target_product` exactly. + + Split out from `create_mesh` so the arithmetic is testable without needing a + particular number of devices attached. + + Args: + sizes: Per-axis sizes, in `MESH_AXIS_NAMES` order. + target_product: What the sizes must multiply to. + kind: `"ICI"` or `"DCN"`; names the network in error messages. + target_desc: A count noun naming what `target_product` counts, such as + `"slice count"`. Phrased so error messages read correctly at any value. + + Returns: + The resolved sizes, in `MESH_AXIS_NAMES` order. + + Raises: + ValueError: If more than one axis is -1, if any size is invalid, or if the + sizes do not multiply to `target_product`. + """ + resolved = list(sizes) + + if resolved.count(-1) > 1: + raise ValueError(f"At most one {kind} mesh axis may be -1 (auto); got {dict(zip(MESH_AXIS_NAMES, resolved))}.") + + invalid = {name: size for name, size in zip(MESH_AXIS_NAMES, resolved) if size < 1 and size != -1} + if invalid: + raise ValueError(f"{kind} mesh axis sizes must be >= 1, or -1 for auto; got {invalid}.") + + if -1 in resolved: + specified = math.prod(size for size in resolved if size != -1) + if target_product % specified: + raise ValueError( + f"Cannot auto-size the -1 {kind} mesh axis: {target_desc} ({target_product}) is not divisible by " + f"{specified}, the product of the specified axes {dict(zip(MESH_AXIS_NAMES, resolved))}." + ) + resolved[resolved.index(-1)] = target_product // specified + + product = math.prod(resolved) + if product != target_product: + raise ValueError( + f"{kind} mesh axes {dict(zip(MESH_AXIS_NAMES, resolved))} have product {product}, " + f"which does not match the {target_desc} ({target_product})." + ) + return tuple(resolved) + + +def _count_slices(devices: Sequence[jax.Device]) -> int: + """Counts how many slices `devices` spans. + + Devices that are not part of a multi-slice job carry no `slice_index` + attribute at all -- true of CPU devices and of single-slice TPU devices -- + so the default folds them into one group. + """ + return len({getattr(device, "slice_index", 0) for device in devices}) + + +def _shape_from_mapping(sizes: Mapping[str, int] | None, kind: str) -> tuple[int, ...]: + """Turns `{axis: size}` into a shape tuple in `MESH_AXIS_NAMES` order. + + Unnamed axes default to 1, so a caller names only the axes it is actually + splitting. Unknown names are rejected rather than ignored -- a typo would + otherwise silently leave that axis at 1 and shard nothing. + """ + sizes = dict(sizes or {}) + unknown = sorted(set(sizes) - set(MESH_AXIS_NAMES)) + if unknown: + raise ValueError(f"Unknown {kind} mesh axes {unknown}; expected names from {MESH_AXIS_NAMES}.") + return tuple(sizes.get(name, 1) for name in MESH_AXIS_NAMES) + + +def create_mesh( + *, + ici: Mapping[str, int] | None = None, + dcn: Mapping[str, int] | None = None, + devices: Sequence[jax.Device] | None = None, +) -> Mesh: + """Builds the device mesh over `MESH_AXIS_NAMES`. + + The mesh is the product of two shapes over the same axes. `ici` spreads the + devices within one slice over the fast interconnect; `dcn` spreads the slices + themselves over the data center network. Each may carry at most one `-1`, + meaning "use whatever is left", and both may do so at once because they + resolve against different totals: devices per slice, and slices. + + Axes left unnamed default to 1, so callers name only what they split:: + + create_mesh(ici={"fsdp": -1}) # every device on fsdp + create_mesh(ici={"fsdp": -1, "tp": 2}) # tp within a slice + create_mesh(ici={"fsdp": -1}, dcn={"dp": -1}) # fsdp in a slice, dp across + + Mappings rather than one keyword per axis: the axis list is expected to grow + (see `MESH_AXIS_NAMES`), and a keyword pair per axis per network would scale + to twenty-odd parameters and churn every call site each time one is added. + + The defaults -- `ici={"fsdp": -1}`, `dcn={"dp": -1}` -- put every device in a + slice on `fsdp` and every slice on `dp`. That split is the usual starting + point because DCN bandwidth comfortably carries a gradient all-reduce but not + the per-layer traffic that FSDP or tensor parallelism generate. On a single + device it degenerates to a mesh where every `PartitionSpec` resolves to + replicated, which is what keeps the single-device code path identical to the + distributed one instead of a special case. + + Device ordering is delegated to `jax.experimental.mesh_utils`, which lays the + mesh out over the detected topology; this function only picks which of its + two helpers to call. Hand-tuned orderings still beat it on some topologies -- + legacy exposes a couple as `custom_mesh` -- which m3 does not offer yet. + https://docs.jax.dev/en/latest/jax.experimental.mesh_utils.html + + The returned mesh must be the same object the training driver uses. A + `Sharding` resolved against a different mesh is silently replaced by full + replication rather than raising. + + Args: + ici: Axis sizes within a slice, e.g. `{"fsdp": -1, "tp": 2}`. Defaults to + `{"fsdp": -1}`. + dcn: Axis sizes across slices. Defaults to `{"dp": -1}`. + devices: Devices to arrange; defaults to `jax.devices()`. + + Returns: + A `Mesh` over `MESH_AXIS_NAMES`. + + Raises: + ValueError: If `devices` is empty, if either mapping names an unknown axis, + has more than one -1, contains an invalid size, or does not multiply to + the device or slice count. + """ + ici_sizes = _shape_from_mapping({"fsdp": -1} if ici is None else ici, "ICI") + dcn_sizes = _shape_from_mapping({"dp": -1} if dcn is None else dcn, "DCN") + + devices = list(jax.devices() if devices is None else devices) + if not devices: + raise ValueError("Cannot build a mesh over zero devices.") + + num_slices = _count_slices(devices) + if len(devices) % num_slices: + raise ValueError(f"{len(devices)} devices do not divide evenly into {num_slices} slices.") + + ici_shape = _resolve_axis_sizes(ici_sizes, len(devices) // num_slices, "ICI", "per-slice device count") + dcn_shape = _resolve_axis_sizes(dcn_sizes, num_slices, "DCN", "slice count") + + if num_slices > 1: + device_array = mesh_utils.create_hybrid_device_mesh(ici_shape, dcn_shape, devices) + else: + # Not `create_hybrid_device_mesh` with an all-ones DCN shape: it requires + # every device to expose `slice_index`, which single-slice TPU and CPU + # devices do not, so it raises rather than degrading to the ICI-only case. + device_array = mesh_utils.create_device_mesh(ici_shape, devices) + + return Mesh(device_array, MESH_AXIS_NAMES) diff --git a/src/maxtext/m3/infra/__init__.py b/src/maxtext/m3/infra/__init__.py new file mode 100644 index 0000000000..817af52b1a --- /dev/null +++ b/src/maxtext/m3/infra/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Post-construction transforms applied from training setup, never from models. + +Planned modules: + * `remat.py` — path-transparent remat wrapper + `maybe_apply_remat`. + * `scan.py` — layer stacking via `jax.lax.scan` + `maybe_apply_scan`. + * `quantization.py` — qwix int8/fp8 rules + `maybe_quantize`. + * `offload.py` — `to_host` for optimizer/parameter offloading. + +Composition order is fixed: +`construct -> remat -> scan -> quantize -> optimizer -> host offload`. + +Model files contain only zero-cost annotations (`checkpoint_name` tags and +sharding specs); they must never import from this package. +""" diff --git a/src/maxtext/m3/models/__init__.py b/src/maxtext/m3/models/__init__.py new file mode 100644 index 0000000000..8cd25ba4db --- /dev/null +++ b/src/maxtext/m3/models/__init__.py @@ -0,0 +1,46 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""m3 model garden: one self-contained directory per model family. + +Each family owns its architecture end to end: + + models/ + qwen3/ + modeling_qwen3.py # attention, MLP, decoder layer, model + +This module holds the single dispatch point (create_model) and the model +registry. There is no shared Decoder, no shared attention class, and no +model-name branching. +""" + +from collections.abc import Callable +from maxtext.m3.models.qwen3.modeling_qwen3 import create_qwen3_model +from maxtext.m3.models.qwen3_vl.modeling_qwen3_vl import create_qwen3_vl_model + + +MODEL_REGISTRY: dict[str, Callable] = { + "qwen3-0.6b": create_qwen3_model, + "qwen3-vl-2b": create_qwen3_vl_model, +} + + +def create_model(config, mesh, **kwargs): + """Dispatches to a registered m3 constructor using driver configuration.""" + if config.model_name not in MODEL_REGISTRY: + raise ValueError( + f"Model {config.model_name!r} does not support the m3 backend (use_m3_model=True). " + f"Registered models: {', '.join(sorted(MODEL_REGISTRY))}." + ) + return MODEL_REGISTRY[config.model_name](config, mesh, **kwargs) diff --git a/src/maxtext/m3/models/diffuser_schedulers/__init__.py b/src/maxtext/m3/models/diffuser_schedulers/__init__.py new file mode 100644 index 0000000000..33af0da8a2 --- /dev/null +++ b/src/maxtext/m3/models/diffuser_schedulers/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Diffusion noise schedulers package.""" diff --git a/src/maxtext/m3/models/diffuser_schedulers/scheduling_flow_match.py b/src/maxtext/m3/models/diffuser_schedulers/scheduling_flow_match.py new file mode 100644 index 0000000000..096c493f1b --- /dev/null +++ b/src/maxtext/m3/models/diffuser_schedulers/scheduling_flow_match.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Continuous-time Flow Match diffusion noise scheduler.""" diff --git a/src/maxtext/m3/models/diffuser_schedulers/scheduling_unipc_multistep.py b/src/maxtext/m3/models/diffuser_schedulers/scheduling_unipc_multistep.py new file mode 100644 index 0000000000..1c4cab0bf6 --- /dev/null +++ b/src/maxtext/m3/models/diffuser_schedulers/scheduling_unipc_multistep.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""UniPC multistep ODE solver diffusion noise scheduler.""" diff --git a/src/maxtext/m3/models/diffuser_schedulers/scheduling_utils.py b/src/maxtext/m3/models/diffuser_schedulers/scheduling_utils.py new file mode 100644 index 0000000000..6ad06b5af9 --- /dev/null +++ b/src/maxtext/m3/models/diffuser_schedulers/scheduling_utils.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Diffusion noise scheduler mixins and output dataclasses.""" diff --git a/src/maxtext/m3/models/qwen3/__init__.py b/src/maxtext/m3/models/qwen3/__init__.py new file mode 100644 index 0000000000..55bbd501f8 --- /dev/null +++ b/src/maxtext/m3/models/qwen3/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the License); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://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. + +"""Qwen3 model family package.""" diff --git a/src/maxtext/m3/models/qwen3/modeling_qwen3.py b/src/maxtext/m3/models/qwen3/modeling_qwen3.py new file mode 100644 index 0000000000..d20d73ff62 --- /dev/null +++ b/src/maxtext/m3/models/qwen3/modeling_qwen3.py @@ -0,0 +1,490 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://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. + +"""Qwen3 dense model in m3 format. + +This module provides a self-contained implementation of the Qwen3 model +family (dense variant) following the m3 modern MaxText architecture: +- Composed directly from Flax NNX built-ins: nnx.Linear, nnx.Einsum, nnx.RMSNorm, nnx.Embed. +- Completely free of legacy maxtext.layers dependencies (AttentionOp, DenseGeneral, + Linen RMSNorm, Linen Embed, quantizations). +- Pure NNX module composition with clean parameter paths matching existing checkpoints. +- Uses m3 core RoPE. +- Designed with clear extension points for future MoE and hybrid variants. +""" + +from typing import Any, Optional +from flax import nnx +import jax +import jax.numpy as jnp +from jax.sharding import Mesh + +from maxtext.common.common_types import Config +from maxtext.m3.core.rope import apply_rope + + +class Qwen3MLP(nnx.Module): + """Qwen3 feed-forward network with SwiGLU activation. + + Submodule and parameter names match standard checkpoint paths: + wi_0 (gate), wi_1 (up), wo (down). + """ + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: nnx.Rngs, + ): + """Initializes Qwen3MLP feed-forward network. + + Args: + config: Model configuration specifying dimensions and dtypes. + mesh: Device mesh used for parallel execution. + rngs: NNX random number generators. + """ + self.config = config + + kw = { + "use_bias": False, + "dtype": config.dtype, + "param_dtype": config.weight_dtype, + "rngs": rngs, + } + self.wi_0 = nnx.Linear( + in_features=config.emb_dim, + out_features=config.mlp_dim, + **kw, + ) + self.wi_1 = nnx.Linear( + in_features=config.emb_dim, + out_features=config.mlp_dim, + **kw, + ) + self.wo = nnx.Linear( + in_features=config.mlp_dim, + out_features=config.emb_dim, + **kw, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + """Applies gated SwiGLU feed-forward projection: wo(silu(wi_0(x)) * wi_1(x)).""" + gate = jax.nn.silu(self.wi_0(x)) + up = self.wi_1(x) + hidden = (gate * up).astype(self.config.dtype) + return self.wo(hidden) + + +class Qwen3Attention(nnx.Module): + """Qwen3 multi-head / grouped-query attention with QK-norm, RoPE, and self-contained attention. + + Submodule and parameter names match checkpoint paths: + query, key, value, out, query_norm, key_norm. + """ + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: nnx.Rngs, + ): + """Initializes Qwen3 Attention module with Q, K, V projections and QK-Norm. + + Args: + config: Model configuration specifying dimensions, head counts, and RoPE settings. + mesh: Device mesh used for parallel execution. + rngs: NNX random number generators. + """ + self.config = config + + self.num_query_heads = config.num_query_heads + self.num_kv_heads = config.num_kv_heads + self.head_dim = config.head_dim + + kw = { + "bias_shape": None, + "dtype": config.dtype, + "param_dtype": config.weight_dtype, + "rngs": rngs, + } + # Q, K, V multi-axis projections + self.query = nnx.Einsum( + "...e,ehd->...hd", + (config.emb_dim, self.num_query_heads, self.head_dim), + **kw, + ) + self.key = nnx.Einsum( + "...e,ehd->...hd", + (config.emb_dim, self.num_kv_heads, self.head_dim), + **kw, + ) + self.value = nnx.Einsum( + "...e,ehd->...hd", + (config.emb_dim, self.num_kv_heads, self.head_dim), + **kw, + ) + self.out = nnx.Einsum( + "...hd,hde->...e", + (self.num_query_heads, self.head_dim, config.emb_dim), + **kw, + ) + + # QK-Norm + norm_kw = { + "epsilon": config.normalization_layer_epsilon, + "dtype": config.dtype, + "param_dtype": config.weight_dtype, + "rngs": rngs, + } + self.query_norm = nnx.RMSNorm( + num_features=self.head_dim, + **norm_kw, + ) + self.key_norm = nnx.RMSNorm( + num_features=self.head_dim, + **norm_kw, + ) + + def _compute_attention( + self, + query: jax.Array, + key: jax.Array, + value: jax.Array, + decoder_positions: Optional[jax.Array] = None, + decoder_segment_ids: Optional[jax.Array] = None, + ) -> jax.Array: + """Computes self-contained scaled dot-product attention with GQA and causal masking.""" + b, t, n_q, d = query.shape + n_kv = key.shape[2] + g = n_q // n_kv + query_g = jnp.reshape(query, (b, t, n_kv, g, d)) + + def _compute_local_block(q_g, k, v, mask=None): + """Computes local attention output, max, and sum for scaled dot-product block.""" + weights = jnp.einsum("btkgd,bskd->bkgts", q_g, k) + if mask is not None: + weights = jnp.where(mask, weights, -1e10) + if self.config.float32_logits: + weights = weights.astype(jnp.float32) + logits = jnp.reshape(weights, (b, n_q, t, -1)) + local_max = jnp.max(logits, axis=-1, keepdims=True) + local_exp = jnp.exp(logits - local_max) + local_sum = jnp.sum(local_exp, axis=-1, keepdims=True) + + local_exp = jnp.reshape(local_exp.astype(v.dtype), (b, n_kv, g, t, -1)) + local_out = jnp.einsum("bkgts,bskd->btkgd", local_exp, v) + local_out = jnp.reshape(local_out, (b, t, n_q, d)) + + local_max = jnp.transpose(local_max, (0, 2, 1, 3)) + local_sum = jnp.transpose(local_sum, (0, 2, 1, 3)) + return local_out, local_max, local_sum + + s = key.shape[1] + q_pos = jnp.arange(t) if decoder_positions is None else decoder_positions + k_pos = jnp.arange(s) if decoder_positions is None else decoder_positions + causal_mask = k_pos[..., None, :] <= q_pos[..., :, None] + if decoder_segment_ids is not None: + seg_mask = decoder_segment_ids[..., :, None] == decoder_segment_ids[..., None, :] + mask = causal_mask & seg_mask + else: + mask = causal_mask + mask = mask[:, None, None, :, :] + out, _, lsum = _compute_local_block(query_g, key, value, mask) + return (out / lsum).astype(self.config.dtype) + + def __call__( + self, + inputs: jax.Array, + decoder_positions: jax.Array, + decoder_segment_ids: Optional[jax.Array] = None, + deterministic: bool = False, + **kwargs, + ) -> jax.Array: + """Executes Qwen3 self-attention forward pass with RoPE.""" + query = self.query(inputs) + key = self.key(inputs) + value = self.value(inputs) + + # QK-Norm before RoPE + if self.config.use_qk_norm: + query = self.query_norm(query) + key = self.key_norm(key) + + # RoPE + query = apply_rope(query, decoder_positions, max_timescale=self.config.rope_max_timescale) + key = apply_rope(key, decoder_positions, max_timescale=self.config.rope_max_timescale) + + # Scale query by 1/sqrt(head_dim) + query_pre_attn_scalar = self.head_dim**-0.5 + query = query * query_pre_attn_scalar + + attn_out = self._compute_attention( + query=query, + key=key, + value=value, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + ) + return self.out(attn_out) + + +class Qwen3DecoderLayer(nnx.Module): + """Qwen3 Transformer decoder layer (dense). + + Pre-norm architecture with self-attention and MLP residual blocks. + Leaves clean extension points for MoE or hybrid variants. + """ + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: nnx.Rngs, + ): + """Initializes pre-norm attention and MLP residual blocks. + + Args: + config: Model configuration specifying dimensions, layers, and dtypes. + mesh: Device mesh used for parallel execution. + rngs: NNX random number generators. + """ + self.config = config + + norm_kw = { + "epsilon": config.normalization_layer_epsilon, + "dtype": config.dtype, + "param_dtype": config.weight_dtype, + "rngs": rngs, + } + self.pre_self_attention_layer_norm = nnx.RMSNorm( + num_features=config.emb_dim, + **norm_kw, + ) + self.self_attention = Qwen3Attention( + config=config, + rngs=rngs, + ) + self.post_self_attention_layer_norm = nnx.RMSNorm( + num_features=config.emb_dim, + **norm_kw, + ) + self.mlp = Qwen3MLP( + config=config, + rngs=rngs, + ) + + def __call__( + self, + inputs: jax.Array, + decoder_segment_ids: Optional[jax.Array] = None, + decoder_positions: Optional[jax.Array] = None, + deterministic: bool = False, + **kwargs, + ) -> jax.Array: + """Applies pre-norm attention and pre-norm MLP blocks with residual connections.""" + # Pre-attention norm & Self-attention + normed_attn_in = self.pre_self_attention_layer_norm(inputs) + attn_out = self.self_attention( + normed_attn_in, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=deterministic, + ) + x = (inputs + attn_out).astype(self.config.dtype) + + # Post-attention norm & MLP + normed_mlp_in = self.post_self_attention_layer_norm(x) + mlp_out = self.mlp(normed_mlp_in) + output = (x + mlp_out).astype(self.config.dtype) + + return output + + +class Qwen3Decoder(nnx.Module): + """Qwen3 Decoder stack composing decoder layers with flat attribute naming.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: nnx.Rngs, + layer_cls: type[nnx.Module] = Qwen3DecoderLayer, + ): + """Initializes the Qwen3 decoder stack of N layers and final layer normalization. + + Args: + config: Model configuration specifying layer count, dimensions, and dtypes. + mesh: Device mesh used for parallel execution. + rngs: NNX random number generators. + layer_cls: Decoder layer constructor for model-family subclasses. + """ + self.config = config + self.num_layers = config.num_decoder_layers + + # Flat attribute names: layers_0, layers_1, ... + for lyr in range(self.num_layers): + layer = layer_cls( + config=config, + rngs=rngs, + ) + setattr(self, f"layers_{lyr}", layer) + + self.decoder_norm = nnx.RMSNorm( + num_features=config.emb_dim, + epsilon=config.normalization_layer_epsilon, + dtype=config.dtype, + param_dtype=config.weight_dtype, + rngs=rngs, + ) + + def apply_output_head( + self, + shared_embedding: Any, + y: jax.Array, + deterministic: bool = False, + ) -> jax.Array: + """Applies final decoder norm and projects hidden states to logits via tied embedding.""" + y = self.decoder_norm(y) + attend_dtype = jnp.float32 if self.config.logits_dot_in_fp32 else self.config.dtype + if self.config.logits_dot_in_fp32: + table = ( + shared_embedding.embedding.value + if hasattr(shared_embedding, "embedding") + else getattr(shared_embedding, "value", shared_embedding) + ) + logits = jnp.dot(y.astype(attend_dtype), jnp.asarray(table, attend_dtype).T) + elif hasattr(shared_embedding, "attend"): + logits = shared_embedding.attend(y) + else: + table = ( + shared_embedding.embedding.value + if hasattr(shared_embedding, "embedding") + else getattr(shared_embedding, "value", shared_embedding) + ) + logits = jnp.dot(y.astype(attend_dtype), jnp.asarray(table, attend_dtype).T) + + if self.config.normalize_embedding_logits: + logits = logits / jnp.sqrt(y.shape[-1]) + if self.config.final_logits_soft_cap: + logits = logits / self.config.final_logits_soft_cap + logits = jnp.tanh(logits) * self.config.final_logits_soft_cap + if self.config.cast_logits_to_fp32: + logits = logits.astype(jnp.float32) + return logits + + def __call__( + self, + shared_embedding: Any, + decoder_input_tokens: jax.Array, + decoder_positions: jax.Array, + decoder_segment_ids: Optional[jax.Array] = None, + deterministic: bool = False, + **kwargs, + ) -> tuple[jax.Array, jax.Array]: + """Iterates through all decoder layers and projects to logits via tied embeddings.""" + y = shared_embedding(decoder_input_tokens) + + for lyr in range(self.num_layers): + layer = getattr(self, f"layers_{lyr}") + y = layer( + y, + decoder_segment_ids=decoder_segment_ids, + decoder_positions=decoder_positions, + deterministic=deterministic, + ) + + logits = self.apply_output_head(shared_embedding, y, deterministic=deterministic) + return logits, y + + +class Qwen3Model(nnx.Module): + """Top-level Qwen3 autoregressive language model in m3 format.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: nnx.Rngs, + **kwargs, + ): + """Initializes the top-level Qwen3 language model. + + Args: + config: Model configuration specifying dimensions, layers, and vocabulary. + mesh: Device mesh used for parallel execution. + rngs: NNX random number generators. + **kwargs: Additional keyword arguments. + """ + self.config = config + + self.token_embedder = nnx.Embed( + num_embeddings=config.vocab_size, + features=config.emb_dim, + dtype=config.dtype, + param_dtype=config.weight_dtype, + rngs=rngs, + ) + self.decoder = Qwen3Decoder( + config=config, + rngs=rngs, + ) + + def logits_from_hidden_states_for_vocab_tiling( + self, hidden_states: jax.Array, deterministic: bool, **kwargs + ) -> jax.Array: + """Projects hidden states to logits for vocab-tiled computation.""" + return self.decoder.apply_output_head( + shared_embedding=self.token_embedder, + y=hidden_states, + deterministic=deterministic, + ) + + def __call__( + self, + decoder_input_tokens: jnp.ndarray, + decoder_positions: jnp.ndarray, + decoder_segment_ids: Optional[jnp.ndarray] = None, + cache: Any = None, + enable_dropout: bool = False, + true_length: Optional[int] = None, + **kwargs, + ) -> jax.Array: + """Executes full Qwen3 forward pass from input token IDs to output logits.""" + deterministic = kwargs.get("deterministic", not enable_dropout) + logits, _ = self.decoder( + shared_embedding=self.token_embedder, + decoder_input_tokens=decoder_input_tokens, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=deterministic, + ) + return logits + + +def create_qwen3_model( + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: Optional[nnx.Rngs] = None, + **kwargs, +) -> Qwen3Model: + """Factory function creating a Qwen3Model instance.""" + if rngs is None: + rngs = nnx.Rngs(0) + return Qwen3Model(config, mesh, rngs=rngs, **kwargs) diff --git a/src/maxtext/m3/models/qwen3_vl/__init__.py b/src/maxtext/m3/models/qwen3_vl/__init__.py new file mode 100644 index 0000000000..9a94bc7e1f --- /dev/null +++ b/src/maxtext/m3/models/qwen3_vl/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://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. + +"""Qwen3-VL multimodal model family package.""" diff --git a/src/maxtext/m3/models/qwen3_vl/modeling_qwen3_vl.py b/src/maxtext/m3/models/qwen3_vl/modeling_qwen3_vl.py new file mode 100644 index 0000000000..9ef6b9dc72 --- /dev/null +++ b/src/maxtext/m3/models/qwen3_vl/modeling_qwen3_vl.py @@ -0,0 +1,1083 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://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. + +"""Qwen3-VL Multimodal Model Family in m3 format. + +Self-contained NNX implementation for multimodal forward passes: +- No imports from maxtext.layers. +- NNX parameters and built-ins preserve existing checkpoint paths. +- Pure RoPE functions: apply_vision_rope (2D) and apply_mrope (3D). +- Qwen3VLDecoder inherits directly from Qwen3Decoder. +- Reuses Qwen3MLP from qwen3. +- Parameter hierarchy matches legacy checkpoints (Qwen3VLVisionEncoder_0, Qwen3VLVisionProjector_0). +""" + +from typing import Any, Optional, Sequence, Union +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh + +from maxtext.common.common_types import Config, MultimodalInput +from maxtext.m3.models.qwen3.modeling_qwen3 import ( + Qwen3MLP, + Qwen3Decoder, +) + + +# --------------------------------------------------------------------------- +# Core Pure Linear Layer +# --------------------------------------------------------------------------- + + +class DenseGeneral(nnx.Module): + """Self-contained linear layer matching legacy checkpoint kernel/bias structure.""" + + def __init__( + self, + in_features_shape: Union[int, tuple[int, ...]], + out_features_shape: Union[int, tuple[int, ...]], + use_bias: bool = True, + dtype: Any = jnp.float32, + weight_dtype: Any = jnp.float32, + rngs: Optional[nnx.Rngs] = None, + ): + self.in_features_shape = (in_features_shape,) if isinstance(in_features_shape, int) else tuple(in_features_shape) + self.out_features_shape = (out_features_shape,) if isinstance(out_features_shape, int) else tuple(out_features_shape) + self.use_bias = use_bias + self.dtype = dtype + self.weight_dtype = weight_dtype + + kernel_shape = self.in_features_shape + self.out_features_shape + bias_shape = self.out_features_shape + + if rngs is not None: + fan_in = np.prod(self.in_features_shape) + init_fn = nnx.initializers.normal(stddev=float(fan_in) ** -0.5) + self.kernel = nnx.Param(init_fn(rngs.params(), kernel_shape, weight_dtype)) + if use_bias: + self.bias = nnx.Param(jnp.zeros(bias_shape, dtype=weight_dtype)) + else: + self.bias = None + else: + self.kernel = nnx.Param(jnp.zeros(kernel_shape, dtype=weight_dtype)) + if use_bias: + self.bias = nnx.Param(jnp.zeros(bias_shape, dtype=weight_dtype)) + else: + self.bias = None + + def __call__(self, x: jax.Array) -> jax.Array: + num_in_axes = len(self.in_features_shape) + x_in_axes = tuple(range(x.ndim - num_in_axes, x.ndim)) + kernel_in_axes = tuple(range(num_in_axes)) + out = jnp.tensordot(x.astype(self.dtype), self.kernel[...].astype(self.dtype), axes=(x_in_axes, kernel_in_axes)) + if self.bias is not None: + out = out + self.bias[...].astype(self.dtype) + return out.astype(self.dtype) + + +# --------------------------------------------------------------------------- +# Pure RoPE Mathematics: Vision 2D RoPE & Language 3D M-RoPE +# --------------------------------------------------------------------------- + + +def generate_block_coords( + video_grid_thw: jax.Array, max_patches: int, merge_size: int = 2 +) -> tuple[jax.Array, jax.Array]: + """Generate row and col coordinates in block-based order for padded video.""" + V_T = video_grid_thw[:, 0:1] + V_H = video_grid_thw[:, 1:2] + V_W = video_grid_thw[:, 2:3] + + merged_w = V_W // merge_size + V_len = V_T * V_H * V_W + + idx = jnp.arange(max_patches, dtype=jnp.int32)[None, :] + is_valid = idx < V_len + + stride_lh = V_H * V_W + safe_stride_lh = jnp.maximum(stride_lh, 1) + s_idx = idx % safe_stride_lh + + intra_col = s_idx % merge_size + intra_row = (s_idx // merge_size) % merge_size + + safe_merged_w = jnp.maximum(merged_w, 1) + block_elements = merge_size * merge_size + block_col = (s_idx // block_elements) % safe_merged_w + block_row = s_idx // (safe_merged_w * block_elements) + + row = jnp.where(is_valid, block_row * merge_size + intra_row, 0) + col = jnp.where(is_valid, block_col * merge_size + intra_col, 0) + + return row, col + + +def apply_vision_rope( + inputs: jax.Array, + num_frames: int, + height: int, + width: int, + spatial_merge_size: int = 2, + rope_theta: float = 10000.0, + valid_grid: Optional[tuple[int, int, int] | jax.Array] = None, + token_mask: Optional[jax.Array] = None, +) -> jax.Array: + """Applies 2D rotary position embeddings to vision patch tokens.""" + is_3d = inputs.ndim == 3 + if is_3d: + inputs = inputs[None, :, :, :] + + batch_size = inputs.shape[0] + head_dim = inputs.shape[-1] + max_patches = num_frames * height * width + + if valid_grid is not None: + if isinstance(valid_grid, (tuple, list)): + valid_grid_arr = jnp.array([valid_grid], dtype=jnp.int32) + elif valid_grid.ndim == 1: + valid_grid_arr = valid_grid[None, :] + else: + valid_grid_arr = valid_grid + else: + valid_grid_arr = jnp.tile( + jnp.array([[num_frames, height, width]], dtype=jnp.int32), + (batch_size, 1), + ) + + row, col = generate_block_coords(valid_grid_arr, max_patches, spatial_merge_size) + + max_hw = max(height, width) + inv_freq = 1.0 / (rope_theta ** (jnp.arange(0, head_dim // 2, 2, dtype=jnp.float32) / (head_dim // 2))) + positions = jnp.arange(max_hw, dtype=jnp.float32) + freq_table = jnp.outer(positions, inv_freq) + + row_freqs = freq_table[row] + col_freqs = freq_table[col] + embeddings = jnp.concatenate([row_freqs, col_freqs], axis=-1) + embeddings = jnp.concatenate([embeddings, embeddings], axis=-1) + + cos_emb = jnp.cos(embeddings) + sin_emb = jnp.sin(embeddings) + + if token_mask is not None: + is_valid = token_mask[:, :, None] + cos_emb = jnp.where(is_valid, cos_emb, 1.0) + sin_emb = jnp.where(is_valid, sin_emb, 0.0) + + cos_emb = cos_emb[:, :, None, :].astype(inputs.dtype) + sin_emb = sin_emb[:, :, None, :].astype(inputs.dtype) + + x1 = inputs[..., : head_dim // 2] + x2 = inputs[..., head_dim // 2 :] + rotated = jnp.concatenate([-x2, x1], axis=-1) + + out = inputs * cos_emb + rotated * sin_emb + if is_3d: + out = out[0] + return out + + +def _apply_interleaved_mrope(freqs: jax.Array, mrope_section: tuple[int, int, int]) -> jax.Array: + freqs_t = freqs[..., 0, :] + for dim_idx, offset in enumerate([1, 2], start=1): + section_size = mrope_section[dim_idx] * 3 + idx = slice(offset, section_size, 3) + freqs_t = freqs_t.at[..., idx].set(freqs[..., dim_idx, idx]) + return freqs_t + + +def apply_mrope( + x: jax.Array, + positions: jax.Array, + mrope_section: tuple[int, int, int] = (24, 20, 20), + rope_theta: float = 5000000.0, +) -> jax.Array: + """Applies 3D Multimodal Rotary Position Embeddings (M-RoPE) matching HuggingFace and MaxText.""" + + b, _, _, d = x.shape + half_dim = d // 2 + + # The preprocessor supplies (batch, sequence, 3); text-only positions are (batch, sequence). + if positions.ndim == 3 and positions.shape[-1] == 1: + positions = jnp.squeeze(positions, axis=-1) + if positions.ndim == 2: + positions = jnp.broadcast_to(positions[..., None], positions.shape + (3,)) + if positions.shape != (b, x.shape[1], 3): + raise ValueError(f"positions must have shape {(b, x.shape[1], 3)}, got {positions.shape}") + + fraction = 2.0 * jnp.arange(0, half_dim, 1, dtype=jnp.float32) / d + timescale = 1.0 * (rope_theta / 1.0) ** fraction + inv_freq = 1.0 / timescale + + freqs = positions[..., None] * inv_freq[None, None, None, :] + freqs = _apply_interleaved_mrope(freqs, mrope_section) + + emb = jnp.concatenate([freqs, freqs], axis=-1) + cos = jnp.cos(emb)[:, :, None, :].astype(x.dtype) + sin = jnp.sin(emb)[:, :, None, :].astype(x.dtype) + + x1 = x[..., :half_dim] + x2 = x[..., half_dim:] + rotated = jnp.concatenate([-x2, x1], axis=-1) + return (x * cos) + (rotated * sin) + + +# --------------------------------------------------------------------------- +# Self-contained Vision Tower Components +# --------------------------------------------------------------------------- + + +class Qwen3VLVisionPatchEmbed(nnx.Module): + """3D convolution-based patch embedding for vision inputs.""" + + def __init__( + self, + config: Config, + dtype: Any = jnp.float32, + weight_dtype: Any = jnp.float32, + rngs: Optional[nnx.Rngs] = None, + ): + self.config = config + self.patch_size = config.patch_size_for_vit + self.temporal_patch_size = config.temporal_patch_size_for_vit + self.in_channels = config.num_channels_for_vit + self.embed_dim = config.hidden_size_for_vit + + kernel_size = (self.temporal_patch_size, self.patch_size, self.patch_size) + + self.proj = nnx.Conv( + in_features=self.in_channels, + out_features=self.embed_dim, + kernel_size=kernel_size, + strides=kernel_size, + use_bias=True, + dtype=dtype, + param_dtype=weight_dtype, + rngs=rngs, + ) + + def __call__( + self, hidden_states: jax.Array, video_mask: Optional[jax.Array] = None + ) -> tuple[jax.Array, Optional[jax.Array]]: + hidden_states = jnp.transpose(hidden_states, (0, 2, 3, 4, 1)) + hidden_states = self.proj(hidden_states) + batch_size = hidden_states.shape[0] + seq_len = hidden_states.shape[1] * hidden_states.shape[2] * hidden_states.shape[3] + hidden_states = hidden_states.reshape(batch_size, seq_len, self.embed_dim) + + attention_mask = None + if video_mask is not None: + mask_patch_elements = self.temporal_patch_size * self.patch_size * self.patch_size + attention_mask = video_mask.reshape(video_mask.shape[0], -1, mask_patch_elements).max(axis=-1).astype(jnp.int32) + + return hidden_states, attention_mask + + +class Qwen3VLVisionPosEmbedInterpolate(nnx.Module): + """Self-contained bilinear interpolation of learned 2D positional embeddings.""" + + def __init__( + self, + num_position_embeddings: int = 2304, + hidden_size: int = 1024, + spatial_merge_size: int = 2, + dtype: Any = jnp.float32, + rngs: Optional[nnx.Rngs] = None, + ): + self.num_position_embeddings = num_position_embeddings + self.hidden_size = hidden_size + self.spatial_merge_size = spatial_merge_size + self.dtype = dtype + self.num_grid_per_side = int(num_position_embeddings**0.5) + + if rngs is not None: + init_fn = nnx.initializers.normal(stddev=self.hidden_size**-0.5) + self.pos_embed = nnx.Param(init_fn(rngs.params(), (self.num_position_embeddings, self.hidden_size), self.dtype)) + else: + self.pos_embed = nnx.Param(jnp.zeros((self.num_position_embeddings, self.hidden_size), dtype=self.dtype)) + + def __call__( + self, + num_frames: int, + height: int, + width: int, + video_grid_thw: Optional[jax.Array] = None, + attention_mask: Optional[jax.Array] = None, + ) -> jax.Array: + if video_grid_thw is not None: + if isinstance(video_grid_thw, (tuple, list)): + video_grid_thw = jnp.array([video_grid_thw], dtype=jnp.int32) + elif video_grid_thw.ndim == 1: + video_grid_thw = video_grid_thw[None, :] + batch_size = video_grid_thw.shape[0] + elif attention_mask is not None: + batch_size = attention_mask.shape[0] + else: + batch_size = 1 + + max_patches = num_frames * height * width + + if video_grid_thw is None: + video_grid_thw = jnp.tile( + jnp.array([[num_frames, height, width]], dtype=jnp.int32), + (batch_size, 1), + ) + if attention_mask is None: + attention_mask = jnp.ones((batch_size, max_patches), dtype=jnp.int32) + + row, col = generate_block_coords(video_grid_thw, max_patches, self.spatial_merge_size) + V_H = video_grid_thw[:, 1:2] + V_W = video_grid_thw[:, 2:3] + row_norm = row / jnp.maximum(V_H - 1, 1) + col_norm = col / jnp.maximum(V_W - 1, 1) + + N = self.num_grid_per_side + table = self.pos_embed[...].reshape(N, N, self.hidden_size) + + y = row_norm * (N - 1) + x = col_norm * (N - 1) + + y0 = jnp.floor(y).astype(jnp.int32) + x0 = jnp.floor(x).astype(jnp.int32) + y1 = jnp.minimum(y0 + 1, N - 1) + x1 = jnp.minimum(x0 + 1, N - 1) + + dy = (y - y0)[:, :, None] + dx = (x - x0)[:, :, None] + + embed_00 = table[y0, x0] + embed_01 = table[y0, x1] + embed_10 = table[y1, x0] + embed_11 = table[y1, x1] + + interpolated = ( + (1.0 - dy) * (1.0 - dx) * embed_00 + (1.0 - dy) * dx * embed_01 + dy * (1.0 - dx) * embed_10 + dy * dx * embed_11 + ) + return interpolated * attention_mask[:, :, None] + + +class Qwen3VLVisionAttn(nnx.Module): + """Self-contained Q/K/V/Out projection matching legacy checkpoint parameter paths.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: Optional[nnx.Rngs] = None, + ): + self.config = config + self.num_heads = config.num_attention_heads_for_vit + self.head_dim = config.hidden_size_for_vit // self.num_heads + self.hidden_size = config.hidden_size_for_vit + + self.query = DenseGeneral( + in_features_shape=self.hidden_size, + out_features_shape=(self.num_heads, self.head_dim), + dtype=config.dtype_mm, + weight_dtype=config.weight_dtype, + rngs=rngs, + ) + self.key = DenseGeneral( + in_features_shape=self.hidden_size, + out_features_shape=(self.num_heads, self.head_dim), + dtype=config.dtype_mm, + weight_dtype=config.weight_dtype, + rngs=rngs, + ) + self.value = DenseGeneral( + in_features_shape=self.hidden_size, + out_features_shape=(self.num_heads, self.head_dim), + dtype=config.dtype_mm, + weight_dtype=config.weight_dtype, + rngs=rngs, + ) + self.out = DenseGeneral( + in_features_shape=(self.num_heads, self.head_dim), + out_features_shape=self.hidden_size, + dtype=config.dtype_mm, + weight_dtype=config.weight_dtype, + rngs=rngs, + ) + + +class Qwen3VLVisionAttention(nnx.Module): + """Vision attention wrapper matching legacy checkpoint hierarchy.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: Optional[nnx.Rngs] = None, + ): + self.config = config + self.num_heads = config.num_attention_heads_for_vit + self.head_dim = config.hidden_size_for_vit // self.num_heads + self.hidden_size = config.hidden_size_for_vit + self.attn = Qwen3VLVisionAttn(config, rngs=rngs) + + def __call__( + self, + hidden_states: jax.Array, + num_frames: int, + height: int, + width: int, + attention_mask: Optional[jax.Array] = None, + valid_grid: Optional[tuple[int, int, int]] = None, + ) -> jax.Array: + q = self.attn.query(hidden_states) # (b, s, num_heads, head_dim) + k = self.attn.key(hidden_states) # (b, s, num_heads, head_dim) + v = self.attn.value(hidden_states) # (b, s, num_heads, head_dim) + + q = apply_vision_rope( + q, + num_frames, + height, + width, + self.config.spatial_merge_size_for_vit, + self.config.rope_theta_for_vit, + valid_grid=valid_grid, + token_mask=attention_mask, + ) + k = apply_vision_rope( + k, + num_frames, + height, + width, + self.config.spatial_merge_size_for_vit, + self.config.rope_theta_for_vit, + valid_grid=valid_grid, + token_mask=attention_mask, + ) + + q = q * (self.head_dim**-0.5) + + score_dtype = jnp.float32 if self.config.float32_logits else q.dtype + scores = jnp.einsum("bshd,bthd->bhst", q.astype(score_dtype), k.astype(score_dtype)) + if attention_mask is not None: + mask = attention_mask[:, None, None, :] + scores = jnp.where(mask, scores, -1e10) + + weights = jax.nn.softmax(scores, axis=-1).astype(v.dtype) + attn_out = jnp.einsum("bhst,bthd->bshd", weights, v) + return self.attn.out(attn_out) + + +class Qwen3VLVisionBlock(nnx.Module): + """Vision transformer block with LayerNorm, Attention, and GELU MLP.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: Optional[nnx.Rngs] = None, + ): + self.config = config + hs = self.config.hidden_size_for_vit + inter_dim = self.config.intermediate_size_for_vit + + self.ln1 = nnx.LayerNorm(num_features=hs, epsilon=config.normalization_layer_epsilon, rngs=rngs) + self.ln2 = nnx.LayerNorm(num_features=hs, epsilon=config.normalization_layer_epsilon, rngs=rngs) + self.attn = Qwen3VLVisionAttention(config=config, rngs=rngs) + + self.mlp = DenseGeneral( + in_features_shape=hs, + out_features_shape=inter_dim, + dtype=config.dtype_mm, + weight_dtype=config.weight_dtype, + rngs=rngs, + ) + self.mlp_out = DenseGeneral( + in_features_shape=inter_dim, + out_features_shape=hs, + dtype=config.dtype_mm, + weight_dtype=config.weight_dtype, + rngs=rngs, + ) + + def __call__( + self, + x: jax.Array, + num_frames: int, + height: int, + width: int, + attention_mask: Optional[jax.Array] = None, + valid_grid: Optional[tuple[int, int, int]] = None, + ) -> jax.Array: + x = x + self.attn( + self.ln1(x), + num_frames=num_frames, + height=height, + width=width, + attention_mask=attention_mask, + valid_grid=valid_grid, + ) + y = self.ln2(x) + y = self.mlp(y) + y = jax.nn.gelu(y) + y = self.mlp_out(y) + return x + y + + +class Qwen3VLVisionPatchMerger(nnx.Module): + """Vision patch merger that spatially merges patches using an MLP.""" + + def __init__( + self, + config: Config, + use_postshuffle_norm: bool = False, + dtype: Any = jnp.float32, + weight_dtype: Any = jnp.float32, + rngs: Optional[nnx.Rngs] = None, + ): + self.config = config + self.use_postshuffle_norm = use_postshuffle_norm + self.dtype = dtype + self.weight_dtype = weight_dtype + + spatial_merge_size = config.spatial_merge_size_for_vit + base_hidden_size = config.hidden_size_for_vit + out_hidden_size = config.out_hidden_size_for_vit + + self.hidden_size = base_hidden_size * (spatial_merge_size**2) + + ln_features = self.hidden_size if use_postshuffle_norm else base_hidden_size + self.ln_q = nnx.LayerNorm( + num_features=ln_features, + epsilon=config.normalization_layer_epsilon, + dtype=dtype, + rngs=rngs, + ) + + self.mlp_0 = DenseGeneral( + in_features_shape=self.hidden_size, + out_features_shape=self.hidden_size, + dtype=dtype, + weight_dtype=weight_dtype, + rngs=rngs, + ) + self.mlp_2 = DenseGeneral( + in_features_shape=self.hidden_size, + out_features_shape=out_hidden_size, + dtype=dtype, + weight_dtype=weight_dtype, + rngs=rngs, + ) + + def __call__(self, hidden: jax.Array) -> jax.Array: + spatial_merge_size = self.config.spatial_merge_size_for_vit + base_hidden_size = self.config.hidden_size_for_vit + tokens_per_block = spatial_merge_size**2 + + batch_size = hidden.shape[0] + seq_len = hidden.shape[1] + num_blocks = seq_len // tokens_per_block + + hidden = hidden.reshape(batch_size, num_blocks, tokens_per_block * base_hidden_size) + + if self.use_postshuffle_norm: + hidden = self.ln_q(hidden) + else: + hidden_unmerged = hidden.reshape(batch_size, seq_len, base_hidden_size) + hidden_unmerged = self.ln_q(hidden_unmerged) + hidden = hidden_unmerged.reshape(batch_size, num_blocks, tokens_per_block * base_hidden_size) + + hidden = self.mlp_0(hidden) + hidden = jax.nn.gelu(hidden) + hidden = self.mlp_2(hidden) + return hidden + + +class Qwen3VLVisionEncoder(nnx.Module): + """Vision encoder with patch embedding, positional interpolation, and blocks.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: Optional[nnx.Rngs] = None, + ): + self.config = config + self.patch_embed = Qwen3VLVisionPatchEmbed(config=config, rngs=rngs) + + num_pos = config.num_position_embeddings_for_vit + hs = config.hidden_size_for_vit + self.spatial_merge_size = config.spatial_merge_size_for_vit + + self.pos_embed_interpolate = Qwen3VLVisionPosEmbedInterpolate( + num_position_embeddings=num_pos, + hidden_size=hs, + spatial_merge_size=self.spatial_merge_size, + rngs=rngs, + ) + + self.depth = config.num_hidden_layers_for_vit + for i in range(self.depth): + setattr(self, f"blocks_{i}", Qwen3VLVisionBlock(config=config, rngs=rngs)) + + self.deep_idx = tuple(config.deepstack_visual_indexes_for_vit) + for i, _ in enumerate(self.deep_idx): + setattr(self, f"merger_{i}", Qwen3VLVisionPatchMerger(config=config, use_postshuffle_norm=True, rngs=rngs)) + + def __call__( + self, + hidden_states: jax.Array, + video_mask: Optional[jax.Array] = None, + video_grid_thw: Optional[Any] = None, + ): + batch_size, _, num_frames, height, width = hidden_states.shape + num_frames = num_frames // self.config.temporal_patch_size_for_vit + height = height // self.config.patch_size_for_vit + width = width // self.config.patch_size_for_vit + attention_mask = None + if video_mask is not None: + mask_patch_elements = ( + self.config.temporal_patch_size_for_vit * self.config.patch_size_for_vit * self.config.patch_size_for_vit + ) + attention_mask = video_mask.reshape(batch_size, -1, mask_patch_elements).max(axis=-1).astype(jnp.int32) + hidden_states = hidden_states.reshape( + -1, + self.config.num_channels_for_vit, + self.config.temporal_patch_size_for_vit, + self.config.patch_size_for_vit, + self.config.patch_size_for_vit, + ) + + x, _ = self.patch_embed(hidden_states) + x = x.reshape(batch_size, -1, self.config.hidden_size_for_vit) + pos = self.pos_embed_interpolate( + num_frames, + height, + width, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + ) + x = x + pos + valid_grid = video_grid_thw + + h_traj = [] + for i in range(self.depth): + blk = getattr(self, f"blocks_{i}") + x = blk( + x, + num_frames=num_frames, + height=height, + width=width, + attention_mask=attention_mask, + valid_grid=valid_grid, + ) + h_traj.append(x) + + deep_feats = [] + for i, idx in enumerate(self.deep_idx): + merger = getattr(self, f"merger_{i}") + deep_feats.append(merger(h_traj[idx])) + + return x, deep_feats + + +class Qwen3VLVisionProjector(nnx.Module): + """Projection layer converting vision encoder output to model embedding space.""" + + def __init__(self, config: Config, *, rngs: Optional[nnx.Rngs] = None): + self.config = config + self.merger = Qwen3VLVisionPatchMerger(config=config, use_postshuffle_norm=False, rngs=rngs) + + def __call__(self, hidden_states: jax.Array) -> jax.Array: + return self.merger(hidden_states) + + +class Qwen3VLVisionTower(nnx.Module): + """Container module holding Qwen3VLVisionEncoder_0 and Qwen3VLVisionProjector_0 to match checkpoint.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: Optional[nnx.Rngs] = None, + ): + self.config = config + self.Qwen3VLVisionEncoder_0 = Qwen3VLVisionEncoder(config=config, rngs=rngs) + self.Qwen3VLVisionProjector_0 = Qwen3VLVisionProjector(config=config, rngs=rngs) + + def __call__( + self, + input_images: Optional[jax.Array] = None, + input_masks: Optional[jax.Array] = None, + video_grid_thw: Optional[Any] = None, + ): + enc_out, deep_feats = self.Qwen3VLVisionEncoder_0( + input_images, + video_mask=input_masks, + video_grid_thw=video_grid_thw, + ) + proj_out = self.Qwen3VLVisionProjector_0(enc_out) + return proj_out, deep_feats + + +# --------------------------------------------------------------------------- +# Qwen3-VL Text Attention & Decoder +# --------------------------------------------------------------------------- + + +class Qwen3VLAttention(nnx.Module): + """Full-sequence Qwen3 text attention with multi-axis projections and 3D MRoPE.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: nnx.Rngs, + ): + self.config = config + self.num_query_heads = config.num_query_heads + self.num_kv_heads = config.num_kv_heads + self.head_dim = config.head_dim + + self.query = DenseGeneral( + in_features_shape=config.emb_dim, + out_features_shape=(self.num_query_heads, self.head_dim), + use_bias=False, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + rngs=rngs, + ) + self.key = DenseGeneral( + in_features_shape=config.emb_dim, + out_features_shape=(self.num_kv_heads, self.head_dim), + use_bias=False, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + rngs=rngs, + ) + self.value = DenseGeneral( + in_features_shape=config.emb_dim, + out_features_shape=(self.num_kv_heads, self.head_dim), + use_bias=False, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + rngs=rngs, + ) + self.out = DenseGeneral( + in_features_shape=(self.num_query_heads, self.head_dim), + out_features_shape=config.emb_dim, + use_bias=False, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + rngs=rngs, + ) + + norm_kw = { + "epsilon": config.normalization_layer_epsilon, + "dtype": config.dtype, + "param_dtype": config.weight_dtype, + "rngs": rngs, + } + self.query_norm = nnx.RMSNorm(num_features=self.head_dim, **norm_kw) + self.key_norm = nnx.RMSNorm(num_features=self.head_dim, **norm_kw) + + def __call__( + self, + inputs: jax.Array, + decoder_positions: Optional[jax.Array] = None, + decoder_segment_ids: Optional[jax.Array] = None, + deterministic: bool = False, + ) -> jax.Array: + s = inputs.shape[1] + query = self.query(inputs) + key = self.key(inputs) + value = self.value(inputs) + + query = self.query_norm(query) + key = self.key_norm(key) + + if decoder_positions is not None: + mrope_section = getattr(self.config, "mrope_section", (24, 20, 20)) + rope_theta = getattr(self.config, "rope_max_timescale", 5000000.0) + query = apply_mrope(query, decoder_positions, mrope_section=mrope_section, rope_theta=rope_theta) + key = apply_mrope(key, decoder_positions, mrope_section=mrope_section, rope_theta=rope_theta) + + query = query * (self.head_dim**-0.5) + + # Group query heads without materializing repeated keys and values. + groups = self.num_query_heads // self.num_kv_heads + query = query.reshape(inputs.shape[0], s, self.num_kv_heads, groups, self.head_dim) + score_dtype = jnp.float32 if self.config.float32_logits else query.dtype + scores = jnp.einsum("bskgd,btkd->bkgst", query.astype(score_dtype), key.astype(score_dtype)) + causal_mask = jnp.tril(jnp.ones((s, s), dtype=bool))[None, :, :] + if decoder_segment_ids is not None: + causal_mask = causal_mask & (decoder_segment_ids[:, :, None] == decoder_segment_ids[:, None, :]) + scores = jnp.where(causal_mask[:, None, None, :, :], scores, -1e10) + weights = jax.nn.softmax(scores, axis=-1).astype(value.dtype) + attn_out = jnp.einsum("bkgst,btkd->bskgd", weights, value) + attn_out = attn_out.reshape(inputs.shape[0], s, self.num_query_heads, self.head_dim) + return self.out(attn_out) + + +class Qwen3VLDecoderLayer(nnx.Module): + """Qwen3-VL Decoder Layer with Qwen3VLAttention and Qwen3MLP.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: nnx.Rngs, + ): + self.config = config + norm_kw = { + "epsilon": config.normalization_layer_epsilon, + "dtype": config.dtype, + "param_dtype": config.weight_dtype, + "rngs": rngs, + } + self.pre_self_attention_layer_norm = nnx.RMSNorm(num_features=config.emb_dim, **norm_kw) + self.post_self_attention_layer_norm = nnx.RMSNorm(num_features=config.emb_dim, **norm_kw) + self.self_attention = Qwen3VLAttention(config, rngs=rngs) + self.mlp = Qwen3MLP(config, rngs=rngs) + + def __call__( + self, + inputs: jax.Array, + decoder_segment_ids: Optional[jax.Array] = None, + decoder_positions: Optional[jax.Array] = None, + deterministic: bool = False, + **kwargs, + ) -> jax.Array: + normed_attn_in = self.pre_self_attention_layer_norm(inputs) + attn_out = self.self_attention( + normed_attn_in, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=deterministic, + ) + x = (inputs + attn_out).astype(self.config.dtype) + + normed_mlp_in = self.post_self_attention_layer_norm(x) + mlp_out = self.mlp(normed_mlp_in) + output = (x + mlp_out).astype(self.config.dtype) + return output + + +def deepstack_process(hidden_states: jax.Array, mask: jax.Array, visual_embeds: jax.Array) -> jax.Array: + """Adds deepstack visual features to hidden states at visual token positions.""" + visual_token_idx = jnp.cumsum(mask, axis=1) - 1 + batch_idx = jnp.arange(hidden_states.shape[0])[:, jnp.newaxis] + visual_embeds_scattered = visual_embeds[batch_idx, visual_token_idx, :] + mask_expanded = mask[:, :, jnp.newaxis] + return hidden_states + jnp.where(mask_expanded, visual_embeds_scattered, 0.0) + + +class Qwen3VLDecoder(Qwen3Decoder): + """Qwen3-VL decoder stack specializing the dense Qwen3 layers.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: nnx.Rngs, + ): + super().__init__(config, mesh, rngs=rngs, layer_cls=Qwen3VLDecoderLayer) + + def __call__( + self, + shared_embedding: Any, + decoder_input_tokens: jax.Array, + decoder_positions: Optional[jax.Array] = None, + decoder_segment_ids: Optional[jax.Array] = None, + deepstack_visual_embeds: Optional[Sequence[jax.Array]] = None, + multimodal_input: Optional[MultimodalInput] = None, + deterministic: bool = False, + ) -> tuple[jax.Array, jax.Array]: + y = shared_embedding(decoder_input_tokens) + visual_mask = None + if multimodal_input is not None and multimodal_input.image_embeddings is not None: + image_pad_token_id = 151655 + visual_mask = decoder_input_tokens == image_pad_token_id + if multimodal_input.image_masks is not None: + valid_count = multimodal_input.image_masks.sum(axis=1, keepdims=True) + visual_mask = visual_mask & (jnp.cumsum(visual_mask, axis=1) <= valid_count) + y = merge_multimodal_embeddings( + text_embeddings=y, + multimodal_embeddings=multimodal_input.image_embeddings, + mask=visual_mask, + feature_mask=multimodal_input.image_masks, + ) + + for lyr in range(self.num_layers): + layer = getattr(self, f"layers_{lyr}") + y = layer( + y, + decoder_segment_ids=decoder_segment_ids, + decoder_positions=decoder_positions, + deterministic=deterministic, + ) + if deepstack_visual_embeds is not None and lyr < len(deepstack_visual_embeds): + feat = deepstack_visual_embeds[lyr] + if feat is not None and multimodal_input is not None and multimodal_input.image_masks is not None: + feat = _compact_visual_features(feat, multimodal_input.image_masks) + if visual_mask is not None and feat is not None: + y = deepstack_process(y, visual_mask, feat) + + logits = self.apply_output_head(shared_embedding, y, deterministic=deterministic) + return logits, y + + +def _compact_visual_features(features: jax.Array, feature_mask: jax.Array) -> jax.Array: + """Move valid features to the front, preserving their preprocessor order.""" + indices = jnp.argsort(~feature_mask.astype(jnp.bool_), axis=1, stable=True) + return jnp.take_along_axis(features, indices[..., None], axis=1) + + +def merge_multimodal_embeddings( + text_embeddings: jax.Array, + multimodal_embeddings: jax.Array, + mask: jax.Array, + feature_mask: Optional[jax.Array] = None, +) -> jax.Array: + """Merges visual tokens into text embeddings at mask positions.""" + batch_size, _, d_model = text_embeddings.shape + flat_mm = multimodal_embeddings.reshape(batch_size, -1, d_model).astype(text_embeddings.dtype) + if feature_mask is not None: + flat_mm = _compact_visual_features(flat_mm, feature_mask) + mask = mask & (jnp.cumsum(mask, axis=1) <= feature_mask.sum(axis=1, keepdims=True)) + + def _merge_single(t_emb, m_emb, m_mask): + cumsum_mask = jnp.cumsum(m_mask != 0) + mask_bool = (m_mask != 0) & (cumsum_mask <= m_emb.shape[0]) + mask_expanded = mask_bool[:, jnp.newaxis] + mm_idx = jnp.clip(cumsum_mask - 1, 0, m_emb.shape[0] - 1) + mm_aligned = m_emb[mm_idx, :] + return jnp.where(mask_expanded, mm_aligned, t_emb) + + return jax.vmap(_merge_single, in_axes=(0, 0, 0))(text_embeddings, flat_mm, mask) + + +class Qwen3VLModel(nnx.Module): + """Top-level self-contained Qwen3-VL multimodal model in m3 format.""" + + def __init__( + self, + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: nnx.Rngs, + **kwargs, + ): + self.config = config + + # 1. Vision Tower matching checkpoint parameter hierarchy + self.vision_encoder = ( + Qwen3VLVisionTower(config=config, rngs=rngs) if getattr(config, "use_multimodal", False) else None + ) + + # 2. Text LLM Token Embedder & Decoder + self.token_embedder = nnx.Embed( + num_embeddings=config.vocab_size, + features=config.emb_dim, + dtype=config.dtype, + param_dtype=config.weight_dtype, + rngs=rngs, + ) + self.decoder = Qwen3VLDecoder( + config=config, + rngs=rngs, + ) + + def logits_from_hidden_states_for_vocab_tiling( + self, hidden_states: jax.Array, deterministic: bool, model_mode: Optional[str] = None, **kwargs + ) -> jax.Array: + return self.decoder.apply_output_head( + shared_embedding=self.token_embedder, + y=hidden_states, + deterministic=deterministic, + ) + + def __call__( + self, + decoder_input_tokens: jax.Array, + decoder_positions: Optional[jax.Array] = None, + decoder_segment_ids: Optional[jax.Array] = None, + encoder_images: Optional[jax.Array] = None, + encoder_image_masks: Optional[jax.Array] = None, + encoder_video_grid_thw: Optional[jax.Array] = None, + enable_dropout: bool = False, + **kwargs, + ) -> jax.Array: + deterministic = kwargs.get("deterministic", not enable_dropout) + + # Process vision inputs if provided + image_embeddings = None + deepstack_feats = None + + if getattr(self.config, "use_multimodal", False) and encoder_images is not None and self.vision_encoder is not None: + image_embeddings, deepstack_feats = self.vision_encoder( + input_images=encoder_images, + input_masks=encoder_image_masks, + video_grid_thw=encoder_video_grid_thw, + ) + + if decoder_positions is None: + decoder_positions = jnp.broadcast_to( + jnp.arange(decoder_input_tokens.shape[1], dtype=jnp.int32)[None, :], + decoder_input_tokens.shape, + ) + + # Keep encoded modality features together at the decoder fusion boundary. + multimodal_input = None + if image_embeddings is not None: + feature_mask = None + if encoder_image_masks is not None: + patch_elements = self.config.temporal_patch_size_for_vit * self.config.patch_size_for_vit**2 + patch_mask = encoder_image_masks.reshape(image_embeddings.shape[0], -1, patch_elements).any(axis=-1) + feature_mask = patch_mask.reshape(image_embeddings.shape[0], -1, self.config.spatial_merge_size_for_vit**2).any( + axis=-1 + ) + multimodal_input = MultimodalInput( + image_embeddings=image_embeddings, + image_masks=feature_mask, + ) + + # Decoder forward pass + logits, _ = self.decoder( + shared_embedding=self.token_embedder, + decoder_input_tokens=decoder_input_tokens, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=deterministic, + deepstack_visual_embeds=deepstack_feats, + multimodal_input=multimodal_input, + ) + return logits + + +def create_qwen3_vl_model( + config: Config, + mesh: Optional[Mesh] = None, + *, + rngs: Optional[nnx.Rngs] = None, + **kwargs, +) -> Qwen3VLModel: + """Factory function creating a Qwen3VLModel instance.""" + if rngs is None: + rngs = nnx.Rngs(0) + return Qwen3VLModel(config, mesh, rngs=rngs, **kwargs) diff --git a/src/maxtext/m3/models/wan/__init__.py b/src/maxtext/m3/models/wan/__init__.py new file mode 100644 index 0000000000..1b9688c2f0 --- /dev/null +++ b/src/maxtext/m3/models/wan/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Wan 2.1 model family package.""" diff --git a/src/maxtext/m3/models/wan/aot_cache.py b/src/maxtext/m3/models/wan/aot_cache.py new file mode 100644 index 0000000000..e2eb6c33a0 --- /dev/null +++ b/src/maxtext/m3/models/wan/aot_cache.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Per-shape Ahead-of-Time (AOT) compiled executable cache.""" diff --git a/src/maxtext/m3/models/wan/autoencoder_kl_wan.py b/src/maxtext/m3/models/wan/autoencoder_kl_wan.py new file mode 100644 index 0000000000..81929c8cf6 --- /dev/null +++ b/src/maxtext/m3/models/wan/autoencoder_kl_wan.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Wan 3D causal Variational Auto-Encoder (VAE).""" diff --git a/src/maxtext/m3/models/wan/configs/wan2.1-t2v-14b.yml b/src/maxtext/m3/models/wan/configs/wan2.1-t2v-14b.yml new file mode 100644 index 0000000000..cf44f45e3e --- /dev/null +++ b/src/maxtext/m3/models/wan/configs/wan2.1-t2v-14b.yml @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +# Wan 2.1 Text-to-Video 14B Configuration Scaffolding. diff --git a/src/maxtext/m3/models/wan/transformer_wan.py b/src/maxtext/m3/models/wan/transformer_wan.py new file mode 100644 index 0000000000..5a26d02e1d --- /dev/null +++ b/src/maxtext/m3/models/wan/transformer_wan.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Wan Diffusion Transformer (DiT) architecture and blocks.""" diff --git a/src/maxtext/m3/models/wan/utils_wan.py b/src/maxtext/m3/models/wan/utils_wan.py new file mode 100644 index 0000000000..956828f54a --- /dev/null +++ b/src/maxtext/m3/models/wan/utils_wan.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""Wan model loading, tensor reshaping, and device transfer utilities.""" diff --git a/src/maxtext/m3/train/__init__.py b/src/maxtext/m3/train/__init__.py new file mode 100644 index 0000000000..1eded54c62 --- /dev/null +++ b/src/maxtext/m3/train/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""m3's own training stack, built on `nnx.Optimizer`. + +Planned modules: + * `state.py` — optimizer construction. + * `loss.py` — cross entropy plus auxiliary losses. + * `step.py` — train/eval step functions. + * `loop.py` — training loop and infrastructure composition site. + +m3 models are drivable by two things: this loop, and the existing +`maxtext.training_engine.maxtext_engine` (which the RL stack targets). Anything +that makes a model only usable by one of the two drivers belongs here or in a +bridge, not in a model file. +""" diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index 2947bc924c..de8b9765f8 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -31,7 +31,6 @@ from maxtext.layers import linears from maxtext.layers import mhc from maxtext.layers import moe -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.linears import Dropout from maxtext.layers.engram import Engram @@ -401,12 +400,6 @@ def __call__( return self.post_process(layer_output, None, None, kv_cache) -DeepSeekDenseLayerToLinen = nnx_wrappers.to_linen_class( - DeepSeekDenseLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - class DeepSeekMoELayer(DeepSeekGenericLayer): """DeepSeek-style MoE layer with Multi-Head Latent Attention. @@ -622,9 +615,3 @@ def mlp_op(self, x, deterministic, *args, **kwargs): x, intermediate_sharding=self.mlp_intermediate_sharding, out_sharding=self.out_sharding ) return self.with_logical_constraint(mlp_lnx), load_balance_loss, moe_bias_updates - - -DeepSeekMoELayerToLinen = nnx_wrappers.to_linen_class( - DeepSeekMoELayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/deepseek4.py b/src/maxtext/models/deepseek4.py index 6181d684e0..a48231381c 100644 --- a/src/maxtext/models/deepseek4.py +++ b/src/maxtext/models/deepseek4.py @@ -17,7 +17,6 @@ from typing import Optional from flax import nnx -import flax.linen as nn from jax.sharding import Mesh from maxtext.common.common_types import Config, AttentionType @@ -25,9 +24,9 @@ from maxtext.layers import attention_compressed from maxtext.layers import initializers from maxtext.layers import moe -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.models import deepseek +from maxtext.utils.sharding import with_logical_constraint from jax.ad_checkpoint import checkpoint_name @@ -72,7 +71,7 @@ def __init__( # DeepSeek V4 applies Hash Routing to the first `config.first_num_hash_layers` layers. # For the unscannable prefix layers, we can safely determine this using `layer_idx`. - # However, for layers inside `nn.scan` blocks, `layer_idx` is a dynamic JAX tracer + # However, for layers inside scanned blocks, `layer_idx` is a dynamic JAX tracer # and cannot be evaluated as a boolean condition. Since all scannable layers occur # after the hash-routed prefix, the scannable block explicitly passes # `is_hash_routing=False` to safely bypass this check. @@ -176,7 +175,7 @@ class DeepSeek4ScannableBlock(nnx.Module): DeepSeek V4 layers alternate `compress_ratio=128` (HCA) and `compress_ratio=4` (CSA) throughout the middle of the network. This block encapsulates one full `[128, 4]` - cycle so it can be perfectly scanned using JAX `nn.scan`. + cycle so it can be perfectly scanned using `jax.lax.scan`. """ def __init__( @@ -228,7 +227,7 @@ def __call__( kv_cache=None, decoder_input_tokens=None, ): - inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + inputs = with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) inputs = checkpoint_name(inputs, "decoder_layer_input") y = inputs @@ -259,14 +258,3 @@ def __call__( ) return y, None - - -DeepSeek4LayerToLinen = nnx_wrappers.to_linen_class( - DeepSeek4DecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - -DeepSeek4ScannableBlockToLinen = nnx_wrappers.to_linen_class( - DeepSeek4ScannableBlock, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/deepseek_batchsplit_fp8.py b/src/maxtext/models/deepseek_batchsplit_fp8.py index 8e8607914e..0eaa6582c2 100644 --- a/src/maxtext/models/deepseek_batchsplit_fp8.py +++ b/src/maxtext/models/deepseek_batchsplit_fp8.py @@ -20,13 +20,14 @@ import math from typing import Any, Sequence -from flax import linen as nn import jax import jax.numpy as jnp from maxtext.kernels import megablox, sort_activations from maxtext.layers import attention_op from maxtext.layers import moe as moe_lib from maxtext.layers import quantizations +from maxtext.utils import max_utils +from maxtext.utils.sharding import logical_to_spec import qwix.pallas as qpl import tokamax @@ -976,10 +977,25 @@ def gmm( wi_gather_axes = [] wo_gather_axes = [] + if max_utils.is_eval(config): + wi_tile_fwd_batch_seq = config.eval_wi_tile_fwd_batch_seq + wi_tile_fwd_embed_dim = config.eval_wi_tile_fwd_embed_dim + wi_tile_fwd_mlp_dim = config.eval_wi_tile_fwd_mlp_dim + wo_tile_fwd_batch_seq = config.eval_wo_tile_fwd_batch_seq + wo_tile_fwd_embed_dim = config.eval_wo_tile_fwd_embed_dim + wo_tile_fwd_mlp_dim = config.eval_wo_tile_fwd_mlp_dim + else: + wi_tile_fwd_batch_seq = config.wi_tile_fwd_batch_seq + wi_tile_fwd_embed_dim = config.wi_tile_fwd_embed_dim + wi_tile_fwd_mlp_dim = config.wi_tile_fwd_mlp_dim + wo_tile_fwd_batch_seq = config.wo_tile_fwd_batch_seq + wo_tile_fwd_embed_dim = config.wo_tile_fwd_embed_dim + wo_tile_fwd_mlp_dim = config.wo_tile_fwd_mlp_dim + wi_tile_size = ( - config.wi_tile_fwd_batch_seq, # m (LHS batch) - config.wi_tile_fwd_embed_dim, # k (contracting) - config.wi_tile_fwd_mlp_dim, # n (RHS batch) + wi_tile_fwd_batch_seq, # m (LHS batch) + wi_tile_fwd_embed_dim, # k (contracting) + wi_tile_fwd_mlp_dim, # n (RHS batch) config.wi_tile_dlhs_batch_seq, # m (LHS batch) config.wi_tile_dlhs_mlp_dim, # k (contracting) config.wi_tile_dlhs_embed_dim, # n (RHS batch) @@ -989,9 +1005,9 @@ def gmm( ) wo_tile_size = ( - config.wo_tile_fwd_batch_seq, # m (LHS batch) - config.wo_tile_fwd_mlp_dim, # k (contracting) - config.wo_tile_fwd_embed_dim, # n (RHS batch) + wo_tile_fwd_batch_seq, # m (LHS batch) + wo_tile_fwd_mlp_dim, # k (contracting) + wo_tile_fwd_embed_dim, # n (RHS batch) config.wo_tile_dlhs_batch_seq, # m (LHS batch) config.wo_tile_dlhs_embed_dim, # k (contracting) config.wo_tile_dlhs_mlp_dim, # n (RHS) @@ -1002,8 +1018,8 @@ def gmm( if config.quantization and config.use_qwix_quantization: gating_pspec, linear_pspec = moe_lib.get_batchsplit_init_kernel_axes() - w0_pspec = nn.logical_to_mesh_axes(gating_pspec) - wo_pspec = nn.logical_to_mesh_axes(linear_pspec) + w0_pspec = logical_to_spec(gating_pspec) + wo_pspec = logical_to_spec(linear_pspec) ignored_axes = ("expert", "tensor") def get_active_sharding_axes(pspec_dim_axes, tensor_dim_index): @@ -1154,8 +1170,8 @@ def process_activations( ) if config.quantization and config.use_qwix_quantization: gating_pspec, linear_pspec = moe_lib.get_batchsplit_init_kernel_axes() - gating_pspec = nn.logical_to_mesh_axes(gating_pspec) - linear_pspec = nn.logical_to_mesh_axes(linear_pspec) + gating_pspec = logical_to_spec(gating_pspec) + linear_pspec = logical_to_spec(linear_pspec) else: gating_pspec = jax.sharding.PartitionSpec(None, None, expert_axis_name) linear_pspec = jax.sharding.PartitionSpec(None, expert_axis_name, None) diff --git a/src/maxtext/models/envy.py b/src/maxtext/models/envy.py index a0cde31910..952fe70525 100644 --- a/src/maxtext/models/envy.py +++ b/src/maxtext/models/envy.py @@ -15,7 +15,6 @@ """Envy MoE decoder layer definition.""" # pylint: disable=arguments-differ, disable=no-name-in-module, missing-function-docstring -from flax import linen as nn from flax import nnx import jax from jax.ad_checkpoint import checkpoint_name @@ -23,7 +22,6 @@ from jax.sharding import Mesh from maxtext.common.common_types import Config, MODEL_MODE_PREFILL from maxtext.layers import initializers -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout @@ -32,6 +30,7 @@ from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils +from maxtext.utils.sharding import with_logical_constraint class EnvyDecoderLayer(nnx.Module): @@ -190,11 +189,11 @@ def __call__( is_scan_carry = True elif isinstance(inputs, tuple): inputs = inputs[0] - inputs = nn.with_logical_constraint(inputs, self.activation_axis_names) + inputs = with_logical_constraint(inputs, self.activation_axis_names) inputs = checkpoint_name(inputs, "decoder_layer_input") lnx = self.pre_self_attention_layer_norm(inputs) - lnx = nn.with_logical_constraint(lnx, self.activation_axis_names) + lnx = with_logical_constraint(lnx, self.activation_axis_names) # Self-attention block attention_lnx, kv_cache = self.self_attention( @@ -209,12 +208,12 @@ def __call__( kv_cache=kv_cache, attention_metadata=attention_metadata, ) - attention_lnx = nn.with_logical_constraint(attention_lnx, self.activation_axis_names) + attention_lnx = with_logical_constraint(attention_lnx, self.activation_axis_names) intermediate_inputs = inputs + attention_lnx # Fully Connected / MLP block hidden_states = self.post_self_attention_layer_norm(intermediate_inputs) - hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names) + hidden_states = with_logical_constraint(hidden_states, self.activation_axis_names) load_balance_loss = None if self.is_moe_layer: @@ -223,11 +222,11 @@ def __call__( mlp_lnx = self.mlp(hidden_states, deterministic=deterministic) if self.config.use_post_ffw_norm: mlp_lnx = self.post_ffw_norm(mlp_lnx) - mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names) + mlp_lnx = with_logical_constraint(mlp_lnx, self.activation_axis_names) layer_output = mlp_lnx + intermediate_inputs layer_output = self.dropout(layer_output, deterministic=deterministic) - layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names) + layer_output = with_logical_constraint(layer_output, self.activation_axis_names) if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) @@ -256,12 +255,6 @@ def update_cache(cache, val): return layer_output, kv_cache -EnvyDecoderLayerToLinen = nnx_wrappers.to_linen_class( - EnvyDecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - class EnvyScannableBlock(nnx.Module): """A repeatable block of inhomogeneous layers for Envy.""" @@ -312,7 +305,7 @@ def __call__( activation_axis_names = ("activation_batch", "prefill_activation_norm_length", "activation_embed") else: activation_axis_names = ("activation_batch", "activation_norm_length", "activation_embed") - inputs = nn.with_logical_constraint(inputs, activation_axis_names) + inputs = with_logical_constraint(inputs, activation_axis_names) inputs = checkpoint_name(inputs, "decoder_layer_input") y = inputs @@ -334,9 +327,3 @@ def __call__( return y, None else: return y - - -EnvyScannableBlockToLinen = nnx_wrappers.to_linen_class( - EnvyScannableBlock, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/gemma.py b/src/maxtext/models/gemma.py index 73489709d7..038a059ab9 100644 --- a/src/maxtext/models/gemma.py +++ b/src/maxtext/models/gemma.py @@ -23,8 +23,6 @@ import jax.numpy as jnp from maxtext.common.common_types import Config -from maxtext.layers import initializers -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock @@ -126,7 +124,6 @@ def __init__( mesh=self.mesh, shard_mode=config.shard_mode, debug_sharding=config.debug_sharding, - extra_stack_level=1, ) def __call__( @@ -199,9 +196,3 @@ def __call__( return layer_output, None else: return layer_output, kv_cache - - -GemmaDecoderLayerToLinen = nnx_wrappers.to_linen_class( - GemmaDecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/gemma2.py b/src/maxtext/models/gemma2.py index 6d0189620a..78c6c84121 100644 --- a/src/maxtext/models/gemma2.py +++ b/src/maxtext/models/gemma2.py @@ -24,8 +24,6 @@ from maxtext.common.common_types import MODEL_MODE_PREFILL, Config from maxtext.layers import attentions -from maxtext.layers import initializers -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock @@ -227,7 +225,6 @@ def __init__( mesh=self.mesh, shard_mode=config.shard_mode, debug_sharding=config.debug_sharding, - extra_stack_level=1, ) def __call__( @@ -350,9 +347,3 @@ def __call__( return layer_output, None else: return layer_output, kv_cache - - -Gemma2DecoderLayerToLinen = nnx_wrappers.to_linen_class( - Gemma2DecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/gemma3.py b/src/maxtext/models/gemma3.py index 106a48e088..04c94b7883 100644 --- a/src/maxtext/models/gemma3.py +++ b/src/maxtext/models/gemma3.py @@ -25,8 +25,6 @@ from maxtext.common.common_types import Config, AttentionType, MODEL_MODE_PREFILL from maxtext.layers import quantizations -from maxtext.layers import nnx_wrappers -from maxtext.layers import initializers from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral, MlpBlock, Dropout from maxtext.layers.normalizations import RMSNorm @@ -188,7 +186,6 @@ def __init__( mesh=self.mesh, shard_mode=config.shard_mode, debug_sharding=config.debug_sharding, - extra_stack_level=1, ) def __call__( @@ -286,12 +283,6 @@ def update_cache(cache, val): return layer_output, kv_cache -Gemma3DecoderLayerToLinen = nnx_wrappers.to_linen_class( - Gemma3DecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - class Gemma3ScannableBlock(nnx.Module): """A repeatable block of Gemma3 decoder layers.""" @@ -339,7 +330,6 @@ def __init__( mesh=self.mesh, shard_mode=config.shard_mode, debug_sharding=config.debug_sharding, - extra_stack_level=1, ) def __call__( @@ -388,12 +378,6 @@ def __call__( return y -Gemma3ScannableBlockToLinen = nnx_wrappers.to_linen_class( - Gemma3ScannableBlock, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - def _posemb_sincos_2d( h: int, w: int, diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py index 82f89d88c2..1fd953ed1d 100644 --- a/src/maxtext/models/gemma4.py +++ b/src/maxtext/models/gemma4.py @@ -20,17 +20,17 @@ from jax.sharding import Mesh import jax.numpy as jnp -from flax import linen as nn from flax import nnx from typing import Optional, Any from maxtext.common.common_types import Config, AttentionType, MODEL_MODE_PREFILL from maxtext.layers import initializers from maxtext.layers import moe -from maxtext.layers import nnx_scan, nnx_wrappers +from maxtext.layers import nnx_scan from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.linears import MlpBlock +from maxtext.utils.sharding import with_logical_constraint import jax.sharding from maxtext.layers.normalizations import RMSNorm @@ -337,11 +337,11 @@ def __call__( is_scan_carry = True elif isinstance(inputs, tuple): inputs = inputs[0] - inputs = nn.with_logical_constraint(inputs, self.activation_axis_names) + inputs = with_logical_constraint(inputs, self.activation_axis_names) inputs = checkpoint_name(inputs, "decoder_layer_input") lnx = self.pre_self_attention_norm(inputs) - lnx = nn.with_logical_constraint(lnx, self.activation_axis_names) + lnx = with_logical_constraint(lnx, self.activation_axis_names) # Gemma4 only applies bidirectional attention in sliding (local) layers, # not in full (global) attention layers. @@ -362,7 +362,7 @@ def __call__( ) if cfg.use_post_attn_norm: attention_lnx = self.post_self_attention_norm(attention_lnx) - attention_lnx = nn.with_logical_constraint(attention_lnx, self.activation_axis_names) + attention_lnx = with_logical_constraint(attention_lnx, self.activation_axis_names) attention_lnx += inputs residual = attention_lnx @@ -383,13 +383,13 @@ def __call__( if cfg.use_post_ffw_norm: mlp_lnx = self.post_ffw_norm(mlp_lnx) - mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names) + mlp_lnx = with_logical_constraint(mlp_lnx, self.activation_axis_names) next_layer_addition = mlp_lnx + residual layer_output = next_layer_addition layer_output = layer_output * jnp.asarray(self.layer_scalar.value, cfg.dtype) - layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names) + layer_output = with_logical_constraint(layer_output, self.activation_axis_names) if getattr(cfg, "record_internal_nn_metrics", False): self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output)) @@ -413,12 +413,6 @@ def update_cache(cache, val): return layer_output, kv_cache -Gemma4DecoderLayerToLinen = nnx_wrappers.to_linen_class( - Gemma4DecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - class Gemma4ScannableBlock(nnx.Module): """A repeatable block of Gemma4 decoder layers, scanning local layers.""" @@ -667,7 +661,7 @@ def __call__( attention_metadata=None, ): cfg = self.config - inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + inputs = with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) inputs = checkpoint_name(inputs, "decoder_layer_input") # Arguments shared by every layer in the block. model_mode differentiates @@ -699,9 +693,3 @@ def __call__( if cfg.scan_layers: return y, None return y - - -Gemma4ScannableBlockToLinen = nnx_wrappers.to_linen_class( - Gemma4ScannableBlock, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/gemma4_small.py b/src/maxtext/models/gemma4_small.py index ca33470bf2..3822327db0 100644 --- a/src/maxtext/models/gemma4_small.py +++ b/src/maxtext/models/gemma4_small.py @@ -19,18 +19,16 @@ from jax.sharding import Mesh import jax.numpy as jnp -from flax import linen as nn from flax import nnx from maxtext.common.common_types import Config, AttentionType, MODEL_MODE_PREFILL -from maxtext.layers import initializers -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral, MlpBlock from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils +from maxtext.utils.sharding import with_logical_constraint # E2B repeats 4 sliding + 1 global (period 5); E4B repeats 5 sliding + 1 global @@ -162,7 +160,7 @@ def __init__( vocab_ple = config.vocab_size_per_layer_input self.embed_tokens_per_layer = nnx.Param( - nn.initializers.normal(stddev=ple_dim**-0.5)( + jax.nn.initializers.normal(stddev=ple_dim**-0.5)( rngs.params(), (vocab_ple, num_layers * ple_dim), config.weight_dtype, @@ -218,12 +216,6 @@ def __call__(self, input_ids: jax.Array, inputs_embeds: jax.Array) -> jax.Array: return out.astype(cfg.dtype) -PLEToLinen = nnx_wrappers.to_linen_class( - Gemma4SmallPLE, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - class Gemma4SmallDecoderLayer(nnx.Module): """Transformer decoder layer for Gemma 4 small (E2B / E4B).""" @@ -426,7 +418,7 @@ def __call__( if isinstance(inputs, tuple): inputs = inputs[0] - inputs = nn.with_logical_constraint(inputs, self.activation_axis_names) + inputs = with_logical_constraint(inputs, self.activation_axis_names) inputs = checkpoint_name(inputs, "decoder_layer_input") # Bidirectional image-token mask is only meaningful in local-sliding @@ -436,7 +428,7 @@ def __call__( residual = inputs h = self.pre_self_attention_norm(inputs) - h = nn.with_logical_constraint(h, self.activation_axis_names) + h = with_logical_constraint(h, self.activation_axis_names) attn_out, kv_cache = self.self_attention( h, @@ -452,14 +444,14 @@ def __call__( shared_value=shared_value, ) attn_out = self.post_self_attention_norm(attn_out) - attn_out = nn.with_logical_constraint(attn_out, self.activation_axis_names) + attn_out = with_logical_constraint(attn_out, self.activation_axis_names) h = residual + attn_out residual = h mlp_in = self.pre_ffw_norm(h) mlp_out = self.mlp(mlp_in, deterministic=deterministic) mlp_out = self.post_ffw_norm(mlp_out) - mlp_out = nn.with_logical_constraint(mlp_out, self.activation_axis_names) + mlp_out = with_logical_constraint(mlp_out, self.activation_axis_names) h = residual + mlp_out if self.per_layer_input_gate is not None and per_layer_input is not None: @@ -469,16 +461,10 @@ def __call__( gated = gate * per_layer_input.astype(cfg.dtype) proj = self.per_layer_projection(gated) proj = self.post_per_layer_input_norm(proj) - proj = nn.with_logical_constraint(proj, self.activation_axis_names) + proj = with_logical_constraint(proj, self.activation_axis_names) h = residual + proj h = h * jnp.asarray(self.layer_scalar.value, cfg.dtype) - h = nn.with_logical_constraint(h, self.activation_axis_names) + h = with_logical_constraint(h, self.activation_axis_names) return h, kv_cache - - -Gemma4SmallDecoderLayerToLinen = nnx_wrappers.to_linen_class( - Gemma4SmallDecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/gpt3.py b/src/maxtext/models/gpt3.py index e711a4f6e2..7dc1fa6b4d 100644 --- a/src/maxtext/models/gpt3.py +++ b/src/maxtext/models/gpt3.py @@ -24,12 +24,11 @@ from jax.ad_checkpoint import checkpoint_name from jax.sharding import Mesh, NamedSharding -from flax import linen as nn from flax import nnx from maxtext.common.common_types import Config, DType, AxisNames, BATCH, LENGTH, EMBED, HEAD, D_KV, Array, MODEL_MODE_TRAIN from maxtext.inference import kvcache -from maxtext.layers import initializers, nnx_wrappers +from maxtext.layers import initializers from maxtext.layers.linears import DenseGeneral, MlpBlock, canonicalize_tuple, normalize_axes from maxtext.layers import quantizations from maxtext.layers import linears @@ -38,6 +37,7 @@ from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_logging from maxtext.utils import max_utils +from maxtext.utils.sharding import with_logical_constraint # ----------------------------------------- # The Normalization Layer specific for GPT3 @@ -54,7 +54,7 @@ def __init__( dtype: Any = jnp.float32, weight_dtype: Any = jnp.float32, kernel_axes: tuple[None | str, ...] = (), - scale_init: Initializer = nn.initializers.zeros, + scale_init: Initializer = jax.nn.initializers.zeros, use_bias: bool = True, reductions_in_fp32: bool = False, parameter_memory_host_offload: bool = False, @@ -310,7 +310,7 @@ def __call__( kv_cache: Array | None = None, attention_metadata: dict[str, Any] | None = None, ): - inputs_q = nn.with_logical_constraint(inputs_q, self.input_axis_names) + inputs_q = with_logical_constraint(inputs_q, self.input_axis_names) if self.fused_qkv: query, key, value = self.qkv_projection(self.qkv_proj, inputs_q) else: @@ -322,11 +322,11 @@ def __call__( query /= depth_scaling # annotate with sharding constraint. - query = nn.with_logical_constraint(query, self.query_axis_names) + query = with_logical_constraint(query, self.query_axis_names) query = checkpoint_name(query, "query_proj") - key = nn.with_logical_constraint(key, self.key_axis_names) + key = with_logical_constraint(key, self.key_axis_names) key = checkpoint_name(key, "key_proj") - value = nn.with_logical_constraint(value, self.value_axis_names) + value = with_logical_constraint(value, self.value_axis_names) value = checkpoint_name(value, "value_proj") cached_values = [None, None] @@ -335,7 +335,7 @@ def __call__( out = self.attention_op(query, key, value, decoder_segment_ids, None, model_mode, cached_values) - out = nn.with_logical_constraint(out, self.out_axis_names) + out = with_logical_constraint(out, self.out_axis_names) # apply output projection, output dim is set to the input dim. out = self.projection(self.out, out) @@ -437,11 +437,11 @@ def __call__( if isinstance(inputs, tuple): inputs = inputs[0] - inputs = nn.with_logical_constraint(inputs, self.activation_axis_names) + inputs = with_logical_constraint(inputs, self.activation_axis_names) inputs = checkpoint_name(inputs, "decoder_layer_input") lnx = self.pre_self_attention_norm(inputs) - lnx = nn.with_logical_constraint(lnx, self.activation_axis_names) + lnx = with_logical_constraint(lnx, self.activation_axis_names) # Self-attention block assert ( @@ -458,15 +458,15 @@ def __call__( attention_metadata=attention_metadata, ) - attention_lnx = nn.with_logical_constraint(attention_lnx, self.activation_axis_names) + attention_lnx = with_logical_constraint(attention_lnx, self.activation_axis_names) attention_lnx += inputs # MLP block. mlp_lnx = self.mlp(attention_lnx, deterministic=deterministic) - mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names) + mlp_lnx = with_logical_constraint(mlp_lnx, self.activation_axis_names) layer_output = attention_lnx + mlp_lnx layer_output = self.dropout(layer_output, deterministic=deterministic) - layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names) + layer_output = with_logical_constraint(layer_output, self.activation_axis_names) if getattr(self.config, "record_internal_nn_metrics", False): self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output)) @@ -481,9 +481,3 @@ def __call__( return layer_output, None else: return layer_output, kv_cache - - -Gpt3DecoderLayerToLinen = nnx_wrappers.to_linen_class( - Gpt3DecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/gpt_oss.py b/src/maxtext/models/gpt_oss.py index 753dfb9568..69a2e59bc3 100644 --- a/src/maxtext/models/gpt_oss.py +++ b/src/maxtext/models/gpt_oss.py @@ -20,7 +20,6 @@ from typing import Optional -from flax import linen as nn from flax import nnx from jax.ad_checkpoint import checkpoint_name import jax @@ -31,12 +30,12 @@ from maxtext.layers import initializers from maxtext.layers import linears from maxtext.layers import moe -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils +from maxtext.utils.sharding import with_logical_constraint # ----------------------------------------- # The Decoder Layer for GPT OSS models @@ -157,11 +156,11 @@ def __call__( elif isinstance(inputs, tuple): inputs = inputs[0] - inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + inputs = with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) inputs = checkpoint_name(inputs, "decoder_layer_input") lnx = self.pre_self_attention_layer_norm(inputs) - lnx = nn.with_logical_constraint(lnx, ("activation_batch", "activation_norm_length", "activation_embed")) + lnx = with_logical_constraint(lnx, ("activation_batch", "activation_norm_length", "activation_embed")) attention_lnx, kv_cache = self.GptOssAttention( lnx, @@ -174,25 +173,25 @@ def __call__( attention_metadata=attention_metadata, ) - attention_lnx = nn.with_logical_constraint( + attention_lnx = with_logical_constraint( attention_lnx, ("activation_batch", "activation_norm_length", "activation_embed") ) intermediate_inputs = inputs + attention_lnx # Fully Connected hidden_states = self.post_self_attention_layer_norm(intermediate_inputs) - hidden_states = nn.with_logical_constraint( + hidden_states = with_logical_constraint( hidden_states, ("activation_batch", "activation_norm_length", "activation_embed") ) load_balance_loss = None mlp_lnx, load_balance_loss, _ = self.GptOssMlp(hidden_states) - mlp_lnx = nn.with_logical_constraint(mlp_lnx, ("activation_batch", "activation_norm_length", "activation_embed")) + mlp_lnx = with_logical_constraint(mlp_lnx, ("activation_batch", "activation_norm_length", "activation_embed")) layer_output = mlp_lnx + intermediate_inputs layer_output = self.dropout(layer_output, deterministic=deterministic) - layer_output = nn.with_logical_constraint( + layer_output = with_logical_constraint( layer_output, ("activation_batch", "activation_norm_length", "activation_embed"), ) @@ -224,18 +223,12 @@ def update_cache(cache, val): return layer_output, kv_cache -GptOssDecoderLayerToLinen = nnx_wrappers.to_linen_class( - GptOssDecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - class GptOssScannableBlock(nnx.Module): """A repeatable block of GPT OSS decoder layers. This block applies multiple decoder layers sequentially, using the attention pattern defined by GPT_OSS_ATTENTION_PATTERN. It's designed to be - used with `nn.scan` for efficient compilation. + used with `jax.lax.scan` for efficient compilation. Attributes: config: Config, MaxText model config @@ -283,7 +276,7 @@ def __call__( ): cfg = self.config - inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + inputs = with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) inputs = checkpoint_name(inputs, "decoder_layer_input") y = inputs for layer_id in range(cfg.inhomogeneous_layer_cycle_interval): @@ -301,9 +294,3 @@ def __call__( attention_metadata=attention_metadata, ) return y, kv_cache - - -GptOssScannableBlockToLinen = nnx_wrappers.to_linen_class( - GptOssScannableBlock, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/llama2.py b/src/maxtext/models/llama2.py index bdc9934ec8..1c94406763 100644 --- a/src/maxtext/models/llama2.py +++ b/src/maxtext/models/llama2.py @@ -24,8 +24,6 @@ from jax.sharding import Mesh from maxtext.common.common_types import Config from maxtext.common.common_types import MODEL_MODE_PREFILL -from maxtext.layers import initializers -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock @@ -33,7 +31,6 @@ from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils from maxtext.utils.sharding import create_sharding, maybe_shard_with_logical -from maxtext.layers.learn_to_init_layer import apply_lti_modification # ----------------------------------------- # The Decoder Layer specific for Llama2 @@ -134,7 +131,6 @@ def __init__( mesh=self.mesh, shard_mode=config.shard_mode, debug_sharding=config.debug_sharding, - extra_stack_level=1, ) def __call__( @@ -224,10 +220,3 @@ def update_cache(cache, val): return layer_output, None else: return layer_output, kv_cache - - -LlamaDecoderLayerToLinen = nnx_wrappers.to_linen_class( - LlamaDecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, - nnx_module_augment_fn=apply_lti_modification, -) diff --git a/src/maxtext/models/llama4.py b/src/maxtext/models/llama4.py index ebbc98d764..e2304ae0f3 100644 --- a/src/maxtext/models/llama4.py +++ b/src/maxtext/models/llama4.py @@ -17,7 +17,6 @@ import math -from flax import linen as nn from flax import nnx import jax from jax import lax @@ -28,7 +27,6 @@ from maxtext.common.common_types import MODEL_MODE_PREFILL from maxtext.layers import initializers from maxtext.layers import linears -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout @@ -37,6 +35,7 @@ from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils +from maxtext.utils.sharding import with_logical_constraint #### Multi modal model implementation @@ -448,11 +447,11 @@ def __call__( is_scan_carry = True elif isinstance(inputs, tuple): inputs = inputs[0] - inputs = nn.with_logical_constraint(inputs, self.activation_axis_names) + inputs = with_logical_constraint(inputs, self.activation_axis_names) inputs = checkpoint_name(inputs, "decoder_layer_input") lnx = self.pre_self_attention_layer_norm(inputs) - lnx = nn.with_logical_constraint(lnx, self.activation_axis_names) + lnx = with_logical_constraint(lnx, self.activation_axis_names) # Self-attention block attention_lnx, kv_cache = self.self_attention( @@ -467,23 +466,23 @@ def __call__( kv_cache=kv_cache, attention_metadata=attention_metadata, ) - attention_lnx = nn.with_logical_constraint(attention_lnx, self.activation_axis_names) + attention_lnx = with_logical_constraint(attention_lnx, self.activation_axis_names) intermediate_inputs = inputs + attention_lnx # Fully Connected hidden_states = self.post_self_attention_layer_norm(intermediate_inputs) - hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names) + hidden_states = with_logical_constraint(hidden_states, self.activation_axis_names) load_balance_loss = None if self.is_moe_layer: mlp_lnx, load_balance_loss, _ = self.moe_block(hidden_states) else: mlp_lnx = self.mlp(hidden_states, deterministic=deterministic) - mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names) + mlp_lnx = with_logical_constraint(mlp_lnx, self.activation_axis_names) layer_output = mlp_lnx + intermediate_inputs layer_output = self.dropout(layer_output, deterministic=deterministic) - layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names) + layer_output = with_logical_constraint(layer_output, self.activation_axis_names) if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) @@ -512,12 +511,6 @@ def update_cache(cache, val): return layer_output, kv_cache -Llama4DecoderLayerToLinen = nnx_wrappers.to_linen_class( - Llama4DecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - class Llama4ScannableBlock(nnx.Module): """A repeatable block given nope_layer_interval and interleave_moe_layer_step.""" @@ -580,7 +573,7 @@ def __call__( cfg = self.config - inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + inputs = with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) inputs = checkpoint_name(inputs, "decoder_layer_input") y = inputs for layer_id in range(cfg.inhomogeneous_layer_cycle_interval): @@ -603,12 +596,6 @@ def __call__( return y -Llama4ScannableBlockToLinen = nnx_wrappers.to_linen_class( - Llama4ScannableBlock, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - class Llama4VisionEncoderLayer(nnx.Module): """Transformer encoder layer for Llama4 vision model.""" diff --git a/src/maxtext/models/mistral.py b/src/maxtext/models/mistral.py index 4dbfdb3256..1711c08ec2 100644 --- a/src/maxtext/models/mistral.py +++ b/src/maxtext/models/mistral.py @@ -24,7 +24,6 @@ import jax.numpy as jnp from jax.sharding import Mesh from maxtext.common.common_types import Config -from maxtext.layers import initializers, nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock @@ -70,7 +69,6 @@ def __init__( mesh=mesh, shard_mode=config.shard_mode, debug_sharding=config.debug_sharding, - extra_stack_level=1, ) self.pre_self_attention_layer_norm = RMSNorm( @@ -221,9 +219,3 @@ def update_cache(cache, val): return layer_output, None else: return layer_output, kv_cache - - -MistralDecoderLayerToLinen = nnx_wrappers.to_linen_class( - MistralDecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/mixtral.py b/src/maxtext/models/mixtral.py index 7b929f15f7..d435cd11b8 100644 --- a/src/maxtext/models/mixtral.py +++ b/src/maxtext/models/mixtral.py @@ -24,7 +24,7 @@ import jax.numpy as jnp from jax.sharding import Mesh from maxtext.common.common_types import Config -from maxtext.layers import initializers, nnx_wrappers +from maxtext.layers import initializers from maxtext.layers import moe from maxtext.layers import quantizations from maxtext.layers.attentions import Attention @@ -70,7 +70,6 @@ def __init__( mesh=mesh, shard_mode=config.shard_mode, debug_sharding=config.debug_sharding, - extra_stack_level=1, ) self.pre_self_attention_layer_norm = RMSNorm( @@ -209,9 +208,3 @@ def __call__( return layer_output, None else: return layer_output, kv_cache - - -MixtralDecoderLayerToLinen = nnx_wrappers.to_linen_class( - MixtralDecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index 41324ac5a3..0311431168 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -22,7 +22,7 @@ import jax.numpy as jnp from jax.sharding import Mesh -from flax import linen as nn +from flax.core import FrozenDict from flax import nnx from maxtext.common.common_types import Config, DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_TRAIN, MultimodalInput @@ -34,6 +34,7 @@ from maxtext.layers.multi_token_prediction import MultiTokenPredictionBlock from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.multimodal import processor as mm_processor +from maxtext.utils import max_logging # ------------------------------------------------------------------------------ # The network: Transformer Definitions @@ -48,12 +49,12 @@ def transformer_as_linen( *, name: str | None = None, ) -> nnx_wrappers.ToLinen: - """Constructs an NNX Transformer wrapped as a Linen module. + """Constructs an NNX Transformer wrapped as a Linen module. Deprecated. Returns a `TransformerLinen` that wraps the NNX-style Transformer so it can be - driven through the Linen init/apply API (checkpoint conversion, AOT compile, - the inference engine). Pure-NNX call sites build `Transformer` directly via - `model_creation_utils.from_config`. + driven through the Linen init/apply API. Nothing in MaxText needs this anymore; + it is kept for callers of `model_creation_utils.from_config` without `rngs`. + Pass `rngs` to `from_config` to get the NNX `Transformer` instead. Args: config (Config): The configuration object specifying model hyperparameters and options. @@ -66,10 +67,14 @@ def transformer_as_linen( Returns: nnx_wrappers.ToLinen: An NNX Transformer wrapped as a Linen module. """ + max_logging.log( + "WARNING: models.transformer_as_linen, and model_creation_utils.from_config without rngs, are deprecated." + " Please pass rngs=nnx.Rngs(...) to from_config to get the NNX Transformer." + ) return TransformerLinen( Transformer, args=(), - kwargs=nn.FrozenDict( + kwargs=FrozenDict( { "mesh": mesh, "config": config, @@ -129,7 +134,7 @@ def __init__( num_features=cfg.emb_dim, dtype=cfg.dtype, attend_dtype=jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype, # for logit training stability - embedding_init=nn.initializers.normal(stddev=1.0), + embedding_init=jax.nn.initializers.normal(stddev=1.0), config=cfg, rngs=rngs, ) @@ -283,6 +288,18 @@ def __call__( bidirectional_mask_video=bidirectional_mask_video, ) + # With mtp_reuse_input_embedding, look up the token embeddings once here; the MTP block shifts them instead of + # looking up the shifted tokens again, and the embedding table gets a single gradient. + mtp_token_embeddings = None + if ( + getattr(self.config, "mtp_reuse_input_embedding", False) + and getattr(self.config, "mtp_num_layers", 0) > 0 + and decoder_input_embeddings is None + and multimodal_input is None + ): + decoder_input_embeddings = self.decoder.embed_tokens(self.token_embedder, decoder_input_tokens, model_mode) + mtp_token_embeddings = decoder_input_embeddings + res = self.decoder( shared_embedding=self.token_embedder, decoder_input_tokens=decoder_input_tokens, @@ -336,6 +353,7 @@ def __call__( decoder_segment_ids=decoder_segment_ids, deterministic=not enable_dropout, model_mode=model_mode, + main_token_embeddings=mtp_token_embeddings, ) if self.config.attention in ("vllm_rpa", "vllm_batched_rpa", "vllm_batched_rpa_long_ctx"): diff --git a/src/maxtext/models/olmo3.py b/src/maxtext/models/olmo3.py index 4b6bc943e2..63587db1cf 100644 --- a/src/maxtext/models/olmo3.py +++ b/src/maxtext/models/olmo3.py @@ -21,7 +21,6 @@ from typing import Optional -from flax import linen as nn from flax import nnx from jax.ad_checkpoint import checkpoint_name import jax @@ -29,15 +28,14 @@ from jax.sharding import Mesh from maxtext.common.common_types import AttentionType, Config from maxtext.layers import attentions -from maxtext.layers import initializers from maxtext.layers import linears -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.linears import MlpBlock from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils +from maxtext.utils.sharding import with_logical_constraint # ----------------------------------------- @@ -169,7 +167,7 @@ def __call__( elif isinstance(inputs, tuple): inputs = inputs[0] - inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + inputs = with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) inputs = checkpoint_name(inputs, "decoder_layer_input") attention_lnx, kv_cache = self.attention( @@ -183,13 +181,13 @@ def __call__( attention_metadata=attention_metadata, ) - attention_lnx = nn.with_logical_constraint( + attention_lnx = with_logical_constraint( attention_lnx, ("activation_batch", "activation_norm_length", "activation_embed") ) # Normalize stream before addition attention_lnx = self.post_self_attention_layer_norm(attention_lnx) - attention_lnx = nn.with_logical_constraint( + attention_lnx = with_logical_constraint( attention_lnx, ("activation_batch", "activation_norm_length", "activation_embed") ) @@ -197,16 +195,16 @@ def __call__( # Fully Connected mlp_lnx = self.mlp(intermediate_inputs) - mlp_lnx = nn.with_logical_constraint(mlp_lnx, ("activation_batch", "activation_norm_length", "activation_embed")) + mlp_lnx = with_logical_constraint(mlp_lnx, ("activation_batch", "activation_norm_length", "activation_embed")) # Normalize stream before addition mlp_lnx = self.post_mlp_layer_norm(mlp_lnx) - mlp_lnx = nn.with_logical_constraint(mlp_lnx, ("activation_batch", "activation_norm_length", "activation_embed")) + mlp_lnx = with_logical_constraint(mlp_lnx, ("activation_batch", "activation_norm_length", "activation_embed")) layer_output = mlp_lnx + intermediate_inputs layer_output = self.dropout(layer_output, deterministic=deterministic) - layer_output = nn.with_logical_constraint( + layer_output = with_logical_constraint( layer_output, ("activation_batch", "activation_norm_length", "activation_embed"), ) @@ -235,18 +233,12 @@ def update_cache(cache, val): return layer_output, kv_cache -Olmo3DecoderLayerToLinen = nnx_wrappers.to_linen_class( - Olmo3DecoderLayer, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) - - class Olmo3ScannableBlock(nnx.Module): """A repeatable block of Olmo 3 decoder layers. This block applies multiple decoder layers sequentially, using the attention pattern defined by OLMO3_ATTENTION_PATTERN. It's designed to be - used with `nn.scan` for efficient compilation. + used with `jax.lax.scan` for efficient compilation. Attributes: config: Config, MaxText model config @@ -295,7 +287,7 @@ def __call__( ): cfg = self.config - inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + inputs = with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) inputs = checkpoint_name(inputs, "decoder_layer_input") y = inputs for layer_id in range(cfg.inhomogeneous_layer_cycle_interval): @@ -318,9 +310,3 @@ def __call__( return y, None else: return y - - -Olmo3ScannableBlockToLinen = nnx_wrappers.to_linen_class( - Olmo3ScannableBlock, - base_metadata_fn=initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/qwen2.py b/src/maxtext/models/qwen2.py index 6b60be79a7..45bc140587 100644 --- a/src/maxtext/models/qwen2.py +++ b/src/maxtext/models/qwen2.py @@ -26,8 +26,6 @@ from flax import nnx from maxtext.common.common_types import Config -from maxtext.layers import initializers as max_initializers -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant @@ -69,7 +67,6 @@ def __init__( mesh=mesh, shard_mode=config.shard_mode, debug_sharding=config.debug_sharding, - extra_stack_level=1, ) # Corresponds to Qwen2's `input_layernorm` @@ -231,9 +228,3 @@ def __call__( return layer_output, None else: return layer_output, kv_cache - - -Qwen2DecoderLayerToLinen = nnx_wrappers.to_linen_class( - Qwen2DecoderLayer, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 84c53cf96a..df80c10710 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -53,7 +53,6 @@ from maxtext.layers import initializers as max_initializers from maxtext.layers import moe from maxtext.layers import nnx_scan -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import PositionalEmbedding, Qwen3OmniMoeVisionPosEmbedInterpolate @@ -696,6 +695,9 @@ def __init__( self.cache = None # No cache for train mode or when inputs_shape not provided block_size = getattr(cfg, "weight_block_size", None) + # A custom rule set can shard the GatedDeltaNet projections apart from the attention projections by mapping + # `embed_gdn`; otherwise their embed dimension follows `embed_attn`. + gdn_embed_axis = "embed_gdn" if any(rule[0] == "embed_gdn" for rule in cfg.logical_axis_rules) else "embed_attn" # Submodule instantiations self.in_proj_qkvz = DenseGeneral( @@ -703,7 +705,7 @@ def __init__( out_features_shape=(self.key_dim * 2 + self.value_dim * 2), dtype=cfg.dtype, weight_dtype=cfg.weight_dtype, - kernel_axes=("embed_attn", "gdn_head"), + kernel_axes=(gdn_embed_axis, "gdn_head"), matmul_precision=cfg.matmul_precision, shard_mode=cfg.shard_mode, block_size=block_size, @@ -715,7 +717,7 @@ def __init__( out_features_shape=(self.num_v_heads * 2), dtype=cfg.dtype, weight_dtype=get_weight_dtype(cfg, "in_proj_ba"), - kernel_axes=("embed_attn", "gdn_head"), + kernel_axes=(gdn_embed_axis, "gdn_head"), matmul_precision=cfg.matmul_precision, shard_mode=cfg.shard_mode, quant=self.quant, @@ -768,7 +770,7 @@ def a_log_init(key, shape, dtype=jnp.float32): out_features_shape=(in_features,), dtype=cfg.dtype, weight_dtype=cfg.weight_dtype, - kernel_axes=("gdn_head", "embed_attn"), + kernel_axes=("gdn_head", gdn_embed_axis), matmul_precision=cfg.matmul_precision, shard_mode=cfg.shard_mode, block_size=block_size, @@ -906,7 +908,7 @@ def __call__( # `mamba_state_indices`, or -- under mamba prefix caching, where # `mamba_state_indices` is None -- a block id looked up in the mamba block # table, which needs `gdn_mamba_block_size`. - mamba_block_size = getattr(cfg, "gdn_mamba_block_size", None) + mamba_block_size = cfg.gdn_mamba_block_size use_paged_state = ( kv_cache is not None and isinstance(kv_cache, tuple) @@ -1858,7 +1860,6 @@ def __init__( mesh=mesh, shard_mode=cfg.shard_mode, debug_sharding=cfg.debug_sharding, - extra_stack_level=1, ) else: self._maybe_shard_with_logical = lambda inputs, *args, **kwargs: inputs @@ -2070,7 +2071,6 @@ def __init__( mesh=mesh, shard_mode=cfg.shard_mode, debug_sharding=cfg.debug_sharding, - extra_stack_level=1, ) else: self.out_sharding = None @@ -2247,7 +2247,6 @@ def __init__( mesh=mesh, shard_mode=config.shard_mode, debug_sharding=config.debug_sharding, - extra_stack_level=1, ) else: self.out_sharding = None @@ -3375,76 +3374,3 @@ def __call__(self, hidden_states: Array) -> Array: hidden_states = jax.nn.gelu(hidden_states) hidden_states = self.proj2(hidden_states) return hidden_states - - -# Vision encoder Linen wrappers -Qwen3OmniMoeVisionPatchMergerToLinen = nnx_wrappers.to_linen_class( - Qwen3OmniMoeVisionPatchMerger, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3OmniMoeVisionMLPToLinen = nnx_wrappers.to_linen_class( - Qwen3OmniMoeVisionMLP, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3OmniMoeVisionPatchEmbedToLinen = nnx_wrappers.to_linen_class( - Qwen3OmniMoeVisionPatchEmbed, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3OmniMoeVisionAttentionToLinen = nnx_wrappers.to_linen_class( - Qwen3OmniMoeVisionAttention, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3OmniMoeVisionBlockToLinen = nnx_wrappers.to_linen_class( - Qwen3OmniMoeVisionBlock, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3OmniMoeVisionEncoderToLinen = nnx_wrappers.to_linen_class( - Qwen3OmniMoeVisionEncoder, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3OmniMoeVisionProjectorToLinen = nnx_wrappers.to_linen_class( - Qwen3OmniMoeVisionProjector, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3DecoderLayerToLinen = nnx_wrappers.to_linen_class( - Qwen3DecoderLayer, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3MoeDecoderLayerToLinen = nnx_wrappers.to_linen_class( - Qwen3MoeDecoderLayer, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3NextDecoderLayerToLinen = nnx_wrappers.to_linen_class( - Qwen3NextDecoderLayer, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3NextScannableBlockToLinen = nnx_wrappers.to_linen_class( - Qwen3NextScannableBlock, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -# Audio encoder Linen wrappers -Qwen3OmniAudioEncoderLayerToLinen = nnx_wrappers.to_linen_class( - Qwen3OmniAudioEncoderLayer, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3OmniAudioEncoderToLinen = nnx_wrappers.to_linen_class( - Qwen3OmniAudioEncoder, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - -Qwen3OmniAudioProjectorToLinen = nnx_wrappers.to_linen_class( - Qwen3OmniAudioProjector, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index 88ef093a19..418fa69f06 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -24,8 +24,7 @@ import jax.numpy as jnp from jax.sharding import Mesh from maxtext.common.common_types import MODEL_MODE_TRAIN, Array, Config, ShardMode, get_weight_dtype -from maxtext.layers import initializers as max_initializers -from maxtext.layers import nnx_wrappers +from maxtext.layers.linears import MlpBlock from maxtext.layers.normalizations import Qwen3NextRMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.models.qwen3 import ( @@ -166,7 +165,6 @@ def __init__( mesh=mesh, shard_mode=cfg.shard_mode, debug_sharding=cfg.debug_sharding, - extra_stack_level=1, ) else: self.out_sharding = None @@ -219,8 +217,22 @@ def __init__( rngs=rngs, ) - # Instantiate our `Qwen3_5SparseMoEBlock`. - self.mlp = Qwen3_5SparseMoEBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) + if getattr(cfg, "num_experts", 1) > 1: + self.mlp = Qwen3_5SparseMoEBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) + else: + self.mlp = MlpBlock( + config=cfg, + mesh=self.mesh, + in_features=cfg.emb_dim, + intermediate_dim=cfg.mlp_dim, + activations=tuple(cfg.mlp_activations), + intermediate_dropout_rate=cfg.dropout_rate, + dtype=cfg.dtype, + weight_dtype=get_weight_dtype(cfg, "mlp"), + quant=self.quant, + model_mode=model_mode, + rngs=rngs, + ) def __call__( self, @@ -276,36 +288,43 @@ def __call__( hidden_states = residual + attention_output hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) - # Prepare for the MoE block by capturing the new residual + # Prepare for the MoE / MLP block by capturing the new residual residual = hidden_states - # Second LayerNorm, applied before the MoE block. + # Second LayerNorm, applied before the MoE / MLP block. hidden_states = self.post_attention_layernorm(hidden_states, out_sharding=self.out_sharding) hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) - # Padding tokens (segment id 0) are excluded from MoE routing when moe_exclude_padding_from_routing. - token_mask = None - if ( - self.config.moe_exclude_padding_from_routing - and decoder_segment_ids is not None - and model_mode == MODEL_MODE_TRAIN - ): - token_mask = decoder_segment_ids != 0 - - # Instantiate and call our `Qwen3_5SparseMoEBlock`. - mlp_output, load_balance_loss = self.mlp( - hidden_states, - deterministic=deterministic, - forced_routed_experts=forced_routed_experts, - intermediate_sharding=self.mlp_intermediate_sharding, - out_sharding=self.out_sharding, - token_mask=token_mask, - ) + # Instantiate and call our `Qwen3_5SparseMoEBlock` or `MlpBlock`. + if getattr(self.config, "num_experts", 1) > 1: + # Padding tokens (segment id 0) are excluded from MoE routing when moe_exclude_padding_from_routing. + token_mask = None + if ( + self.config.moe_exclude_padding_from_routing + and decoder_segment_ids is not None + and model_mode == MODEL_MODE_TRAIN + ): + token_mask = decoder_segment_ids != 0 + mlp_output, load_balance_loss = self.mlp( + hidden_states, + deterministic=deterministic, + forced_routed_experts=forced_routed_experts, + intermediate_sharding=self.mlp_intermediate_sharding, + out_sharding=self.out_sharding, + token_mask=token_mask, + ) - # We sow the load balancing loss so it can be collected and added to the total loss - # during training. - if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: - self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) + # We sow the load balancing loss so it can be collected and added to the total loss + # during training. + if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: + self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) + else: + mlp_output = self.mlp( + hidden_states, + deterministic=deterministic, + intermediate_sharding=self.mlp_intermediate_sharding, + out_sharding=self.out_sharding, + ) # Final residual connection (after the MoE block) mlp_output = self._maybe_shard_with_logical(mlp_output, self.activation_axis_names) @@ -315,15 +334,3 @@ def __call__( self.activation_axis_names, ) return layer_output, new_kv_cache - - -Qwen3_5DecoderLayerToLinen = nnx_wrappers.to_linen_class( - Qwen3_5DecoderLayer, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) - - -Qwen3_5ScannableBlockToLinen = nnx_wrappers.to_linen_class( - Qwen3_5ScannableBlock, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/qwen3_custom.py b/src/maxtext/models/qwen3_custom.py index 8b1dadbaab..7380acff78 100644 --- a/src/maxtext/models/qwen3_custom.py +++ b/src/maxtext/models/qwen3_custom.py @@ -22,7 +22,6 @@ from maxtext.common.common_types import Config from maxtext.layers import initializers as max_initializers from maxtext.layers import moe -from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.layers.attentions import Attention @@ -256,9 +255,3 @@ def __call__( return layer_output, None else: return layer_output, kv_cache - - -Qwen3CustomMoeDecoderLayerToLinen = nnx_wrappers.to_linen_class( - Qwen3CustomMoeDecoderLayer, - base_metadata_fn=max_initializers.variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/simple_layer.py b/src/maxtext/models/simple_layer.py index ac4eb915a8..17be3ec72d 100644 --- a/src/maxtext/models/simple_layer.py +++ b/src/maxtext/models/simple_layer.py @@ -19,8 +19,7 @@ from flax import nnx from maxtext.common.common_types import Config, ShardMode -from maxtext.layers import quantizations, nnx_wrappers -from maxtext.layers.initializers import variable_to_logically_partitioned +from maxtext.layers import quantizations from maxtext.utils.sharding import create_sharding @@ -78,12 +77,6 @@ def __call__( return jnp.dot(inputs, self.weights.astype(inputs.dtype), out_sharding=self.out_sharding) -SimpleDecoderLayerToLinen = nnx_wrappers.to_linen_class( - SimpleDecoderLayer, - base_metadata_fn=variable_to_logically_partitioned, -) - - class SimpleMlpDecoderLayer(nnx.Module): """Decoder layer consisting of [embed,mlp] followed by an [mlp,embed] matmul.""" @@ -142,9 +135,3 @@ def __call__( if self.config.scan_layers: return output, None return output - - -SimpleMlpDecoderLayerToLinen = nnx_wrappers.to_linen_class( - SimpleMlpDecoderLayer, - base_metadata_fn=variable_to_logically_partitioned, -) diff --git a/src/maxtext/models/weaver.py b/src/maxtext/models/weaver.py new file mode 100644 index 0000000000..439f5cb337 --- /dev/null +++ b/src/maxtext/models/weaver.py @@ -0,0 +1,2483 @@ +# 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. + +r"""Self-contained Weaver Mixture-of-Transformers architecture. + +This module implements the Weaver architecture from top to bottom in one +place without routing through a shared ``Attention`` or ``Decoder`` class: + +1. ``WeaverConfig``: Typed dataclass configuration. +2. Sequence packing & stream helpers (``WeaverPackingMetadata``, + ``build_weaver_packing_metadata``, ``unpack_streams``, + ``reinterleave_streams``). +3. Sample-fenced causal & bidirectional attention masks + Pallas Splash masks. +4. 3D M-RoPE rotary position embeddings (``rotate_half``, + ``apply_rotary_pos_emb``, ``compute_3d_mrope_cos_sin``). +5. Dual-pathway attention kernels & ``WeaverDualAttention`` module using pure + ``flax.nnx`` leaf layers (``nnx.Linear``, ``nnx.RMSNorm``). +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 + ======================================================================== + + und_seq [U, M] gen_seq [G, M] + │ │ + ├──────────────┐ ┌──────────────┤ + │ input_layernorm input_layernorm_moe_gen + │ │ │ │ + │ └──── WeaverDualAttention ──────┘ + │ (causal und | full gen) + │ ┌────────────────┬──────────────┐ + │ │ │ │ + ▼ ▼ ▼ ▼ + residual_und = und_seq + und_attn residual_gen = gen_seq + gen_attn + │ │ + post_attention_layernorm post_attention_layernorm_moe_gen + │ │ + mlp mlp_moe_gen + │ │ + ▼ ▼ + und_out = residual_und + mlp_out gen_out = residual_gen + mlp_out + [U, M] [G, M] +""" + +from __future__ import annotations + +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 +import jax +from jax import sharding +from jax.ad_checkpoint import checkpoint_name +import jax.numpy as jnp +from maxtext.kernels.attention import jax_flash_attention +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 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 + qk_norm_for_text: bool = True + qk_norm_for_diffusion: bool = True + 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 = 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") + qk_norm_kernel_axes: tuple[None | str, ...] = () + mlp_up_kernel_axes: tuple[None | str, ...] = ("embed", "mlp") + mlp_down_kernel_axes: tuple[None | str, ...] = ("mlp", "embed") + norm_kernel_axes: tuple[None | str, ...] = ("norm",) + dtype: Any = jnp.float32 + weight_dtype: Any = jnp.float32 + + def __post_init__(self) -> None: + if self.num_attention_heads % self.num_key_value_heads != 0: + raise ValueError( + "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( + kernel_init: Callable[..., jax.Array] | None, +) -> nnx.Initializer: + """Adapts a 3-arg NNX initializer or 5-arg NdInitializer to a 3-arg NNX initializer.""" + 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: + 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, + *, + use_bias: bool = False, + dtype: Any = jnp.float32, + weight_dtype: Any = jnp.float32, + kernel_init: Callable[..., jax.Array] | None = None, + kernel_axes: tuple[None | str, ...] = (), + sharding_hook: Callable[[Any, str], Any] | None = None, + hook_name: str = "linear_kernel_init", + rngs: nnx.Rngs, +) -> nnx.Linear: + """Constructs a pure ``nnx.Linear`` with logical-axis annotations.""" + base_init = _adapt_kernel_init(kernel_init) + 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, + ) + bias_axes = kernel_axes[-1:] if kernel_axes else () + bias_init = nnx.with_metadata( + nnx.initializers.zeros_init(), + out_sharding=bias_axes, + kernel_axes=bias_axes, + eager_sharding=False, + ) + layer = nnx.Linear( + in_features=in_features, + out_features=out_features, + use_bias=use_bias, + dtype=dtype, + param_dtype=weight_dtype, + kernel_init=meta_init, + bias_init=bias_init, + rngs=rngs, + ) + layer.kernel_axes = kernel_axes + return layer + + +def _make_rmsnorm( + num_features: int, + *, + epsilon: float = 1e-6, + dtype: Any = jnp.float32, + weight_dtype: Any = jnp.float32, + kernel_axes: tuple[None | str, ...] = (), + sharding_hook: Callable[[Any, str], Any] | None = None, + hook_name: str = "norm_scale_init", + rngs: nnx.Rngs, +) -> nnx.RMSNorm: + """Constructs a pure ``nnx.RMSNorm`` with logical-axis annotations.""" + base_scale_init: Any = nnx.initializers.ones_init() + if sharding_hook is not None: + hooked = sharding_hook(1, hook_name) + if callable(hooked): + base_scale_init = hooked + meta_scale_init = nnx.with_metadata( + base_scale_init, + out_sharding=kernel_axes, + kernel_axes=kernel_axes, + eager_sharding=False, + ) + norm = nnx.RMSNorm( + num_features=num_features, + epsilon=epsilon, + dtype=dtype, + param_dtype=weight_dtype, + scale_init=meta_scale_init, + rngs=rngs, + ) + norm.kernel_axes = kernel_axes + return norm + + +# ============================================================================== +# 1. Packing Metadata and Stream Helpers +# ============================================================================== + + +@dataclasses.dataclass(frozen=True) +class WeaverPackingMetadata: + """Boundary metadata and stream indexing for packed multimodal sequences. + + Attributes: + causal_q_offsets: Cumulative offsets for understanding (UND) queries, shape + [num_samples + 1]. Defines sample boundaries in the UND stream. + full_q_offsets: Cumulative offsets for generation (GEN) queries, shape + [num_samples + 1]. Defines sample boundaries in the GEN stream. + sample_kv_offsets: Cumulative offsets for total tokens per sample in the + interleaved sequence, shape [num_samples + 1]. Defines sample boundaries + for the combined KV context. + packed_und_token_indexes: Global token indices belonging to the UND stream, + shape [num_und_tokens]. + packed_gen_token_indexes: Global token indices belonging to the GEN stream, + shape [num_gen_tokens]. + max_causal_len: Maximum number of UND tokens across all samples. + max_full_len: Maximum number of GEN tokens across all samples. + max_sample_len: Maximum number of total tokens across all samples. + num_samples: Number of batch instances B packed into the sequence. + total_tokens: Total sequence length N_total = N_und + N_gen. + num_und_tokens: Total number of UND tokens across all samples. + num_gen_tokens: Total number of GEN tokens across all samples. + """ + + causal_q_offsets: jax.Array + full_q_offsets: jax.Array + sample_kv_offsets: jax.Array + packed_und_token_indexes: jax.Array + packed_gen_token_indexes: jax.Array + max_causal_len: int + max_full_len: int + max_sample_len: int + num_samples: int + total_tokens: int + num_und_tokens: int + num_gen_tokens: int + + +def build_weaver_packing_metadata( + sample_und_lens: Sequence[int], + sample_gen_lens: Sequence[int], + packed_und_token_indexes: jax.Array | None = None, + packed_gen_token_indexes: jax.Array | None = None, +) -> WeaverPackingMetadata: + """Constructs WeaverPackingMetadata from per-sample token counts. + + If index arrays are not provided, assumes default contiguous layout per + sample: each sample b has its UND tokens followed by its GEN tokens. + + Args: + sample_und_lens: Number of UND tokens for each sample in the batch. + sample_gen_lens: Number of GEN tokens for each sample in the batch. + packed_und_token_indexes: Optional explicit global indices for UND tokens. + packed_gen_token_indexes: Optional explicit global indices for GEN tokens. + + Returns: + A fully populated WeaverPackingMetadata dataclass. + """ + num_samples = len(sample_und_lens) + if len(sample_gen_lens) != num_samples: + raise ValueError( + f"Length of sample_und_lens ({num_samples}) does not match " f"sample_gen_lens ({len(sample_gen_lens)})" + ) + + und_lens = list(sample_und_lens) + gen_lens = list(sample_gen_lens) + sample_lens = [u + g for u, g in zip(und_lens, gen_lens)] + + causal_q_offsets_list = [0] + full_q_offsets_list = [0] + sample_kv_offsets_list = [0] + + for u, g, s in zip(und_lens, gen_lens, sample_lens): + causal_q_offsets_list.append(causal_q_offsets_list[-1] + u) + full_q_offsets_list.append(full_q_offsets_list[-1] + g) + sample_kv_offsets_list.append(sample_kv_offsets_list[-1] + s) + + total_tokens = sample_kv_offsets_list[-1] + num_und_tokens = causal_q_offsets_list[-1] + num_gen_tokens = full_q_offsets_list[-1] + + if packed_und_token_indexes is None or packed_gen_token_indexes is None: + und_idx_list = [] + gen_idx_list = [] + current_offset = 0 + for u, g in zip(und_lens, gen_lens): + und_idx_list.extend(range(current_offset, current_offset + u)) + gen_idx_list.extend(range(current_offset + u, current_offset + u + g)) + current_offset += u + g + packed_und_token_indexes = jnp.array(und_idx_list, dtype=jnp.int32) + packed_gen_token_indexes = jnp.array(gen_idx_list, dtype=jnp.int32) + else: + packed_und_token_indexes = jnp.asarray(packed_und_token_indexes, dtype=jnp.int32) + packed_gen_token_indexes = jnp.asarray(packed_gen_token_indexes, dtype=jnp.int32) + + max_causal = max(und_lens) if und_lens else 0 + max_full = max(gen_lens) if gen_lens else 0 + max_sample = max(sample_lens) if sample_lens else 0 + + return WeaverPackingMetadata( + causal_q_offsets=jnp.array(causal_q_offsets_list, dtype=jnp.int32), + full_q_offsets=jnp.array(full_q_offsets_list, dtype=jnp.int32), + sample_kv_offsets=jnp.array(sample_kv_offsets_list, dtype=jnp.int32), + packed_und_token_indexes=packed_und_token_indexes, + packed_gen_token_indexes=packed_gen_token_indexes, + max_causal_len=max_causal, + max_full_len=max_full, + max_sample_len=max_sample, + num_samples=num_samples, + total_tokens=total_tokens, + num_und_tokens=num_und_tokens, + num_gen_tokens=num_gen_tokens, + ) + + +def unpack_streams( + packed_tokens: jax.Array, + packed_und_token_indexes: jax.Array, + packed_gen_token_indexes: jax.Array, + use_ragged_ops: bool = False, +) -> tuple[jax.Array, jax.Array]: + """Unpacks a single 1D packed sequence into understanding and generation streams. + + Args: + packed_tokens: Tensor of shape (N_total, D_model) containing all packed + tokens. + packed_und_token_indexes: 1D array of indices for UND tokens. + packed_gen_token_indexes: 1D array of indices for GEN tokens. + use_ragged_ops: Whether to use TPU SparseCore ragged gather kernels if + available. Default is False (standard JAX indexing). + + Returns: + und_tokens: Tensor of shape (N_und, D_model). + gen_tokens: Tensor of shape (N_gen, D_model). + """ + del use_ragged_ops + und_tokens = packed_tokens[packed_und_token_indexes] + gen_tokens = packed_tokens[packed_gen_token_indexes] + return und_tokens, gen_tokens + + +def reinterleave_streams( + und_tokens: jax.Array, + gen_tokens: jax.Array, + packed_und_token_indexes: jax.Array, + packed_gen_token_indexes: jax.Array, + total_tokens: int, + use_ragged_ops: bool = False, +) -> jax.Array: + """Re-interleaves sliced understanding and generation outputs into the global packed order. + + Args: + und_tokens: Tensor of shape (N_und, D_model) from the understanding stream. + gen_tokens: Tensor of shape (N_gen, D_model) from the generation stream. + packed_und_token_indexes: 1D array of indices where UND tokens belong in the + packed sequence. + packed_gen_token_indexes: 1D array of indices where GEN tokens belong in the + packed sequence. + total_tokens: Total sequence length N_total. + use_ragged_ops: Whether to use TPU SparseCore ragged scatter kernels if + available. Default is False (standard JAX dynamic tensor scatter). + + Returns: + A single merged tensor of shape (N_total, D_model) matching the original + global layout. + """ + del use_ragged_ops + trailing_shape = und_tokens.shape[1:] + output_tokens = jnp.zeros((total_tokens, *trailing_shape), dtype=und_tokens.dtype) + if und_tokens.shape[0] > 0: + output_tokens = output_tokens.at[packed_und_token_indexes].set(und_tokens) + if gen_tokens.shape[0] > 0: + output_tokens = output_tokens.at[packed_gen_token_indexes].set(gen_tokens) + return output_tokens + + +# ============================================================================== +# 2. Attention Mask Construction +# ============================================================================== + + +def build_causal_understanding_mask( + causal_q_offsets: jax.Array, + num_und_tokens: int, +) -> jax.Array: + """Builds sample-specific lower-triangular causal attention mask for Kernel 1. + + Cross-sample interference is strictly prohibited via offset fences: + tokens can only attend to earlier tokens within the same sample. + + Args: + causal_q_offsets: Array of shape [num_samples + 1] giving cumulative UND + token counts. + num_und_tokens: Total number of UND tokens N_und. + + Returns: + A boolean array of shape [num_und_tokens, num_und_tokens], where True + indicates valid attention and False indicates masked-out positions. + """ + if num_und_tokens == 0: + return jnp.zeros((0, 0), dtype=jnp.bool_) + + indices = jnp.arange(num_und_tokens) + sample_ids = jnp.sum(indices[:, None] >= causal_q_offsets[None, 1:], axis=-1) + + same_sample = sample_ids[:, None] == sample_ids[None, :] + is_causal = indices[:, None] >= indices[None, :] + return same_sample & is_causal + + +def build_full_generative_mask( + full_q_offsets: jax.Array, + sample_kv_offsets: jax.Array, + num_gen_tokens: int, + total_tokens: int, +) -> jax.Array: + """Builds bidirectional full generative attention mask for Kernel 2. + + Enables generation queries to attend to all preceding context (both + textual/UND and visual/GEN tokens) within the same sample while strictly + prohibiting cross-sample interference. + + Args: + full_q_offsets: Array of shape [num_samples + 1] giving cumulative GEN query + counts. + sample_kv_offsets: Array of shape [num_samples + 1] giving cumulative total + token counts. + num_gen_tokens: Total number of GEN query tokens N_gen. + total_tokens: Total sequence length N_total (combined UND + GEN). + + Returns: + A boolean array of shape [num_gen_tokens, total_tokens], where True + indicates valid attention and False indicates masked-out positions. + """ + if num_gen_tokens == 0 or total_tokens == 0: + return jnp.zeros((num_gen_tokens, total_tokens), dtype=jnp.bool_) + + q_indices = jnp.arange(num_gen_tokens) + kv_indices = jnp.arange(total_tokens) + + q_sample_ids = jnp.sum(q_indices[:, None] >= full_q_offsets[None, 1:], axis=-1) + kv_sample_ids = jnp.sum(kv_indices[:, None] >= sample_kv_offsets[None, 1:], axis=-1) + + return q_sample_ids[:, None] == kv_sample_ids[None, :] + + +def build_causal_understanding_splash_mask( + causal_q_offsets: jax.Array, + num_und_tokens: int, +): + """Builds a Pallas Splash Attention Mask for Kernel 1 (Causal Understanding).""" + dense_mask = build_causal_understanding_mask(causal_q_offsets, num_und_tokens) + return splash_attention_mask.NumpyMask(np.asarray(dense_mask, dtype=np.bool_)) + + +def build_full_generative_splash_mask( + full_q_offsets: jax.Array, + sample_kv_offsets: jax.Array, + num_gen_tokens: int, + total_tokens: int, +): + """Builds a Pallas Splash Attention Mask for Kernel 2 (Full Generative).""" + dense_mask = build_full_generative_mask(full_q_offsets, sample_kv_offsets, num_gen_tokens, total_tokens) + return splash_attention_mask.NumpyMask(np.asarray(dense_mask, dtype=np.bool_)) + + +def compile_weaver_splash_mask( + mask_array: jax.Array | np.ndarray, + block_q: int = 128, + block_kv: int = 128, +): + """Compiles a dense boolean mask into Pallas Splash Attention MaskInfo.""" + mask_np = np.asarray(mask_array, dtype=np.bool_) + q_len, kv_len = mask_np.shape + if q_len % block_q != 0 or kv_len % block_kv != 0: + raise ValueError( + f"Sequence lengths ({q_len}, {kv_len}) must be multiples of block sizes" + f" ({block_q}, {block_kv}) for Splash Attention compilation." + ) + return splash_attention_mask_info.process_dynamic_mask(mask_np, (block_q, block_kv)) + + +# ============================================================================== +# 3. Rotary Position Embedding Helpers (3D M-RoPE) +# ============================================================================== + + +def rotate_half(x: jax.Array) -> jax.Array: + """Rotates half the hidden dimensions of the input tensor.""" + half_dim = x.shape[-1] // 2 + x1 = x[..., :half_dim] + x2 = x[..., half_dim:] + return jnp.concatenate((-x2, x1), axis=-1) + + +def apply_rotary_pos_emb( + x: jax.Array, + cos: jax.Array, + sin: jax.Array, + unsqueeze_dim: int | None = 1, +) -> jax.Array: + """Applies Rotary Position Embedding (RoPE / 3D M-RoPE) to an activation tensor.""" + if unsqueeze_dim is not None: + cos = jnp.expand_dims(cos, axis=unsqueeze_dim) + sin = jnp.expand_dims(sin, axis=unsqueeze_dim) + + x_f32 = x.astype(jnp.float32) + cos_f32 = cos.astype(jnp.float32) + sin_f32 = sin.astype(jnp.float32) + + rotated = (x_f32 * cos_f32) + (rotate_half(x_f32) * sin_f32) + return rotated.astype(x.dtype) + + +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 = 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: + raise ValueError(f"mrope_section {mrope_section} sum must equal head_dim // 2" f" ({head_dim // 2})") + + freq_dim = head_dim // 2 + inv_freq = 1.0 / (rope_theta ** (jnp.arange(0, freq_dim, 1, dtype=jnp.float32) * 2 / head_dim)) + + pos_temporal = position_ids_3d[:, 0, None].astype(jnp.float32) + freqs = pos_temporal * inv_freq[None, :] + + pos_height = position_ids_3d[:, 1, None].astype(jnp.float32) + pos_width = position_ids_3d[:, 2, None].astype(jnp.float32) + + h_limit = mrope_section[1] * 3 + w_limit = mrope_section[2] * 3 + + h_indices = jnp.arange(1, min(h_limit, freq_dim), 3) + w_indices = jnp.arange(2, min(w_limit, freq_dim), 3) + + freqs = freqs.at[:, h_indices].set(pos_height * inv_freq[h_indices]) + freqs = freqs.at[:, w_indices].set(pos_width * inv_freq[w_indices]) + + emb = jnp.concatenate([freqs, freqs], axis=-1) + return jnp.cos(emb), jnp.sin(emb) + + +# ============================================================================== +# 4. Dual-Attention Functional Kernels +# ============================================================================== + + +def causal_understanding_attention( + q_und: jax.Array, + k_und: jax.Array, + v_und: jax.Array, + causal_q_offsets: jax.Array, + 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 + if n_und == 0: + return jnp.zeros_like(q_und) + + num_kv_heads = k_und.shape[1] + if num_heads != num_kv_heads: + repeat_factor = num_heads // num_kv_heads + k_und = jnp.repeat(k_und, repeat_factor, axis=1) + v_und = jnp.repeat(v_und, repeat_factor, axis=1) + + if scale is None: + scale = 1.0 / math.sqrt(head_dim) + + if attention_kernel == "flash": + if n_und % block_q != 0: + raise ValueError( + "FlashAttention requires sequence length divisible by block_q" + f" ({block_q}), but got N_und={n_und}. Use" + " attention_kernel='dot_product' for unaligned lengths." + ) + mask = build_causal_understanding_mask(causal_q_offsets, n_und) + q_scaled = q_und * scale + q_reshaped = q_scaled.transpose(1, 0, 2)[jnp.newaxis, ...] + k_reshaped = k_und.transpose(1, 0, 2)[jnp.newaxis, ...] + v_reshaped = v_und.transpose(1, 0, 2)[jnp.newaxis, ...] + flash_out = jax_flash_attention.flash_attention_block_masked( + q=q_reshaped, + k=k_reshaped, + v=v_reshaped, + segment_ids=None, + block_q=block_q, + block_kv=block_kv, + mask=mask, + mask_value=mask_value, + ) + out = jnp.squeeze(flash_out, axis=0).transpose(1, 0, 2) + return out.astype(q_und.dtype) + + elif attention_kernel == "splash": + if n_und % block_q != 0: + raise ValueError( + "Splash Attention requires sequence length divisible by block size" + f" ({block_q}), but got N_und={n_und}. Use" + " attention_kernel='dot_product' for unaligned lengths." + ) + splash_mask = build_causal_understanding_splash_mask(causal_q_offsets, n_und) + sa_config = dataclasses.replace( + splash_attention_kernel.SplashConfig.get_default(), + block_q=block_q, + block_kv=block_kv, + ) + splash_kernel_fn = splash_attention_kernel.make_splash_mha_single_device( + mask=splash_mask, + config=sa_config, + downcast_smem_data=False, + partial_mask_blocks_dtype=np.int32, + ) + q_scaled = q_und * scale + q_reshaped = q_scaled.transpose(1, 0, 2) + k_reshaped = k_und.transpose(1, 0, 2) + v_reshaped = v_und.transpose(1, 0, 2) + splash_out = splash_kernel_fn(q_reshaped, k_reshaped, v_reshaped) + out = splash_out.transpose(1, 0, 2) + return out.astype(q_und.dtype) + + elif attention_kernel == "dot_product": + scores = jnp.einsum("qhd,khd->hqk", q_und, k_und) * scale + mask = build_causal_understanding_mask(causal_q_offsets, n_und) + scores = jnp.where(mask[None, :, :], scores, mask_value) + attn_weights = jax.nn.softmax(scores.astype(jnp.float32), axis=-1).astype(q_und.dtype) + return jnp.einsum("hqk,khd->qhd", attn_weights, v_und) + else: + raise ValueError( + f"Unsupported attention_kernel '{attention_kernel}'. " "Expected one of ['dot_product', 'flash', 'splash']." + ) + + +def full_generative_attention( + q_gen: jax.Array, + k_all: jax.Array, + v_all: jax.Array, + full_q_offsets: jax.Array, + sample_kv_offsets: jax.Array, + 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 + total_tokens = k_all.shape[0] + if n_gen == 0 or total_tokens == 0: + return jnp.zeros_like(q_gen) + + num_kv_heads = k_all.shape[1] + if num_heads != num_kv_heads: + repeat_factor = num_heads // num_kv_heads + k_all = jnp.repeat(k_all, repeat_factor, axis=1) + v_all = jnp.repeat(v_all, repeat_factor, axis=1) + + if scale is None: + scale = 1.0 / math.sqrt(head_dim) + + if attention_kernel == "flash": + if n_gen % block_q != 0 or total_tokens % block_kv != 0: + raise ValueError( + "FlashAttention requires sequence lengths divisible by block sizes" + f" (block_q={block_q}, block_kv={block_kv}), but got N_gen={n_gen}," + f" N_total={total_tokens}. Use attention_kernel='dot_product' for" + " unaligned lengths." + ) + mask = build_full_generative_mask(full_q_offsets, sample_kv_offsets, n_gen, total_tokens) + q_scaled = q_gen * scale + q_reshaped = q_scaled.transpose(1, 0, 2)[jnp.newaxis, ...] + k_reshaped = k_all.transpose(1, 0, 2)[jnp.newaxis, ...] + v_reshaped = v_all.transpose(1, 0, 2)[jnp.newaxis, ...] + flash_out = jax_flash_attention.flash_attention_block_masked( + q=q_reshaped, + k=k_reshaped, + v=v_reshaped, + segment_ids=None, + block_q=block_q, + block_kv=block_kv, + mask=mask, + mask_value=mask_value, + ) + out = jnp.squeeze(flash_out, axis=0).transpose(1, 0, 2) + return out.astype(q_gen.dtype) + + elif attention_kernel == "splash": + if n_gen % block_q != 0 or total_tokens % block_kv != 0: + raise ValueError( + "Splash Attention requires sequence lengths divisible by block sizes" + f" (block_q={block_q}, block_kv={block_kv}), but got N_gen={n_gen}," + f" N_total={total_tokens}. Use attention_kernel='dot_product' for" + " unaligned lengths." + ) + splash_mask = build_full_generative_splash_mask(full_q_offsets, sample_kv_offsets, n_gen, total_tokens) + sa_config = dataclasses.replace( + splash_attention_kernel.SplashConfig.get_default(), + block_q=block_q, + block_kv=block_kv, + ) + splash_kernel_fn = splash_attention_kernel.make_splash_mha_single_device( + mask=splash_mask, + config=sa_config, + downcast_smem_data=False, + partial_mask_blocks_dtype=np.int32, + ) + q_scaled = q_gen * scale + q_reshaped = q_scaled.transpose(1, 0, 2) + k_reshaped = k_all.transpose(1, 0, 2) + v_reshaped = v_all.transpose(1, 0, 2) + splash_out = splash_kernel_fn(q_reshaped, k_reshaped, v_reshaped) + out = splash_out.transpose(1, 0, 2) + return out.astype(q_gen.dtype) + + elif attention_kernel == "dot_product": + scores = jnp.einsum("qhd,khd->hqk", q_gen, k_all) * scale + mask = build_full_generative_mask(full_q_offsets, sample_kv_offsets, n_gen, total_tokens) + scores = jnp.where(mask[None, :, :], scores, mask_value) + attn_weights = jax.nn.softmax(scores.astype(jnp.float32), axis=-1).astype(q_gen.dtype) + return jnp.einsum("hqk,khd->qhd", attn_weights, v_all) + else: + raise ValueError( + f"Unsupported attention_kernel '{attention_kernel}'. " "Expected one of ['dot_product', 'flash', 'splash']." + ) + + +# ============================================================================== +# 5. Weaver Dual-Attention Module (Pure Flax NNX) +# ============================================================================== + + +class WeaverDualAttention(nnx.Module): + """Dual-pathway packed attention for Weaver using pure ``flax.nnx`` layers. + + Composes ``nnx.Linear`` and ``nnx.RMSNorm`` directly with zero wrapper layer + classes. + """ + + def __init__( + self, + dim: int | WeaverConfig, + num_heads: int | None = None, + num_kv_heads: int | None = None, + head_dim: int | None = None, + *, + qk_norm_for_text: bool = True, + qk_norm_for_diffusion: bool = True, + use_und_k_norm_for_gen: bool = False, + use_bias: bool = False, + 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"), + qk_norm_kernel_axes: tuple[None | str, ...] = (), + sharding_hook: Callable[[Any, str], Any] | None = None, + mesh: sharding.Mesh | None = None, + dtype: Any = jnp.float32, + weight_dtype: Any = jnp.float32, + kernel_init: Callable[..., jax.Array] | None = None, + rngs: nnx.Rngs, + ): + """Initializes ``WeaverDualAttention`` from explicit arguments or ``WeaverConfig``.""" + del mesh + if isinstance(dim, WeaverConfig): + cfg = dim + dim = cfg.hidden_size + num_heads = cfg.num_attention_heads + num_kv_heads = cfg.num_key_value_heads + head_dim = cfg.head_dim + qk_norm_for_text = cfg.qk_norm_for_text + qk_norm_for_diffusion = cfg.qk_norm_for_diffusion + use_und_k_norm_for_gen = cfg.use_und_k_norm_for_gen + use_bias = cfg.attention_bias + 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 + qk_norm_kernel_axes = cfg.qk_norm_kernel_axes + dtype = cfg.dtype + weight_dtype = cfg.weight_dtype + + if num_heads is None or num_kv_heads is None or head_dim is None: + raise ValueError("num_heads, num_kv_heads, and head_dim are required when dim is an int.") + + if attention_kernel not in ("dot_product", "flash", "splash"): + raise ValueError( + f"Unsupported attention_kernel '{attention_kernel}'. " + "Supported kernels are ['dot_product', 'flash', 'splash']." + ) + + 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 + self.head_dim = head_dim + self.qk_norm_for_text = qk_norm_for_text + self.qk_norm_for_diffusion = qk_norm_for_diffusion + self.use_und_k_norm_for_gen = use_und_k_norm_for_gen + self.q_kernel_axes = q_kernel_axes + self.kv_kernel_axes = kv_kernel_axes + self.o_kernel_axes = o_kernel_axes + self.qk_norm_kernel_axes = qk_norm_kernel_axes + self.sharding_hook = sharding_hook or (lambda x, name: x) + self.dtype = dtype + self.weight_dtype = weight_dtype + self.scaling = 1.0 / math.sqrt(head_dim) + + # Understanding Pathway Projections (pure nnx.Linear) + self.q_proj = _make_linear( + dim, + num_heads * head_dim, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=q_kernel_axes, + sharding_hook=sharding_hook, + hook_name="q_proj_kernel_init", + rngs=rngs, + ) + self.k_proj = _make_linear( + dim, + num_kv_heads * head_dim, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=kv_kernel_axes, + sharding_hook=sharding_hook, + hook_name="k_proj_kernel_init", + rngs=rngs, + ) + self.v_proj = _make_linear( + dim, + num_kv_heads * head_dim, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=kv_kernel_axes, + sharding_hook=sharding_hook, + hook_name="v_proj_kernel_init", + rngs=rngs, + ) + self.o_proj = _make_linear( + num_heads * head_dim, + dim, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=o_kernel_axes, + sharding_hook=sharding_hook, + hook_name="o_proj_kernel_init", + rngs=rngs, + ) + + # Understanding QK Normalization (pure nnx.RMSNorm) + if qk_norm_for_text: + self.q_norm = _make_rmsnorm( + head_dim, + epsilon=rms_norm_eps, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_axes=qk_norm_kernel_axes, + sharding_hook=sharding_hook, + hook_name="qk_norm_scale_init", + rngs=rngs, + ) + self.k_norm = _make_rmsnorm( + head_dim, + epsilon=rms_norm_eps, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_axes=qk_norm_kernel_axes, + sharding_hook=sharding_hook, + hook_name="qk_norm_scale_init", + rngs=rngs, + ) + else: + self.q_norm = None + self.k_norm = None + + # Generation Pathway Projections (pure nnx.Linear) + self.q_proj_gen = _make_linear( + dim, + num_heads * head_dim, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=q_kernel_axes, + sharding_hook=sharding_hook, + hook_name="q_proj_gen_kernel_init", + rngs=rngs, + ) + self.k_proj_gen = _make_linear( + dim, + num_kv_heads * head_dim, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=kv_kernel_axes, + sharding_hook=sharding_hook, + hook_name="k_proj_gen_kernel_init", + rngs=rngs, + ) + self.v_proj_gen = _make_linear( + dim, + num_kv_heads * head_dim, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=kv_kernel_axes, + sharding_hook=sharding_hook, + hook_name="v_proj_gen_kernel_init", + rngs=rngs, + ) + self.o_proj_gen = _make_linear( + num_heads * head_dim, + dim, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=o_kernel_axes, + sharding_hook=sharding_hook, + hook_name="o_proj_gen_kernel_init", + rngs=rngs, + ) + + # Generation QK Normalization (pure nnx.RMSNorm) + if qk_norm_for_diffusion: + self.q_norm_gen = _make_rmsnorm( + head_dim, + epsilon=rms_norm_eps, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_axes=qk_norm_kernel_axes, + sharding_hook=sharding_hook, + hook_name="qk_norm_scale_init", + rngs=rngs, + ) + self.k_norm_gen = _make_rmsnorm( + head_dim, + epsilon=rms_norm_eps, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_axes=qk_norm_kernel_axes, + sharding_hook=sharding_hook, + hook_name="qk_norm_scale_init", + rngs=rngs, + ) + else: + self.q_norm_gen = None + self.k_norm_gen = None + + # Cross-Attention UND K Normalization (pure nnx.RMSNorm) + if use_und_k_norm_for_gen and qk_norm_for_diffusion and not qk_norm_for_text: + self.k_norm_und_for_gen = _make_rmsnorm( + head_dim, + epsilon=rms_norm_eps, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_axes=qk_norm_kernel_axes, + sharding_hook=sharding_hook, + hook_name="qk_norm_scale_init", + rngs=rngs, + ) + else: + self.k_norm_und_for_gen = None + + @classmethod + def from_config( + cls, + config: WeaverConfig, + *, + rngs: nnx.Rngs, + sharding_hook: Callable[[Any, str], Any] | None = None, + kernel_init: Callable[..., jax.Array] | None = None, + ) -> WeaverDualAttention: + """Constructs ``WeaverDualAttention`` from a ``WeaverConfig``.""" + return cls(config, sharding_hook=sharding_hook, kernel_init=kernel_init, rngs=rngs) + + @overload + def __call__( + self, + tokens: jax.Array, + metadata: WeaverPackingMetadata, + cos: jax.Array | None = None, + sin: jax.Array | None = None, + *, + reinterleave: Literal[True] = True, + ) -> jax.Array: + ... + + @overload + def __call__( + self, + tokens: tuple[jax.Array, jax.Array], + metadata: WeaverPackingMetadata, + cos: jax.Array | None = None, + sin: jax.Array | None = None, + *, + reinterleave: bool = False, + ) -> tuple[jax.Array, jax.Array]: + ... + + @overload + def __call__( + self, + tokens: jax.Array, + metadata: WeaverPackingMetadata, + cos: jax.Array | None = None, + sin: jax.Array | None = None, + *, + reinterleave: Literal[False], + ) -> tuple[jax.Array, jax.Array]: + ... + + def __call__( + self, + tokens: jax.Array | tuple[jax.Array, jax.Array], + metadata: WeaverPackingMetadata, + cos: jax.Array | None = None, + sin: jax.Array | None = None, + *, + reinterleave: bool = True, + ) -> jax.Array | tuple[jax.Array, jax.Array]: + """Executes dual-pathway attention on understanding and generation tokens. + + Input shapes and layouts: + tokens: + * Packed mode: 2D array of shape ``[N_total, dim]`` in global packed + order, where ``N_total = N_und + N_gen``. + * Unpacked mode: Tuple ``(und_tokens, gen_tokens)`` with shapes + ``[N_und, dim]`` and ``[N_gen, dim]``. + metadata: ``WeaverPackingMetadata`` containing cumulative sample + boundaries (``causal_q_offsets``: ``[B + 1]``, ``full_q_offsets``: + ``[B + 1]``, ``sample_kv_offsets``: ``[B + 1]``) and 1D index arrays + (``packed_und_token_indexes``: ``[N_und]``, + ``packed_gen_token_indexes``: ``[N_gen]``). + cos: Optional 3D M-RoPE cosine table of shape ``[N_total, head_dim]``. + sin: Optional 3D M-RoPE sine table of shape ``[N_total, head_dim]``. + reinterleave: When True and ``tokens`` is a 2D array, re-interleaves the + outputs back into global packed order. + + Output shapes and layouts: + * When ``reinterleave=True`` and ``tokens`` is 2D: Transformed packed array + of shape ``[N_total, dim]`` in global packed order. + * When ``reinterleave=False`` or ``tokens`` is a tuple: Tuple + ``(und_out, gen_out)`` of shapes ``[N_und, dim]`` and ``[N_gen, dim]``. + """ + hook = self.sharding_hook + if isinstance(tokens, tuple): + und_tokens, gen_tokens = tokens + should_reinterleave = False + else: + und_tokens, gen_tokens = unpack_streams( + tokens, + metadata.packed_und_token_indexes, + metadata.packed_gen_token_indexes, + use_ragged_ops=self.use_ragged_ops, + ) + should_reinterleave = reinterleave + + n_und = und_tokens.shape[0] + n_gen = gen_tokens.shape[0] + total_tokens = metadata.total_tokens + + # 1. Linear Projections + if n_und > 0: + q_und = self.q_proj(und_tokens).reshape(n_und, self.num_heads, self.head_dim) + k_und = self.k_proj(und_tokens).reshape(n_und, self.num_kv_heads, self.head_dim) + v_und = self.v_proj(und_tokens).reshape(n_und, self.num_kv_heads, self.head_dim) + else: + q_und = jnp.zeros((0, self.num_heads, self.head_dim), dtype=self.dtype) + k_und = jnp.zeros((0, self.num_kv_heads, self.head_dim), dtype=self.dtype) + v_und = jnp.zeros((0, self.num_kv_heads, self.head_dim), dtype=self.dtype) + + if n_gen > 0: + q_gen = self.q_proj_gen(gen_tokens).reshape(n_gen, self.num_heads, self.head_dim) + k_gen = self.k_proj_gen(gen_tokens).reshape(n_gen, self.num_kv_heads, self.head_dim) + v_gen = self.v_proj_gen(gen_tokens).reshape(n_gen, self.num_kv_heads, self.head_dim) + else: + q_gen = jnp.zeros((0, self.num_heads, self.head_dim), dtype=self.dtype) + k_gen = jnp.zeros((0, self.num_kv_heads, self.head_dim), dtype=self.dtype) + v_gen = jnp.zeros((0, self.num_kv_heads, self.head_dim), dtype=self.dtype) + + # 2. QK Normalization + if self.q_norm is not None and n_und > 0: + q_und = self.q_norm(q_und) + if self.k_norm is not None and n_und > 0: + k_und = self.k_norm(k_und) + + if self.q_norm_gen is not None and n_gen > 0: + q_gen = self.q_norm_gen(q_gen) + if self.k_norm_gen is not None and n_gen > 0: + k_gen = self.k_norm_gen(k_gen) + + # 3. Apply 3D M-RoPE + k_und_pre_rope = k_und + if cos is not None and sin is not None: + if n_und > 0: + cos_und = cos[metadata.packed_und_token_indexes] + sin_und = sin[metadata.packed_und_token_indexes] + q_und = apply_rotary_pos_emb(q_und, cos_und, sin_und, unsqueeze_dim=1) + k_und = apply_rotary_pos_emb(k_und, cos_und, sin_und, unsqueeze_dim=1) + if n_gen > 0: + cos_gen = cos[metadata.packed_gen_token_indexes] + sin_gen = sin[metadata.packed_gen_token_indexes] + q_gen = apply_rotary_pos_emb(q_gen, cos_gen, sin_gen, unsqueeze_dim=1) + k_gen = apply_rotary_pos_emb(k_gen, cos_gen, sin_gen, unsqueeze_dim=1) + + q_und = checkpoint_name(hook(q_und, "query"), "query_proj") + k_und = checkpoint_name(hook(k_und, "key"), "key_proj") + v_und = checkpoint_name(hook(v_und, "value"), "value_proj") + q_gen = checkpoint_name(hook(q_gen, "gen_query"), "gen_query_proj") + k_gen = checkpoint_name(hook(k_gen, "gen_key"), "gen_key_proj") + v_gen = checkpoint_name(hook(v_gen, "gen_value"), "gen_value_proj") + + # 4. Cross-Attention UND K Normalization (if configured) + if self.k_norm_und_for_gen is not None and n_und > 0: + k_und_norm_for_gen = self.k_norm_und_for_gen(k_und_pre_rope) + if cos is not None and sin is not None: + k_und_for_gen = apply_rotary_pos_emb(k_und_norm_for_gen, cos_und, sin_und, unsqueeze_dim=1) + else: + k_und_for_gen = k_und_norm_for_gen + else: + k_und_for_gen = k_und + + # 5. Assemble Interleaved K_all and V_all for Kernel 2 + k_all = jnp.zeros((total_tokens, self.num_kv_heads, self.head_dim), dtype=self.dtype) + v_all = jnp.zeros((total_tokens, self.num_kv_heads, self.head_dim), dtype=self.dtype) + if n_und > 0: + k_all = k_all.at[metadata.packed_und_token_indexes].set(k_und_for_gen) + v_all = v_all.at[metadata.packed_und_token_indexes].set(v_und) + if n_gen > 0: + k_all = k_all.at[metadata.packed_gen_token_indexes].set(k_gen) + v_all = v_all.at[metadata.packed_gen_token_indexes].set(v_gen) + + # 6. Kernel 1: Causal Understanding Attention + if n_und > 0: + und_attn_context = causal_understanding_attention( + q_und=q_und, + k_und=k_und, + v_und=v_und, + 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) + und_out = hook(self.o_proj(und_attn_flat), "post_attn") + else: + und_out = jnp.zeros((0, self.dim), dtype=self.dtype) + + # 7. Kernel 2: Full Generative Attention + if n_gen > 0: + gen_attn_context = full_generative_attention( + q_gen=q_gen, + k_all=k_all, + v_all=v_all, + full_q_offsets=metadata.full_q_offsets, + 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) + gen_out = hook(self.o_proj_gen(gen_attn_flat), "gen_post_attn") + else: + gen_out = jnp.zeros((0, self.dim), dtype=self.dtype) + + if should_reinterleave: + return reinterleave_streams( + und_tokens=und_out, + gen_tokens=gen_out, + 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 und_out, gen_out + + +WeaverAttention = WeaverDualAttention + + +# ============================================================================== +# 6. Feed-Forward Network (Pure Flax NNX) +# ============================================================================== + + +class WeaverMLP(nnx.Module): + """Feed-forward network for a single Weaver tower using pure ``nnx.Linear``. + + Supports two activation variants: + * ``silu`` - gated SwiGLU: ``down(silu(gate(x)) * up(x))`` (weaver-mini) + * ``relu2`` - squared ReLU: ``down(relu(up(x)) ** 2)`` (weaver-max) + """ + + def __init__( + self, + hidden_size: int | WeaverConfig, + intermediate_size: int | None = None, + *, + hidden_act: str = "silu", + use_bias: bool = False, + up_kernel_axes: tuple[None | str, ...] = ("embed", "mlp"), + down_kernel_axes: tuple[None | str, ...] = ("mlp", "embed"), + sharding_hook: Callable[[Any, str], Any] | None = None, + mesh: sharding.Mesh | None = None, + dtype: Any = jnp.float32, + weight_dtype: Any = jnp.float32, + kernel_init: Callable[..., jax.Array] | None = None, + rngs: nnx.Rngs, + ): + """Initializes the Weaver feed-forward block.""" + del mesh + if isinstance(hidden_size, WeaverConfig): + cfg = hidden_size + hidden_size = cfg.hidden_size + intermediate_size = cfg.intermediate_size + hidden_act = cfg.hidden_act + up_kernel_axes = cfg.mlp_up_kernel_axes + down_kernel_axes = cfg.mlp_down_kernel_axes + dtype = cfg.dtype + weight_dtype = cfg.weight_dtype + + if intermediate_size is None: + raise ValueError("intermediate_size is required when hidden_size is an int.") + + if hidden_act not in ("silu", "relu2"): + raise ValueError(f"Weaver only supports `hidden_act` values 'silu' and 'relu2', got {hidden_act!r}.") + + self.hidden_act = hidden_act + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.up_kernel_axes = up_kernel_axes + self.down_kernel_axes = down_kernel_axes + self.sharding_hook = sharding_hook or (lambda x, name: x) + + self.gate_proj = ( + _make_linear( + hidden_size, + intermediate_size, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=up_kernel_axes, + sharding_hook=sharding_hook, + hook_name="gate_proj_kernel_init", + rngs=rngs, + ) + if hidden_act == "silu" + else None + ) + self.up_proj = _make_linear( + hidden_size, + intermediate_size, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=up_kernel_axes, + sharding_hook=sharding_hook, + hook_name="up_proj_kernel_init", + rngs=rngs, + ) + self.down_proj = _make_linear( + intermediate_size, + hidden_size, + use_bias=use_bias, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + kernel_axes=down_kernel_axes, + sharding_hook=sharding_hook, + hook_name="down_proj_kernel_init", + rngs=rngs, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + """Applies the feed-forward transformation. + + Input shape and layout: + x: Activation array of shape ``[N, hidden_size]``, where ``N`` is the + number of tokens in the tower (e.g. ``N_und`` or ``N_gen``). + + Output shape and layout: + Transformed activations of shape ``[N, hidden_size]``. + """ + hook = self.sharding_hook + if self.hidden_act == "relu2": + hidden = jnp.square(jax.nn.relu(self.up_proj(x))) + else: + hidden = jax.nn.silu(self.gate_proj(x)) * self.up_proj(x) + hidden = checkpoint_name(hook(hidden, "mlpwi"), "mlpwi") + return hook(self.down_proj(hidden), "post_mlp") + + +# ============================================================================== +# 7. Dual-Pathway Mixture-of-Transformers Decoder Layer (Pure Flax NNX) +# ============================================================================== + + +class WeaverMoTDecoderLayer(nnx.Module): + """Dual-pathway Weaver Mixture-of-Transformers decoder layer. + + Encapsulates two parallel transformer towers that share one fused + dual-attention step: + + * Understanding pathway: ``input_layernorm``, ``to_q`` / ``to_k`` / ``to_v`` + / ``to_out``, ``norm_q`` / ``norm_k``, ``post_attention_layernorm``, + ``mlp``. + * Generation pathway: ``input_layernorm_moe_gen``, ``add_q_proj`` / + ``add_k_proj`` / ``add_v_proj`` / ``to_add_out``, ``norm_added_q`` / + ``norm_added_k``, ``post_attention_layernorm_moe_gen``, ``mlp_moe_gen``. + """ + + def __init__( + self, + hidden_size: int | WeaverConfig, + head_dim: int | None = None, + num_attention_heads: int | None = None, + num_key_value_heads: int | None = None, + intermediate_size: int | None = None, + *, + attention_bias: bool = False, + rms_norm_eps: float = 1e-6, + hidden_act: str = "silu", + qk_norm_for_text: bool = True, + qk_norm_for_diffusion: bool = True, + 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"), + qk_norm_kernel_axes: tuple[None | str, ...] = (), + mlp_up_kernel_axes: tuple[None | str, ...] = ("embed", "mlp"), + mlp_down_kernel_axes: tuple[None | str, ...] = ("mlp", "embed"), + norm_kernel_axes: tuple[None | str, ...] = ("norm",), + sharding_hook: Callable[[Any, str], Any] | None = None, + mesh: sharding.Mesh | None = None, + dtype: Any = jnp.float32, + weight_dtype: Any = jnp.float32, + kernel_init: Callable[..., jax.Array] | None = None, + rngs: nnx.Rngs, + ): + """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 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 + num_key_value_heads = cfg.num_key_value_heads + intermediate_size = cfg.intermediate_size + attention_bias = cfg.attention_bias + rms_norm_eps = cfg.rms_norm_eps + hidden_act = cfg.hidden_act + qk_norm_for_text = cfg.qk_norm_for_text + qk_norm_for_diffusion = cfg.qk_norm_for_diffusion + 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 + qk_norm_kernel_axes = cfg.qk_norm_kernel_axes + mlp_up_kernel_axes = cfg.mlp_up_kernel_axes + mlp_down_kernel_axes = cfg.mlp_down_kernel_axes + norm_kernel_axes = cfg.norm_kernel_axes + dtype = cfg.dtype + weight_dtype = cfg.weight_dtype + + if head_dim is None or num_attention_heads is None or num_key_value_heads is None or intermediate_size is None: + raise ValueError( + "head_dim, num_attention_heads, num_key_value_heads, and intermediate_size " + "are required when hidden_size is an int." + ) + + self.hidden_size = hidden_size + self.head_dim = head_dim + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.intermediate_size = intermediate_size + self.rms_norm_eps = rms_norm_eps + self.hidden_act = hidden_act + self.dtype = dtype + self.sharding_hook = sharding_hook or (lambda x, name: x) + + self.self_attn = WeaverDualAttention( + dim=hidden_size, + num_heads=num_attention_heads, + num_kv_heads=num_key_value_heads, + head_dim=head_dim, + qk_norm_for_text=qk_norm_for_text, + qk_norm_for_diffusion=qk_norm_for_diffusion, + use_und_k_norm_for_gen=use_und_k_norm_for_gen, + use_bias=attention_bias, + 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, + qk_norm_kernel_axes=qk_norm_kernel_axes, + sharding_hook=sharding_hook, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + rngs=rngs, + ) + + def _norm() -> nnx.RMSNorm: + return _make_rmsnorm( + hidden_size, + epsilon=rms_norm_eps, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_axes=norm_kernel_axes, + sharding_hook=sharding_hook, + hook_name="norm_scale_init", + rngs=rngs, + ) + + def _mlp() -> WeaverMLP: + return WeaverMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + hidden_act=hidden_act, + up_kernel_axes=mlp_up_kernel_axes, + down_kernel_axes=mlp_down_kernel_axes, + sharding_hook=sharding_hook, + dtype=dtype, + weight_dtype=weight_dtype, + kernel_init=kernel_init, + rngs=rngs, + ) + + # Understanding tower. + self.input_layernorm = _norm() + self.post_attention_layernorm = _norm() + self.mlp = _mlp() + + # Generation tower. + self.input_layernorm_moe_gen = _norm() + self.post_attention_layernorm_moe_gen = _norm() + self.mlp_moe_gen = _mlp() + + @classmethod + def from_config( + cls, + config: WeaverConfig | Any, + *, + rngs: nnx.Rngs, + sharding_hook: Callable[[Any, str], Any] | None = None, + kernel_init: Callable[..., jax.Array] | None = None, + ) -> WeaverMoTDecoderLayer: + """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) + + # --------------------------------------------------------------------- + # Canonical aliases onto the shared attention submodule. + # --------------------------------------------------------------------- + + @property + def to_q(self) -> nnx.Linear: + """Understanding query projection (``self_attn.q_proj``).""" + return self.self_attn.q_proj + + @property + def to_k(self) -> nnx.Linear: + """Understanding key projection (``self_attn.k_proj``).""" + return self.self_attn.k_proj + + @property + def to_v(self) -> nnx.Linear: + """Understanding value projection (``self_attn.v_proj``).""" + return self.self_attn.v_proj + + @property + def to_out(self) -> nnx.Linear: + """Understanding output projection (``self_attn.o_proj``).""" + return self.self_attn.o_proj + + @property + def norm_q(self) -> nnx.RMSNorm | None: + """Understanding query normalization (``self_attn.q_norm``).""" + return self.self_attn.q_norm + + @property + def norm_k(self) -> nnx.RMSNorm | None: + """Understanding key normalization (``self_attn.k_norm``).""" + return self.self_attn.k_norm + + @property + def add_q_proj(self) -> nnx.Linear: + """Generation query projection (``self_attn.q_proj_gen``).""" + return self.self_attn.q_proj_gen + + @property + def add_k_proj(self) -> nnx.Linear: + """Generation key projection (``self_attn.k_proj_gen``).""" + return self.self_attn.k_proj_gen + + @property + def add_v_proj(self) -> nnx.Linear: + """Generation value projection (``self_attn.v_proj_gen``).""" + return self.self_attn.v_proj_gen + + @property + def to_add_out(self) -> nnx.Linear: + """Generation output projection (``self_attn.o_proj_gen``).""" + return self.self_attn.o_proj_gen + + @property + def norm_added_q(self) -> nnx.RMSNorm | None: + """Generation query normalization (``self_attn.q_norm_gen``).""" + return self.self_attn.q_norm_gen + + @property + def norm_added_k(self) -> nnx.RMSNorm | None: + """Generation key normalization (``self_attn.k_norm_gen``).""" + return self.self_attn.k_norm_gen + + @property + def k_norm_und_for_gen(self) -> nnx.RMSNorm | None: + """Understanding key normalization used by generation queries.""" + return self.self_attn.k_norm_und_for_gen + + def __call__( + self, + und_seq: jax.Array, + gen_seq: jax.Array, + metadata: WeaverPackingMetadata, + cos: jax.Array | None = None, + sin: jax.Array | None = None, + ) -> tuple[jax.Array, jax.Array]: + """Runs one dual-pathway Mixture-of-Transformers decoder layer step. + + Input shapes and layouts: + und_seq: Understanding token sequence of shape ``[N_und, hidden_size]``. + gen_seq: Generation token sequence of shape ``[N_gen, hidden_size]``. + metadata: ``WeaverPackingMetadata`` containing cumulative sample + boundaries (``causal_q_offsets``: ``[B + 1]``, ``full_q_offsets``: + ``[B + 1]``, ``sample_kv_offsets``: ``[B + 1]``) and 1D index arrays + (``packed_und_token_indexes``: ``[N_und]``, + ``packed_gen_token_indexes``: ``[N_gen]``). + cos: Optional 3D M-RoPE cosine table of shape ``[N_total, head_dim]`` in + global packed order (``N_total = N_und + N_gen``). + sin: Optional 3D M-RoPE sine table of shape ``[N_total, head_dim]`` in + global packed order (``N_total = N_und + N_gen``). + + Output shapes and layouts: + Tuple ``(und_out, gen_out)`` where: + * und_out: Transformed understanding tokens of shape + ``[N_und, hidden_size]``. + * gen_out: Transformed generation tokens of shape + ``[N_gen, hidden_size]``. + """ + hook = self.sharding_hook + + # 1. Pre-attention normalization (independent per tower). + und_norm = hook(self.input_layernorm(und_seq), "attn_input") + gen_norm = hook(self.input_layernorm_moe_gen(gen_seq), "attn_input") + + # 2. Fused dual attention. + und_attn_out, gen_attn_out = self.self_attn( + (und_norm, gen_norm), + metadata, + cos, + sin, + ) + + # 3. Attention residual. + residual_und = hook(und_seq + und_attn_out, "post_attn") + residual_gen = hook(gen_seq + gen_attn_out, "gen_post_attn") + + # 4. Pre-MLP normalization, feed-forward, and MLP residual. + mlp_in_und = hook(self.post_attention_layernorm(residual_und), "mlp_input") + mlp_in_gen = hook(self.post_attention_layernorm_moe_gen(residual_gen), "mlp_input") + mlp_out_und = self.mlp(mlp_in_und) + 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/src/maxtext/multimodal/__init__.py b/src/maxtext/multimodal/__init__.py index 2237c9162e..442f251e68 100644 --- a/src/maxtext/multimodal/__init__.py +++ b/src/maxtext/multimodal/__init__.py @@ -4,10 +4,12 @@ # 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 +# 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. + +"""Multimodal VAE and processor utilities package.""" diff --git a/src/maxtext/multimodal/processor.py b/src/maxtext/multimodal/processor.py index c0ed59c153..5bfafb86ee 100644 --- a/src/maxtext/multimodal/processor.py +++ b/src/maxtext/multimodal/processor.py @@ -42,9 +42,9 @@ "qwen3.5-397b-a17b-fp8": ("qwen3_5", "qwen3_5"), # Stitched model "maxtext-omni-gemma3-qwen3": ("gemma3", "qwen3"), - # Cosmos - "cosmos3-nano-reasoner": ("qwen3_vl", "qwen3"), - "cosmos3-super-reasoner": ("qwen3_vl", "qwen3"), + # Weaver + "weaver-mini": ("qwen3_vl", "qwen3"), + "weaver-max": ("qwen3_vl", "qwen3"), } @@ -364,7 +364,7 @@ def downsample_video_mask_to_tokens(video_mask, config): """Routes video-mask reduction to the model-specific multimodal processor.""" if video_mask is None: return None - if config.model_name.startswith(("qwen3", "cosmos3")): + if config.model_name.startswith(("qwen3", "weaver")): from maxtext.multimodal.processor_qwen3_omni import ( # pylint: disable=import-outside-toplevel downsample_video_mask_to_tokens as downsample_qwen3_video_mask, ) diff --git a/src/maxtext/multimodal/processor_vae.py b/src/maxtext/multimodal/processor_vae.py new file mode 100644 index 0000000000..8b754537dd --- /dev/null +++ b/src/maxtext/multimodal/processor_vae.py @@ -0,0 +1,15 @@ +# Copyright 2023–2025 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. + +"""VAE image and video preprocessors and postprocessors for diffusion pipelines.""" diff --git a/src/maxtext/trainers/diloco/diloco.py b/src/maxtext/trainers/diloco/diloco.py index 7a9e733e78..9664257a8e 100644 --- a/src/maxtext/trainers/diloco/diloco.py +++ b/src/maxtext/trainers/diloco/diloco.py @@ -210,8 +210,13 @@ def build_streaming_diloco_train_step( config: pyconfig.HyperParameters, train_step: Callable[[Any, Batch, PRNGKey], tuple[Any, Metrics]], mesh: jax.sharding.Mesh | None = None, + outer_params_shardings: PyTree | None = None, ) -> Callable[[DiLoCoTrainState, Batch, PRNGKey], tuple[DiLoCoTrainState, Metrics]]: """Convert a local state and train step into streaming DiLoCo train step.""" + if getattr(config, "diloco_bucketize_non_scanned", False) and outer_params_shardings is None: + # The manipulator is built from tracers, which carry no sharding; without explicit shardings the fragmenter + # could split a sharded axis and make XLA gather whole leaves at every sync. + raise ValueError("diloco_bucketize_non_scanned requires outer_params_shardings for SPMD streaming DiLoCo.") outer_optimizer = optax.sgd( config.diloco_outer_lr, momentum=config.diloco_outer_momentum, @@ -233,7 +238,8 @@ def streaming_diloco_train_step(state: DiLoCoTrainState, batch: Batch, prng: PRN step=new_step, ) - manipulator = diloco_utils.FragmentedTreeManipulator.create(state.params, config) + # state.params are tracers here, so the fragmenter gets their shardings explicitly. + manipulator = diloco_utils.FragmentedTreeManipulator.create(state.params, config, shardings=outer_params_shardings) # Step 1: Run the synchronization logic if we hit a sync step is_sync_step = (new_step > 0) & (new_step % steps_between_syncs == 0) @@ -280,6 +286,7 @@ def build_diloco_train_step( config: pyconfig.HyperParameters, train_step: Callable[[Any, Batch, PRNGKey], tuple[Any, Metrics]], mesh: jax.sharding.Mesh | None = None, + outer_params_shardings: PyTree | None = None, ) -> Callable[[DiLoCoTrainState, Batch, PRNGKey], tuple[DiLoCoTrainState, Metrics]]: """Convert a local state and train step into DiLoCo-compatible versions. @@ -288,9 +295,11 @@ def build_diloco_train_step( train_step: A local train step. This will be executed independently within each replica. mesh: The mesh for sharding. + outer_params_shardings: Shardings of the outer parameters (`DiLoCoTrainState.params`). Streaming DiLoCo uses + them to split bucketized leaves only along unsharded axes; required when `diloco_bucketize_non_scanned` is set. """ if config.enable_streaming_diloco: - return build_streaming_diloco_train_step(config, train_step, mesh=mesh) + return build_streaming_diloco_train_step(config, train_step, mesh=mesh, outer_params_shardings=outer_params_shardings) return build_vanilla_diloco_train_step(config, train_step, mesh=mesh) diff --git a/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh b/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh index 505e2eff72..6d5f50660e 100755 --- a/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh +++ b/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh @@ -5,31 +5,40 @@ # - src/maxtext/trainers/diloco/scripts/run_moe.sh (Qwen3-30B MoE & DiLoCo hyperparameters, SparseCore flags) # - MyStuff/scripts/pretrain/run_olmo_streaming_diloco_v5p-128.sh (OLMo grain dataset pipeline & gcsfuse setup) # +# This script builds a container image with the local MaxText workspace code and submits +# the workload via Cluster Toolkit (gcluster); see docs/run_maxtext/run_maxtext_via_cluster_toolkit.md. +# TPU reservations are configured on the cluster's node pools, so there is no per-workload +# reservation flag. The OLMo dataset bucket is mounted with the GKE gcsfuse CSI driver +# (gcluster --mount), so the cluster must have the GcsFuseCsiDriver add-on enabled. +# # ------------------------------------------------------------------------------ # Example Invocations: # ------------------------------------------------------------------------------ # 1. Multi-Slice OLMo Grain DiLoCo Training (2x v5p-128, Qwen3-30B-A3B): -# XPK_CLUSTER="mlperf-v5p" XPK_ZONE="europe-west4-b" XPK_PROJECT="cloud-tpu-multipod-dev" \ -# XPK_DEVICE_TYPE="v5p-128" XPK_NUM_SLICES="2" RUN_NAME="qw3-olmo-dlco-01" WORKLOAD_NAME="qw3-olmo-01" \ -# BASE_OUTPUT_DIRECTORY="gs://chriszuo-maxtext-logs" OLMO_GCS_BASE="gs://chriszuo-maxtext-datasets" \ -# XPK_RESERVATION="cloudtpu-20240716121201-595617744" \ +# CLUSTER="your-cluster-name" LOCATION="your-cluster-location" PROJECT="your-project-id" \ +# DEVICE_TYPE="v5p-128" NUM_SLICES="2" RUN_NAME="qw3-olmo-dlco-01" WORKLOAD_NAME="qw3-olmo-01" \ +# BASE_OUTPUT_DIRECTORY="gs://your-bucket/maxtext-logs" OLMO_GCS_BASE="gs://your-bucket/datasets" \ # bash src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh # ------------------------------------------------------------------------------ set -euo pipefail # -------------------------- Cluster & Topology Settings -------------------------- -: "${XPK_PROJECT:=cloud-tpu-multipod-dev}" -: "${XPK_ZONE:=europe-west4-b}" -: "${XPK_CLUSTER:=mlperf-v5p}" -: "${XPK_RESERVATION:=}" -: "${XPK_DEVICE_TYPE:=v5p-128}" -: "${XPK_NUM_SLICES:=2}" -: "${XPK_PRIORITY:=medium}" -: "${XPK_MAX_RESTARTS:=50}" +# Required: GKE cluster name, project, and cluster location +# (LOCATION is the cluster's region for regional clusters, or its zone for zonal ones). +for var in CLUSTER PROJECT LOCATION; do + if [ -z "${!var:-}" ]; then + echo "Error: ${var} is not set. Please set CLUSTER, PROJECT and LOCATION (e.g. CLUSTER=your-cluster-name PROJECT=your-project-id LOCATION=your-cluster-location)." + exit 1 + fi +done +: "${DEVICE_TYPE:=v5p-128}" +: "${NUM_SLICES:=2}" +: "${PRIORITY:=medium}" +: "${MAX_RESTARTS:=50}" # -------------------------- Docker & Storage -------------------------- -: "${XPK_DOCKER_IMAGE:=gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:2026-07-17}" +: "${DOCKER_IMAGE_BASE:=us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_pre_training:latest}" if [ -z "${BASE_OUTPUT_DIRECTORY:-}" ]; then echo "Error: BASE_OUTPUT_DIRECTORY is not set. Please set it as an environment variable (e.g. export BASE_OUTPUT_DIRECTORY=gs://your-bucket/maxtext-logs)." exit 1 @@ -66,7 +75,11 @@ fi # -------------------------- Training & Optimizer Hyperparameters -------------------------- : "${RUN_NAME:=qwen3-olmo-dlco-$(date +%m%d%H%M)}" -: "${WORKLOAD_NAME:=qw3-olmo-$(date +%m%d%H%M)}" +: "${WORKLOAD_NAME:=qw3-olmo-$(date +%m%d%H%M)}" # gcluster limits workload names to 28 characters +if [ "${#WORKLOAD_NAME}" -gt 28 ]; then + echo "Error: WORKLOAD_NAME '${WORKLOAD_NAME}' exceeds the 28-character limit for gcluster workload names." + exit 1 +fi : "${WARMUP_STEPS:=2000}" : "${TARGET_GLOBAL_BATCH:=512}" @@ -93,8 +106,8 @@ TOTAL_DATASET_STEPS=$(( TOTAL_INSTANCES / TARGET_GLOBAL_BATCH )) # Determine batch size per device -DEVICE_NUM=$(echo "${XPK_DEVICE_TYPE}" | cut -d'-' -f2) -TOTAL_DEVICES=$(( DEVICE_NUM * XPK_NUM_SLICES )) +DEVICE_NUM=$(echo "${DEVICE_TYPE}" | cut -d'-' -f2) +TOTAL_DEVICES=$(( DEVICE_NUM * NUM_SLICES )) : "${PER_DEVICE_BATCH_SIZE:=$(( TARGET_GLOBAL_BATCH / TOTAL_DEVICES ))}" # Resolve container-side index path @@ -139,34 +152,56 @@ LIBTPU_INIT_ARGS=" \ --xla_tpu_enable_3d_reduce_scatter_decomposer=false " ## -------------------------- Command Execution -------------------------- -CMD="set -euo pipefail; mkdir -p /deps/src/src/dependencies/scripts && ( [ -f /deps/src/src/dependencies/scripts/setup_gcsfuse.sh ] || ln -sf /deps/src/dependencies/scripts/setup_gcsfuse.sh /deps/src/src/dependencies/scripts/setup_gcsfuse.sh ); mkdir -p /deps/src/src/maxtext && ( [ -d /deps/src/src/maxtext/configs ] || ln -sf /deps/src/maxtext/configs /deps/src/src/maxtext/configs ); if [[ \"\${MOUNT_GCSFUSE:-1}\" = \"1\" ]]; then bucket=\"\${GCS_BASE#gs://}\"; bucket=\"\${bucket%/}\"; bash /deps/src/dependencies/scripts/setup_gcsfuse.sh DATASET_GCS_BUCKET=\"\${bucket}\" MOUNT_PATH=\"\${LOCAL_MOUNT}\"; fi; export PYTHONPATH=/app/src:\${PYTHONPATH:-}; cd /app/src && unset XLA_FLAGS; export LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\"; LOAD_PARAM_ARG=\"\"; if [ -n \"\${LOAD_FULL_STATE_PATH}\" ]; then LOAD_PARAM_ARG=\"load_full_state_path=\${LOAD_FULL_STATE_PATH}\"; elif [ -n \"\${LOAD_PARAMETERS_PATH}\" ]; then LOAD_PARAM_ARG=\"load_parameters_path=\${LOAD_PARAMETERS_PATH}\"; fi; python3 -m maxtext.trainers.pre_train.train maxtext/configs/base.yml run_name=\"\${RUN_NAME}\" save_config_to_gcs=true base_output_directory=\"\${OUTPUT_DIR}\" \${LOAD_PARAM_ARG} dataset_type=olmo_grain olmo_index_path=\"\${INDEX_PATH}\" olmo_path_remap_from=\"\${GCS_BASE}\" olmo_path_remap_to=\"\${LOCAL_MOUNT}\" olmo_apply_ngram_filter=True data_shuffle_seed=\"\${DATA_SEED}\" model_name=\"\${MODEL_NAME}\" tokenizer_type=\"\${TOKENIZER_TYPE}\" tokenizer_path=\"\${TOKENIZER_PATH}\" per_device_batch_size=\"\${PER_DEVICE_BATCH_SIZE}\" max_target_length=\"\${MAX_TARGET_LENGTH}\" learning_rate=\"\${LEARNING_RATE}\" learning_rate_schedule_steps=\"\${LR_SCHEDULE_STEPS}\" learning_rate_final_fraction=\"\${COSINE_FINAL_FRAC}\" warmup_steps_fraction=\"\${WARMUP_FRAC}\" adam_b1=\"\${ADAM_B1}\" adam_b2=\"\${ADAM_B2}\" adam_eps=\"\${ADAM_EPS}\" adam_weight_decay=\"\${ADAM_WD}\" gradient_clipping_threshold=\"\${GRAD_CLIP}\" z_loss_multiplier=\"\${Z_LOSS}\" load_balance_loss_weight=\"\${LOAD_BALANCE_LOSS_WEIGHT}\" float32_gate_logits=\"\${FLOAT32_GATE_LOGITS}\" enable_diloco=true enable_streaming_diloco=\"\${ENABLE_STREAMING_DILOCO}\" dcn_diloco_parallelism=\"\${XPK_NUM_SLICES}\" diloco_sync_period=\"\${DILOCO_SYNC_PERIOD}\" diloco_outer_lr=\"\${DILOCO_OUTER_LR}\" diloco_outer_momentum=\"\${DILOCO_OUTER_MOMENTUM}\" num_diloco_fragments=\"\${DILOCO_NUM_FRAGMENTS}\" use_sequential_layers=\"\${DILOCO_USE_SEQUENTIAL_LAYERS}\" jax_distributed_initialization_timeout=1200 steps=\"\${STEPS}\"" +CMD="set -euo pipefail; mkdir -p /deps/src/src/maxtext && ( [ -d /deps/src/src/maxtext/configs ] || ln -sf /deps/src/maxtext/configs /deps/src/src/maxtext/configs ); export PYTHONPATH=/app/src:\${PYTHONPATH:-}; cd /app/src && unset XLA_FLAGS; export LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\"; LOAD_PARAM_ARG=\"\"; if [ -n \"\${LOAD_FULL_STATE_PATH}\" ]; then LOAD_PARAM_ARG=\"load_full_state_path=\${LOAD_FULL_STATE_PATH}\"; elif [ -n \"\${LOAD_PARAMETERS_PATH}\" ]; then LOAD_PARAM_ARG=\"load_parameters_path=\${LOAD_PARAMETERS_PATH}\"; fi; python3 -m maxtext.trainers.pre_train.train maxtext/configs/base.yml run_name=\"\${RUN_NAME}\" save_config_to_gcs=true base_output_directory=\"\${OUTPUT_DIR}\" \${LOAD_PARAM_ARG} dataset_type=olmo_grain olmo_index_path=\"\${INDEX_PATH}\" olmo_path_remap_from=\"\${GCS_BASE}\" olmo_path_remap_to=\"\${LOCAL_MOUNT}\" olmo_apply_ngram_filter=True data_shuffle_seed=\"\${DATA_SEED}\" model_name=\"\${MODEL_NAME}\" tokenizer_type=\"\${TOKENIZER_TYPE}\" tokenizer_path=\"\${TOKENIZER_PATH}\" per_device_batch_size=\"\${PER_DEVICE_BATCH_SIZE}\" max_target_length=\"\${MAX_TARGET_LENGTH}\" learning_rate=\"\${LEARNING_RATE}\" learning_rate_schedule_steps=\"\${LR_SCHEDULE_STEPS}\" learning_rate_final_fraction=\"\${COSINE_FINAL_FRAC}\" warmup_steps_fraction=\"\${WARMUP_FRAC}\" adam_b1=\"\${ADAM_B1}\" adam_b2=\"\${ADAM_B2}\" adam_eps=\"\${ADAM_EPS}\" adam_weight_decay=\"\${ADAM_WD}\" gradient_clipping_threshold=\"\${GRAD_CLIP}\" z_loss_multiplier=\"\${Z_LOSS}\" load_balance_loss_weight=\"\${LOAD_BALANCE_LOSS_WEIGHT}\" float32_gate_logits=\"\${FLOAT32_GATE_LOGITS}\" enable_diloco=true enable_streaming_diloco=\"\${ENABLE_STREAMING_DILOCO}\" dcn_diloco_parallelism=\"\${NUM_SLICES}\" diloco_sync_period=\"\${DILOCO_SYNC_PERIOD}\" diloco_outer_lr=\"\${DILOCO_OUTER_LR}\" diloco_outer_momentum=\"\${DILOCO_OUTER_MOMENTUM}\" num_diloco_fragments=\"\${DILOCO_NUM_FRAGMENTS}\" use_sequential_layers=\"\${DILOCO_USE_SEQUENTIAL_LAYERS}\" jax_distributed_initialization_timeout=1200 steps=\"\${STEPS}\"" + +# -------------------------- Image Build -------------------------- +# The training command runs from /app/src, so build an image that layers the local +# workspace onto the base image at /app (same approach as the sibling SPMD DiLoCo scripts). +MY_IMAGE="gcr.io/${PROJECT}/$(whoami)-runner:${WORKLOAD_NAME}" +echo "Building docker image containing local changes..." +docker build -t "${MY_IMAGE}" -f - . </dev/null || true +EOF + +echo "Pushing image ${MY_IMAGE}..." +docker push "${MY_IMAGE}" + +# -------------------------- OLMo Dataset Mount -------------------------- +# Mount OLMO_GCS_BASE at OLMO_LOCAL_MOUNT using the GKE gcsfuse CSI driver. If OLMO_GCS_BASE +# includes a sub-directory (gs://bucket/dir), only that directory is mounted. +OLMO_GCS_PATH="${OLMO_GCS_BASE#gs://}" +OLMO_GCS_PATH="${OLMO_GCS_PATH%/}" +OLMO_GCS_BUCKET="${OLMO_GCS_PATH%%/*}" +OLMO_MOUNT_OPTIONS="implicit-dirs" +if [[ "${OLMO_GCS_PATH}" == */* ]]; then + OLMO_MOUNT_OPTIONS+=",only-dir=${OLMO_GCS_PATH#*/}" +fi -echo "Submitting Qwen3-30B-A3B + OLMo DiLoCo training workload using XPK..." +echo "Submitting Qwen3-30B-A3B + OLMo DiLoCo training workload using Cluster Toolkit (gcluster)..." echo "Workload Name: ${WORKLOAD_NAME}" echo "Run Name: ${RUN_NAME}" echo "Model: ${MODEL_NAME}" echo "Dataset: olmo_grain (${OLMO_INDEX_PATH_IN_CONTAINER})" -echo "Topology: ${XPK_DEVICE_TYPE} x ${XPK_NUM_SLICES} slices" - -XPK_ARGS=( - --cluster "${XPK_CLUSTER}" - --workload "${WORKLOAD_NAME}" - --project "${XPK_PROJECT}" - --zone "${XPK_ZONE}" - --tpu-type "${XPK_DEVICE_TYPE}" - --num-slices "${XPK_NUM_SLICES}" - --priority "${XPK_PRIORITY}" - --max-restarts "${XPK_MAX_RESTARTS}" - --base-docker-image "${XPK_DOCKER_IMAGE}" - --script-dir "$(pwd)" - --command "export HF_TOKEN='${HF_TOKEN}'; export INDEX_PATH='${OLMO_INDEX_PATH_IN_CONTAINER}'; export GCS_BASE='${OLMO_GCS_BASE}'; export LOCAL_MOUNT='${OLMO_LOCAL_MOUNT}'; export OUTPUT_DIR='${BASE_OUTPUT_DIRECTORY}'; export RUN_NAME='${RUN_NAME}'; export LOAD_FULL_STATE_PATH='${LOAD_FULL_STATE_PATH}'; export LOAD_PARAMETERS_PATH='${LOAD_PARAMETERS_PATH}'; export MODEL_NAME='${MODEL_NAME}'; export TOKENIZER_TYPE='${TOKENIZER_TYPE}'; export TOKENIZER_PATH='${TOKENIZER_PATH}'; export MAX_TARGET_LENGTH='${MAX_TARGET_LENGTH}'; export STEPS='${STEPS}'; export WARMUP_STEPS='${WARMUP_STEPS}'; export LR_SCHEDULE_STEPS='${LR_SCHEDULE_STEPS}'; export LEARNING_RATE='${LEARNING_RATE}'; export COSINE_FINAL_FRAC='${COSINE_FINAL_FRAC}'; export ADAM_B1='${ADAM_B1}'; export ADAM_B2='${ADAM_B2}'; export ADAM_EPS='${ADAM_EPS}'; export ADAM_WD='${ADAM_WD}'; export GRAD_CLIP='${GRAD_CLIP}'; export Z_LOSS='${Z_LOSS}'; export LOAD_BALANCE_LOSS_WEIGHT='${LOAD_BALANCE_LOSS_WEIGHT}'; export FLOAT32_GATE_LOGITS='${FLOAT32_GATE_LOGITS}'; export ENABLE_STREAMING_DILOCO='${ENABLE_STREAMING_DILOCO}'; export DILOCO_SYNC_PERIOD='${DILOCO_SYNC_PERIOD}'; export DILOCO_OUTER_LR='${DILOCO_OUTER_LR}'; export DILOCO_OUTER_MOMENTUM='${DILOCO_OUTER_MOMENTUM}'; export DILOCO_NUM_FRAGMENTS='${DILOCO_NUM_FRAGMENTS}'; export DILOCO_USE_SEQUENTIAL_LAYERS='${DILOCO_USE_SEQUENTIAL_LAYERS}'; export DATA_SEED='${DATA_SEED}'; export MOUNT_GCSFUSE=1; export XPK_NUM_SLICES='${XPK_NUM_SLICES}'; export PER_DEVICE_BATCH_SIZE='${PER_DEVICE_BATCH_SIZE}'; export WARMUP_FRAC='${WARMUP_FRAC}'; ${CMD}" +echo "Topology: ${DEVICE_TYPE} x ${NUM_SLICES} slices" + +GCLUSTER_ARGS=( + --cluster "${CLUSTER}" + --name "${WORKLOAD_NAME}" + --project "${PROJECT}" + --location "${LOCATION}" + --compute-type "${DEVICE_TYPE}" + --num-slices "${NUM_SLICES}" + --priority "${PRIORITY}" + --restarts "${MAX_RESTARTS}" + --image "${MY_IMAGE}" + --mount "gs://${OLMO_GCS_BUCKET};${OLMO_LOCAL_MOUNT};ro;options=${OLMO_MOUNT_OPTIONS}" + --command "export HF_TOKEN='${HF_TOKEN}'; export INDEX_PATH='${OLMO_INDEX_PATH_IN_CONTAINER}'; export GCS_BASE='${OLMO_GCS_BASE}'; export LOCAL_MOUNT='${OLMO_LOCAL_MOUNT}'; export OUTPUT_DIR='${BASE_OUTPUT_DIRECTORY}'; export RUN_NAME='${RUN_NAME}'; export LOAD_FULL_STATE_PATH='${LOAD_FULL_STATE_PATH}'; export LOAD_PARAMETERS_PATH='${LOAD_PARAMETERS_PATH}'; export MODEL_NAME='${MODEL_NAME}'; export TOKENIZER_TYPE='${TOKENIZER_TYPE}'; export TOKENIZER_PATH='${TOKENIZER_PATH}'; export MAX_TARGET_LENGTH='${MAX_TARGET_LENGTH}'; export STEPS='${STEPS}'; export WARMUP_STEPS='${WARMUP_STEPS}'; export LR_SCHEDULE_STEPS='${LR_SCHEDULE_STEPS}'; export LEARNING_RATE='${LEARNING_RATE}'; export COSINE_FINAL_FRAC='${COSINE_FINAL_FRAC}'; export ADAM_B1='${ADAM_B1}'; export ADAM_B2='${ADAM_B2}'; export ADAM_EPS='${ADAM_EPS}'; export ADAM_WD='${ADAM_WD}'; export GRAD_CLIP='${GRAD_CLIP}'; export Z_LOSS='${Z_LOSS}'; export LOAD_BALANCE_LOSS_WEIGHT='${LOAD_BALANCE_LOSS_WEIGHT}'; export FLOAT32_GATE_LOGITS='${FLOAT32_GATE_LOGITS}'; export ENABLE_STREAMING_DILOCO='${ENABLE_STREAMING_DILOCO}'; export DILOCO_SYNC_PERIOD='${DILOCO_SYNC_PERIOD}'; export DILOCO_OUTER_LR='${DILOCO_OUTER_LR}'; export DILOCO_OUTER_MOMENTUM='${DILOCO_OUTER_MOMENTUM}'; export DILOCO_NUM_FRAGMENTS='${DILOCO_NUM_FRAGMENTS}'; export DILOCO_USE_SEQUENTIAL_LAYERS='${DILOCO_USE_SEQUENTIAL_LAYERS}'; export DATA_SEED='${DATA_SEED}'; export NUM_SLICES='${NUM_SLICES}'; export PER_DEVICE_BATCH_SIZE='${PER_DEVICE_BATCH_SIZE}'; export WARMUP_FRAC='${WARMUP_FRAC}'; ${CMD}" ) -if [ -n "${XPK_RESERVATION:-}" ] && [ "${XPK_RESERVATION}" != "NONE" ]; then - XPK_ARGS+=(--reservation "${XPK_RESERVATION}") -fi - -xpk workload create "${XPK_ARGS[@]}" +gcluster job submit "${GCLUSTER_ARGS[@]}" echo "Qwen3-30B-A3B + OLMo DiLoCo workload submission complete!" diff --git a/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh b/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh index bf964e1c5a..fa6cabcc37 100644 --- a/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh +++ b/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh @@ -3,7 +3,8 @@ # SPMD Streaming DiLoCo Multi-Slice Training Runner # ============================================================================== # This script builds a container image with local MaxText workspace code and submits -# a multi-slice SPMD Streaming DiLoCo training workload on GKE TPU clusters via XPK. +# a multi-slice SPMD Streaming DiLoCo training workload on GKE TPU clusters via +# Cluster Toolkit (gcluster). See docs/run_maxtext/run_maxtext_via_cluster_toolkit.md. # # Reference Paper: # "Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch" @@ -14,39 +15,43 @@ # - Fragment: Model weights partitioned into K disjoint subsets (num_diloco_fragments). # - Streaming Sync: Communicates 1 fragment per inner step to overlap DCN all-reduce with compute. # +# Note: TPU reservations are configured on the cluster's node pools when the cluster +# is created, so there is no per-workload reservation flag. +# # ------------------------------------------------------------------------------ # Example Invocations: # ------------------------------------------------------------------------------ # 1. Standard Multi-Slice Pretraining (2x v5p-128, Qwen3-8B): -# CLUSTER="mlperf-v5p" ZONE="europe-west4-b" PROJECT="cloud-tpu-multipod-dev" \ +# CLUSTER="your-cluster-name" LOCATION="your-cluster-location" PROJECT="your-project-id" \ # DEVICE_TYPE="v5p-128" NUM_SLICES="2" RUNNAME="dlco-qwen8b-01" \ -# BASE_OUTPUT_DIRECTORY="gs://chriszuo-maxtext-logs" DATASET_PATH="gs://chriszuo-maxtext-datasets" \ -# RESERVATION="cloudtpu-20240716121201-595617744" \ +# BASE_OUTPUT_DIRECTORY="gs://your-bucket/maxtext-logs" DATASET_PATH="gs://your-bucket/maxtext-datasets" \ # STEPS="1000" CHECKPOINT_PERIOD="100" DILOCO_SYNC_PERIOD="37" DILOCO_NUM_FRAGMENTS="37" \ # bash src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh # -# 2. Fast Smoke / Integration Test Run (2x v5p-16, short workload name < 20 chars): -# CLUSTER="mlperf-v5p" ZONE="europe-west4-b" PROJECT="cloud-tpu-multipod-dev" \ -# DEVICE_TYPE="v5p-16" NUM_SLICES="2" RUNNAME="dlco-smk-01" XPK_WORKLOAD="dlco-smk-01" \ -# BASE_OUTPUT_DIRECTORY="gs://chriszuo-maxtext-logs" DATASET_PATH="gs://chriszuo-maxtext-datasets" \ -# RESERVATION="cloudtpu-20240716121201-595617744" \ +# 2. Fast Smoke / Integration Test Run (2x v5p-16, workload name <= 28 chars): +# CLUSTER="your-cluster-name" LOCATION="your-cluster-location" PROJECT="your-project-id" \ +# DEVICE_TYPE="v5p-16" NUM_SLICES="2" RUNNAME="dlco-smk-01" WORKLOAD_NAME="dlco-smk-01" \ +# BASE_OUTPUT_DIRECTORY="gs://your-bucket/maxtext-logs" DATASET_PATH="gs://your-bucket/maxtext-datasets" \ # STEPS="20" CHECKPOINT_PERIOD="10" DILOCO_SYNC_PERIOD="3" DILOCO_NUM_FRAGMENTS="3" \ # bash src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh # # 3. Automatic Checkpoint Resumption (Reuse same RUNNAME, zero explicit path flags): -# CLUSTER="mlperf-v5p" ZONE="europe-west4-b" PROJECT="cloud-tpu-multipod-dev" \ -# DEVICE_TYPE="v5p-16" NUM_SLICES="2" RUNNAME="dlco-smk-01" XPK_WORKLOAD="dlco-resm-01" \ -# BASE_OUTPUT_DIRECTORY="gs://chriszuo-maxtext-logs" DATASET_PATH="gs://chriszuo-maxtext-datasets" \ -# RESERVATION="cloudtpu-20240716121201-595617744" \ +# CLUSTER="your-cluster-name" LOCATION="your-cluster-location" PROJECT="your-project-id" \ +# DEVICE_TYPE="v5p-16" NUM_SLICES="2" RUNNAME="dlco-smk-01" WORKLOAD_NAME="dlco-resm-01" \ +# BASE_OUTPUT_DIRECTORY="gs://your-bucket/maxtext-logs" DATASET_PATH="gs://your-bucket/maxtext-datasets" \ # STEPS="40" CHECKPOINT_PERIOD="10" LOAD_FULL_STATE_PATH="" \ # bash src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh # ============================================================================== set -e -# cluster -CLUSTER="${CLUSTER:-mlperf-v5p}" -PROJECT="${PROJECT:-cloud-tpu-multipod-dev}" -ZONE="${ZONE:-europe-west4-b}" +# Required: GKE cluster name, project, and cluster location +# (LOCATION is the cluster's region for regional clusters, or its zone for zonal ones). +for var in CLUSTER PROJECT LOCATION; do + if [ -z "${!var:-}" ]; then + echo "Error: ${var} is not set. Please set CLUSTER, PROJECT and LOCATION (e.g. CLUSTER=your-cluster-name PROJECT=your-project-id LOCATION=your-cluster-location)." + exit 1 + fi +done # specify resource NUM_SLICES="${NUM_SLICES:-2}" @@ -54,9 +59,13 @@ DEVICE_TYPE="${DEVICE_TYPE:-v5p-128}" # command RUNNAME="${RUNNAME:-spmd-dlco-$(date +%H%M)}" -XPK_WORKLOAD="${XPK_WORKLOAD:-$RUNNAME}" -DOCKER_IMAGE_BASE="${DOCKER_IMAGE_BASE:-gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:latest}" -MY_IMAGE="gcr.io/${PROJECT}/$(whoami)-runner:${XPK_WORKLOAD}" +WORKLOAD_NAME="${WORKLOAD_NAME:-$RUNNAME}" # gcluster limits workload names to 28 characters +if [ "${#WORKLOAD_NAME}" -gt 28 ]; then + echo "Error: WORKLOAD_NAME '${WORKLOAD_NAME}' exceeds the 28-character limit for gcluster workload names." + exit 1 +fi +DOCKER_IMAGE_BASE="${DOCKER_IMAGE_BASE:-us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_pre_training:latest}" +MY_IMAGE="gcr.io/${PROJECT}/$(whoami)-runner:${WORKLOAD_NAME}" if [ -z "${BASE_OUTPUT_DIRECTORY:-}" ]; then echo "Error: BASE_OUTPUT_DIRECTORY is not set. Please set it as an environment variable (e.g. export BASE_OUTPUT_DIRECTORY=gs://your-bucket/maxtext-logs)." @@ -167,23 +176,20 @@ EOF echo "Pushing image ${MY_IMAGE}..." docker push "${MY_IMAGE}" -# 2. Create the workload directly using xpk -echo "Creating workload: ${XPK_WORKLOAD}" -XPK_ARGS=( - --workload "${XPK_WORKLOAD}" - --docker-image "${MY_IMAGE}" +# 2. Submit the workload using Cluster Toolkit (gcluster) +echo "Submitting workload: ${WORKLOAD_NAME}" +GCLUSTER_ARGS=( + --name "${WORKLOAD_NAME}" + --image "${MY_IMAGE}" --command "${CMD}" --num-slices "${NUM_SLICES}" --priority "${PRIORITY:-medium}" - --enable-debug-logs + --restarts "${MAX_RESTARTS:-0}" + --verbose --cluster "${CLUSTER}" - --tpu-type "${DEVICE_TYPE}" + --compute-type "${DEVICE_TYPE}" --project "${PROJECT}" - --zone "${ZONE}" + --location "${LOCATION}" ) -if [ -n "${RESERVATION:-}" ] && [ "${RESERVATION}" != "NONE" ]; then - XPK_ARGS+=(--reservation "${RESERVATION}") -fi - -xpk workload create "${XPK_ARGS[@]}" +gcluster job submit "${GCLUSTER_ARGS[@]}" diff --git a/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_dlco_moe_c4.sh b/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_dlco_moe_c4.sh index 694c7a6a38..fcddd80a06 100755 --- a/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_dlco_moe_c4.sh +++ b/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_dlco_moe_c4.sh @@ -8,27 +8,39 @@ set -e # 4. Post-Sync Loss Spike Severity (Delta L_sync / diloco/post_sync_loss_spike_severity) # 5. Post-Sync Expert Utilization Entropy (EUE / diloco/post_sync_expert_utilization_entropy) # +# The workload is submitted via Cluster Toolkit (gcluster); see +# docs/run_maxtext/run_maxtext_via_cluster_toolkit.md. TPU reservations are configured +# on the cluster's node pools, so there is no per-workload reservation flag. +# # ------------------------------------------------------------------------------ # Example Invocations: # ------------------------------------------------------------------------------ # 1. Standard Multi-Slice MoE DiLoCo Training (2x v5p-128, Qwen3-30B-A3B on C4): -# CLUSTER="mlperf-v5p" ZONE="europe-west4-b" PROJECT="cloud-tpu-multipod-dev" \ +# CLUSTER="your-cluster-name" LOCATION="your-cluster-location" PROJECT="your-project-id" \ # DEVICE_TYPE="v5p-128" NUM_SLICES="2" RUNNAME="dlco-moe-01" \ -# BASE_OUTPUT_DIRECTORY="gs://chriszuo-maxtext-logs" DATASET_PATH="gs://chriszuo-maxtext-datasets" \ -# RESERVATION="cloudtpu-20240716121201-595617744" \ +# BASE_OUTPUT_DIRECTORY="gs://your-bucket/maxtext-logs" DATASET_PATH="gs://your-bucket/maxtext-datasets" \ # STEPS="100" DILOCO_SYNC_PERIOD="49" DILOCO_NUM_FRAGMENTS="49" \ # bash src/maxtext/trainers/diloco/scripts/run_spmd_streaming_dlco_moe_c4.sh # ------------------------------------------------------------------------------ -CLUSTER="${CLUSTER:-mlperf-v5p}" -PROJECT="${PROJECT:-cloud-tpu-multipod-dev}" -ZONE="${ZONE:-europe-west4-b}" +# Required: GKE cluster name, project, and cluster location +# (LOCATION is the cluster's region for regional clusters, or its zone for zonal ones). +for var in CLUSTER PROJECT LOCATION; do + if [ -z "${!var:-}" ]; then + echo "Error: ${var} is not set. Please set CLUSTER, PROJECT and LOCATION (e.g. CLUSTER=your-cluster-name PROJECT=your-project-id LOCATION=your-cluster-location)." + exit 1 + fi +done NUM_SLICES="${NUM_SLICES:-2}" DEVICE_TYPE="${DEVICE_TYPE:-v5p-128}" -RUNNAME="${RUNNAME:-dlco-moe-$(date +%H%M)}" -DOCKER_IMAGE_BASE="${DOCKER_IMAGE_BASE:-gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:latest}" +RUNNAME="${RUNNAME:-dlco-moe-$(date +%H%M)}" # also used as the gcluster workload name +if [ "${#RUNNAME}" -gt 28 ]; then + echo "Error: RUNNAME '${RUNNAME}' exceeds the 28-character limit for gcluster workload names." + exit 1 +fi +DOCKER_IMAGE_BASE="${DOCKER_IMAGE_BASE:-us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_pre_training:latest}" MY_IMAGE="gcr.io/${PROJECT}/$(whoami)-runner:${RUNNAME}" if [ -z "${BASE_OUTPUT_DIRECTORY:-}" ]; then @@ -151,24 +163,21 @@ EOF echo "Pushing docker image: ${MY_IMAGE}" docker push "${MY_IMAGE}" -echo "Submitting MoE DiLoCo training workload using XPK..." -XPK_ARGS=( - --workload "${RUNNAME}" - --docker-image "${MY_IMAGE}" +echo "Submitting MoE DiLoCo training workload using Cluster Toolkit (gcluster)..." +GCLUSTER_ARGS=( + --name "${RUNNAME}" + --image "${MY_IMAGE}" --command "${CMD}" --num-slices "${NUM_SLICES}" --priority "${PRIORITY:-medium}" - --enable-debug-logs + --restarts "${MAX_RESTARTS:-0}" + --verbose --cluster "${CLUSTER}" - --tpu-type "${DEVICE_TYPE}" + --compute-type "${DEVICE_TYPE}" --project "${PROJECT}" - --zone "${ZONE}" + --location "${LOCATION}" ) -if [ -n "${RESERVATION:-}" ] && [ "${RESERVATION}" != "NONE" ]; then - XPK_ARGS+=(--reservation "${RESERVATION}") -fi - -xpk workload create "${XPK_ARGS[@]}" +gcluster job submit "${GCLUSTER_ARGS[@]}" echo "MoE DiLoCo workload submission complete!" diff --git a/src/maxtext/trainers/diloco/utils/fragmenter.py b/src/maxtext/trainers/diloco/utils/fragmenter.py index 2bd6c976ae..843a108ad4 100644 --- a/src/maxtext/trainers/diloco/utils/fragmenter.py +++ b/src/maxtext/trainers/diloco/utils/fragmenter.py @@ -17,15 +17,99 @@ A **fragment** is a disjoint subset of the model parameter PyTree (e.g., embeddings/head or a block of decoder layers). While standard DiLoCo synchronizes the entire model at once, Streaming DiLoCo pipelines cross-island communication by synchronizing one fragment per inner step to overlap inter-cluster communication with computation. + +Fragment 0 holds the non-scanned parameters; fragments 1..N-1 each hold a subset of the scanned decoder layers. +With `diloco_bucketize_non_scanned`, non-scanned matrices (ndim >= 2, e.g. the embedding table and the output +projection) are instead split across the layer fragments along their longest unsharded axis that has at least N-1 +indices, so that every fragment carries a similar number of bytes. There is no size threshold: any such matrix is +split, whatever its size. A matrix whose long axes are all sharded stays whole in fragment 0. Bucketization applies to +every streaming DiLoCo runner that builds its fragments with this class. """ import re -from typing import Any +from typing import Any, Iterator, NamedTuple import jax import jax.numpy as jnp import numpy as np +from maxtext.utils import max_logging + +# Suffixes of flat-fragment keys that carry a slice of a bucketized non-scanned leaf. +BUCKET_KEY_SUFFIX = "#bucket" +REMAINDER_KEY_SUFFIX = "#remainder" + + +class BucketSpec(NamedTuple): + """Split of a non-scanned leaf along one of its unsharded axes across the layer fragments. + + Layer fragment `f` (1-based) carries indices `[(f - 1) * chunk_size, f * chunk_size)` of `axis` (the leaf's longest + unsharded axis, e.g. vocab for both the embedding table and the output head under FSDP); the trailing `remainder` + indices stay in fragment 0. + """ + + chunk_size: int + remainder: int + axis: int = 0 + + +def _slice_axis(leaf: Any, axis: int, start: int, size: int) -> Any: + """Static slice `[start, start + size)` along `axis`; also accepts `jax.ShapeDtypeStruct`.""" + if isinstance(leaf, jax.ShapeDtypeStruct): + shape = list(leaf.shape) + shape[axis] = size + return jax.ShapeDtypeStruct(tuple(shape), leaf.dtype, sharding=getattr(leaf, "sharding", None)) + return jax.lax.slice_in_dim(leaf, start, start + size, axis=axis) + + +def _is_sharding_leaf(x: Any) -> bool: + return x is None or isinstance(x, jax.sharding.Sharding) + + +def _sharding_leaves(kvs: list[tuple[Any, Any]], shardings: Any) -> list[Any]: + """Per-leaf shardings: from `shardings` if given, else from the leaves themselves (None for tracers). + + Args: + kvs: `(keypath, leaf)` pairs of the parameter tree, as given by `jax.tree_util.tree_flatten_with_path`. + shardings: Optional PyTree of shardings with the structure of the parameter tree. + """ + if shardings is not None: + sharding_kvs = jax.tree_util.tree_flatten_with_path(shardings, is_leaf=_is_sharding_leaf)[0] + if len(sharding_kvs) != len(kvs): + raise ValueError(f"shardings has {len(sharding_kvs)} leaves but the parameter tree has {len(kvs)}.") + for (param_path, _), (sharding_path, _) in zip(kvs, sharding_kvs): + param_key, sharding_key = jax.tree_util.keystr(param_path), jax.tree_util.keystr(sharding_path) + if param_key != sharding_key: + raise ValueError( + f"shardings does not have the structure of the parameter tree: found {sharding_key} in shardings where" + f" the parameter tree has {param_key}." + ) + return [s for _, s in sharding_kvs] + # Tracers carry no sharding (accessing it raises), so traced leaves count as unknown -> unsharded. + return [None if isinstance(leaf, jax.core.Tracer) else getattr(leaf, "sharding", None) for _, leaf in kvs] + + +def _axis_shard_counts(shape: tuple[int, ...], sharding: Any) -> list[int]: + """Number of shards along each axis of `shape` under `sharding`; 1 everywhere if it is not a `NamedSharding`. + + Entries of the PartitionSpec that are neither mesh axis names nor tuples of them (e.g. `UNCONSTRAINED`) give 0, so + such an axis is never chosen for splitting. + """ + if not isinstance(sharding, jax.sharding.NamedSharding): + return [1] * len(shape) + spec = tuple(sharding.spec) + (None,) * (len(shape) - len(sharding.spec)) + counts = [] + for entry in spec[: len(shape)]: + if entry is None: + counts.append(1) + elif isinstance(entry, str): + counts.append(int(sharding.mesh.shape[entry])) + elif isinstance(entry, tuple): + counts.append(int(np.prod([sharding.mesh.shape[a] for a in entry]))) + else: + counts.append(0) + return counts + class FragmentedTreeManipulator: """For Streaming DiLoCo: Partitions and manipulates fragments of a JAX PyTree, supporting scanned layers.""" @@ -37,18 +121,31 @@ def __init__( num_fragments: int, param_scan_axis: int = 0, leaf_keystrs: list[str] | None = None, - fragment_to_leaf_indices: dict[int, list[int]] | None = None, + bucketized_leaves: dict[str, BucketSpec] | None = None, ): self.keypath_to_is_scanned = keypath_to_is_scanned self.fragment_to_layer_indices = fragment_to_layer_indices self.num_fragments = num_fragments self.param_scan_axis = param_scan_axis self.leaf_keystrs = leaf_keystrs or [] - self.fragment_to_leaf_indices = fragment_to_leaf_indices or {} + self.bucketized_leaves = bucketized_leaves or {} + + @property + def num_layer_fragments(self) -> int: + return self.num_fragments - 1 @classmethod - def create(cls, params_tree: Any, config: Any) -> "FragmentedTreeManipulator": - """Creates a FragmentedTreeManipulator from the parameters PyTree and configuration.""" + def create(cls, params_tree: Any, config: Any, shardings: Any = None) -> "FragmentedTreeManipulator": + """Creates a FragmentedTreeManipulator from the parameters PyTree and configuration. + + Args: + params_tree: Parameter PyTree (concrete arrays, `jax.ShapeDtypeStruct`s or tracers). + config: Config with the DiLoCo fragment settings. + shardings: Optional PyTree of `NamedSharding`s with the structure of `params_tree`. Only used to pick the split + axis of bucketized leaves. When omitted, the shardings of concrete (non-traced) leaves are used; leaves + without a `NamedSharding` are treated as unsharded. Callers that build the manipulator inside a traced + function (SPMD streaming DiLoCo) must pass it, since tracers carry no sharding. + """ kvs, _ = jax.tree_util.tree_flatten_with_path(params_tree) num_layers = config.num_decoder_layers @@ -101,12 +198,36 @@ def create(cls, params_tree: Any, config: Any) -> "FragmentedTreeManipulator": ) keypath_to_is_scanned[keystr] = is_scanned - fragment_to_leaf_indices = {} - for f in range(num_fragments): - if f == 0: - fragment_to_leaf_indices[0] = [i for i, k in enumerate(leaf_keystrs) if not keypath_to_is_scanned.get(k, False)] - else: - fragment_to_leaf_indices[f] = [i for i, k in enumerate(leaf_keystrs) if keypath_to_is_scanned.get(k, False)] + # Matrices with an unsharded axis long enough to give every layer fragment at least one index are split along + # the longest such axis (vocab for both the embedding table (vocab, embed) and the output head (embed, vocab) + # when vocab is unsharded). Slicing a sharded axis would make XLA gather the whole leaf on every fragment, so a + # leaf whose long axes are all sharded stays whole in fragment 0, as do vectors (e.g. norm scales). + bucketized_leaves = {} + if getattr(config, "diloco_bucketize_non_scanned", False): + sharding_leaves = _sharding_leaves(kvs, shardings) + if not any(isinstance(s, jax.sharding.NamedSharding) for s in sharding_leaves): + max_logging.warning( + "DiLoCo bucketization found no NamedSharding for any parameter (the given `shardings` are all None, or" + " none were given and the leaves are tracers, host arrays or not placed with a NamedSharding); every axis" + " is treated as unsharded, so a sharded axis may be split." + ) + for (keypath, v), sharding in zip(kvs, sharding_leaves): + keystr = jax.tree_util.keystr(keypath) + shape = tuple(getattr(v, "shape", ())) + if keypath_to_is_scanned[keystr] or len(shape) < 2: + continue + long_axes = [a for a in range(len(shape)) if shape[a] >= num_transformer_fragments] + shard_counts = _axis_shard_counts(shape, sharding) + candidates = [a for a in long_axes if shard_counts[a] == 1] + if candidates: + axis = max(candidates, key=shape.__getitem__) + bucketized_leaves[keystr] = BucketSpec(*divmod(shape[axis], num_transformer_fragments), axis=axis) + max_logging.log(f"DiLoCo: bucketizing {keystr} {shape} -> {bucketized_leaves[keystr]}") + elif long_axes: + max_logging.log( + f"DiLoCo: not bucketizing {keystr} {shape}: every axis with >= {num_transformer_fragments} indices is" + f" sharded (shards per axis {shard_counts}); it stays whole in fragment 0." + ) return cls( keypath_to_is_scanned=keypath_to_is_scanned, @@ -114,73 +235,70 @@ def create(cls, params_tree: Any, config: Any) -> "FragmentedTreeManipulator": num_fragments=num_fragments, param_scan_axis=param_scan_axis, leaf_keystrs=leaf_keystrs, - fragment_to_leaf_indices=fragment_to_leaf_indices, + bucketized_leaves=bucketized_leaves, ) - def get_flat_fragment(self, tree: Any, fragment_idx: int, has_replica_dim: bool = False) -> dict[str, Any]: - """Extracts a flat dictionary containing parameters for the specified fragment index.""" - flat_frag = {} + def _iter_keyed_leaves(self, tree: Any) -> Iterator[tuple[str, Any]]: + """Yields `(keystr, leaf)`, reusing the cached key strings when the tree matches the construction tree.""" leaves = jax.tree_util.tree_leaves(tree) - if len(leaves) == len(self.leaf_keystrs): - leaf_indices = self.fragment_to_leaf_indices.get(fragment_idx, []) - if fragment_idx == 0: - for idx in leaf_indices: - flat_frag[self.leaf_keystrs[idx]] = leaves[idx] - return flat_frag - - raw_indices = self.fragment_to_layer_indices.get(fragment_idx, (fragment_idx - 1,)) - layer_indices = tuple(int(x) for x in raw_indices) - is_contiguous = len(layer_indices) > 0 and ( - list(layer_indices) == list(range(layer_indices[0], layer_indices[-1] + 1)) - ) + yield from zip(self.leaf_keystrs, leaves) + else: + for keypath, leaf in jax.tree_util.tree_flatten_with_path(tree)[0]: + yield jax.tree_util.keystr(keypath), leaf - for idx in leaf_indices: - keystr = self.leaf_keystrs[idx] - v = leaves[idx] - axis = self.param_scan_axis + 1 if has_replica_dim and v.ndim > self.param_scan_axis + 1 else self.param_scan_axis - if isinstance(v, jax.ShapeDtypeStruct): - new_shape = list(v.shape) - new_shape[axis] = len(layer_indices) - shd = getattr(v, "sharding", None) - flat_frag[keystr] = jax.ShapeDtypeStruct(tuple(new_shape), v.dtype, sharding=shd) - elif is_contiguous: - slc = [slice(None)] * v.ndim - slc[axis] = slice(layer_indices[0], layer_indices[-1] + 1) - flat_frag[keystr] = v[tuple(slc)] + def _scan_axis(self, leaf: Any, has_replica_dim: bool) -> int: + return self.param_scan_axis + 1 if has_replica_dim and leaf.ndim > self.param_scan_axis + 1 else self.param_scan_axis + + def _bucket_rows(self, keystr: str, fragment_idx: int) -> tuple[int, int] | None: + """Returns `(start, size)` along the split axis of bucketized leaf `keystr` carried by `fragment_idx`, if any.""" + spec = self.bucketized_leaves[keystr] + if fragment_idx > 0: + return (fragment_idx - 1) * spec.chunk_size, spec.chunk_size + if spec.remainder: + return self.num_layer_fragments * spec.chunk_size, spec.remainder + return None + + def _bucket_axis(self, keystr: str, has_replica_dim: bool) -> int: + return self.bucketized_leaves[keystr].axis + (1 if has_replica_dim else 0) + + def _layer_indices(self, fragment_idx: int) -> tuple[tuple[int, ...], bool]: + """Returns the layer indices of `fragment_idx` and whether they are contiguous.""" + layer_indices = tuple(int(x) for x in self.fragment_to_layer_indices.get(fragment_idx, (fragment_idx - 1,))) + is_contiguous = len(layer_indices) > 0 and ( + list(layer_indices) == list(range(layer_indices[0], layer_indices[-1] + 1)) + ) + return layer_indices, is_contiguous + + def _check_fragment_idx(self, fragment_idx: int) -> None: + if not 0 <= fragment_idx < self.num_fragments: + raise ValueError(f"fragment_idx ({fragment_idx}) must be in [0, {self.num_fragments}).") + + def get_flat_fragment(self, tree: Any, fragment_idx: int, has_replica_dim: bool = False) -> dict[str, Any]: + """Extracts a flat dictionary containing parameters for the specified fragment index. + + Keys are leaf key strings; row slices of bucketized leaves use `BUCKET_KEY_SUFFIX` (layer fragments) or + `REMAINDER_KEY_SUFFIX` (fragment 0). + """ + self._check_fragment_idx(fragment_idx) + flat_frag = {} + layer_indices, is_contiguous = self._layer_indices(fragment_idx) if fragment_idx > 0 else ((), False) + for keystr, v in self._iter_keyed_leaves(tree): + if self.keypath_to_is_scanned.get(keystr, False): + if fragment_idx == 0: + continue + axis = self._scan_axis(v, has_replica_dim) + if isinstance(v, jax.ShapeDtypeStruct) or is_contiguous: + flat_frag[keystr] = _slice_axis(v, axis, layer_indices[0], len(layer_indices)) else: flat_frag[keystr] = jnp.take(v, np.array(layer_indices, dtype=np.int32), axis=axis) - return flat_frag - - # Fallback if tree structure does not match leaf count - kvs, _ = jax.tree_util.tree_flatten_with_path(tree) - for k, v in kvs: - keystr = jax.tree_util.keystr(k) - is_scanned = self.keypath_to_is_scanned.get(keystr, False) - if fragment_idx == 0: - if not is_scanned: - flat_frag[keystr] = v - else: - if is_scanned: - raw_indices = self.fragment_to_layer_indices.get(fragment_idx, (fragment_idx - 1,)) - layer_indices = tuple(int(x) for x in raw_indices) - is_contiguous = len(layer_indices) > 0 and ( - list(layer_indices) == list(range(layer_indices[0], layer_indices[-1] + 1)) - ) - axis = ( - self.param_scan_axis + 1 if has_replica_dim and v.ndim > self.param_scan_axis + 1 else self.param_scan_axis - ) - if isinstance(v, jax.ShapeDtypeStruct): - new_shape = list(v.shape) - new_shape[axis] = len(layer_indices) - shd = getattr(v, "sharding", None) - flat_frag[keystr] = jax.ShapeDtypeStruct(tuple(new_shape), v.dtype, sharding=shd) - elif is_contiguous: - slc = [slice(None)] * v.ndim - slc[axis] = slice(layer_indices[0], layer_indices[-1] + 1) - flat_frag[keystr] = v[tuple(slc)] - else: - flat_frag[keystr] = jnp.take(v, np.array(layer_indices, dtype=np.int32), axis=axis) + elif keystr in self.bucketized_leaves: + rows = self._bucket_rows(keystr, fragment_idx) + if rows is not None: + suffix = BUCKET_KEY_SUFFIX if fragment_idx > 0 else REMAINDER_KEY_SUFFIX + flat_frag[keystr + suffix] = _slice_axis(v, self._bucket_axis(keystr, has_replica_dim), *rows) + elif fragment_idx == 0: + flat_frag[keystr] = v return flat_frag def apply_flat_fragment( @@ -190,78 +308,35 @@ def apply_flat_fragment( flat_fragment: dict[str, Any], has_replica_dim: bool = False, ) -> Any: - """Merges a flat fragment dictionary back into the full parameters PyTree structure.""" - leaves, treedef = jax.tree_util.tree_flatten(tree) - new_leaves = list(leaves) - - if len(leaves) == len(self.leaf_keystrs): - leaf_indices = self.fragment_to_leaf_indices.get(fragment_idx, []) - if fragment_idx == 0: - for idx in leaf_indices: - keystr = self.leaf_keystrs[idx] - if keystr in flat_fragment: - new_leaves[idx] = flat_fragment[keystr] - return jax.tree_util.tree_unflatten(treedef, new_leaves) - - raw_indices = self.fragment_to_layer_indices.get(fragment_idx, (fragment_idx - 1,)) - layer_indices = tuple(int(x) for x in raw_indices) - is_contiguous = len(layer_indices) > 0 and ( - list(layer_indices) == list(range(layer_indices[0], layer_indices[-1] + 1)) - ) + """Merges a flat fragment dictionary (as produced by `get_flat_fragment`) back into the full PyTree. - for idx in leaf_indices: - keystr = self.leaf_keystrs[idx] - if keystr not in flat_fragment: - continue - v = leaves[idx] - axis = self.param_scan_axis + 1 if has_replica_dim and v.ndim > self.param_scan_axis + 1 else self.param_scan_axis - if isinstance(v, jax.ShapeDtypeStruct): - continue - if is_contiguous: - slc = [slice(None)] * v.ndim - slc[axis] = slice(layer_indices[0], layer_indices[-1] + 1) - new_leaves[idx] = v.at[tuple(slc)].set(flat_fragment[keystr]) - else: - new_leaves[idx] = v.at[ - tuple(slice(None) if i != axis else np.array(layer_indices, dtype=np.int32) for i in range(v.ndim)) - ].set(flat_fragment[keystr]) - return jax.tree_util.tree_unflatten(treedef, new_leaves) - - # Fallback if tree structure does not match leaf count - kvs, treedef = jax.tree_util.tree_flatten_with_path(tree) - new_kvs = [] - for k, v in kvs: - keystr = jax.tree_util.keystr(k) - is_scanned = self.keypath_to_is_scanned.get(keystr, False) - if fragment_idx == 0: - if not is_scanned and keystr in flat_fragment: - new_kvs.append(flat_fragment[keystr]) - else: - new_kvs.append(v) - else: - if is_scanned and keystr in flat_fragment: - raw_indices = self.fragment_to_layer_indices.get(fragment_idx, (fragment_idx - 1,)) - layer_indices = tuple(int(x) for x in raw_indices) - is_contiguous = len(layer_indices) > 0 and ( - list(layer_indices) == list(range(layer_indices[0], layer_indices[-1] + 1)) - ) - axis = ( - self.param_scan_axis + 1 if has_replica_dim and v.ndim > self.param_scan_axis + 1 else self.param_scan_axis - ) - if isinstance(v, jax.ShapeDtypeStruct): - new_v = v - elif is_contiguous: - slc = [slice(None)] * v.ndim - slc[axis] = slice(layer_indices[0], layer_indices[-1] + 1) - new_v = v.at[tuple(slc)].set(flat_fragment[keystr]) + Slices written into scanned or bucketized leaves are cast to the dtype of the leaf they are written into. + """ + self._check_fragment_idx(fragment_idx) + _, treedef = jax.tree_util.tree_flatten(tree) + layer_indices, is_contiguous = self._layer_indices(fragment_idx) if fragment_idx > 0 else ((), False) + new_leaves = [] + for keystr, v in self._iter_keyed_leaves(tree): + if self.keypath_to_is_scanned.get(keystr, False): + if fragment_idx > 0 and keystr in flat_fragment and not isinstance(v, jax.ShapeDtypeStruct): + axis = self._scan_axis(v, has_replica_dim) + update = jnp.asarray(flat_fragment[keystr], dtype=v.dtype) + if is_contiguous: + v = jax.lax.dynamic_update_slice_in_dim(v, update, layer_indices[0], axis=axis) else: - new_v = v.at[ - tuple(slice(None) if i != axis else np.array(layer_indices, dtype=np.int32) for i in range(v.ndim)) - ].set(flat_fragment[keystr]) - new_kvs.append(new_v) - else: - new_kvs.append(v) - return jax.tree_util.tree_unflatten(treedef, new_kvs) + index = tuple(slice(None) if i != axis else np.array(layer_indices, dtype=np.int32) for i in range(v.ndim)) + v = v.at[index].set(update) + elif keystr in self.bucketized_leaves: + suffix = BUCKET_KEY_SUFFIX if fragment_idx > 0 else REMAINDER_KEY_SUFFIX + rows = self._bucket_rows(keystr, fragment_idx) + if rows is not None and keystr + suffix in flat_fragment and not isinstance(v, jax.ShapeDtypeStruct): + axis = self._bucket_axis(keystr, has_replica_dim) + update = jnp.asarray(flat_fragment[keystr + suffix], dtype=v.dtype) + v = jax.lax.dynamic_update_slice_in_dim(v, update, rows[0], axis=axis) + elif fragment_idx == 0 and keystr in flat_fragment: + v = flat_fragment[keystr] + new_leaves.append(v) + return jax.tree_util.tree_unflatten(treedef, new_leaves) def get_streaming_schedule(config: Any) -> tuple[int, int]: diff --git a/src/maxtext/trainers/post_train/distillation/README.md b/src/maxtext/trainers/post_train/distillation/README.md index 768098099d..8e980bcee7 100644 --- a/src/maxtext/trainers/post_train/distillation/README.md +++ b/src/maxtext/trainers/post_train/distillation/README.md @@ -6,8 +6,11 @@ operational — for concepts, see [`knowledge_distillation.md`](../../../../../docs/tutorials/posttraining/knowledge_distillation.md) and [`post_training_index.md`](../../../../../docs/tutorials/post_training_index.md). -Canonical launcher: [`scripts/run_distill_xpk.sh`](scripts/run_distill_xpk.sh) -(see its header for all env vars). +Canonical launchers: +- Cluster Toolkit (recommended): [`scripts/run_distill_ctk.sh`](scripts/run_distill_ctk.sh) +- Legacy XPK: [`scripts/run_distill_xpk.sh`](scripts/run_distill_xpk.sh) + +See each script's header for all supported environment variables, and see [`knowledge_distillation.md`](../../../../../docs/tutorials/posttraining/knowledge_distillation.md#multi-host-on-gke-via-cluster-toolkit) for step-by-step Cluster Toolkit usage. ## 1. Pick a config @@ -34,11 +37,11 @@ HBM — ≥ v5p-16 in practice; a v5p-8 only fits with bf16 weights or a shrunken student (see below): ```bash -source /bin/activate +source /bin/activate PYTHONPATH=$PWD/src python -m maxtext.trainers.post_train.distillation.train_distill \ src/maxtext/configs/post_train/distillation.yml \ run_name=local_smoke \ - base_output_directory=gs:///distill_smoke \ + base_output_directory=gs:///distill_smoke \ steps=5 ``` @@ -55,8 +58,8 @@ pip install git+https://github.com/AI-Hypercomputer/xpk.git # Kubeconfig (use --dns-endpoint; IP endpoints are often stale). # Use --zone for zonal clusters, --region for regional ones. -gcloud container clusters get-credentials \ - --zone= --project= --dns-endpoint +gcloud container clusters get-credentials \ + --zone= --project= --dns-endpoint # Verify RBAC in the default namespace: kubectl auth can-i create roles --namespace=default # must print: yes @@ -72,7 +75,7 @@ The flow is: build the MaxText base → `prep_image` rebuilds `$XPK_BASE_IMAGE` bakes the workspace `./src` into a runner image and pushes it to GCR. ```bash -export XPK_PROJECT= +export XPK_PROJECT= export XPK_BASE_IMAGE=maxtext_base_image # local tag prep_image rebuilds export XPK_RUNNER_IMAGE_NAME=maxtext_base_image # short name pushed under gcr.io/$XPK_PROJECT/ # XPK_RUNNER_IMAGE_TAG defaults to ${USER}-distill; override (or set USER) @@ -98,13 +101,13 @@ layout and xpk's crane-overlay fallback. ## 5. Submit ```bash -export XPK_CLUSTER= -export XPK_PROJECT= -export XPK_ZONE= +export XPK_CLUSTER= +export XPK_PROJECT= +export XPK_ZONE= export XPK_DEVICE_TYPE=tpu7x-4x4x4 export XPK_BASE_IMAGE=gcr.io/${XPK_PROJECT}/${XPK_RUNNER_IMAGE_NAME}:${XPK_RUNNER_IMAGE_TAG:-${USER}-distill} # slash → --docker-image auto-selected -export XPK_BASE_OUTPUT_DIR=gs:///distillation -export XPK_RUN_NAME= # default: distill_run; set per experiment +export XPK_BASE_OUTPUT_DIR=gs:///distillation +export XPK_RUN_NAME= # default: distill_run; set per experiment # to scope checkpoints + TB under # ${XPK_BASE_OUTPUT_DIR}/${XPK_WORKLOAD}/${XPK_RUN_NAME}/ diff --git a/src/maxtext/trainers/post_train/distillation/distillation_utils.py b/src/maxtext/trainers/post_train/distillation/distillation_utils.py index da8b02cfed..2f56f9a44b 100644 --- a/src/maxtext/trainers/post_train/distillation/distillation_utils.py +++ b/src/maxtext/trainers/post_train/distillation/distillation_utils.py @@ -509,7 +509,7 @@ def compute_loss( # 2. Gather Student unnormalized logits at the Teacher's exact Top-K indices s_logits_sparse = jnp.take_along_axis( s_logits, teacher_output.top_k_indices, axis=-1 # pyrefly: ignore[bad-argument-type] - ) # pyrefly: ignore[bad-argument-type] + ) # 3. Normalize Student probabilities only over the exact same Top-K subset log_s_T_sparse = jax.nn.log_softmax(s_logits_sparse / temperature, axis=-1) @@ -561,7 +561,7 @@ def compute_loss( feature_loss = beta_feature * self.feature_loss_fn( # pyrefly: ignore[not-callable] s_features_sliced, t_features_sliced, mask - ) # pyrefly: ignore[not-callable] + ) total_loss = base_logit_loss + feature_loss @@ -733,7 +733,7 @@ def maybe_restore( # pyrefly: ignore[bad-override] target_model = getattr(model, "student_model", model) step, custom_metadata = super().maybe_restore( - model=target_model, # pyrefly: ignore[bad-argument-type] + model=target_model, optimizer=optimizer, restore_only_lora_params=restore_only_lora_params, ) @@ -780,4 +780,4 @@ def close(self) -> None: super().close() if getattr(self, "_checkpointer", None) is not None: if hasattr(self._checkpointer, "close"): - self._checkpointer.close() # pytype: disable=attribute-error + self._checkpointer.close() diff --git a/src/maxtext/trainers/post_train/distillation/scripts/run_distill_ctk.sh b/src/maxtext/trainers/post_train/distillation/scripts/run_distill_ctk.sh new file mode 100755 index 0000000000..9de444cd76 --- /dev/null +++ b/src/maxtext/trainers/post_train/distillation/scripts/run_distill_ctk.sh @@ -0,0 +1,460 @@ +#!/bin/bash +# +# 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. +# +# Reference launcher for MaxText distillation on a GKE TPU cluster via +# Cluster Toolkit (gcluster). +# Treat this as a starting template — copy it, adapt the env vars and +# `--command` body for your cluster + run, and submit from CI / a tmux / screen. +# +# The script expects a base image at $DOCKER_IMAGE (or $CTK_BASE_IMAGE). `prep_image` +# (below) builds the MaxText TPU post-training Docker image as described in +# https://maxtext.readthedocs.io/page/tutorials/build_maxtext.html#tpu-post-training-docker-image +# (equivalent to running `build_maxtext_docker_image WORKFLOW=post-training` from +# an activated MaxText virtual environment). Post-training dependencies such as +# Tunix come from the pins in src/dependencies/extra_deps/post_train_github_deps.txt. +# Then bake your local `./src` into a runner image pushed to GCR via `upload_runner`. +# +# Usage: +# bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_ctk.sh prep_image # one-time image build +# bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_ctk.sh upload_runner # bake workspace + push to GCR +# bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_ctk.sh submit # fire-and-forget job submission +# bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_ctk.sh monitor # stream logs for the last submit +# bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_ctk.sh cleanup # cancel/delete the last workload +# bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_ctk.sh resume_until_done # auto-retry loop for long jobs +# +# Reading logs on-demand: +# - GCP Logs Explorer URL that `gcluster job submit` prints at submit time. +# - `gcluster job logs --main-only=false` +# - `kubectl logs -c jax-tpu --tail=100` for a quick ad-hoc look. +# - `kubectl logs --previous` for a crashed container's last lines. +# +# Watch pod state changes (restarts, failures) in the background: +# nohup kubectl get pods -l jobset.sigs.k8s.io/jobset-name=$CTK_WORKLOAD --watch \ +# > ~/distill-events.log 2>&1 & +# disown +# # Then: grep -iE 'error|crashloop|restart' ~/distill-events.log +# +# Restart policy (three layers): +# L1 pod-level GKE restarts a crashed pod automatically; the trainer resumes +# from the latest checkpoint via `maybe_restore` — no action needed. +# L2 workload-level when the whole JobSet terminates, `resume_until_done` cancels + re-submits +# with the same CTK_WORKLOAD/BASE_OUTPUT_DIRECTORY (checkpoint resume, again). +# L3 human `resume_until_done` exits non-zero after MAX_RETRIES; page yourself from there. +# +# REQUIRED env vars (no defaults; the script exits with an error if unset): +# GKE_CLUSTER (or CTK_CLUSTER) GKE cluster name +# PROJECT_ID (or CTK_PROJECT) GCP project hosting the cluster +# LOCATION (or CTK_ZONE / ZONE) cluster region or zone (e.g. europe-west4, us-central1-a) +# COMPUTE_TYPE (or CTK_COMPUTE_TYPE) e.g. tpu7x-standard-4t, ct5p-hightpu-4t, ct6e-hightpu-4t +# TOPOLOGY (or CTK_TOPOLOGY) e.g. 4x4x4, 16x16 +# BASE_OUTPUT_DIRECTORY (or CTK_BASE_OUTPUT_DIR) +# GCS prefix for run outputs (each run writes to +# ${BASE_OUTPUT_DIRECTORY}/${CTK_WORKLOAD}/) +# +# OPTIONAL env vars (with defaults): +# DOCKER_IMAGE (or CTK_BASE_IMAGE) +# default: maxtext_base_image. A slash in the name +# (e.g. gcr.io/...) uses `--image` (pull from registry); +# otherwise uses `--base-image` with `--build-context=.`. +# Prefer the registry path after `upload_runner`. +# CTK_WORKLOAD default: d-${USER:0:8}-${RANDOM} (~14 chars max). +# Must be <= 28 chars (or <= 22 chars with Pathways) and a valid DNS label. +# CTK_PRIORITY default: medium +# NUM_SLICES (or CTK_NUM_SLICES) +# default: 1 +# CTK_DISTILL_CONFIG default: src/maxtext/configs/post_train/distillation.yml +# RUN_NAME (or CTK_RUN_NAME) +# default: distill_run — passed as MaxText run_name; becomes +# the subdir under base_output_directory where checkpoints +# and TB logs land (...${OUTPUT_DIR}/${RUN_NAME}/...). +# resume_until_done lists this subdir to find the latest step. +# CTK_USE_GCSFUSE default: 1 — mount CTK_DATASET_BUCKET via GKE GCS FUSE CSI +# (`--mount`) and point grain at the local mount path. +# Set to 0 to bypass gcsfuse and read directly from gs://. +# CTK_DATASET_BUCKET default: maxtext-dataset +# CTK_DATASET_SUBPATH default: array-record/climbmix/*.arrayrecord +# The script always sets grain_train_files from these +# two, overriding the YAML in both modes. +# CTK_HF_CACHE_DIR default: /dev/shm/hf — HF `datasets` Arrow cache dir. +# STEPS_OVERRIDE default: empty — yml `steps` is used unless set +# CHECKPOINT_PERIOD_OVERRIDE default: empty — yml `checkpoint_period` is used +# MAX_RETRIES default: 10 — only used by resume_until_done +# STUDENT_CKPT_PATH default: empty — overrides student_overrides.load_parameters_path if set +# TEACHER_CKPT_PATH default: empty — overrides teacher_overrides.load_parameters_path if set +# TOKENIZER_PATH default: empty — overrides tokenizer_path if set +# HF_TOKEN default: empty — overrides hf_access_token if set +# +# Feature-mapping / distillation loss hyperparameters (always passed to the +# trainer; override yml values). Defaults enable feature mapping on the first +# 8 layers. Set DISTILL_BETA=0.0 to disable feature mapping. +# DISTILL_ALPHA default: 0.5 +# DISTILL_TEMPERATURE default: 1.0 +# DISTILL_BETA default: 1.0 (>0 enables feature-map loss; +# requires scan_layers=True) +# DISTILL_LAYER_INDICES default: [0,1,2,3,4,5,6,7] (no spaces inside brackets) +# +# upload_runner env vars: +# CTK_RUNNER_IMAGE_NAME default: maxtext_base_image — GCR short name. +# CTK_RUNNER_IMAGE_TAG default: ${USER}-distill — per-user tag avoids +# clobbering shared :latest. Pushes to +# gcr.io/$PROJECT_ID/$CTK_RUNNER_IMAGE_NAME:$CTK_RUNNER_IMAGE_TAG. + +set -euo pipefail +MODE="${1:-submit}" + +# -------------------------- required env -------------------------- +require_env() { + local missing=() + for v in "$@"; do + [ -z "${!v:-}" ] && missing+=("$v") + done + if [ "${#missing[@]}" -gt 0 ]; then + echo "ERROR: required env vars not set: ${missing[*]}" >&2 + echo "See this script's header for descriptions and example values." >&2 + exit 1 + fi +} + +# -------------------------- normalize aliases -------------------------- +export GKE_CLUSTER="${GKE_CLUSTER:-${CTK_CLUSTER:-}}" +export PROJECT_ID="${PROJECT_ID:-${CTK_PROJECT:-}}" +export LOCATION="${LOCATION:-${CTK_LOCATION:-${CTK_ZONE:-${ZONE:-}}}}" +export COMPUTE_TYPE="${COMPUTE_TYPE:-${CTK_COMPUTE_TYPE:-}}" +export TOPOLOGY="${TOPOLOGY:-${CTK_TOPOLOGY:-}}" +export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-${CTK_BASE_OUTPUT_DIR:-}}" + +# -------------------------- defaults -------------------------- +: "${CTK_BASE_IMAGE:=${DOCKER_IMAGE:-maxtext_base_image}}" +: "${CTK_WORKLOAD:=d-${USER:0:8}-${RANDOM}}" +: "${CTK_PRIORITY:=medium}" +: "${CTK_NUM_SLICES:=${NUM_SLICES:-1}}" +: "${CTK_DISTILL_CONFIG:=src/maxtext/configs/post_train/distillation.yml}" +: "${CTK_RUN_NAME:=${RUN_NAME:-distill_run}}" +: "${CTK_USE_GCSFUSE:=1}" +: "${CTK_DATASET_BUCKET:=maxtext-dataset}" +: "${CTK_DATASET_SUBPATH:=array-record/climbmix/*.arrayrecord}" +: "${CTK_HF_CACHE_DIR:=/dev/shm/hf}" +: "${MAX_RETRIES:=10}" + +# Feature-mapping / distillation loss hyperparameters. +: "${DISTILL_ALPHA:=0.5}" +: "${DISTILL_TEMPERATURE:=1.0}" +: "${DISTILL_BETA:=1.0}" +: "${DISTILL_LAYER_INDICES:=[0,1,2,3,4,5,6,7]}" + +OUTPUT_DIR="${BASE_OUTPUT_DIRECTORY:-}" +OUTPUT_DIR="${OUTPUT_DIR%/}/${CTK_WORKLOAD}" +LAST_WORKLOAD_FILE="${CTK_LAST_WORKLOAD_FILE:-${HOME}/.ctk_last_workload}" + +# CLI overrides for the trainer. +extra_cli="distill_alpha=${DISTILL_ALPHA} \ +distill_temperature=${DISTILL_TEMPERATURE} \ +distill_beta=${DISTILL_BETA} \ +distill_layer_indices=${DISTILL_LAYER_INDICES}" +if [ -n "${STEPS_OVERRIDE:-}" ]; then + extra_cli="$extra_cli learning_rate_schedule_steps=${STEPS_OVERRIDE} steps=${STEPS_OVERRIDE}" +fi +if [ -n "${CHECKPOINT_PERIOD_OVERRIDE:-}" ]; then + extra_cli="$extra_cli checkpoint_period=${CHECKPOINT_PERIOD_OVERRIDE}" +fi +if [ -n "${STUDENT_CKPT_PATH:-}" ]; then + extra_cli="$extra_cli student_overrides.load_parameters_path=${STUDENT_CKPT_PATH}" +fi +if [ -n "${TEACHER_CKPT_PATH:-}" ]; then + extra_cli="$extra_cli teacher_overrides.load_parameters_path=${TEACHER_CKPT_PATH}" +fi +if [ -n "${TOKENIZER_PATH:-}" ]; then + extra_cli="$extra_cli tokenizer_path=${TOKENIZER_PATH} tokenizer_type=huggingface" +fi +if [ -n "${HF_TOKEN:-}" ]; then + extra_cli="$extra_cli hf_access_token=${HF_TOKEN}" +fi + +# Build grain_train_files (configs leave it empty); pick GCS FUSE mount or direct gs://. +mount_args=() +if [ "$CTK_USE_GCSFUSE" = "1" ]; then + mount_args=(--mount="gs://${CTK_DATASET_BUCKET};/tmp/gcsfuse;ro") + grain_files_override="grain_train_files=/tmp/gcsfuse/${CTK_DATASET_SUBPATH}" +else + grain_files_override="grain_train_files=gs://${CTK_DATASET_BUCKET}/${CTK_DATASET_SUBPATH}" +fi + +# Optional: stage the YAML from GCS instead of baking via upload_runner. +yaml_prelude="" +if [ -n "${CTK_YAML_GCS:-}" ]; then + yaml_prelude="gcloud storage cp \"${CTK_YAML_GCS}\" \"${CTK_DISTILL_CONFIG}\";" +fi + +# Optional: stage HF tokenizer files from GCS for models whose tokenizer isn't +# baked into the image (e.g. gpt-oss). +tokenizer_prelude="" +if [ -n "${CTK_TOKENIZER_GCS:-}" ] && [ -n "${CTK_TOKENIZER_LOCAL:-}" ]; then + tokenizer_prelude="mkdir -p \"${CTK_TOKENIZER_LOCAL}\" && gcloud storage rsync \"${CTK_TOKENIZER_GCS}\" \"${CTK_TOKENIZER_LOCAL}\";" +fi + +# Default v7x XLA flags. The default vmem limit (32 MB) is too small for +# tokamax splash backward; we need ≥60 MB. +default_libtpu_args="--xla_tpu_scoped_vmem_limit_kib=61440 \ +--xla_tpu_enable_all_experimental_scheduler_features=true \ +--xla_tpu_enable_scheduler_memory_pressure_tracking=true \ +--xla_tpu_host_transfer_overlap_limit=24 \ +--xla_tpu_aggressive_opt_barrier_removal=ENABLED \ +--xla_lhs_prioritize_async_depth_over_stall=ENABLED \ +--xla_tpu_enable_ag_backward_pipelining=true \ +--xla_should_allow_loop_variant_parameter_in_chain=ENABLED \ +--xla_should_add_loop_invariant_op_in_chain=ENABLED \ +--xla_max_concurrent_host_send_recv=100 \ +--xla_tpu_scheduler_percent_shared_memory_limit=100 \ +--xla_latency_hiding_scheduler_rerun=2" +libtpu_init_args=$(printf '%s' "${CTK_LIBTPU_INIT_ARGS:-$default_libtpu_args}" | tr -s '[:space:]' ' ') + +# -------------------------- configure_cluster -------------------------- +configure_cluster() { + if ! command -v gcluster >/dev/null 2>&1; then + echo "ERROR: gcluster CLI not found in PATH. Install Cluster Toolkit first." >&2 + exit 1 + fi + gcloud config set project "$PROJECT_ID" >/dev/null + gcloud container clusters get-credentials "$GKE_CLUSTER" \ + --location "$LOCATION" \ + --project "$PROJECT_ID" >/dev/null + gcluster job config set project "$PROJECT_ID" + gcluster job config set cluster "$GKE_CLUSTER" + gcluster job config set location "$LOCATION" +} + +# -------------------------- prep_image -------------------------- +# Builds the MaxText TPU post-training Docker image following +# https://maxtext.readthedocs.io/page/tutorials/build_maxtext.html#tpu-post-training-docker-image. +# `build_maxtext_docker_image` always produces a local image named +# `maxtext_base_image`; it is retagged as $CTK_BASE_IMAGE if that differs. +# Must be run from the MaxText repo root with the MaxText virtual environment +# activated (it provides the `build_maxtext_docker_image` console script). +prep_image() { + if ! command -v build_maxtext_docker_image >/dev/null 2>&1; then + echo "ERROR: build_maxtext_docker_image not found in PATH. Activate the MaxText virtual environment first;" >&2 + echo " see https://maxtext.readthedocs.io/page/tutorials/build_maxtext.html" >&2 + exit 1 + fi + echo "== building TPU post-training image -> ${CTK_BASE_IMAGE} ==" + # Run under sudo like the other docker calls in this script; keep PATH so the + # venv console script (and its python) is still found. + sudo env "PATH=$PATH" build_maxtext_docker_image WORKFLOW=post-training + if [ "$CTK_BASE_IMAGE" != "maxtext_base_image" ]; then + sudo docker tag maxtext_base_image "$CTK_BASE_IMAGE" + fi + # Sanity check: verify the installed shard_input carries the upstream fix. + sudo docker run --rm "$CTK_BASE_IMAGE" python -c " +import inspect, tunix +from tunix.sft import sharding_utils +src = inspect.getsource(sharding_utils.shard_input) +assert 'is_fully_addressable' in src, 'tunix install does not contain the shard_input fix' +print(f'tunix {tunix.__version__}: shard_input fix present.') +" +} + +# -------------------------- upload_runner -------------------------- +upload_runner() { + : "${CTK_RUNNER_IMAGE_NAME:=maxtext_base_image}" + : "${CTK_RUNNER_IMAGE_TAG:=${USER}-distill}" + local target="gcr.io/${PROJECT_ID}/${CTK_RUNNER_IMAGE_NAME}:${CTK_RUNNER_IMAGE_TAG}" + echo "== upload_runner -> ${target} ==" + if ! sudo docker image inspect "$CTK_BASE_IMAGE" >/dev/null 2>&1; then + echo "ERROR: base image $CTK_BASE_IMAGE not found locally. Run prep_image first." >&2 + exit 1 + fi + local runner_local="${CTK_BASE_IMAGE}__runner" + sudo docker build --no-cache \ + --build-arg "BASEIMAGE=${CTK_BASE_IMAGE}" \ + --build-arg "PACKAGE_DIR=src" \ + -f src/dependencies/dockerfiles/maxtext_runner.Dockerfile \ + -t "$runner_local" . + sudo docker tag "$runner_local" "$target" + sudo docker push "$target" + echo "Pushed: $target" +} + +# -------------------------- submit -------------------------- +submit_workload() { + configure_cluster + + echo "Workload: $CTK_WORKLOAD" + echo "Cluster: $GKE_CLUSTER ($PROJECT_ID, $LOCATION)" + echo "Compute/Topo: $COMPUTE_TYPE / $TOPOLOGY x ${CTK_NUM_SLICES} slice(s)" + echo "Image: $CTK_BASE_IMAGE" + echo "Output dir: $OUTPUT_DIR" + echo "Config: $CTK_DISTILL_CONFIG" + [ -n "$extra_cli" ] && echo "Overrides: $extra_cli" + + local image_args=() + if [[ "$CTK_BASE_IMAGE" == *"/"* ]]; then + image_args=(--image="$CTK_BASE_IMAGE") + else + image_args=(--base-image="$CTK_BASE_IMAGE" --build-context=.) + fi + echo "Image args: ${image_args[*]}" + + gcluster job submit \ + "${image_args[@]}" \ + --name="$CTK_WORKLOAD" \ + --priority="$CTK_PRIORITY" \ + --compute-type="$COMPUTE_TYPE" \ + --topology="$TOPOLOGY" \ + --num-slices="$CTK_NUM_SLICES" \ + "${mount_args[@]}" \ + --command="export PYTHONPATH=/deps/src:/app/src; \ +export BASE_OUTPUT_DIRECTORY=${OUTPUT_DIR}; \ +export LIBTPU_INIT_ARGS='${libtpu_init_args}'; \ +export TMPDIR=/dev/shm; export JAX_COMPILATION_CACHE_DIR=/dev/shm/jax_cache; \ +export HF_HOME=${CTK_HF_CACHE_DIR}; export HF_DATASETS_CACHE=${CTK_HF_CACHE_DIR}/datasets; mkdir -p ${CTK_HF_CACHE_DIR}/datasets; \ +${yaml_prelude} \ +${tokenizer_prelude} \ +python3 -m maxtext.trainers.post_train.distillation.train_distill ${CTK_DISTILL_CONFIG} \ + run_name=${CTK_RUN_NAME} \ + base_output_directory=\$BASE_OUTPUT_DIRECTORY \ + ${grain_files_override} \ + ${extra_cli} \ + save_checkpoint_on_completion=True" + + echo "$CTK_WORKLOAD" > "$LAST_WORKLOAD_FILE" +} + +# -------------------------- monitor -------------------------- +monitor_workload() { + local workload="${2:-$(cat "$LAST_WORKLOAD_FILE" 2>/dev/null || true)}" + if [ -z "$workload" ]; then + echo "ERROR: no workload to monitor (none recorded in $LAST_WORKLOAD_FILE)." >&2 + exit 1 + fi + echo "Monitoring $workload" + gcluster job list || true + kubectl get jobset -l "gcluster.google.com/workload=${workload}" || true + + # Wait for any pod to reach Running, but bail if all pods finish without ever running. + while true; do + local phases + phases=$(kubectl get pods -l "jobset.sigs.k8s.io/jobset-name=${workload}" \ + -o jsonpath='{.items[*].status.phase}' 2>/dev/null || echo "") + if echo "$phases" | grep -q Running; then break; fi + if [ -n "$phases" ] && ! echo "$phases" | grep -qE "Pending|ContainerCreating"; then + echo "No Running pod; phases: $phases" + local one + one=$(kubectl get pods -l "jobset.sigs.k8s.io/jobset-name=${workload}" \ + -o name 2>/dev/null | head -1) + [ -n "$one" ] && kubectl logs "$one" --all-containers 2>&1 | tail -40 + exit 1 + fi + echo "waiting for Running (phases: ${phases:-none})..."; sleep 15 + done + + local n + n=$(kubectl get pods -l "jobset.sigs.k8s.io/jobset-name=${workload}" -o name | wc -l) + local max=$((n * 4)) # raise above 2 x num_pods + local out_log="${HOME}/training-${workload}.log" + echo "Streaming logs from $n pods (--max-log-requests=${max}) -> ${out_log}" + kubectl logs -f -l "jobset.sigs.k8s.io/jobset-name=${workload}" \ + --all-containers --max-log-requests="$max" --prefix \ + 2>&1 | tee "$out_log" +} + +# -------------------------- cleanup -------------------------- +cleanup_workload() { + local workload="${2:-$(cat "$LAST_WORKLOAD_FILE" 2>/dev/null || true)}" + if [ -z "$workload" ]; then + echo "ERROR: no workload to clean up (none recorded in $LAST_WORKLOAD_FILE)." >&2 + exit 1 + fi + echo "Canceling workload $workload" + gcluster job cancel "$workload" || kubectl delete jobset "$workload" --ignore-not-found +} + +# -------------------------- resume_until_done -------------------------- +resume_until_done() { + if [ -z "${STEPS_OVERRIDE:-}" ]; then + echo "ERROR: STEPS_OVERRIDE must be set so the loop knows when to stop." >&2 + exit 1 + fi + local target="$STEPS_OVERRIDE" + local retry=0 + + while [ "$retry" -lt "$MAX_RETRIES" ]; do + echo "=== resume attempt $((retry + 1)) / $MAX_RETRIES (target steps: $target) ===" + submit_workload + + while true; do + sleep 60 + local terminal + terminal=$(kubectl get jobset "$CTK_WORKLOAD" \ + -o jsonpath='{.status.terminalState}' 2>/dev/null || echo "") + if [ -n "$terminal" ]; then + echo "Workload $CTK_WORKLOAD reached terminal state: $terminal" + break + fi + done + + # latest checkpoint step on disk (subdirs named after the step number) + local last_step + last_step=$(gcloud storage ls "${OUTPUT_DIR}/${CTK_RUN_NAME}/checkpoints/" 2>/dev/null \ + | grep -oE '/[0-9]+/$' | tr -d '/' | sort -n | tail -1) + last_step=${last_step:-0} + echo "Latest checkpoint step on disk: ${last_step}" + + if [ "$last_step" -ge "$target" ]; then + echo "Reached target step ${target}. Done." + return 0 + fi + + # Free the workload name so we can resubmit with the same name. + gcluster job cancel "$CTK_WORKLOAD" >/dev/null 2>&1 \ + || kubectl delete jobset "$CTK_WORKLOAD" --ignore-not-found >/dev/null 2>&1 \ + || true + + retry=$((retry + 1)) + echo "Resubmitting from step ${last_step} (attempt $((retry + 1))/${MAX_RETRIES}) in 60s..." + sleep 60 + done + + echo "ERROR: max_retries=${MAX_RETRIES} reached without completing ${target} steps." >&2 + return 1 +} + +# -------------------------- dispatch -------------------------- +case "$MODE" in + prep_image) + prep_image + ;; + upload_runner) + require_env PROJECT_ID # GCR target; default gcloud project is usually wrong here. + upload_runner + ;; + submit|resume_until_done) + require_env GKE_CLUSTER PROJECT_ID LOCATION COMPUTE_TYPE TOPOLOGY BASE_OUTPUT_DIRECTORY + case "$MODE" in + submit) submit_workload ;; + resume_until_done) resume_until_done ;; + esac + ;; + monitor) + monitor_workload "$@" + ;; + cleanup|delete) + cleanup_workload "$@" + ;; + *) + echo "Unknown mode: $MODE (use prep_image|upload_runner|submit|monitor|cleanup|resume_until_done)" >&2 + exit 1 + ;; +esac diff --git a/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh b/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh index fc0a31e934..1a8c2723b2 100644 --- a/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh +++ b/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh @@ -18,17 +18,15 @@ # Treat this as a starting template — copy it, adapt the env vars and # `--command` body for your cluster + run, and submit from CI / a tmux / screen. # -# The script expects a base image at $XPK_BASE_IMAGE. Build the MaxText base -# with: -# sudo bash src/dependencies/scripts/docker_build_dependency_image.sh \ -# MODE=stable WORKFLOW=post-training -# -# Then layer tunix on top via `prep_image` (below). We install tunix WITH deps -# (it needs google-metrax + kagglehub at runtime) and then force-reinstall -# jax/jaxlib/libtpu back to versions compatible with the base image's libtpu. +# The script expects a base image at $XPK_BASE_IMAGE. `prep_image` (below) +# builds the MaxText TPU post-training Docker image as described in +# https://maxtext.readthedocs.io/page/tutorials/build_maxtext.html#tpu-post-training-docker-image +# (equivalent to running `build_maxtext_docker_image WORKFLOW=post-training` from +# an activated MaxText virtual environment). Post-training dependencies such as +# Tunix come from the pins in src/dependencies/extra_deps/post_train_github_deps.txt. # # Usage: -# bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh prep_image # one-time image layering +# bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh prep_image # one-time image build # bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh upload_runner # bake workspace + push to GCR # bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh submit # fire-and-forget; returns in ~60s # bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh monitor # stream logs for the last submit @@ -103,16 +101,6 @@ # requires scan_layers=True) # DISTILL_LAYER_INDICES default: [0,1,2,3,4,5,6,7] (no spaces inside brackets) # -# Image pinning (used by prep_image): -# TUNIX_SOURCE pip-installable spec for tunix. -# default: git+https://github.com/google/tunix@348959d18a4a09c75e58a7d49aec9d8b0eb4a8b6 -# Use "google-tunix==" once a pypi release ships with the -# multi-host shard_input fix. -# JAX_PIN default: 0.10.0 — version to pin back after tunix deps resolve. -# Must be ≥ 0.10.0 (tunix's flax dep imports jax.extend.core.Effect). -# JAXLIB_PIN default: 0.10.0 -# LIBTPU_PIN default: 0.0.39 -# # upload_runner env vars: # XPK_RUNNER_IMAGE_NAME default: maxtext_base_image — GCR short name. # XPK_RUNNER_IMAGE_TAG default: ${USER}-distill — per-user tag avoids @@ -130,8 +118,8 @@ MODE="${1:-submit}" # -------------------------- required env -------------------------- # Collects all missing vars so the user fixes them in one pass, not N runs. -# Not called for `prep_image` — that mode only needs XPK_BASE_IMAGE + TUNIX_SOURCE -# + JAX_PIN family, all of which have defaults. +# Not called for `prep_image` — that mode only needs XPK_BASE_IMAGE, which has a +# default. require_env() { local missing=() for v in "$@"; do @@ -163,12 +151,6 @@ require_env() { : "${DISTILL_BETA:=1.0}" : "${DISTILL_LAYER_INDICES:=[0,1,2,3,4,5,6,7]}" -# Image pinning (used by prep_image). -: "${TUNIX_SOURCE:=git+https://github.com/google/tunix@a8d70582f1e2f1fb65973210989e0e148b5ef7ad}" -: "${JAX_PIN:=0.10.0}" -: "${JAXLIB_PIN:=0.10.0}" -: "${LIBTPU_PIN:=0.0.39}" - # Computed at top-level so both submit_workload and resume_until_done can read it. # `${:-}` keeps `set -u` happy for `prep_image`, which doesn't need XPK_BASE_OUTPUT_DIR. OUTPUT_DIR="${XPK_BASE_OUTPUT_DIR:-}" @@ -233,32 +215,25 @@ default_libtpu_args="--xla_tpu_scoped_vmem_limit_kib=61440 \ libtpu_init_args=$(printf '%s' "${XPK_LIBTPU_INIT_ARGS:-$default_libtpu_args}" | tr -s '[:space:]' ' ') # -------------------------- prep_image -------------------------- -# Adds tunix and repins jax/libtpu on top of $XPK_BASE_IMAGE, then retags -# the result as $XPK_BASE_IMAGE (the original tag is overwritten). -# Safe to re-run: pins are force-reinstalled, but layers accumulate. -# To reset to a clean base, rebuild via docker_build_dependency_image.sh. +# Builds the MaxText TPU post-training Docker image following +# https://maxtext.readthedocs.io/page/tutorials/build_maxtext.html#tpu-post-training-docker-image. +# `build_maxtext_docker_image` always produces a local image named +# `maxtext_base_image`; it is retagged as $XPK_BASE_IMAGE if that differs. +# Must be run from the MaxText repo root with the MaxText virtual environment +# activated (it provides the `build_maxtext_docker_image` console script). prep_image() { - echo "== layering on ${XPK_BASE_IMAGE} ==" - echo " tunix : ${TUNIX_SOURCE}" - echo " jax : ${JAX_PIN}" - echo " jaxlib : ${JAXLIB_PIN}" - echo " libtpu : ${LIBTPU_PIN}" - if ! sudo docker image inspect "$XPK_BASE_IMAGE" >/dev/null 2>&1; then - echo "ERROR: base image $XPK_BASE_IMAGE not found locally. Build it first:" >&2 - echo " sudo bash src/dependencies/scripts/docker_build_dependency_image.sh MODE=stable WORKFLOW=post-training" >&2 + if ! command -v build_maxtext_docker_image >/dev/null 2>&1; then + echo "ERROR: build_maxtext_docker_image not found in PATH. Activate the MaxText virtual environment first;" >&2 + echo " see https://maxtext.readthedocs.io/page/tutorials/build_maxtext.html" >&2 exit 1 fi - local tmp; tmp=$(mktemp -d) - cat > "$tmp/Dockerfile" < ${XPK_BASE_IMAGE} ==" + # Run under sudo like the other docker calls in this script; keep PATH so the + # venv console script (and its python) is still found. + sudo env "PATH=$PATH" build_maxtext_docker_image WORKFLOW=post-training + if [ "$XPK_BASE_IMAGE" != "maxtext_base_image" ]; then + sudo docker tag maxtext_base_image "$XPK_BASE_IMAGE" + fi # Sanity check: verify the installed shard_input carries the upstream fix. sudo docker run --rm "$XPK_BASE_IMAGE" python -c " import inspect, tunix @@ -270,7 +245,7 @@ print(f'tunix {tunix.__version__}: shard_input fix present.') } # -------------------------- upload_runner -------------------------- -# Bakes ./src into the layered image and pushes to +# Bakes ./src into the post-training image and pushes to # gcr.io/$XPK_PROJECT/$XPK_RUNNER_IMAGE_NAME:$XPK_RUNNER_IMAGE_TAG. # Does the build/tag/push inline rather than calling docker_upload_runner.sh, # because that script hardcodes :latest and would clobber the shared tag. diff --git a/src/maxtext/trainers/post_train/distillation/train_distill.py b/src/maxtext/trainers/post_train/distillation/train_distill.py index f8c604a5d4..324e330169 100644 --- a/src/maxtext/trainers/post_train/distillation/train_distill.py +++ b/src/maxtext/trainers/post_train/distillation/train_distill.py @@ -442,12 +442,12 @@ def _post_process_train_step(self, aux: dict[str, tuple[jax.Array, jax.Array]]) for name, value in aux.items(): if name not in self._buffered_train_metrics.additional_metrics: - self._buffered_train_metrics.additional_metrics[name] = ( # pyrefly: ignore[unsupported-operation] + self._buffered_train_metrics.additional_metrics[name] = ( [], distillation_utils.weighted_mean, - ) # pyrefly: ignore[unsupported-operation] + ) - self._buffered_train_metrics.additional_metrics[name][0].append(value) # pyrefly: ignore[bad-argument-type] + self._buffered_train_metrics.additional_metrics[name][0].append(value) # Compact per-step summary: only the metrics that change run-to-run, and # only those that are nonzero (so MoE / feature-distill terms hide when off). @@ -639,6 +639,7 @@ def build_training_components( skip_first_n_steps=student_config.skip_first_n_steps_for_profiler, profiler_steps=student_config.profiler_steps, set_profile_options=False, + enable_continuous_profiling=student_config.enable_continuous_profiling, ) metrics_logging_options = metrics_logger.MetricsLoggerOptions( @@ -914,15 +915,15 @@ def main(argv: Sequence[str]) -> None: teacher_argv = [argv[0], argv[1]] teacher_config = pyconfig.initialize(teacher_argv, **teacher_overrides) - # Batch shape (per_device_batch_size / max_target_length / gradient_accumulation_steps) + # Batch layout (per_device_batch_size / max_target_length / gradient_accumulation_steps / packing) # must be set at the YAML top level — not inside *_overrides — since student and # teacher share the input pipeline. - for batch_field in ("per_device_batch_size", "max_target_length", "gradient_accumulation_steps"): + for batch_field in ("per_device_batch_size", "max_target_length", "gradient_accumulation_steps", "packing"): s_val = getattr(student_config, batch_field) t_val = getattr(teacher_config, batch_field) if s_val != t_val: raise ValueError( - f"Distillation batch shape mismatch on '{batch_field}': " + f"Distillation batch layout mismatch on '{batch_field}': " f"student={s_val} vs teacher={t_val}. The teacher consumes batches from the " f"student-driven input pipeline, so these must agree. Set '{batch_field}' at " f"the YAML top level (not inside *_overrides) so both configs inherit it." diff --git a/src/maxtext/trainers/post_train/dpo/train_dpo.py b/src/maxtext/trainers/post_train/dpo/train_dpo.py index f071ae9f05..3836cb1348 100644 --- a/src/maxtext/trainers/post_train/dpo/train_dpo.py +++ b/src/maxtext/trainers/post_train/dpo/train_dpo.py @@ -89,6 +89,7 @@ def get_tunix_config(mt_config: MaxTextConfig) -> DPOTrainingConfig: skip_first_n_steps=mt_config.skip_first_n_steps_for_profiler, profiler_steps=mt_config.profiler_steps, set_profile_options=set_profile_options, + enable_continuous_profiling=mt_config.enable_continuous_profiling, ) return DPOTrainingConfig( diff --git a/src/maxtext/trainers/post_train/hooks.py b/src/maxtext/trainers/post_train/hooks.py index bc8666a1c5..8dd5f4d1ad 100644 --- a/src/maxtext/trainers/post_train/hooks.py +++ b/src/maxtext/trainers/post_train/hooks.py @@ -162,7 +162,7 @@ def on_eval_step_end(self, train_ctx: peft_trainer.PeftTrainer, eval_loss: float self.eval_metadata["eval_step_count"] != 0 ), "BaseTrainingHooks.on_eval_step_start() must be called before BaseTrainingHooks.on_eval_step_end()" - avg_loss = eval_loss / self.eval_metadata["eval_step_count"] + avg_loss = eval_loss if self.eval_loss_is_preaveraged() else eval_loss / self.eval_metadata["eval_step_count"] metrics = { "scalar": { "eval/total_loss": eval_loss, @@ -195,6 +195,10 @@ def on_eval_step_end(self, train_ctx: peft_trainer.PeftTrainer, eval_loss: float if avg_loss <= self.config.target_eval_loss: raise exceptions.StopTraining(f"Target loss {self.config.target_eval_loss=} is achieved.") + def eval_loss_is_preaveraged(self) -> bool: + """Whether Tunix supplies a cross-batch weighted mean to the eval hook.""" + return False + @abc.abstractmethod def get_total_weights(self, batch) -> jax.Array: """Calculate the number of non-padded tokens in the batch.""" @@ -215,9 +219,12 @@ def load_next_train_batch(self, train_ctx: peft_trainer.PeftTrainer): # pylint: """Loads the next batch of data for training.""" try: self.train_batch = self.train_data_loader.load_next_batch() - except Exception as e: # pylint: disable=broad-exception-caught - max_logging.log(f"Exception in load_next_train_batch: {str(e)}") - self.train_batch = None + except exceptions.StopTraining as e: + if isinstance(e.__cause__ or e.__context__, StopIteration): + max_logging.log(f"Training data exhausted: {str(e)}") + self.train_batch = None + else: + raise return self.train_batch @override @@ -232,7 +239,7 @@ def load_next_eval_batch(self, train_ctx: peft_trainer.PeftTrainer): self.eval_batch = None else: self.eval_batch = next(self.eval_data_iterator) - except Exception as e: # pylint: disable=broad-exception-caught - max_logging.log(f"Exception in load_next_eval_batch: {str(e)}") + except StopIteration as e: + max_logging.log(f"Evaluation data exhausted: {str(e)}") self.eval_batch = None return self.eval_batch diff --git a/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_26b_rl.sh b/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_26b_rl.sh index 679ed30185..6ea3069a9c 100644 --- a/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_26b_rl.sh +++ b/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_26b_rl.sh @@ -1,14 +1,13 @@ #!/bin/bash # This script launches a Reinforcement Learning (RL) training workload for the -# gemma4-26b model on a GKE cluster using XPK. +# gemma4-26b model on a GKE cluster using Cluster Toolkit (gcluster). set -e # --- Environment Setup --- -if ! pip show xpk &> /dev/null; then - echo "xpk not found in the environment. Please install it by running:" - echo "uv pip install -e .[runner] --resolution=lowest" +if ! command -v gcluster &> /dev/null; then + echo "gcluster not found in the environment. Please install Cluster Toolkit before running this script." exit 1 fi @@ -19,14 +18,15 @@ export ZONE="${ZONE:-}" # Zone where your v7x cluster is deployed export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-}" # GCS bucket path for outputs (e.g., gs://my-bucket/outputs) export DOCKER_IMAGE="${DOCKER_IMAGE:-}" # Full path to the Docker image you pushed (e.g., us-docker.pkg.dev///:) export MAXTEXT_CKPT_PATH="${MAXTEXT_CKPT_PATH:-}" # GCS path of the MaxText checkpoint you want to fine-tune from (e.g., gs://my-bucket/checkpoints/maxtext-ckpt) -export TPU_TYPE="tpu7x-4x4x4" # v7x (Ironwood) 4x4x4 = 64-chip slice (tpu7x-standard-4t, 4 chips/VM) -export WORKLOAD_NAME="rl-$(date +%Y%m%d-%H%M)" +export COMPUTE_TYPE="${COMPUTE_TYPE:-tpu7x-standard-4t}" # v7x (Ironwood) 4 chips/VM +export TOPOLOGY="${TOPOLOGY:-4x4x4}" # v7x (Ironwood) 4x4x4 = 64-chip slice +export RUN_NAME="${RUN_NAME:-rl-$(date +%Y%m%d-%H%M)}" # Pathways component images, pinned to specific versions (mirrors the reference config). -# Note: --server-image is used for BOTH the Pathways resource-manager server and the +# Note: PATHWAYS_SERVER_IMAGE is used for BOTH the Pathways resource-manager server and the # workers; in the reference config the worker and pathwaysServer images are identical. -export PATHWAYS_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server@sha256:6f5031086c3d603f275763c933e85e587e9332be1642fff515e9ecb52374950b" -export PATHWAYS_PROXY_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server@sha256:8ffcc48a7bbd093a9071e0842ed7055d09d140d75f2b6c0a98bdb76cac77c44e" +export PATHWAYS_SERVER_IMAGE="${PATHWAYS_SERVER_IMAGE:-us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server@sha256:6f5031086c3d603f275763c933e85e587e9332be1642fff515e9ecb52374950b}" +export PATHWAYS_PROXY_SERVER_IMAGE="${PATHWAYS_PROXY_SERVER_IMAGE:-us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server@sha256:8ffcc48a7bbd093a9071e0842ed7055d09d140d75f2b6c0a98bdb76cac77c44e}" # --- Variable Validation --- if [ -z "$PROJECT_ID" ]; then @@ -95,7 +95,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \ model_name=gemma4-26b \ tokenizer_type=huggingface \ tokenizer_path=google/gemma-4-26B-A4B \ -run_name=$WORKLOAD_NAME \ +run_name=$RUN_NAME \ async_scheduling=True \ base_output_directory=$BASE_OUTPUT_DIRECTORY \ chips_per_vm=4 \ @@ -162,17 +162,24 @@ vllm_hf_overrides='{architectures: [\"MaxTextForCausalLM\"]}' \ vllm_additional_config='{\"maxtext_config\": {\"model_name\": \"gemma4-26b\", \"model_call_mode\": \"inference\", \"enable_dp_attention\": false, \"allow_split_physical_axes\": true, \"log_config\": false, \"weight_dtype\": \"bfloat16\", \"prefuse_moe_weights\": true}}'" # Workload Creation -xpk workload create-pathways \ - --cluster=$CLUSTER_NAME \ - --project=$PROJECT_ID \ - --zone=$ZONE \ - --priority=medium \ - --max-restarts=0 \ - --tpu-type=$TPU_TYPE \ +# Configure target cluster in gcluster job config (gcluster job submit reads project/cluster/location from job config) +gcluster job config set project "$PROJECT_ID" +gcluster job config set cluster "$CLUSTER_NAME" +gcluster job config set location "$ZONE" + +# Submit Pathways workload (see `gcluster job submit --help` for --pathways-* flag options) +gcluster job submit \ + --image="${DOCKER_IMAGE}" \ + --name="${RUN_NAME}" \ + --pathways \ + --compute-type="${COMPUTE_TYPE}" \ + --topology="${TOPOLOGY}" \ --num-slices=1 \ - --docker-image="${DOCKER_IMAGE}" \ - --workload="${WORKLOAD_NAME}" \ - --server-image="${PATHWAYS_SERVER_IMAGE}" \ - --proxy-server-image="${PATHWAYS_PROXY_SERVER_IMAGE}" \ - --custom-pathways-proxy-server-args='${XLA_FLAGS}' \ + --priority=medium \ + --restarts=0 \ + --pathways-gcs-location="${BASE_OUTPUT_DIRECTORY}" \ + --pathways-server-image="${PATHWAYS_SERVER_IMAGE}" \ + --pathways-worker-image="${PATHWAYS_SERVER_IMAGE}" \ + --pathways-proxy-server-image="${PATHWAYS_PROXY_SERVER_IMAGE}" \ + --pathways-proxy-args="${XLA_FLAGS}" \ --command="${MAXTEXT_COMMAND}" diff --git a/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_e4b_rl.sh b/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_e4b_rl.sh index 14cec0fced..0a639b0ce3 100644 --- a/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_e4b_rl.sh +++ b/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_e4b_rl.sh @@ -1,14 +1,13 @@ #!/bin/bash # This script launches a Reinforcement Learning (RL) training workload for the -# gemma4-e4b model on a GKE cluster using XPK. +# gemma4-e4b model on a GKE cluster using Cluster Toolkit (gcluster). set -e # --- Environment Setup --- -if ! pip show xpk &> /dev/null; then - echo "xpk not found in the environment. Please install it by running:" - echo "uv pip install -e .[runner] --resolution=lowest" +if ! command -v gcluster &> /dev/null; then + echo "gcluster not found in the environment. Please install Cluster Toolkit before running this script." exit 1 fi @@ -19,14 +18,15 @@ export ZONE="${ZONE:-}" # Zone where your v6e cluster is deployed export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-}" # GCS bucket path for outputs (e.g., gs://my-bucket/outputs) export DOCKER_IMAGE="${DOCKER_IMAGE:-}" # Full path to the Docker image you pushed (e.g., us-docker.pkg.dev///:) export MAXTEXT_CKPT_PATH="${MAXTEXT_CKPT_PATH:-}" # GCS path of the MaxText checkpoint you want to fine-tune from (e.g., gs://my-bucket/checkpoints/maxtext-ckpt) -export TPU_TYPE="v6e-32" -export WORKLOAD_NAME="rl-$(date +%Y%m%d-%H%M)" +export COMPUTE_TYPE="${COMPUTE_TYPE:-ct6e-standard-4t}" # v6e (Trillium) 4 chips/VM +export TOPOLOGY="${TOPOLOGY:-4x8}" # v6e (Trillium) 4x8 = 32-chip slice +export RUN_NAME="${RUN_NAME:-rl-$(date +%Y%m%d-%H%M)}" # Pathways component images, pinned to specific versions (mirrors the reference config). -# Note: --server-image is used for BOTH the Pathways resource-manager server and the +# Note: PATHWAYS_SERVER_IMAGE is used for BOTH the Pathways resource-manager server and the # workers; in the reference config the worker and pathwaysServer images are identical. -export PATHWAYS_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server@sha256:d03068b39a8a2fab0621086ccb7c9445ce17ad34f1520159ca4ecd395346a162" -export PATHWAYS_PROXY_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server@sha256:6342a3cae2818d2f887d396e9ae4156b3316f79fef186d71ff16d535b44e1724" +export PATHWAYS_SERVER_IMAGE="${PATHWAYS_SERVER_IMAGE:-us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server@sha256:d03068b39a8a2fab0621086ccb7c9445ce17ad34f1520159ca4ecd395346a162}" +export PATHWAYS_PROXY_SERVER_IMAGE="${PATHWAYS_PROXY_SERVER_IMAGE:-us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server@sha256:6342a3cae2818d2f887d396e9ae4156b3316f79fef186d71ff16d535b44e1724}" # --- Variable Validation --- if [ -z "$PROJECT_ID" ]; then @@ -95,7 +95,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \ model_name=gemma4-e4b \ tokenizer_type=huggingface \ tokenizer_path=google/gemma-4-E4B \ -run_name=$WORKLOAD_NAME \ +run_name=$RUN_NAME \ async_scheduling=True \ base_output_directory=$BASE_OUTPUT_DIRECTORY \ chips_per_vm=4 \ @@ -164,17 +164,24 @@ vllm_hf_overrides='{architectures: [\"MaxTextForCausalLM\"]}' \ vllm_additional_config='{\"maxtext_config\": {\"model_name\": \"gemma4-e4b\", \"model_call_mode\": \"inference\", \"enable_dp_attention\": true, \"allow_split_physical_axes\": true, \"log_config\": false, \"weight_dtype\": \"bfloat16\", \"prefuse_moe_weights\": true}}'" # Workload Creation -xpk workload create-pathways \ - --cluster=$CLUSTER_NAME \ - --project=$PROJECT_ID \ - --zone=$ZONE \ - --priority=medium \ - --max-restarts=0 \ - --tpu-type=$TPU_TYPE \ +# Configure target cluster in gcluster job config (gcluster job submit reads project/cluster/location from job config) +gcluster job config set project "$PROJECT_ID" +gcluster job config set cluster "$CLUSTER_NAME" +gcluster job config set location "$ZONE" + +# Submit Pathways workload (see `gcluster job submit --help` for --pathways-* flag options) +gcluster job submit \ + --image="${DOCKER_IMAGE}" \ + --name="${RUN_NAME}" \ + --pathways \ + --compute-type="${COMPUTE_TYPE}" \ + --topology="${TOPOLOGY}" \ --num-slices=1 \ - --docker-image="${DOCKER_IMAGE}" \ - --workload="${WORKLOAD_NAME}" \ - --server-image="${PATHWAYS_SERVER_IMAGE}" \ - --proxy-server-image="${PATHWAYS_PROXY_SERVER_IMAGE}" \ - --custom-pathways-proxy-server-args='${XLA_FLAGS}' \ + --priority=medium \ + --restarts=0 \ + --pathways-gcs-location="${BASE_OUTPUT_DIRECTORY}" \ + --pathways-server-image="${PATHWAYS_SERVER_IMAGE}" \ + --pathways-worker-image="${PATHWAYS_SERVER_IMAGE}" \ + --pathways-proxy-server-image="${PATHWAYS_PROXY_SERVER_IMAGE}" \ + --pathways-proxy-args="${XLA_FLAGS}" \ --command="${MAXTEXT_COMMAND}" diff --git a/src/maxtext/trainers/post_train/rl/scripts/run_gptoss_20b_rl.sh b/src/maxtext/trainers/post_train/rl/scripts/run_gptoss_20b_rl.sh index 37aad38b30..c073ee4d49 100644 --- a/src/maxtext/trainers/post_train/rl/scripts/run_gptoss_20b_rl.sh +++ b/src/maxtext/trainers/post_train/rl/scripts/run_gptoss_20b_rl.sh @@ -1,25 +1,26 @@ #!/bin/bash # This script launches a Reinforcement Learning (RL) training workload for the -# GPT-OSS 20B model on a GKE cluster using XPK. +# GPT-OSS 20B model on a GKE cluster using Cluster Toolkit (gcluster). set -e # --- Environment Setup --- -if ! pip show xpk &> /dev/null; then - echo "xpk not found in the environment. Please install maxtext[runner] before running this script." +if ! command -v gcluster &> /dev/null; then + echo "gcluster not found in the environment. Please install Cluster Toolkit before running this script." exit 1 fi # --- Environment Variables --- export PROJECT_ID="${PROJECT_ID:-}" # GCP project ID where the Ironwood cluster is deployed -export CLUSTER_NAME="${CLUSTER_NAME:-}" # Name of your Ironwood cluster +export CLUSTER_NAME="${CLUSTER_NAME:-}" # Name of your cluster export ZONE="${ZONE:-}" # Zone where your Ironwood cluster is deployed export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-}" # GCS bucket path for outputs (e.g., gs://my-bucket/outputs) export DOCKER_IMAGE="${DOCKER_IMAGE:-}" # Full path to the Docker image you pushed (e.g., us-docker.pkg.dev///:) export MAXTEXT_CKPT_PATH="${MAXTEXT_CKPT_PATH:-}" # GCS path of the MaxText checkpoint you want to fine-tune from (e.g., gs://my-bucket/checkpoints/maxtext-ckpt) -export TPU_TYPE="v5p-64" -export WORKLOAD_NAME="rl-$(date +%Y%m%d-%H%M)" +export COMPUTE_TYPE="${COMPUTE_TYPE:-ct5p-hightpu-4t}" # v5p 4 chips/VM +export TOPOLOGY="${TOPOLOGY:-2x4x4}" # v5p 2x4x4 = 32-chip (64-core) slice +export RUN_NAME="${RUN_NAME:-rl-$(date +%Y%m%d-%H%M)}" # --- Variable Validation --- if [ -z "$PROJECT_ID" ]; then @@ -95,14 +96,14 @@ MAXTEXT_COMMAND="MODEL_IMPL_TYPE=flax_nnx \ model_name=gpt-oss-20b \ tokenizer_path=unsloth/gpt-oss-20b-BF16 \ load_parameters_path=${MAXTEXT_CKPT_PATH} \ - run_name=${WORKLOAD_NAME} \ + run_name=${RUN_NAME} \ base_output_directory=${BASE_OUTPUT_DIRECTORY} \ checkpoint_storage_use_ocdbt=False \ checkpoint_storage_use_zarr3=False \ rollout_data_parallelism=2 \ rollout_tensor_parallelism=8 \ hbm_utilization_vllm=0.8 \ - batch_size=8 profiler=xplane \ + profiler=xplane \ profiler_steps=2 \ num_batches=500 \ base_emb_dim=2880 \ @@ -118,16 +119,22 @@ MAXTEXT_COMMAND="MODEL_IMPL_TYPE=flax_nnx \ " # Workload Creation -xpk workload create-pathways \ - --cluster=$CLUSTER_NAME \ - --project=$PROJECT_ID \ - --zone=$ZONE \ - --priority=medium \ - --max-restarts=0 \ - --tpu-type=$TPU_TYPE \ +# Configure target cluster in gcluster job config (gcluster job submit reads project/cluster/location from job config) +gcluster job config set project "$PROJECT_ID" +gcluster job config set cluster "$CLUSTER_NAME" +gcluster job config set location "$ZONE" + +# Submit Pathways workload (see `gcluster job submit --help` for --pathways-* flag options) +gcluster job submit \ + --image="${DOCKER_IMAGE}" \ + --name="${RUN_NAME}" \ + --pathways \ + --compute-type="${COMPUTE_TYPE}" \ + --topology="${TOPOLOGY}" \ --num-slices=1 \ - --docker-image="${DOCKER_IMAGE}" \ - --workload="${WORKLOAD_NAME}" \ - --custom-pathways-proxy-server-args='${XLA_FLAGS}' \ + --priority=medium \ + --restarts=0 \ + --pathways-gcs-location="${BASE_OUTPUT_DIRECTORY}" \ + --pathways-proxy-args="${XLA_FLAGS}" \ --command="${MAXTEXT_COMMAND}" diff --git a/src/maxtext/trainers/post_train/rl/scripts/run_qwen3_30b_rl.sh b/src/maxtext/trainers/post_train/rl/scripts/run_qwen3_30b_rl.sh index ac75b0eeb6..b2244b5a28 100644 --- a/src/maxtext/trainers/post_train/rl/scripts/run_qwen3_30b_rl.sh +++ b/src/maxtext/trainers/post_train/rl/scripts/run_qwen3_30b_rl.sh @@ -1,13 +1,13 @@ #!/bin/bash # This script launches a Reinforcement Learning (RL) training workload for the -# Qwen3-30B-A3B model on a GKE cluster using XPK. +# Qwen3-30B-A3B model on a GKE cluster using Cluster Toolkit (gcluster). set -e # --- Environment Setup --- -if ! pip show xpk &> /dev/null; then - echo "xpk not found in the environment. Please install maxtext[runner] before running this script." +if ! command -v gcluster &> /dev/null; then + echo "gcluster not found in the environment. Please install Cluster Toolkit before running this script." exit 1 fi @@ -18,8 +18,9 @@ export ZONE="${ZONE:-}" # Zone where your Ironwood cluster is deployed export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-}" # GCS bucket path for outputs (e.g., gs://my-bucket/outputs) export DOCKER_IMAGE="${DOCKER_IMAGE:-}" # Full path to the Docker image you pushed (e.g., us-docker.pkg.dev///:) export MAXTEXT_CKPT_PATH="${MAXTEXT_CKPT_PATH:-}" # GCS path of the MaxText checkpoint you want to fine-tune from (e.g., gs://my-bucket/checkpoints/maxtext-ckpt) -export TPU_TYPE="tpu7x-128" -export WORKLOAD_NAME="rl-$(date +%Y%m%d-%H%M)" +export COMPUTE_TYPE="${COMPUTE_TYPE:-tpu7x-standard-4t}" # v7x (Ironwood) 4 chips/VM +export TOPOLOGY="${TOPOLOGY:-4x4x4}" # v7x (Ironwood) 4x4x4 = 64-chip (128-core) slice +export RUN_NAME="${RUN_NAME:-rl-$(date +%Y%m%d-%H%M)}" # --- Variable Validation --- if [ -z "$PROJECT_ID" ]; then @@ -93,7 +94,7 @@ ENABLE_PATHWAYS_PERSISTENCE=1 \ python3 -m maxtext.trainers.post_train.rl.train_rl \ model_name=qwen3-30b-a3b-base \ tokenizer_path=Qwen/Qwen3-30B-A3B-Base \ -run_name=$WORKLOAD_NAME \ +run_name=$RUN_NAME \ async_scheduling=True \ base_output_directory=$BASE_OUTPUT_DIRECTORY \ chips_per_vm=8 \ @@ -138,15 +139,21 @@ vllm_hf_overrides='{architectures: [\"MaxTextForCausalLM\"]}' \ vllm_additional_config='{\"maxtext_config\": {\"model_name\": \"qwen3-30b-a3b\", \"model_call_mode\": \"inference\", \"enable_dp_attention\": false, \"allow_split_physical_axes\": true, \"log_config\": false, \"weight_dtype\": \"bfloat16\", \"prefuse_moe_weights\": true}}'" # Workload Creation -xpk workload create-pathways \ - --cluster=$CLUSTER_NAME \ - --project=$PROJECT_ID \ - --zone=$ZONE \ - --priority=medium \ - --max-restarts=0 \ - --tpu-type=$TPU_TYPE \ +# Configure target cluster in gcluster job config (gcluster job submit reads project/cluster/location from job config) +gcluster job config set project "$PROJECT_ID" +gcluster job config set cluster "$CLUSTER_NAME" +gcluster job config set location "$ZONE" + +# Submit Pathways workload (see `gcluster job submit --help` for --pathways-* flag options) +gcluster job submit \ + --image="${DOCKER_IMAGE}" \ + --name="${RUN_NAME}" \ + --pathways \ + --compute-type="${COMPUTE_TYPE}" \ + --topology="${TOPOLOGY}" \ --num-slices=1 \ - --docker-image="${DOCKER_IMAGE}" \ - --workload="${WORKLOAD_NAME}" \ - --custom-pathways-proxy-server-args='${XLA_FLAGS}' \ + --priority=medium \ + --restarts=0 \ + --pathways-gcs-location="${BASE_OUTPUT_DIRECTORY}" \ + --pathways-proxy-args="${XLA_FLAGS}" \ --command="${MAXTEXT_COMMAND}" diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index 74daaea68f..21a751a01f 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/src/maxtext/trainers/post_train/rl/train_rl.py @@ -391,7 +391,12 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments log_dir=trainer_config.tensorboard_dir, skip_first_n_steps=trainer_config.skip_first_n_steps_for_profiler, profiler_steps=trainer_config.profiler_steps, + # Skip setting tracer levels. set_profile_options=False, + **_kwargs_supported_by( + profiler.ProfilerOptions, + enable_continuous_profiling=trainer_config.enable_continuous_profiling, + ), ) # Parse vllm_additional_config @@ -487,6 +492,10 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments # to support off-policy filtering and multi-iteration training. **({"return_logprobs": True} if trainer_config.rl.use_agentic_rollout else {}), **get_rollout_kwargs_for_parallelism(sampler_config, len(sampler_devices)), + **_kwargs_supported_by( + base_rollout.RolloutConfig, + rollout_vllm_free_kv_cache_during_weight_sync=trainer_config.free_kv_cache_during_weight_sync, + ), ), ) diff --git a/src/maxtext/trainers/post_train/sft/hooks.py b/src/maxtext/trainers/post_train/sft/hooks.py index 9f30a3d6d8..ed26790d53 100644 --- a/src/maxtext/trainers/post_train/sft/hooks.py +++ b/src/maxtext/trainers/post_train/sft/hooks.py @@ -29,7 +29,15 @@ class SFTTrainingHooks(BaseTrainingHooks): @override def get_total_weights(self, batch) -> jax.Array: """Calculate the number of non-padded tokens in the batch.""" - return jnp.sum(batch["targets_segmentation"] != 0) + loss_weights = batch.get("targets_loss_mask") + if loss_weights is None: + loss_weights = batch["targets_segmentation"] + return jnp.sum(loss_weights != 0) + + @override + def eval_loss_is_preaveraged(self) -> bool: + """Whether the configured objective supplies a preaveraged eval loss.""" + return getattr(self.config, "training_objective", "causal_lm") == "block_diffusion" class SFTDataHooks(BaseDataHooks): diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py index 274c66230c..d0532ac8ae 100644 --- a/src/maxtext/trainers/post_train/sft/train_sft.py +++ b/src/maxtext/trainers/post_train/sft/train_sft.py @@ -50,10 +50,12 @@ from orbax import checkpoint as ocp +from tunix.sft import diffusion as tunix_diffusion_sft from tunix.sft import metrics_logger, peft_trainer, profiler from maxtext.optimizers import optimizers from maxtext.configs import pyconfig +from maxtext.integration.tunix import diffusion_sft as maxtext_diffusion_sft from maxtext.trainers.pre_train.train import loss_fn from maxtext.common.goodput import ( GoodputEvent, @@ -225,6 +227,7 @@ def get_tunix_config(mt_config): skip_first_n_steps=mt_config.skip_first_n_steps_for_profiler, profiler_steps=mt_config.profiler_steps, set_profile_options=set_profile_options, + enable_continuous_profiling=mt_config.enable_continuous_profiling, ) return peft_trainer.TrainingConfig( @@ -255,29 +258,64 @@ def use_maxtext_loss_function(trainer, mt_config): The trainer configured with the MaxText loss function. """ - def loss_func( - model, - inputs, - inputs_position, - inputs_segmentation, - targets, - targets_position, - targets_segmentation, - ): - data = { - "inputs": inputs, - "inputs_position": inputs_position, - "inputs_segmentation": inputs_segmentation, - "targets": targets, - "targets_position": targets_position, - "targets_segmentation": targets_segmentation, - } - return loss_fn(model, mt_config, data, dropout_rng=None, params=None, is_train=True) - - trainer = trainer.with_loss_fn(loss_func, has_aux=True) + def make_loss_func(is_train): + def loss_func( + model, + inputs, + inputs_position, + inputs_segmentation, + targets, + targets_position, + targets_segmentation, + ): + data = { + "inputs": inputs, + "inputs_position": inputs_position, + "inputs_segmentation": inputs_segmentation, + "targets": targets, + "targets_position": targets_position, + "targets_segmentation": targets_segmentation, + } + return loss_fn(model, mt_config, data, dropout_rng=None, params=None, is_train=is_train) + + return loss_func + + # Tunix PeftTrainer.with_loss_fn() defaults `eval_loss_fn` to the same callable + # as `loss_fn`. We override `eval_loss_fn` with `is_train=False` so that evaluation: + # - Disables dropout to ensure deterministic eval loss / perplexity. + # - Excludes auxiliary Multi-Token Prediction (MTP) loss from the reported eval loss. + # - Prevents updating running sparsity statistics (`batch_stats`) during eval. + trainer = trainer.with_loss_fn(make_loss_func(is_train=True), has_aux=True) + trainer.eval_loss_fn = make_loss_func(is_train=False) return trainer +def configure_training_objective(trainer, mt_config): + """Configures causal or target-aligned diffusion SFT.""" + if getattr(mt_config, "training_objective", "causal_lm") != "block_diffusion": + return use_maxtext_loss_function(trainer, mt_config) + max_logging.log("Configuring Tunix target-aligned block-diffusion SFT adapter.") + return tunix_diffusion_sft.configure_diffusion_sft( + trainer, + maxtext_diffusion_sft.create_batch_adapter(mt_config), + maxtext_diffusion_sft.create_target_aligned_logits_fn(mt_config), + ) + + +def _create_trainer(model, optimizer, tunix_config, mt_config): + """Creates the trainer implementation required by the selected objective. + + Causal SFT uses ``MaxTextPeftTrainer`` so its MaxText loss can run through + the explicit NNX split/merge train step above. Block-diffusion SFT instead + uses Tunix's base ``PeftTrainer`` because ``configure_diffusion_sft`` installs + a Tunix ``LossOutput`` with an unreduced weighted loss and denominator-aware + gradient accumulation, which the MaxText-specific train step does not handle. + """ + if getattr(mt_config, "training_objective", "causal_lm") == "block_diffusion": + return peft_trainer.PeftTrainer(model, optimizer, tunix_config) + return MaxTextPeftTrainer(model, optimizer, tunix_config) + + def validate_config(config): """Validates the configuration parameters for SFT training.""" if config.optimizer_memory_host_offload: @@ -286,10 +324,28 @@ def validate_config(config): "SFT path because the underlying Tunix PeftTrainer does not support " "host offloading of the optimizer state." ) + if getattr(config, "training_objective", "causal_lm") == "block_diffusion": + unsupported_auxiliary_losses = { + name: value + for name, value in ( + ("z_loss_multiplier", getattr(config, "z_loss_multiplier", 0.0)), + ("load_balance_loss_weight", getattr(config, "load_balance_loss_weight", 0.0)), + ("mtp_num_layers", getattr(config, "mtp_num_layers", 0)), + ("indexer_loss_scaling_factor", getattr(config, "indexer_loss_scaling_factor", 0.0)), + ) + if value != 0 + } + if unsupported_auxiliary_losses: + settings = ", ".join(f"{name}={value!r}" for name, value in unsupported_auxiliary_losses.items()) + raise ValueError( + "Block-diffusion SFT currently supports only weighted cross-entropy; " + f"set unsupported auxiliary-loss settings to zero: {settings}" + ) def setup_trainer_state(mt_config, goodput_recorder=None): """Set up prerequisites for training loop.""" + validate_config(mt_config) tunix_config = get_tunix_config(mt_config) with maybe_record_goodput(goodput_recorder, GoodputEvent.TPU_INIT): @@ -316,10 +372,10 @@ def setup_trainer_state(mt_config, goodput_recorder=None): nnx.pop(model, nnx.Intermediate) if mt_config.lora.lora_restore_path: lora_utils.restore_lora_from_path(model, mt_config) - trainer = MaxTextPeftTrainer(model, optimizer, tunix_config) + trainer = _create_trainer(model, optimizer, tunix_config, mt_config) trainer.with_training_hooks(training_hooks) trainer.with_data_hooks(data_hooks) - trainer = use_maxtext_loss_function(trainer, mt_config) + trainer = configure_training_objective(trainer, mt_config) return trainer, mesh @@ -327,9 +383,9 @@ def setup_trainer_state(mt_config, goodput_recorder=None): def train_model(mt_config, trainer, mesh): """Runs the SFT training loop in Tunix.""" with jax.set_mesh(mesh), logical_axis_rules(mt_config.logical_axis_rules): - # Disable NNX graph caching for MoE models (where experts > 1) to allow - # necessary dynamic metadata synchronization during forward passes (e.g., in jax.lax.scan). - enable_nnx_cache = mt_config.num_experts <= 1 + is_block_diffusion = getattr(mt_config, "training_objective", "causal_lm") == "block_diffusion" + records_internal_metrics = bool(getattr(mt_config, "record_internal_nn_metrics", False)) + enable_nnx_cache = mt_config.num_experts <= 1 and not (is_block_diffusion and records_internal_metrics) trainer.train( trainer.data_hooks.train_data_iterator, diff --git a/src/maxtext/trainers/pre_train/scripts/olmo/README.md b/src/maxtext/trainers/pre_train/scripts/olmo/README.md index f12221c4d6..f2fd36dbdd 100644 --- a/src/maxtext/trainers/pre_train/scripts/olmo/README.md +++ b/src/maxtext/trainers/pre_train/scripts/olmo/README.md @@ -8,29 +8,100 @@ global batch ≈4M tokens, peak LR 3e-4 cosine to 0.1×, 2k warmup, β=(0.9, 0.9 ## Scripts -| Script | Purpose | -|---|---| -| `run_olmo3_7b_stage1.sh` | Direct host launcher. Used by the XPK wrapper inside the runner image; also usable standalone for single-host smoke tests. | +| Script | Purpose | +| ------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `run_olmo3_7b_stage1.sh` | Direct host launcher. Used by the XPK wrapper inside the runner image; also usable standalone for single-host smoke tests. | | `xpk_olmo3_7b_stage1.sh` | XPK wrapper for multi-host TPU clusters. Modes: `submit` / `monitor` / `resume_until_done` (deletes + resubmits on Kueue preemption until a target step is reached). | Both scripts are env-var driven. Header comments in each enumerate required vs. optional env. -## Quick start (multi-host TPU via XPK) +## Quick start (multi-host TPU via Cluster Toolkit) + +```bash +export PROJECT_ID= +export GKE_CLUSTER= +export LOCATION= +export RUN_NAME= +export GCS_BUCKET= +export BASE_OUTPUT_DIRECTORY=gs://${GCS_BUCKET?}/olmo/runs +export COMPUTE_TYPE= +export TOPOLOGY= +export TOTAL_DEVICES=512 # set to match device-type * 2 (Ironwood has 2 JAX devices/chip) +export DOCKER_IMAGE=gcr.io/${PROJECT_ID?}/maxtext-olmo3:latest # see "Runner image" in Prerequisites below + +export OLMO_INDEX_PATH=/tmp/olmo-data/olmo/indices/olmo_index_seq8192.json +export OLMO_GCS_BASE=gs://${GCS_BUCKET?}/ +export LOAD_PARAMETERS_PATH=gs://${GCS_BUCKET?}/olmo/checkpoints/stage1-step0/0/items +export HF_TOKEN= + +gcloud config set project ${PROJECT_ID?} +gcloud container clusters get-credentials ${GKE_CLUSTER?} \ + --location ${LOCATION?} \ + --project ${PROJECT_ID?} +gcluster job config set project ${PROJECT_ID?} +gcluster job config set cluster ${GKE_CLUSTER?} +gcluster job config set location ${LOCATION?} + +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?} \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --mount="gs://${GCS_BUCKET?};/tmp/olmo-data;ro" \ + --command="INDEX_PATH=${OLMO_INDEX_PATH?} \ +GCS_BASE=${OLMO_GCS_BASE?} \ +LOCAL_MOUNT=/tmp/olmo-data \ +OUTPUT_DIR=${BASE_OUTPUT_DIRECTORY?} \ +RUN_NAME=${RUN_NAME?} \ +LOAD_PARAMETERS_PATH=${LOAD_PARAMETERS_PATH?} \ +TOTAL_DEVICES=${TOTAL_DEVICES?} \ +HF_TOKEN=${HF_TOKEN?} \ +MOUNT_GCSFUSE=0 \ +VENV_PATH=/__skip_venv__ \ +bash src/maxtext/trainers/pre_train/scripts/olmo/run_olmo3_7b_stage1.sh" +``` + +Set the data, checkpoint, and launcher environment variables before invoking `run_olmo3_7b_stage1.sh`. For smoke testing without tokenized OLMo data, see the single-host or synthetic data section below. + +## Quick start (single-host / smoke test) + +```bash +source ${MAXTEXT_ROOT?}/maxtext_venv/bin/activate + +export GCS_BUCKET= +export RUN_NAME= + +INDEX_PATH=/tmp/olmo-data/olmo/indices/olmo_index_seq8192.json \ +GCS_BASE=gs://${GCS_BUCKET?}/ \ +LOCAL_MOUNT=/tmp/olmo-data \ +OUTPUT_DIR=gs://${GCS_BUCKET?}/olmo/runs \ +LOAD_PARAMETERS_PATH=gs://${GCS_BUCKET?}/olmo/checkpoints/stage1-step0/0/items \ +HF_SECRETS=~/.hf_token.sh \ +RUN_NAME=${RUN_NAME?} \ +STEPS=50 WARMUP_STEPS=10 CHECKPOINT_PERIOD=50 \ +bash src/maxtext/trainers/pre_train/scripts/olmo/run_olmo3_7b_stage1.sh +``` + +## Legacy quick start (multi-host TPU via XPK) + +For new deployments, use the Cluster Toolkit version above. The XPK wrapper +section is retained for existing environments. ```bash source ~/.hf_token.sh -export XPK_CLUSTER= -export XPK_PROJECT= -export XPK_ZONE= +export XPK_CLUSTER= +export XPK_PROJECT= +export XPK_ZONE= +export GCS_BUCKET= export XPK_DEVICE_TYPE=tpu7x-4x8x8 # or tpu7x-4x4x4 for a smoke run export XPK_TOTAL_DEVICES=512 # set to match device-type * 2 (Ironwood has 2 JAX devices/chip) -export XPK_BASE_OUTPUT_DIR=gs:///olmo/runs +export XPK_BASE_OUTPUT_DIR=gs://${GCS_BUCKET?}/olmo/runs export XPK_RUN_NAME=olmo3_7b_stage1 export OLMO_INDEX_PATH=/tmp/olmo-data/olmo/indices/olmo_index_seq8192.json -export OLMO_GCS_BASE=gs:/// -export LOAD_PARAMETERS_PATH=gs:///olmo/checkpoints/stage1-step0/0/items +export OLMO_GCS_BASE=gs://${GCS_BUCKET?}/ +export LOAD_PARAMETERS_PATH=gs://${GCS_BUCKET?}/olmo/checkpoints/stage1-step0/0/items bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh submit bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh monitor @@ -43,22 +114,6 @@ STEPS_OVERRIDE=1414078 \ horizon. The wrapper passes `LIBTPU_INIT_ARGS` (Ironwood XLA flags) and the full MaxText perf flag set automatically — no manual override needed. -## Quick start (single-host / smoke test) - -```bash -source $MAXTEXT_ROOT/maxtext_venv/bin/activate - -INDEX_PATH=/tmp/olmo-data/olmo/indices/olmo_index_seq8192.json \ -GCS_BASE=gs:/// \ -LOCAL_MOUNT=/tmp/olmo-data \ -OUTPUT_DIR=gs:///olmo/runs \ -LOAD_PARAMETERS_PATH=gs:///olmo/checkpoints/stage1-step0/0/items \ -HF_SECRETS=~/.hf_token.sh \ -RUN_NAME=olmo3_7b_stage1 \ -STEPS=50 WARMUP_STEPS=10 CHECKPOINT_PERIOD=50 \ -bash src/maxtext/trainers/pre_train/scripts/olmo/run_olmo3_7b_stage1.sh -``` - ## Prerequisites **Tokenized data corpus.** AI2 maintains the OLMo-mix manifests in @@ -74,14 +129,13 @@ underlying preprocessed `.npy` tokens live under AI2's `s3://ai2-llm/` bucket (see the `base_dir` arg on `DataMix.build` in [`OLMo-core`](https://github.com/allenai/OLMo-core/blob/main/src/olmo_core/data/mixes/__init__.py)); mirror them into your own GCS bucket with -[`tools/data_generation/download_olmo_data_to_gcs.py`](../../../../../tools/data_generation/download_olmo_data_to_gcs.py) +[`tools/data_generation/download_olmo_data_to_gcs.py`](../../../../../../tools/data_generation/download_olmo_data_to_gcs.py) (reads a manifest, pulls from AI2's source, uploads to `--gcs-dest`). **Data index.** Once the corpus is mirrored, build the index with -[`tools/data_generation/build_olmo_npy_index.py`](../../../../../tools/data_generation/build_olmo_npy_index.py) +[`tools/data_generation/build_olmo_npy_index.py`](../../../../../../tools/data_generation/build_olmo_npy_index.py) against the same manifest + sequence length, then upload the resulting JSON -to GCS. Mount your bucket read-only via gcsfuse inside the pod — the XPK -wrapper does this automatically (`MOUNT_GCSFUSE=1`). +to GCS. Mount your bucket read-only inside the pod using `--mount="gs://${GCS_BUCKET?};/tmp/olmo-data;ro"` with `MOUNT_GCSFUSE=0` when submitting via `gcluster job submit` (since non-Pathways `gcluster` containers run unprivileged and use the GKE GCS FUSE CSI driver), or via `MOUNT_GCSFUSE=1` in the legacy XPK wrapper. See [`docs/guides/data_input_pipeline/olmo_grain.md`](../../../../../../docs/guides/data_input_pipeline/olmo_grain.md) for the full data-pipeline reference. @@ -93,7 +147,7 @@ for the full data-pipeline reference. This puts MaxText at exactly the same weights as the PyTorch reference, so the loss curve can be overlaid against AI2's published WandB curve. AI2 also publishes intermediate checkpoints at every 1000-step boundary on the -[`allenai/Olmo-3-1025-7B`](https://huggingface.co/allenai/Olmo-3-1025-7B/refs) +[`allenai/Olmo-3-1025-7B`](https://huggingface.co/allenai/Olmo-3-1025-7B) HF repo (`stage1-step0`, `stage1-step1000`, …, `stage1-step1413814`); the same procedure works for any of them — just swap the `--revision` flag. @@ -101,7 +155,7 @@ same procedure works for any of them — just swap the `--revision` flag. peak ~26 GB RAM; needs `huggingface_hub` auth via `HF_TOKEN`): ```bash - export HF_TOKEN=... + export HF_TOKEN= python -m maxtext.checkpoint_conversion.to_maxtext \ model_name=olmo3-7b-pt scan_layers=True \ --revision=stage1-step0 \ @@ -109,20 +163,22 @@ same procedure works for any of them — just swap the `--revision` flag. ``` The script writes an Orbax checkpoint under - `//0/items/` (default `` is `/tmp/maxtext`; + `//0/items/` (default `` is `/tmp/maxtext`; override with `--base_output_directory=...`). 2. Upload the converted checkpoint to GCS so all pods can read it: ```bash - gsutil -m cp -r /0/items gs:///olmo/checkpoints/stage1-step0/0/items + gcloud storage cp -r /0/items gs://${GCS_BUCKET?}/olmo/checkpoints/stage1-step0/0/items ``` -3. Point the launcher at it via `LOAD_PARAMETERS_PATH`: +3. Point the launcher at it via `LOAD_PARAMETERS_PATH` when submitting with `gcluster job submit` (or `xpk_olmo3_7b_stage1.sh` for legacy XPK): ```bash - export LOAD_PARAMETERS_PATH=gs:///olmo/checkpoints/stage1-step0/0/items - bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh submit + export LOAD_PARAMETERS_PATH=gs://${GCS_BUCKET?}/olmo/checkpoints/stage1-step0/0/items + # Then submit via Cluster Toolkit (see "Quick start (multi-host TPU via Cluster Toolkit)" above) + # Or for legacy XPK: + # bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh submit ``` AI2's published checkpoints are parameters only — Adam state is freshly @@ -141,25 +197,25 @@ tracks the AI2-init curve within data-shuffle noise). ```bash unset LOAD_PARAMETERS_PATH -bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh submit +# Submit via Cluster Toolkit without LOAD_PARAMETERS_PATH, or for legacy XPK: +# bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh submit ``` -**Runner image** (XPK only). Build + push once: +**Runner image** (multi-host GKE). Build + push once: ```bash sudo bash src/dependencies/scripts/docker_build_dependency_image.sh \ MODE=stable WORKFLOW=pre-training sudo bash src/dependencies/scripts/docker_upload_runner.sh \ - CLOUD_IMAGE_NAME=maxtext-olmo3 PROJECT=$XPK_PROJECT + CLOUD_IMAGE_NAME=maxtext-olmo3 PROJECT=${PROJECT_ID?} ``` -Override the resulting image with `XPK_DOCKER_IMAGE` (defaults to -`gcr.io/${XPK_PROJECT}/maxtext-olmo3:latest`). +Use `gcr.io/${PROJECT_ID?}/maxtext-olmo3:latest` as `DOCKER_IMAGE` for Cluster Toolkit (or override `XPK_DOCKER_IMAGE` for legacy XPK, which defaults to `gcr.io/${XPK_PROJECT?}/maxtext-olmo3:latest`). ## Resume -Keep `XPK_RUN_NAME` and `XPK_BASE_OUTPUT_DIR` stable across submissions. -Orbax picks up the latest checkpoint under `${OUTPUT_DIR}/${RUN_NAME}/checkpoints/`; +Keep `RUN_NAME` and `BASE_OUTPUT_DIRECTORY` (or `XPK_RUN_NAME` and `XPK_BASE_OUTPUT_DIR` for legacy XPK) stable across submissions. +Orbax picks up the latest checkpoint under `${BASE_OUTPUT_DIRECTORY?}/${RUN_NAME?}/checkpoints/`; the OLMo grain sampler resumes its data position via stateless `initial_step = step × per-host-batch` (no Grain-iterator-state in the checkpoint). @@ -168,10 +224,10 @@ checkpoint). workload on preemption, and exits when checkpoint step ≥ `STEPS_OVERRIDE` or after `MAX_RETRIES` failed attempts. Tune for preemption-prone clusters: -| Env var | Default | Use | -|---|---|---| -| `MAX_RETRIES` | 10 | Cap on resubmit attempts before giving up. | -| `RETRY_BACKOFF_SECONDS` | 300 | Sleep between resubmits. Bump if Kueue admission is slow. | +| Env var | Default | Use | +| ----------------------- | ------- | --------------------------------------------------------- | +| `MAX_RETRIES` | 10 | Cap on resubmit attempts before giving up. | +| `RETRY_BACKOFF_SECONDS` | 300 | Sleep between resubmits. Bump if Kueue admission is slow. | ## Loss-overlay validation diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 04bdb1f6ab..cb5ed14683 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -23,6 +23,7 @@ import functools import gc import os +import time from absl import app @@ -39,6 +40,7 @@ import jax import jax.numpy as jnp +from jax.experimental import multihost_utils from jax.sharding import NamedSharding from flax import nnx, traverse_util @@ -56,6 +58,7 @@ # pylint: disable=too-many-positional-arguments from maxtext.layers.multi_token_prediction import calculate_mtp_acceptance_rate, calculate_mtp_loss, mtp_acceptance, mtp_losses from maxtext.layers.attention_mla import indexer_losses +from maxtext.layers import moe from maxtext.common import checkpointing, profiler from maxtext.common.goodput import ( GoodputEvent, @@ -82,6 +85,12 @@ from maxtext.utils.gradient_accumulation import gradient_accumulation_loss_and_grad from maxtext.utils.vocabulary_tiling import vocab_tiling_nnx_loss +try: + # lineage_adapter is Google-internal and excluded from the open-source export. + from maxtext.experimental.lineage import lineage_adapter # pylint: disable=g-import-not-at-top +except ImportError: + lineage_adapter = None + class EncoderKwargs(TypedDict, total=False): """Multimodal encoder arguments forwarded to the model.""" @@ -148,6 +157,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr else: for k, v in data.items(): data[k] = v[: config.micro_batch_size_to_eval_on, :] + if lineage_adapter is not None and lineage_adapter.is_native(config): + return lineage_adapter.loss_and_aux(model, config, data, is_train=is_train) # Only forward the kwarg when router replay is actually in use, so models # and adapters whose __call__ predates the feature keep working. forced_routing_kwargs = ( @@ -203,6 +214,9 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr intermediates = nnx.pop(model, nnx.Intermediate) intermediate_outputs = intermediates.to_pure_dict() + # defer_small_all_reduces: the MoE layers emitted local partial expert counts; reduce the stacked arrays here, once, + # after the layer loop (a no-op when the flag is off). + intermediate_outputs = moe.finalize_deferred_intermediates(intermediate_outputs, config) # Store them under the collection name so calculate_mtp_loss and # calculate_mtp_acceptance_rate find them at the path they expect. @@ -213,9 +227,12 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr if indexer_losses_state is not None: intermediate_outputs["indexer_losses"] = indexer_losses_state.to_pure_dict() - if (config.use_indexer and not config.indexer_sparse_training) and is_train: + is_dense_warmup = config.use_indexer and config.indexer_loss_scaling_factor > 0.0 and not config.indexer_sparse_training + if is_dense_warmup and is_train: # In Dense Warm-up stage, we skip main model loss calculation for efficiency. # The main model parameters are frozen and only the indexer is trained via KL divergence. + # Gated on indexer_loss_scaling_factor > 0 like the logits skip in nnx_decoders.py; otherwise + # there is no indexer loss either and the total objective would be zero. xent_sum = 0.0 total_z_loss = 0.0 elif config.num_vocab_tiling > 1: @@ -271,7 +288,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr # Zero1+GA to reduce communication overhead. # EPS was used to avoid division by zero, but it's not needed when gradient # accumulation is enabled since there's no division. - if config.gradient_accumulation_steps > 1 and not config.use_tunix_gradient_accumulation: + manual_gradient_accumulation = config.gradient_accumulation_steps > 1 and not config.use_tunix_gradient_accumulation + if manual_gradient_accumulation: loss = xent_sum else: # When using Tunix gradient accumulation, we revert to standard normalization. @@ -283,11 +301,24 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr # We keep z-loss normalized by total_weights. total_z_loss = total_z_loss / (total_weights + EPS) + # The auxiliary losses below (MTP, indexer, MoE load balance) are already + # normalized means. Under manual gradient accumulation the objective is the + # unnormalized xent_sum and the accumulated gradient is divided by the token + # count of the whole global batch, so each auxiliary term is scaled by this + # microbatch's token count to keep its gradient weight at 1/GA per microbatch + # (exactly so when microbatches carry equal token counts). The aux values + # themselves are returned unscaled for logging, and eval_step's loss is left + # as before. + def _add_aux_loss(total, aux_loss): + if manual_gradient_accumulation and is_train: + return total + aux_loss * jnp.asarray(total_weights, jnp.float32) + return total + aux_loss + # Calculate and Add MTP Loss mtp_loss = 0.0 if config.mtp_num_layers > 0 and is_train: mtp_loss = calculate_mtp_loss(intermediate_outputs, config) - loss += mtp_loss + loss = _add_aux_loss(loss, mtp_loss) # Calculate and add auxiliary Indexer loss indexer_loss = 0.0 @@ -296,7 +327,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr indexer_losses_list = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "indexer_loss") if indexer_losses_list: indexer_loss = jnp.mean(jnp.concatenate([jnp.atleast_1d(x) for x in indexer_losses_list])) - loss += indexer_loss # Injects loss into scalar objective to drive backward gradients for indexer weights. + # Injects loss into scalar objective to drive backward gradients for indexer weights. + loss = _add_aux_loss(loss, indexer_loss) else: max_logging.debug("No Indexer loss found. Defaulting to 0.0.") @@ -305,8 +337,11 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr if config.num_experts > 1: moe_lb_losses = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "moe_lb_loss") if moe_lb_losses: - moe_lb_loss = jnp.mean(jnp.concatenate(moe_lb_losses)) - loss += moe_lb_loss + lb_losses = jnp.concatenate(moe_lb_losses) + # Megatron-LM (MoEAuxLossAutoScaler) adds every MoE layer's aux loss to the objective, so routers + # that use its seq_aux_loss sum over layers (decoder + MTP); other routers keep the average. + moe_lb_loss = jnp.sum(lb_losses) if moe.uses_megatron_seq_aux_loss(config) else jnp.mean(lb_losses) + loss = _add_aux_loss(loss, moe_lb_loss) else: max_logging.debug("\nNo MoE load balance loss found. Defaulting to 0.0.") @@ -362,14 +397,83 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr } ) has_moe_overflow = jnp.bool_(False) - if config.retry_when_tokens_dropped: - moe_overflow_flags = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "moe_has_overflow") + if config.moe_dropless_fallback == "step": + moe_overflow_flags = maxtext_utils.collect_intermediates_by_suffix( + intermediate_outputs, "moe_has_overflow", ravel=False + ) if moe_overflow_flags: has_moe_overflow = jnp.any(jnp.stack([jnp.any(x) for x in moe_overflow_flags])) aux["has_moe_overflow"] = has_moe_overflow + if config.moe_dropless_fallback == "layer": + fallback_flags = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "moe_dropless_fallback") + aux["moe_dropless_fallback_layers"] = sum( + (jnp.sum(jnp.asarray(x).astype(jnp.int32)) for x in fallback_flags), jnp.int32(0) + ) + if getattr(config, "log_required_ragged_buffer_factor", False): + # Probe: per-MoE-layer minimum ragged_buffer_factor that would have avoided drops (RoutedMoE sows it next to + # moe_has_overflow). Leaf order of collect_intermediates_by_suffix; scanned layers contribute one entry each. + required_rbf_values = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "moe_required_rbf") + if required_rbf_values: + aux["moe_required_rbf"] = jnp.concatenate([x.astype(jnp.float32) for x in required_rbf_values]) + if getattr(config, "moe_log_max_load_ratio", False): + # One value per MoE layer (decoder + MTP), each already max-reduced over token chunks, + # expert shards and EP groups. A step overflows iff the max exceeds ragged_buffer_factor. + load_ratios = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "moe_max_load_ratio") + if load_ratios: + load_ratios = jnp.concatenate(load_ratios) + aux["moe_max_load_ratio"] = jnp.max(load_ratios) + aux["moe_max_load_ratio_mean"] = jnp.mean(load_ratios) return loss, aux +def _rollback_on_overflow(new_state, old_leaves, has_moe_overflow): + """Returns the pre-step state leaves instead of `new_state` if the MoE ragged buffer overflowed.""" + if not isinstance(has_moe_overflow, jax.core.Tracer) and not bool(has_moe_overflow): + return new_state + + def _select(new, old): + dtype = getattr(new, "dtype", None) + if dtype is not None and jnp.issubdtype(dtype, jax.dtypes.extended): + # NNX never updates typed PRNG keys in place; their stream counts are rolled back. + return new + return jnp.where(has_moe_overflow, old, new) + + new_leaves, treedef = jax.tree_util.tree_flatten(new_state) + if len(new_leaves) != len(old_leaves): + raise ValueError( + "moe_dropless_fallback='step': cannot roll back the step -- the updated state has " + f"{len(new_leaves)} leaves but the input state had {len(old_leaves)}." + ) + return jax.tree_util.tree_unflatten(treedef, [_select(n, o) for n, o in zip(new_leaves, old_leaves)]) + + +def step_diagnostics(raw_grads, has_moe_overflow, bias_values, bias_updates): + """log_step_diagnostics scalars that are not already in the step metrics. + + Returns max |raw gradient| over all parameters (pre-clip), the routed-bias checksum (sum over MoE layers of the sum + of the updated router bias, float32), the number of nonzero routed-bias update entries of this step and the MoE + overflow flag, each as a float32 scalar. The gradient and parameter L2 norms are logged from the existing + learning/raw_grad_norm, learning/grad_norm and learning/param_norm. + """ + leaves = [ + g for g in jax.tree_util.tree_leaves(raw_grads) if hasattr(g, "dtype") and jnp.issubdtype(g.dtype, jnp.inexact) + ] + max_abs = jnp.max(jnp.stack([jnp.max(jnp.abs(g)).astype(jnp.float32) for g in leaves])) if leaves else jnp.float32(0) + checksum = jnp.float32(0) + for b in bias_values: + checksum = checksum + jnp.sum(jnp.asarray(b, jnp.float32)) + nonzero = jnp.float32(0) + for u in bias_updates: + nonzero = nonzero + jnp.sum(jnp.asarray(u) != 0).astype(jnp.float32) + overflow = jnp.asarray(False if has_moe_overflow is None else has_moe_overflow).astype(jnp.float32) + return { + "diag/max_abs_grad": max_abs, + "diag/moe_bias_checksum": checksum, + "diag/moe_bias_update_nonzero": nonzero, + "diag/moe_overflow": overflow, + } + + def _find_gate_bias(module: nnx.Module | None) -> nnx.Variable | None: """Finds the router gate bias parameter in a module graph.""" if module is None: @@ -399,6 +503,7 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat del dropout_rng # unused for NNX (kept for jit signature parity) # pylint: disable=too-many-nested-blocks # --- Per-path initialization --- + input_state_leaves = jax.tree_util.tree_leaves(state) if config.moe_dropless_fallback == "step" else None state = nnx.merge(model, state) # reconstruct TrainStateNNX loss_model, loss_params, loss_rng = state.model, None, None @@ -494,6 +599,8 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): mtp_loss = aux.get("mtp_loss", 0.0) new_opt_state = None bias_metrics = {} + # log_step_diagnostics: the updated router biases and this step's bias updates (see step_diagnostics). + diag_bias_values, diag_bias_updates = [], [] if config.gradient_clipping_threshold > 0: grads = maxtext_utils.apply_gradient_clipping(raw_grads, None, config.gradient_clipping_threshold) @@ -522,7 +629,10 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family # pylint: disable=too-many-nested-blocks - if config.routed_bias and config.routed_bias_update_rate > 0.0: + if lineage_adapter is not None and lineage_adapter.is_native(config): + diag_bias_values = list(aux.get("diag_bias_values", ())) + diag_bias_updates = list(aux.get("diag_bias_updates", ())) + elif config.routed_bias and config.routed_bias_update_rate > 0.0: if getattr(config, "model_name", "").startswith("deepseek4"): max_logging.log("DeepSeek V4: Applying auxiliary-loss-free routing bias via pure NNX MoEBiasVar.") flat_intermediates = traverse_util.flatten_dict(aux.get("intermediate_outputs", {})) @@ -548,6 +658,8 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): if getattr(config, "log_moe_bias_norms", False): bias_metrics[f"learning/moe_bias_before_norm_{name_prefix}"] = jnp.linalg.norm(node.bias.value) node.bias.value = node.bias.value + jnp.array(update_val) + diag_bias_values.append(node.bias.value) + diag_bias_updates.append(update_val) if getattr(config, "log_moe_bias_norms", False): bias_metrics[f"learning/moe_bias_update_norm_{name_prefix}"] = jnp.linalg.norm(jnp.array(update_val)) else: @@ -555,8 +667,11 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): # The update from the scan is (num_moe_layers, num_experts) and must be transposed. decoder_layer = getattr(new_state.model.decoder, "moe_layers", new_state.model.decoder) decoder_bias = _find_gate_bias(decoder_layer) - if decoder_bias is not None: - decoder_bias.value = decoder_bias.value + jnp.array(moe_bias_updates[0]) + if decoder_bias is not None and moe_bias_updates is not None: + decoder_update = jnp.array(moe_bias_updates[0]) + decoder_bias.value = decoder_bias.value + decoder_update + diag_bias_values.append(decoder_bias.value) + diag_bias_updates.append(decoder_update) # 2. Update auxiliary MTP MoE layers (if enabled). # Unlike the main decoder, each MTP layer is an individual un-scanned layer @@ -566,7 +681,10 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): mtp_layer = getattr(new_state.model.mtp_block, f"mtp_layer_{i + 1}", None) mtp_bias = _find_gate_bias(mtp_layer) if mtp_bias is not None: - mtp_bias.value = mtp_bias.value + jnp.array(update) + mtp_update = jnp.array(update) + mtp_bias.value = mtp_bias.value + mtp_update + diag_bias_values.append(mtp_bias.value) + diag_bias_updates.append(mtp_update) lm_loss = xent_sum / (total_weights + EPS) scalar_metrics = { @@ -587,6 +705,11 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): "learning/te_moe_recv_capacity_per_rank": aux["te_moe_recv_capacity_per_rank"], } ) + if config.moe_dropless_fallback == "layer": + scalar_metrics["learning/moe_dropless_fallback_layers"] = aux["moe_dropless_fallback_layers"] + if "moe_max_load_ratio" in aux: + scalar_metrics["learning/moe_max_load_ratio"] = aux["moe_max_load_ratio"] + scalar_metrics["learning/moe_max_load_ratio_mean"] = aux["moe_max_load_ratio_mean"] scalar_metrics.update(bias_metrics) if config.use_qk_clip: new_state = qk_clip_utils.apply_qk_clip_nnx(new_state, intermediate_outputs, config) @@ -601,6 +724,13 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): model_params = nnx.state(new_state.model, nnx.Param) scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(model_params) + if getattr(config, "log_step_diagnostics", False): + scalar_metrics.update(step_diagnostics(raw_grads, has_moe_overflow, diag_bias_values, diag_bias_updates)) + if "learning/raw_grad_norm" not in scalar_metrics: # optimizer_memory_host_offload skips the norms above + scalar_metrics["learning/raw_grad_norm"] = max_utils.l2norm_pytree(raw_grads) + scalar_metrics["learning/grad_norm"] = max_utils.l2norm_pytree(grads) + scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(nnx.state(new_state.model, nnx.Param)) + # Surface skip-step rejections as a TB metric. The skip-step optimizer stores # is_skipped in its opt_state; read it back off the optimizer just updated in place. if config.skip_step_on_spikes: @@ -612,17 +742,23 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): "scalar": scalar_metrics, "scalars": {}, } - if config.retry_when_tokens_dropped: - metrics["has_moe_overflow"] = ( # pyrefly: ignore[bad-assignment] + if config.moe_dropless_fallback == "step": + metrics["has_moe_overflow"] = ( has_moe_overflow if has_moe_overflow is not None else jnp.bool_(False) # pyrefly: ignore[bad-assignment] ) + if getattr(config, "log_required_ragged_buffer_factor", False) and "moe_required_rbf" in aux: + # Top-level (not "scalar"): a per-layer vector, read by training_loop_iteration for the REQUIRED_RBF line. + metrics["moe_required_rbf"] = aux["moe_required_rbf"] if getattr(config, "record_internal_nn_metrics", False): record_activation_metrics(metrics, intermediate_outputs, config) # Drop Intermediates (e.g. sowed max_logits for QK-Clip) and the MTP sown # vars (mtp_losses/mtp_acceptance) before returning. They're absent from # state_mesh_shardings and would cause a leaf-count / structure mismatch. - return nnx.state(new_state, nnx.Not(nnx.Intermediate)), metrics + out_state = nnx.state(new_state, nnx.Not(nnx.Intermediate)) + if input_state_leaves is not None and has_moe_overflow is not None: + out_state = _rollback_on_overflow(out_state, input_state_leaves, has_moe_overflow) + return out_state, metrics def eval_step(model, config, state, data, dropout_rng=None): @@ -654,12 +790,25 @@ def eval_step(model, config, state, data, dropout_rng=None): "evaluation/mtp_acceptance_rate_percent": mtp_acceptance_rate, }, } - if config.retry_when_tokens_dropped: + if config.moe_dropless_fallback == "step": metrics["has_moe_overflow"] = aux.get("has_moe_overflow", False) + if config.moe_dropless_fallback == "layer": + metrics["scalar"]["evaluation/moe_dropless_fallback_layers"] = aux["moe_dropless_fallback_layers"] return metrics +def _maybe_log_required_rbf(config, step, metrics, program): + """log_required_ragged_buffer_factor probe: one REQUIRED_RBF line per train program run (fetches the vector).""" + if not getattr(config, "log_required_ragged_buffer_factor", False) or "moe_required_rbf" not in metrics: + return + per_layer = np.asarray(jax.device_get(metrics["moe_required_rbf"]), dtype=np.float32).ravel() + if per_layer.size == 0: + return + per_layer_str = ",".join(f"{v:.4f}" for v in per_layer) + max_logging.log(f"REQUIRED_RBF step={step} max={per_layer.max():.4f} per_layer=[{per_layer_str}] program={program}") + + def training_loop_iteration( jax_device_state: dict[str, Any], python_vars: dict[str, Any], @@ -672,6 +821,7 @@ def training_loop_iteration( mesh = jax_device_state["mesh"] p_train_step = jax_device_state["p_train_step"] p_train_step_dropless = jax_device_state.get("p_train_step_dropless", None) + p_train_step_first_phase = jax_device_state.get("p_train_step_first_phase", None) p_eval_step = jax_device_state["p_eval_step"] p_eval_step_dropless = jax_device_state.get("p_eval_step_dropless", None) @@ -719,26 +869,34 @@ def training_loop_iteration( step_rng_args = () with maybe_record_goodput(recorder, GoodputEvent.STEP, step): with jax.set_mesh(mesh), logical_axis_rules(logical_axis_rules_for_train): - if config.retry_when_tokens_dropped and p_train_step_dropless is not None: - candidate_state, metrics = p_train_step(state, example_batch, *step_rng_args) - if bool(metrics.get("has_moe_overflow")): + in_first_phase = config.retry_dropless_first_steps > 0 and step - start_step < config.retry_dropless_first_steps + if in_first_phase and p_train_step_dropless is not None and p_train_step_first_phase is None: + # First-phase schedule: the steps right after (re)start overflow the ragged buffer most often, + # so run them directly with the dropless program instead of attempt + discard + replay. + max_logging.log(f"Step {step}: dropless first-phase program") + state, metrics = p_train_step_dropless(state, example_batch, *step_rng_args) + _maybe_log_required_rbf(config, step, metrics, "dropless") + else: + # With first_phase_ragged_buffer_factor > 0, first-phase steps attempt the first-phase program (larger + # finite ragged buffer) instead of the normal one; an overflow still replays with the dropless program. + if in_first_phase and p_train_step_first_phase is not None: max_logging.log( - f"Step {step}: MoE ragged buffer overflow detected! " - f"Discarding candidate state and replaying step with dropless buffer..." - ) - # Explicitly deallocate device buffers held by candidate_state before replaying - # with p_train_step_dropless to avoid pinning two ~13.4 GB model states in HBM. - jax.tree_util.tree_map( - lambda x: x.delete() if hasattr(x, "delete") else None, - candidate_state, + f"Step {step}: first-phase program (ragged_buffer_factor={config.first_phase_ragged_buffer_factor})" ) - del candidate_state - gc.collect() - state, metrics = p_train_step_dropless(state, example_batch, *step_rng_args) + p_attempt, attempt_name = p_train_step_first_phase, "first_phase" else: - state = candidate_state - else: - state, metrics = p_train_step(state, example_batch, *step_rng_args) + p_attempt, attempt_name = p_train_step, "normal" + state, metrics = p_attempt(state, example_batch, *step_rng_args) + _maybe_log_required_rbf(config, step, metrics, attempt_name) + if ( + config.moe_dropless_fallback == "step" + and p_train_step_dropless is not None + and bool(metrics.get("has_moe_overflow")) + ): + max_logging.log(f"Step {step}: MoE ragged buffer overflow detected! Replaying step with dropless buffer...") + # The attempt rolled its update back, so `state` is still the pre-step state. + state, metrics = p_train_step_dropless(state, example_batch, *step_rng_args) + _maybe_log_required_rbf(config, step, metrics, "dropless") step_time_delta = datetime.datetime.now() - last_step_completion last_step_completion = datetime.datetime.now() @@ -765,6 +923,13 @@ def training_loop_iteration( ) if ran_eval: assert eval_data_iterator + prefill_thread = python_vars.get("eval_prefill_thread") + if prefill_thread is not None and prefill_thread.is_alive(): + # Let the background cache pass commit before building the eval iterator (never a race in practice: + # the first eval is tens of steps after run_start). + wait_t = time.perf_counter() + prefill_thread.join(timeout=600) + max_logging.log(f"eval cache prefill: waited {time.perf_counter() - wait_t:.3f} s before eval") # Explicitly reset the eval iterator and counters before starting the eval loop if hasattr(eval_data_iterator, "reset"): eval_data_iterator.reset() @@ -781,12 +946,10 @@ def training_loop_iteration( eval_batch = jax.device_put( eval_batch, sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules_for_eval) ) - if 0 < eval_steps <= eval_step_count: - break with jax.set_mesh(mesh), logical_axis_rules(logical_axis_rules_for_eval): eval_metrics = p_eval_step(state, eval_batch, *step_rng_args) if ( - config.retry_when_tokens_dropped + config.moe_dropless_fallback == "step" and p_eval_step_dropless is not None and bool(eval_metrics.get("has_moe_overflow")) ): @@ -801,6 +964,11 @@ def training_loop_iteration( eval_metrics, eval_step_count, step_time_delta=eval_step_time_delta, is_training=False ) eval_step_count += 1 + # Stop before fetching another batch: the extra fetch would reshard a batch that + # is never used, and the data loading hosts may already be out of data while the + # placeholder iterators of the other hosts keep going. + if 0 < eval_steps <= eval_step_count: + break prof.maybe_deactivate_profiler(step, state) @@ -817,8 +985,191 @@ def training_loop_iteration( return metrics -def train_loop(config, recorder, state=None): - """Main Training loop.""" +def build_dropless_graphdef(config, jit_model, state): + """Graphdef of the dropless retry program: RoutedMoE force_dropless + token chunks + barrier, decoder full remat.""" + reconstructed = nnx.merge(jit_model, state) + for _, module in nnx.iter_graph(reconstructed): + if type(module).__name__ == "RoutedMoE": + module.force_dropless = True + module.num_moe_token_chunks = getattr(config, "retry_num_moe_token_chunks", 2) + module.moe_chunk_barrier = True + elif hasattr(module, "get_remat_policy"): # NNXDecoder; the old "Decoder" name matched nothing + module.remat_policy_override = "full" + jit_model_dropless, _ = nnx.split(reconstructed) + del reconstructed, _ + gc.collect() + return jit_model_dropless + + +def build_eval_graphdef(config, jit_model, state): + """Graphdef for the eval program with RoutedMoE ragged_buffer_factor_override = config.eval_ragged_buffer_factor. + + None when eval_ragged_buffer_factor <= 0. In that case the eval program uses the normal graphdef, where + RoutedMoE.get_ragged_buffer_factor() reads eval_ragged_buffer_factor under eval logical axis rules (<= 0 is the + dropless worst-case buffer) and ragged_buffer_factor otherwise. With a value > 0 the override applies to the eval + program whether or not eval runs under custom logical axis rules. + """ + if config.eval_ragged_buffer_factor <= 0: + return None + reconstructed = nnx.merge(jit_model, state) + n_moe = 0 + for _, module in nnx.iter_graph(reconstructed): + if type(module).__name__ == "RoutedMoE": + module.ragged_buffer_factor_override = float(config.eval_ragged_buffer_factor) + n_moe += 1 + max_logging.log(f"eval graphdef: ragged_buffer_factor_override={config.eval_ragged_buffer_factor} on {n_moe} RoutedMoE") + jit_model_eval, _ = nnx.split(reconstructed) + del reconstructed, _ + gc.collect() + return jit_model_eval + + +def build_first_phase_graphdef(config, jit_model, state): + """Graphdef of the first-phase train program: the normal graphdef with RoutedMoE ragged_buffer_factor_override = + config.first_phase_ragged_buffer_factor. None when first_phase_ragged_buffer_factor is 0 (off).""" + if getattr(config, "first_phase_ragged_buffer_factor", 0.0) <= 0: + return None + reconstructed = nnx.merge(jit_model, state) + n_moe = 0 + for _, module in nnx.iter_graph(reconstructed): + if type(module).__name__ == "RoutedMoE": + module.ragged_buffer_factor_override = float(config.first_phase_ragged_buffer_factor) + n_moe += 1 + max_logging.log( + f"first-phase graphdef: ragged_buffer_factor_override={config.first_phase_ragged_buffer_factor} " + f"on {n_moe} RoutedMoE" + ) + jit_model_first_phase, _ = nnx.split(reconstructed) + del reconstructed, _ + gc.collect() + return jit_model_first_phase + + +def aot_compile_step(p_step, lower_args, compiler_options, prefix): + """Lowers and compiles p_step ahead of time (warms the executable cache) and prints its memory stats.""" + compiled = p_step.lower(*lower_args).compile(compiler_options=compiler_options) + max_utils.print_compiled_memory_stats(compiled.memory_analysis(), prefix=prefix) + return compiled + + +def make_synthetic_batch(config, shaped_batch, seed=0): + """Materializes a batch with the keys, shapes, dtypes and shardings of `shaped_batch` (ShapeDtypeStructs) without + touching the dataset: token ids are drawn from jax.random.PRNGKey(seed), positions are an arange over the sequence + axis, segmentations and other integer fields are ones, float fields are zeros. Built inside one jit whose + out_shardings are the structs' shardings, so the result is a global array like the loader's device_put batch.""" + names = sorted(shaped_batch) + seq_len = int(getattr(config, "max_target_length", 0)) + + def _seq_axis(shape): + axes = [i for i, n in enumerate(shape) if n == seq_len and i > 0] + return axes[0] if axes else len(shape) - 1 + + def _gen(key): + out = {} + keys = jax.random.split(key, len(names)) + for sub_key, name in zip(keys, names): + shape, dtype = shaped_batch[name].shape, shaped_batch[name].dtype + if name in ("inputs", "targets"): + out[name] = jax.random.randint(sub_key, shape, 0, max(int(config.vocab_size), 1), dtype=dtype) + elif name.endswith("_position"): + out[name] = jax.lax.broadcasted_iota(dtype, shape, _seq_axis(shape)) + elif jnp.issubdtype(dtype, jnp.floating): + out[name] = jnp.zeros(shape, dtype) + else: + out[name] = jnp.ones(shape, dtype) + return out + + out_shardings = {name: getattr(shaped_batch[name], "sharding", None) for name in names} + return jax.jit(_gen, out_shardings=out_shardings)(jax.random.PRNGKey(seed)) + + +def _array_leaf_ids(tree): + return {id(leaf) for leaf in jax.tree_util.tree_leaves(tree) if isinstance(leaf, jax.Array)} + + +def _delete_outputs(outputs, keep_ids): + """Frees the device buffers of a warmup program's outputs, except leaves that are the kept (real) state's arrays + (jit may forward an unchanged input to its output).""" + for leaf in jax.tree_util.tree_leaves(outputs): + if isinstance(leaf, jax.Array) and id(leaf) not in keep_ids and not leaf.is_deleted(): + leaf.delete() + + +def warmup_programs( + programs, state, mesh, train_rules, eval_rules, train_batch, eval_batch, rng_args=(), copy_train_state=False +): + """warmup_programs_in_init: executes each program once through its jit (the same dispatch path as the loop, so + step 0 hits the jit cache and the executable is loaded) under the loop's mesh / logical-axis-rule contexts. + + programs: list of (name, kind, p_step) with kind "train" or "eval"; entries whose p_step is None are skipped. + The real `state` is never modified: eval programs and non-donating train programs cannot change their inputs; + when copy_train_state is True (train programs donate the state) each train program gets a fresh copy of it. + Every output (new state, metrics, routed-bias updates) is discarded. Returns {name: wall seconds}. + """ + keep_ids = _array_leaf_ids(state) + timings = {} + for name, kind, p_step in programs: + if p_step is None: + continue + rules, batch = (train_rules, train_batch) if kind == "train" else (eval_rules, eval_batch) + if batch is None: + max_logging.log(f"warmup_programs_in_init: skipping {name} (no synthetic {kind} batch)") + continue + start = time.perf_counter() + with jax.set_mesh(mesh), logical_axis_rules(rules): + arg_state = state + if kind == "train" and copy_train_state: + arg_state = jax.tree_util.tree_map(lambda x: x.copy() if isinstance(x, jax.Array) else x, state) + outputs = p_step(arg_state, batch, *rng_args) + jax.block_until_ready(outputs) + timings[name] = time.perf_counter() - start + _delete_outputs(outputs, keep_ids) + if arg_state is not state: + _delete_outputs(arg_state, keep_ids) + del outputs, arg_state + max_logging.log(f"warmup_programs_in_init: {name} ({kind}) {timings[name]:.3f} s") + gc.collect() + return timings + + +def start_run_clock(config, start_step, warmup_fn=None, barrier_fn=None, python_vars=None): + """Logs init_print, then (warmup_programs_in_init) runs warmup_fn and a cross-host barrier, then init_stop, + run_start and block_start. With warmup_fn None the sequence is the original init_print, init_stop, run_start, + block_start. mllog writes only on jax.process_index() == 0; the barrier makes every host leave init together, so + run_start (process 0) sits right before every host's first step.""" + mllog_utils.init_print(config) + if warmup_fn is not None: + start = time.perf_counter() + warmup_fn() + warmed = time.perf_counter() + if barrier_fn is None: + barrier_fn = lambda: multihost_utils.sync_global_devices("warmup_programs_in_init") # pylint: disable=unnecessary-lambda-assignment + barrier_fn() + done = time.perf_counter() + max_logging.log( + f"warmup_programs_in_init: warmup {warmed - start:.3f} s, barrier wait {done - warmed:.3f} s " + f"(process {jax.process_index()})" + ) + if python_vars is not None: + python_vars["last_step_completion"] = datetime.datetime.now() + mllog_utils.init_stop() + mllog_utils.run_start() + mllog_utils.block_start(config, start_step) + + +def train_loop(config, recorder, state=None, elastic_retry_cache=None): + """Main Training loop. + + Args: + config: Config object. + recorder: Goodput recorder. + state: Optional initial train state. + elastic_retry_cache: Optional `elastic_utils.RetryCache`. An elastic retry on the same slices reuses the setup and the + compiled steps cached in it instead of building them again. + + Returns: + The final train state. + """ ( init_rng, checkpoint_manager, @@ -831,7 +1182,7 @@ def train_loop(config, recorder, state=None): rampup_manager, eval_data_iterator, state, - ) = train_utils.setup_train_loop(config, recorder) + ) = train_utils.setup_train_loop(config, recorder, elastic_retry_cache=elastic_retry_cache) # Throttling is applied only if configured (dcn_bandwidth_limit is set). # The default flag value is empty, meaning no throttling is applied by default. @@ -841,24 +1192,18 @@ def train_loop(config, recorder, state=None): train_utils.validate_completed_steps(start_step, config.steps) jit_model_dropless = None + jit_model_eval = None + jit_model_first_phase = None if config.enable_diloco: # state is the DiLoCoTrainState; `model` is already the TrainStateNNX graphdef the inner step needs. jit_model = model else: jit_model, state = nnx.split(state) - if config.retry_when_tokens_dropped: - reconstructed = nnx.merge(jit_model, state) - for _, module in nnx.iter_graph(reconstructed): - if type(module).__name__ == "RoutedMoE": - module.force_dropless = True - module.num_moe_token_chunks = getattr(config, "retry_num_moe_token_chunks", 2) - module.moe_chunk_barrier = True - elif type(module).__name__ == "Decoder": - module.remat_policy_override = "full" - jit_model_dropless, _ = nnx.split(reconstructed) - del reconstructed, _ - gc.collect() + if config.moe_dropless_fallback == "step": + jit_model_dropless = build_dropless_graphdef(config, jit_model, state) + jit_model_first_phase = build_first_phase_graphdef(config, jit_model, state) + jit_model_eval = build_eval_graphdef(config, jit_model, state) if config.enable_diloco: # DiLoCoTrainState.params already holds the param shardings the inner step needs; @@ -867,18 +1212,36 @@ def train_loop(config, recorder, state=None): else: params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings) - p_train_step, p_eval_step = train_utils.jit_train_and_eval_step( - config, - jit_model, - mesh, - state, - state_mesh_shardings, - train_step, - eval_step, - eval_data_iterator, - params_shardings, + # Jit the train and eval steps, or reuse the ones compiled by the previous elastic attempt. A pause/resume retry + # comes back on the same slices, so the compiled programs still fit and the retrace, lowering and compile are + # skipped. After a replica resize the cache has been emptied, so the steps are compiled again for the new mesh. + p_train_step, p_eval_step = elastic_utils.get_or_build( + elastic_retry_cache, + "train_and_eval_step", + functools.partial( + train_utils.jit_train_and_eval_step, + config, + jit_model, + mesh, + state, + state_mesh_shardings, + train_step, + eval_step, + eval_data_iterator, + params_shardings, + ), ) + if jit_model_eval is not None and p_eval_step is not None: + # eval_ragged_buffer_factor > 0: the eval program uses its own graphdef (RoutedMoE ragged_buffer_factor_override). + p_eval_step = train_utils.jit_eval_step( + config, + jit_model_eval, + state_mesh_shardings, + sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules_for_eval), + eval_step, + ) + p_train_step_dropless = None p_eval_step_dropless = None if jit_model_dropless is not None: @@ -894,9 +1257,24 @@ def train_loop(config, recorder, state=None): params_shardings=params_shardings, ) + p_train_step_first_phase = None + if jit_model_first_phase is not None: + p_train_step_first_phase, _ = train_utils.jit_train_and_eval_step( + config, + jit_model_first_phase, + mesh, + state, + state_mesh_shardings, + train_step, + params_shardings=params_shardings, + ) + # Do not enter the legacy `mesh` context manager here: the training loop calls # p_train_step without it, and the mismatch in jit's tracing-cache key would # cause train_step to be traced and compiled a second time on the first step. + precompiled_train = False + precompiled_eval = False + shaped_eval_batch = None with jax.set_mesh(mesh), logical_axis_rules(config.logical_axis_rules): data_sharding = sharding.get_input_data_sharding(config, mesh) shaped_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding) @@ -919,15 +1297,24 @@ def train_loop(config, recorder, state=None): compiled = p_train_step.lower(*lower_args).compile(compiler_options=compiler_options) compiled_stats = compiled.memory_analysis() max_utils.print_compiled_memory_stats(compiled_stats, prefix="train") - + precompiled_train = True + if p_train_step_dropless is not None: + # Precompile the dropless retry program here so no compile happens inside the timed loop. + aot_compile_step(p_train_step_dropless, lower_args, compiler_options, prefix="train_dropless") + if p_train_step_first_phase is not None: + aot_compile_step(p_train_step_first_phase, lower_args, compiler_options, prefix="train_first_phase") + + if p_eval_step is not None: + # Also used by warm_eval_input_reshard_before_run_start, which applies with a pre-compiled train step too. + with jax.set_mesh(mesh), logical_axis_rules(config.logical_axis_rules_for_eval): + data_sharding_eval = sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules_for_eval) + shaped_eval_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding_eval, is_eval=True) # Ahead-of-time compile the evaluation step alongside the training step to # warm up the XLA executable cache and avoid JIT compilation pause on the # first eval step. if p_eval_step is not None and config.compiled_trainstep_file == "" and not jax.config.jax_enable_pgle: with jax.set_mesh(mesh), logical_axis_rules(config.logical_axis_rules_for_eval): compiler_options = max_utils.parse_libtpu_flags_to_dict(config.compile_xla_flags) - data_sharding_eval = sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules_for_eval) - shaped_eval_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding_eval, is_eval=True) if config.enable_diloco: eval_lower_args = (state, shaped_eval_batch, init_rng) else: @@ -935,6 +1322,9 @@ def train_loop(config, recorder, state=None): compiled_eval = p_eval_step.lower(*eval_lower_args).compile(compiler_options=compiler_options) compiled_eval_stats = compiled_eval.memory_analysis() max_utils.print_compiled_memory_stats(compiled_eval_stats, prefix="eval") + precompiled_eval = True + if p_eval_step_dropless is not None: + aot_compile_step(p_eval_step_dropless, eval_lower_args, compiler_options, prefix="eval_dropless") prof = profiler.Profiler(config, offset_step=start_step) metric_logger_instance = metric_logger.MetricLogger( config=config, learning_rate_schedule=learning_rate_schedule, start_step=start_step @@ -957,6 +1347,7 @@ def train_loop(config, recorder, state=None): "state_mesh_shardings": state_mesh_shardings, "p_train_step": p_train_step, "p_train_step_dropless": p_train_step_dropless, + "p_train_step_first_phase": p_train_step_first_phase, "p_eval_step": p_eval_step, "p_eval_step_dropless": p_eval_step_dropless, "model": model, @@ -996,15 +1387,73 @@ def train_loop(config, recorder, state=None): "dump_hlo_upload_all": config.dump_hlo_upload_all, } + warmup_fn = None + if getattr(config, "warmup_programs_in_init", False): + if not (precompiled_train or precompiled_eval): + max_logging.log("warmup_programs_in_init: no precompiled program (compiled_trainstep_file or AutoPGLE); skipped") + else: + warmup_list = [] + if precompiled_train: + warmup_list += [ + ("train", "train", p_train_step), + ("train_first_phase", "train", p_train_step_first_phase), + ] + if config.warmup_skip_train_dropless: + if p_train_step_dropless is not None: + max_logging.log("warmup_programs_in_init: skipping train_dropless (warmup_skip_train_dropless=True)") + else: + warmup_list.append(("train_dropless", "train", p_train_step_dropless)) + if precompiled_eval: + warmup_list += [("eval", "eval", p_eval_step), ("eval_dropless", "eval", p_eval_step_dropless)] + # Same shardings as the loop's batches: the loader's device_put sharding for train, the eval rules for eval. + train_batch_sharding = getattr(data_loader, "input_data_shardings", None) or data_sharding + train_shaped = { + k: jax.ShapeDtypeStruct(v.shape, v.dtype, sharding=train_batch_sharding) for k, v in shaped_batch.items() + } + warmup_train_batch = make_synthetic_batch(config, train_shaped, seed=0) if precompiled_train else None + warmup_eval_batch = make_synthetic_batch(config, shaped_eval_batch, seed=1) if precompiled_eval else None + # train_step always donates the state (maxtext_utils.get_functional_train_with_signature, including under + # moe_dropless_fallback="step", which rolls back in-graph), so each warmed train program gets a copy. + copy_train_state = True + # pylint: disable=not-callable + warmup_rng_args = (jax.jit(jax.random.fold_in)(init_rng, start_step),) if config.enable_diloco else () + + def _run_warmup(): + warmup_programs( + warmup_list, + state, + mesh, + config.logical_axis_rules, + config.logical_axis_rules_for_eval, + warmup_train_batch, + warmup_eval_batch, + rng_args=warmup_rng_args, + copy_train_state=copy_train_state, + ) + + warmup_fn = _run_warmup + _job_completed_gracefully = False te_moe_overflow_window = [] try: + train_utils.prepare_before_run_start( + config, + state, + learning_rate_schedule, + start_step, + mesh, + data_loader, + shaped_batch, + shaped_eval_batch=shaped_eval_batch, + ) python_vars["last_step_completion"] = datetime.datetime.now() - mllog_utils.init_print(config) - mllog_utils.init_stop() - mllog_utils.run_start() - mllog_utils.block_start(config, start_step) + start_run_clock(config, start_step, warmup_fn=warmup_fn, python_vars=python_vars) + + # Inside the timed region: overlap the eval pipeline's file reads with the first training steps. + python_vars["eval_prefill_thread"] = None + if config.eval_cache_prefill_in_background and eval_data_iterator is not None: + python_vars["eval_prefill_thread"] = train_utils.start_eval_cache_prefill(eval_data_iterator) # Using while loop to allow for potential dynamic 'steps' adjustment in future while python_vars["step"] < immutable_data["steps"]: @@ -1112,10 +1561,10 @@ def initialize(argv: Sequence[str]) -> tuple[pyconfig.HyperParameters, Any]: return config, recorder -def run(config, recorder): +def run(config, recorder, elastic_retry_cache=None): """Run the job given hyperparameters and utilities.""" with (max_utils.maybe_get_transformer_engine_context(config),): - train_loop(config, recorder) + train_loop(config, recorder, elastic_retry_cache=elastic_retry_cache) def get_train_func(config, recorder, argv): @@ -1123,24 +1572,38 @@ def get_train_func(config, recorder, argv): if config.elastic_enabled: max_logging.log("Elastic utils: Elastic training enabled.") - def on_elastic_event(): - elastic_utils.record_elastic_event_start(recorder, config) + # Setup that an elastic retry on the same slices can reuse: the `initialize` result, the mesh, the abstract + # model and the jitted steps. `elastic_retry` turns it on for the run and empties it when the slices change, + # so a pause/resume retry reuses everything and a replica resize rebuilds everything. It is shared by + # `elastic_retry` and the attempts it runs, and is passed down to `train_loop` and `setup_train_loop`. It + # can't live in the config: the config built for an attempt is itself one of the cached objects. + elastic_retry_cache = elastic_utils.RetryCache() + + def on_elastic_event(scale_up): + elastic_utils.record_elastic_event_start(recorder, scale_up=scale_up) def on_slices_ready(): elastic_utils.record_elastic_wait_end_and_reinit_start(recorder) def elastic_train_wrapper(argv: Sequence[str]) -> None: """Wrapper for elastic training initializes variables and runs the train loop.""" - elastic_config, elastic_recorder = initialize(argv) + # Reuse the config and goodput recorder from the previous attempt if it ran on the same slices. After a + # replica resize `initialize` runs again so the config picks up the new slice count. + elastic_config, elastic_recorder = elastic_utils.get_or_build( + elastic_retry_cache, "initialize", functools.partial(initialize, argv) + ) run( elastic_config, elastic_recorder, + elastic_retry_cache=elastic_retry_cache, ) + # Every attempt goes through `elastic_train_wrapper`; the cache decides how much of it is skipped. train_func = elastic_utils.elastic_retry( config, callback_fn=on_elastic_event, pre_callback_fn=on_slices_ready, + retry_cache=elastic_retry_cache, )(functools.partial(elastic_train_wrapper, argv=argv)) else: # Use the already initialized variables diff --git a/src/maxtext/trainers/pre_train/train_compile.py b/src/maxtext/trainers/pre_train/train_compile.py index f2ff3b3f08..6b237eca53 100644 --- a/src/maxtext/trainers/pre_train/train_compile.py +++ b/src/maxtext/trainers/pre_train/train_compile.py @@ -41,7 +41,6 @@ from maxtext.common import train_state_nnx from maxtext.common.common_types import ShardMode from maxtext.configs import pyconfig -from maxtext.models import models from maxtext.optimizers import optimizers from maxtext.trainers.diloco import diloco from maxtext.trainers.pre_train import train @@ -55,8 +54,6 @@ # pylint: disable=too-many-positional-arguments -Transformer = models.transformer_as_linen - def validate_config(config): """Validates the config is is setup correctly to compile, returning a useful error message if not.""" @@ -173,7 +170,7 @@ def create_train_state_fn(): # Collect NNX activation shardings via an abstract forward pass (must run # after get_abstract_state, which only traces __init__). if config.debug_sharding: - _collect_nnx_activation_shardings(_create_model_partial, config, topology_mesh) # pyrefly: ignore[unbound-name] + _collect_nnx_activation_shardings(_create_model_partial, config, topology_mesh) return ( shaped_train_args, @@ -346,7 +343,9 @@ def main(argv: Sequence[str]) -> None: # Wrap train_step with diloco train_step_partial = functools.partial(train.train_step, model, config, inner_state_shardings, params_shardings) - train_step_fn = diloco.build_diloco_train_step(config, train_step_partial) + train_step_fn = diloco.build_diloco_train_step( + config, train_step_partial, outer_params_shardings=state_mesh_shardings.params + ) # For DiLoCo, the train_step_fn is already fully wrapped and takes (state, batch, prng) func_to_compile = train_step_fn diff --git a/src/maxtext/training_engine/README.md b/src/maxtext/training_engine/README.md index 58b2797dff..d3da9b81a9 100644 --- a/src/maxtext/training_engine/README.md +++ b/src/maxtext/training_engine/README.md @@ -344,6 +344,7 @@ sharding at construction (for example Tokamax ring attention), while under `jax. | `parameter_memory_host_offload` | not supported; raises | | `shard_optimizer_over_data` (Zero-1) | shards the optimizer state over the `data` axis inside `update`; requires `shard_mode=explicit` | | `skip_step_on_nan` (default true), `skip_step_on_spikes`, `max_grad_norm_spike` | a non-finite or spiking step leaves the state unchanged and records `step_skipped` | +| `trainable_parameters_mask` | parameters it does not match are frozen. `fwd_bwd` stops their gradient inside the function it differentiates, so it returns zeros for them without computing a weight gradient, or reducing one inside `fwd_bwd`. The zeros keep their slots in the accumulator (and in `update`'s all-reduce, when that is deferred), so kernel signatures and checkpoints are unchanged. `update` zeroes them again before the gradient norm, clipping and the optimizer, so the norm and clipping cover the trainable parameters only. On CPU the trainable gradients, the norm and the trained weights are bitwise those of computing the frozen gradients and discarding them (`maxtext_engine_frozen_grad_test.py`) | | `retry_when_tokens_dropped`, `retry_num_moe_token_chunks` | a micro-batch whose MoE ragged buffer overflowed is accumulated as zero and rerun through `fwd_bwd_dropless`, built on a copy of the model with dropless MoE layers (that many token chunks, as a rematerialized scan) and full remat; compiled on a background thread from the first micro-batch. Its flag is read after the next micro-batch is dispatched, or in `update` for the last; `update` records `moe_dropless_replays`. `eval_step` and `model_scope`/`fwd_only` run the dropless graph too | | `enable_checkpointing` | `false` disables saving and restoring the train state; `load_parameters_path` still loads weights | | `gradient_accumulation_steps` | read only by the standalone script; Tunix chooses the number of micro-batches itself | @@ -391,6 +392,7 @@ The engine's tests are in `tests/post_training/unit/` and run on CPU; they requi | `maxtext_engine_model_build_test.py` | model construction with and without a mesh set by the caller | | `maxtext_engine_checkpoint_test.py` | resuming from a mid-step checkpoint through Orbax reproduces the uninterrupted step | | `maxtext_engine_grad_cast_test.py` | `cast_grads_after_all_reduce`: which gradients leave `fwd_bwd` uncast, and that moving their cast changes no number on CPU | +| `maxtext_engine_frozen_grad_test.py` | `trainable_parameters_mask`: `fwd_bwd` has no weight-gradient matmul for a frozen kernel and returns zero for it, while the other gradients, the gradient norms and the trained weights stay bitwise those of computing it, under fsdp, under the deferred all-reduce and on the eager path; a stop on a trainable kernel is caught | | `maxtext_engine_prepared_input{,_dropless,_sharded}_test.py` | `prepare_fwd_bwd_input` and `wait_for_device`: a prepared micro-batch trains bit-identically on one device, through Tunix's `TrainerWorker` staging, under the dropless replay and final dropless programs, and on a four-device mesh with Tunix's payload; the wait leaves the in-flight queue and the metrics writes as they were | | `maxtext_engine_stage_before_update_test.py` | the order Tunix uses to stage before a deferred update (`fwd_bwd(last)`, `prepare_fwd_bwd_input(next)`, `wait_for_device()`, `update()`) under the dropless replay and final dropless programs, directly and through Tunix's `TrainerWorker`: bit-identical to the base order, including a rerun of the last micro-batch inside `update()` | | `maxtext_engine_{constructor,data_parallel,xaot,packing,profiling}_test.py`, `router_replay_engine_test.py` | construction, data parallelism, ahead-of-time compilation, sequence packing, profiling, router replay | @@ -404,3 +406,7 @@ The engine's tests are in `tests/post_training/unit/` and run on CPU; they requi of 0, are unaffected. - `pre_train/train.py` sums micro-batch gradients in the dtype of the parameters it differentiates; the engine sums them in `grad_accumulation_dtype`, which defaults to `grad_dtype`. +- With `trainable_parameters_mask`, `pre_train/train.py` computes the frozen parameters' gradients and + includes them in the gradient norm and clipping; only its optimizer zeroes their updates. The engine + leaves them out of the norm and clipping and does not compute them, which is why stopping their + gradient changes no number there. diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 0f1d0ddc1b..987de28389 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -1530,6 +1530,38 @@ def _to_grad_dtype(self, x: Any) -> Any: return x return x.astype(self._config.grad_dtype) + def _freeze_mask_for(self, tree: Any) -> Any: + """Returns which leaves of `tree` `trainable_parameters_mask` freezes, or None if it freezes none. + + A pytree of Python bools with `tree`'s structure, read off the leaves' paths alone, so `tree` can + be the parameters or a gradient tree shaped like them. `_compile_for_batch` caches the parameters' + mask in `_freeze_mask`; while that cache is empty -- on the eager path, or for a kernel traced + after `_invalidate_pure_state` -- the mask is computed from `tree` instead. `_fwd_bwd_kernel` and + `_update_kernel` both read it here, so the leaves whose gradient `fwd_bwd` stops are exactly the + leaves `update` zeroes. + """ + # Read once: a background compile can trace a kernel while the main thread empties the cache + # (`_disable_pure_state`), and a second read could then return None and skip the stop. + freeze_mask = self._freeze_mask + if freeze_mask is not None: + return freeze_mask + if self._freeze_mask_fn is None: + return None + return self._freeze_mask_fn(tree) + + def _stop_frozen_gradients(self, params: Any) -> Any: + """Returns `params` with `jax.lax.stop_gradient` on every leaf `trainable_parameters_mask` freezes. + + Takes effect only inside the function being differentiated: `_fwd_bwd_kernel`'s `diff_wrapper` + calls it on the parameters it differentiates. See the comment there for why it changes no number. + The leaves stay in the gradient tree, as zeros, so every kernel's signature, the accumulator's + layout and the checkpoint format are unchanged. + """ + freeze_mask = self._freeze_mask_for(params) + if freeze_mask is None: + return params + return jax.tree.map(lambda x, frozen: jax.lax.stop_gradient(x) if frozen else x, params, freeze_mask) + def _fwd_bwd_kernel(self, params, rest, batch, graphdef=None): """Executes a single forward and backward pass. @@ -1543,7 +1575,8 @@ def _fwd_bwd_kernel(self, params, rest, batch, graphdef=None): Returns: `(primary_loss, aux_metrics, new_rest, grads, denominator)`: this micro-batch's gradients, not yet normalized, in the accumulation dtype except the leaves in `_uncast_grad_mask`, which - keep the parameters' dtype; and the loss denominator to normalize them by. Under + keep the parameters' dtype, and exactly zero for the parameters `trainable_parameters_mask` + freezes; and the loss denominator to normalize them by. Under `retry_when_tokens_dropped`, followed by `has_moe_overflow`: whether any MoE layer's ragged buffer dropped tokens, in which case the gradients and denominator are zero and `fwd_bwd` reruns the micro-batch once the flag reaches the host. @@ -1551,6 +1584,19 @@ def _fwd_bwd_kernel(self, params, rest, batch, graphdef=None): loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn def diff_wrapper(p, r, b): + # No gradient for frozen parameters. `update()` zeroes the gradient of every leaf that + # `trainable_parameters_mask` freezes before the norm, clipping and the optimizer see it, so + # nothing uses that gradient; stopped here, inside the differentiated function, it is not computed + # either (on `value_and_grad`'s arguments the stop would change nothing). JAX emits no + # weight-gradient op for those leaves -- a custom-VJP kernel still runs its whole backward, and + # XLA drops the cotangent nobody reads where it can -- and `fwd_bwd` returns a constant zero for + # them, with nothing to reduce-scatter or all-reduce inside it. The gradient through a frozen + # weight into the activations, and on to the trainable parameters, comes from the same ops as + # before, so the rest is the same computation: bitwise equal on CPU + # (`maxtext_engine_frozen_grad_test.py`), though on TPU XLA may fuse the smaller program + # differently. The MLPerf RL recipes freeze the MoE router (`routed_experts/gate/kernel`), so this + # drops the router's weight gradient and its reduce-scatter from every MoE layer. + p = self._stop_frozen_gradients(p) mdl = nnx.merge(graphdef if graphdef is not None else self._model_graphdef, p, r, copy=True) if self._gen_model_input_fn is not None: # A gen_model_input_fn maps a payload to the loss fn's *keyword arguments* -- see @@ -1789,10 +1835,12 @@ def normalize(g): # Cast to `grad_dtype` once, after the sum and the division, as `train.py` does with its # accumulated `raw_grads`. A no-op unless `grad_accumulation_dtype` differs from it. grads = jax.tree.map(self._to_grad_dtype, grads) - freeze_mask = self._freeze_mask - if freeze_mask is None and self._freeze_mask_fn is not None: - freeze_mask = self._freeze_mask_fn(accumulated_grads) + freeze_mask = self._freeze_mask_for(accumulated_grads) if freeze_mask is not None: + # `fwd_bwd` already returns zeros for these leaves (`_stop_frozen_gradients`). Zeroing them again + # costs nothing and still matters: an accumulator restored from an intra-step checkpoint written + # by an older engine, or under another mask, can hold non-zero gradients for them, and it keeps + # frozen parameters out of the norm and clipping below by construction. def _apply_freeze(g, is_frozen): return jnp.zeros_like(g) if is_frozen else g @@ -1903,6 +1951,12 @@ def _needs_recompile(self, signature: Any, previous: Any) -> bool: def _prepare_batch(self, payload: Any) -> Any: """Maps a payload to the inputs the loss function is called with.""" + if getattr(payload, "segment_ids", None) is not None and self._config.use_gdn_kernel and not self._config.packing: + raise ValueError( + "The payload carries segment_ids, which Tunix sets only on sequence-packed rows, but packing=False: " + "with use_gdn_kernel=True the GDN kernel then ignores those document boundaries in training. " + "Set packing=True." + ) if self._gen_model_input_fn is not None: if dataclasses.is_dataclass(payload) and "metadata" in payload.__dataclass_fields__: payload = dataclasses.replace(payload, metadata={}) diff --git a/src/maxtext/training_engine/micro_step_profiler.py b/src/maxtext/training_engine/micro_step_profiler.py index 9dc407af44..3de95ec3e7 100644 --- a/src/maxtext/training_engine/micro_step_profiler.py +++ b/src/maxtext/training_engine/micro_step_profiler.py @@ -69,6 +69,10 @@ def __init__(self, config: pyconfig.HyperParameters) -> None: "TPU advanced profiling enabled: %s", self.profiling_options.advanced_configuration, ) + if config.enable_continuous_profiling: + advanced = dict(getattr(self.profiling_options, "advanced_configuration", None) or {}) + advanced["enable_continuous_profiling"] = True + self.profiling_options.advanced_configuration = advanced self.profile_active: bool = False # profiler_period is the distance between successive profiling window, # profiler_steps=3, profiler_period=4 gives [1,3], [5,7], [9,11]. diff --git a/src/maxtext/utils/accelerator_to_spec_map.py b/src/maxtext/utils/accelerator_to_spec_map.py index 9732d19a7f..aca5056022 100644 --- a/src/maxtext/utils/accelerator_to_spec_map.py +++ b/src/maxtext/utils/accelerator_to_spec_map.py @@ -15,8 +15,9 @@ """Static map of TPU names such as v4-8 to properties such as chip layout.""" """ !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! -IF YOU MODIFY THIS FILE YOU SHOULD ALSO ADD CORRESPONDING MODIFICATIONS TO -UserFacingNameToSystemCharacteristics in xpk/xpk.py !!!!! """ +IF YOU MODIFY THIS FILE YOU SHOULD ALSO CHECK THAT THE LAUNCHERS SUPPORT THE SAME +ACCELERATORS: AcceleratorShorthandMap in Cluster Toolkit (pkg/config/hardware.go) and, +for legacy XPK, src/xpk/core/system_characteristics.py !!!!! """ from dataclasses import dataclass diff --git a/src/maxtext/utils/diloco_sharding.py b/src/maxtext/utils/diloco_sharding.py index 354002e91e..b80ce343bc 100644 --- a/src/maxtext/utils/diloco_sharding.py +++ b/src/maxtext/utils/diloco_sharding.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Mesh sharding, batch dimension reshaping, and metric extraction utilities for SPMD DiLoCo.""" +"""Mesh sharding, batch dimension reshaping, and metric extraction utilities for DiLoCo.""" from collections.abc import Sequence @@ -20,6 +20,7 @@ import jax import jax.numpy as jnp from jaxtyping import PyTree +import numpy as np def add_diloco_to_sharding(pytree: PyTree) -> PyTree: @@ -103,3 +104,43 @@ def extract_per_island_metrics(metrics: PyTree, num_diloco_replicas: int) -> PyT loss_i = drjax.reduce_sum(metrics["scalar"]["learning/loss"] * mask_i) default_metrics["scalar"][f"learning/loss_island_{i}"] = loss_i return default_metrics + + +def split_mesh_along_axis(mesh: jax.sharding.Mesh, axis_name: str = "diloco") -> list[jax.sharding.Mesh]: + """Splits `mesh` into one submesh per index of `axis_name`, each without that axis. + + Submesh `i` holds the devices at position `i` along `axis_name`, with the remaining axes (and their axis types) in + their original order, so any PartitionSpec that does not mention `axis_name` is valid on every submesh. + """ + if axis_name not in mesh.axis_names: + raise ValueError(f"Axis {axis_name!r} not found in mesh axis names {mesh.axis_names}.") + axis_index = mesh.axis_names.index(axis_name) + keep = [i for i in range(len(mesh.axis_names)) if i != axis_index] + axis_names = tuple(mesh.axis_names[i] for i in keep) + axis_types = tuple(mesh.axis_types[i] for i in keep) + return [ + jax.sharding.Mesh(np.take(mesh.devices, i, axis=axis_index), axis_names, axis_types=axis_types) + for i in range(mesh.shape[axis_name]) + ] + + +def remove_mesh_axis_from_rules( + logical_axis_rules: Sequence[tuple[str, str | Sequence[str] | None]], axis_name: str = "diloco" +) -> list[tuple[str, str | tuple[str, ...] | None]]: + """Drops `axis_name` from the physical axes of every logical axis rule. + + Rules that map only to `axis_name` are removed, so the logical axis falls through to the next matching rule, which + is what a mesh built by `split_mesh_along_axis` needs. + """ + new_rules = [] + for logical_axis, physical_axes in logical_axis_rules: + if isinstance(physical_axes, str): + if physical_axes == axis_name: + continue + elif physical_axes is not None: + stripped = tuple(ax for ax in physical_axes if ax != axis_name) + if physical_axes and not stripped: + continue + physical_axes = stripped + new_rules.append((logical_axis, physical_axes)) + return new_rules diff --git a/src/maxtext/utils/elastic_utils.py b/src/maxtext/utils/elastic_utils.py index 823d79054e..d3e43c121a 100644 --- a/src/maxtext/utils/elastic_utils.py +++ b/src/maxtext/utils/elastic_utils.py @@ -15,7 +15,10 @@ """Utility functions for Elastic Training.""" from collections import Counter +import dataclasses +import functools from types import SimpleNamespace +from typing import Any, Callable import jax from maxtext.utils import gcs_utils @@ -52,10 +55,19 @@ def record_slice_state(recorder, active_slices_override: int | None = None) -> N ) -def record_elastic_event_start(recorder, config) -> None: - """Records start of an elastic scale up event.""" +def record_elastic_event_start(recorder, scale_up: bool) -> None: + """Records the start of an elastic event. + + Args: + recorder: Goodput recorder. + scale_up: True if the attempt was interrupted by `maybe_elastic_scale_up` because new slices are available, False + if it failed because a slice went down. + """ global pending_elastic_event_type - event_type = "elastic_scale_up" if is_scale_up_event(config) else "elastic_slice_down" + if scale_up: + event_type = "elastic_scale_up" + else: + event_type = "elastic_slice_down" pending_elastic_event_type = event_type if recorder and hasattr(recorder, "record_elastic_wait_start_time"): recorder.record_elastic_wait_start_time(event_type=event_type) @@ -111,7 +123,7 @@ def should_use_elastic(config) -> bool: return config is not None and elastic_enabled(config) -def clean_up_checkpoints(checkpoint_dir: str): +def clean_up_incomplete_checkpoints(checkpoint_dir: str): """Cleans up incomplete checkpoints after an elastic event.""" max_logging.log("Elastic utils: Checking for incomplete checkpoint after an elastic event...") checkpoint_dir = gcs_utils.add_trailing_slash(checkpoint_dir) @@ -152,6 +164,24 @@ def ensure_elastic_manager_initialized(config): elastic_manager = manager.Manager() +def is_pause_resume(config) -> bool: + """Returns whether every attempt waits for all the slices in the jobset, i.e. pause/resume rather than resize. + + In pause/resume every attempt comes back on the same slices, so the compiled steps still fit and the setup of the + previous attempt can be reused. `elastic_min_slice_count` is -1 (wait for every slice) or equals the number of + slices in the jobset. `total_slice_count` counts every slice in the jobset, live or not; `config.num_slices` can't + be used here because it only counts the live slices. + + Args: + config: Config object. + """ + ensure_elastic_manager_initialized(config) + assert elastic_manager is not None + if config.elastic_min_slice_count == -1: + return True + return config.elastic_min_slice_count == elastic_manager.total_slice_count + + def get_local_batch_size(config) -> int: """Returns the local batch size based on the config.""" return config.per_device_batch_size * get_devices_per_host(config) @@ -201,19 +231,139 @@ def wrapper(): return wrapper -def elastic_retry(config, callback_fn=None, pre_callback_fn=None): +class RetryCache: + """Setup objects that a retry on the same slices can reuse, e.g. the config, mesh and compiled steps. + + `elastic_retry` turns the cache on for the duration of one elastic run and empties it whenever the active slices + change, so an entry is only ever returned to an attempt on the same slices as the attempt that stored it. + + In practice only pause/resume (see `is_pause_resume`) gets cache hits: every attempt runs on the same slices, so + the config, mesh and compiled steps of the previous attempt are reused. In replica-resize mode every elastic event + changes the active slices, so the cache is emptied before the next attempt and everything is rebuilt. Every + elastic run still gets a cache, so both modes share one code path. + + Don't cache the train state or other large arrays: they would keep the previous attempt's memory alive. The one + deliberate exception is `snapshotter`, which is not a cache entry: it is set once for the run, outlives every + attempt and every slice change, and is reset by the elastic event callback rather than by this class. + """ + + def __init__(self): + # Whether `get` and `put` do anything. Only `elastic_retry` changes it: it is set to True when the decorated + # function starts and back to False in its `finally` block, once the run is over. While it is False, `get` + # always misses and `put` is a no-op, so a `train_loop` that is not running under `elastic_retry` never caches. + self._enabled = False + self._slices: frozenset[int] | None = None + self._values: dict[str, Any] = {} + # Host-memory snapshots of the train state for `elastic_backup_kind: snapshot`. Not wired up yet: it stays None + # until snapshot save and restore land, and `reset` and `clear_if_slices_changed` never touch it. + self.snapshotter: Any = None + + @property + def enabled(self) -> bool: + return self._enabled + + def reset(self, enabled: bool) -> None: + """Empties the cache and turns it on or off.""" + self._enabled = enabled + self._slices = None + self._values.clear() + + def clear_if_slices_changed(self, active_slices: frozenset[int]) -> None: + """Empties the cache if this attempt runs on different slices than the previous one. + + Args: + active_slices: The slices the attempt about to start runs on. + """ + if self._slices != active_slices: + self._values.clear() + self._slices = active_slices + + def contains(self, name: str) -> bool: + """Returns whether an earlier attempt on the same slices cached an object under `name`. + + Use this rather than `get(name) is None` to test for a hit: a cached value may itself be None. + """ + return self._enabled and name in self._values + + def get(self, name: str) -> Any: + """Returns the object cached under `name` by an earlier attempt on the same slices, or None. + + None is also returned when the cache is off or the entry is missing; use `contains` to tell these apart. + """ + if not self._enabled: + return None + return self._values.get(name) + + def put(self, name: str, value: Any) -> None: + """Caches `value` under `name`, so that a retry on the same slices can reuse it. A no-op while the cache is off. + + Args: + name: The name to cache `value` under. + value: The object to reuse, for example a mesh or a jitted step. + """ + if self._enabled: + self._values[name] = value + + +def get_or_build(retry_cache: RetryCache | None, name: str, build_fn: Callable[[], Any]) -> Any: + """Returns the object cached under `name`, or builds it with `build_fn` and caches it for the next attempt. + + Args: + retry_cache: The cache of the current elastic run, or None when training is not elastic. With None, `build_fn` + is simply called. + name: The name the object is cached under. + build_fn: Builds the object when there is no cached one. + """ + if retry_cache is None: + return build_fn() + if retry_cache.contains(name): + return retry_cache.get(name) + value = build_fn() + retry_cache.put(name, value) + return value + + +@dataclasses.dataclass +class _AttemptOutcome: + """Why the last attempt of an elastic run ended, for the elastic event callback. + + pathwaysutils resets `available_inactive_slices` before it runs the elastic event callback, so the callback can't + tell a scale-up from a slice-down by looking at the manager. The attempt itself records it here on the way out. + """ + + scale_up: bool = False + + +def elastic_retry(config, callback_fn=None, pre_callback_fn=None, retry_cache=None): """Decorator for elastic retry. If an elastic event occurs, the decorator will retry the decorated function up to `config.elastic_max_retries` times. Before each retry, it cleans up partial checkpoints by calling - `clean_up_checkpoints`. If `callback_fn` is provided, it is - called after `clean_up_checkpoints`. + `clean_up_incomplete_checkpoints`. If `callback_fn` is provided, it is + called after `clean_up_incomplete_checkpoints`. + + The decorator covers both ways of coming back from an elastic event: + + * Pause/resume: `elastic_min_slice_count` equals the number of slices in the jobset, or is -1, which means the + same. Every attempt waits for all the slices, so it always comes back on the same slices and the compiled steps + cached in `retry_cache` are not compiled again. + * Replica resize: `elastic_min_slice_count` is smaller than the number of slices in the jobset. An attempt can + come back on fewer slices (slice-down) or on more (scale-up). The retry cache is emptied whenever the slices + change, so the next attempt rebuilds the config, mesh and compiled steps for the slices it has. + + The device memory of a failed attempt is freed by JAX once the attempt has unwound and nothing references its + arrays any more; nothing is deleted explicitly. Args: config: Config object. - callback_fn: Optional callback function to be called after - `clean_up_checkpoints` on an elastic event. + callback_fn: Optional callback called after `clean_up_incomplete_checkpoints` on an elastic event, as + `callback_fn(scale_up=...)`: True if the attempt was interrupted by `maybe_elastic_scale_up` because new slices + are available, False if a slice went down. + pre_callback_fn: Optional callback function to be called before each attempt, once the slices are ready. + retry_cache: Optional `RetryCache` shared with the decorated function. It is turned on for the duration of the + elastic run and emptied whenever the active slices change. The caller owns it rather than `config` because + the config built for an attempt is itself one of the cached objects. Without it nothing is reused. Returns: A decorator for elastic retry. @@ -240,21 +390,83 @@ def cleanup_iterators_and_checkpoints(): cleanup_all_iterators() except Exception as e: max_logging.log(f"Failed to cleanup iterators during elastic scale up: {e}") - clean_up_checkpoints(config.checkpoint_dir) - - if callback_fn is None: - effective_callback = cleanup_iterators_and_checkpoints + clean_up_incomplete_checkpoints(config.checkpoint_dir) + + outcome = _AttemptOutcome() + + def effective_callback(): + scale_up = outcome.scale_up + # `scale_up` only decides the goodput label of this event (`elastic_scale_up` vs `elastic_slice_down`). Only + # `attempt` sets it, and only to True when the attempt raised `ScaleUpSignalError`; a slice-down leaves it + # untouched. Clear it here so it describes just the event being reported: otherwise a slice-down that follows a + # scale-up in the same run would be reported to goodput as a second scale-up. + outcome.scale_up = False + cleanup_iterators_and_checkpoints() + if callback_fn is not None: + callback_fn(scale_up=scale_up) + + def effective_pre_callback(): + if retry_cache is not None: + # Objects built for other slices are stale after a resize. + retry_cache.clear_if_slices_changed(frozenset(elastic_manager.active_slice_indices)) + if pre_callback_fn is not None: + pre_callback_fn() + + if config.elastic_min_slice_count == -1: + minimum_slice_count = None # Wait for every slice in the jobset. else: - effective_callback = chain_callbacks(cleanup_iterators_and_checkpoints, callback_fn) + minimum_slice_count = config.elastic_min_slice_count + if minimum_slice_count is not None and minimum_slice_count > elastic_manager.total_slice_count: + raise ValueError( + f"elastic_min_slice_count ({minimum_slice_count}) is larger than the number of slices in the jobset" + f" ({elastic_manager.total_slice_count})." + ) + if not is_pause_resume(config) and config.dataset_type == "grain" and not config.grain_use_elastic_iterator: + # The regular grain iterator checkpoints one state file per host, which can't be restored onto a different + # number of hosts. Only `ElasticIterator`, whose state is a single global position, survives a resize. + raise ValueError( + "Replica resize (elastic_min_slice_count smaller than the number of slices in the jobset) restores the data" + " iterator onto a different number of hosts, which only grain's ElasticIterator supports. Set" + " grain_use_elastic_iterator=True." + ) - return elastic_manager.elastic_retry( + pathways_retry = elastic_manager.elastic_retry( max_retries=config.elastic_max_retries, timeout=config.elastic_timeout_seconds, - minimum_slice_count=None if config.elastic_min_slice_count == -1 else config.elastic_min_slice_count, - pre_callback=pre_callback_fn, + minimum_slice_count=minimum_slice_count, + pre_callback=effective_pre_callback, on_elastic_event_callback=effective_callback, ) + def decorator(func): + @functools.wraps(func) + def attempt(*args, **kwargs): + try: + return func(*args, **kwargs) + except manager.ScaleUpSignalError: + # pathwaysutils runs the elastic event callback next; let it know this was a scale-up. + outcome.scale_up = True + raise + + retried_func = pathways_retry(attempt) + + @functools.wraps(func) + def wrapper(*args, **kwargs): + if retry_cache is not None: + # The cache lives for one elastic run, so a later run in the same process starts empty. + retry_cache.reset(enabled=True) + try: + return retried_func(*args, **kwargs) + finally: + if retry_cache is not None: + retry_cache.reset(enabled=False) + # A scale-up that ended the run (retries exhausted) must not label the next run's first event. + outcome.scale_up = False + + return wrapper + + return decorator + def is_scale_up_event(config) -> bool: """Returns whether a scale up event is detected.""" diff --git a/src/maxtext/utils/generate_param_only_checkpoint.py b/src/maxtext/utils/generate_param_only_checkpoint.py index a58ac10b15..aa979786ec 100644 --- a/src/maxtext/utils/generate_param_only_checkpoint.py +++ b/src/maxtext/utils/generate_param_only_checkpoint.py @@ -333,7 +333,7 @@ def _generate_lora_decode_checkpoints_nnx(config, mesh): ) lora_config, lora_state, lora_state_annotations = lora_utils.setup_initial_lora_state( - model, None, tx, config, rng, mesh, checkpoint_manager, lora_adapter_path + None, tx, config, mesh, checkpoint_manager, lora_adapter_path ) _possibly_unroll_lora_params_nnx(config, lora_state, lora_state_annotations, mesh) diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index be6fb2600f..d410256424 100644 --- a/src/maxtext/utils/globals.py +++ b/src/maxtext/utils/globals.py @@ -88,14 +88,15 @@ "qwen3.5-35b-a3b": "Qwen/Qwen3.5-35B-A3B", "qwen3.5-35b-a3b-fp8": "Qwen/Qwen3.5-35B-A3B-FP8", "qwen3.5-35b-fp8": "Qwen/Qwen3.5-35B-A3B-FP8", + "qwen3.5-9b": "Qwen/Qwen3.5-9B", "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistral-7b": "mistralai/Mistral-7B-v0.1", "mixtral-8x22b": "mistralai/Mixtral-8x22B-Instruct-v0.1", "olmo3-7b": "allenai/Olmo-3-7B-Instruct", "olmo3-7b-pt": "allenai/Olmo-3-1025-7B", "olmo3-32b": "allenai/Olmo-3-32B-Think", - "cosmos3-nano-reasoner": "nvidia/Cosmos3-Nano", - "cosmos3-super-reasoner": "nvidia/Cosmos3-Super", + "weaver-mini": "Qwen/Qwen3-VL-4B-Instruct", + "weaver-max": "Qwen/Qwen3-VL-8B-Instruct", # "default" is not HF model, but adding to to avoid confusing warning about tokenizer_path "default": os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers/tokenizer.llama2"), } diff --git a/src/maxtext/utils/gradient_accumulation.py b/src/maxtext/utils/gradient_accumulation.py index 19edc06c99..5386a51488 100644 --- a/src/maxtext/utils/gradient_accumulation.py +++ b/src/maxtext/utils/gradient_accumulation.py @@ -14,6 +14,8 @@ """Functions for gradient accumulation (GA)""" +import functools + import jax import jax.numpy as jnp from jax.experimental.xla_metadata import set_xla_metadata @@ -21,6 +23,7 @@ from flax import nnx from maxtext.common.common_types import ShardMode +from maxtext.layers.moe import load_balance_updates_from_counts from maxtext.utils.sharding import maybe_shard_with_name @@ -102,10 +105,28 @@ def convert_to_bf16(param): ga_params = params ga_params = jax.tree.map(_maybe_shard_with_name, ga_params, ga_params_shardings) + + # Under manual GA, _loss_fn returns the token-summed xent (plus aux terms scaled by the microbatch token count), so + # differentiating it directly seeds the backward pass with 1 where the GA=1 loss (xent_sum / W) seeds it with 1/W. + # Every backward cotangent is then about W times its GA=1 value. A static (fixed-range) fp8 calibration of the + # backward cotangents (e.g. bwd_quantization_calibration_method=fixed,0.01) is sized for the GA=1 magnitude and + # saturates, which distorts the gradient. Scale the differentiated loss by 1/S, with S the static global-batch + # token capacity, so the cotangents carry the GA=1 magnitude, and fold S back into the final division below. + # The accumulated gradient is mathematically unchanged: sum_i grad(L_i) / W. + use_tunix_ga = getattr(config, "use_tunix_gradient_accumulation", False) + token_capacity = getattr(config, "global_batch_size_to_train_on", 0) * getattr(config, "max_target_length", 0) + loss_scale = 1.0 / float(token_capacity) if (token_capacity > 0 and not use_tunix_ga) else 1.0 + + # functools.wraps keeps _loss_fn's signature visible, which nnx.value_and_grad needs to resolve is_train=. + @functools.wraps(_loss_fn) + def _scaled_loss_fn(*args, **kwargs): + loss, aux = _loss_fn(*args, **kwargs) + return loss * loss_scale, aux + if is_nnx: - grad_func = nnx.value_and_grad(_loss_fn, argnums=0, has_aux=True) + grad_func = nnx.value_and_grad(_scaled_loss_fn, argnums=0, has_aux=True) else: - grad_func = jax.value_and_grad(_loss_fn, argnums=4, has_aux=True) + grad_func = jax.value_and_grad(_scaled_loss_fn, argnums=4, has_aux=True) def accumulate_gradient(acc_grad_and_loss, data): ga_params = acc_grad_and_loss["ga_params"] @@ -179,19 +200,46 @@ def reshape_to_microbatch_accumulations(batch_arr): unreduced_shardings = jax.tree.map(update_sharding_for_unreduced, params_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, unreduced_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, params_shardings) - divisor = ( - config.gradient_accumulation_steps if getattr(config, "use_tunix_gradient_accumulation", False) else denominator - ) + divisor = config.gradient_accumulation_steps if use_tunix_ga else denominator * loss_scale raw_grads = jax.tree_util.tree_map( - lambda arr: jnp.where(has_weights, arr / divisor, jnp.zeros_like(arr)), + lambda arr: jnp.where(has_weights, (arr / divisor).astype(arr.dtype), jnp.zeros_like(arr)), raw_grads, ) scanned_aux = aux aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux) # pytype: disable=module-attr + # The auxiliary losses are per-microbatch means; the step's value is their mean over microbatches (as in `loss` + # above), not their sum, so the logged learning/{moe_lb,indexer,mtp}_loss match GA=1. + for key in ("moe_lb_loss", "indexer_loss", "mtp_loss"): + if key in aux: + aux[key] = aux[key] / config.gradient_accumulation_steps + if getattr(config, "routed_bias", False) and getattr(config, "routed_bias_update_rate", 0.0) > 0.0: + # RoutedMoE emits raw per-microbatch int32 expert counts when gradient_accumulation_steps > 1; + # convert the counts summed across microbatches into a single full-batch sign() update. + def finalize_bias(counts): + if isinstance(counts, jax.Array) and jnp.issubdtype(counts.dtype, jnp.integer): + return load_balance_updates_from_counts(counts, config.num_experts, config.routed_bias_update_rate) + return counts + + for key in ("moe_bias_updates", "mtp_moe_bias_updates"): + if aux.get(key) is not None: + aux[key] = jax.tree.map(finalize_bias, aux[key]) + if "intermediate_outputs" in aux: + aux["intermediate_outputs"] = jax.tree_util.tree_map_with_path( + lambda path, val: ( + finalize_bias(val) if path and getattr(path[-1], "key", None) == "moe_bias_updates" else val + ), + aux["intermediate_outputs"], + ) if "te_moe_capacity_overflow" in scanned_aux: aux["te_moe_capacity_overflow"] = jnp.any(scanned_aux["te_moe_capacity_overflow"], axis=0) aux["te_moe_max_total_recv_tokens"] = jnp.max(scanned_aux["te_moe_max_total_recv_tokens"], axis=0) aux["te_moe_recv_capacity_per_rank"] = jnp.min(scanned_aux["te_moe_recv_capacity_per_rank"], axis=0) + if "moe_required_rbf" in scanned_aux: + # log_required_ragged_buffer_factor probe: the step needs the largest factor over microbatches. + aux["moe_required_rbf"] = jnp.max(scanned_aux["moe_required_rbf"], axis=0) + if "moe_max_load_ratio" in scanned_aux: + aux["moe_max_load_ratio"] = jnp.max(scanned_aux["moe_max_load_ratio"], axis=0) + aux["moe_max_load_ratio_mean"] = jnp.mean(scanned_aux["moe_max_load_ratio_mean"], axis=0) if is_nnx: nnx.update(model, grad_and_loss["rest_state"]) diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py index 9111cf2ebd..1d306ae7e3 100644 --- a/src/maxtext/utils/lora_utils.py +++ b/src/maxtext/utils/lora_utils.py @@ -186,20 +186,17 @@ def load_adapter(config, base_abstract_state_params, adapter_config_path, adapte return lora_params, lora_config -def setup_initial_lora_state(model, data_iterator, tx, config, rng, mesh, checkpoint_manager, lora_adapter_path): +def setup_initial_lora_state(data_iterator, tx, config, mesh, checkpoint_manager, lora_adapter_path): """Initialize the LoRA train state and optionally restore it from a checkpoint. - On the NNX path, `model` is unused; the abstract state is built from - `model_creation_utils.create_nnx_abstract_model` and `lora_state.params` - follows the NNX shape. On the Linen path the existing `{"params": ...}` - tree shape is preserved. + The abstract state is built from + `model_creation_utils.create_nnx_abstract_model`, so `lora_state.params` + follows the NNX shape. Args: - model: Linen `nn.Module` used on the Linen path; ignored on NNX. data_iterator: Data iterator passed through to `load_state_if_possible`. tx: Optax gradient transformation for the optimizer. config: Top-level MaxText config. - rng: PRNG key used for the Linen init. mesh: JAX device mesh. checkpoint_manager: Orbax `CheckpointManager` for the adapter. lora_adapter_path: Path to the adapter directory containing diff --git a/src/maxtext/utils/max_utils.py b/src/maxtext/utils/max_utils.py index 48b548ff23..0ffae2d012 100644 --- a/src/maxtext/utils/max_utils.py +++ b/src/maxtext/utils/max_utils.py @@ -31,6 +31,7 @@ from etils import epath import flax +from flax.core.spmd import get_logical_axis_rules import jax from pathlib import Path from contextlib import contextmanager @@ -368,8 +369,8 @@ def initialize_jax_for_gpu(raw_keys): def initialize_jax_for_cpu(raw_keys): """Jax distributed initialize for CPUs. Includes retries until the coordinator is ready.""" coordinator_ip_address = get_coordinator_ip_address() - coordinator_address = coordinator_ip_address + ":1234" # JAX coordinator port used in XPK - # Env variables to be set in XPK or otherwise + coordinator_address = coordinator_ip_address + ":1234" # JAX coordinator port (the port XPK exposes for CPU workloads) + # These env variables must be set by the launcher (XPK sets them for CPU workloads) job_index = int(os.environ.get("JOB_INDEX")) # pyrefly: ignore[bad-argument-type] job_completion_index = int(os.environ.get("JOB_COMPLETION_INDEX")) # pyrefly: ignore[bad-argument-type] processes_in_job = int(os.environ.get("PROCESSES_IN_JOB")) # pyrefly: ignore[bad-argument-type] @@ -1324,7 +1325,19 @@ def maybe_bootstrap_te_moe(config, mesh, shaped_batch): ep_size = mesh.shape.get(ep_axis, 1) fsdp_size = mesh.shape.get(fsdp_axis, 1) - batch_size, sequence_length = shaped_batch["inputs"].shape[:2] + loaded_batch_size, sequence_length = shaped_batch["inputs"].shape[:2] + ga_steps = config.gradient_accumulation_steps + if loaded_batch_size % ga_steps != 0: + raise ValueError(f"TE MoE EP loaded batch size ({loaded_batch_size}) must be divisible by GA steps ({ga_steps}).") + # Bootstrap for one model invocation, not the full accumulated batch. + # Cover both the loaded microbatch and the model-initialization batch. + batch_size = max(loaded_batch_size // ga_steps, config.micro_batch_size_to_train_on) + if config.eval_interval > 0: + batch_size = max(batch_size, config.micro_batch_size_to_eval_on) + if batch_size <= 0 or batch_size % (fsdp_size * ep_size) != 0: + raise ValueError( + f"TE MoE EP per-call batch size ({batch_size}) must be positive and divisible by FSDP * EP ({fsdp_size * ep_size})." + ) if config.num_experts % ep_size != 0: raise ValueError(f"num_experts={config.num_experts} must be divisible by EP size={ep_size}.") @@ -1401,6 +1414,28 @@ def maybe_pad(inputs, tile_size): inputs = jax.lax.pad( inputs, jnp.array(0.0, dtype=inputs.dtype), - [(0, padding_amount, 0), (0, 0, 0)], + [(0, padding_amount, 0)] + [(0, 0, 0)] * (inputs.ndim - 1), ) return inputs, padding_amount + + +def is_eval(config) -> bool: + """Returns True when evaluation-specific kernel tile sizes (`eval_*`) should be used. + + Evaluation tile sizes (`eval_*`) take effect when evaluation runs under a + custom logical mesh/sharding rule (`logical_axis_rules_for_eval != + logical_axis_rules`). Because `eval_step` executes inside + `logical_axis_rules(config.logical_axis_rules_for_eval)`, comparing the + active logical axis rules against `logical_axis_rules_for_eval` reliably + distinguishes evaluation from training without relying on batch size or + tensor sharding. + + Args: + config: MaxText config object. + + Returns: + True if executing under evaluation logical axis rules. + """ + eval_rules = getattr(config, "logical_axis_rules_for_eval", None) + train_rules = getattr(config, "logical_axis_rules", None) + return bool(eval_rules and eval_rules != train_rules and get_logical_axis_rules() == eval_rules) diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 219ee0d9ad..c9ac2fb43d 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -36,6 +36,12 @@ ) from maxtext.configs import pyconfig from maxtext.configs import types + +try: + # lineage_adapter is Google-internal and excluded from the open-source export. + from maxtext.experimental.lineage import lineage_adapter +except ImportError: + lineage_adapter = None from maxtext.multimodal import processor as mm_processor from maxtext.trainers.diloco import diloco from maxtext.trainers.diloco import utils as diloco_utils @@ -99,10 +105,7 @@ def get_functional_train_with_signature( in_shardings = (state_mesh_shardings, data_sharding) # State, batch out_shardings = (state_mesh_shardings, None) # State, metrics static_argnums = () # We partial out the static argnums of model and config - if getattr(config, "retry_when_tokens_dropped", False) is True: - donate_argnums = () # Preserve state so it can be replayed if an overflow occurs - else: - donate_argnums = 0 # This is the index of the state - we allow the compiler to make use of this memory. + donate_argnums = 0 # This is the index of the state - we allow the compiler to make use of this memory. return functional_train, in_shardings, out_shardings, static_argnums, donate_argnums @@ -805,7 +808,7 @@ def get_dense_moe_layers(config): num_moe_layers = config.num_decoder_layers // config.interleave_moe_layer_step num_dense_layers = config.num_decoder_layers - num_moe_layers return num_dense_layers, num_moe_layers - elif config.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, DecoderBlockType.DEEPSEEK4): + elif config.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.DEEPSEEK4): return 0, config.num_decoder_layers elif config.decoder_block == DecoderBlockType.DEFAULT: raise ValueError("Unsupported decoder block for dense/MoE layer calculation") @@ -1372,6 +1375,7 @@ def calculate_tflops_training_per_device(config, log=True): gate_flops = 2 * config.per_device_batch_size * config.max_target_length * config.emb_dim * config.num_experts if config.decoder_block in ( DecoderBlockType.DEEPSEEK, + DecoderBlockType.DEEPSEEK4, DecoderBlockType.LLAMA4, DecoderBlockType.QWEN3_NEXT, DecoderBlockType.GEMMA4, @@ -1517,7 +1521,7 @@ def get_nested_value(dictionary, nested_key, default=None): return current_level -def collect_intermediates_by_suffix(intermediate_outputs, *suffix_keys: str) -> list: +def collect_intermediates_by_suffix(intermediate_outputs, *suffix_keys: str, ravel: bool = True) -> list: """Collects intermediate leaf values whose dict-key path ends with suffix_keys. Works regardless of model architecture (scanned, scannable blocks, or standard), @@ -1527,9 +1531,11 @@ def collect_intermediates_by_suffix(intermediate_outputs, *suffix_keys: str) -> intermediate_outputs: The intermediates dict returned by model.apply(). *suffix_keys: One or more key names forming the expected path suffix, e.g. ``("moe_lb_loss",)`` or ``("self_attention", "indexer_loss")``. + ravel: Flatten each leaf to 1-D. Pass False for leaves that are sharded + across the mesh, where flattening would force a reshard collective. Returns: - A list of 1-D JAX arrays, one per matching leaf (already ravelled). + A list of JAX arrays, one per matching leaf (1-D unless ``ravel=False``). """ suffix = tuple(suffix_keys) n = len(suffix) @@ -1537,7 +1543,7 @@ def collect_intermediates_by_suffix(intermediate_outputs, *suffix_keys: str) -> for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs): path_keys = tuple(k.key for k in path if hasattr(k, "key")) if len(path_keys) >= n and path_keys[-n:] == suffix: - values.append(jnp.ravel(val)) + values.append(jnp.ravel(val) if ravel else val) return values @@ -1751,10 +1757,11 @@ def setup_initial_state( # Initialization with axis_rules(config.logical_axis_rules): + native = lineage_adapter is not None and lineage_adapter.is_native(config) restored, raw_params = checkpointing.load_state_if_possible( checkpoint_manager, data_iterator, - config.load_parameters_path, + "" if native else config.load_parameters_path, config.load_full_state_path, config.checkpoint_storage_concurrent_gb, unboxed_abstract_state, @@ -1768,6 +1775,18 @@ def setup_initial_state( expansion_factor_real_data=config.expansion_factor_real_data, maxtext_config=config, ) + if native and restored is None and raw_params is None and config.load_parameters_path: + # The Lineage-native model restores MaxText-layout params against its own abstract target. + assert lineage_adapter is not None + raw_params = checkpointing.load_params_from_path( + config.load_parameters_path, + lineage_adapter.maxtext_restore_target(config, mesh), + config.checkpoint_storage_concurrent_gb, + use_ocdbt=config.checkpoint_storage_use_ocdbt, + use_zarr3=config.checkpoint_storage_use_zarr3, + enable_single_replica_ckpt_restoring=bool(config.enable_single_replica_ckpt_restoring), + colocated_python_checkpointing=bool(config.enable_single_controller and config.colocated_python_checkpointing), + ) # Partial or fully restored was_restored = bool(restored is not None or raw_params is not None) @@ -1782,7 +1801,11 @@ def setup_initial_state( in_shardings=None, out_shardings=state_mesh_shardings, )() - if raw_params: + if raw_params and native: + # Lineage-native model: convert the MaxText-layout params into the Lineage layout. + assert lineage_adapter is not None + lineage_adapter.restore_params(state["model"], raw_params, config, mesh) + elif raw_params: # Params-only load (base model weights): overlay restored weights, keep init for everything else. target_model = ( state["model"] @@ -2088,7 +2111,17 @@ def create_device_mesh(config, devices=None): allow_split_physical_axes = config.allow_split_physical_axes if config.allow_split_physical_axes else False - if num_slices > 1: + if getattr(config, "use_lineage", False) and lineage_adapter is not None: + dcn_parallelism = ( + max_utils.fill_unspecified_mesh_axes(config.dcn_parallelism.copy(), num_slices, "DCN") if num_slices > 1 else None + ) + mesh = lineage_adapter.create_device_mesh( + ici_parallelism, + devices, + dcn_parallelism=dcn_parallelism, + allow_split_physical_axes=allow_split_physical_axes, + ) + elif num_slices > 1: dcn_parallelism = config.dcn_parallelism.copy() dcn_parallelism = max_utils.fill_unspecified_mesh_axes(dcn_parallelism, num_slices, "DCN") diff --git a/src/maxtext/utils/mllog_utils.py b/src/maxtext/utils/mllog_utils.py index 28bf48029d..80bb3f2165 100644 --- a/src/maxtext/utils/mllog_utils.py +++ b/src/maxtext/utils/mllog_utils.py @@ -14,6 +14,7 @@ """Utils for MLPerf submission compliance.""" import atexit +import math import os import threading import time @@ -198,6 +199,41 @@ def _axis_product(config, *names) -> int: return product +def _rule_parallelism(config, logical_axis) -> int: + """Returns the mesh size (ICI x DCN) that `logical_axis` is sharded over.""" + sizes = { + axis: max(ici, 1) * max(dcn, 1) + for axis, ici, dcn in zip(config.mesh_axes, config.ici_parallelism, config.dcn_parallelism) + } + axes = dict(config.logical_axis_rules).get(logical_axis) or () + if isinstance(axes, str): + axes = (axes,) + return math.prod(sizes.get(axis, 1) for axis in axes) + + +# Bits per element of the MLPerf pre-approved numerical formats, used to pick the lowest of several. +_PRECISION_BITS = { + "fp64": 64, + "fp32": 32, + "tf32": 19, + "fp16": 16, + "bfloat16": 16, + "fp8": 8, + "int8": 8, + "uint8": 8, + "mxfp6": 6, + "nvfp4": 4, + "mxfp4": 4, + "int4": 4, + "uint4": 4, +} + + +def _lowest_precision(*precisions: str) -> str: + """Returns the precision with the fewest bits (the first on ties); unknown formats rank lowest, so they surface.""" + return min(precisions, key=lambda p: _PRECISION_BITS.get(p, 0)) + + def init_print(config): """Logs the static submission and hyperparameter events for compliance checking.""" setup_mllog(config) @@ -235,22 +271,37 @@ def init_print(config): mllogger.event("target_accuracy", config.target_eval_loss) # MLPerf v6.1 mandatory precision, parallelism, micro-batch size, and config filename disclosure. + use_lineage = getattr(config, "use_lineage", False) + tensor_parallelism = _axis_product( + config, + "ici_tensor_parallelism", + "dcn_tensor_parallelism", + "ici_tensor_sequence_parallelism", + "dcn_tensor_sequence_parallelism", + ) + expert_parallelism = _axis_product(config, "ici_expert_parallelism", "dcn_expert_parallelism") + quantization = getattr(config, "quantization", None) + token_all_gather_quantized = quantization and getattr(config, "moe_quantize_token_all_gather", False) + if use_lineage: + # Lineage runs on a physical mesh where every ici_*_parallelism field is 1. + # Its MLA projections are head-sharded across the `activation_length` + # (TensorCore) axis while attention sees the full sequence: TP + SP, not + # CP. fp8_full also quantizes the EP token all-gather. + expert_parallelism = _rule_parallelism(config, "exp") + tensor_parallelism = _rule_parallelism(config, "activation_length") + quantization = quantization or config.lineage_quantization + token_all_gather_quantized = True dtype_str = _mllog_precision(getattr(config, "dtype", "bfloat16")) - linear_prec = _mllog_precision(getattr(config, "quantization", None), fallback=dtype_str) - comm_prec = _mllog_precision(getattr(config, "grad_dtype", None), fallback=dtype_str) + linear_prec = _mllog_precision(quantization, fallback=dtype_str) + comm_precisions = [_mllog_precision(getattr(config, "grad_dtype", None), fallback=dtype_str)] + if token_all_gather_quantized and expert_parallelism > 1: + # The EP token all-gather sends tokens in the GMM activation qtype. + comm_precisions.append(linear_prec) + comm_prec = _lowest_precision(*comm_precisions) mllogger.event(mllog.constants.LOWEST_NUMERICAL_PRECISION_IN_LINEAR, linear_prec) mllogger.event(mllog.constants.LOWEST_NUMERICAL_PRECISION_IN_ATTN, dtype_str) mllogger.event(mllog.constants.LOWEST_NUMERICAL_PRECISION_IN_COMM, comm_prec) - mllogger.event( - mllog.constants.TENSOR_PARALLELISM, - _axis_product( - config, - "ici_tensor_parallelism", - "dcn_tensor_parallelism", - "ici_tensor_sequence_parallelism", - "dcn_tensor_sequence_parallelism", - ), - ) + mllogger.event(mllog.constants.TENSOR_PARALLELISM, tensor_parallelism) mllogger.event( mllog.constants.PIPELINE_PARALLELISM, _axis_product(config, "ici_pipeline_parallelism", "dcn_pipeline_parallelism"), @@ -259,19 +310,13 @@ def init_print(config): mllog.constants.CONTEXT_PARALLELISM, _axis_product(config, "ici_context_parallelism", "dcn_context_parallelism"), ) - mllogger.event( - mllog.constants.EXPERT_PARALLELISM, - _axis_product(config, "ici_expert_parallelism", "dcn_expert_parallelism"), - ) + mllogger.event(mllog.constants.EXPERT_PARALLELISM, expert_parallelism) # TPU v7x exposes 2 JAX devices (TensorCores) per chip, while system descriptions count chips. mllogger.event( mllog.constants.MICRO_BATCH_SIZE, max(1, int(round(getattr(config, "per_device_batch_size", 1) * 2))), ) - mllogger.event( - mllog.constants.CONFIG_FILENAME, - getattr(config, "mllog_config_filename", "") or "config.yml", - ) + mllogger.event(mllog.constants.CONFIG_FILENAME, f"{config.model_name}.yml") def init_stop(): diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index bad16d6086..1b9c8d1c9c 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -42,7 +42,6 @@ from flax import nnx from flax.core.meta import Partitioned from flax.core.spmd import logical_axis_rules -import flax.linen as nn from huggingface_hub import get_token import jax import jax.numpy as jnp @@ -52,6 +51,7 @@ from maxtext.configs import pyconfig from maxtext.integration.tunix.tunix_adapter import TunixMaxTextAdapter from maxtext.integration.vllm.convert_utils import _partition_size +from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.models import models from maxtext.utils import max_logging @@ -59,6 +59,12 @@ import numpy as np from orbax import checkpoint as ocp +try: + # lineage_adapter is Google-internal and excluded from the open-source export. + from maxtext.experimental.lineage import lineage_adapter # pylint: disable=g-import-not-at-top +except ImportError: + lineage_adapter = None + try: from orbax.checkpoint.metadata import ArrayMetadata as _OrbaxArrayMetadata @@ -220,6 +226,19 @@ def _align_checkpoint_to_model_shapes(ckpt_arr, model_arr, logical_axes=None): return jax.device_put(result, model_arr.sharding) +def _intersect_with_metadata(target, meta_tree): + """Returns the sub-tree of `target` whose paths exist in checkpoint `meta_tree`.""" + if not hasattr(target, "items"): + return target + out = {} + for k, v in target.items(): + if hasattr(meta_tree, "items") and k in meta_tree: + sub_tree = _intersect_with_metadata(v, meta_tree[k]) + if not hasattr(sub_tree, "items") or sub_tree: + out[k] = sub_tree + return out + + def _fuse_moe_weights(ckpt_tree, model_arrays_tree, config=None): """Fuse separate wi_0/wi_1 checkpoint entries into a single wi when model uses fused layout. @@ -598,7 +617,7 @@ def from_config( model_mode: str = MODEL_MODE_TRAIN, rngs: None = None, quant_mode_str: str = "train", -) -> nn.Module: +) -> nnx_wrappers.ToLinen: ... @@ -623,9 +642,12 @@ def from_config( model_mode: str = MODEL_MODE_TRAIN, rngs: nnx.Rngs | None = None, quant_mode_str: str = "train", -) -> nn.Module | models.Transformer: +) -> nnx_wrappers.ToLinen | models.Transformer: """Load a pretrained MaxText model from checkpoint. + With `rngs` this returns the NNX `Transformer`. Without it, it returns the + Linen-wrapped model from `models.transformer_as_linen`, which is deprecated. + `quant_mode_str` is one of "train", "convert", "serve" — controls the AQT quantization mode at model construction time. NNX layers freeze their param shape on `quant_mode_str` (e.g. SERVE skips the full-precision kernel), @@ -641,6 +663,8 @@ def from_config( def get_transformer_model(config, mesh, quant, model_mode: str = MODEL_MODE_TRAIN, rngs: nnx.Rngs | None = None): """Returns the transformer model based on the configuration.""" + if rngs is not None and lineage_adapter is not None and lineage_adapter.is_native(config): + return lineage_adapter.LineageTransformer(config, mesh, quant=quant, rngs=rngs, model_mode=model_mode) if rngs is not None: return models.Transformer(config, mesh, quant=quant, rngs=rngs, model_mode=model_mode) else: @@ -651,6 +675,11 @@ def create_model( config, mesh, model_mode: str = MODEL_MODE_TRAIN, rngs: nnx.Rngs | None = None, *, quant_mode_str: str = "train" ): """Instantiates and returns the model object, sharded across the mesh.""" + if config.use_m3_model: + from maxtext.m3 import models as m3_models # pylint: disable=import-outside-toplevel + + return m3_models.create_model(config, mesh, rngs=rngs, quant_mode_str=quant_mode_str) + # Model definition quant = quantizations.configure_quantization(config, quant_mode_str=quant_mode_str) model = get_transformer_model(config, mesh, quant, model_mode=model_mode, rngs=rngs) @@ -1158,15 +1187,23 @@ def _adjust_target_for_moe_fusion(target, meta_tree, is_nnx): target_for_restore, metadata.item_metadata.tree["params"]["params"], False ) + # Linen checkpoints keep non-trainable variables (e.g. `Tid2EidVar`) in + # collections beside "params"; restore the model leaves they hold too. + ckpt_collections = metadata.item_metadata.tree["params"] item_to_restore = {"params": {"params": target_for_restore}} - base_restore_args = ocp.checkpoint_utils.construct_restore_args(target_for_restore) + for col, col_meta in ckpt_collections.items(): + if col != "params": + col_target = _intersect_with_metadata(target_for_restore, col_meta) + if col_target: + item_to_restore["params"][col] = col_target restore_args = { "params": { - "params": _fix_restore_args_for_shape_mismatch( - base_restore_args, - metadata.item_metadata.tree["params"]["params"], + col: _fix_restore_args_for_shape_mismatch( + ocp.checkpoint_utils.construct_restore_args(col_target), + ckpt_collections[col], mesh, ) + for col, col_target in item_to_restore["params"].items() } } else: diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py index ec0b1a6281..a36041d4e3 100644 --- a/src/maxtext/utils/muon_utils.py +++ b/src/maxtext/utils/muon_utils.py @@ -29,7 +29,6 @@ import sys from typing import Optional, Tuple -import flax.linen as nn from flax import nnx import jax from maxtext.configs import pyconfig @@ -219,22 +218,11 @@ def _print_structure_debug(abstract_param, muon_weight_dimension_numbers): """Prints the model structure and the resulting Muon config.""" def get_leaf_info(leaf): - # For linen: - # Access the shape from the inner ShapeDtypeStruct and names from the wrapper - # Return a new tree with the same structure containing only shapes/names - if isinstance(leaf, nn.LogicallyPartitioned): - return {"shape": leaf.value.shape, "names": leaf.names} - # For nnx: - # Only return the shape because it doesn't have a wrapper. - elif isinstance(leaf, jax.ShapeDtypeStruct): + if isinstance(leaf, jax.ShapeDtypeStruct): return {"shape": leaf.shape} return {"shape": "N/A"} - info_tree = jax.tree_util.tree_map( - get_leaf_info, - abstract_param, - is_leaf=lambda x: isinstance(x, (nn.LogicallyPartitioned, jax.ShapeDtypeStruct)), - ) + info_tree = jax.tree_util.tree_map(get_leaf_info, abstract_param, is_leaf=lambda x: isinstance(x, jax.ShapeDtypeStruct)) print(f"\n=== Model Structure ===\n{info_tree}") print(f"\n=== Muon Dimension Numbers ===\n{muon_weight_dimension_numbers}") print("\nIs this reasonable?") diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index 0fbaba1197..ce006d55f3 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -15,11 +15,12 @@ # pylint: disable=line-too-long, disable=bare-except, consider-using-generator """Utils that are only interesting to MaxText and sharding related.""" +import collections from collections.abc import Iterable import inspect # for debugging only from pathlib import Path -from flax import linen as nn, nnx +from flax import nnx from flax.core.spmd import get_logical_axis_rules as flax_get_logical_axis_rules import jax from jax.core import Tracer @@ -55,6 +56,71 @@ def get_logical_axis_rules(): return flax_get_logical_axis_rules() +class _UnassignedAxis: + """Marks a logical dimension that no rule has mapped to a mesh axis yet.""" + + def __repr__(self): + return "UnassignedAxis" + + def __bool__(self): + return False + + +_UNASSIGNED_AXIS = _UnassignedAxis() + + +def _mesh_axes_free(new_assignment, existing_assignments): + """True if none of the mesh axes in `new_assignment` is used in `existing_assignments`.""" + new = set(jax.tree_util.tree_leaves(new_assignment)) + existing = set(jax.tree_util.tree_leaves(existing_assignments)) + new.discard(P.UNCONSTRAINED) + new.discard(None) + return not existing.intersection(new) + + +def logical_to_spec(logical_names, rules=None): + """Maps logical dimension names to a PartitionSpec using the logical axis rules. + + Rules are (logical name, mesh axes) pairs in priority order. A dimension takes + the first matching rule whose mesh axes no other dimension uses yet; + dimensions no rule assigns are left unsharded. Without `rules`, uses the ones + set by `logical_axis_rules`. Same behavior as Flax's + `flax.linen.logical_to_mesh_axes`. + """ + if logical_names is None: + return None + if rules is None: + rules = get_logical_axis_rules() + counts = collections.Counter(logical_names) + # None and special values like PartitionSpec.UNCONSTRAINED may repeat. + dups = tuple(k for k, v in counts.items() if v > 1 and isinstance(k, str)) + if dups: + raise ValueError(f"Unsupported: Dimensions {dups} occur more than once in array names.") + if not isinstance(rules, (tuple, list)): + raise ValueError("Unknown axis rule specification type.") + result = [(_UNASSIGNED_AXIS if isinstance(name, str) else name) for name in logical_names] + for rule_logical_name, rule_mesh_axes in rules: + if rule_logical_name in logical_names: + pos = logical_names.index(rule_logical_name) + if _mesh_axes_free(rule_mesh_axes, result) and result[pos] is _UNASSIGNED_AXIS: + result[pos] = rule_mesh_axes + return P(*(None if x is _UNASSIGNED_AXIS else x for x in result)) + + +def with_logical_constraint(x, logical_axes): + """Constrains `x` to `logical_axes` when a global mesh is set; otherwise returns `x`. + + Port of `flax.linen.with_logical_constraint` for one array and the rules set by + `logical_axis_rules`. Like the Flax call it only applies under a global mesh + (`jax.set_mesh`); MaxEngine enters its mesh with `with mesh:`, so there it does + nothing. Use `maybe_shard_with_logical` for a constraint that always applies. + """ + rules = get_logical_axis_rules() + if not rules or logical_axes is None or jax.sharding.get_abstract_mesh().empty: + return x + return jax.lax.with_sharding_constraint(x, logical_to_spec(logical_axes, rules)) + + def _get_sharding_desc(inputs, extra_stack_level): """Get the inputs sharding description using inspect module""" frame = inspect.currentframe() @@ -368,10 +434,33 @@ def nnx_construct_named_sharding(abs_var_state: nnx.State, mesh) -> nnx.State: ) +class CompoundLogicalAxis(tuple): + """Logical name of one dimension sharded over several logical axes, e.g. CompoundLogicalAxis(("exp", "embed_moe")). + + The dimension is sharded over the concatenation of the mesh axes each name resolves to. Only this type is resolved + this way: a plain tuple entry keeps Flax's meaning (already-physical mesh axes, passed through unchanged). + """ + + +# NNX walks every tuple as a pytree node, so JAX must flatten this subclass too (it treats unregistered ones as leaves). +jax.tree_util.register_pytree_node( + CompoundLogicalAxis, lambda x: (tuple(x), None), lambda _, children: CompoundLogicalAxis(children) +) + + def logical_to_mesh_axes(logical_names, mesh, rules=None): """Remove size one mesh axes given logical names.""" - tensor_spec = nn.logical_to_mesh_axes(logical_names, rules=rules) - return remove_size_one_mesh_axis(tensor_spec, mesh) + if logical_names is None or not any(isinstance(name, CompoundLogicalAxis) for name in logical_names): + return remove_size_one_mesh_axis(logical_to_spec(logical_names, rules=rules), mesh) + groups = [tuple(name) if isinstance(name, CompoundLogicalAxis) else (name,) for name in logical_names] + # Resolve all names in one call so a mesh axis is still used by at most one logical name. + flat_spec = list(logical_to_spec(tuple(n for group in groups for n in group), rules=rules)) + grouped = [] + for group in groups: + resolved, flat_spec = flat_spec[: len(group)], flat_spec[len(group) :] + axes = [ax for item in resolved if item is not None for ax in ((item,) if isinstance(item, str) else item)] + grouped.append(None if not axes else axes[0] if len(axes) == 1 else tuple(axes)) + return remove_size_one_mesh_axis(P(*grouped), mesh) def logical_to_mesh(tree, mesh, rules=None): diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index e208bcbf30..8899d5cbc9 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -16,11 +16,16 @@ """Utils that are only interesting for training in MaxText.""" import subprocess +import threading +import time import jax +import numpy as np import optax import functools import orbax.checkpoint.pathways as ocp_pathways +from jax.experimental import multihost_utils + from flax import nnx from flax.core.spmd import logical_axis_rules @@ -30,9 +35,11 @@ from maxtext.common.common_types import ReorderStrategy from maxtext.common.data_loader import create_dataloader from maxtext.common.goodput import GoodputEvent, maybe_record_goodput +from maxtext.input_pipeline import multihost_dataloading from maxtext.optimizers import optimizers from maxtext.trainers.diloco import diloco from maxtext.utils import diloco_sharding +from maxtext.utils import elastic_utils from maxtext.utils import lora_utils from maxtext.utils import max_logging from maxtext.utils import max_utils @@ -42,6 +49,114 @@ from maxtext.utils.rampup_batch import create_rampup_manager +def prepare_before_run_start( + config, + state, + learning_rate_schedule, + start_step, + mesh, + data_loader, + shaped_batch, + shaped_eval_batch=None, +): + """Optional setup that runs before init_stop/run_start and never touches the dataset. + + MLPerf starts the clock before any part of the system touches the dataset, so everything here uses only the + model state, the config and synthetic (all-zero) inputs. The flags are independent. + """ + timings = [] + + if config.warm_input_reshard_before_run_start: + # The loader builds each host batch as a global array sharded over all mesh axes + # (multihost_dataloading._form_global_array) and then device_puts it to the input sharding, which compiles a + # reshard program on first use. Warm that path with an all-zero synthetic batch (no training data). + t = time.perf_counter() + local_rows = config.global_batch_size_to_load // jax.process_count() + + def _synthetic_global(leaf): + local = np.zeros((local_rows,) + tuple(leaf.shape[1:]), dtype=leaf.dtype) + return multihost_dataloading._form_global_array((), local, mesh) # pylint: disable=protected-access + + warm = jax.device_put(jax.tree.map(_synthetic_global, shaped_batch), data_loader.input_data_shardings) + jax.block_until_ready(warm) + del warm + timings.append(f"input reshard warm {time.perf_counter() - t:.3f} s") + + if config.warm_eval_input_reshard_before_run_start and shaped_eval_batch is not None: + # The eval loop forms each host batch as a global array over all mesh axes + # (MultiHostDataLoadIterator._get_next_batch_sharded -> _form_global_array) and then device_puts it to the eval + # input sharding (train.py eval loop). Warm that path with an all-zero synthetic eval batch (no eval data). + t = time.perf_counter() + local_rows_eval = config.global_batch_size_to_load_eval // jax.process_count() + + def _synthetic_global_eval(leaf): + local = np.zeros((local_rows_eval,) + tuple(leaf.shape[1:]), dtype=leaf.dtype) + return multihost_dataloading._form_global_array((), local, mesh) # pylint: disable=protected-access + + eval_input_sharding = sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules_for_eval) + warm_eval = jax.device_put(jax.tree.map(_synthetic_global_eval, shaped_eval_batch), eval_input_sharding) + jax.block_until_ready(warm_eval) + del warm_eval + timings.append(f"eval input reshard warm {time.perf_counter() - t:.3f} s") + + if config.block_state_before_run_start: + # Wait for the restored/initialized train state so it is not paid for inside step 0. + t = time.perf_counter() + jax.block_until_ready(state) + timings.append(f"state ready {time.perf_counter() - t:.3f} s") + # The metric logger evaluates learning_rate_schedule(step) eagerly after every step; the first call compiles a + # handful of tiny eager ops. Do it once here so step 0 does not pay for it. + t = time.perf_counter() + jax.block_until_ready(learning_rate_schedule(start_step)) + timings.append(f"lr schedule warm {time.perf_counter() - t:.3f} s") + # Make all hosts log run_start together instead of process 0 starting the clock while others are still in init. + t = time.perf_counter() + multihost_utils.sync_global_devices("pre_run_start") + timings.append(f"barrier {time.perf_counter() - t:.3f} s") + + if timings: + max_logging.log("Pre-run_start setup: " + ", ".join(timings)) + + +def start_eval_cache_prefill(eval_data_iterator): + """eval_cache_prefill_in_background: iterate the eval tf.data pipeline once on a daemon thread. + + The c4_mlperf eval dataset ends in `.cache()`; tf.data's in-memory cache belongs to the dataset object and is + shared by every iterator created from it, but it is only committed by an iterator that reads the dataset to the + end. Doing that full pass here (right after run_start, overlapped with the first training steps) means the + first eval's `reset()` iterator reads from memory instead of paying the per-host eval file reads serially. + Returns the thread (joined before the first eval) or None when the loader is not a tf.data pipeline. + """ + # Unwrap view wrappers (e.g. _ReorderedDataIterator.data_iterator) down to the MultiHostDataLoadIterator, whose + # `.dataloader` is the tf.data Dataset. Hosts that load no real eval data hold a PlaceHolderDataIterator: skip. + inner = eval_data_iterator + for _ in range(8): + if hasattr(inner, "dataloader") or not hasattr(inner, "data_iterator"): + break + inner = inner.data_iterator + ds = getattr(inner, "dataloader", None) + if ds is None or not hasattr(ds, "as_numpy_iterator"): + max_logging.log( + f"eval_cache_prefill_in_background: eval loader {type(eval_data_iterator).__name__} -> {type(inner).__name__} " + "is not a tf.data pipeline; skipped" + ) + return None + + def _run(): + t = time.perf_counter() + n = 0 + try: + for _ in ds.as_numpy_iterator(): + n += 1 + max_logging.log(f"eval cache prefill: {n} batches in {time.perf_counter() - t:.3f} s") + except Exception as e: # pylint: disable=broad-except + max_logging.log(f"eval cache prefill failed ({e!r}); the first eval will read the eval files itself") + + thread = threading.Thread(target=_run, name="eval_cache_prefill", daemon=True) + thread.start() + return thread + + def create_training_optimizer(config, model, mesh=None): """Creates the optimizer and learning rate schedule.""" learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) @@ -53,9 +168,9 @@ def create_training_optimizer(config, model, mesh=None): def create_checkpoint_manager(config, mesh, init_state_fn): """Creates the init_rng, optimizer, learning rate schedule, and checkpoint manager.""" # pass in model for muon - # `setup_checkpoint_logger` only emits a deprecation warning now (Orbax v1 logs - # internally) and always returns None; we still pass it through for API parity. - logger = checkpointing.setup_checkpoint_logger(config) # pylint: disable=assignment-from-no-return + # Orbax checkpoint Cloud Logger (None unless `enable_checkpoint_cloud_logger`), + # used by ML Goodput to measure checkpoint save/restore badput. + logger = checkpointing.setup_checkpoint_logger(config) if config.enable_multi_tier_checkpointing: checkpoint_manager = emergency_checkpointing.create_replicator_checkpoint_manager( config.local_checkpoint_directory, @@ -100,6 +215,7 @@ def create_checkpoint_manager(config, mesh, init_state_fn): config.checkpoint_todelete_subdir, config.checkpoint_todelete_full_path, config.checkpoint_storage_target_data_file_size_bytes, + colocated_python_checkpointing=config.enable_single_controller and config.colocated_python_checkpointing, ) # Use Colocated Python checkpointing dispatchers optimization (Single Controller only). @@ -196,7 +312,13 @@ def jit_train_and_eval_step( """Returns a JIT-compiled train and eval step function.""" if config.enable_diloco: train_step_partial = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings) - train_step = diloco.build_diloco_train_step(config, train_step_partial, mesh=mesh) + if not isinstance(state_mesh_shardings, diloco.DiLoCoTrainState): + raise TypeError( + f"enable_diloco requires state_mesh_shardings to be a DiLoCoTrainState, got {type(state_mesh_shardings)}." + ) + train_step = diloco.build_diloco_train_step( + config, train_step_partial, mesh=mesh, outer_params_shardings=state_mesh_shardings.params + ) data_sharding_for_train = sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules) data_sharding_for_eval = sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules_for_eval) p_train_step = jit_train_step( @@ -233,13 +355,18 @@ def _reorder_data_iterator_for_loader(reorder_fn, data_iterator): return _ReorderedDataIterator(reorder_fn, data_iterator) -def setup_train_loop(config, recorder, devices=None): +def setup_train_loop(config, recorder, devices=None, elastic_retry_cache=None): """Set up prerequisites for the training loop - checkpoint_manager, PRNG keys, Mesh, Model and optimizer. Set up data iterator and tokenizer, initialize the model. - Args: config recorder + Args: + config: + recorder: + devices: + elastic_retry_cache: Optional `elastic_utils.RetryCache`. An elastic retry on the same slices reuses the mesh and + abstract model cached in it instead of building them again. Returns: init_rng: @@ -258,10 +385,19 @@ def setup_train_loop(config, recorder, devices=None): with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT): init_rng = jax.random.PRNGKey(config.init_weights_seed) - mesh = maxtext_utils.get_mesh_from_config(config, devices) + # Build the mesh, or take the one from the previous elastic attempt if it ran on the same slices. + # `elastic_retry_cache` is None outside elastic training, in which case the mesh is simply built. + mesh = elastic_utils.get_or_build( + elastic_retry_cache, "mesh", functools.partial(maxtext_utils.get_mesh_from_config, config, devices) + ) context_parallel_size = mesh.shape.get(config.context_sharding, 1) - # Create abstract NNX model. - _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices) + # Create abstract NNX model. Like the mesh, it is reused across elastic attempts on the same slices: it holds + # only shapes and the graphdef, no device arrays, so it survives the cleanup after an elastic event. + _create_model_partial, model = elastic_utils.get_or_build( + elastic_retry_cache, + "abstract_model", + functools.partial(model_creation_utils.create_nnx_abstract_model, config, mesh, devices), + ) learning_rate_schedule, tx = create_training_optimizer(config, model, mesh=mesh) # The train state is wrapped in the TrainStateNNX module. diff --git a/src/maxtext/utils/vocabulary_tiling.py b/src/maxtext/utils/vocabulary_tiling.py index 91fed35bee..97ebffc446 100644 --- a/src/maxtext/utils/vocabulary_tiling.py +++ b/src/maxtext/utils/vocabulary_tiling.py @@ -112,7 +112,6 @@ def vocab_tiling_nnx_loss(model, hidden_states, data, config, is_train): maybe_shard_with_name, shard_mode=config.shard_mode, debug_sharding=config.debug_sharding, - extra_stack_level=1, ) def _reshape(inputs, out_shape, out_sharding): diff --git a/src/maxtext/version.py b/src/maxtext/version.py index e2d3cad75e..5515747a03 100644 --- a/src/maxtext/version.py +++ b/src/maxtext/version.py @@ -15,7 +15,7 @@ """Version and package metadata for MaxText.""" __author__ = "Google LLC" -__version__ = "0.2.4" +__version__ = "0.2.5" __description__ = ( "MaxText is a high performance, highly scalable, open-source LLM written in" " pure Python/Jax and targeting Google Cloud TPUs and GPUs for training and" diff --git a/tests/post_training/unit/maxtext_engine_frozen_grad_test.py b/tests/post_training/unit/maxtext_engine_frozen_grad_test.py new file mode 100644 index 0000000000..5ede9104bb --- /dev/null +++ b/tests/post_training/unit/maxtext_engine_frozen_grad_test.py @@ -0,0 +1,402 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://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. + +"""`trainable_parameters_mask`: `fwd_bwd` computes no gradient for a frozen parameter, and no number changes. + +`update()` zeroes the gradient of every parameter the mask freezes before the norm, clipping and the +optimizer see it. `fwd_bwd` also stops that gradient inside the function it differentiates +(`_stop_frozen_gradients`), so it is never computed. Each test compares the engine with the same engine +whose `_stop_frozen_gradients` is the identity: the program that computes the frozen gradient and zeroes +it in `update()` alone. Both rigs clip every update, so a change to any gradient or to the norm would +reach the trained weights. + +On a tiny real decoder with one MLP kernel frozen, trained by AdamW: +- the lowered `fwd_bwd` has one `dot_general` fewer, shaped like the frozen kernel, and none added; +- `fwd_bwd` returns exactly zero for the frozen kernel, where the reference computes a non-zero gradient; +- every other gradient, the running sums, each update's gradient norm and the trained weights are + bitwise equal: under fsdp, under the deferred data-parallel all-reduce, and on the eager path, which + does not cache the mask; +- control: a stop that also hits one trainable kernel is caught, and the comparison names that kernel. +These run in a subprocess with four CPU devices, for the reason given in +`maxtext_engine_data_parallel_test.py`: a sharded mesh needs more than one device, and the device count is +fixed before this file is imported. + +On the one-layer Qwen3.5 MoE rig of the dropless tests, with the MLPerf RL recipes' own mask, which freezes +the router: the same comparison through Tunix's router replay payload and loss. It runs on one device, in +the pytest process, as that rig's own tests do. +""" + +import collections +import contextlib +import functools +import os +import re +import subprocess +import sys +import tempfile +from typing import Any +import unittest +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized +from flax import nnx +import jax +from maxtext.configs import pyconfig +from maxtext.models import models +from maxtext.training_engine import maxtext_engine +from maxtext.utils import maxtext_utils +import numpy as np +import pytest + +from tests.post_training.unit.maxtext_engine_dropless_replay_test import _cfg as _moe_rig_cfg +from tests.post_training.unit.maxtext_engine_dropless_replay_test import _payload as _moe_rig_payload +from tests.utils.test_helpers import get_test_config_path + +# training_engine imports tunix, so these tests need the post-training dependency bundle. +pytestmark = [pytest.mark.post_training, pytest.mark.cpu_only] + +_REQUIRED_DEVICES = 4 +_SENTINEL = "MAXTEXT_ENGINE_FROZEN_GRAD_TESTS_PASSED" +_RAN = re.compile(rf"{_SENTINEL} ran=(\d+)") + +# The kernel the mask freezes, and the trainable kernel the control stops by mistake. Paths are written +# the way `trainable_parameters_mask` matches them (`optimizers.get_path_mask_fn`). +_FROZEN = "layers_0/mlp/wi_0/kernel" +_WRONGLY_STOPPED = "layers_1/mlp/wo/kernel" +# The mask lists what trains, so freezing one kernel takes a negative lookahead. +_MASK = f"^(?!.*{_FROZEN}).*" +# The MLPerf RL recipes' mask, verbatim (`TRAINABLE_PARAMETERS_MASK` in Tunix's `mlperf_base.sh`). +_ROUTER = "routed_experts/gate/kernel" +_ROUTER_MASK = "^(?!.*routed_experts/gate/kernel).*" +# Below every update's gradient norm (checked), so clipping scales each update by the norm and a change +# to the norm would change the trained weights. +_CLIP = 0.5 + +# fsdp is the kind of mesh the RL recipes train on. The deferred all-reduce returns the gradients +# `unreduced`, so a stopped leaf's zero has to come out with that tag too. +_FSDP = {"ici_fsdp_parallelism": _REQUIRED_DEVICES, "ici_data_parallelism": 1} +_DEFERRED = {"ici_fsdp_parallelism": 1, "ici_data_parallelism": _REQUIRED_DEVICES, "shard_mode": "explicit"} + +# Every engine's `base_output_directory`, removed when the process exits. +_OUTPUT_DIR = tempfile.TemporaryDirectory(prefix="engine_frozen_grad_") # pylint: disable=consider-using-with + + +@pytest.mark.post_training +@pytest.mark.cpu_only +def test_frozen_gradients_on_a_four_device_cpu_mesh(): + """Runs `FrozenGradientTest` in a child process with four CPU devices.""" + env = os.environ.copy() + env["XLA_FLAGS"] = f"{env.get('XLA_FLAGS', '')} --xla_force_host_platform_device_count={_REQUIRED_DEVICES}".strip() + env["JAX_PLATFORMS"] = "cpu" + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) + env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root + + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + ran = _RAN.search(result.stdout) + assert ran, f"the child did not report a completed run\n{report}" + assert int(ran.group(1)) > 0, f"every test in the child skipped\n{report}" + + +def _config(**overrides) -> pyconfig.HyperParameters: + """A tiny real decoder, float32 throughout, trained by AdamW with clipping, with `_FROZEN` frozen.""" + fields = { + "model_name": "default", + "run_name": "engine_frozen_grad_test", + "base_output_directory": _OUTPUT_DIR.name, + "enable_checkpointing": False, + "convert_checkpoint_if_possible": False, + "skip_jax_distributed_system": True, + "enable_tensorboard": False, + "record_internal_nn_metrics": False, + "enable_dropout": False, + "init_weights_seed": 0, + "dtype": "float32", + "weight_dtype": "float32", + "grad_dtype": "float32", + "remat_policy": "none", + "scan_layers": False, + "attention": "dot_product", + "ici_tensor_parallelism": 1, + "per_device_batch_size": 1, + "vocab_size": 128, + "base_emb_dim": 64, + "base_mlp_dim": 128, + "base_num_decoder_layers": 2, + "base_num_query_heads": 4, + "base_num_kv_heads": 4, + "head_dim": 16, + "max_target_length": 32, + "opt_type": "adamw", + "learning_rate": 0.01, + "gradient_clipping_threshold": _CLIP, + "warmup_steps_fraction": 0.0, + "learning_rate_final_fraction": 1.0, + "profiler_steps": 0, + "trainable_parameters_mask": _MASK, + **_FSDP, + } + fields.update(overrides) + argv = ["maxtext_engine_frozen_grad_test.py", get_test_config_path("base.yml")] + return pyconfig.initialize(argv + [f"{k}={v}" for k, v in fields.items()]) + + +def _batch(cfg: pyconfig.HyperParameters, seed: int) -> dict[str, np.ndarray]: + """A `train.py` loss batch as host arrays, which the engine places itself.""" + rows, seq = int(cfg.micro_batch_size_to_train_on), cfg.max_target_length + tokens = np.random.default_rng(seed).integers(1, cfg.vocab_size, size=(rows, seq)).astype(np.int32) + ones = np.ones((rows, seq), dtype=np.int32) + return { + "inputs": tokens, + "targets": np.roll(tokens, -1, axis=-1), + "inputs_position": np.tile(np.arange(seq, dtype=np.int32), (rows, 1)), + "inputs_segmentation": ones, + "targets_segmentation": ones, + } + + +def _path(path) -> str: + """A leaf's path the way `trainable_parameters_mask` matches it, e.g. `decoder/layers_0/mlp/wi_0/kernel/value`.""" + return jax.tree_util.keystr(path, simple=True, separator="/") + + +def _bits(tree) -> dict[str, np.ndarray]: + """Every leaf's raw bits by path, so comparing two trees is bitwise.""" + out = {} + for path, leaf in jax.tree_util.tree_leaves_with_path(tree): + value = np.asarray(jax.device_get(leaf)) + out[_path(path)] = value.view(np.uint16 if value.dtype.itemsize == 2 else np.uint32) + return out + + +def _no_stop(engine, params): + """The reference `_stop_frozen_gradients`: stops nothing, so `fwd_bwd` computes the frozen gradient too.""" + del engine + return params + + +_REAL_STOP = maxtext_engine.MaxTextTrainingEngine._stop_frozen_gradients # pylint: disable=protected-access + + +def _also_stopping(pattern: str): + """A deliberately wrong `_stop_frozen_gradients`: the real one, plus a stop on every leaf whose path has `pattern`.""" + + def stop(engine, params): + return jax.tree_util.tree_map_with_path( + lambda path, x: jax.lax.stop_gradient(x) if pattern in _path(path) else x, _REAL_STOP(engine, params) + ) + + return stop + + +def _with_stop(stop): + """Replaces `MaxTextTrainingEngine._stop_frozen_gradients` by `stop` while active; None leaves the engine as it is.""" + if stop is None: + return contextlib.nullcontext() + return mock.patch.object(maxtext_engine.MaxTextTrainingEngine, "_stop_frozen_gradients", stop) + + +def _run(engine: maxtext_engine.MaxTextTrainingEngine, steps: list[list[Any]]) -> dict[str, Any]: + """Trains one update per entry of `steps`, over that entry's micro-batches, and closes the engine. + + Returns, as raw bits by leaf path: `first`, the first micro-batch's gradients as `fwd_bwd` returned + them; `sums`, the running sum before each update; `norms`, each update's gradient norm; `weights`, the + trained parameters; `initial`, the parameters before training. Also `frozen`, the paths the mask froze; + `clip`, the clipping threshold; `deferred`, whether the all-reduce was deferred; and `mask_cached`, + whether the engine had cached the freeze mask. + """ + params = nnx.state(engine.state.model, nnx.Param) + initial = _bits(params) + mask = engine._freeze_mask_for(params) # pylint: disable=protected-access + frozen = sorted(_path(path) for path, is_frozen in jax.tree_util.tree_leaves_with_path(mask) if is_frozen) + first, sums, norms = None, [], [] + for micro_batches in steps: + for micro_batch in micro_batches: + engine.fwd_bwd(micro_batch) + if first is None: + first = _bits(engine._reduced_accumulated_grads()) # pylint: disable=protected-access + sums.append(_bits(engine._reduced_accumulated_grads())) # pylint: disable=protected-access + engine.update() + norms.append(np.asarray(engine.get_metrics(clear_cache=True).scalar_metrics["gradient_norm"]).reshape(-1)) + run = { + "first": first, + "sums": sums, + "norms": norms, + "weights": _bits(nnx.state(engine.state.model, nnx.Param)), + "initial": initial, + "frozen": frozen, + "clip": engine._config.gradient_clipping_threshold, # pylint: disable=protected-access + "deferred": engine._unreduced_grad_shardings is not None, # pylint: disable=protected-access + "mask_cached": engine._freeze_mask is not None, # pylint: disable=protected-access + } + engine.close() + return run + + +def _train(stop=None, *, eager: bool = False, steps: int = 2, micro_batches: int = 2, **overrides) -> dict[str, Any]: + """`_run` on the tiny decoder, with `_stop_frozen_gradients` replaced by `stop` for the whole run, tracing included.""" + with _with_stop(stop): + cfg = _config(**overrides) + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=maxtext_utils.get_mesh_from_config(cfg)) + if not eager: + engine.compile(_batch(cfg, 0)) + seeds = np.arange(steps * micro_batches).reshape(steps, micro_batches) + 1 + return _run(engine, [[_batch(cfg, int(seed)) for seed in row] for row in seeds]) + + +@functools.cache +def _reference(eager: bool = False, steps: int = 2, mesh: str = "fsdp") -> dict[str, Any]: + """`_train` with nothing stopped, the program before `_stop_frozen_gradients`; cached, as several tests share it.""" + return _train(_no_stop, eager=eager, steps=steps, **(_DEFERRED if mesh == "deferred" else _FSDP)) + + +def _fwd_bwd_dots(stop=None) -> collections.Counter[str]: + """The `dot_general` ops in `fwd_bwd` as lowered, before XLA could remove any, counted by result type.""" + with _with_stop(stop): + cfg = _config() + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=maxtext_utils.get_mesh_from_config(cfg)) + text = engine._lower_kernels(_batch(cfg, 0))["fwd_bwd"].as_text() # pylint: disable=protected-access + engine.close() + dots = [line for line in text.splitlines() if "stablehlo.dot_general" in line] + result_types = collections.Counter(re.search(r"->\s*(tensor<[^>]*>)", line).group(1) for line in dots) + assert sum(result_types.values()) == len(dots), "a dot_general line without a result type" + return result_types + + +def _train_moe_rig(stop=None) -> dict[str, Any]: + """`_run` on the dropless tests' one-layer Qwen3.5 MoE rig, with the recipes' router mask. + + Through Tunix's router replay payload and loss, and `compile(None)` as Tunix's `TrainerWorker` calls + it, so the kernels compile on the first `fwd_bwd`. Two updates of two micro-batches. + """ + with _with_stop(stop): + cfg = _moe_rig_cfg( + trainable_parameters_mask=_ROUTER_MASK, + gradient_clipping_threshold=_CLIP, + base_output_directory=_OUTPUT_DIR.name, + ) + mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + model = models.Transformer(config=cfg, mesh=mesh, quant=None, model_mode="train", rngs=nnx.Rngs(0)) + with mock.patch.object(maxtext_engine.model_creation_utils, "from_pretrained", return_value=model): + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + engine.with_gen_model_input_fn(maxtext_engine.router_replay_gen_model_input_fn) + engine.with_loss_fn(maxtext_engine.make_router_replay_loss_fn(cfg), has_aux=True) + engine.compile(None) + return _run(engine, [[_moe_rig_payload(10), _moe_rig_payload(20)], [_moe_rig_payload(30), _moe_rig_payload(40)]]) + + +class _FrozenGradientChecks: + """The comparison every test here makes, for a `unittest.TestCase`.""" + + def _differing(self, got: dict[str, np.ndarray], want: dict[str, np.ndarray]) -> list[str]: + """The paths whose bits differ between two `_bits` dicts over the same leaves.""" + self.assertEqual(set(got), set(want)) + return sorted(path for path in want if not np.array_equal(got[path], want[path])) + + def _assert_frozen_gradient_skipped_and_nothing_else_changed( + self, run: dict[str, Any], reference: dict[str, Any], frozen_path: str + ): + """`run` returned zero for the gradient of the leaf at `frozen_path`, and matches `reference` bitwise elsewhere.""" + # The mask froze exactly that one leaf, the same in both runs. + self.assertLen(run["frozen"], 1) + (frozen,) = run["frozen"] + self.assertIn(frozen_path, frozen) + self.assertEqual(reference["frozen"], run["frozen"]) + # Not vacuous: the reference computes a non-zero gradient for the frozen leaf ... + self.assertTrue(np.any(reference["first"][frozen]), "the frozen leaf has no gradient, so this checks nothing") + # ... where the engine returns exactly zero (every bit 0, i.e. +0.0). + self.assertFalse(np.any(run["first"][frozen]), "fwd_bwd computed a gradient for the frozen leaf") + # Every other gradient is bitwise unchanged, in the first micro-batch and in the running sums. + self.assertEqual(self._differing(run["first"], reference["first"]), [frozen], "first micro-batch's gradients") + for step, (got, want) in enumerate(zip(run["sums"], reference["sums"], strict=True)): + self.assertEqual(self._differing(got, want), [frozen], f"running sum before update {step}") + # So is everything the update computes from them: the norm, which clips every update, ... + for step, (got, want) in enumerate(zip(run["norms"], reference["norms"], strict=True)): + self.assertGreater(float(want[0]), run["clip"], f"update {step} was not clipped, so its norm reached no weight") + np.testing.assert_array_equal(got.view(np.uint32), want.view(np.uint32), err_msg=f"gradient norm, update {step}") + # ... and every trained weight. The frozen leaf did not move. + self.assertEqual(self._differing(run["weights"], reference["weights"]), [], "trained weights") + np.testing.assert_array_equal(run["weights"][frozen], run["initial"][frozen], err_msg="the frozen leaf moved") + + +class FrozenGradientTest(_FrozenGradientChecks, parameterized.TestCase): + """The tiny decoder with one kernel frozen, against the same engine computing its gradient.""" + + __test__ = False # collected only via the subprocess entry point at the top of this file. + + def test_fwd_bwd_drops_the_frozen_kernels_weight_gradient(self): + """One `dot_general` fewer, shaped like the frozen kernel, and none added: its weight gradient is never computed. + + Its gradient through to the activations is still computed, since layer 0's input depends on trainable + parameters, so of the kernel's three matmuls (forward, input gradient, weight gradient) one goes. + """ + reference, engine = _fwd_bwd_dots(_no_stop), _fwd_bwd_dots() + removed, added = reference - engine, engine - reference + self.assertEqual(dict(added), {}, f"dot_general ops the stop added, by result type; reference {dict(reference)}") + self.assertEqual(sum(removed.values()), 1, f"dot_general ops the stop removed, by result type: {dict(removed)}") + # The [emb, mlp] kernel's weight gradient, in either layout (JAX's transpose rule may emit it [mlp, emb]). + cfg = _config() + kernel_shaped = {f"tensor<{cfg.emb_dim}x{cfg.mlp_dim}xf32>", f"tensor<{cfg.mlp_dim}x{cfg.emb_dim}xf32>"} + self.assertIn(next(iter(removed)), kernel_shaped, "the dot_general the stop removed is not shaped like the kernel") + + @parameterized.named_parameters(("fsdp", "fsdp"), ("deferred_all_reduce", "deferred")) + def test_frozen_gradient_is_skipped_and_nothing_else_changes(self, mesh): + run = _train(**(_DEFERRED if mesh == "deferred" else _FSDP)) + self.assertTrue(run["mask_cached"], "the compiled path did not cache the mask") + self.assertEqual(run["deferred"], mesh == "deferred", "the deferral is not what this case is about") + self._assert_frozen_gradient_skipped_and_nothing_else_changed(run, _reference(mesh=mesh), _FROZEN) + + def test_eager_path_computes_the_mask_and_stops_the_same_gradient(self): + """Without `compile()` the mask is never cached; `_freeze_mask_for` computes it from the parameters.""" + run = _train(eager=True, steps=1) + self.assertFalse(run["mask_cached"], "the mask was cached, so the eager fallback did not run") + self._assert_frozen_gradient_skipped_and_nothing_else_changed(run, _reference(eager=True, steps=1), _FROZEN) + + def test_control_a_stop_on_a_trainable_kernel_is_caught(self): + """Control: the same comparison, against a stop that also hits one trainable kernel, fails and names it.""" + reference = _reference(mesh="fsdp") + wrong = _train(_also_stopping(_WRONGLY_STOPPED)) + (frozen,) = wrong["frozen"] + (wrongly_stopped,) = [path for path in reference["first"] if _WRONGLY_STOPPED in path] + # The wrongly stopped kernel's gradient is zero, and it is the only other one that changed ... + self.assertFalse(np.any(wrong["first"][wrongly_stopped])) + self.assertEqual(self._differing(wrong["first"], reference["first"]), sorted([frozen, wrongly_stopped])) + # ... so the norm differs, and with it every update. + self.assertFalse(np.array_equal(wrong["norms"][0], reference["norms"][0])) + with self.assertRaisesRegex(AssertionError, re.escape(wrongly_stopped)): + self._assert_frozen_gradient_skipped_and_nothing_else_changed(wrong, reference, _FROZEN) + + +class FrozenRouterTest(_FrozenGradientChecks, absltest.TestCase): + """The MLPerf RL recipes' router freeze on a one-layer Qwen3.5 MoE, against the same engine computing it.""" + + def test_the_router_gradient_is_skipped_and_nothing_else_changes(self): + run, reference = _train_moe_rig(), _train_moe_rig(_no_stop) + self.assertTrue(run["mask_cached"], "the compiled path did not cache the mask") + self._assert_frozen_gradient_skipped_and_nothing_else_changed(run, reference, _ROUTER) + + +if __name__ == "__main__": + if jax.device_count() < _REQUIRED_DEVICES: + raise SystemExit( + f"needs {_REQUIRED_DEVICES} devices, got {jax.device_count()}; run this through pytest, which sets " + f"XLA_FLAGS=--xla_force_host_platform_device_count={_REQUIRED_DEVICES}" + ) + _result = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(FrozenGradientTest)) + if not _result.wasSuccessful(): + sys.exit(1) + print(f"{_SENTINEL} ran={_result.testsRun - len(_result.skipped)}") diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py index 674b7dc280..e021c5dbf1 100644 --- a/tests/post_training/unit/maxtext_engine_test.py +++ b/tests/post_training/unit/maxtext_engine_test.py @@ -1774,6 +1774,28 @@ def metadata(self): prepared_prop = t_no_gen._prepare_batch(prop_payload) self.assertIn("tokens", prepared_prop) + def test_prepare_batch_rejects_packed_payload_for_unpacked_gdn_kernel(self): + """With packing=False the GDN kernel ignores segment IDs in training, so a packed payload must be refused.""" + + def payload(segment_ids=None): + return datatypes.RLTrainerPayload( + prompt_ids=jnp.zeros((1, 0), jnp.int32), + prompt_mask=jnp.zeros((1, 0), jnp.int32), + completion_ids=jnp.zeros((1, 4), jnp.int32), + completion_mask=jnp.ones((1, 4), jnp.int32), + advantages=jnp.zeros((1,)), + segment_ids=segment_ids, + ) + + packed = payload(jnp.array([[1, 1, 2, 2]], jnp.int32)) + unpacked_gdn = maxtext_engine.MaxTextTrainingEngine(self.setup_config(use_gdn_kernel=True, packing=False)) + with self.assertRaisesRegex(ValueError, "Set packing=True"): + unpacked_gdn._prepare_batch(packed) + self.assertIn("completion_ids", unpacked_gdn._prepare_batch(payload())) + for overrides in ({"use_gdn_kernel": True, "packing": True}, {"use_gdn_kernel": False, "packing": False}): + engine = maxtext_engine.MaxTextTrainingEngine(self.setup_config(**overrides)) + self.assertIn("segment_ids", engine._prepare_batch(packed)) + def test_prepare_batch_prevents_recompilation_on_metadata_change(self): """Payload metadata changes must not alter the Treedef or trigger recompilation.""" t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) diff --git a/tests/post_training/unit/train_distill_test.py b/tests/post_training/unit/train_distill_test.py index 845014db32..ca12de10b8 100644 --- a/tests/post_training/unit/train_distill_test.py +++ b/tests/post_training/unit/train_distill_test.py @@ -1134,17 +1134,19 @@ def test_main_offline_mode_skips_teacher_loading( mock_student_cfg.save_checkpoint_on_completion = False mock_student_cfg.logical_axis_rules = [] - # main() validates that student/teacher share batch shape — set explicit + # main() validates that student/teacher share batch layout — set explicit # equal scalars on both mocks so the assertion passes. mock_student_cfg.per_device_batch_size = 1 mock_student_cfg.max_target_length = 16 mock_student_cfg.gradient_accumulation_steps = 1 + mock_student_cfg.packing = True mock_teacher_cfg = mock.Mock() mock_teacher_cfg.vocab_size = 32000 mock_teacher_cfg.per_device_batch_size = 1 mock_teacher_cfg.max_target_length = 16 mock_teacher_cfg.gradient_accumulation_steps = 1 + mock_teacher_cfg.packing = True mock_pyconfig_init.side_effect = [mock_global, mock_student_cfg, mock_teacher_cfg] # 2. Model Loading @@ -1244,17 +1246,19 @@ def test_main_online_mode_loads_teacher( mock_student_cfg.save_checkpoint_on_completion = False mock_student_cfg.logical_axis_rules = [] - # main() validates that student/teacher share batch shape — set explicit + # main() validates that student/teacher share batch layout — set explicit # equal scalars on both mocks so the assertion passes. mock_student_cfg.per_device_batch_size = 1 mock_student_cfg.max_target_length = 16 mock_student_cfg.gradient_accumulation_steps = 1 + mock_student_cfg.packing = True mock_teacher_cfg = mock.Mock() mock_teacher_cfg.vocab_size = 32000 mock_teacher_cfg.per_device_batch_size = 1 mock_teacher_cfg.max_target_length = 16 mock_teacher_cfg.gradient_accumulation_steps = 1 + mock_teacher_cfg.packing = True mock_pyconfig_init.side_effect = [mock_global, mock_student_cfg, mock_teacher_cfg] mock_student_model = mock.Mock() @@ -1278,6 +1282,24 @@ def test_main_online_mode_loads_teacher( self.assertIs(model_bundle.student_model, mock_student_model) self.assertIs(model_bundle.teacher_model, mock_teacher_model) + @mock.patch("maxtext.configs.pyconfig.initialize") + def test_main_rejects_teacher_packing_that_differs_from_the_student(self, mock_pyconfig_init): + """The teacher consumes the student's batches, so a packing mismatch is refused before any model loads.""" + mock_global = mock.Mock( + student_overrides={}, + teacher_overrides={"load_parameters_path": "gs://ckpt"}, + offline_data_dir=None, + base_output_directory="", + run_name="", + ) + batch_layout = {"per_device_batch_size": 1, "max_target_length": 16, "gradient_accumulation_steps": 1} + mock_student_cfg = mock.Mock(packing=True, **batch_layout) + mock_teacher_cfg = mock.Mock(packing=False, **batch_layout) + mock_pyconfig_init.side_effect = [mock_global, mock_student_cfg, mock_teacher_cfg] + + with self.assertRaisesRegex(ValueError, "mismatch on 'packing'"): + train_distill.main(["train_distill.py", "config.yml"]) + @mock.patch.object(distillation_utils, "calculate_distillation_tflops_per_device", return_value=(0.0, 0.0, 0.0)) def test_student_freeze_param_filter(self, _mock_tflops): """Verifies that student_freeze_param_filter correctly freezes specified parameters.""" diff --git a/tests/unit/cp_as_ep_dense_no_cp_test.py b/tests/unit/cp_as_ep_dense_no_cp_test.py new file mode 100644 index 0000000000..0e37ead1d6 --- /dev/null +++ b/tests/unit/cp_as_ep_dense_no_cp_test.py @@ -0,0 +1,413 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The cp-as-ep-dense-no-cp and cp-as-ep-dense-fsdp custom mesh rules on a CPU mesh. + +Both rules are cp-as-ep with the dense weights ('embed', 'embed_attn', 'q_lora', 'kv_lora') taken off the context axis: +cp-as-ep-dense-no-cp shards them over fsdp x expert, cp-as-ep-dense-fsdp over fsdp only, and both route the dense +weights' gradients through a reduce-scatter over every token axis ('embed_grad', linears.route_dense_weight_grad). +Only the sharding changes, so the model must compute the same function. + +The model is a tiny hybrid Qwen3.5 (GatedDeltaNet kernel path, gated full attention, routed and shared-expert MoE) on +16 CPU devices, fsdp2 x context4 x expert2, with sequence-sharded GatedDeltaNet context parallelism and ring of experts +over the expert-parallel group context x expert, as in the Qwen3.5-397B configuration. Every arm runs with the same +parameter values (the cp-as-ep initialization, placed on each rule's shardings). The checks run in a subprocess so the +forced device count takes effect before JAX initializes, with JAX's persistent compilation cache off. + + - Shardings: exactly the parameters whose embed dimension is 'embed' or 'embed_attn' move (attention and + shared-expert projections, the shared-expert gate and the router), each losing 'context' (and 'expert' under + cp-as-ep-dense-fsdp) and nothing else; the GatedDeltaNet projections ('embed_gdn') and every other parameter keep + their cp-as-ep sharding. + - Collectives: the compiled program all-gathers over fsdp x expert (4 devices) or fsdp (2) where cp-as-ep gathers + over all 16. + - Numerics: the loss is bit-identical to cp-as-ep, and every parameter gradient agrees to reduction order, + max|new - old| <= 2**-6 * max|old| per leaf: the activations and the gradient partial sums are bfloat16 (the ragged + sort supports no other activation dtype), so a changed reduction order moves a gradient by bfloat16 roundings. A + float32 variant (dense MLP instead of the MoE block) must agree to 1e-5. + - Positive controls, each of which must fail the numerics gate: + the rule without its RoutedMoE change, which keeps ("context", "expert") as the expert-parallel axes for these + rules as for cp-as-ep: expert parallelism over 'expert' alone while the expert weights are sharded over + context x expert; + one reshared weight (the first attention query kernel) scaled by 1 + 1e-3. +""" + +# pylint: disable=protected-access + +import collections +import contextlib +import dataclasses +import os +import re +import subprocess +import sys +from unittest import mock + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.configs import pyconfig +from maxtext.layers import moe +from maxtext.utils import maxtext_utils +from maxtext.utils import globals as maxtext_globals + +_BASE_CONFIG_PATH = os.path.join(maxtext_globals.MAXTEXT_CONFIGS_DIR, "base.yml") +_FSDP, _CONTEXT, _EXPERT = 2, 4, 2 +_NUM_DEVICES = _FSDP * _CONTEXT * _EXPERT +# Four context ranks of two 64-token GatedDeltaNet kernel chunks each. +_SEQ = 512 +# Reduction-order gate on every gradient leaf: max|new - old| <= tolerance * max|old|. The gradient partial sums are +# bfloat16, so a changed reduction order moves an element by bfloat16 roundings of the partial sums: four bfloat16 +# ulps (2**-6) at the leaf's largest magnitude, log2 of the 16 summands (1.5 measured). The float32 variant gets a +# float32 reduction-order tolerance. +_BF16_TOL = 2.0**-6 +_F32_TOL = 1e-5 +# The float32 variant: a dense MLP instead of the MoE block (the ragged sort only supports bfloat16 activations). +_F32_DENSE = {"num_experts": 1, "num_experts_per_tok": 1, "dtype": "float32"} +_REFERENCE_RULE = "cp-as-ep" +# rule -> the mesh axes of the dense weights' embed dimension. +_RULES = {"cp-as-ep-dense-no-cp": ("fsdp", "expert"), "cp-as-ep-dense-fsdp": ("fsdp",)} +_PASS_TOKEN = "CP_AS_EP_DENSE_NO_CP_CHECKS_PASSED" + + +@pytest.mark.cpu_only +def test_cp_as_ep_dense_no_cp_on_cpu_mesh(): + """Runs `main` below in a child process with 16 CPU devices.""" + env = os.environ.copy() + env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" --xla_force_host_platform_device_count={_NUM_DEVICES}" + env["JAX_PLATFORMS"] = "cpu" + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + + +def _config(rule, **overrides): + """The tiny hybrid Qwen3.5 on the fsdp2 x context4 x expert2 mesh under `rule`, with `overrides` applied.""" + argv = [ + None, + _BASE_CONFIG_PATH, + "run_name=cp_as_ep_dense_no_cp_test", + "model_name=qwen3.5-397b-a17b", + "override_model_config=true", + "hardware=cpu", + f"custom_mesh_and_rule={rule}", + f"ici_fsdp_parallelism={_FSDP}", + f"ici_context_parallelism={_CONTEXT}", + f"ici_expert_parallelism={_EXPERT}", + "ici_tensor_parallelism=1", + "context_parallel_load_balance=False", + "attention=dot_product", + "use_gdn_kernel=true", + "gdn_cp_mode=seq", + "sparse_matmul=true", + "megablox=true", + # Tokamax gmm v2 is TPU-only; every arm runs the same GMM, so CPU uses megablox. + "use_tokamax_gmm=false", + "use_gmm_v2=false", + "use_ring_of_experts=true", + "use_ragged_sort=true", + "use_custom_sort_vjp=false", + "ragged_buffer_factor=2.0", + "num_moe_token_chunks=2", + "ring_of_experts_local_routing=true", + # The ragged sort only supports bfloat16 activations. + "dtype=bfloat16", + "weight_dtype=float32", + "scan_layers=True", + "remat_policy=custom", + "decoder_layer_input=device", + "context=device", + "packing=false", + "enable_checkpointing=false", + "enable_dropout=false", + "dataset_type=synthetic", + "base_num_decoder_layers=8", # two blocks of [3 GatedDeltaNet + 1 full attention] + "base_emb_dim=64", + "base_num_query_heads=4", + "base_num_kv_heads=2", + "head_dim=64", + "mrope_section=[3,3,2]", # rotary_dim / 2 = 64 * 0.25 / 2 + "base_mlp_dim=32", + "base_moe_mlp_dim=32", + "num_experts=16", + "num_experts_per_tok=2", + "shared_experts=1", + # The GatedDeltaNet kernel path needs head dims that are multiples of 128 and gdn_chunk_size=64. + "gdn_key_head_dim=128", + "gdn_value_head_dim=128", + "gdn_num_key_heads=2", + "gdn_num_value_heads=4", + "gdn_chunk_size=64", + "vocab_size=384", + f"max_target_length={_SEQ}", + "per_device_batch_size=0.5", + "use_multimodal=false", + "skip_jax_distributed_system=True", + "enable_tensorboard=False", + ] + argv = [a for a in argv if a is None or a.split("=", 1)[0] not in overrides] + return pyconfig.initialize(argv + [f"{key}={value}" for key, value in overrides.items()]) + + +def _inputs(cfg): + """Random tokens in one segment per row, and a fixed random cotangent on the logits. + + The cotangent makes the backward per-token sensitive: at random init the cross-entropy is nearly flat. + """ + batch = cfg.micro_batch_size_to_train_on + tokens = jax.random.randint(jax.random.PRNGKey(7), (batch, _SEQ + 1), 0, cfg.vocab_size) + ones = jnp.ones((batch, _SEQ), jnp.int32) + data = { + "inputs": tokens[:, :-1], + "targets": tokens[:, 1:], + "inputs_position": jnp.broadcast_to(jnp.arange(_SEQ, dtype=jnp.int32), (batch, _SEQ)), + "inputs_segmentation": ones, + "targets_segmentation": ones, + } + cotangent = jax.random.normal(jax.random.PRNGKey(3), (batch, _SEQ, cfg.vocab_size), jnp.float32) + return data, cotangent + + +@dataclasses.dataclass +class _Result: + loss: np.ndarray + grads: dict[str, np.ndarray] + specs: dict[str, tuple] # parameter -> the PartitionSpec of its sharding + values: dict[str, np.ndarray] # parameter -> its value + hlo: str + + +def _run(cfg, data, cotangent, values=None, control=None) -> _Result: + """Loss and its parameter gradients; `values` replaces the initial parameter values.""" + # pylint: disable=import-outside-toplevel + from flax import nnx + from flax.linen import partitioning as nn_partitioning + from maxtext.common.common_types import MODEL_MODE_TRAIN + from maxtext.utils import maxtext_utils_nnx, model_creation_utils + + mesh = maxtext_utils.get_mesh_from_config(cfg) + with control or contextlib.nullcontext(): + # As the trainer builds it: the parameters are created directly on the rule's shardings. + with nn_partitioning.axis_rules(cfg.logical_axis_rules): + create_model, abstract_model = model_creation_utils.create_nnx_abstract_model( + cfg, mesh, model_mode=MODEL_MODE_TRAIN, rng_key=jax.random.PRNGKey(0) + ) + model = maxtext_utils_nnx.create_nnx_sharded_model(abstract_model, create_model, mesh=mesh) + graphdef, params, rest = nnx.split(model, nnx.Param, ...) + keyed = jax.tree_util.tree_leaves_with_path(params) + specs = {jax.tree_util.keystr(path): tuple(leaf.sharding.spec) for path, leaf in keyed} + if values is not None: + assert values.keys() == specs.keys() + params = jax.tree_util.tree_map_with_path( + lambda path, leaf: jax.device_put(values[jax.tree_util.keystr(path)].astype(leaf.dtype), leaf.sharding), + params, + ) + + def loss_fn(params): + merged = nnx.merge(graphdef, params, rest, copy=True) + logits = merged( + decoder_input_tokens=data["inputs"], + decoder_positions=data["inputs_position"], + decoder_segment_ids=data["inputs_segmentation"], + enable_dropout=False, + decoder_target_tokens=data["targets"], + decoder_target_mask=data["targets_segmentation"], + ) + return jnp.sum(logits.astype(jnp.float32) * cotangent) + + with nn_partitioning.axis_rules(cfg.logical_axis_rules), jax.set_mesh(mesh): + compiled = jax.jit(jax.value_and_grad(loss_fn)).lower(params).compile() + loss, grads = compiled(params) + as_np = lambda tree: {jax.tree_util.keystr(p): np.asarray(x) for p, x in jax.tree_util.tree_leaves_with_path(tree)} + return _Result(np.asarray(loss), as_np(grads), specs, as_np(params), compiled.as_text()) + + +# --------------------------------------------------------------------------- +# Gates. +# --------------------------------------------------------------------------- + + +def _numerics(label, result, reference, tol=_BF16_TOL) -> bool: + """Bit-identical loss and every gradient leaf within reduction order of the reference; prints the report.""" + assert result.grads.keys() == reference.grads.keys() + loss_equal = bool(np.array_equal(result.loss, reference.loss)) + worst, worst_name, worst_l2, not_bit_equal, grads_ok = 0.0, "", 0.0, 0, True + for name, ref in reference.grads.items(): + new, ref = result.grads[name].astype(np.float64), ref.astype(np.float64) + diff, scale = float(np.max(np.abs(new - ref))), float(np.max(np.abs(ref))) + # Every comparison is False on NaN, so a NaN fails. + grads_ok &= bool(diff <= tol * scale) and bool(np.isfinite(new).all()) + rel = diff / scale if scale > 0 else diff + if rel > worst or not worst_name: + worst, worst_name = rel, name + worst_l2 = float(np.linalg.norm(new - ref)) / max(float(np.linalg.norm(ref)), 1e-30) + not_bit_equal += int(not np.array_equal(new, ref)) + passed = loss_equal and grads_ok + print( + f"{label}: numerics {'PASS' if passed else 'FAIL'}: loss bit-identical={loss_equal}" + f" (d={float(result.loss) - float(reference.loss):.6e}); gradients within {tol:.3g} x max|ref|={grads_ok}," + f" worst max|d|/max|ref|={worst:.3e} (||d||/||ref||={worst_l2:.3e}) at {worst_name}," + f" {not_bit_equal} of {len(reference.grads)} leaves not bit-identical", + flush=True, + ) + return passed + + +def _mesh_axes(entry): + """The mesh axes of one PartitionSpec entry, as a tuple.""" + if entry is None: + return () + return tuple(entry) if isinstance(entry, (tuple, list)) else (entry,) + + +def _sharding_moves(rule, result, reference) -> bool: + """Exactly the dense weights move, each from cp-as-ep's fsdp x context x expert to the rule's axes.""" + dense_axes, moved, wrong = _RULES[rule], [], [] + for name, ref_spec in reference.specs.items(): + spec = result.specs[name] + if spec == ref_spec: + continue + moved.append(name) + # The only change allowed: one dimension sharded over fsdp x context x expert becomes the rule's dense axes. + changes = [(_mesh_axes(a), _mesh_axes(b)) for a, b in zip(ref_spec, spec) if a != b] + if len(spec) != len(ref_spec) or changes != [(("fsdp", "context", "expert"), dense_axes)]: + wrong.append((name, ref_spec, spec)) + required = ( + "['query']", + "['key']", + "['value']", + "['out']", + "shared_expert']['wi_0", + "shared_expert']['wo", + "shared_expert_gate", + "routed_experts']['gate", + ) + missing = [r for r in required if not any(r in name for name in moved)] + kept = ("token_embedder", "routed_experts']['wi", "routed_experts']['wo", "in_proj_qkvz", "in_proj_ba", "out_proj") + forbidden = [n for n in moved if any(k in n for k in kept)] + print(f"{rule}: {len(moved)} of {len(reference.specs)} parameters change sharding:", flush=True) + for name in moved: + print(f" {name}: {reference.specs[name]} -> {result.specs[name]}", flush=True) + if wrong or missing or forbidden: + print(f"{rule}: unexpected sharding changes {wrong}; required but unmoved {missing}; moved {forbidden}", flush=True) + return not (wrong or missing or forbidden) + + +# The result shape is one token, or a parenthesised tuple of them when XLA combined several collectives into one. +_COLLECTIVE = re.compile(r"=\s*(?:\([^=]*?\)|\S+)\s+(all-gather|reduce-scatter|all-reduce|all-to-all)(?:-start)?\(") + + +def _group_size(line) -> int: + """The replica-group size of one HLO collective instruction (explicit, iota or mesh form).""" + if "replica_groups={}" in line: # one group of every device + return _NUM_DEVICES + match = re.search(r"replica_groups=\{(\{[^}]*\})", line) + if match: + return len(match.group(1).strip("{}").split(",")) + match = re.search(r"replica_groups=\[(\d+),(\d+)\]<=", line) + if match: + return int(match.group(2)) + match = re.search(r"replica_groups=mesh\[([^\]]*)\](?:,\s*device_ids=\((?:[^()]|\([^()]*\))*\))?\s*\{([^}]*)\}", line) + if match: + sizes = dict(re.findall(r"'(\w+)'=(\d+)", match.group(1))) + return int(np.prod([int(sizes[a]) for a in re.findall(r"'(\w+)'", match.group(2))])) + return -1 + + +def _collectives(hlo_text) -> collections.Counter: + """(collective, replica-group size) -> count, over the compiled program.""" + counts = collections.Counter() + for line in hlo_text.splitlines(): + match = _COLLECTIVE.search(line) + if match: + counts[(match.group(1), _group_size(line))] += 1 + return counts + + +# --------------------------------------------------------------------------- +# Positive controls. +# --------------------------------------------------------------------------- + + +@contextlib.contextmanager +def _moe_without_rule_change(): + """RoutedMoE as before this change: expert parallelism over 'expert' alone under the new rules.""" + real = moe.RoutedMoE.__init__ + + def init(self, *args, **kwargs): + real(self, *args, **kwargs) + assert self._expert_parallelism_name == ("context", "expert"), self._expert_parallelism_name + self._expert_parallelism_name = "expert" + + with mock.patch.object(moe.RoutedMoE, "__init__", init): + yield + + +def main(): + assert jax.device_count() == _NUM_DEVICES, jax.devices() + jax.config.update("jax_enable_compilation_cache", False) + reference_cfg = _config(_REFERENCE_RULE) + data, cotangent = _inputs(reference_cfg) + reference = _run(reference_cfg, data, cotangent) + print(f"{_REFERENCE_RULE}: collectives {sorted(_collectives(reference.hlo).items())}", flush=True) + checks, results = {}, {} + for rule, dense_axes in _RULES.items(): + result = results[rule] = _run(_config(rule), data, cotangent, values=reference.values) + gathers = _collectives(result.hlo) + print(f"{rule}: collectives {sorted(gathers.items())}", flush=True) + dense_group = int(np.prod([{"fsdp": _FSDP, "expert": _EXPERT}[a] for a in dense_axes])) + checks[f"{rule}: exactly the dense weights move, off context"] = _sharding_moves(rule, result, reference) + # Every all-gather over all devices that goes away comes back over the dense weights' smaller group. + ref_gathers = _collectives(reference.hlo) + gone = ref_gathers[("all-gather", _NUM_DEVICES)] - gathers[("all-gather", _NUM_DEVICES)] + checks[f"{rule}: {_NUM_DEVICES}-device all-gathers become {dense_group}-device ones"] = ( + 0 < gone <= gathers[("all-gather", dense_group)] - ref_gathers[("all-gather", dense_group)] + # No all-gather this test cannot size beyond those of cp-as-ep. + and gathers[("all-gather", -1)] <= ref_gathers[("all-gather", -1)] + ) + checks[f"{rule}: numerics"] = _numerics(rule, result, reference) + + # The bfloat16 differences are reduction order: with float32 activations they shrink to float32 roundings. + reference32 = _run(_config(_REFERENCE_RULE, **_F32_DENSE), data, cotangent) + result32 = _run(_config("cp-as-ep-dense-no-cp", **_F32_DENSE), data, cotangent, values=reference32.values) + checks["float32 dense variant: cp-as-ep-dense-no-cp numerics"] = _numerics( + "float32 dense cp-as-ep-dense-no-cp", result32, reference32, tol=_F32_TOL + ) + + # Positive controls. Each must fail the numerics gate (or fail to run); `controls` records whether it was caught. + controls = {} + rule = "cp-as-ep-dense-no-cp" + try: + wrong = _run(_config(rule), data, cotangent, values=reference.values, control=_moe_without_rule_change()) + controls["the rule without its RoutedMoE change"] = not _numerics("control: no RoutedMoE change", wrong, reference) + except Exception as error: # pylint: disable=broad-exception-caught + print(f"control: no RoutedMoE change: raised {type(error).__name__}: {str(error)[:400]}", flush=True) + controls["the rule without its RoutedMoE change"] = True + query = next(name for name in reference.values if re.search(r"\['query'\].*\['kernel'\]", name)) + perturbed = dict(reference.values, **{query: reference.values[query] * np.float32(1 + 1e-3)}) + wrong = _run(_config(rule), data, cotangent, values=perturbed) + controls[f"{query} scaled by 1 + 1e-3"] = not _numerics("control: perturbed query", wrong, reference) + + for name, ok in {**checks, **{f"positive control caught: {k}": v for k, v in controls.items()}}.items(): + print(f" [{'ok' if ok else 'FAIL'}] {name}", flush=True) + failed = [name for name, ok in checks.items() if not ok] + undetected = [name for name, caught in controls.items() if not caught] + if failed or undetected: + print(f"checks that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True) + sys.exit(1) + print(_PASS_TOKEN, flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/gdn_bwd_pallas_test.py b/tests/unit/gdn_bwd_pallas_test.py index ee91c2b33e..6f201ff581 100644 --- a/tests/unit/gdn_bwd_pallas_test.py +++ b/tests/unit/gdn_bwd_pallas_test.py @@ -1582,12 +1582,10 @@ def test_bwd_block_specs_and_symbolic_zero_dht_dh0_omission(self): orig_kernel = bwd_api.pallas_gdn_bwd_kernel def wrapped_kernel(*args, **kwargs): - recorded_calls.append( - ( - kwargs.get("d_recurrent_state") is not None, - bool(kwargs.get("return_dh0", False)), - ) - ) + recorded_calls.append(( + kwargs.get("d_recurrent_state") is not None, + bool(kwargs.get("return_dh0", False)), + )) return orig_kernel(*args, **kwargs) with mock.patch.object(bwd_api, "pallas_gdn_bwd_kernel", side_effect=wrapped_kernel): @@ -2783,6 +2781,375 @@ def operand_dtypes(eqn): self.assertLen(f32_dots, 17) self.assertEqual({e.params["precision"] for e in f32_dots}, {highest}) + def test_chunk_size_128_fwd_bwd_and_remat_policies(self): + """Verifies chunk_size=128 forward + backward across bf16/fp32, packing, initial states, and remat.""" + is_tpu = jax.default_backend() == "tpu" + chunk_size = 128 + dims = _packing_dims() + batch_size, seq_len = (2 if is_tpu else 1), 2 * chunk_size + key = jax.random.PRNGKey(2026) + qkv, b, a, cw, cb, al, dt, up = _packing_inputs(key, batch_size, seq_len, dims) + cs, rs = _packing_states(jax.random.PRNGKey(2027), batch_size, dims) + seg = jnp.zeros((batch_size, seq_len), dtype=jnp.int32) + # Row 0: doc 1 across chunk boundary, doc 2 with trailing 0-pad + seg = seg.at[0, : chunk_size + 20].set(1) + seg = seg.at[0, chunk_size + 20 : 2 * chunk_size - 12].set(2) + if batch_size > 1: + # Row 1: leading 0-pad, doc 1, doc 2 + seg = seg.at[1, 8 : chunk_size - 8].set(1) + seg = seg.at[1, chunk_size - 8 :].set(2) + + def ref_loss(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_): + o, (ncs, nrs) = _call_pure_jax(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, seg, dims, chunk_size) + loss = jnp.sum(o * up) + 0.1 * jnp.sum(ncs) + 0.2 * jnp.sum(nrs) + return loss, (o, ncs, nrs) + + ref_vg = jax.value_and_grad(ref_loss, argnums=tuple(range(9)), has_aux=True) + if is_tpu: + ref_vg = jax.jit(ref_vg) + (_, ref_aux), ref_grads = ref_vg(qkv, b, a, cw, cb, al, dt, cs, rs) + + cp_state_policy = jax.checkpoint_policies.save_only_these_names("gdn_cp_state") + cases = ( + [ + (act_dtype, tol, policy) + for act_dtype, tol in ((jnp.float32, 1.5e-2), (jnp.bfloat16, 4e-2)) + for policy in (None, jax.checkpoint_policies.nothing_saveable, cp_state_policy) + ] + if is_tpu + else [ + (jnp.float32, 2e-3, None), + (jnp.bfloat16, 4e-2, cp_state_policy), + ] + ) + + for act_dtype, tol, policy in cases: + qkv_d = qkv.astype(act_dtype) + b_d = b.astype(act_dtype) + a_d = a.astype(act_dtype) + cw_d = cw.astype(act_dtype) + cb_d = cb.astype(act_dtype) + up_d = up.astype(act_dtype) + + def kern_step(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, _act_dtype=act_dtype): + return gdn_bwd_pallas.gdn_decoupled_conv1d( + qkv_, + b_, + a_, + cw_, + cb_, + al_, + dt_, + cs_, + rs_, + dims["num_k_heads"], + dims["num_v_heads"], + dims["head_k_dim"], + dims["head_v_dim"], + dims["conv_kernel_size"], + chunk_size, + True, + _act_dtype, + None, + seg, + ) + + step_fn = jax.checkpoint(kern_step, policy=policy) if policy is not None else kern_step + + def kern_loss(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, _step=step_fn, _up_d=up_d): + o, (ncs, nrs) = _step(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_) + loss = ( + jnp.sum(o.astype(jnp.float32) * _up_d.astype(jnp.float32)) + + 0.1 * jnp.sum(ncs.astype(jnp.float32)) + + 0.2 * jnp.sum(nrs) + ) + return loss, (o, ncs, nrs) + + kern_vg = jax.value_and_grad(kern_loss, argnums=tuple(range(9)), has_aux=True) + if is_tpu: + kern_vg = jax.jit(kern_vg) + (_, got_aux), got_grads = kern_vg(qkv_d, b_d, a_d, cw_d, cb_d, al, dt, cs, rs) + for name, got, exp in zip(("out", "next_conv_state", "next_recurrent_state"), got_aux, ref_aux): + np.testing.assert_allclose( + np.asarray(got, dtype=np.float32), + np.asarray(exp, dtype=np.float32), + rtol=tol, + atol=tol, + err_msg=f"chunk_size={chunk_size} dtype={act_dtype} policy={policy}: {name}", + ) + for name, got, exp in zip(_GRAD_NAMES + ("d_conv_state", "d_recurrent_state"), got_grads, ref_grads): + np.testing.assert_allclose( + np.asarray(got, dtype=np.float32), + np.asarray(exp, dtype=np.float32), + rtol=tol, + atol=tol, + err_msg=f"chunk_size={chunk_size} dtype={act_dtype} policy={policy}: {name}", + ) + + def test_chunk_size_128_head_group_size_64_v_heads(self): + """Verifies chunk_size=128 with num_k_heads=16, num_v_heads=64 for both bf16 and fp32 activations + packing.""" + if jax.default_backend() != "tpu": + self.skipTest("TPU-specific VMEM head_group_size test at chunk_size=128") + from maxtext.kernels.gdn import compute_gdn as local_compute_gdn # pylint: disable=import-outside-toplevel + + self.assertIsNone(local_compute_gdn.default_head_group_size(64, 16, 64)) + self.assertEqual(local_compute_gdn.default_head_group_size(128, 16, 64), 8) + self.assertEqual(local_compute_gdn.default_head_group_size(128, 16, 64, jnp.bfloat16), 8) + self.assertEqual(local_compute_gdn.default_head_group_size(128, 16, 64, jnp.float32), 2) + self.assertIsNone(local_compute_gdn.default_head_group_size(128, 8, 32, jnp.bfloat16)) + + chunk_size = 128 + dims = {"num_k_heads": 16, "num_v_heads": 64, "head_k_dim": 128, "head_v_dim": 128, "conv_kernel_size": 4} + batch_size, seq_len = 1, 256 + key = jax.random.PRNGKey(3030) + qkv, b, a, cw, cb, al, dt, up = _packing_inputs(key, batch_size, seq_len, dims) + cs, rs = _packing_states(jax.random.PRNGKey(3031), batch_size, dims) + seg = jnp.ones((batch_size, seq_len), dtype=jnp.int32) + seg = seg.at[0, :8].set(0) + seg = seg.at[0, chunk_size + 20 : 2 * chunk_size - 6].set(2) + seg = seg.at[0, 2 * chunk_size - 6 :].set(0) + + def ref_loss(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_): + o, (ncs, nrs) = _call_pure_jax(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, seg, dims, chunk_size) + loss = jnp.sum(o * up) + 0.1 * jnp.sum(ncs) + 0.2 * jnp.sum(nrs) + return loss, (o, ncs, nrs) + + (_, ref_aux), ref_grads = jax.jit(jax.value_and_grad(ref_loss, argnums=tuple(range(9)), has_aux=True))( + qkv, b, a, cw, cb, al, dt, cs, rs + ) + + for act_dtype, tol in ((jnp.bfloat16, 4e-2), (jnp.float32, 2e-2)): + qkv_d, b_d, a_d, cw_d, cb_d, up_d = (x.astype(act_dtype) for x in (qkv, b, a, cw, cb, up)) + + def kern_loss(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, _act_dtype=act_dtype, _up_d=up_d): + o, (ncs, nrs) = gdn_bwd_pallas.gdn_decoupled_conv1d( + qkv_, + b_, + a_, + cw_, + cb_, + al_, + dt_, + cs_, + rs_, + dims["num_k_heads"], + dims["num_v_heads"], + dims["head_k_dim"], + dims["head_v_dim"], + dims["conv_kernel_size"], + chunk_size, + True, + _act_dtype, + None, + seg, + ) + loss = ( + jnp.sum(o.astype(jnp.float32) * _up_d.astype(jnp.float32)) + + 0.1 * jnp.sum(ncs.astype(jnp.float32)) + + 0.2 * jnp.sum(nrs) + ) + return loss, (o, ncs, nrs) + + (_, got_aux), got_grads = jax.jit(jax.value_and_grad(kern_loss, argnums=tuple(range(9)), has_aux=True))( + qkv_d, b_d, a_d, cw_d, cb_d, al, dt, cs, rs + ) + for name, got, exp in zip(("out", "next_conv_state", "next_recurrent_state"), got_aux, ref_aux): + np.testing.assert_allclose( + np.asarray(got, dtype=np.float32), + np.asarray(exp, dtype=np.float32), + rtol=tol, + atol=tol, + err_msg=f"64 V heads chunk_size=128 dtype={act_dtype}: {name}", + ) + for name, got, exp in zip(_GRAD_NAMES + ("d_conv_state", "d_recurrent_state"), got_grads, ref_grads): + np.testing.assert_allclose( + np.asarray(got, dtype=np.float32), + np.asarray(exp, dtype=np.float32), + rtol=tol, + atol=tol, + err_msg=f"64 V heads chunk_size=128 dtype={act_dtype}: {name}", + ) + + def test_chunk_size_128_head_cp_packed_with_initial_states(self): + """Verifies head-sharded CP (gdn_cp_mode='head') at chunk_size=128 with packed segment_ids and initial states.""" + devices = jax.devices() + if len(devices) < 2: + self.skipTest("Requires at least 2 devices for head-CP") + is_tpu = jax.default_backend() == "tpu" + chunk_size = 128 + dims = _packing_dims() + cp_size = 2 + local_dims = { + **dims, + "num_k_heads": dims["num_k_heads"] // cp_size, + "num_v_heads": dims["num_v_heads"] // cp_size, + } + batch_size, seq_len = (2 if is_tpu else 1), 2 * chunk_size + key = jax.random.PRNGKey(5050) + qkv, b, a, cw, cb, al, dt, up = _packing_inputs(key, batch_size, seq_len, dims) + cs, rs = _packing_states(jax.random.PRNGKey(5051), batch_size, dims) + seg = jnp.ones((batch_size, seq_len), dtype=jnp.int32) + seg = seg.at[0, :10].set(0) + seg = seg.at[0, chunk_size + 14 : 2 * chunk_size - 6].set(2) + seg = seg.at[0, 2 * chunk_size - 6 :].set(0) + if batch_size > 1: + seg = seg.at[1, :chunk_size].set(0) + seg = seg.at[1, chunk_size + 8 :].set(2) + + q_size = dims["num_k_heads"] * dims["head_k_dim"] + k_size = dims["num_k_heads"] * dims["head_k_dim"] + v_size = dims["num_v_heads"] * dims["head_v_dim"] + + def split_qkv_heads(x, axis): + q, k, v = jnp.split(x, [q_size, q_size + k_size], axis=axis) + q_s = q.reshape(*q.shape[:axis], cp_size, q_size // cp_size, *q.shape[axis + 1 :]) + k_s = k.reshape(*k.shape[:axis], cp_size, k_size // cp_size, *k.shape[axis + 1 :]) + v_s = v.reshape(*v.shape[:axis], cp_size, v_size // cp_size, *v.shape[axis + 1 :]) + return jnp.concatenate([q_s, k_s, v_s], axis=axis + 1) + + def merge_qkv_heads(x, axis): + lq = q_size // cp_size + lk = k_size // cp_size + q_s, k_s, v_s = jnp.split(x, [lq, lq + lk], axis=axis + 1) + q = q_s.reshape(*q_s.shape[:axis], q_size, *q_s.shape[axis + 2 :]) + k = k_s.reshape(*k_s.shape[:axis], k_size, *k_s.shape[axis + 2 :]) + v = v_s.reshape(*v_s.shape[:axis], v_size, *v_s.shape[axis + 2 :]) + return jnp.concatenate([q, k, v], axis=axis) + + mesh = jax.sharding.Mesh(np.array(devices[:cp_size]), ("context",)) + ps_h2 = jax.sharding.PartitionSpec(None, None, "context", None) + ps_cw = jax.sharding.PartitionSpec(None, None, "context", None) + ps_h0 = jax.sharding.PartitionSpec("context", None) + ps_h0_1d = jax.sharding.PartitionSpec("context", None) + ps_rs = jax.sharding.PartitionSpec(None, "context", None, None) + ps_rep = jax.sharding.PartitionSpec() + + def loss_single(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_): + o, (ncs, nrs) = _call_pure_jax(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, seg, dims, chunk_size) + return jnp.sum(o * up) + 0.13 * jnp.sum(ncs) + 0.19 * jnp.sum(nrs), (o, ncs, nrs) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=(ps_h2, ps_h2, ps_h2, ps_cw, ps_h0, ps_h0_1d, ps_h0_1d, ps_h2, ps_rs, ps_rep), + out_specs=(ps_h2, ps_h2, ps_rs), + check_vma=False, + ) + def run_head_cp(qkv_h, b_h, a_h, cw_h, cb_h, al_h, dt_h, cs_h, rs_h, seg_rep): + o_h, (ncs_h, nrs_h) = _call_kernel_api( + qkv_h.squeeze(2), + b_h.squeeze(2), + a_h.squeeze(2), + cw_h.squeeze(2), + cb_h.squeeze(0), + al_h.squeeze(0), + dt_h.squeeze(0), + cs_h.squeeze(2), + rs_h, + seg_rep, + local_dims, + chunk_size, + None, + ) + return jnp.expand_dims(o_h, 2), jnp.expand_dims(ncs_h, 2), nrs_h + + def loss_head_cp(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_): + nv_local = local_dims["num_v_heads"] + o_h, ncs_h, nrs = run_head_cp( + split_qkv_heads(qkv_, 2), + b_.reshape(batch_size, seq_len, cp_size, nv_local), + a_.reshape(batch_size, seq_len, cp_size, nv_local), + split_qkv_heads(cw_, 2), + split_qkv_heads(cb_, 0), + al_.reshape(cp_size, nv_local), + dt_.reshape(cp_size, nv_local), + split_qkv_heads(cs_, 2), + rs_, + seg, + ) + o = o_h.reshape(batch_size, seq_len, dims["num_v_heads"], dims["head_v_dim"]) + ncs = merge_qkv_heads(ncs_h, 2) + return jnp.sum(o * up) + 0.13 * jnp.sum(ncs) + 0.19 * jnp.sum(nrs), (o, ncs, nrs) + + args = (qkv, b, a, cw, cb, al, dt, cs, rs) + vg_s = jax.value_and_grad(loss_single, argnums=tuple(range(9)), has_aux=True) + vg_h = jax.value_and_grad(loss_head_cp, argnums=tuple(range(9)), has_aux=True) + if is_tpu: + vg_s = jax.jit(vg_s) + vg_h = jax.jit(vg_h) + (_, aux_s), grads_s = vg_s(*args) + (_, aux_h), grads_h = vg_h(*args) + tol = 2e-2 if is_tpu else 2e-3 + for name, got, exp in zip(("out", "next_conv_state", "next_recurrent_state"), aux_h, aux_s): + np.testing.assert_allclose(np.asarray(got), np.asarray(exp), rtol=tol, atol=tol, err_msg=f"head-CP c=128: {name}") + for name, got, exp in zip(_GRAD_NAMES + ("d_conv_state", "d_recurrent_state"), grads_h, grads_s): + np.testing.assert_allclose(np.asarray(got), np.asarray(exp), rtol=tol, atol=tol, err_msg=f"head-CP c=128: {name}") + + def test_chunk_size_128_seq_cp_packed_with_initial_states(self): + """Verifies seq-CP (cp_axis_name='context') at chunk_size=128 with on-the-fly t_inv recomputation.""" + devices = jax.devices() + if len(devices) < 2: + self.skipTest("Requires at least 2 devices for seq-CP") + is_tpu = jax.default_backend() == "tpu" + chunk_size = 128 + dims = _packing_dims() + batch_size, seq_len = (2 if is_tpu else 1), 4 * chunk_size # 2 chunks per rank across 2 CP ranks + key = jax.random.PRNGKey(4040) + qkv, b, a, cw, cb, al, dt, up = _packing_inputs(key, batch_size, seq_len, dims) + cs, rs = _packing_states(jax.random.PRNGKey(4041), batch_size, dims) + seg = jnp.ones((batch_size, seq_len), dtype=jnp.int32) + seg = seg.at[0, :12].set(0) + seg = seg.at[0, chunk_size + 16 : 3 * chunk_size - 8].set(2) + seg = seg.at[0, 3 * chunk_size - 8 :].set(0) + if batch_size > 1: + seg = seg.at[1, : 2 * chunk_size].set(0) + seg = seg.at[1, 2 * chunk_size + 10 :].set(2) + + mesh = jax.sharding.Mesh(np.array(devices[:2]), ("context",)) + ps_seq = jax.sharding.PartitionSpec(None, "context", None) + ps_seg = jax.sharding.PartitionSpec(None, "context") + ps_rep = jax.sharding.PartitionSpec() + ps_up = jax.sharding.PartitionSpec(None, "context", None, None) + + def loss_single(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_): + o, (ncs, nrs) = _call_pure_jax(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, seg, dims, chunk_size) + return jnp.sum(o * up) + 0.13 * jnp.sum(ncs) + 0.19 * jnp.sum(nrs), (o, ncs, nrs) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=(ps_seq, ps_seq, ps_seq, ps_rep, ps_rep, ps_rep, ps_rep, ps_rep, ps_rep, ps_seg, ps_up), + out_specs=(ps_up, ps_rep, ps_rep), + check_vma=False, + ) + def run_cp(qkv_s, b_s, a_s, cw_, cb_, al_, dt_, cs_, rs_, seg_s, up_s): + del up_s + o, (ncs, nrs) = _call_kernel_api(qkv_s, b_s, a_s, cw_, cb_, al_, dt_, cs_, rs_, seg_s, dims, chunk_size, "context") + return o, ncs, nrs + + # Also exercise custom remat saving only gdn_cp_state on TPU to verify chunk_states != None and t_inv_fwd == None. + run_cp_fn = ( + jax.checkpoint(run_cp, policy=jax.checkpoint_policies.save_only_these_names("gdn_cp_state")) if is_tpu else run_cp + ) + + def loss_cp(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_): + o, ncs, nrs = run_cp_fn(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, seg, up) + return jnp.sum(o * up) + 0.13 * jnp.sum(ncs) + 0.19 * jnp.sum(nrs), (o, ncs, nrs) + + args = (qkv, b, a, cw, cb, al, dt, cs, rs) + vg_s = jax.value_and_grad(loss_single, argnums=tuple(range(9)), has_aux=True) + vg_c = jax.value_and_grad(loss_cp, argnums=tuple(range(9)), has_aux=True) + if is_tpu: + vg_s = jax.jit(vg_s) + vg_c = jax.jit(vg_c) + (_, aux_s), grads_s = vg_s(*args) + (_, aux_c), grads_c = vg_c(*args) + tol = 2e-2 if is_tpu else 2e-3 + for name, got, exp in zip(("out", "next_conv_state", "next_recurrent_state"), aux_c, aux_s): + np.testing.assert_allclose(np.asarray(got), np.asarray(exp), rtol=tol, atol=tol, err_msg=f"CP c=128: {name}") + for name, got, exp in zip(_GRAD_NAMES + ("d_conv_state", "d_recurrent_state"), grads_c, grads_s): + np.testing.assert_allclose(np.asarray(got), np.asarray(exp), rtol=tol, atol=tol, err_msg=f"CP c=128: {name}") + def _find_eqns(jaxpr, primitive_name): """Returns every `primitive_name` equation in `jaxpr`, including those in nested (tuples of) sub-jaxprs.""" diff --git a/tests/unit/gdn_conv1d_dense_token_stack_test.py b/tests/unit/gdn_conv1d_dense_token_stack_test.py new file mode 100644 index 0000000000..da223e6194 --- /dev/null +++ b/tests/unit/gdn_conv1d_dense_token_stack_test.py @@ -0,0 +1,922 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://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. + +"""gdn_conv1d_dense_token_stack: the Pallas GDN forward kernel's Conv1D output built from dense per-token views. + +The option changes how compute_conv1d.causal_conv1d assembles a value the kernel computes either way: each token's +float32 row is viewed as [dim / 1024, 8, 128] and the views are concatenated, instead of jnp.stack(axis=1) of the +rows. Every element lands where jnp.stack puts it; only the layout the TPU compiler gives the value changes. + +The forward kernel cannot run on CPU (neither Pallas interpreter supports its emit_pipeline DMAs), so: + - Values (in this process): causal_conv1d, the one function whose code changes, gives the same bits with the + option as without it at the kernel's in-VMEM shapes (widths 1024 and 12288, packed segment ids on/off, conv + bias on/off, full and partial chunks), and so does the SiLU the kernel applies next. A token-minor regrouping + is caught, so the comparison can fail. Without the option the rows go through jnp.stack itself (the same + jaxpr), so the default program is the base program. + - Programs, traced for a TPU target on this CPU host with no TPU or libtpu (runtime_utils.target_platform follows + the abstract mesh). In this process: fused_conv1d_gdn hands the option to its per-sequence kernel only, and its + batched (decode) kernel is the same program with and without it. In a child process (see main): every + forward-kernel call carries the dense regrouping exactly when the option is on: wrapper.fused_conv1d_gdn; + gdn_decoupled_conv1d's primal, its forward rule and the rule's replay under remat; both sequence-sharded + context-parallel passes, under the forward rule and (forward only) under the primal; head-sharded context + parallelism; and model_runner.run_gdn_kernel_layer reading the config key. The option composes with + gdn_conv1d_input_in_qkv_dtype, and the backward rule accepts every positional form. Deliberately wrong + versions must FAIL every check they break: an option the conv ignores, one the forward rule drops, and one the + primal drops. + - Guards: the per-call width check and the config validator reject the option where it would do nothing, and the + validator's copy of model_runner's head-vs-sequence context-parallel rule agrees with model_runner. +""" + +# pylint: disable=protected-access + +import collections +import functools +import itertools +import os +import re +import subprocess +import sys +import types + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax.extend import core as jax_core +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.common.common_types import ShardMode +from maxtext.configs import pyconfig +from maxtext.configs import types as config_types +from maxtext.kernels.gdn import compute_conv1d +from maxtext.kernels.gdn import config as gdn_config +from maxtext.kernels.gdn import model_runner +from maxtext.kernels.gdn import wrapper +from maxtext.kernels.gdn.gdn_bwd import api as gdn_api +from maxtext.kernels.gdn.gdn_bwd import cp_gdn +from tests.utils.test_helpers import get_test_config_path + +_PASS_TOKEN = "GDN_CONV1D_DENSE_TOKEN_STACK_CHECKS_PASSED" +_CHUNK = 64 +_KERNEL = 4 +_HEAD_DIM = 128 + + +def _bits(x): + """The array's bytes, so equality means bit-identical (NaN-safe, -0.0 != 0.0).""" + return np.asarray(jax.device_get(x)).tobytes() + + +def _gdn_config(dim, num_v_heads, has_seg_ids, dense): + return gdn_config.GDNConfig( + mode=gdn_config.GDNMode.PER_SEQ, + dtypes=gdn_config.Dtypes(jnp.bfloat16, jnp.bfloat16, jnp.float32, jnp.float32, jnp.float32), + batch_size=_CHUNK, + dim_size=dim, + kernel_size=_KERNEL, + tile_size=_CHUNK, + num_kq_heads=num_v_heads // 2, + num_v_heads=num_v_heads, + kq_head_dim=_HEAD_DIM, + v_head_dim=_HEAD_DIM, + has_seg_ids=has_seg_ids, + conv_dense_token_stack=dense, + ) + + +def _conv_inputs(dim, num_v_heads, seed): + """causal_conv1d's in-VMEM operands for one PER_SEQ tile: lhs = [halo | chunk tokens], weights, bias, b lanes.""" + rng = np.random.default_rng(seed) + lhs = jnp.asarray(rng.standard_normal((1, _CHUNK + _KERNEL - 1, 1, dim)).astype(np.float32)) + weight = jnp.asarray((rng.standard_normal((_KERNEL, 1, dim)) * 0.5).astype(np.float32)) + bias = jnp.asarray(rng.standard_normal((dim,)).astype(np.float32)) + # b's lanes num_v_heads (signed active segment of each token: two documents, then padding) and num_v_heads + 1 + # (chunk header: the halo's segment ids, then the previous chunk's segment). + b_vreg = np.zeros((1, _CHUNK, 1, 128), np.float32) + b_vreg[0, :, 0, num_v_heads] = np.concatenate([np.full(21, 1.0), np.full(30, 2.0), np.full(_CHUNK - 51, -2.0)]) + b_vreg[0, :_KERNEL, 0, num_v_heads + 1] = 1.0 + return lhs, weight, bias, jnp.asarray(b_vreg) + + +def _stack_counts(jaxpr, dim): + """(dense, stacked): how many times `jaxpr` (sub-jaxprs included) assembles a [1, chunk, 1, dim] Conv1D output from + dense views (a concatenate to [chunk * dim / 1024, 8, 128]) or by stacking (a stack, or a concatenate, whose result + is [1, chunk, 1, dim]; jax 0.11 traces jnp.stack as its own primitive).""" + dense_shape, stack_shape = (_CHUNK * dim // 1024, 8, 128), (1, _CHUNK, 1, dim) + dense = stacked = 0 + for eqn in jaxpr.eqns: + if eqn.primitive.name in ("concatenate", "stack"): + shape = tuple(eqn.outvars[0].aval.shape) + dense += eqn.primitive.name == "concatenate" and shape == dense_shape + stacked += shape == stack_shape + for sub in jax_core.jaxprs_in_params(eqn.params): + sub_dense, sub_stacked = _stack_counts(sub, dim) + dense, stacked = dense + sub_dense, stacked + sub_stacked + return dense, stacked + + +def _forward_kernels(closed_jaxpr): + """(name, kernel jaxpr) of every forward-kernel call (a pallas_call named fused_conv1d_gdn*), in trace order.""" + found = [] + + def visit(jaxpr): + for eqn in jaxpr.eqns: + if eqn.primitive.name == "pallas_call" and str(eqn.params.get("name", "")).startswith("fused_conv1d_gdn"): + found.append((str(eqn.params["name"]), eqn.params["jaxpr"])) + continue + for sub in jax_core.jaxprs_in_params(eqn.params): + visit(sub) + + visit(closed_jaxpr.jaxpr) + return found + + +def _token_minor_stack(out_list): + """Deliberately wrong regrouping: interleaves the tokens inside each group of 8 x 128 registers.""" + row = out_list[0] + views = jnp.stack([o.reshape(-1, 8, 128) for o in out_list], axis=1) + return views.reshape(row.shape[0], len(out_list), *row.shape[1:]) + + +class StackConv1dTokensTest(parameterized.TestCase): + """stack_conv1d_tokens(dense=True) equals jnp.stack(axis=1) bit for bit, and rejects rows it cannot regroup.""" + + @parameterized.named_parameters( + ("chunk64_width1024", 64, 1024), + ("chunk64_width12288", 64, 12288), + ("chunk3_width2048", 3, 2048), + ) + def test_dense_equals_stack(self, num_tokens, dim): + rng = np.random.default_rng(num_tokens + dim) + rows = [jnp.asarray(rng.standard_normal((1, 1, dim)).astype(np.float32)) for _ in range(num_tokens)] + ref = jnp.stack(rows, axis=1) + for got in ( + compute_conv1d.stack_conv1d_tokens(rows, True), + jax.jit(compute_conv1d.stack_conv1d_tokens, static_argnums=1)(rows, True), + ): + self.assertEqual(got.shape, ref.shape) + self.assertEqual(got.dtype, ref.dtype) + self.assertEqual(_bits(got), _bits(ref)) + self.assertEqual(_bits(compute_conv1d.stack_conv1d_tokens(rows, False)), _bits(ref)) + # Positive control: a regrouping that moves elements is caught by the same comparison. (With one register per + # token, width 1024, the token-minor interleave is the same order, so it only moves elements at wider rows.) + if dim > 1024: + self.assertNotEqual(_bits(_token_minor_stack(rows)), _bits(ref)) + + def test_default_is_jnp_stack(self): + """dense=False traces to jnp.stack(axis=1) itself, so with the option off causal_conv1d is the base program. + + Equal values or shapes would not show this: a rewrite of the default path that keeps the values (a concatenate of + expanded rows, say) is a different program. + """ + rows = [jnp.zeros((1, 1, 2048), jnp.float32) for _ in range(_CHUNK)] + ref = str(jax.make_jaxpr(lambda r: jnp.stack(r, axis=1))(rows)) + self.assertEqual(str(jax.make_jaxpr(lambda r: compute_conv1d.stack_conv1d_tokens(r, False))(rows)), ref) + # Positive control: the comparison sees a program change that keeps the values. + self.assertNotEqual(str(jax.make_jaxpr(lambda r: compute_conv1d.stack_conv1d_tokens(r, True))(rows)), ref) + + @parameterized.named_parameters( + ("two_sequences", (2, 1, 1024), jnp.float32), + ("partial_register", (1, 1, 1536), jnp.float32), + ("bfloat16_rows", (1, 1, 2048), jnp.bfloat16), + ) + def test_rejects_rows_it_cannot_regroup(self, shape, dtype): + rows = [jnp.zeros(shape, dtype)] * 4 + with self.assertRaisesRegex(ValueError, "one sequence per tile and float32 rows"): + compute_conv1d.stack_conv1d_tokens(rows, True) + self.assertEqual(compute_conv1d.stack_conv1d_tokens(rows, False).shape, (shape[0], 4) + shape[1:]) + + +_CONV_CASES = [ + { + "testcase_name": f"width{dim}_seg{int(seg)}_bias{int(bias)}_real{real}", + "dim": dim, + "num_v_heads": nv, + "has_seg_ids": seg, + "with_bias": bias, + "real": real, + } + for dim, nv in ((1024, 4), (12288, 64)) + for seg in (False, True) + for bias in (False, True) + for real in (_CHUNK, 37) +] + + +class CausalConv1dTest(parameterized.TestCase): + """causal_conv1d with the option gives the same bits; and the option changes the traced conv.""" + + @parameterized.named_parameters(*_CONV_CASES) + def test_bit_identical(self, dim, num_v_heads, has_seg_ids, with_bias, real): + lhs, weight, bias, b_vreg = _conv_inputs(dim, num_v_heads, seed=dim + 2 * has_seg_ids + with_bias) + + def run(dense): + out, state = compute_conv1d.causal_conv1d( + real_sizes=jnp.array([real], jnp.int32), + lhs=lhs, + conv_weight=weight, + conv_bias=bias if with_bias else None, + cfg=_gdn_config(dim, num_v_heads, has_seg_ids, dense), + b_vreg=b_vreg if has_seg_ids else None, + ) + return out, state, jax.nn.silu(out) # inner_kernel applies the SiLU to the conv output next + + for ref, got in zip(jax.jit(lambda: run(False))(), jax.jit(lambda: run(True))()): + self.assertEqual(ref.shape, got.shape) + self.assertEqual(ref.dtype, got.dtype) + self.assertEqual(_bits(ref), _bits(got)) + + def test_option_changes_the_traced_conv(self): + """Positive control for the value test: the option must change the program, not only keep the values.""" + dim, num_v_heads = 2048, 8 + lhs, weight, bias, _ = _conv_inputs(dim, num_v_heads, seed=0) + + def stacks(dense): + cfg = _gdn_config(dim, num_v_heads, False, dense) + jaxpr = jax.make_jaxpr(lambda: compute_conv1d.causal_conv1d(jnp.array([_CHUNK]), lhs, weight, bias, cfg))() + return _stack_counts(jaxpr.jaxpr, dim) + + self.assertEqual(stacks(True), (1, 0)) + self.assertEqual(stacks(False), (0, 1)) + + +class KernelIdentityTest(absltest.TestCase): + """The option is left out of the kernel name and metadata, so the program outside the kernel body is unchanged.""" + + def test_name_and_metadata_unchanged(self): + off, on = _gdn_config(1024, 4, True, False), _gdn_config(1024, 4, True, True) + self.assertEqual(off.get_kernel_name(), on.get_kernel_name()) + self.assertEqual(off.get_metadata(), on.get_metadata()) + self.assertNotIn("conv_dense_token_stack", off.get_metadata()) + self.assertFalse(gdn_config.GDNConfig.__dataclass_fields__["conv_dense_token_stack"].default) + + +def _tpu_target(shape=(), axes=()): + """A context in which traces target TPU7x as in an AOT compile: no TPU and no libtpu needed.""" + device = jax.sharding.AbstractDevice(device_kind="TPU7x", num_cores=1, platform="tpu") + mesh = jax.sharding.AbstractMesh( + shape, axes, axis_types=(jax.sharding.AxisType.Auto,) * len(axes), abstract_device=device + ) + return jax.sharding.use_abstract_mesh(mesh), mesh + + +def _fused_args(dim, num_v_heads, seq=128): + a = jax.ShapeDtypeStruct + return ( + a((seq, dim), jnp.bfloat16), + a((seq, num_v_heads), jnp.bfloat16), + a((seq, num_v_heads), jnp.bfloat16), + a((2, _KERNEL - 1, dim), jnp.bfloat16), + a((2, num_v_heads, _HEAD_DIM, _HEAD_DIM), jnp.float32), + a((dim, 1, _KERNEL), jnp.float32), + None, + a((num_v_heads,), jnp.float32), + a((num_v_heads,), jnp.float32), + a((2,), jnp.int32), + a((1,), jnp.int32), + a((3,), jnp.int32), + a((1,), jnp.int32), + ) + + +def _trace_fused(num_k_heads, num_v_heads, dense, *, prefill_only=True): + """wrapper.fused_conv1d_gdn traced for a TPU7x target. + + Prefill-only (as gdn_decoupled_conv1d calls it) runs the per-sequence kernel alone; prefill_only=False (mixed + mode) runs the batched (decode) kernel first, at a decode tile of 1 because _fused_args has one sequence. + """ + dim = (2 * num_k_heads + num_v_heads) * _HEAD_DIM + ctx, _ = _tpu_target() + fn = functools.partial( + wrapper.fused_conv1d_gdn, + n_kq=num_k_heads, + n_v=num_v_heads, + d_k=_HEAD_DIM, + d_v=_HEAD_DIM, + kernel_size=_KERNEL, + compute_precision=jnp.dtype(jnp.float32), + mixed_tile_size=_CHUNK, + is_prefill_only=prefill_only, + conv_dense_token_stack=dense, + ) + with ctx: + return jax.make_jaxpr(fn)(*_fused_args(dim, num_v_heads)) + + +class WidthCheckTest(absltest.TestCase): + """fused_conv1d_gdn rejects the option for a Conv1D width that is not a whole number of float32 registers.""" + + def test_rejects_width_512(self): + _trace_fused(1, 2, False) # width 512; the default path is unchanged + with self.assertRaisesRegex(ValueError, r"requires the Conv1D width \(512\) to be a multiple of 1024"): + _trace_fused(1, 2, True) + + def test_accepts_width_1024(self): + _trace_fused(2, 4, True) + + +class BatchedKernelTest(absltest.TestCase): + """fused_conv1d_gdn hands the option to its per-sequence kernel only; the batched (decode) kernel keeps jnp.stack.""" + + def test_batched_kernel_does_not_take_the_option(self): + # Mixed mode runs the batched kernel, then the per-sequence kernel. With one sequence the batched kernel's tile + # holds one sequence, so its one-token rows are the only batched rows the dense regrouping could take. + kernels = {dense: _forward_kernels(_trace_fused(2, 4, dense, prefill_only=False)) for dense in (False, True)} + names = [name for name, _ in kernels[True]] + self.assertEqual(names, [name for name, _ in kernels[False]]) + self.assertLen(names, 2) + self.assertTrue(names[0].startswith("fused_conv1d_gdn_batched_b1_"), names) + self.assertTrue(names[1].startswith("fused_conv1d_gdn_per_seq_"), names) + (_, batched_off), (_, per_seq_off) = kernels[False] + (_, batched_on), (_, per_seq_on) = kernels[True] + self.assertEqual(str(batched_on), str(batched_off)) + self.assertEqual(_stack_counts(per_seq_on, 1024), (1, 0)) + self.assertEqual(_stack_counts(per_seq_off, 1024), (0, 1)) + # Positive control: the same comparison tells the per-sequence kernel's two programs apart. + self.assertNotEqual(str(per_seq_on), str(per_seq_off)) + + +def _validator_cfg(**overrides): + """The fields validate_gdn_config reads, for a Qwen3.5-397B-shaped GDN kernel path that passes it.""" + fields = { + "decoder_block": config_types.DecoderBlockType.QWEN3_5, + "ici_context_parallelism": 1, + "dcn_context_parallelism": 1, + "ici_context_usp_ulysses_parallelism": 1, + "dcn_context_usp_ulysses_parallelism": 1, + "use_gdn_kernel": True, + "ici_tensor_parallelism": 1, + "dcn_tensor_parallelism": 1, + "gdn_chunk_size": 64, + "gdn_key_head_dim": 128, + "gdn_value_head_dim": 128, + "gdn_num_key_heads": 16, + "gdn_num_value_heads": 64, + "use_qk_norm_in_gdn": True, + "max_target_length": 4096, + "gdn_conv1d_input_in_qkv_dtype": False, + "gdn_conv1d_dense_token_stack": False, + "gdn_cp_mode": "auto", + } + fields.update(overrides) + return types.SimpleNamespace(**fields) + + +class ValidatorTest(parameterized.TestCase): + """The validator rejects the key where it would do nothing; the default config is untouched.""" + + @parameterized.named_parameters( + ("no_cp", {}), + ("head_cp4", {"ici_context_parallelism": 4, "gdn_cp_mode": "head"}), # 3072 per rank + ("auto_cp2_head_sharded", {"ici_context_parallelism": 2}), # 6144 per rank + ("auto_cp4_seq_sharded", {"ici_context_parallelism": 4, "gdn_num_key_heads": 2, "gdn_num_value_heads": 4}), + # 2048: model_runner runs auto CP 4 sequence-sharded even with 4 key heads (head-sharded would be 512 per rank). + ( + "auto_cp4_seq_sharded_four_key_heads", + {"ici_context_parallelism": 4, "gdn_num_key_heads": 4, "gdn_num_value_heads": 8}, + ), + ) + def test_accepted(self, overrides): + cfg = _validator_cfg(gdn_conv1d_dense_token_stack=True, **overrides) + config_types.MaxTextConfig.validate_gdn_config(cfg, is_rl_config=False) + + def test_head_sharding_agrees_with_model_runner(self): + """The validator's copy of model_runner._resolve_gdn_cp_mode's head-vs-sequence rule agrees with it. + + configs/types.py cannot import the kernel module, so the validator repeats the rule that decides whether one + forward-kernel call convolves every head or one rank's heads. Over every GDN context-parallel mode, degrees 1-8 on + each context axis, and head counts and head dims that model_runner runs, the validator must accept the key exactly + when the width one forward-kernel call convolves, as model_runner shards it, is a multiple of 1024. + """ + seen = collections.Counter() + for cp_mode, cp, axis, (nk, nv), head_dim in itertools.product( + ("auto", "head", "seq"), + range(1, 9), + ("ici_context_parallelism", "ici_context_usp_ulysses_parallelism", "dcn_context_parallelism"), + ((1, 2), (2, 4), (3, 6), (4, 8), (6, 12), (8, 16), (16, 32), (16, 64), (32, 64)), + (128, 256), + ): + fields = { + axis: cp, + "gdn_cp_mode": cp_mode, + "gdn_num_key_heads": nk, + "gdn_num_value_heads": nv, + "gdn_key_head_dim": head_dim, + "gdn_value_head_dim": head_dim, + "max_target_length": 840 * _CHUNK, # a whole number of chunks per rank at every degree 1-8 + } + try: + head_sharded, _ = model_runner._resolve_gdn_cp_mode(_validator_cfg(**fields), cp, nk) + except ValueError: + continue # model_runner refuses this config itself + ranks = cp if head_sharded else 1 + width = 2 * (nk // ranks) * head_dim + (nv // ranks) * head_dim + try: + config_types.MaxTextConfig.validate_gdn_config( + _validator_cfg(gdn_conv1d_dense_token_stack=True, **fields), is_rl_config=False + ) + accepted = True + except ValueError as e: + self.assertIn("gdn_conv1d_dense_token_stack=True requires the width", str(e)) + accepted = False + self.assertEqual(accepted, width % 1024 == 0, f"{fields}: model_runner head_sharded={head_sharded}") + seen[(head_sharded, accepted)] += 1 + # Head- and sequence-sharded configs, each both accepted and rejected: the comparison can fail either way. + self.assertLen(seen, 4, seen) + + @parameterized.named_parameters( + ("no_gdn_layers", {"decoder_block": config_types.DecoderBlockType.LLAMA2}, "a decoder with Gated DeltaNet"), + ("no_gdn_kernel", {"use_gdn_kernel": False}, "use_gdn_kernel=True"), + ("width_512", {"gdn_num_key_heads": 1, "gdn_num_value_heads": 2}, r"the width .* \(the whole Conv1D: .* = 512\)"), + ( + "head_cp8_rank_width_1536", + {"ici_context_parallelism": 8, "gdn_cp_mode": "head"}, + r"the width .* \(one rank's heads under head-sharded GDN context parallelism of degree 8: .* = 1536\)", + ), + ( + "auto_cp2_rank_width_512", + {"ici_context_parallelism": 2, "gdn_num_key_heads": 2, "gdn_num_value_heads": 4}, + r"the width .* \(one rank's heads under head-sharded GDN context parallelism of degree 2: .* = 512\)", + ), + ) + def test_rejected(self, overrides, message): + cfg = _validator_cfg(gdn_conv1d_dense_token_stack=True, **overrides) + with self.assertRaisesRegex(ValueError, f"gdn_conv1d_dense_token_stack=True requires {message}"): + config_types.MaxTextConfig.validate_gdn_config(cfg, is_rl_config=False) + # The same config with the key off passes. + config_types.MaxTextConfig.validate_gdn_config(_validator_cfg(**overrides), is_rl_config=False) + + def test_pyconfig(self): + def cfg(**kwargs): + return pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + run_name="gdn_conv1d_dense_token_stack", + enable_checkpointing=False, + log_config=False, + skip_jax_distributed_system=True, + **kwargs, + ) + + self.assertFalse(cfg().gdn_conv1d_dense_token_stack) + self.assertFalse(config_types.Qwen3Next().gdn_conv1d_dense_token_stack) + on = cfg(model_name="qwen3.5-397b-a17b", use_gdn_kernel=True, gdn_conv1d_dense_token_stack=True) + self.assertTrue(on.gdn_conv1d_dense_token_stack) + with self.assertRaisesRegex(ValueError, "requires a decoder with Gated DeltaNet layers"): + cfg(gdn_conv1d_dense_token_stack=True) + + +@pytest.mark.cpu_only +def test_programs_in_a_child_process(): + """Runs `main` below in a fresh CPU process (the program checks trace TPU kernels on this host).""" + env = os.environ.copy() + env["JAX_PLATFORMS"] = "cpu" + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + print(result.stdout) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + + +# ---------------------------------------------------------------------------------------------------------------------- +# Program checks (main). Each forward-kernel call is a pallas_call named fused_conv1d_gdn*; _stack_counts reads which +# way its kernel jaxpr assembles the Conv1D output. + + +def _forward_kernel_stacks(closed_jaxpr, dim): + """[_stack_counts] of every fused_conv1d_gdn kernel in the jaxpr, in trace order.""" + return [_stack_counts(kernel, dim) for _, kernel in _forward_kernels(closed_jaxpr)] + + +def _bf16_conv_inputs(closed_jaxpr, shape): + """bf16 pad / concatenate results of `shape` outside the Pallas kernels (gdn_conv1d_input_in_qkv_dtype's marker).""" + found = 0 + + def visit(jaxpr): + nonlocal found + for eqn in jaxpr.eqns: + if eqn.primitive.name == "pallas_call": + continue + if eqn.primitive.name in ("pad", "concatenate"): + aval = eqn.outvars[0].aval + found += tuple(aval.shape) == tuple(shape) and aval.dtype == jnp.bfloat16 + for sub in jax_core.jaxprs_in_params(eqn.params): + visit(sub) + + visit(closed_jaxpr.jaxpr) + return found + + +class _Value: + """A parameter stand-in: .value and [...] like an nnx Param.""" + + def __init__(self, value): + self.value = value + + def __getitem__(self, idx): + return self.value[idx] + + +def _trace_model_runner(mode, nk, nv, *, dense, qkv_dtype=False, remat=True, seg=True, seq=512, grad=True): + """model_runner.run_gdn_kernel_layer on a stand-in GDN layer, traced for a TPU7x target. + + grad=True traces value_and_grad, so the custom VJP's forward rule makes the forward-kernel calls; grad=False traces + the forward pass alone (as a forward-only RL pass runs it), so its primal makes them. + """ + cp = {"none": 1, "seq4": 4, "head2": 2}[mode] + if mode == "none": + ctx, mesh = _tpu_target() + layer_mesh = None + else: + ctx, mesh = _tpu_target((cp,), ("context",)) + layer_mesh = mesh + cfg = types.SimpleNamespace( + ici_context_parallelism=cp, + dcn_context_parallelism=1, + ici_context_usp_ulysses_parallelism=1, + dcn_context_usp_ulysses_parallelism=1, + gdn_cp_mode={"head2": "head", "seq4": "seq"}.get(mode, "auto"), + gdn_cp_matmul_precision="highest", + shard_mode=ShardMode.AUTO, + logical_axis_rules=[], + gdn_conv_kernel_dim=_KERNEL, + gdn_chunk_size=_CHUNK, + use_qk_norm_in_gdn=True, + context_parallel_load_balance=False, + gdn_conv1d_input_in_qkv_dtype=qkv_dtype, + gdn_conv1d_dense_token_stack=dense, + ) + key_dim, conv_dim = nk * _HEAD_DIM, (2 * nk + nv) * _HEAD_DIM + + def loss(params, x, aux): + def body(params, x): + layer = types.SimpleNamespace( + config=cfg, + mesh=layer_mesh, + num_k_heads=nk, + num_v_heads=nv, + head_k_dim=_HEAD_DIM, + head_v_dim=_HEAD_DIM, + key_dim=key_dim, + v_heads_per_k_head=nv // nk, + conv1d=types.SimpleNamespace(kernel=_Value(params["conv_w"]), bias=None), + A_log=_Value(params["a_log"]), + dt_bias=_Value(params["dt_bias"]), + ) + out, next_cs, next_rs = model_runner.run_gdn_kernel_layer( + layer=layer, + query=x["query"], + key=x["key"], + value=x["value_raw"].reshape(1, seq, nv, _HEAD_DIM), + value_raw=x["value_raw"], + b=x["b"], + a=x["a"], + active_cache=None, + decoder_segment_ids=aux["seg"] if seg else None, + model_mode="train", + state_dtype=jnp.float32, + ) + return ( + jnp.sum(out.astype(jnp.float32) * aux["do"]) + + 0.5 * jnp.sum(next_cs.astype(jnp.float32)) + + 0.25 * jnp.sum(next_rs.astype(jnp.float32)) + ) + + if remat: + body = jax.checkpoint(body, policy=jax.checkpoint_policies.nothing_saveable) + return body(params, x) + + a = jax.ShapeDtypeStruct + params = { + "conv_w": a((_KERNEL, 1, conv_dim), jnp.float32), + "a_log": a((nv,), jnp.float32), + "dt_bias": a((nv,), jnp.float32), + } + x = { + "query": a((1, seq, nk, _HEAD_DIM), jnp.bfloat16), + "key": a((1, seq, nk, _HEAD_DIM), jnp.bfloat16), + "value_raw": a((1, seq, nk, (nv // nk) * _HEAD_DIM), jnp.bfloat16), + "b": a((1, seq, nv), jnp.bfloat16), + "a": a((1, seq, nv), jnp.bfloat16), + } + aux = {"seg": a((1, seq), jnp.int32), "do": a((1, seq, nv, _HEAD_DIM), jnp.float32)} + fn = jax.value_and_grad(loss, argnums=(0, 1)) if grad else loss + with ctx: + return jax.make_jaxpr(fn)(params, x, aux) + + +def _trace_decoupled(dense, *, grad, remat, qkv_dtype=False, nk=2, nv=4, seq=128): + """gdn_decoupled_conv1d for a TPU7x target: the primal alone, or value_and_grad (forward rule), optionally remat.""" + dim = (2 * nk + nv) * _HEAD_DIM + op = functools.partial( + gdn_api.gdn_decoupled_conv1d, + num_k_heads=nk, + num_v_heads=nv, + head_k_dim=_HEAD_DIM, + head_v_dim=_HEAD_DIM, + conv_kernel_size=_KERNEL, + chunk_size=_CHUNK, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + conv_input_in_qkv_dtype=qkv_dtype, + conv_dense_token_stack=dense, + ) + + def loss(qkv, b, a, cw, alog, dtb, do): + def body(qkv, b, a, cw, alog, dtb): + out, _ = op( + qkv=qkv, + b=b, + a=a, + conv_weight=cw, + conv_bias=None, + a_log=alog, + dt_bias=dtb, + conv_state=None, + recurrent_state=None, + ) + return jnp.sum(out.astype(jnp.float32) * do) + + if remat: + body = jax.checkpoint(body, policy=jax.checkpoint_policies.nothing_saveable) + return body(qkv, b, a, cw, alog, dtb) + + a = jax.ShapeDtypeStruct + args = ( + a((1, seq, dim), jnp.bfloat16), + a((1, seq, nv), jnp.bfloat16), + a((1, seq, nv), jnp.bfloat16), + a((_KERNEL, 1, dim), jnp.float32), + a((nv,), jnp.float32), + a((nv,), jnp.float32), + a((1, seq, nv, _HEAD_DIM), jnp.float32), + ) + fn = jax.value_and_grad(loss, argnums=tuple(range(6))) if grad else loss + ctx, _ = _tpu_target() + with ctx: + return jax.make_jaxpr(fn)(*args), dim + + +def _check(label, ok, detail=""): + print(f"{label}: [{'ok' if ok else 'FAIL'}] {detail}", flush=True) + return ok + + +def _all_dense(stacks, expected_calls): + return len(stacks) == expected_calls and all(s == (1, 0) for s in stacks) + + +def _all_stacked(stacks, expected_calls): + return len(stacks) == expected_calls and all(s == (0, 1) for s in stacks) + + +# The cases of _option_reaches_every_forward_kernel whose forward-kernel calls go through the custom VJP's forward +# rule (traced under value_and_grad) or through its primal (traced forward only). A deliberately wrong version in main +# that drops the option in one of the two must break every case in its group. +_FORWARD_RULE_CASES = ( + "decoupled_forward_rule", + "decoupled_forward_rule_remat", + "runner_no_cp", + "runner_seq_cp4", + "runner_head_cp2", + "runner_seq_cp4_qwen35_397b_heads", +) +_PRIMAL_CASES = ("decoupled_primal", "runner_seq_cp4_forward_only") + + +def _option_reaches_every_forward_kernel(names=None): + """Every forward-kernel call carries the dense regrouping iff the option is on; returns (ok, {case: (ok, detail)}). + + `names` runs only those cases (the deliberately wrong versions in main use it). + """ + results = {} + + def want(name): + return names is None or name in names + + # wrapper.fused_conv1d_gdn: the per-sequence kernel (the only one a prefill-only call runs). + for nk, nv in ((2, 4), (16, 64)): + dim = (2 * nk + nv) * _HEAD_DIM + name = f"fused_conv1d_gdn_width{dim}" + if want(name): + on = _forward_kernel_stacks(_trace_fused(nk, nv, True), dim) + off = _forward_kernel_stacks(_trace_fused(nk, nv, False), dim) + results[name] = (_all_dense(on, 1) and _all_stacked(off, 1), f"on={on} off={off}") + # gdn_decoupled_conv1d: the primal (1 call), the forward rule (1), the forward rule and its remat replay (2). + for name, grad, remat, calls in ( + ("decoupled_primal", False, False, 1), + ("decoupled_forward_rule", True, False, 1), + ("decoupled_forward_rule_remat", True, True, 2), + ): + if want(name): + jaxpr_on, dim = _trace_decoupled(True, grad=grad, remat=remat) + jaxpr_off, _ = _trace_decoupled(False, grad=grad, remat=remat) + on, off = _forward_kernel_stacks(jaxpr_on, dim), _forward_kernel_stacks(jaxpr_off, dim) + results[name] = (_all_dense(on, calls) and _all_stacked(off, calls), f"on={on} off={off}") + # model_runner reading the config key, under remat. With value_and_grad (the custom VJP's forward rule): no context + # parallelism (1 call + replay), sequence-sharded CP 4 (2 passes + their replays), head-sharded CP 2 (1 call per + # rank + replay), and the production head counts (16 key, 64 value heads) under sequence-sharded CP 4. Forward + # only (the custom VJP's primal): sequence-sharded CP 4 (2 passes, no replay). + for name, mode, nk, nv, seq, grad, calls in ( + ("runner_no_cp", "none", 4, 8, 512, True, 2), + ("runner_seq_cp4", "seq4", 4, 8, 512, True, 4), + ("runner_head_cp2", "head2", 4, 8, 512, True, 2), + ("runner_seq_cp4_qwen35_397b_heads", "seq4", 16, 64, 256, True, 4), + ("runner_seq_cp4_forward_only", "seq4", 4, 8, 512, False, 2), + ): + if want(name): + dim = (2 * nk + nv) * _HEAD_DIM // (2 if mode == "head2" else 1) + on = _forward_kernel_stacks(_trace_model_runner(mode, nk, nv, dense=True, seq=seq, grad=grad), dim) + off = _forward_kernel_stacks(_trace_model_runner(mode, nk, nv, dense=False, seq=seq, grad=grad), dim) + results[name] = (_all_dense(on, calls) and _all_stacked(off, calls), f"on={on} off={off}") + return all(ok for ok, _ in results.values()), results + + +def _composes_with_input_in_qkv_dtype(): + """All four combinations: each key's own marker is present iff that key is on.""" + ok = True + for qkv_dtype in (False, True): + for dense in (False, True): + jaxpr, dim = _trace_decoupled(dense, grad=True, remat=False, qkv_dtype=qkv_dtype) + stacks = _forward_kernel_stacks(jaxpr, dim) + narrow = _bf16_conv_inputs(jaxpr, (1, 128 + _KERNEL - 1, dim)) + good = (_all_dense(stacks, 1) if dense else _all_stacked(stacks, 1)) and (narrow > 0) == qkv_dtype + ok &= _check( + f"compose qkv_dtype={qkv_dtype} dense={dense}", good, f"forward kernels={stacks} bf16 conv inputs={narrow}" + ) + return ok + + +def _backward_rule_positional_forms(): + """The backward rule traces the same program from every positional form custom_vjp or a caller may use.""" + nk, nv, seq = 2, 4, 128 + dim = (2 * nk + nv) * _HEAD_DIM + nondiff = (nk, nv, _HEAD_DIM, _HEAD_DIM, _KERNEL, _CHUNK, True, jnp.float32) + prec = cp_gdn.DEFAULT_PRECISION + forms = { + 10: lambda res, cot: gdn_api._gdn_decoupled_conv1d_bwd(*nondiff, res, cot), + 11: lambda res, cot: gdn_api._gdn_decoupled_conv1d_bwd(*nondiff, None, res, cot), + 12: lambda res, cot: gdn_api._gdn_decoupled_conv1d_bwd(*nondiff, None, prec, res, cot), + 13: lambda res, cot: gdn_api._gdn_decoupled_conv1d_bwd(*nondiff, None, prec, False, res, cot), + 14: lambda res, cot: gdn_api._gdn_decoupled_conv1d_bwd(*nondiff, None, prec, False, True, res, cot), + } + + def program(form): + def f(qkv, b, a, cw, alog, dtb, do): + _, res = gdn_api._gdn_decoupled_conv1d_fwd(qkv, b, a, cw, None, alog, dtb, None, None, *nondiff) + return forms[form](res, (do, (None, None))) + + a = jax.ShapeDtypeStruct + args = ( + a((1, seq, dim), jnp.bfloat16), + a((1, seq, nv), jnp.bfloat16), + a((1, seq, nv), jnp.bfloat16), + a((_KERNEL, 1, dim), jnp.float32), + a((nv,), jnp.float32), + a((nv,), jnp.float32), + a((1, seq, nv, _HEAD_DIM), jnp.float32), + ) + ctx, _ = _tpu_target() + with ctx: + return str(jax.make_jaxpr(f)(*args)) + + ref = program(14) + ok = True + for form in (10, 11, 12, 13): + try: + same = program(form) == ref + except Exception as e: # pylint: disable=broad-exception-caught + same = False + print(f" form {form} raised {type(e).__name__}: {e}", flush=True) + ok &= _check(f"backward rule, {form} positional arguments", same) + return ok + + +_BODY = re.compile(r"\\22body\\22: \\22[A-Za-z0-9+/=]+\\22") + + +def _lowers_to_mosaic(): + """Lowered for a TPU7x target (Pallas -> Mosaic, as an AOT compile does), the option changes the kernel body only.""" + ok = True + for nk, nv in ((2, 4), (16, 64)): + texts = {} + for dense in (False, True): + ctx, _ = _tpu_target() + fn = functools.partial( + wrapper.fused_conv1d_gdn, + n_kq=nk, + n_v=nv, + d_k=_HEAD_DIM, + d_v=_HEAD_DIM, + kernel_size=_KERNEL, + compute_precision=jnp.dtype(jnp.float32), + mixed_tile_size=_CHUNK, + is_prefill_only=True, + conv_dense_token_stack=dense, + ) + with ctx: + lowered = jax.jit(fn).trace(*_fused_args((2 * nk + nv) * _HEAD_DIM, nv)).lower(lowering_platforms=("tpu",)) + texts[dense] = lowered.as_text(debug_info=False) + names = {d: re.findall(r'kernel_name = "([^"]+)"', t) for d, t in texts.items()} + good = ( + len(names[False]) == 1 + and names[False] == names[True] + and names[False][0].startswith("fused_conv1d_gdn") + and len(_BODY.findall(texts[True])) == 1 + and texts[False] != texts[True] + and _BODY.sub("BODY", texts[False]) == _BODY.sub("BODY", texts[True]) + ) + ok &= _check(f"mosaic lowering width {(2 * nk + nv) * _HEAD_DIM}", good, f"kernels={names[True]}") + return ok + + +def _with_wrong_version(obj, attr, wrong, names): + """The results of the forward-kernel cases `names` (None: all) with `obj.attr` replaced by a wrong version.""" + real = getattr(obj, attr) + # jax.jit caches traces (fused_conv1d_gdn is jitted): clear them so the patched function is traced, and again + # afterwards so no later check reuses a trace of it. + jax.clear_caches() + setattr(obj, attr, wrong) + try: + return _option_reaches_every_forward_kernel(names)[1] + finally: + setattr(obj, attr, real) + jax.clear_caches() + + +def _ignored_option(names): + """Deliberately wrong: the conv ignores the option (always jnp.stack).""" + real = compute_conv1d.stack_conv1d_tokens + + def ignore(out_list, dense): + del dense + return real(out_list, False) + + return _with_wrong_version(compute_conv1d, "stack_conv1d_tokens", ignore, names) + + +def _option_dropped_by_forward_rule(names): + """Deliberately wrong: the custom VJP's forward rule (and so its remat replay) drops the option.""" + real_fwd = gdn_api.gdn_decoupled_conv1d.fwd + + def drop(*args): + args = list(args) + args[21] = False # conv_dense_token_stack + return real_fwd(*args) + + return _with_wrong_version(gdn_api.gdn_decoupled_conv1d, "fwd", drop, names) + + +def _option_dropped_by_primal(names): + """Deliberately wrong: the custom VJP's primal (every forward-only call, with or without CP) drops the option.""" + real_fun = gdn_api.gdn_decoupled_conv1d.fun + + @functools.wraps(real_fun) # custom_vjp binds keyword arguments against the primal's signature + def drop(*args): + args = list(args) + args[21] = False # conv_dense_token_stack + return real_fun(*args) + + return _with_wrong_version(gdn_api.gdn_decoupled_conv1d, "fun", drop, names) + + +def main(): + ok = True + reached, details = _option_reaches_every_forward_kernel() + for name, (good, detail) in details.items(): + ok &= _check(name, good, detail) + ok &= _check("option reaches every forward kernel", reached) + ok &= _lowers_to_mosaic() + ok &= _composes_with_input_in_qkv_dtype() + ok &= _backward_rule_positional_forms() + # Each deliberately wrong version must make every case it breaks FAIL. A case that passes, or that does not exist + # (a misspelt name), counts as a miss. + controls = { + "control_ignored_option": (_ignored_option, tuple(details)), + "control_option_dropped_by_forward_rule": (_option_dropped_by_forward_rule, _FORWARD_RULE_CASES), + "control_option_dropped_by_primal": (_option_dropped_by_primal, _PRIMAL_CASES), + } + for name, (wrong_version, must_break) in controls.items(): + results = wrong_version(must_break) + missed = [case for case in must_break if results.get(case, (True, ""))[0]] + detail = f"detected by all {len(must_break)} cases it breaks" if not missed else f"NOT detected by {missed}" + ok &= _check(name, not missed, detail) + if not ok: + sys.exit(1) + print(_PASS_TOKEN, flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py b/tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py new file mode 100644 index 0000000000..0735d75ffa --- /dev/null +++ b/tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py @@ -0,0 +1,304 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""gdn_conv1d_input_in_qkv_dtype: the GDN Conv1D + SiLU with its halo-extended input in qkv's dtype. + +The option computes the same float32 values in the same order (it only changes which tensors are materialized), so +op by op every output is bit-identical. Compiled, the forward outputs and the input / conv-state gradients stay +bit-identical, but the conv weight gradient is a float32 reduction over batch x sequence whose summation order XLA +may change when it fuses the tap's bf16 -> float32 cast into it (measured on CPU: a few ulp, in the configuration +that recomputes the pre-activation without packing or a conv state). The option must also actually change the +traced program; a test that only compared values would pass if the switch were ignored, so the value checks are +paired with checks that the two programs differ. +""" + +# pylint: disable=protected-access + +import types + +from absl.testing import absltest +from absl.testing import parameterized +import jax +import jax.numpy as jnp +from jax.extend import core as jax_core +import numpy as np + +from maxtext.configs import types as config_types +from maxtext.kernels.gdn import gdn_bwd_pallas +from maxtext.kernels.gdn.gdn_bwd import compute_conv1d_bwd + +_KERNEL_SIZE = 4 + + +def _bits(x): + """The array's bytes, so equality means bit-identical (NaN-safe, -0.0 != 0.0).""" + return np.asarray(jax.device_get(x)).tobytes() + + +def _conv_inputs(seed, *, batch=2, seq_len=24, dim=16, with_state, with_bias, with_segments, dtype=jnp.bfloat16): + """Seeded Conv1D + SiLU inputs: (qkv, conv_weight, conv_bias, conv_state, dy, segment_ids, conv_halo_seg).""" + keys = jax.random.split(jax.random.PRNGKey(seed), 6) + qkv = jax.random.normal(keys[0], (batch, seq_len, dim), jnp.float32).astype(dtype) + conv_weight = jax.random.normal(keys[1], (_KERNEL_SIZE, 1, dim), jnp.float32) + conv_bias = jax.random.normal(keys[2], (dim,), jnp.float32) if with_bias else None + conv_state = ( + jax.random.normal(keys[3], (batch, _KERNEL_SIZE - 1, dim), jnp.float32).astype(dtype) if with_state else None + ) + dy = jax.random.normal(keys[4], (batch, seq_len, dim), jnp.float32) + segment_ids = conv_halo_seg = None + if with_segments: + # Two packed documents and trailing padding (segment 0) per row. + row = np.concatenate([np.full(9, 1), np.full(11, 2), np.zeros(seq_len - 20)]).astype(np.int32) + segment_ids = jnp.asarray(np.stack([row] * batch)) + if with_state: + conv_halo_seg = jnp.ones((batch, _KERNEL_SIZE - 1), jnp.int32) + return qkv, conv_weight, conv_bias, conv_state, dy, segment_ids, conv_halo_seg + + +def _pad_or_concat_out_dtypes(jaxpr, shape): + """Dtypes of every pad / concatenate in `jaxpr` (sub-jaxprs included) whose output has `shape`.""" + found = [] + + def visit(jx): + for eqn in jx.eqns: + if eqn.primitive.name in ("pad", "concatenate"): + aval = eqn.outvars[0].aval + if tuple(aval.shape) == tuple(shape): + found.append(aval.dtype) + for sub in jax_core.jaxprs_in_params(eqn.params): + visit(sub) + + visit(jaxpr.jaxpr) + return found + + +_CASES = [ + {"testcase_name": f"state{s}_bias{b}_seg{g}", "with_state": s, "with_bias": b, "with_segments": g} + for s in (False, True) + for b in (False, True) + for g in (False, True) +] + + +class GdnConv1dInputInQkvDtypeTest(parameterized.TestCase): + + @parameterized.named_parameters(*_CASES) + def test_forward_bit_identical(self, with_state, with_bias, with_segments): + qkv, w, bias, state, _, seg, halo = _conv_inputs( + 1, with_state=with_state, with_bias=with_bias, with_segments=with_segments + ) + + def fwd(narrow): + return compute_conv1d_bwd.conv1d_silu_fwd( + qkv, + w, + bias, + _KERNEL_SIZE, + conv_state=state, + segment_ids=seg, + conv_halo_seg=halo, + conv_input_in_qkv_dtype=narrow, + ) + + ref_out, ref_act = jax.jit(lambda: fwd(False))() + out, act = jax.jit(lambda: fwd(True))() + self.assertEqual(ref_out.dtype, out.dtype) + self.assertEqual(ref_act.dtype, act.dtype) + self.assertEqual(_bits(ref_out), _bits(out)) + self.assertEqual(_bits(ref_act), _bits(act)) + + @parameterized.named_parameters(*_CASES) + def test_backward_bit_identical(self, with_state, with_bias, with_segments): + qkv, w, bias, state, dy, seg, halo = _conv_inputs( + 2, with_state=with_state, with_bias=with_bias, with_segments=with_segments + ) + conv_out, _ = compute_conv1d_bwd.conv1d_silu_fwd( + qkv, w, bias, _KERNEL_SIZE, conv_state=state, segment_ids=seg, conv_halo_seg=halo + ) + for cached_conv_out in (None, conv_out): + for return_d_conv_state in (False, True): + + def bwd(narrow, cached_conv_out=cached_conv_out, return_d_conv_state=return_d_conv_state): + return compute_conv1d_bwd.conv1d_silu_bwd( + qkv, + w, + bias, + dy, + _KERNEL_SIZE, + conv_out=cached_conv_out, + conv_state=state, + return_d_conv_state=return_d_conv_state, + segment_ids=seg, + conv_halo_seg=halo, + conv_input_in_qkv_dtype=narrow, + ) + + # Op by op: the same values in the same order, so every output is bit-identical. + ref, got = bwd(False), bwd(True) + self.assertEqual(len(ref), len(got)) + for r, g in zip(ref, got): + if r is None: + self.assertIsNone(g) + continue + self.assertEqual(r.dtype, g.dtype) + self.assertEqual(_bits(r), _bits(g)) + # Compiled: dx, db and d_conv_state stay bit-identical; dw (index 1) is a float32 reduction whose order XLA + # may change with the fusion, so it is compared to float32 summation-order tolerance. + ref = jax.jit(lambda bwd=bwd: bwd(False))() + got = jax.jit(lambda bwd=bwd: bwd(True))() + for i, (r, g) in enumerate(zip(ref, got)): + if r is None: + self.assertIsNone(g) + continue + self.assertEqual(r.dtype, g.dtype) + if i == 1: + np.testing.assert_allclose(np.asarray(g), np.asarray(r), rtol=2e-5, atol=1e-6) + else: + self.assertEqual(_bits(r), _bits(g)) + + @parameterized.named_parameters( + {"testcase_name": "no_state", "with_state": False}, + {"testcase_name": "state", "with_state": True}, + ) + def test_conv_input_is_narrow_only_when_enabled(self, with_state): + """Positive control: the default path builds the [B, S + K - 1, D] input in float32, the option in bf16.""" + qkv, w, bias, state, dy, _, _ = _conv_inputs(3, with_state=with_state, with_bias=True, with_segments=False) + shape = (qkv.shape[0], qkv.shape[1] + _KERNEL_SIZE - 1, qkv.shape[2]) + for fn in ("fwd", "bwd"): + + def run(narrow, fn=fn): + if fn == "fwd": + return compute_conv1d_bwd.conv1d_silu_fwd( + qkv, w, bias, _KERNEL_SIZE, conv_state=state, conv_input_in_qkv_dtype=narrow + ) + return compute_conv1d_bwd.conv1d_silu_bwd( + qkv, w, bias, dy, _KERNEL_SIZE, conv_state=state, conv_input_in_qkv_dtype=narrow + ) + + off = _pad_or_concat_out_dtypes(jax.make_jaxpr(lambda run=run: run(False))(), shape) + on = _pad_or_concat_out_dtypes(jax.make_jaxpr(lambda run=run: run(True))(), shape) + # The backward also pads dz (float32) to the same shape; only the conv input's buffer changes dtype. + self.assertNotIn(jnp.dtype(jnp.bfloat16), off, fn) + self.assertEqual(on.count(jnp.dtype(jnp.bfloat16)), 1, fn) + self.assertEqual(off.count(jnp.dtype(jnp.float32)), on.count(jnp.dtype(jnp.float32)) + 1, fn) + + def test_scope_sets_the_default(self): + """conv_input_in_qkv_dtype=None follows conv_input_in_qkv_dtype_scope(); the default is off.""" + qkv, w, bias, state, _, _, _ = _conv_inputs(4, with_state=True, with_bias=True, with_segments=False) + + def jaxpr_text(narrow): + return str( + jax.make_jaxpr( + lambda: compute_conv1d_bwd.conv1d_silu_fwd( + qkv, w, bias, _KERNEL_SIZE, conv_state=state, conv_input_in_qkv_dtype=narrow + ) + )() + ) + + off, on = jaxpr_text(False), jaxpr_text(True) + self.assertNotEqual(off, on) + self.assertEqual(jaxpr_text(None), off) + with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(True): + self.assertEqual(jaxpr_text(None), on) + self.assertEqual(jaxpr_text(False), off) # an explicit argument wins over the scope + with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(False): + self.assertEqual(jaxpr_text(None), off) + self.assertEqual(jaxpr_text(None), on) + self.assertEqual(jaxpr_text(None), off) + + def test_gdn_decoupled_conv1d_grads_bit_identical(self): + """Through the custom VJP: the same loss and gradients, and the option reaches the backward rule's trace.""" + gdn_bwd_pallas.ensure_cpu_interpret_registered() + is_tpu = jax.default_backend() == "tpu" + batch, chunk_size, num_chunks = 1, (64 if is_tpu else 32), 2 + seq_len = chunk_size * num_chunks + num_k_heads, num_v_heads = 1, 2 + head_k_dim = head_v_dim = 128 if is_tpu else 64 + dim = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + keys = jax.random.split(jax.random.PRNGKey(5), 8) + qkv = jax.random.normal(keys[0], (batch, seq_len, dim), jnp.float32).astype(jnp.bfloat16) + b = jax.random.normal(keys[1], (batch, seq_len, num_v_heads), jnp.float32) + a = jax.random.normal(keys[2], (batch, seq_len, num_v_heads), jnp.float32) + conv_weight = jax.random.normal(keys[3], (_KERNEL_SIZE, 1, dim), jnp.float32) + conv_bias = jax.random.normal(keys[4], (dim,), jnp.float32) + a_log = jax.random.normal(keys[5], (num_v_heads,), jnp.float32) + dt_bias = jax.random.normal(keys[6], (num_v_heads,), jnp.float32) + do = jax.random.normal(keys[7], (batch, seq_len, num_v_heads, head_v_dim), jnp.float32) + args = (qkv, b, a, conv_weight, conv_bias, a_log, dt_bias) + + def loss(narrow, qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = gdn_bwd_pallas.gdn_decoupled_conv1d( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=None, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=_KERNEL_SIZE, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + conv_input_in_qkv_dtype=narrow, + ) + return jnp.sum(out.astype(jnp.float32) * do) + + ref_val, ref_grads = jax.value_and_grad(lambda *x: loss(False, *x), argnums=tuple(range(7)))(*args) + val, grads = jax.value_and_grad(lambda *x: loss(True, *x), argnums=tuple(range(7)))(*args) + self.assertEqual(_bits(ref_val), _bits(val)) + for i, (r, g) in enumerate(zip(ref_grads, grads)): + self.assertEqual(r.dtype, g.dtype) + if i == 3: # conv_weight: a float32 reduction over batch x sequence (see the module docstring) + np.testing.assert_allclose(np.asarray(g), np.asarray(r), rtol=2e-5, atol=1e-6) + else: + self.assertEqual(_bits(r), _bits(g)) + + # The local forward convolves a float32 copy of qkv either way, so a bf16 [B, S + K - 1, D] conv input can only + # come from the backward rule: present with the option, absent without it. + shape = (batch, seq_len + _KERNEL_SIZE - 1, dim) + grad_off = jax.make_jaxpr(jax.grad(lambda *x: loss(False, *x), argnums=tuple(range(7))))(*args) + grad_on = jax.make_jaxpr(jax.grad(lambda *x: loss(True, *x), argnums=tuple(range(7))))(*args) + self.assertNotIn(jnp.dtype(jnp.bfloat16), _pad_or_concat_out_dtypes(grad_off, shape)) + self.assertIn(jnp.dtype(jnp.bfloat16), _pad_or_concat_out_dtypes(grad_on, shape)) + + def test_requires_gdn_kernel(self): + def cfg(**overrides): + fields = { + "decoder_block": config_types.DecoderBlockType.QWEN3_5, + "ici_context_parallelism": 1, + "dcn_context_parallelism": 1, + "ici_context_usp_ulysses_parallelism": 1, + "dcn_context_usp_ulysses_parallelism": 1, + "use_gdn_kernel": False, + "gdn_conv1d_input_in_qkv_dtype": False, + } + fields.update(overrides) + return types.SimpleNamespace(**fields) + + config_types.MaxTextConfig.validate_gdn_config(cfg(), is_rl_config=False) + with self.assertRaisesRegex(ValueError, "gdn_conv1d_input_in_qkv_dtype=True requires use_gdn_kernel=True"): + config_types.MaxTextConfig.validate_gdn_config(cfg(gdn_conv1d_input_in_qkv_dtype=True), is_rl_config=False) + + def test_config_default_is_off(self): + self.assertIs(config_types.Qwen3Next.model_fields["gdn_conv1d_input_in_qkv_dtype"].default, False) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/gdn_cp_head_sharded_test.py b/tests/unit/gdn_cp_head_sharded_test.py index c6d19b09d8..10681602f6 100644 --- a/tests/unit/gdn_cp_head_sharded_test.py +++ b/tests/unit/gdn_cp_head_sharded_test.py @@ -69,6 +69,7 @@ def create_gdn_config( use_gdn_kernel=True, enable_gdn_sequence_packing=enable_gdn_sequence_packing, gdn_cp_mode=gdn_cp_mode, + gdn_conv1d_input_in_qkv_dtype=False, load_balance_loss_weight=0.0, scan_layers=False, using_pipeline_parallelism=False, diff --git a/tests/unit/gdn_cp_test.py b/tests/unit/gdn_cp_test.py index 94394fcbd0..851578d63a 100644 --- a/tests/unit/gdn_cp_test.py +++ b/tests/unit/gdn_cp_test.py @@ -372,7 +372,7 @@ def test_cp_gdn_kernel_module_primitives(self): np.testing.assert_allclose(m_loc_konly, m_loc, rtol=1e-6, atol=1e-6) # Independently verify s_ext (initial_state=0) and m_loc @ h0 + s_ext (initial_state=h0) - # against pure_jax_decoupled_conv1d_gdn so neither m_loc nor s_ext is compared with itself. + # against pure_jax_decoupled_conv1d_gdn so neither m_loc nor s_ext is compared with itself (C6). _, (_, s_ext_ref) = gdn_bwd_pallas.pure_jax_decoupled_conv1d_gdn( qkv=qkv_raw, b=b, @@ -448,6 +448,240 @@ def test_cp_gdn_kernel_module_primitives(self): np.testing.assert_allclose(dm_uncached, dm_cached, rtol=1e-5, atol=1e-5) np.testing.assert_allclose(ds_uncached, ds_cached, rtol=1e-5, atol=1e-5) + # Verify on-the-fly t_inv recomputation (t_inv=None) used at chunk_size=128 matches explicit t_inv. + m_loc_no_tinv, s_ext_no_tinv = cp_gdn.compose_local_from_t_inv( + qkv_conv=qkv_conv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + t_inv=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + np.testing.assert_allclose(m_loc_no_tinv, m_loc, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(s_ext_no_tinv, s_ext, rtol=1e-5, atol=1e-5) + + dm_no_tinv, ds_no_tinv = cp_gdn.compose_bwd_local_from_t_inv( + qkv_conv=qkv_conv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + t_inv=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + np.testing.assert_allclose(dm_no_tinv, dm_uncached, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(ds_no_tinv, ds_uncached, rtol=1e-5, atol=1e-5) + + # Verify on-the-fly t_inv recomputation with packed segment_ids (including intra-chunk reset + padding). + seg = jnp.ones((batch, seq_len), dtype=jnp.int32) + seg = seg.at[0, :3].set(0) + seg = seg.at[0, chunk_size + 5 : 2 * chunk_size + 4].set(2) + seg = seg.at[0, -4:].set(0) + init_seg = jnp.array([1], dtype=jnp.int32) + qkv_conv_seg, _, t_inv_seg = gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv_raw, + b=b, + a=a, + conv_weight=cw_id, + conv_bias=None, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=4, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + segment_ids=seg, + init_seg=init_seg, + ) + m_seg_ref, s_seg_ref = cp_gdn.compose_local_from_t_inv( + qkv_conv=qkv_conv_seg, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + t_inv=t_inv_seg, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + segment_ids=seg, + init_seg=init_seg, + ) + m_seg_no_tinv, s_seg_no_tinv = cp_gdn.compose_local_from_t_inv( + qkv_conv=qkv_conv_seg, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + t_inv=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + segment_ids=seg, + init_seg=init_seg, + ) + np.testing.assert_allclose(m_seg_no_tinv, m_seg_ref, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(s_seg_no_tinv, s_seg_ref, rtol=1e-5, atol=1e-5) + + dm_seg_ref, ds_seg_ref = cp_gdn.compose_bwd_local_from_t_inv( + qkv_conv=qkv_conv_seg, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + t_inv=t_inv_seg, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + segment_ids=seg, + init_seg=init_seg, + ) + dm_seg_no_tinv, ds_seg_no_tinv = cp_gdn.compose_bwd_local_from_t_inv( + qkv_conv=qkv_conv_seg, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + t_inv=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + segment_ids=seg, + init_seg=init_seg, + ) + np.testing.assert_allclose(dm_seg_no_tinv, dm_seg_ref, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(ds_seg_no_tinv, ds_seg_ref, rtol=1e-5, atol=1e-5) + + # 2. Verify chunk_size=128 with t_inv=None and packed segment_ids. + batch, chunk_size, num_chunks = 1, 128, 2 + seq_len = num_chunks * chunk_size + num_k_heads, num_v_heads = 2, 4 + head_k_dim, head_v_dim = 32, 32 + dim_size = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + k0, k1, k2, k3, k4, k5 = jax.random.split(jax.random.PRNGKey(228), 6) + qkv_raw = jax.random.normal(k0, (batch, seq_len, dim_size), jnp.float32) * 0.1 + b = jax.random.normal(k1, (batch, seq_len, num_v_heads), jnp.float32) * 0.2 + a = jax.random.normal(k2, (batch, seq_len, num_v_heads), jnp.float32) * 0.2 + a_log = jnp.log(jax.random.uniform(k3, (num_v_heads,), minval=0.1, maxval=1.0)) + dt_bias = jax.random.normal(k4, (num_v_heads,), jnp.float32) * 0.1 + do = jax.random.normal(k5, (batch, seq_len, num_v_heads, head_v_dim), jnp.float32) * 0.1 + seg128 = jnp.ones((batch, seq_len), dtype=jnp.int32).at[0, chunk_size + 12 :].set(2) + cw_id = jnp.zeros((4, 1, dim_size), jnp.float32).at[3, 0, :].set(1.0) + qkv_conv128, _, t_inv128 = gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv_raw, + b=b, + a=a, + conv_weight=cw_id, + conv_bias=None, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=4, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + segment_ids=seg128, + ) + m128_ref, s128_ref = cp_gdn.compose_local_from_t_inv( + qkv_conv=qkv_conv128, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + t_inv=t_inv128, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + segment_ids=seg128, + ) + m128_got, s128_got = cp_gdn.compose_local_from_t_inv( + qkv_conv=qkv_conv128, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + t_inv=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + segment_ids=seg128, + ) + np.testing.assert_allclose(m128_got, m128_ref, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(s128_got, s128_ref, rtol=1e-5, atol=1e-5) + + dm128_ref, ds128_ref = cp_gdn.compose_bwd_local_from_t_inv( + qkv_conv=qkv_conv128, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + t_inv=t_inv128, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + segment_ids=seg128, + ) + dm128_got, ds128_got = cp_gdn.compose_bwd_local_from_t_inv( + qkv_conv=qkv_conv128, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + t_inv=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + segment_ids=seg128, + ) + np.testing.assert_allclose(dm128_got, dm128_ref, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(ds128_got, ds128_ref, rtol=1e-5, atol=1e-5) + if __name__ == "__main__": absltest.main() diff --git a/tests/unit/gdn_kernel_segment_ids_test.py b/tests/unit/gdn_kernel_segment_ids_test.py new file mode 100644 index 0000000000..daa1b19aee --- /dev/null +++ b/tests/unit/gdn_kernel_segment_ids_test.py @@ -0,0 +1,131 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests which segment IDs the GDN layer hands to the Pallas kernel.""" + +import unittest +from unittest import mock + +from flax import nnx +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +import numpy as np + +from maxtext.common.common_types import MODEL_MODE_PREFILL, MODEL_MODE_TRAIN +from maxtext.configs import pyconfig +from maxtext.kernels.gdn import gdn_bwd_pallas +from maxtext.kernels.gdn import model_runner +from maxtext.models import qwen3 +from tests.utils.test_helpers import get_test_config_path + +_BATCH, _SEQ_LEN = 2, 128 + + +def _gdn_kernel_config(packing): + """A minimal Qwen3-Next config on the GDN Pallas kernel path, sized for one CPU device.""" + return pyconfig.initialize( + [ + None, + get_test_config_path(), + "run_name=gdn_kernel_segment_ids_test", + "dtype=float32", + "weight_dtype=float32", + "decoder_block=qwen3_next", + "attention=dot_product", + "base_emb_dim=64", + "base_num_query_heads=2", + "base_num_kv_heads=2", + "head_dim=32", + "gdn_num_key_heads=1", + "gdn_num_value_heads=2", + "gdn_key_head_dim=128", + "gdn_value_head_dim=128", + "gdn_conv_kernel_dim=4", + "gdn_chunk_size=64", + f"max_target_length={_SEQ_LEN}", + "use_gdn_kernel=True", + f"packing={packing}", + ], + skip_jax_distributed_system=True, + ) + + +def _layer(packing, model_mode=MODEL_MODE_TRAIN): + cfg = _gdn_kernel_config(packing) + mesh = Mesh(np.array(jax.devices()[:1]).reshape([1] * len(cfg.mesh_axes)), cfg.mesh_axes) + return qwen3.Qwen3NextGatedDeltaNet( + config=cfg, + mesh=mesh, + model_mode=model_mode, + inputs_shape=(_BATCH, _SEQ_LEN, cfg.emb_dim), + rngs=nnx.Rngs(0), + ) + + +def _hidden_states(layer): + return jax.random.normal(jax.random.PRNGKey(0), (_BATCH, _SEQ_LEN, layer.config.emb_dim), jnp.float32) + + +def _run(layer, segment_ids, model_mode=MODEL_MODE_TRAIN): + """Runs the layer, returning its output and the segment IDs the kernel was called with.""" + with mock.patch.object(model_runner, "gdn_decoupled_conv1d", wraps=model_runner.gdn_decoupled_conv1d) as kernel: + out, _ = layer(_hidden_states(layer), model_mode=model_mode, decoder_segment_ids=segment_ids) + kernel.assert_called_once() + return np.asarray(out), kernel.call_args.kwargs["segment_ids"] + + +def _unpacked_segment_ids(): + """One sequence per row: row 0 is left- and right-padded, row 1 right-padded.""" + seg = np.zeros((_BATCH, _SEQ_LEN), np.int32) + seg[0, 8:108] = 1 + seg[1, :96] = 1 + return jnp.asarray(seg) + + +class GdnKernelSegmentIdsTest(unittest.TestCase): + + def setUp(self): + super().setUp() + gdn_bwd_pallas.ensure_cpu_interpret_registered() + + def test_unpacked_training_skips_segment_path(self): + seg = _unpacked_segment_ids() + out_unpacked, kernel_seg_unpacked = _run(_layer(packing=False), seg) + out_packed, kernel_seg_packed = _run(_layer(packing=True), seg) + self.assertIsNone(kernel_seg_unpacked) + self.assertIsNotNone(kernel_seg_packed) + # With one sequence per row the segment path is pure overhead: same output either way, up to fp32 + # rounding (the two kernel variants order some arithmetic differently; ~6e-6 on CPU). + np.testing.assert_allclose(out_unpacked, out_packed, rtol=1e-4, atol=1e-4) + + def test_packed_training_keeps_document_boundaries(self): + seg = np.ones((_BATCH, _SEQ_LEN), np.int32) + seg[:, 50:] = 2 + layer = _layer(packing=True) + out_two_docs, kernel_seg = _run(layer, jnp.asarray(seg)) + out_one_doc, _ = _run(layer, jnp.ones_like(jnp.asarray(seg))) + self.assertIsNotNone(kernel_seg) + # The second document must start from a fresh state, so only tokens after the boundary change. + np.testing.assert_allclose(out_two_docs[:, :50], out_one_doc[:, :50], rtol=1e-5, atol=1e-5) + self.assertGreater(np.max(np.abs(out_two_docs[:, 50:] - out_one_doc[:, 50:])), 1e-3) + + def test_unpacked_prefill_keeps_segment_ids(self): + # Prefill carries its states forward, and the segment path ends the conv state at the last valid token. + _, kernel_seg = _run(_layer(packing=False, model_mode=MODEL_MODE_PREFILL), _unpacked_segment_ids(), MODEL_MODE_PREFILL) + self.assertIsNotNone(kernel_seg) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/megablox_gmm_silu_test.py b/tests/unit/megablox_gmm_silu_test.py new file mode 100644 index 0000000000..df23bc29db --- /dev/null +++ b/tests/unit/megablox_gmm_silu_test.py @@ -0,0 +1,149 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TPU tests: the SwiGLU in the gmm_v2 epilogue (megablox ops.gmm_silu, moe_wi_gmm_fused_swiglu) against the unfused GMM. + +gmm_v2(fuse_act="silu", return_pre_act=True) returns (act, gate, up) from one accumulator: gate and up must be the +unfused gmm_v2 output's halves bit for bit (they are rounded the same way), every row outside the computed groups must +be zero in all three, and act must be silu(gate) * up computed from the float32 accumulator (close to, not equal to, +the same activation of the bf16-rounded halves: that is the numerics change). The act-only call must give the same +act. ops.gmm_silu's VJP is the unfused one: its gradients must equal, value for value, the VJP of the unfused path +(ops.gmm, then silu(gate) * up on its bf16 output). Groups are skewed, with an empty group, padding rows past the +last group, boundaries inside a sublane, and an expert-parallel shard (group_offset). + +Pallas TPU kernels: these run only on a TPU. On CPU the moe.py wiring of the key is covered by +moe_wi_gmm_fused_swiglu_test.py with a jnp stand-in for the kernel. +""" + +import functools + +from absl.testing import absltest +from absl.testing import parameterized +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.kernels.megablox import ops +from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_gmm_kernel as gmm_v2 + +pytestmark = pytest.mark.tpu_only # Pallas TPU kernels. + +M, K, N = 1024, 512, 512 # rhs is [groups, K, 2 * N] = [gate | up] +GROUP_SIZES = np.array([300, 0, 517, 133], np.int32) # sum 950 < M: padding rows; group ends inside sublanes +TILING = (256, 512, 1024, 256, 512, 512, 256, 512, 512) # gmm's 9-tuple; gmm_silu halves the forward tile_n + + +def _rand(key, shape, scale=1.0): + return (jax.random.normal(key, shape, jnp.float32) * scale).astype(jnp.bfloat16) + + +def _reference_accumulator(lhs, rhs, group_sizes, group_offset): + """f32 lhs @ rhs[expert of row] for the rows of the local groups [group_offset, group_offset + groups), else 0.""" + ends = jnp.cumsum(group_sizes) + local = jnp.sum(jnp.arange(lhs.shape[0])[:, None] >= ends[None, :], axis=1) - group_offset + out = jnp.zeros((lhs.shape[0], rhs.shape[-1]), jnp.float32) + for e in range(rhs.shape[0]): + out = out + jnp.where((local == e)[:, None], lhs.astype(jnp.float32) @ rhs[e].astype(jnp.float32), 0.0) + return out + + +def _f32(x): + return np.asarray(x.astype(jnp.float32)) + + +class GmmV2PreActTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + self.assertEqual(jax.devices()[0].platform, "tpu") + + @parameterized.product(group_offset=[0, 2], tile_k=[256, 512]) + def test_pre_act_halves_are_the_unfused_output(self, group_offset, tile_k): + k0, k1 = jax.random.split(jax.random.key(group_offset * 7 + tile_k)) + lhs = _rand(k0, (M, K)) + rhs = _rand(k1, (len(GROUP_SIZES) - group_offset, K, 2 * N), 0.05) + gs, go = jnp.asarray(GROUP_SIZES), jnp.array([group_offset], jnp.int32) + unfused = jax.jit( + functools.partial( + gmm_v2.gmm_v2, tile_info=gmm_v2.TileSizes(256, tile_k, 512), preferred_element_type=jnp.bfloat16 + ) + )(lhs, rhs, gs, group_offset=go) + fused = functools.partial( + gmm_v2.gmm_v2, tile_info=gmm_v2.TileSizes(256, tile_k, 256), preferred_element_type=jnp.bfloat16, fuse_act="silu" + ) + act, gate, up = jax.jit(functools.partial(fused, return_pre_act=True))(lhs, rhs, gs, group_offset=go) + act_only = jax.jit(fused)(lhs, rhs, gs, group_offset=go) + + np.testing.assert_array_equal(_f32(gate), _f32(unfused[:, :N])) + np.testing.assert_array_equal(_f32(up), _f32(unfused[:, N:])) + np.testing.assert_array_equal(_f32(act_only), _f32(act)) + # Rows outside the local groups: zero in every output. + computed = np.zeros(M, bool) + start = int(np.sum(GROUP_SIZES[:group_offset])) + computed[start : start + int(np.sum(GROUP_SIZES[group_offset:]))] = True + for out in (act, gate, up): + self.assertTrue(np.all(_f32(out)[~computed] == 0)) + # act is silu on the f32 accumulator, rounded once to bf16. + acc = _reference_accumulator(lhs, rhs, gs, group_offset) + want = (jax.nn.silu(acc[:, :N]) * acc[:, N:]).astype(jnp.bfloat16) + np.testing.assert_allclose(_f32(act), _f32(want), rtol=1e-2, atol=1e-3) + + +class GmmSiluTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + self.assertEqual(jax.devices()[0].platform, "tpu") + + @parameterized.parameters(0, 2) + def test_forward_and_vjp_match_the_unfused_path(self, group_offset): + k0, k1, k2 = jax.random.split(jax.random.key(group_offset), 3) + lhs = _rand(k0, (M, K)) + rhs = _rand(k1, (len(GROUP_SIZES) - group_offset, K, 2 * N), 0.05) + cot = _rand(k2, (M, N)) + gs, go = jnp.asarray(GROUP_SIZES), jnp.array([group_offset], jnp.int32) + + def unfused(lhs, rhs): + out = ops.gmm( + lhs, + rhs, + gs, + preferred_element_type=jnp.bfloat16, + tiling=TILING, + group_offset=go, + use_tokamax_backend=True, + use_gmm_v2=True, + ) + return (jax.nn.silu(out[:, :N]) * out[:, N:]).astype(jnp.bfloat16) + + def fused(lhs, rhs): + return ops.gmm_silu(lhs, rhs, gs, preferred_element_type=jnp.bfloat16, tiling=TILING, group_offset=go) + + @jax.jit + def run(lhs, rhs, cot): + y_u, vjp_u = jax.vjp(unfused, lhs, rhs) + y_f, vjp_f = jax.vjp(fused, lhs, rhs) + return y_u, vjp_u(cot), y_f, vjp_f(cot) + + y_u, (dl_u, dr_u), y_f, (dl_f, dr_f) = run(lhs, rhs, cot) + # Forward: the f32-accumulator activation, close to the unfused one. + np.testing.assert_allclose(_f32(y_f), _f32(y_u), rtol=2e-2, atol=2e-3) + # Backward: the unfused backward on the same bf16 gate and up, so the same values. + np.testing.assert_array_equal(_f32(dl_f), _f32(dl_u)) + np.testing.assert_array_equal(_f32(dr_f), _f32(dr_u)) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/moe_glue_rewrites_test.py b/tests/unit/moe_glue_rewrites_test.py new file mode 100644 index 0000000000..93ae001d70 --- /dev/null +++ b/tests/unit/moe_glue_rewrites_test.py @@ -0,0 +1,602 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Three exact rewrites of MoE glue ops, each against the plain program, on CPU. + + - moe_unsplit_prefused_wi: the prefused wi = [w0 | w1] stays one tensor through the fp8 fake-quant, the FSDP + all-gather and the prefetch barrier into the up-projection GEMM (no split at the layer entry, no concatenate). + - moe_ragged_sort_unique_pair_argsort: the ring-of-experts ragged sort's argsorts as unstable two-operand sorts of + unique (key, index) pairs (`argsort_unique_pairs`) instead of stable argsorts. + - moe_prefused_swiglu_vjp: the activation on the whole [gate | up] output as one custom-VJP op + (`moe._prefused_ffn_activation`) whose backward writes d[gate | up] as one array. + +None changes a value, so everything must be bit-identical. + +What a CPU run can and cannot reach: the SparseCore gather-reduce (and with it the validity sort that +moe_ragged_sort_unique_pair_argsort also rewrites) runs only on a TPU target; on CPU `ragged_gather_reduce` takes its +jnp reference fallback, which does not sort. `PreprocessTest` therefore calls `_preprocess` directly, the pure-JAX +part of that path. The ragged sort's own argsorts, the weight layout and the activation are plain JAX and run in the +layer below as on TPU (the GMM is megablox in interpret mode; tokamax gmm v2 is TPU-only). + +In-process checks: `argsort_unique_pairs` equals the stable argsort (ties, booleans, every axis, the index dtype); +`_preprocess` with and without it agree; `_prefused_ffn_activation` gives the split activation's value bit for bit and +its VJP value for value (SwiGLU and the clipped GPT-OSS form, bfloat16 and float32). A -0.0 cotangent is where the +two can differ in bits: XLA:CPU turns it into +0.0 on the split path (an add of zero pads) and the fused concatenate +keeps it, so the gradient bits may differ only in the sign of exact zeros, which the test pins down. The validators +reject each key for each requirement it does not meet, naming that requirement; that includes moe_unsplit_prefused_wi +with tensor parallelism, which shards the expert MLP dimension. + +The layer check runs as a subprocess with 8 CPU devices: the shrunk Qwen3.5 routed MoE of +`moe_roe_local_routing_test.py` with prefuse_moe_weights (16 experts, top 4, cp-as-ep, ragged sort, local routing, +two token chunks). Pre-registered: each key alone and all three together give an output, a loss and every gradient +bit-identical to all three off, in the plain and the pipelined chunk loop, with fp8 fake-quant and padding excluded +from routing, with forced ids and a dropless buffer, and with moe_routing_index_math="count". Two-sided: + - Structure, so a switch that silently does nothing, or does only part of its job, cannot pass: + moe_unsplit_prefused_wi traces fewer concatenates (and with fake-quant fewer barriers); + moe_ragged_sort_unique_pair_argsort traces exactly two two-key sorts per token chunk (the permutation and its + inverse) and the plain program none, and every ragged_gather_reduce call gets unique_pair_argsort=True (a spy: + CPU runs the reference fallback, which does not sort); moe_prefused_swiglu_vjp traces one more custom_vjp call + per token chunk. + - The trace-time guard behind the validator: on a mesh with tensor parallelism 2 (default axis rules, expert 4 x + tensor 2), a config that claims moe_unsplit_prefused_wi past the validator must make the layer raise rather + than compute with a sharded [w0 | w1] (it would give the wrong output), while the plain layer runs. + - The bench's MoE settings (expert-weight prefetch, row-major reduce-scatter, routing maps saved on device, + ragged_buffer_factor 2.75, fp8 fake-quant, padding excluded) with all three keys in one case. + - Deliberately wrong versions must FAIL: a fused activation whose backward swaps d(gate) and d(up) (the output + is unchanged, the gradients are not), and a pair argsort that breaks ties by decreasing index, under a buffer + that drops tokens (the tie order decides which tokens a full buffer drops, so the output changes). +""" + +import contextlib +import os +import re +import subprocess +import sys +from unittest import mock + +from absl.testing import parameterized +from flax import nnx +from flax.linen import partitioning as nn_partitioning +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +import numpy as np +import pytest + +from maxtext.configs import pyconfig +from maxtext.kernels.ragged import ragged_gather_reduce_v2 +from maxtext.kernels.ragged import ragged_sort +from maxtext.layers import moe +from maxtext.layers.initializers import nd_dense_init +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + +_NUM_DEVICES = 8 +_NUM_EXPERTS = 16 +_TOP_K = 4 +_EMB = 256 +_SEQ = 64 +_PASS_TOKEN = "MOE_GLUE_REWRITES_CHECKS_PASSED" +_XLA_FLAGS = f"--xla_force_host_platform_device_count={_NUM_DEVICES} --xla_allow_excess_precision=false" +_KEYS = ("moe_unsplit_prefused_wi", "moe_ragged_sort_unique_pair_argsort", "moe_prefused_swiglu_vjp") + + +@pytest.mark.cpu_only +class ArgsortUniquePairsTest(parameterized.TestCase): + """`argsort_unique_pairs` is the stable argsort.""" + + @parameterized.named_parameters( + ("int_ties_1d", (257,), jnp.int32, -1), + ("int_ties_2d_last", (5, 96), jnp.int32, -1), + ("int_ties_2d_first", (96, 5), jnp.int32, 0), + ("bool_2d", (8, 100), jnp.bool_, 1), + ) + def test_matches_stable_argsort(self, shape, dtype, axis): + keys = jax.random.randint(jax.random.PRNGKey(0), shape, 0, 7).astype(dtype) # many ties + want = jnp.argsort(keys, axis=axis, stable=True) + got = ragged_gather_reduce_v2.argsort_unique_pairs(keys, axis=axis) + self.assertEqual(got.dtype, want.dtype) + np.testing.assert_array_equal(np.asarray(got), np.asarray(want)) + + def test_wrong_tie_order_is_detected(self): + """Two-sided: the same comparison fails for a pair sort that breaks ties by decreasing index.""" + keys = jnp.array([1, 0, 1, 0, 1], jnp.int32) + wrong = _argsort_ties_reversed(keys) + self.assertFalse(np.array_equal(np.asarray(wrong), np.asarray(jnp.argsort(keys, stable=True)))) + + +@pytest.mark.cpu_only +class PreprocessTest(parameterized.TestCase): + """The SparseCore gather-reduce's validity sort with and without the pair argsort. + + Only the SparseCore path runs `_preprocess`; on a CPU target `ragged_gather_reduce` takes its reference fallback, + so this pure-JAX function is called directly. + """ + + @parameterized.named_parameters(("one_partition", 96, 4, 1, 16, 16), ("v7x_like", 1024, 8, 16, 16, 48)) + def test_unique_pair_argsort_matches_argsort(self, rows, group, partitions, lanes, chunk): + for density in (0.0, 0.4, 1.0): + mask = jax.random.uniform(jax.random.PRNGKey(rows), (rows,)) < density + want = ragged_gather_reduce_v2._preprocess(mask, group, partitions, lanes, chunk) # pylint: disable=protected-access + got = ragged_gather_reduce_v2._preprocess( # pylint: disable=protected-access + mask, group, partitions, lanes, chunk, unique_pair_argsort=True + ) + for a, b in zip(want, got): + self.assertEqual(a.dtype, b.dtype) + np.testing.assert_array_equal(np.asarray(b), np.asarray(a)) + + +def _swiglu(gate, up): + return (jax.nn.silu(gate) * up).astype(gate.dtype) + + +def _gpt_oss(gate, up, limit=7.0): + gate = jnp.clip(gate, max=limit) + up = jnp.clip(up, -limit, limit) + return (gate * jax.nn.sigmoid(gate * 1.702) * (up + 1)).astype(gate.dtype) + + +@pytest.mark.cpu_only +class PrefusedFfnActivationTest(parameterized.TestCase): + """`_prefused_ffn_activation` gives the split activation's value bit for bit and its VJP value for value.""" + + @parameterized.product(act=("swiglu", "gpt_oss"), dtype=("bfloat16", "float32")) + def test_value_and_vjp_match_split(self, act, dtype): + act_fn = {"swiglu": _swiglu, "gpt_oss": _gpt_oss}[act] + h01 = (jax.random.normal(jax.random.PRNGKey(1), (64, 2 * 48), jnp.float32) * 4).astype(dtype) + cot = jax.random.normal(jax.random.PRNGKey(2), (64, 48), jnp.float32).astype(dtype) + cot = cot.at[0].set(-0.0) # -0.0 cotangents must come back with the sign the split path gives them + + def split(h): + n = h.shape[-1] // 2 + return act_fn(h[:, :n], h[:, n:]) + + @jax.jit + def both(h, c): + y_split, vjp_split = jax.vjp(split, h) + y_fused, vjp_fused = jax.vjp(lambda h_: moe._prefused_ffn_activation(act_fn, h_), h) # pylint: disable=protected-access + return y_split, vjp_split(c)[0], y_fused, vjp_fused(c)[0] + + y_split, d_split, y_fused, d_fused = both(h01, cot) + np.testing.assert_array_equal(np.asarray(y_fused).view(np.uint8), np.asarray(y_split).view(np.uint8)) + d_split, d_fused = np.asarray(d_split), np.asarray(d_fused) + # Equal values everywhere. Bits may differ only in the sign of an exact zero: the split path's gradient is the + # transpose of two slices, add(pad(d_gate), pad(d_up)), and XLA:CPU computes -0.0 + 0.0 = +0.0 there, while the + # fused backward's concatenate keeps -0.0. + np.testing.assert_array_equal(d_fused, d_split) + bits_differ = d_fused.view(np.uint16 if d_fused.itemsize == 2 else np.uint32) != d_split.view( + np.uint16 if d_split.itemsize == 2 else np.uint32 + ) + # Zeros come from the -0.0 cotangent row and, in the clipped GPT-OSS form, from clipped entries (derivative 0). + self.assertTrue(np.all(d_fused[bits_differ] == 0) and np.all(d_split[bits_differ] == 0)) + + +# A mesh with tensor parallelism: the default logical axis rules (cp-as-ep has no tensor axis), expert 4 x tensor 2. +_TP_MESH = { + "custom_mesh_and_rule": "", + "ici_context_parallelism": 1, + "ici_expert_parallelism": 4, + "ici_tensor_parallelism": 2, + "per_device_batch_size": 0.5, +} +_OFF_PREFUSED_PATH = ( + ("prefuse", {"prefuse_moe_weights": False}, "prefuse_moe_weights=True"), + ("dense", {"sparse_matmul": False}, "sparse_matmul=True"), + ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"), + ("vllm_rpa", {"attention": "vllm_rpa"}, "an attention other than vllm_rpa*"), +) +_OFF_RING_RAGGED_SORT = ( + ("dense", {"sparse_matmul": False}, "sparse_matmul=True"), + ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"), + ( + "no_ring", + {"use_ring_of_experts": False, "num_moe_token_chunks": 1, "ring_of_experts_local_routing": False}, + "use_ring_of_experts=True", + ), + ("no_ragged_sort", {"use_ragged_sort": False}, "use_ragged_sort=True"), +) +_REJECTED = ( + [(f"pair_argsort_{name}", "moe_ragged_sort_unique_pair_argsort", o, m) for name, o, m in _OFF_RING_RAGGED_SORT] + + [ + ( + "pair_argsort_with_count", + "moe_ragged_sort_unique_pair_argsort", + {"moe_routing_index_math": "count"}, + "moe_routing_index_math='sort'", + ) + ] + + [(f"unsplit_wi_{name}", "moe_unsplit_prefused_wi", o, m) for name, o, m in _OFF_PREFUSED_PATH] + + [ + (f"unsplit_wi_{name}", "moe_unsplit_prefused_wi", o, m) + for name, o, m in ( + ("quantization", {"quantization": "int8"}, "no weight quantization"), + ("qwix", {"use_qwix_quantization": True}, "no weight quantization"), + ( + "two_stage_ag", + {"moe_fsdp_use_two_stage_all_gather": True}, + "moe_fsdp_use_two_stage_all_gather=False", + ), + ("flat_fsdp", {"moe_flat_fsdp_weights": True}, "moe_flat_fsdp_weights=False"), + ("shard_exp_on_fsdp", {"shard_exp_on_fsdp": True}, "shard_exp_on_fsdp=False"), + ("2d_fsdp", {"use_2d_fsdp_sharding": True}, "use_2d_fsdp_sharding=False"), + ("shard_embed", {"shard_embed_moe_on_fsdp": True}, "shard_embed_moe_on_fsdp=False"), + ("chunk_wgrad", {"moe_accumulate_chunk_wgrad": True}, "moe_accumulate_chunk_wgrad=False"), + ("emb_chunks", {"num_moe_emb_chunks": 2}, "num_moe_emb_chunks=0"), + ("tensor_parallel", _TP_MESH, "an unsharded expert MLP dimension"), + ( + "tensor_sequence_parallel", + {**_TP_MESH, "ici_tensor_parallelism": 1, "ici_tensor_sequence_parallelism": 2}, + "an unsharded expert MLP dimension", + ), + ( + "autoregressive_parallel", + {**_TP_MESH, "ici_tensor_parallelism": 1, "ici_autoregressive_parallelism": 2}, + "an unsharded expert MLP dimension", + ), + ) + ] + + [(f"swiglu_vjp_{name}", "moe_prefused_swiglu_vjp", o, m) for name, o, m in _OFF_PREFUSED_PATH] + + [ + (f"swiglu_vjp_{name}", "moe_prefused_swiglu_vjp", o, m) + for name, o, m in ( + ("mlp_bias", {"mlp_bias": True}, "mlp_bias=False"), + ("emb_chunks", {"num_moe_emb_chunks": 2}, "num_moe_emb_chunks=0"), + ("minimal_remat", {"remat_policy": "minimal"}, "remat_policy=minimal saves it as mlpwi_0"), + ("offloaded_remat", {"remat_policy": "minimal_offloaded"}, "remat_policy=minimal_offloaded saves it as"), + ("context_remat", {"remat_policy": "minimal_with_context"}, "remat_policy=minimal_with_context saves it as"), + ("flash_remat", {"remat_policy": "minimal_flash"}, "remat_policy=minimal_flash saves it as"), + ( + "quantization_remat", + {"remat_policy": "minimal_with_quantization"}, + "remat_policy=minimal_with_quantization saves it as", + ), + ( + "context_quantization_remat", + {"remat_policy": "minimal_with_context_and_quantization"}, + "remat_policy=minimal_with_context_and_quantization saves it as", + ), + ("saved_gate", {"remat_policy": "custom", "mlpwi_0": "device"}, "remat_policy=custom saves it as mlpwi_0"), + ("saved_moe_gate", {"remat_policy": "custom", "moe_mlpwi_0": "device"}, "saves it as moe_mlpwi_0"), + ("offloaded_moe_up", {"remat_policy": "custom", "moe_mlpwi_1": "offload"}, "saves it as moe_mlpwi_1"), + ) + ] +) + + +class GlueRewritesConfigTest(parameterized.TestCase): + """Each key is rejected for each requirement it does not meet, and accepted on the path it rewrites.""" + + @parameterized.named_parameters(*_REJECTED) + def test_rejected(self, key, overrides, requirement): + overrides = {"prefuse_moe_weights": True, **overrides} + with self.assertRaisesRegex(ValueError, f"{key}=True requires: .*{re.escape(requirement)}"): + _cfg(**{key: True}, **overrides) + + def test_accepted(self): + cfg = _cfg(prefuse_moe_weights=True, remat_policy="custom", **{k: True for k in _KEYS}) + self.assertTrue(all(getattr(cfg, k) for k in _KEYS)) + + def test_accepted_where_the_policy_saves_only_other_tensors(self): + """Policies that save no routed-expert gate/up tensor: full, a dense-only name, the shared expert's mlpwi_1.""" + for overrides in ({"remat_policy": "full"}, {"remat_policy": "custom", "mlpwi_1": "device", "mlpwo": "device"}): + self.assertTrue(_cfg(prefuse_moe_weights=True, moe_prefused_swiglu_vjp=True, **overrides).moe_prefused_swiglu_vjp) + + +def _cfg(**overrides): + """The shrunk Qwen3.5 routed-MoE config on the cp-as-ep mesh (moe_roe_local_routing_test.py's, local routing).""" + kwargs = { + "run_name": "moe_glue_rewrites", + "enable_checkpointing": False, + "log_config": False, + "skip_jax_distributed_system": True, + "model_name": "qwen3.5-397b-a17b", + "override_model_config": True, + "num_experts": _NUM_EXPERTS, + "num_experts_per_tok": _TOP_K, + "base_emb_dim": _EMB, + "base_mlp_dim": 256, + "base_moe_mlp_dim": 256, + "max_target_length": _SEQ, + "max_prefill_predict_length": _SEQ, + "per_device_batch_size": 0.25, + "dtype": "bfloat16", + "weight_dtype": "float32", + "float32_gate_logits": True, + "megablox": True, + "sparse_matmul": True, + # Tokamax gmm v2 is TPU-only; both programs share whichever GMM runs, so CPU uses megablox. + "use_tokamax_gmm": False, + "use_gmm_v2": False, + "use_ring_of_experts": True, + "use_ragged_sort": True, + "use_custom_sort_vjp": False, + "ragged_buffer_factor": 2.0, + "num_moe_token_chunks": 2, + "custom_mesh_and_rule": "cp-as-ep", + "ici_fsdp_parallelism": 1, + "ici_context_parallelism": 4, + "ici_expert_parallelism": 2, + "context_parallel_load_balance": False, + "ring_of_experts_local_routing": True, + } + kwargs.update(overrides) + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs) + + +def _argsort_ties_reversed(keys, axis=-1): + """Deliberately wrong: a two-operand pair sort that breaks ties by DEcreasing index.""" + dim = axis % keys.ndim + iota = jax.lax.broadcasted_iota(jnp.int32, keys.shape, dim) + _, rev = jax.lax.sort((keys, -iota), dimension=dim, is_stable=False, num_keys=2) + return -rev + + +@pytest.mark.cpu_only +def test_glue_rewrites_match_plain_program_on_cpu_mesh(): + """Runs `main` below in a child process with 8 CPU devices.""" + env = os.environ.copy() + env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" {_XLA_FLAGS}" + env["JAX_PLATFORMS"] = "cpu" + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + print(result.stdout) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + + +def _build(cfg, mesh): + return moe.RoutedMoE( + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.moe_mlp_dim, + weight_dtype=cfg.weight_dtype, + dtype=cfg.dtype, + rngs=nnx.Rngs(params=0), + ) + + +def _inputs(cfg, mode): + """The layer input, forced expert ids (padded and duplicate slots) and a token mask, per `mode`.""" + batch = int(cfg.per_device_batch_size * _NUM_DEVICES) + x = jax.random.normal(jax.random.PRNGKey(7), (batch, _SEQ, _EMB), jnp.float32) + forced = token_mask = None + if "forced" in mode: + keys = jax.random.split(jax.random.PRNGKey(11), batch * _SEQ) + ids = jax.vmap(lambda k: jax.random.permutation(k, _NUM_EXPERTS)[:_TOP_K])(keys).reshape(batch, _SEQ, _TOP_K) + pos = jnp.arange(_SEQ)[None, :, None] + ids = jnp.where((pos % 5 == 0) & (jnp.arange(_TOP_K) == _TOP_K - 1), -1, ids) + ids = jnp.where((pos % 7 == 3) & (jnp.arange(_TOP_K) == 1), ids[..., :1], ids) + forced = ids.astype(jnp.int32) + if "padmask" in mode: + lens = jnp.array([_SEQ - 5 * (b + 1) for b in range(batch)]) + token_mask = (jnp.arange(_SEQ)[None, :] < lens[:, None]).astype(jnp.int32) + return x, forced, token_mask + + +def _eqns(jaxpr): + """Every equation of a jaxpr, including those of nested jaxprs.""" + for eqn in jaxpr.eqns: + yield eqn + for param in eqn.params.values(): + for sub in param if isinstance(param, (list, tuple)) else (param,): + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + if hasattr(sub, "eqns"): + yield from _eqns(sub) + + +def _structure(jaxpr): + eqns = list(_eqns(jaxpr)) + return { + "concatenate": sum(e.primitive.name == "concatenate" for e in eqns), + "barrier": sum(e.primitive.name == "optimization_barrier" for e in eqns), + "pair_sort": sum(e.primitive.name == "sort" and e.params.get("num_keys") == 2 for e in eqns), + "custom_vjp": sum(e.primitive.name.startswith("custom_vjp_call") for e in eqns), + } + + +def _run(cfg, mesh, params, x, forced, token_mask): + """Loss, output, gradients, the structure of the traced forward, and the unique_pair_argsort flag of every + ragged_gather_reduce call traced.""" + pair_flags = [] + real_gather_reduce = ragged_sort.ragged_gather_reduce + + def spy(*args, **kwargs): + pair_flags.append(kwargs.get("unique_pair_argsort", False)) + return real_gather_reduce(*args, **kwargs) + + with ( + mock.patch.object(ragged_sort, "ragged_gather_reduce", spy), + jax.set_mesh(mesh), + nn_partitioning.axis_rules(cfg.logical_axis_rules), + ): + graphdef, _, rest = nnx.split(_build(cfg, mesh), nnx.Param, ...) + + def loss_fn(params, x): + out, lb_loss, _ = nnx.merge(graphdef, params, rest)(x, forced_routed_experts=forced, token_mask=token_mask) + loss = jnp.mean(out.astype(jnp.float32) ** 2) + if lb_loss is not None: + loss = loss + lb_loss.astype(jnp.float32) + return loss, out + + structure = _structure(jax.make_jaxpr(loss_fn)(params, x).jaxpr) + (loss, out), grads = jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(params, x) + structure["gather_reduce_pair_flags"] = sorted(set(pair_flags)) + return jax.device_get((loss, out, grads)), structure + + +def _bit_identical(label, old, new): + """Prints the per-leaf report; True iff the output, the loss and every gradient are bit-identical.""" + (loss_old, out_old, grads_old), (loss_new, out_new, grads_new) = old, new + named = [("output", out_old, out_new), ("loss", loss_old, loss_new)] + [ + (jax.tree_util.keystr(path), a, b) + for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_old)[0], jax.tree_util.tree_leaves(grads_new)) + ] + passed, lines = True, [] + for name, a, b in named: + a, b = np.asarray(a).astype(np.float32), np.asarray(b).astype(np.float32) + same = bool(np.array_equal(a, b)) and bool(np.isfinite(b).all()) and float(np.max(np.abs(a))) > 0 + passed &= same + lines.append(f" [{'ok' if same else 'FAIL'}] {name}: bit-identical={same} max|d|={np.max(np.abs(a - b)):.3e}") + print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True) + print("\n".join(lines), flush=True) + return passed + + +def _case(label, keys, overrides, mode="real", control=None): + """Runs the plain program and the one with `keys` on, on the same parameters; `control` wraps the second run.""" + overrides = {"prefuse_moe_weights": True, **overrides} + cfg_old = _cfg(**overrides) + cfg_new = _cfg(**overrides, **{k: True for k in keys}) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg_old), cfg_old.mesh_axes) + x, forced, token_mask = _inputs(cfg_old, mode) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_old.logical_axis_rules): + _, params, _ = nnx.split(_build(cfg_old, mesh), nnx.Param, ...) + old, s_old = _run(cfg_old, mesh, params, x, forced, token_mask) + with control if control is not None else contextlib.nullcontext(): + new, s_new = _run(cfg_new, mesh, params, x, forced, token_mask) + n_chunks = cfg_new.num_moe_token_chunks + want = {} + if "moe_unsplit_prefused_wi" in keys: + want["concatenate"] = s_new["concatenate"] < s_old["concatenate"] + if cfg_new.fp8_moe_fake_quant: + want["barrier"] = s_new["barrier"] < s_old["barrier"] + if "moe_ragged_sort_unique_pair_argsort" in keys: + # Both argsorts of every chunk's ragged sort (the forward traced here), and the backward's gather-reduce flag. + want["pair_sort"] = s_new["pair_sort"] == 2 * n_chunks and s_old["pair_sort"] == 0 + flags = (s_new["gather_reduce_pair_flags"], s_old["gather_reduce_pair_flags"]) + want["gather_reduce_pair_flags"] = flags == ([True], [False]) + if "moe_prefused_swiglu_vjp" in keys: + want["custom_vjp"] = s_new["custom_vjp"] == s_old["custom_vjp"] + n_chunks + structure = all(want.values()) + print(f"{label}: [{'ok' if structure else 'FAIL'}] structure plain={s_old} new={s_new} checks={want}", flush=True) + return _bit_identical(label, old, new) and structure + + +def _swapped_activation_cotangents(): + """Deliberately wrong: the fused activation's backward returns [d(up) | d(gate)].""" + real = moe._prefused_ffn_activation # pylint: disable=protected-access + + def wrong(act_fn, h01): + n = h01.shape[-1] // 2 + swapped = jnp.concatenate([h01[..., n:], h01[..., :n]], axis=-1) + # Same forward value (act_fn reads the halves in the right order), gradient of the halves exchanged. + return real(lambda up, gate: act_fn(gate, up), jax.lax.stop_gradient(swapped) + (h01 - jax.lax.stop_gradient(h01))) + + return mock.patch.object(moe, "_prefused_ffn_activation", wrong) + + +def _ties_reversed_pair_argsort(): + return mock.patch.object(ragged_sort, "argsort_unique_pairs", _argsort_ties_reversed) + + +class _ConfigOverride: + """`config` with some attributes replaced: reaches the layer's trace-time checks past the config validator.""" + + def __init__(self, config, **overrides): + self._config = config + self._overrides = overrides + + def __getattr__(self, name): # only called for names that are not set on the instance + if name.startswith("__") or "_config" not in self.__dict__: + raise AttributeError(name) + return self._overrides[name] if name in self._overrides else getattr(self._config, name) + + +def _unsplit_wi_tensor_parallel_guard(): + """On a tensor-parallel mesh the layer refuses moe_unsplit_prefused_wi at trace time; the plain layer runs.""" + cfg = _cfg(prefuse_moe_weights=True, **_TP_MESH) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + x, forced, token_mask = _inputs(cfg, "real") + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + _, params, _ = nnx.split(_build(cfg, mesh), nnx.Param, ...) + (loss, _, _), _ = _run(cfg, mesh, params, x, forced, token_mask) + plain_runs = bool(np.isfinite(loss)) + try: + _run(_ConfigOverride(cfg, moe_unsplit_prefused_wi=True), mesh, params, x, forced, token_mask) + except ValueError as e: + raised = "the MLP dimension sharded 2-way in the expert shard_map" in str(e) + else: + raised = False + ok = plain_runs and raised + print( + f"unsplit_wi_tensor_parallel_guard: [{'ok' if ok else 'FAIL'}] mesh {dict(mesh.shape)}: plain layer runs:" + f" {plain_runs}; moe_unsplit_prefused_wi past the validator raises at trace time: {raised}", + flush=True, + ) + return ok + + +# The subset of the Qwen3.5-397B RL trainer benchmark's MoE settings that this file needs (the cases above do not +# all set them). +_BENCH_FLAGS = { + "moe_chunk_pipeline": True, + "moe_expert_weight_prefetch": "forward", + "ring_of_experts_row_major_reduce_scatter": True, + "moe_routing_maps": "device", + "ragged_buffer_factor": 2.75, + "fp8_moe_fake_quant": True, + "moe_exclude_padding_from_routing": True, +} + + +def main(): + assert jax.device_count() == _NUM_DEVICES, jax.devices() + wi, pairs, swiglu = _KEYS + pipelined = {"moe_chunk_pipeline": True} + fake_quant = {**pipelined, "fp8_moe_fake_quant": True, "moe_exclude_padding_from_routing": True} + results = { + "unsplit_wi_plain_loop": _case("unsplit_wi_plain_loop", (wi,), {}), + "unsplit_wi_fake_quant": _case("unsplit_wi_fake_quant", (wi,), fake_quant, mode="padmask"), + "pair_argsort_pipelined": _case("pair_argsort_pipelined", (pairs,), pipelined), + "pair_argsort_dropping": _case( + "pair_argsort_dropping", (pairs,), {**pipelined, "ragged_buffer_factor": 1.0}, mode="forced" + ), + "swiglu_vjp_pipelined": _case("swiglu_vjp_pipelined", (swiglu,), pipelined), + "all_three_plain_loop": _case("all_three_plain_loop", _KEYS, {}), + "all_three_fake_quant": _case("all_three_fake_quant", _KEYS, fake_quant, mode="padmask"), + "all_three_dropless": _case("all_three_dropless", _KEYS, {"ragged_buffer_factor": -1.0}, mode="forced"), + "wi_and_swiglu_with_count": _case( + "wi_and_swiglu_with_count", (wi, swiglu), {**pipelined, "moe_routing_index_math": "count"} + ), + "all_three_bench_flags": _case("all_three_bench_flags", _KEYS, _BENCH_FLAGS, mode="padmask"), + "unsplit_wi_tensor_parallel_guard": _unsplit_wi_tensor_parallel_guard(), + } + controls = { + "control_swapped_activation_cotangents": _case( + "control_swapped_activation_cotangents", (swiglu,), pipelined, control=_swapped_activation_cotangents() + ), + "control_ties_reversed_under_drops": _case( + "control_ties_reversed_under_drops", + (pairs,), + {**pipelined, "ragged_buffer_factor": 1.0}, + mode="forced", + control=_ties_reversed_pair_argsort(), + ), + } + print("SUMMARY", {**results, **controls}, flush=True) + failed = [k for k, v in results.items() if not v] + undetected = [k for k, v in controls.items() if v] + if failed or undetected: + print(f"cases that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True) + sys.exit(1) + print(_PASS_TOKEN, flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/moe_roe_packed_routing_gather_test.py b/tests/unit/moe_roe_packed_routing_gather_test.py new file mode 100644 index 0000000000..da3be34079 --- /dev/null +++ b/tests/unit/moe_roe_packed_routing_gather_test.py @@ -0,0 +1,370 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""moe_roe_packed_routing_gather against the per-chunk routing gathers of the pipelined chunk loop, on a CPU mesh. + +With ring_of_experts_local_routing and moe_chunk_pipeline, every token chunk routes its own tokens and all-gathers +its top-k weights and expert ids: two all-gathers per chunk. The key routes all of the layer's local tokens once and +all-gathers weights and ids as ONE float32 [batch, seq, 2 * top_k] array, which each chunk slices. Routing is per +token, so each slice holds what the chunk would have computed. + +In-process checks: `pack_routing_for_all_gather` / `unpack_routing` round-trip routing weights and expert ids bit +for bit, ids up to 4095 and the -1 padding id included (the layer below has 16 experts, and bfloat16 holds every +integer up to 256, so only ids beyond 256 can tell an exact float32 packing from a lossy one), and the packed array +is float32. The validator rejects the key for each requirement it does not meet, naming it. + +Runs as a subprocess so the forced device count takes effect before JAX initializes: the shrunk Qwen3.5 routed MoE +of `moe_roe_local_routing_test.py` (16 experts, top 4, 8 CPU devices under cp-as-ep, ragged sort, local routing, +two token chunks, megablox in interpret mode), with the pipelined loop in both runs. Not reachable on CPU: the +all-gather pinned to a SparseCore (moe_pin_sparse_core_all_gathers / moe_pin_sparse_core_ep_all_gathers), a TPU-only +compute_on placement; the packed gather shares it with the per-chunk gathers, and the Qwen3.5-397B RL trainer +benchmark does not pin (it measured +0.363 s per micro-batch slower with the pin). + +Pre-registered comparisons, per case: + - Bit-identical output, loss and every gradient with float32 routing weights (float32_weight_sum, the default) and + no load-balance loss: real routing, padding excluded from routing, forced ids with a buffer that drops tokens, + saved routing maps, with moe_routing_index_math="count", and with the benchmark's MoE settings (expert-weight + prefetch, row-major reduce-scatter, routing maps on device, ragged_buffer_factor 2.75, prefused weights with + fp8 fake-quant, padding excluded). + - With the load-balance loss on, the loss and the output stay bit-identical (the per-chunk loss is computed from + this shard's own ids of the chunk, as without the key) and the gradients agree to float32 reduction order, + max |new - old| <= 1e-5 * max |old|: the router gradient sums the top-k and the loss terms in another order. + - bfloat16 routing weights (float32_weight_sum=False): the output is bit-identical; their cotangent is + reduce-scattered in float32 instead of bfloat16, so the gradients agree to relative L2 < 1e-2. +Two-sided: + - Structure: the plain loop must all-gather two [..., top_k] arrays per chunk and the packed one none of those + but exactly one [..., 2 * top_k] float32 array; under moe_routing_maps=device that array is tagged + moe_routing_maps, under remat it is untagged, and nothing is ever tagged moe_route. + - Deliberately wrong versions must FAIL: a packed gather whose result is shifted along the sequence (each chunk + gets another chunk's routing), and the packed path with its load-balance loss dropped. +""" + +import contextlib +import os +import re +import subprocess +import sys +from unittest import mock + +from absl.testing import parameterized +from flax import nnx +from flax.linen import partitioning as nn_partitioning +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +import numpy as np +import pytest + +from maxtext.configs import pyconfig +from maxtext.layers import moe +from maxtext.layers.initializers import nd_dense_init +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + +_NUM_DEVICES = 8 +_NUM_EXPERTS = 16 +_TOP_K = 4 +_EMB = 256 +_SEQ = 64 +_PASS_TOKEN = "MOE_ROE_PACKED_ROUTING_GATHER_CHECKS_PASSED" +_XLA_FLAGS = f"--xla_force_host_platform_device_count={_NUM_DEVICES} --xla_allow_excess_precision=false" +_REL_TOL = 1e-5 +_BF16_REL_NORM_TOL = 1e-2 + + +def _cfg(**overrides): + """The shrunk Qwen3.5 routed-MoE config on the cp-as-ep mesh, local routing, pipelined chunk loop.""" + kwargs = { + "run_name": "moe_roe_packed_routing_gather", + "enable_checkpointing": False, + "log_config": False, + "skip_jax_distributed_system": True, + "model_name": "qwen3.5-397b-a17b", + "override_model_config": True, + "num_experts": _NUM_EXPERTS, + "num_experts_per_tok": _TOP_K, + "base_emb_dim": _EMB, + "base_mlp_dim": 256, + "base_moe_mlp_dim": 256, + "max_target_length": _SEQ, + "max_prefill_predict_length": _SEQ, + "per_device_batch_size": 0.25, + "dtype": "bfloat16", + "weight_dtype": "float32", + "float32_gate_logits": True, + "megablox": True, + "sparse_matmul": True, + # Tokamax gmm v2 is TPU-only; both runs share whichever GMM runs, so CPU uses megablox. + "use_tokamax_gmm": False, + "use_gmm_v2": False, + "use_ring_of_experts": True, + "use_ragged_sort": True, + "use_custom_sort_vjp": False, + "ragged_buffer_factor": 2.0, + "num_moe_token_chunks": 2, + "custom_mesh_and_rule": "cp-as-ep", + "ici_fsdp_parallelism": 1, + "ici_context_parallelism": 4, + "ici_expert_parallelism": 2, + "context_parallel_load_balance": False, + "ring_of_experts_local_routing": True, + "moe_chunk_pipeline": True, + } + kwargs.update(overrides) + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs) + + +class PackRoutingTest(parameterized.TestCase): + """The packed routing array carries weights and expert ids exactly, ids beyond bfloat16's integers included.""" + + @parameterized.named_parameters(("float32_weights", jnp.float32), ("bfloat16_weights", jnp.bfloat16)) + def test_round_trip_is_exact(self, weights_dtype): + k0, k1 = jax.random.split(jax.random.PRNGKey(0)) + weights = jax.random.uniform(k0, (2, 64, 10), jnp.float32).astype(weights_dtype) + # Every id from -1 (padding) to 4095, so the odd ids above 256 that bfloat16 cannot hold are all there. + ids = (jax.random.permutation(k1, 2 * 64 * 10) % 4097 - 1).reshape(2, 64, 10).astype(jnp.int32) + packed = moe.pack_routing_for_all_gather(weights, ids) + self.assertEqual(packed.dtype, jnp.float32) + self.assertEqual(packed.shape, (2, 64, 20)) + got_weights, got_ids = moe.unpack_routing(packed, weights.dtype, ids.dtype) + self.assertEqual((got_weights.dtype, got_ids.dtype), (weights.dtype, ids.dtype)) + np.testing.assert_array_equal(np.asarray(got_ids), np.asarray(ids)) + np.testing.assert_array_equal(np.asarray(got_weights).view(np.uint8), np.asarray(weights).view(np.uint8)) + + +class PackedRoutingGatherConfigTest(parameterized.TestCase): + """The key is rejected for each requirement it does not meet; the load-balance loss is allowed.""" + + @parameterized.named_parameters( + ("dense", {"sparse_matmul": False}, "sparse_matmul=True"), + ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"), + ("no_local_routing", {"ring_of_experts_local_routing": False}, "ring_of_experts_local_routing=True"), + ("no_pipeline", {"moe_chunk_pipeline": False}, "moe_chunk_pipeline=True"), + ("random_routing", {"use_random_routing": True}, "use_random_routing=False"), + ) + def test_rejected(self, overrides, requirement): + with self.assertRaisesRegex(ValueError, f"moe_roe_packed_routing_gather=True requires: .*{re.escape(requirement)}"): + _cfg(moe_roe_packed_routing_gather=True, **overrides) + + def test_accepted_with_load_balance_loss(self): + cfg = _cfg(moe_roe_packed_routing_gather=True, load_balance_loss_weight=0.01) + self.assertTrue(cfg.moe_roe_packed_routing_gather) + + +@pytest.mark.cpu_only +def test_packed_routing_gather_matches_per_chunk_gathers_on_cpu_mesh(): + """Runs `main` below in a child process with 8 CPU devices.""" + env = os.environ.copy() + env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" {_XLA_FLAGS}" + env["JAX_PLATFORMS"] = "cpu" + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + print(result.stdout) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + + +def _build(cfg, mesh): + return moe.RoutedMoE( + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.moe_mlp_dim, + weight_dtype=cfg.weight_dtype, + dtype=cfg.dtype, + rngs=nnx.Rngs(params=0), + ) + + +def _inputs(cfg, mode): + """The layer input, forced expert ids (padded and duplicate slots) and a token mask, per `mode`.""" + batch = int(cfg.per_device_batch_size * _NUM_DEVICES) + x = jax.random.normal(jax.random.PRNGKey(7), (batch, _SEQ, _EMB), jnp.float32) + forced = token_mask = None + if "forced" in mode: + keys = jax.random.split(jax.random.PRNGKey(11), batch * _SEQ) + ids = jax.vmap(lambda k: jax.random.permutation(k, _NUM_EXPERTS)[:_TOP_K])(keys).reshape(batch, _SEQ, _TOP_K) + pos = jnp.arange(_SEQ)[None, :, None] + ids = jnp.where((pos % 5 == 0) & (jnp.arange(_TOP_K) == _TOP_K - 1), -1, ids) + ids = jnp.where((pos % 7 == 3) & (jnp.arange(_TOP_K) == 1), ids[..., :1], ids) + forced = ids.astype(jnp.int32) + if "padmask" in mode: + lens = jnp.array([_SEQ - 5 * (b + 1) for b in range(batch)]) + token_mask = (jnp.arange(_SEQ)[None, :] < lens[:, None]).astype(jnp.int32) + return x, forced, token_mask + + +def _eqns(jaxpr): + """Every equation of a jaxpr, including those of nested jaxprs.""" + for eqn in jaxpr.eqns: + yield eqn + for param in eqn.params.values(): + for sub in param if isinstance(param, (list, tuple)) else (param,): + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + if hasattr(sub, "eqns"): + yield from _eqns(sub) + + +def _structure(jaxpr): + """Routing all-gathers by last dim, and the checkpoint names on [..., 2 * top_k] arrays and anywhere.""" + eqns = list(_eqns(jaxpr)) + gathers = [tuple(v.aval.shape) for e in eqns if e.primitive.name == "all_gather" for v in e.outvars] + names = [(e.params.get("name"), tuple(e.invars[0].aval.shape)) for e in eqns if e.primitive.name == "name"] + return { + "topk_gathers": sum(s[-1:] == (_TOP_K,) for s in gathers), + "packed_gathers": sum(s[-1:] == (2 * _TOP_K,) for s in gathers), + "packed_tags": sorted({n for n, s in names if s[-1:] == (2 * _TOP_K,)}), + "moe_route_tags": sum(n == "moe_route" for n, _ in names), + } + + +def _run(cfg, mesh, params, x, forced, token_mask): + """Loss, output, gradients, and the structure of the traced forward.""" + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + graphdef, _, rest = nnx.split(_build(cfg, mesh), nnx.Param, ...) + + def loss_fn(params, x): + out, lb_loss, _ = nnx.merge(graphdef, params, rest)(x, forced_routed_experts=forced, token_mask=token_mask) + loss = jnp.mean(out.astype(jnp.float32) ** 2) + if lb_loss is not None: + loss = loss + lb_loss.astype(jnp.float32) + return loss, out + + structure = _structure(jax.make_jaxpr(loss_fn)(params, x).jaxpr) + (loss, out), grads = jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(params, x) + return jax.device_get((loss, out, grads)), structure + + +def _compare(label, old, new, mode): + """Prints the per-leaf report. mode: "bitwise", "reduction_order" (gradients) or "bf16" (gradients).""" + (loss_old, out_old, grads_old), (loss_new, out_new, grads_new) = old, new + named = [("output", out_old, out_new), ("loss", loss_old, loss_new)] + [ + (jax.tree_util.keystr(path), a, b) + for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_old)[0], jax.tree_util.tree_leaves(grads_new)) + ] + passed, lines = True, [] + for name, a, b in named: + a, b = np.asarray(a).astype(np.float32), np.asarray(b).astype(np.float32) + diff, scale = float(np.max(np.abs(b - a))), float(np.max(np.abs(a))) + rel_l2 = float(np.linalg.norm(b - a)) / max(float(np.linalg.norm(a)), 1e-30) + same = bool(np.array_equal(a, b)) + if mode == "bitwise" or name in ("output", "loss"): + ok = same + elif mode == "reduction_order": + ok = diff <= _REL_TOL * scale + else: + ok = rel_l2 < _BF16_REL_NORM_TOL + ok = ok and bool(np.isfinite(b).all()) and scale > 0 + passed &= ok + lines.append( + f" [{'ok' if ok else 'FAIL'}] {name}: bit-identical={same} max|d|={diff:.3e} max|old|={scale:.3e}" + f" rel_l2={rel_l2:.3e}" + ) + print(f"{label}: {'PASS' if passed else 'FAIL'} ({mode})", flush=True) + print("\n".join(lines), flush=True) + return passed + + +def _case(label, overrides, mode="real", compare="bitwise", control=None): + """The per-chunk routing gathers against the packed one, on the same parameters; `control` wraps the packed run.""" + cfg_old = _cfg(**overrides) + cfg_new = _cfg(**overrides, moe_roe_packed_routing_gather=True) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg_old), cfg_old.mesh_axes) + x, forced, token_mask = _inputs(cfg_old, mode) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_old.logical_axis_rules): + _, params, _ = nnx.split(_build(cfg_old, mesh), nnx.Param, ...) + old, s_old = _run(cfg_old, mesh, params, x, forced, token_mask) + with control if control is not None else contextlib.nullcontext(): + new, s_new = _run(cfg_new, mesh, params, x, forced, token_mask) + n_chunks = cfg_new.num_moe_token_chunks + want_tags = ["moe_routing_maps"] if cfg_new.moe_routing_maps != "remat" else [] + structure = ( + s_old["topk_gathers"] == 2 * n_chunks + and s_old["packed_gathers"] == 0 + and s_new["topk_gathers"] == 0 + and s_new["packed_gathers"] == 1 + and s_new["packed_tags"] == want_tags + and s_old["moe_route_tags"] == s_new["moe_route_tags"] == 0 + ) + print(f"{label}: [{'ok' if structure else 'FAIL'}] structure plain={s_old} packed={s_new}", flush=True) + return _compare(label, old, new, compare) and structure + + +def _shifted_packed_gather(): + """Deliberately wrong: the packed routing array comes back rolled along the sequence by one chunk.""" + real = jax.lax.all_gather + + def wrong(x, axis_name, **kwargs): + out = real(x, axis_name, **kwargs) + if out.ndim == 3 and out.shape[-1] == 2 * _TOP_K: + # Half the LOCAL sequence: inside the shard_map the sequence is sharded over the context axis. + out = jnp.roll(out, out.shape[1] // 2, axis=1) + return out + + return mock.patch.object(jax.lax, "all_gather", wrong) + + +def _dropped_load_balance_loss(): + """Deliberately wrong: the packed path without its load-balance loss (pmean of the loss returns zeros).""" + return mock.patch.object(jax.lax, "pmean", lambda x, axis_name, **kwargs: jnp.zeros_like(x)) + + +# The subset of the Qwen3.5-397B RL trainer benchmark's MoE settings that this file needs (the cases above do not +# all set them). +_BENCH_FLAGS = { + "moe_expert_weight_prefetch": "forward", + "ring_of_experts_row_major_reduce_scatter": True, + "moe_routing_maps": "device", + "ragged_buffer_factor": 2.75, + "prefuse_moe_weights": True, + "fp8_moe_fake_quant": True, + "moe_exclude_padding_from_routing": True, +} + + +def main(): + assert jax.device_count() == _NUM_DEVICES, jax.devices() + padded = {"moe_exclude_padding_from_routing": True} + lb = {"load_balance_loss_weight": 0.01} + results = { + "real": _case("real", {}), + "padding_excluded": _case("padding_excluded", padded, mode="padmask"), + "forced_dropping": _case("forced_dropping", {"ragged_buffer_factor": 1.0}, mode="forced"), + "saved_routing_maps": _case("saved_routing_maps", {"moe_routing_maps": "device"}), + "with_count": _case("with_count", {**padded, "moe_routing_index_math": "count"}, mode="padmask"), + "load_balance_loss": _case("load_balance_loss", lb, compare="reduction_order"), + "bf16_routing_weights": _case("bf16_routing_weights", {"float32_weight_sum": False}, compare="bf16"), + "bench_flags": _case("bench_flags", _BENCH_FLAGS, mode="padmask"), + } + controls = { + "control_shifted_packed_gather": _case("control_shifted_packed_gather", {}, control=_shifted_packed_gather()), + "control_dropped_load_balance_loss": _case( + "control_dropped_load_balance_loss", lb, compare="reduction_order", control=_dropped_load_balance_loss() + ), + } + print("SUMMARY", {**results, **controls}, flush=True) + failed = [k for k, v in results.items() if not v] + undetected = [k for k, v in controls.items() if v] + if failed or undetected: + print(f"cases that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True) + sys.exit(1) + print(_PASS_TOKEN, flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/moe_router_grad_from_wo_input_test.py b/tests/unit/moe_router_grad_from_wo_input_test.py new file mode 100644 index 0000000000..17b0ac4d6f --- /dev/null +++ b/tests/unit/moe_router_grad_from_wo_input_test.py @@ -0,0 +1,834 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""moe_router_grad_from_wo_input: the down projection and the unsort as one op, against the two apart, on CPU. + +`ragged_sort.ring_ragged_unsort_down_projection` runs the down-projection GEMM and then the forward of +`ring_ragged_unsort` (the same helpers), so the output must be bit-identical. Its backward does not read the GEMM's +output: it takes the GEMM's own VJP at the input w * h with the unweighted cotangent g, giving u = g wo^T and +dwo = (w * h)^T g, then dh = w * u and the routing-weight gradient (= ). Each gradient keeps one +rounding to the activation dtype, at a different point, so the gradients agree with the two ops apart to bf16 +rounding, not bit for bit. + +What a CPU run can and cannot reach: the ragged gather and gather-reduce take their jnp reference fallbacks off a TPU +target (the SparseCore kernels, and the validity sort of the gather-reduce, are TPU-only), and the layer's GMM is +megablox in interpret mode. The fused op's own math (the custom VJP, the weighting, the dot, the masks, the slot-order +return under both index maps and both buffering modes) is plain JAX and runs here as on TPU. + +Pre-registered, all in a subprocess with 8 CPU devices: + - The op on a 4-device expert-parallel mesh, with a jnp grouped matmul as the down projection, the full buffer and + a packed buffer that drops tokens, moe_routing_index_math "sort" and "count": the output bit-identical to the + two ops apart, every gradient (dx, dw, dwi, dwo) within relative L2 1e-2 of theirs. With a forward tie: the tie + comes back unchanged, its cotangent passes through, and the traced forward ties it to the GEMM output with an + optimization barrier. "count" without the forward permutation must raise. + - Precision, the op alone (h an input, so only the backward rounds): against the same function written densely in + float32 (slot by slot, plain autodiff, no ragged kernels; the ragged gather-reduce returns bf16, so the two-ops + path cannot serve as a float32 reference) at the same bf16-representable inputs, each fused gradient (dh, dw, + dwo) must be no further than _PRECISION_RATIO = 1.10 times as far as the two-ops gradient, on 6 seeds x {full + buffer "sort", dropping packed buffer "count"}. The routing weights sit at the bottom of a binade + ([0.25, 0.3)), where bf16 rounds them worst, so a backward that rounds them is visible. Measured over the 12 + runs: dh 0.996-1.005, dwo 0.992-1.003, dw 0.941-1.065 (dw, one value per slot, is the noisiest), so 1.10 sits + above the spread. A backward that rounds the routing weight to bf16 measures 1.31-1.34 on dh and dwo on every + seed, and one that shrinks the routing-weight gradient by 0.5% measures 3.0-3.3 on dw; both must fail. + Both paths share the unsort helpers, so a defect there moves both errors together and leaves the ratio at 1; + each path's error, and the forward output's, must therefore also stay within _PRECISION_MAX_ERR = 1e-2 of the + float32 reference (measured at most 2.4e-3). + - The shrunk Qwen3.5 routed MoE of moe_roe_local_routing_test.py (16 experts, top 4, cp-as-ep, ragged sort, local + routing, two token chunks): output and loss bit-identical to the key off, every gradient within relative L2 1e-2 + (measured <= 4.2e-3): plain and pipelined chunk loop, dropless, padding excluded with "count", saved routing + maps, the prefused-weight keys, and the benchmark's MoE settings with every compatible MoE key on against all + off (moe_routing_index_math="count", moe_unsplit_prefused_wi, moe_prefused_swiglu_vjp, + moe_roe_packed_routing_gather; expert-weight prefetch, row-major reduce-scatter, routing maps on device, + ragged_buffer_factor 2.75, fp8 fake-quant, padding excluded). A 0.5% error in the routing-weight gradient sits + inside the layer's 1e-2 tolerance; the precision check above is what catches it. Structure: the compiled + jax.grad(jax.checkpoint(loss)) has exactly num_moe_token_chunks GEMM loops fewer (the down-projection GEMM is no + longer recomputed in the backward), the pipelined forward keeps the same number of optimization barriers (tie 2 + moves into the op), and every chunk but the last reads the next chunk's routed tokens from the op's tie output. + - On an expert 4 x tensor 2 mesh the layer runs without the key and raises at trace time when the key reaches it + past the validator. With moe_accumulate_chunk_wgrad past the validator, the plain and the pipelined chunk loop + both raise (the fused op does not return the weights). + - Deliberately wrong versions must FAIL: the fused backward with the routing weight dropped from the input and + weight gradients (at the op and in the layer). + - The validator rejects the key for each requirement it does not meet, naming it. +""" + +import contextlib +import os +import re +import subprocess +import sys +from unittest import mock + +from absl.testing import parameterized +from flax import nnx +from flax.linen import partitioning as nn_partitioning +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +from jax.sharding import PartitionSpec as P +import numpy as np +import pytest + +from maxtext.configs import pyconfig +from maxtext.kernels.ragged import ragged_sort +from maxtext.layers import moe +from maxtext.layers.initializers import nd_dense_init +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + +_NUM_DEVICES = 8 +_NUM_EXPERTS = 16 +_TOP_K = 4 +_EMB = 256 +_SEQ = 64 +_PASS_TOKEN = "MOE_ROUTER_GRAD_FROM_WO_INPUT_CHECKS_PASSED" +_XLA_FLAGS = f"--xla_force_host_platform_device_count={_NUM_DEVICES} --xla_allow_excess_precision=false" +_GRAD_REL_L2 = 1e-2 +# The fused backward's error against a float32 run, over the two-ops backward's error (see the module docstring). +_PRECISION_RATIO = 1.10 +# The largest relative L2 error either path's output or gradients may have against the float32 run. +_PRECISION_MAX_ERR = 1e-2 +_PRECISION_SEEDS = range(6) + + +def _cfg(**overrides): + """The shrunk Qwen3.5 routed-MoE config on the cp-as-ep mesh (moe_roe_local_routing_test.py's, local routing).""" + kwargs = { + "run_name": "moe_router_grad_from_wo_input", + "enable_checkpointing": False, + "log_config": False, + "skip_jax_distributed_system": True, + "model_name": "qwen3.5-397b-a17b", + "override_model_config": True, + "num_experts": _NUM_EXPERTS, + "num_experts_per_tok": _TOP_K, + "base_emb_dim": _EMB, + "base_mlp_dim": 256, + "base_moe_mlp_dim": 256, + "max_target_length": _SEQ, + "max_prefill_predict_length": _SEQ, + "per_device_batch_size": 0.25, + "dtype": "bfloat16", + "weight_dtype": "float32", + "float32_gate_logits": True, + "megablox": True, + "sparse_matmul": True, + # Tokamax gmm v2 is TPU-only; both runs share whichever GMM runs, so CPU uses megablox. + "use_tokamax_gmm": False, + "use_gmm_v2": False, + "use_ring_of_experts": True, + "use_ragged_sort": True, + "use_custom_sort_vjp": False, + "ragged_buffer_factor": 2.0, + "num_moe_token_chunks": 2, + "custom_mesh_and_rule": "cp-as-ep", + "ici_fsdp_parallelism": 1, + "ici_context_parallelism": 4, + "ici_expert_parallelism": 2, + "context_parallel_load_balance": False, + "ring_of_experts_local_routing": True, + } + kwargs.update(overrides) + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs) + + +# A mesh with tensor parallelism: the default logical axis rules (cp-as-ep has no tensor axis), expert 4 x tensor 2. +_TP_MESH = { + "custom_mesh_and_rule": "", + "ici_context_parallelism": 1, + "ici_expert_parallelism": 4, + "ici_tensor_parallelism": 2, + "per_device_batch_size": 0.5, +} + + +class RouterGradFromWoInputConfigTest(parameterized.TestCase): + """The key is rejected for each requirement it does not meet.""" + + @parameterized.named_parameters( + ("dense", {"sparse_matmul": False}, "sparse_matmul=True"), + ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"), + ( + "no_ring", + {"use_ring_of_experts": False, "num_moe_token_chunks": 1, "ring_of_experts_local_routing": False}, + "use_ring_of_experts=True", + ), + ("no_ragged_sort", {"use_ragged_sort": False}, "use_ragged_sort=True"), + ("tensor_parallel", _TP_MESH, "no tensor parallelism"), + ("mlp_bias", {"mlp_bias": True}, "mlp_bias=False"), + ("quantization", {"quantization": "int8"}, "no weight quantization"), + ("tc_ragged_sort", {"moe_tc_ragged_sort": True}, "moe_tc_ragged_sort=False"), + ("chunk_wgrad", {"moe_accumulate_chunk_wgrad": True}, "moe_accumulate_chunk_wgrad=False"), + ("prequant_unsort", {"moe_bwd_prequant_before_unsort": True}, "moe_bwd_prequant_before_unsort=False"), + ("quantized_token_ag", {"moe_quantize_token_all_gather": True}, "moe_quantize_token_all_gather=False"), + ("emb_chunks", {"num_moe_emb_chunks": 2}, "num_moe_emb_chunks=0"), + ("minimal_remat", {"remat_policy": "minimal"}, "remat_policy=minimal saves it as mlpwo"), + ("save_dot_remat", {"remat_policy": "save_dot_except_mlpwi"}, "remat_policy=save_dot_except_mlpwi saves it"), + ("saved_mlpwo", {"remat_policy": "custom", "mlpwo": "device"}, "remat_policy=custom saves it as mlpwo"), + ("saved_moe_mlpwo", {"remat_policy": "custom", "moe_mlpwo": "offload"}, "saves it as moe_mlpwo"), + ) + def test_rejected(self, overrides, requirement): + with self.assertRaisesRegex(ValueError, f"moe_router_grad_from_wo_input=True requires: .*{re.escape(requirement)}"): + _cfg(moe_router_grad_from_wo_input=True, **overrides) + + def test_accepted(self): + self.assertTrue(_cfg(moe_router_grad_from_wo_input=True, moe_chunk_pipeline=True).moe_router_grad_from_wo_input) + # A policy that saves only other tensors (here the dense MLP input of the shared expert) is fine. + cfg = _cfg(moe_router_grad_from_wo_input=True, remat_policy="custom", mlpwi_1="device") + self.assertTrue(cfg.moe_router_grad_from_wo_input) + + +@pytest.mark.cpu_only +def test_router_grad_from_wo_input_on_cpu_mesh(): + """Runs `main` below in a child process with 8 CPU devices.""" + env = os.environ.copy() + env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" {_XLA_FLAGS}" + env["JAX_PLATFORMS"] = "cpu" + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + print(result.stdout) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + + +# --------------------------------------------------------------------------------------------------- the op alone +_OP_EP, _OP_EXPERTS, _OP_TOPK, _OP_TOKENS, _OP_D, _OP_F = 4, 8, 2, 64, 128, 64 +# A packed buffer smaller than the busiest shard's share of the 128 slots (32 on average), so tokens are dropped. +_OP_DROPPING_BUFFER = 24 +# The precision check's op: more slots (1024, 256 per shard on average), so each error estimate is tight. +_PRECISION_TOKENS = 512 +_PRECISION_DROPPING_BUFFER = 192 + + +def _grouped_matmul(h, w, group_sizes, expert_assignments, group_offset): + """A down projection with its own (autodiff) VJP: row r uses expert (global group of r) - group_offset.""" + del expert_assignments + ends = jnp.cumsum(group_sizes) + local = jnp.sum(jnp.arange(h.shape[0])[:, None] >= ends[None, :], axis=1) - group_offset + out = jnp.zeros((h.shape[0], w.shape[-1]), jnp.float32) + for e in range(w.shape[0]): + out = out + jnp.where((local == e)[:, None], h.astype(jnp.float32) @ w[e].astype(jnp.float32), 0.0) + return out.astype(h.dtype) + + +def _op_inputs(seed=0, num_tokens=_OP_TOKENS): + """Tokens, top-k expert ids, routing weights, wi, wo, the output cotangent and per-shard down-projection inputs.""" + k = jax.random.split(jax.random.PRNGKey(seed), 7) + x = jax.random.normal(k[0], (num_tokens, _OP_D), jnp.float32) + ids = jnp.stack([jax.random.permutation(kk, _OP_EXPERTS)[:_OP_TOPK] for kk in jax.random.split(k[1], num_tokens)]) + w = jax.nn.softmax(jax.random.normal(k[2], (num_tokens, _OP_TOPK)), axis=-1).reshape(-1) + wi = jax.random.normal(k[3], (_OP_EXPERTS, _OP_D, _OP_F)) * 0.1 + wo = jax.random.normal(k[4], (_OP_EXPERTS, _OP_F, _OP_D)) * 0.1 + cot = jax.random.normal(k[5], (num_tokens, _OP_D), jnp.float32) + # For the precision check: the down-projection input of every slot, on every shard. + h = jax.random.normal(k[6], (_OP_EP, num_tokens * _OP_TOPK, _OP_F), jnp.float32) + return x, ids.astype(jnp.int32), w, wi, wo, cot, h + + +def _expert_operands(gs, buffer_rows, num_slots, local): + """The grouped matmul's operands for this shard's buffer: (group sizes, None, group offset).""" + shard = jax.lax.axis_index("ep") + if buffer_rows < num_slots: + # Packed buffer: this shard's groups from row 0, clipped to the buffer. + ends = jnp.minimum(jnp.cumsum(jax.lax.dynamic_slice_in_dim(gs, shard * local, local)), buffer_rows) + return (ends - jnp.concatenate([jnp.zeros((1,), ends.dtype), ends[:-1]]), None, 0) + return (gs, None, shard * local) + + +def _op_run(buffer_size, index_math, fused, tie=False): + """bf16 run on a 4-device EP mesh; returns (output, dx, dw, dwi, dwo) summed / gathered over the shards, and the + forward-tie checks (tie unchanged, cotangent passed through).""" + x, ids, w, wi, wo, cot, _ = _op_inputs() + ep, local = _OP_EP, _OP_EXPERTS // _OP_EP + mesh = Mesh(np.array(jax.devices()[:ep]), ("ep",)) + wi_s = wi.reshape(ep, local, _OP_D, _OP_F).astype(jnp.bfloat16) + wo_s = wo.reshape(ep, local, _OP_F, _OP_D).astype(jnp.bfloat16) + tie_value = jnp.arange(8.0, dtype=jnp.float32) + + def body(x, w, wi, wo, cot): + wi, wo = wi[0], wo[0] + + def f(x, w, wi, wo, t): + sorted_x, gs, revert, perm = ragged_sort.ring_ragged_sort( + x, ids, _OP_EXPERTS, _OP_TOPK, "ep", ep, buffer_size=buffer_size, index_math=index_math + ) + operands = _expert_operands(gs, sorted_x.shape[0], _OP_TOKENS * _OP_TOPK, local) + h = jax.nn.silu(_grouped_matmul(sorted_x, wi, *operands).astype(jnp.float32)).astype(sorted_x.dtype) + perm_arg = perm if index_math == "count" else None + if fused: + return ragged_sort.ring_ragged_unsort_down_projection( + h, + wo, + _grouped_matmul, + operands, + gs, + revert, + _OP_TOPK, + local, + "ep", + topk_weights=w, + topk_argsort_indices=perm_arg, + forward_tie=t if tie else None, + index_math=index_math, + ) + y = _grouped_matmul(h, wo, *operands) + out = ragged_sort.ring_ragged_unsort( + y, gs, revert, _OP_TOPK, local, "ep", topk_weights=w, topk_argsort_indices=perm_arg, index_math=index_math + ) + return (out, t) if tie else out + + def g(x, w, wi, wo, t): + res = f(x, w, wi, wo, t) + return res if tie else (res, t) + + (out, t_out), vjp = jax.vjp(g, x, w, wi, wo, tie_value) + gx, gw, gwi, gwo, gt = vjp((cot.astype(out.dtype), 2.0 * t_out)) + return out[None], gx[None], gw[None], gwi[None], gwo[None], t_out[None], gt[None] + + sm = jax.shard_map( + body, mesh=mesh, in_specs=(P(), P(), P("ep"), P("ep"), P()), out_specs=(P("ep"),) * 7, check_vma=False + ) + args = (x.astype(jnp.bfloat16), w, wi_s, wo_s, cot) + outs = [np.asarray(jnp.asarray(a).astype(jnp.float32)) for a in jax.jit(sm)(*args)] + out, gx, gw, gwi, gwo, t_out, gt = outs + arrays = [ + out.sum(0), + gx.sum(0), + gw.sum(0), + gwi.reshape(_OP_EXPERTS, _OP_D, _OP_F), + gwo.reshape(_OP_EXPERTS, _OP_F, _OP_D), + ] + tie_ok = bool(np.array_equal(t_out[0], np.asarray(tie_value))) + tie_ok &= bool(np.array_equal(gt[0], 2.0 * np.asarray(tie_value))) + barrier = None + if fused and tie: + jaxpr = jax.make_jaxpr(sm)(*args).jaxpr + barrier = any(_barrier_ties_shapes(eqn) for eqn in _eqns(jaxpr)) + return arrays, tie_ok, barrier + + +def _barrier_ties_shapes(eqn): + """An optimization barrier whose operands include the tie (shape (8,)) next to a [rows, D] GEMM output.""" + if eqn.primitive.name != "optimization_barrier": + return False + shapes = [tuple(v.aval.shape) for v in eqn.invars] + return (8,) in shapes and any(len(s) == 2 and s[-1] == _OP_D for s in shapes) + + +def _rel(value, reference): + """Relative L2 distance of `value` from `reference`.""" + return float(np.linalg.norm(value - reference) / max(np.linalg.norm(reference), 1e-30)) + + +def _op_case(label, buffer_size, index_math, tie=False, control=None): + """The fused op against the two ops apart; `control` wraps the fused run.""" + base, _, _ = _op_run(buffer_size, index_math, False, tie) + with control if control is not None else contextlib.nullcontext(): + new, tie_ok, barrier = _op_run(buffer_size, index_math, True, tie) + passed = bool(np.array_equal(new[0], base[0])) + lines = [f" [{'ok' if passed else 'FAIL'}] output bit-identical={passed}"] + for name, a, b in zip(("dx", "dw", "dwi", "dwo"), base[1:], new[1:]): + rel = _rel(b, a) + ok = rel < _GRAD_REL_L2 and bool(np.isfinite(b).all()) + passed &= ok + lines.append(f" [{'ok' if ok else 'FAIL'}] {name}: rel_l2 vs two ops={rel:.2e}") + if tie: + passed &= tie_ok and bool(barrier) + lines.append( + f" [{'ok' if tie_ok and barrier else 'FAIL'}] forward tie unchanged, cotangent passed: {tie_ok};" + f" barrier ties it to the GEMM output: {barrier}" + ) + print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True) + print("\n".join(lines), flush=True) + return passed + + +def _dense_unsort_down_projection(h, w, wo, ids, gs, revert, shard, local): + """The fused op's function written densely, slot by slot (no ragged kernels, no rounding in float32). + + Slot i (token i // top_k) reads buffer row `revert[i]` (packed buffer: minus this shard's window start, and only + below the buffer's row count), multiplies it by its expert's down projection, and adds it, weighted, to its token. + """ + num_slots, buffer_rows = revert.shape[0], h.shape[0] + local_expert = ids.reshape(-1) - shard * local + is_local = (local_expert >= 0) & (local_expert < local) + start = jnp.sum(jnp.where(jnp.arange(gs.shape[0]) < shard * local, gs, 0)) + row = revert if buffer_rows >= num_slots else revert - start + limit = buffer_rows if buffer_rows >= num_slots else jnp.minimum(jnp.sum(jnp.where(is_local, 1, 0)), buffer_rows) + valid = is_local & (row >= 0) & (row < limit) + h_slot = jnp.where(valid[:, None], h[jnp.where(valid, row, 0)], 0.0) + y_slot = jnp.einsum("nf,nfd->nd", h_slot, wo[jnp.clip(local_expert, 0, local - 1)]) + weighted = jnp.where(valid[:, None], (w[:, None] * y_slot), 0.0) + return jnp.sum(weighted.reshape(num_slots // _OP_TOPK, _OP_TOPK, -1), axis=1) + + +def _op_backward_run(seed, buffer_size, index_math, path): + """The op (path "fused"), the two ops apart ("two_ops", both bf16) or the dense float32 function ("reference"), + differentiated at (h, w, wo) with the cotangent g on a 4-device EP mesh. h is an input, so the only roundings are + the backward's own (and, in the two-ops path, the forward GEMM output's). Returns (dh, dw, dwo), summed or gathered + over the shards, after each shard's forward output. The routing weights are in [0.25, 0.3), the bottom of a binade, + where bf16 rounds them worst.""" + x, ids, _, _, wo, cot, h = _op_inputs(seed, _PRECISION_TOKENS) + w = jax.random.uniform(jax.random.PRNGKey(1000 + seed), (_PRECISION_TOKENS * _OP_TOPK,), jnp.float32, 0.25, 0.3) + ep, local = _OP_EP, _OP_EXPERTS // _OP_EP + mesh = Mesh(np.array(jax.devices()[:ep]), ("ep",)) + # The same bf16-representable inputs on every path; the reference only computes in float32. + dtype = jnp.float32 if path == "reference" else jnp.bfloat16 + wo_s = wo.reshape(ep, local, _OP_F, _OP_D).astype(jnp.bfloat16).astype(dtype) + h_s = h.astype(jnp.bfloat16).astype(dtype) + cot = cot.astype(jnp.bfloat16).astype(dtype) + + def body(h, w, wo, cot): + h, wo = h[0], wo[0] + sorted_x, gs, revert, perm = ragged_sort.ring_ragged_sort( + x.astype(jnp.bfloat16), ids, _OP_EXPERTS, _OP_TOPK, "ep", ep, buffer_size=buffer_size, index_math=index_math + ) + rows = sorted_x.shape[0] + operands = _expert_operands(gs, rows, _PRECISION_TOKENS * _OP_TOPK, local) + perm_arg = perm if index_math == "count" else None + + def f(h, w, wo): + if path == "reference": + return _dense_unsort_down_projection(h, w, wo, ids, gs, revert, jax.lax.axis_index("ep"), local) + if path == "fused": + return ragged_sort.ring_ragged_unsort_down_projection( + h, + wo, + _grouped_matmul, + operands, + gs, + revert, + _OP_TOPK, + local, + "ep", + topk_weights=w, + topk_argsort_indices=perm_arg, + index_math=index_math, + ) + y = _grouped_matmul(h, wo, *operands) + return ragged_sort.ring_ragged_unsort( + y, gs, revert, _OP_TOPK, local, "ep", topk_weights=w, topk_argsort_indices=perm_arg, index_math=index_math + ) + + out, vjp = jax.vjp(f, h[:rows], w, wo) + dh, dw, dwo = vjp(cot.astype(out.dtype)) + return out[None], dh[None], dw[None], dwo[None] + + sm = jax.shard_map(body, mesh=mesh, in_specs=(P("ep"), P(), P("ep"), P()), out_specs=(P("ep"),) * 4, check_vma=False) + out, dh, dw, dwo = [np.asarray(jnp.asarray(a).astype(jnp.float32)) for a in jax.jit(sm)(h_s, w, wo_s, cot)] + return out, dh, dw.sum(0), dwo.reshape(_OP_EXPERTS, _OP_F, _OP_D) + + +def _op_precision_case(label, buffer_size, index_math, control=None): + """Fused vs two-ops error against a float32 run, over _PRECISION_SEEDS; `control` wraps the fused runs.""" + passed, lines, ratios = True, [], [] + for seed in _PRECISION_SEEDS: + reference = _op_backward_run(seed, buffer_size, index_math, "reference") + two_ops = _op_backward_run(seed, buffer_size, index_math, "two_ops") + with control if control is not None else contextlib.nullcontext(): + fused = _op_backward_run(seed, buffer_size, index_math, "fused") + parts = [] + for name, r, a, b in zip(("out", "dh", "dw", "dwo"), reference, two_ops, fused): + err_two_ops, err_fused = _rel(a, r), _rel(b, r) + ratio = err_fused / err_two_ops + ok = ratio <= _PRECISION_RATIO and bool(np.isfinite(b).all()) and err_two_ops > 0 + ok &= max(err_two_ops, err_fused) <= _PRECISION_MAX_ERR + passed &= ok + ratios.append(ratio) + parts.append(f"{name} {err_two_ops:.2e}/{err_fused:.2e}={ratio:.3f}{'' if ok else ' FAIL'}") + lines.append(f" seed {seed}: error vs float32, two ops / fused = ratio: {'; '.join(parts)}") + print( + f"{label}: {'PASS' if passed else 'FAIL'} (ratio <= {_PRECISION_RATIO}, errors <= {_PRECISION_MAX_ERR};" + f" over the seeds min {min(ratios):.3f} max {max(ratios):.3f})", + flush=True, + ) + print("\n".join(lines), flush=True) + return passed + + +def _count_needs_forward_permutation(): + try: + _op_run_count_without_permutation() + except ValueError as e: + raised = "needs topk_argsort_indices" in str(e) + else: + raised = False + print(f"count without the forward permutation: [{'ok' if raised else 'FAIL'}] the fused backward raises", flush=True) + return raised + + +def _op_run_count_without_permutation(): + """Differentiates the fused op under "count" without passing the forward permutation (must raise).""" + x, ids, w, _, wo, cot, _ = _op_inputs() + mesh = Mesh(np.array(jax.devices()[:2]), ("ep",)) + + def body(x, wo, cot): + def f(x, wo): + sorted_x, gs, revert, _ = ragged_sort.ring_ragged_sort(x, ids, _OP_EXPERTS, _OP_TOPK, "ep", 2, index_math="count") + return ragged_sort.ring_ragged_unsort_down_projection( + sorted_x[:, :_OP_F], + wo[:4], + _grouped_matmul, + (gs, None, 0), + gs, + revert, + _OP_TOPK, + 4, + "ep", + topk_weights=w, + index_math="count", + ) + + out, vjp = jax.vjp(f, x, wo) + return vjp(cot.astype(out.dtype))[1] + + jax.shard_map(body, mesh=mesh, in_specs=(P(), P(), P()), out_specs=P(), check_vma=False)( + x.astype(jnp.bfloat16), wo.astype(jnp.bfloat16), cot + ) + + +def _fan_out_changed(weight=None, slot_value=None): + """Deliberately wrong: the fused op's backward fan-out with its routing weights or its per-slot result changed. + + ring_ragged_unsort uses the same helper, so only wrap the fused run with this. + """ + real = ragged_sort._ring_unsort_fan_out # pylint: disable=protected-access + + def wrong(*args, **kwargs): + gathered, weight_for_sorted, mask, to_slot_order = real(*args, **kwargs) + if weight is not None: + weight_for_sorted = weight(weight_for_sorted) + if slot_value is not None: + return gathered, weight_for_sorted, mask, lambda per_row: to_slot_order(slot_value(per_row)) + return gathered, weight_for_sorted, mask, to_slot_order + + return mock.patch.object(ragged_sort, "_ring_unsort_fan_out", wrong) + + +def _weight_dropped_from_backward(): + """Deliberately wrong: the fused backward's fan-out reports every routing weight as 1.""" + return _fan_out_changed(weight=jnp.ones_like) + + +def _weight_rounded_to_bf16(): + """Deliberately wrong: the fused backward uses the routing weight rounded to bf16.""" + return _fan_out_changed(weight=lambda w: w.astype(jnp.bfloat16).astype(jnp.float32)) + + +def _router_grad_shrunk(): + """Deliberately wrong: the fused backward's routing-weight gradient 0.5% too small.""" + return _fan_out_changed(slot_value=lambda dot: dot * 0.995) + + +# --------------------------------------------------------------------------------------------------- the layer +def _build(cfg, mesh): + return moe.RoutedMoE( + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.moe_mlp_dim, + weight_dtype=cfg.weight_dtype, + dtype=cfg.dtype, + rngs=nnx.Rngs(params=0), + ) + + +def _inputs(cfg, mode): + """The layer input, forced expert ids and a token mask, per `mode`.""" + batch = int(cfg.per_device_batch_size * _NUM_DEVICES) + x = jax.random.normal(jax.random.PRNGKey(7), (batch, _SEQ, _EMB), jnp.float32) + forced = token_mask = None + if "forced" in mode: + keys = jax.random.split(jax.random.PRNGKey(11), batch * _SEQ) + ids = jax.vmap(lambda k: jax.random.permutation(k, _NUM_EXPERTS)[:_TOP_K])(keys).reshape(batch, _SEQ, _TOP_K) + forced = ids.astype(jnp.int32) + if "padmask" in mode: + lens = jnp.array([_SEQ - 5 * (b + 1) for b in range(batch)]) + token_mask = (jnp.arange(_SEQ)[None, :] < lens[:, None]).astype(jnp.int32) + return x, forced, token_mask + + +def _eqns(jaxpr): + for eqn in jaxpr.eqns: + yield eqn + for param in eqn.params.values(): + for sub in param if isinstance(param, (list, tuple)) else (param,): + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + if hasattr(sub, "eqns"): + yield from _eqns(sub) + + +def _compiled_gemm_loops(hlo_text): + """Number of while instructions in a compiled HLO module. + + The GMMs are megablox kernels in interpret mode, each a loop over its tiles; identical GEMMs share one body + computation, so the loop instructions (one per call), not the dots in the bodies, count the GEMMs. + """ + return sum(" while(" in line for line in hlo_text.splitlines()) + + +def _sub_jaxprs(eqn): + for param in eqn.params.values(): + for sub in param if isinstance(param, (list, tuple)) else (param,): + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + if hasattr(sub, "eqns"): + yield sub + + +def _tie_reads(jaxpr): + """(fused ops with a forward tie, how many of them have their tie output read later), over all sub-jaxprs. + + A fused op with a tie is a custom_vjp call with two outputs whose primal holds an optimization barrier and the + unsort's ragged_gather_reduce (barrier_tie's own custom_vjp holds only the barrier). Its second output is the next + chunk's routed tokens after the barrier: if nothing reads it, that chunk is no longer ordered after this chunk's down + projection. + """ + ties = read = 0 + for i, eqn in enumerate(jaxpr.eqns): + call = eqn.params.get("call_jaxpr") + call = getattr(call, "jaxpr", call) + if ( + eqn.primitive.name.startswith("custom_vjp_call") + and len(eqn.outvars) == 2 + and hasattr(call, "eqns") + and any(e.primitive.name == "optimization_barrier" for e in _eqns(call)) + and any(e.params.get("name") == "ragged_gather_reduce" for e in _eqns(call)) + ): + ties += 1 + tie_out = eqn.outvars[-1] + read += any(tie_out in later.invars for later in jaxpr.eqns[i + 1 :]) or tie_out in jaxpr.outvars + for sub in _sub_jaxprs(eqn): + sub_ties, sub_read = _tie_reads(sub) + ties, read = ties + sub_ties, read + sub_read + return ties, read + + +def _run(cfg, mesh, params, x, forced, token_mask): + """Loss, output, gradients; forward barriers and fused-op ties; GEMM loops of the compiled rematerialized + backward.""" + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + graphdef, _, rest = nnx.split(_build(cfg, mesh), nnx.Param, ...) + + def loss_fn(params, x): + out, lb_loss, _ = nnx.merge(graphdef, params, rest)(x, forced_routed_experts=forced, token_mask=token_mask) + loss = jnp.mean(out.astype(jnp.float32) ** 2) + if lb_loss is not None: + loss = loss + lb_loss.astype(jnp.float32) + return loss, out + + forward = jax.make_jaxpr(loss_fn)(params, x).jaxpr + barriers = sum(e.primitive.name == "optimization_barrier" for e in _eqns(forward)) + ties = _tie_reads(forward) + + def linear_loss(params, x): + # Linear in the layer output, as behind the decoder's residual add: the layer's own backward does not need its + # output, so only the backward's own reads of the down-projection output keep that GEMM in the replay. + out, _, _ = nnx.merge(graphdef, params, rest)(x, forced_routed_experts=forced, token_mask=token_mask) + return jnp.sum(out.astype(jnp.float32) * jax.random.normal(jax.random.PRNGKey(3), out.shape, jnp.float32)) + + remat = jax.grad(jax.checkpoint(linear_loss, policy=jax.checkpoint_policies.nothing_saveable), argnums=(0, 1)) + loops = _compiled_gemm_loops(jax.jit(remat).lower(params, x).compile().as_text()) + (loss, out), grads = jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(params, x) + return jax.device_get((loss, out, grads)), barriers, ties, loops + + +def _layer_case(label, overrides, mode="real", control=None, keys=None, same_barriers=True): + """The layer with `keys` off and on (default: this key alone), on the same parameters; `control` wraps the run + with them on. `same_barriers`: the forward keeps its optimization barriers (False where other keys change them).""" + keys = {"moe_router_grad_from_wo_input": True} if keys is None else keys + cfg_old = _cfg(**overrides) + cfg_new = _cfg(**overrides, **keys) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg_old), cfg_old.mesh_axes) + x, forced, token_mask = _inputs(cfg_old, mode) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_old.logical_axis_rules): + _, params, _ = nnx.split(_build(cfg_old, mesh), nnx.Param, ...) + old, barriers_old, ties_old, loops_old = _run(cfg_old, mesh, params, x, forced, token_mask) + with control if control is not None else contextlib.nullcontext(): + new, barriers_new, ties_new, loops_new = _run(cfg_new, mesh, params, x, forced, token_mask) + (loss_old, out_old, grads_old), (loss_new, out_new, grads_new) = old, new + n_chunks = cfg_new.num_moe_token_chunks + # One GEMM loop fewer per chunk in the rematerialized backward: the down projection is not replayed. + loops_ok = loops_old - loops_new == n_chunks + barriers_ok = barriers_new == barriers_old or not same_barriers + # Pipelined: every chunk but the last ties its down projection to the next chunk's routed tokens, which are read. + want_ties = (n_chunks - 1, n_chunks - 1) if cfg_new.moe_chunk_pipeline else (0, 0) + ties_ok = ties_old == (0, 0) and ties_new == want_ties + passed = loops_ok and barriers_ok and ties_ok + lines = [ + f" [{'ok' if loops_ok else 'FAIL'}] compiled remat-backward GEMM loops off/on={loops_old}/{loops_new}" + f" (want {n_chunks} fewer)", + f" [{'ok' if barriers_ok else 'FAIL'}] forward barriers off/on={barriers_old}/{barriers_new}" + f"{'' if same_barriers else ' (not compared)'}", + f" [{'ok' if ties_ok else 'FAIL'}] fused-op ties (count, read) off/on={ties_old}/{ties_new} (want {want_ties})", + ] + for name, a, b in (("output", out_old, out_new), ("loss", loss_old, loss_new)): + same = bool(np.array_equal(np.asarray(a, np.float32), np.asarray(b, np.float32))) + passed &= same + lines.append(f" [{'ok' if same else 'FAIL'}] {name}: bit-identical={same}") + for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_old)[0], jax.tree_util.tree_leaves(grads_new)): + a, b = np.asarray(a, np.float32), np.asarray(b, np.float32) + rel = _rel(b, a) + ok = rel < _GRAD_REL_L2 and bool(np.isfinite(b).all()) + passed &= ok + lines.append(f" [{'ok' if ok else 'FAIL'}] {jax.tree_util.keystr(path)}: rel_l2={rel:.2e}") + print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True) + print("\n".join(lines), flush=True) + return passed + + +class _ConfigOverride: + """`config` with some attributes replaced: reaches the layer's trace-time checks past the config validator.""" + + def __init__(self, config, **overrides): + self._config = config + self._overrides = overrides + + def __getattr__(self, name): # only called for names that are not set on the instance + if name.startswith("__") or "_config" not in self.__dict__: + raise AttributeError(name) + return self._overrides[name] if name in self._overrides else getattr(self._config, name) + + +def _tensor_parallel_guard(): + """On a tensor-parallel mesh the layer refuses the key at trace time; without it the layer runs.""" + cfg = _cfg(**_TP_MESH) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + x, forced, token_mask = _inputs(cfg, "real") + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + _, params, _ = nnx.split(_build(cfg, mesh), nnx.Param, ...) + (loss, _, _), _, _, _ = _run(cfg, mesh, params, x, forced, token_mask) + plain_runs = bool(np.isfinite(loss)) + try: + _run(_ConfigOverride(cfg, moe_router_grad_from_wo_input=True), mesh, params, x, forced, token_mask) + except ValueError as e: + raised = "does not support tensor parallelism in the MoE layer" in str(e) + else: + raised = False + ok = plain_runs and raised + print( + f"tensor_parallel_guard: [{'ok' if ok else 'FAIL'}] mesh {dict(mesh.shape)}: plain layer runs: {plain_runs};" + f" the key past the validator raises at trace time: {raised}", + flush=True, + ) + return ok + + +# The subset of the Qwen3.5-397B RL trainer benchmark's MoE settings that this file needs (the cases above do not +# all set them). +_BENCH_FLAGS = { + "moe_chunk_pipeline": True, + "moe_expert_weight_prefetch": "forward", + "ring_of_experts_row_major_reduce_scatter": True, + "moe_routing_maps": "device", + "ragged_buffer_factor": 2.75, + "prefuse_moe_weights": True, + "fp8_moe_fake_quant": True, + "moe_exclude_padding_from_routing": True, +} +# Every MoE performance key that composes with the others on the megablox path the CPU runs (the fused SiLU GMM is +# tokamax-gmm_v2-only; moe_wi_gmm_fused_swiglu_test.py stacks it with the others using kernel stand-ins). +_ALL_KEYS = { + "moe_routing_index_math": "count", + "moe_unsplit_prefused_wi": True, + "moe_prefused_swiglu_vjp": True, + "moe_roe_packed_routing_gather": True, + "moe_router_grad_from_wo_input": True, +} + + +def _chunk_wgrad_guard(): + """With moe_accumulate_chunk_wgrad past the validator, both chunk loops refuse the key at trace time.""" + results = [] + for loop, overrides in (("plain", {}), ("pipelined", {"moe_chunk_pipeline": True})): + cfg = _cfg(**overrides) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + x, forced, token_mask = _inputs(cfg, "real") + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + _, params, _ = nnx.split(_build(cfg, mesh), nnx.Param, ...) + override = _ConfigOverride(cfg, moe_router_grad_from_wo_input=True, moe_accumulate_chunk_wgrad=True) + try: + _run(override, mesh, params, x, forced, token_mask) + except NotImplementedError as e: + raised = "does not support moe_accumulate_chunk_wgrad" in str(e) + else: + raised = False + results.append(raised) + print(f"chunk_wgrad_guard: [{'ok' if raised else 'FAIL'}] {loop} loop raises at trace time: {raised}", flush=True) + return all(results) + + +def main(): + assert jax.device_count() == _NUM_DEVICES, jax.devices() + pipelined = {"moe_chunk_pipeline": True} + dropping, precision_dropping = _OP_DROPPING_BUFFER, _PRECISION_DROPPING_BUFFER + results = { + "op_full_buffer_sort": _op_case("op_full_buffer_sort", None, "sort"), + "op_full_buffer_count": _op_case("op_full_buffer_count", None, "count"), + "op_dropping_buffer_sort": _op_case("op_dropping_buffer_sort", dropping, "sort"), + "op_dropping_buffer_count": _op_case("op_dropping_buffer_count", dropping, "count"), + "op_forward_tie": _op_case("op_forward_tie", dropping, "sort", tie=True), + "op_count_guard": _count_needs_forward_permutation(), + "op_precision_full_buffer_sort": _op_precision_case("op_precision_full_buffer_sort", None, "sort"), + "op_precision_dropping_buffer_count": _op_precision_case( + "op_precision_dropping_buffer_count", precision_dropping, "count" + ), + "layer_plain_loop": _layer_case("layer_plain_loop", {}), + "layer_pipelined": _layer_case("layer_pipelined", pipelined), + "layer_dropless": _layer_case("layer_dropless", {"ragged_buffer_factor": -1.0}, mode="forced"), + "layer_count_padding_excluded": _layer_case( + "layer_count_padding_excluded", + {**pipelined, "moe_routing_index_math": "count", "moe_exclude_padding_from_routing": True}, + mode="padmask", + ), + "layer_saved_routing_maps": _layer_case("layer_saved_routing_maps", {**pipelined, "moe_routing_maps": "device"}), + "layer_prefused_keys": _layer_case( + "layer_prefused_keys", + {**pipelined, "prefuse_moe_weights": True, "moe_unsplit_prefused_wi": True, "moe_prefused_swiglu_vjp": True}, + ), + "layer_bench_flags_all_keys": _layer_case( + "layer_bench_flags_all_keys", _BENCH_FLAGS, mode="padmask", keys=_ALL_KEYS, same_barriers=False + ), + "tensor_parallel_guard": _tensor_parallel_guard(), + "chunk_wgrad_guard": _chunk_wgrad_guard(), + } + controls = { + "control_op_weight_dropped": _op_case( + "control_op_weight_dropped", dropping, "sort", control=_weight_dropped_from_backward() + ), + "control_op_precision_weight_rounded_to_bf16": _op_precision_case( + "control_op_precision_weight_rounded_to_bf16", None, "sort", control=_weight_rounded_to_bf16() + ), + "control_op_precision_router_grad_shrunk": _op_precision_case( + "control_op_precision_router_grad_shrunk", precision_dropping, "count", control=_router_grad_shrunk() + ), + "control_layer_weight_dropped": _layer_case( + "control_layer_weight_dropped", pipelined, control=_weight_dropped_from_backward() + ), + } + print("SUMMARY", {**results, **controls}, flush=True) + failed = [k for k, v in results.items() if not v] + undetected = [k for k, v in controls.items() if v] + if failed or undetected: + print(f"cases that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True) + sys.exit(1) + print(_PASS_TOKEN, flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/moe_routing_index_math_test.py b/tests/unit/moe_routing_index_math_test.py new file mode 100644 index 0000000000..80fa35b372 --- /dev/null +++ b/tests/unit/moe_routing_index_math_test.py @@ -0,0 +1,513 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""moe_routing_index_math="count" against the default "sort", on CPU. + +"count" computes the ring-of-experts ragged sort's routing index maps without sorts where it can: the expert +permutation by one packed single-operand sort (`packed_stable_sort`), the inverse permutation by counting on this +shard's slots only (`ragged_sort._local_slot_positions`, -1 elsewhere), the router-weight gradient back to slot order +by a masked gather, the gather-reduce's validity sort packed into one uint32 operand (`_preprocess(packed=True)`), and +the router's top-k indices of non-forced tokens by k argmax passes (`moe.topk_indices_by_argmax`). The permutation is +the same, so everything must be bit-identical to "sort". + +What a CPU run can and cannot reach: + - The gather-reduce's validity sort runs only on the SparseCore path. Off a TPU target `ragged_gather_reduce` takes + its jnp reference fallback, which does not sort, so `packed_preprocess` changes nothing in the layer runs below. + `PackedPreprocessTest` therefore calls `_preprocess` directly, the pure-JAX part of the SparseCore path. + - Likewise the ragged gather is the jnp fallback here. Everything else ("count"'s permutation, its counted inverse, + the masked weight-gradient gather, the argmax top-k) is plain JAX and runs in the layer below exactly as on TPU. + +In-process checks: + - `_local_slot_positions` equals argsort(argsort(ids)) - (this shard's window start) on this shard's slots and -1 + elsewhere, for every shard, with token counts that are not a multiple of the block, the padding id, and tokens + that repeat an expert. + - `topk_indices_by_argmax` equals `jax.lax.top_k`'s indices, ties included. + - `_preprocess(packed=True)` equals `_preprocess(packed=False)`, every output. + - The config validator rejects "count" off the ring-of-experts ragged sort, one case per requirement. + +The layer check runs as a subprocess, so the forced device count takes effect before JAX initializes: the shrunk +Qwen3.5 routed MoE of `moe_roe_local_routing_test.py` (16 experts, top 4, 8 CPU devices under cp-as-ep, ragged sort, +local routing, two token chunks, megablox in interpret mode). Pre-registered: the forward output, the loss and every +gradient bit-identical to "sort", in each case: the plain and the pipelined chunk loop, padding excluded from routing +(every token takes the forced-routing branch, so the argmax top-k runs), forced ids with a buffer that drops tokens +(asserted to overflow), a dropless buffer, saved routing maps, and a bfloat16 router. Two-sided: + - Structure: "count" must trace fewer `sort` equations than "sort" in the gradient and no argsort at all (a sort + of more than one operand: jnp.argsort sorts its keys with an iota), while "sort" traces argsorts; and with + padding excluded no `top_k`. So a switch that silently fell back, completely or for one map, cannot pass. + - Wiring of the SparseCore-only validity sort: CPU cannot run it, but the flag that selects it reaches + `ragged_gather_reduce` on every target. A spy records it: every call under "count" must pass + packed_preprocess=True and every call under "sort" False. + - The bench's MoE settings (expert-weight prefetch, row-major reduce-scatter, routing maps saved on device, + ragged_buffer_factor 2.75, prefused weights with fp8 fake-quant, padding excluded) in one case. + - A deliberately wrong counted inverse (every local position shifted by one) must FAIL the comparison. + - "count" without the forward permutation must raise in the unsort backward instead of argsorting a partial inverse. +""" + +import contextlib +import os +import re +import subprocess +import sys +from unittest import mock + +from absl.testing import parameterized +from flax import nnx +from flax.linen import partitioning as nn_partitioning +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +from jax.sharding import PartitionSpec as P +import numpy as np +import pytest + +from maxtext.configs import pyconfig +from maxtext.kernels.ragged import ragged_gather_reduce_v2 +from maxtext.kernels.ragged import ragged_sort +from maxtext.layers import moe +from maxtext.layers.initializers import nd_dense_init +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + +_NUM_DEVICES = 8 +_NUM_EXPERTS = 16 +_TOP_K = 4 +_EMB = 256 +_SEQ = 64 +_PASS_TOKEN = "MOE_ROUTING_INDEX_MATH_CHECKS_PASSED" +_XLA_FLAGS = f"--xla_force_host_platform_device_count={_NUM_DEVICES} --xla_allow_excess_precision=false" + + +def _expected_local_positions(ids, num_experts, shard, ep_size): + """Reference: each slot's rank in the stable sort by expert id, relative to this shard's window; -1 off-shard.""" + flat = np.asarray(ids).reshape(-1) + rank = np.argsort(np.argsort(flat, kind="stable"), kind="stable") + local_num_experts = num_experts // ep_size + start = shard * local_num_experts + local = (flat >= start) & (flat < start + local_num_experts) + window_start = int(np.sum(flat < start)) + return np.where(local, rank - window_start, -1).astype(np.int32) + + +@pytest.mark.cpu_only +class LocalSlotPositionsTest(parameterized.TestCase): + """`_local_slot_positions` is the inverse permutation restricted to this shard's window.""" + + @parameterized.named_parameters( + # (num_tokens, topk, num_experts, ep_size, block): one exact block, several blocks with a partial last one, + # fewer tokens than a block, and the production shape's ratios at a small scale. + ("one_block", 128, 4, 16, 8, 128), + ("partial_blocks", 300, 4, 16, 8, 64), + ("short", 5, 2, 8, 2, 128), + ("wide_topk", 200, 10, 64, 4, 32), + ) + def test_matches_double_argsort(self, num_tokens, topk, num_experts, ep_size, block): + rng = np.random.default_rng(num_tokens) + # Ids in [0, num_experts]: num_experts is the padding id (never local). Distinct per token, as top-k gives. + ids = np.stack([rng.permutation(num_experts + 1)[:topk] for _ in range(num_tokens)]).astype(np.int32) + for shard in range(ep_size): + got = np.asarray(ragged_sort._local_slot_positions(jnp.asarray(ids), num_experts, shard, ep_size, block)) # pylint: disable=protected-access + np.testing.assert_array_equal(got, _expected_local_positions(ids, num_experts, shard, ep_size), err_msg=f"{shard=}") + + def test_repeated_expert_within_a_token(self): + """A token that picks one expert twice: the stable order puts its earlier slot first.""" + ids = np.array([[3, 3, 1], [1, 3, 3], [16, 16, 0], [0, 1, 3]], np.int32) + for shard in range(8): + got = np.asarray(ragged_sort._local_slot_positions(jnp.asarray(ids), 16, shard, 8, 2)) # pylint: disable=protected-access + np.testing.assert_array_equal(got, _expected_local_positions(ids, 16, shard, 8), err_msg=f"{shard=}") + + def test_wrong_reference_is_detected(self): + """Two-sided: the comparison above fails on a deliberately wrong expectation (an unstable rank).""" + ids = np.array([[1, 0], [1, 0], [0, 1]], np.int32) + got = np.asarray(ragged_sort._local_slot_positions(jnp.asarray(ids), 2, 0, 1, 2)) # pylint: disable=protected-access + flat = ids.reshape(-1) + unstable_rank = np.argsort(np.argsort(-np.arange(flat.size) + flat * 100)) # ties broken by DEcreasing index + self.assertFalse(np.array_equal(got, unstable_rank.astype(np.int32))) + + +@pytest.mark.cpu_only +class TopkIndicesByArgmaxTest(parameterized.TestCase): + """`topk_indices_by_argmax` returns `jax.lax.top_k`'s indices.""" + + @parameterized.named_parameters(("distinct", False), ("ties", True)) + def test_matches_top_k(self, ties): + scores = jax.random.uniform(jax.random.PRNGKey(0), (4, 33, 64), jnp.float32) + if ties: + # Few distinct values, so most rows tie at the top-k boundary; and exact zeros, like underflowed softmax. + scores = jnp.round(scores * 4) / 4 + want = jax.lax.top_k(scores, 10)[1] + got = moe.topk_indices_by_argmax(scores, 10) + np.testing.assert_array_equal(np.asarray(got), np.asarray(want)) + + def test_on_softmax_probabilities(self): + logits = jax.random.normal(jax.random.PRNGKey(1), (2, 16, 512), jnp.float32) * 30.0 # many probs underflow to 0 + probs = jax.nn.softmax(logits, axis=-1) + np.testing.assert_array_equal( + np.asarray(moe.topk_indices_by_argmax(probs, 10)), np.asarray(jax.lax.top_k(probs, 10)[1]) + ) + + +@pytest.mark.cpu_only +class PackedPreprocessTest(parameterized.TestCase): + """The SparseCore gather-reduce's validity sort, packed and unpacked. + + Only the SparseCore path runs `_preprocess`; on a CPU target `ragged_gather_reduce` takes its reference fallback, + so this pure-JAX function is called directly. + """ + + @parameterized.named_parameters( + # (rows, reduce_group_size, num_row_partitions, num_simd_lanes, row_chunk_size) + ("one_partition", 96, 4, 1, 16, 16), + ("v7x_like", 1024, 8, 16, 16, 48), + ("padded_chunks", 640, 10, 8, 8, 24), + ) + def test_packed_matches_argsort(self, rows, group, partitions, lanes, chunk): + for density in (0.0, 0.3, 0.9, 1.0): + mask = jax.random.uniform(jax.random.PRNGKey(rows), (rows,)) < density + want = ragged_gather_reduce_v2._preprocess(mask, group, partitions, lanes, chunk) # pylint: disable=protected-access + got = ragged_gather_reduce_v2._preprocess(mask, group, partitions, lanes, chunk, packed=True) # pylint: disable=protected-access + for name, a, b in zip(("sorted_by_validity", "rows_per_partition", "group_mask"), want, got): + self.assertEqual(a.dtype, b.dtype, name) + np.testing.assert_array_equal(np.asarray(b), np.asarray(a), err_msg=f"{name} {density=}") + + +class RoutingIndexMathConfigTest(parameterized.TestCase): + """moe_routing_index_math="count" is rejected where no ring-of-experts ragged sort computes routing index maps.""" + + @parameterized.named_parameters( + # One case per requirement of the validator; each must be named in the error. + ("no_ragged_sort", {"use_ragged_sort": False}, "use_ragged_sort=True"), + ( + "no_ring_of_experts", + {"use_ring_of_experts": False, "num_moe_token_chunks": 1, "ring_of_experts_local_routing": False}, + "use_ring_of_experts=True", + ), + ("dense_matmul", {"sparse_matmul": False}, "sparse_matmul=True"), + ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"), + ) + def test_rejected_off_the_ring_ragged_sort(self, overrides, requirement): + with self.assertRaisesRegex(ValueError, f"moe_routing_index_math='count' requires: .*{re.escape(requirement)}"): + _cfg(moe_routing_index_math="count", **overrides) + + def test_accepted_on_the_ring_ragged_sort(self): + self.assertEqual(_cfg(moe_routing_index_math="count").moe_routing_index_math, "count") + + def test_names_every_unmet_requirement(self): + """The error names all unmet requirements, not only the first.""" + with self.assertRaises(ValueError) as raised: + _cfg(moe_routing_index_math="count", te_moe_block=True, use_ragged_sort=False) + self.assertIn("te_moe_block=False", str(raised.exception)) + self.assertIn("use_ragged_sort=True", str(raised.exception)) + + +def _cfg(**overrides): + """The shrunk Qwen3.5 routed-MoE config on the cp-as-ep mesh (moe_roe_local_routing_test.py's, local routing).""" + kwargs = { + "run_name": "moe_routing_index_math", + "enable_checkpointing": False, + "log_config": False, + "skip_jax_distributed_system": True, + "model_name": "qwen3.5-397b-a17b", + "override_model_config": True, + "num_experts": _NUM_EXPERTS, + "num_experts_per_tok": _TOP_K, + "base_emb_dim": _EMB, + "base_mlp_dim": 256, + "base_moe_mlp_dim": 256, + "max_target_length": _SEQ, + "max_prefill_predict_length": _SEQ, + "per_device_batch_size": 0.25, + "dtype": "bfloat16", + "weight_dtype": "float32", + "float32_gate_logits": True, + "megablox": True, + "sparse_matmul": True, + # Tokamax gmm v2 is TPU-only; both settings share whichever GMM runs, so CPU uses megablox. + "use_tokamax_gmm": False, + "use_gmm_v2": False, + "use_ring_of_experts": True, + "use_ragged_sort": True, + "use_custom_sort_vjp": False, + "ragged_buffer_factor": 2.0, + "num_moe_token_chunks": 2, + "custom_mesh_and_rule": "cp-as-ep", + "ici_fsdp_parallelism": 1, + "ici_context_parallelism": 4, + "ici_expert_parallelism": 2, + "context_parallel_load_balance": False, + "ring_of_experts_local_routing": True, + } + kwargs.update(overrides) + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs) + + +@pytest.mark.cpu_only +def test_count_matches_sort_on_cpu_mesh(): + """Runs `main` below in a child process with 8 CPU devices.""" + env = os.environ.copy() + env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" {_XLA_FLAGS}" + env["JAX_PLATFORMS"] = "cpu" + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + print(result.stdout) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + + +def _build(cfg, mesh): + return moe.RoutedMoE( + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.moe_mlp_dim, + weight_dtype=cfg.weight_dtype, + dtype=cfg.dtype, + rngs=nnx.Rngs(params=0), + ) + + +def _inputs(cfg, mode): + """The layer input, forced expert ids (padded and duplicate slots) and a token mask, per `mode`.""" + batch = int(cfg.per_device_batch_size * _NUM_DEVICES) + x = jax.random.normal(jax.random.PRNGKey(7), (batch, _SEQ, _EMB), jnp.float32) + x = x if cfg.float32_gate_logits else x.astype(jnp.bfloat16) + forced = token_mask = None + if "forced" in mode: + keys = jax.random.split(jax.random.PRNGKey(11), batch * _SEQ) + ids = jax.vmap(lambda k: jax.random.permutation(k, _NUM_EXPERTS)[:_TOP_K])(keys).reshape(batch, _SEQ, _TOP_K) + pos = jnp.arange(_SEQ)[None, :, None] + ids = jnp.where((pos % 5 == 0) & (jnp.arange(_TOP_K) == _TOP_K - 1), -1, ids) + ids = jnp.where((pos % 7 == 3) & (jnp.arange(_TOP_K) == 1), ids[..., :1], ids) + forced = ids.astype(jnp.int32) + if "padmask" in mode: + # Each sequence ends in a run of padding of a different length. + lens = jnp.array([_SEQ - 5 * (b + 1) for b in range(batch)]) + token_mask = (jnp.arange(_SEQ)[None, :] < lens[:, None]).astype(jnp.int32) + return x, forced, token_mask + + +def _count(jaxpr, name): + """Number of `name` equations in a jaxpr, including nested jaxprs.""" + total = 0 + for eqn in jaxpr.eqns: + total += eqn.primitive.name == name + for param in eqn.params.values(): + for sub in param if isinstance(param, (list, tuple)) else (param,): + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + if hasattr(sub, "eqns"): + total += _count(sub, name) + return total + + +def _count_argsorts(jaxpr): + """Number of `sort` equations with more than one operand: an argsort sorts its keys together with an iota. + + "count" sorts one packed key; the forced-routing duplicate check's `jnp.sort` is single-operand too. + """ + total = 0 + for eqn in jaxpr.eqns: + total += eqn.primitive.name == "sort" and len(eqn.invars) > 1 + for param in eqn.params.values(): + for sub in param if isinstance(param, (list, tuple)) else (param,): + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None) + if hasattr(sub, "eqns"): + total += _count_argsorts(sub) + return total + + +def _run(cfg, mesh, params, x, forced, token_mask): + """Runs the layer's value_and_grad; returns (loss, output, gradients), the overflow flag and the counts. + + The counts are the gradient's sort / top_k equations and the packed_preprocess flag of every + ragged_gather_reduce call traced. + """ + preprocess_flags = [] + real_gather_reduce = ragged_sort.ragged_gather_reduce + + def spy(*args, **kwargs): + preprocess_flags.append(kwargs.get("packed_preprocess", False)) + return real_gather_reduce(*args, **kwargs) + + with ( + mock.patch.object(ragged_sort, "ragged_gather_reduce", spy), + jax.set_mesh(mesh), + nn_partitioning.axis_rules(cfg.logical_axis_rules), + ): + graphdef, _, rest = nnx.split(_build(cfg, mesh), nnx.Param, ...) + + def loss_fn(params, x): + layer = nnx.merge(graphdef, params, rest) + out, lb_loss, _ = layer(x, forced_routed_experts=forced, token_mask=token_mask) + sown = nnx.state(layer, nnx.Intermediate).to_pure_dict() + flags = maxtext_utils.collect_intermediates_by_suffix(sown, "moe_has_overflow") + loss = jnp.mean(out.astype(jnp.float32) ** 2) + if lb_loss is not None: + loss = loss + lb_loss.astype(jnp.float32) + return loss, (out, jnp.any(jnp.concatenate([jnp.ravel(f) for f in flags]))) + + grad_fn = jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True) + jaxpr = jax.make_jaxpr(grad_fn)(params, x).jaxpr + (loss, (out, overflow)), grads = jax.jit(grad_fn)(params, x) + counts = { + "sort": _count(jaxpr, "sort"), + "argsort": _count_argsorts(jaxpr), + "top_k": _count(jaxpr, "top_k"), + "packed_preprocess": sorted(set(preprocess_flags)), + } + return jax.device_get((loss, out, grads)), bool(overflow), counts + + +def _bit_identical(label, old, new): + """Prints the per-leaf report; True iff the output, the loss and every gradient are bit-identical.""" + (loss_old, out_old, grads_old), (loss_new, out_new, grads_new) = old, new + named = [("output", out_old, out_new), ("loss", loss_old, loss_new)] + [ + (jax.tree_util.keystr(path), a, b) + for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_old)[0], jax.tree_util.tree_leaves(grads_new)) + ] + passed, lines = True, [] + for name, a, b in named: + a, b = np.asarray(a).astype(np.float32), np.asarray(b).astype(np.float32) + same = bool(np.array_equal(a, b)) and bool(np.isfinite(b).all()) and float(np.max(np.abs(a))) > 0 + passed &= same + lines.append(f" [{'ok' if same else 'FAIL'}] {name}: bit-identical={same} max|d|={np.max(np.abs(a - b)):.3e}") + print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True) + print("\n".join(lines), flush=True) + return passed + + +def _case(label, overrides, mode="real", control=None, expect_overflow=False): + """Runs "sort" and "count" on the same parameters; `control` is applied to the "count" run only.""" + cfg_sort = _cfg(**overrides) + cfg_count = _cfg(**overrides, moe_routing_index_math="count") + mesh = Mesh(maxtext_utils.create_device_mesh(cfg_sort), cfg_sort.mesh_axes) + x, forced, token_mask = _inputs(cfg_sort, mode) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_sort.logical_axis_rules): + _, params, _ = nnx.split(_build(cfg_sort, mesh), nnx.Param, ...) + old, overflow_old, counts_old = _run(cfg_sort, mesh, params, x, forced, token_mask) + with control if control is not None else contextlib.nullcontext(): + new, overflow_new, counts_new = _run(cfg_count, mesh, params, x, forced, token_mask) + structure = counts_new["sort"] < counts_old["sort"] + # "count" sorts only the packed key (one operand); "sort" argsorts with an iota operand. + structure &= counts_new["argsort"] == 0 < counts_old["argsort"] + # The SparseCore validity-sort flag, on every gather-reduce call (CPU runs the reference fallback, which ignores it). + structure &= counts_new["packed_preprocess"] == [True] and counts_old["packed_preprocess"] == [False] + if token_mask is not None: + structure &= counts_new["top_k"] == 0 < counts_old["top_k"] + overflow_ok = overflow_old == overflow_new == expect_overflow + print( + f"{label}: [{'ok' if structure else 'FAIL'}] sort/top_k equations sort={counts_old} count={counts_new};" + f" [{'ok' if overflow_ok else 'FAIL'}] overflow sort={overflow_old} count={overflow_new}" + f" (want {expect_overflow})", + flush=True, + ) + return _bit_identical(label, old, new) and structure and overflow_ok + + +def _shifted_local_positions(): + """Deliberately wrong: every local slot's counted position is one too far.""" + real = ragged_sort._local_slot_positions # pylint: disable=protected-access + + def wrong(*args, **kwargs): + pos = real(*args, **kwargs) + return jnp.where(pos >= 0, pos + 1, pos) + + return mock.patch.object(ragged_sort, "_local_slot_positions", wrong) + + +def _count_needs_forward_permutation(): + """The unsort backward of "count" raises without the forward permutation (its inverse is partial).""" + mesh = Mesh(np.array(jax.devices()[:2]), ("expert",)) + num_tokens, topk, num_experts = 8, 2, 4 + ids = jnp.asarray(np.stack([np.random.default_rng(t).permutation(num_experts)[:topk] for t in range(num_tokens)])) + + def unsort_loss(tokens): + def body(tokens, ids): + _, group_sizes, revert, _ = ragged_sort.ring_ragged_sort( + tokens, ids, num_experts, topk, "expert", 2, index_math="count" + ) + out = ragged_sort.ring_ragged_unsort( + tokens.repeat(topk, axis=0), + group_sizes, + revert, + topk, + num_experts // 2, + "expert", + topk_weights=jnp.ones((num_tokens * topk,), jnp.float32), + index_math="count", + ) + return jnp.sum(out.astype(jnp.float32)) + + return jax.shard_map(body, mesh=mesh, in_specs=(P(), P()), out_specs=P(), check_vma=False)(tokens, ids) + + try: + jax.grad(unsort_loss)(jnp.ones((num_tokens, 128), jnp.bfloat16)) + except ValueError as e: + raised = "needs topk_argsort_indices" in str(e) + else: + raised = False + print(f"count without the forward permutation: [{'ok' if raised else 'FAIL'}] the unsort backward raises", flush=True) + return raised + + +# The subset of the Qwen3.5-397B RL trainer benchmark's MoE settings that this file needs (the cases above do not +# all set them). +_BENCH_FLAGS = { + "moe_chunk_pipeline": True, + "moe_expert_weight_prefetch": "forward", + "ring_of_experts_row_major_reduce_scatter": True, + "moe_routing_maps": "device", + "ragged_buffer_factor": 2.75, + "prefuse_moe_weights": True, + "fp8_moe_fake_quant": True, + "moe_exclude_padding_from_routing": True, +} + + +def main(): + assert jax.device_count() == _NUM_DEVICES, jax.devices() + pipelined = {"moe_chunk_pipeline": True} + padded = {**pipelined, "moe_exclude_padding_from_routing": True} + results = { + "plain_loop": _case("plain_loop", {}), + "pipelined": _case("pipelined", pipelined), + "padding_excluded": _case("padding_excluded", padded, mode="padmask"), + "forced_dropping": _case( + "forced_dropping", {**padded, "ragged_buffer_factor": 1.0}, mode="forced_padmask", expect_overflow=True + ), + "dropless": _case("dropless", {"ragged_buffer_factor": -1.0}, mode="forced"), + "saved_routing_maps": _case("saved_routing_maps", {**pipelined, "moe_routing_maps": "device"}), + "bf16_router": _case("bf16_router", {**padded, "float32_gate_logits": False}, mode="padmask"), + "bench_flags": _case("bench_flags", _BENCH_FLAGS, mode="padmask"), + "guard": _count_needs_forward_permutation(), + } + controls = { + "control_shifted_local_positions": _case( + "control_shifted_local_positions", pipelined, control=_shifted_local_positions() + ), + } + print("SUMMARY", {**results, **controls}, flush=True) + failed = [k for k, v in results.items() if not v] + undetected = [k for k, v in controls.items() if v] + if failed or undetected: + print(f"cases that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True) + sys.exit(1) + print(_PASS_TOKEN, flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/moe_wi_gmm_fused_swiglu_test.py b/tests/unit/moe_wi_gmm_fused_swiglu_test.py new file mode 100644 index 0000000000..09076e8853 --- /dev/null +++ b/tests/unit/moe_wi_gmm_fused_swiglu_test.py @@ -0,0 +1,554 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""moe_wi_gmm_fused_swiglu on CPU: megablox ops.gmm_silu's own code, its wiring in the MoE layer, and the validator. + +With the key, the up projection calls megablox `gmm_silu` (the tokamax gmm_v2 kernel with the SwiGLU in its epilogue) +instead of `gmm` followed by the activation. The kernels are Pallas TPU kernels and do not run on CPU; their own test, +megablox_gmm_silu_test.py, is TPU-only (not yet run on TPU). Two CPU checks cover everything around them: + +ops.gmm_silu (in-process, `GmmSiluVjpTest`). Only the Pallas entry `gmm_v2.gmm_v2` and gmm's TPU backward +`ops._gmm_bwd` are replaced: the first by a jnp grouped matmul that honours fuse_act and return_pre_act and records the +tile it is given, the second by a stand-in that records how it is called and passes its cotangent through as the input +gradient. The rest is the real ops.py: gmm_silu, its custom VJP, the forward tile (tiling[2] halved: the fused kernel's +tile_n counts the columns of each half), the residuals, and the backward, which must hand gmm's backward [d gate | +d up] of silu(gate) * up on the bf16 halves, bit for bit, with this op's DLHS options bound to the right keyword +arguments. + +The layer (a subprocess, so the forced device count applies): the shrunk Qwen3.5 routed MoE of +moe_roe_local_routing_test.py with prefuse_moe_weights on 8 CPU devices (cp-as-ep, ragged sort, local routing, two +token chunks). `gmm` is replaced by a grouped matmul over the same group sizes and group offset, and `gmm_silu` by the +unfused composition of that same stand-in, i.e. silu(gate) * up of its bf16 output. The layer with the key must then +reproduce the layer without it bit for bit (the output, the loss and every gradient): that checks what the key +changes in moe.py (the rows and the routing operands handed to the fused GMM, the padding to the forward tile and its +removal, the whole [w0 | w1] weight, and the skipped activation), not the kernel. The fused kernel's real numerics +differ by rounding (silu reads the f32 accumulator). Cases: the plain and the pipelined chunk loop, a dropless buffer +(where each shard's group offset is not 0), padding excluded from routing, the whole prefused wi +(moe_unsplit_prefused_wi), with it and the other MoE keys (moe_routing_index_math="count", +moe_roe_packed_routing_gather), and with those and the benchmark's MoE settings +(expert-weight prefetch, row-major reduce-scatter, routing maps on device, ragged_buffer_factor 2.75, fp8 fake-quant, +moe_gmm_v2_dlhs_transpose_rhs). Two-sided: the fused path must trace exactly one gmm_silu call per chunk and the plain +one none; gmm_silu must get the routing operands, the wi tiling the plain path gives gmm, and the DLHS options of the +config (moe_gmm_v2_dlhs_transpose_rhs is True in the benchmark case and False elsewhere); stand-ins that swap the gate +and up halves, or that ignore the group offset, must FAIL. + +Key off, not compared with anything this change touches: the full-buffer ring-of-experts ragged sort (dropless), whose +GMM routing operands get_gmm_operands_for_local_experts now builds for both paths, must match the ring of experts +without the ragged sort (use_ragged_sort=False), and a GMM that ignores the group offset must fail that comparison. + +The validator rejects the key for each requirement it does not meet, naming it. +""" + +import os +import re +import subprocess +import sys +from unittest import mock + +from absl.testing import parameterized +from flax import nnx +from flax.linen import partitioning as nn_partitioning +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +import numpy as np +import pytest + +from maxtext.configs import pyconfig +from maxtext.kernels import megablox +from maxtext.kernels.megablox import ops as megablox_ops +from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_gmm_kernel as gmm_v2 +from maxtext.layers import moe +from maxtext.layers.initializers import nd_dense_init +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + +_NUM_DEVICES = 8 +_NUM_EXPERTS = 16 +_TOP_K = 4 +_EMB = 256 +_SEQ = 64 +_PASS_TOKEN = "MOE_WI_GMM_FUSED_SWIGLU_CHECKS_PASSED" +_XLA_FLAGS = f"--xla_force_host_platform_device_count={_NUM_DEVICES} --xla_allow_excess_precision=false" +# The ragged-sort and plain ring paths combine the same expert outputs, possibly in another order: bf16 rounding. +_UNSORTED_REL_L2 = 1e-2 + + +def _cfg(**overrides): + """The shrunk Qwen3.5 routed-MoE config with prefused weights and the tokamax gmm_v2 path the key requires.""" + kwargs = { + "run_name": "moe_wi_gmm_fused_swiglu", + "enable_checkpointing": False, + "log_config": False, + "skip_jax_distributed_system": True, + "model_name": "qwen3.5-397b-a17b", + "override_model_config": True, + "num_experts": _NUM_EXPERTS, + "num_experts_per_tok": _TOP_K, + "base_emb_dim": _EMB, + "base_mlp_dim": 256, + "base_moe_mlp_dim": 256, + "max_target_length": _SEQ, + "max_prefill_predict_length": _SEQ, + "per_device_batch_size": 0.25, + "dtype": "bfloat16", + "weight_dtype": "float32", + "float32_gate_logits": True, + "megablox": True, + "sparse_matmul": True, + "use_tokamax_gmm": True, + "use_gmm_v2": True, + "prefuse_moe_weights": True, + "use_ring_of_experts": True, + "use_ragged_sort": True, + "use_custom_sort_vjp": False, + "ragged_buffer_factor": 2.0, + "num_moe_token_chunks": 2, + "custom_mesh_and_rule": "cp-as-ep", + "ici_fsdp_parallelism": 1, + "ici_context_parallelism": 4, + "ici_expert_parallelism": 2, + "context_parallel_load_balance": False, + "ring_of_experts_local_routing": True, + } + kwargs.update(overrides) + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs) + + +class FusedSwigluConfigTest(parameterized.TestCase): + """The key is rejected for each requirement of the fused gmm_v2 epilogue it does not meet.""" + + @parameterized.named_parameters( + ("not_prefused", {"prefuse_moe_weights": False}, "prefuse_moe_weights=True"), + ("dense", {"sparse_matmul": False}, "sparse_matmul=True"), + ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"), + ("vllm_rpa", {"attention": "vllm_rpa"}, "an attention other than vllm_rpa*"), + ("no_gmm_v2", {"use_gmm_v2": False}, "use_tokamax_gmm=True and use_gmm_v2=True"), + ("no_tokamax", {"use_tokamax_gmm": False}, "use_tokamax_gmm=True and use_gmm_v2=True"), + ("quantization", {"quantization": "int8"}, "no weight quantization"), + ("gelu", {"mlp_activations": ["gelu", "linear"]}, "mlp_activations=['silu', 'linear']"), + ("gpt_oss", {"decoder_block": "gpt_oss"}, "a decoder_block other than gpt_oss"), + ("clipped", {"mlp_activations_limit": 7.0}, "mlp_activations_limit <= 0 (no clipping)"), + ("mlp_bias", {"mlp_bias": True}, "mlp_bias=False"), + ("emb_chunks", {"num_moe_emb_chunks": 2}, "num_moe_emb_chunks=0"), + ("flat_fsdp", {"moe_flat_fsdp_weights": True}, "moe_flat_fsdp_weights=False"), + ("tc_3d_gmm", {"moe_tc_ragged_3d_gmm": True}, "moe_tc_ragged_3d_gmm=False"), + ("quantized_token_ag", {"moe_quantize_token_all_gather": True}, "moe_quantize_token_all_gather=False"), + ("chunk_wgrad", {"moe_accumulate_chunk_wgrad": True}, "moe_accumulate_chunk_wgrad=False"), + ("minimal_remat", {"remat_policy": "minimal"}, "remat_policy=minimal saves it as mlpwi_0"), + ("saved_gate", {"remat_policy": "custom", "mlpwi_0": "device"}, "remat_policy=custom saves it as mlpwi_0"), + ("offloaded_moe_up", {"remat_policy": "custom", "moe_mlpwi_1": "offload"}, "saves it as moe_mlpwi_1"), + ) + def test_rejected(self, overrides, requirement): + with self.assertRaisesRegex(ValueError, f"moe_wi_gmm_fused_swiglu=True requires: .*{re.escape(requirement)}"): + _cfg(moe_wi_gmm_fused_swiglu=True, **overrides) + + def test_fused_activation_rejected_next_to_it(self): + with self.assertRaisesRegex(ValueError, "moe_prefused_swiglu_vjp=True requires: .*moe_wi_gmm_fused_swiglu=False"): + _cfg(moe_wi_gmm_fused_swiglu=True, moe_prefused_swiglu_vjp=True) + + def test_accepted(self): + self.assertTrue(_cfg(moe_wi_gmm_fused_swiglu=True, moe_unsplit_prefused_wi=True).moe_wi_gmm_fused_swiglu) + # remat_policy=full saves nothing, and the shared expert's mlpwi_1 is not a routed-expert tensor. + for overrides in ({"remat_policy": "full"}, {"remat_policy": "custom", "mlpwi_1": "device"}): + self.assertTrue(_cfg(moe_wi_gmm_fused_swiglu=True, **overrides).moe_wi_gmm_fused_swiglu) + + +_OP_M, _OP_K, _OP_N = 64, 256, 128 # the op case: lhs [M, K], rhs [groups, K, 2N]; K == 2N, so a pass-through dlhs fits +_OP_GROUP_SIZES = np.array([20, 0, 30, 10], np.int32) # sum 60 < M: padding rows past the last group +_OP_TILING = (64, 256, 256, 64, 256, 256, 64, 256, 256) # gmm's 9-tuple; the fused forward tile is (64, 256, 128) + + +def _grouped_f32(lhs, rhs, group_sizes, group_offset): + """f32 lhs @ rhs[local expert of row]: rows of the groups [group_offset, group_offset + groups), others 0.""" + ends = jnp.cumsum(group_sizes) + local = jnp.sum(jnp.arange(lhs.shape[0])[:, None] >= ends[None, :], axis=1) - jnp.reshape(group_offset, ()) + out = jnp.zeros((lhs.shape[0], rhs.shape[-1]), jnp.float32) + for e in range(rhs.shape[0]): + out = out + jnp.where((local == e)[:, None], lhs.astype(jnp.float32) @ rhs[e].astype(jnp.float32), 0.0) + return out + + +@pytest.mark.cpu_only +class GmmSiluVjpTest(parameterized.TestCase): + """ops.gmm_silu around stand-ins for the TPU kernel and gmm's TPU backward (see the module docstring).""" + + @parameterized.product( + group_offset=(0, 1), + # (use_dlhs_transpose_rhs, fuse_dlhs_scale, fuse_3d_dlhs_ps): each alone, so a swapped binding shows. + dlhs_options=((False, False, False), (True, False, False), (False, True, False), (False, False, True)), + heuristic=(False, True), + ) + def test_custom_vjp(self, group_offset, dlhs_options, heuristic): + dlhs_transpose, fuse_scale, fuse_3d = dlhs_options + seen = {"tile_n": [], "return_pre_act": [], "bwd": []} + + def kernel_stand_in( + lhs, + rhs, + group_sizes, + tile_info, + preferred_element_type, + group_offset=None, + fuse_act=None, + return_pre_act=False, + **kwargs, + ): + del kwargs + self.assertEqual(fuse_act, "silu") + seen["tile_n"].append(tile_info if callable(tile_info) else tile_info.tile_n) + seen["return_pre_act"].append(return_pre_act) + acc = _grouped_f32(lhs, rhs, group_sizes, 0 if group_offset is None else group_offset) + n = acc.shape[-1] // 2 + act = (jax.nn.silu(acc[:, :n]) * acc[:, n:]).astype(preferred_element_type) + if not return_pre_act: + return act + return act, acc[:, :n].astype(preferred_element_type), acc[:, n:].astype(preferred_element_type) + + def gmm_bwd_stand_in(**kwargs): # keyword arguments only: ops._gmm_silu_bwd must name every one + seen["bwd"].append(kwargs) + lhs, rhs = kwargs["residual"][:2] + return kwargs["grad"][0].astype(lhs.dtype), jnp.zeros_like(rhs), None, None, None, None, None + + k0, k1, k2 = jax.random.split(jax.random.key(group_offset), 3) + lhs = jax.random.normal(k0, (_OP_M, _OP_K), jnp.float32).astype(jnp.bfloat16) + rhs = (jax.random.normal(k1, (len(_OP_GROUP_SIZES) - group_offset, _OP_K, 2 * _OP_N)) * 0.05).astype(jnp.bfloat16) + cotangent = jax.random.normal(k2, (_OP_M, _OP_N), jnp.float32).astype(jnp.bfloat16) + group_sizes, offset = jnp.asarray(_OP_GROUP_SIZES), jnp.array([group_offset], jnp.int32) + + def gmm_silu(l, r): + return megablox_ops.gmm_silu( + l, + r, + group_sizes, + preferred_element_type=jnp.bfloat16, + tiling=_OP_TILING, + group_offset=offset, + use_gmm_v2_heuristic_tiling=heuristic, + use_dlhs_transpose_rhs=dlhs_transpose, + fuse_dlhs_scale=fuse_scale, + fuse_3d_dlhs_ps=fuse_3d, + ) + + with ( + mock.patch.object(gmm_v2, "gmm_v2", kernel_stand_in), + mock.patch.object(megablox_ops, "_gmm_bwd", gmm_bwd_stand_in), + ): + gmm_silu(lhs, rhs) # not differentiated: the activation alone + _, vjp = jax.vjp(gmm_silu, lhs, rhs) + dlhs, _ = vjp(cotangent) + + # The forward tile: tiling[2] halved (or gmm_v2's heuristic), every call; act alone undifferentiated, act + gate + + # up differentiated. + self.assertEqual(seen["tile_n"], [gmm_v2.calculate_tiling if heuristic else _OP_TILING[2] // 2] * 2) + self.assertEqual(seen["return_pre_act"], [False, True]) + # The backward: gmm's backward gets [d gate | d up] of silu(gate) * up on the bf16 halves, bit for bit. + acc = _grouped_f32(lhs, rhs, group_sizes, offset) + gate, up = acc[:, :_OP_N].astype(jnp.bfloat16), acc[:, _OP_N:].astype(jnp.bfloat16) + _, swiglu_vjp = jax.vjp(lambda g, u: (jax.nn.silu(g) * u).astype(g.dtype), gate, up) + want = jnp.concatenate(swiglu_vjp(cotangent), axis=-1) + np.testing.assert_array_equal(np.asarray(dlhs, np.float32), np.asarray(want, np.float32)) + self.assertLen(seen["bwd"], 1) + bwd = seen["bwd"][0] + np.testing.assert_array_equal(np.asarray(bwd["grad"][0], np.float32), np.asarray(want, np.float32)) + self.assertEqual(bwd["grad"][1:], (None, None)) + # Its residuals are this op's operands, and its options are this op's (the plain bf16 tokamax gmm_v2 GEMM). + res_lhs, res_rhs, res_group_sizes, res_offset, *res_rest = bwd["residual"] + for got, want_operand in ((res_lhs, lhs), (res_rhs, rhs), (res_group_sizes, group_sizes), (res_offset, offset)): + np.testing.assert_array_equal(np.asarray(got), np.asarray(want_operand)) + self.assertEqual(tuple(res_rest), (None, False, False, None)) + expected_options = { + "lhs_dtype": jnp.bfloat16, + "rhs_dtype": jnp.bfloat16, + "preferred_element_type": jnp.bfloat16, + "tiling": _OP_TILING, + "transpose_rhs": False, + "quantization_rule": None, + "use_tokamax_backend": True, + "use_gmm_v2": True, + "use_gmm_v2_heuristic_tiling": heuristic, + "return_rhs": False, + "return_lhs": False, + "fuse_dlhs_scale": fuse_scale, + "use_dlhs_transpose_rhs": dlhs_transpose, + "out_is_3d": False, + "fuse_3d_dlhs_ps": fuse_3d, + "tc_unsort_cfg": None, + } + self.assertEqual({name: bwd[name] for name in expected_options}, expected_options) + + def test_odd_forward_tile_is_rejected(self): + lhs = jnp.zeros((_OP_M, _OP_K), jnp.bfloat16) + rhs = jnp.zeros((len(_OP_GROUP_SIZES), _OP_K, 2 * _OP_N), jnp.bfloat16) + with self.assertRaisesRegex(ValueError, "must halve to a multiple of 128"): + megablox_ops.gmm_silu(lhs, rhs, jnp.asarray(_OP_GROUP_SIZES), jnp.bfloat16, (64, 256, 192) + _OP_TILING[3:]) + + +@pytest.mark.cpu_only +def test_fused_swiglu_wiring_on_cpu_mesh(): + """Runs `main` below in a child process with 8 CPU devices.""" + env = os.environ.copy() + env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" {_XLA_FLAGS}" + env["JAX_PLATFORMS"] = "cpu" + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + print(result.stdout) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + + +def _grouped_matmul(lhs, rhs, group_sizes, group_offset): + """f32 lhs @ rhs[local expert of row]: rows of the groups [group_offset, group_offset + groups), others 0.""" + ends = jnp.cumsum(group_sizes) + local = jnp.sum(jnp.arange(lhs.shape[0])[:, None] >= ends[None, :], axis=1) - group_offset + out = jnp.zeros((lhs.shape[0], rhs.shape[-1]), jnp.float32) + for e in range(rhs.shape[0]): + out = out + jnp.where((local == e)[:, None], lhs.astype(jnp.float32) @ rhs[e].astype(jnp.float32), 0.0) + return out + + +def _gmm_stand_in(lhs, rhs, group_sizes, preferred_element_type=jnp.float32, tiling=None, group_offset=None, **kwargs): + """jnp stand-in for megablox gmm (the TPU kernel), for the calls the layer makes here.""" + del tiling + for name in ("partial_sum", "return_lhs", "return_rhs"): + assert not kwargs.get(name), name + offset = 0 if group_offset is None else group_offset + return _grouped_matmul(lhs, rhs, group_sizes, offset).astype(preferred_element_type) + + +def _gmm_silu_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset=None, **kwargs): + """megablox gmm_silu as the unfused composition of the gmm stand-in: silu(gate) * up of its rounded output.""" + del kwargs + out = _gmm_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset) + n = out.shape[-1] // 2 + return (jax.nn.silu(out[:, :n]) * out[:, n:]).astype(preferred_element_type) + + +def _halves_swapped_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset=None, **kwargs): + """Deliberately wrong: treats the up half as the gate and the gate half as up.""" + del kwargs + out = _gmm_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset) + n = out.shape[-1] // 2 + return (jax.nn.silu(out[:, n:]) * out[:, :n]).astype(preferred_element_type) + + +def _offset_ignored_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset=None, **kwargs): + """Deliberately wrong: ignores the group offset, i.e. computes the first shard's experts on every shard.""" + del group_offset + return _gmm_silu_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, None, **kwargs) + + +def _build(cfg, mesh): + return moe.RoutedMoE( + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.moe_mlp_dim, + weight_dtype=cfg.weight_dtype, + dtype=cfg.dtype, + rngs=nnx.Rngs(params=0), + ) + + +def _inputs(cfg, mode): + batch = int(cfg.per_device_batch_size * _NUM_DEVICES) + x = jax.random.normal(jax.random.PRNGKey(7), (batch, _SEQ, _EMB), jnp.float32) + token_mask = None + if mode == "padmask": + lens = jnp.array([_SEQ - 5 * (b + 1) for b in range(batch)]) + token_mask = (jnp.arange(_SEQ)[None, :] < lens[:, None]).astype(jnp.int32) + return x, token_mask + + +def _run(cfg, mesh, params, x, token_mask, gmm_silu, gmm=None): + """Loss, output and gradients with the stand-ins; the keyword arguments of every gmm_silu call traced, and the + tiling of the plain path's up-projection gmm calls ([K, 2N] weight).""" + silu_calls, wi_tilings = [], [] + gmm = gmm or _gmm_stand_in + + def recorded_gmm_silu(*args, **kwargs): + silu_calls.append(kwargs) + return gmm_silu(*args, **kwargs) + + def recorded_gmm(lhs, rhs, *args, **kwargs): + if rhs.shape[-1] == 2 * cfg.moe_mlp_dim: + wi_tilings.append(tuple(kwargs["tiling"])) + return gmm(lhs, rhs, *args, **kwargs) + + with ( + mock.patch.object(megablox, "gmm", recorded_gmm), + mock.patch.object(megablox, "gmm_silu", recorded_gmm_silu), + jax.set_mesh(mesh), + nn_partitioning.axis_rules(cfg.logical_axis_rules), + ): + graphdef, _, rest = nnx.split(_build(cfg, mesh), nnx.Param, ...) + + def loss_fn(params, x): + out, lb_loss, _ = nnx.merge(graphdef, params, rest)(x, token_mask=token_mask) + loss = jnp.mean(out.astype(jnp.float32) ** 2) + if lb_loss is not None: + loss = loss + lb_loss.astype(jnp.float32) + return loss, out + + jax.make_jaxpr(loss_fn)(params, x) + traced = list(silu_calls) + (loss, out), grads = jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(params, x) + return jax.device_get((loss, out, grads)), traced, sorted(set(wi_tilings)) + + +def _bit_identical(label, old, new): + """Prints the per-leaf report; True iff the output, the loss and every gradient are bit-identical.""" + (loss_old, out_old, grads_old), (loss_new, out_new, grads_new) = old, new + named = [("output", out_old, out_new), ("loss", loss_old, loss_new)] + [ + (jax.tree_util.keystr(path), a, b) + for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_old)[0], jax.tree_util.tree_leaves(grads_new)) + ] + passed, lines = True, [] + for name, a, b in named: + a, b = np.asarray(a).astype(np.float32), np.asarray(b).astype(np.float32) + same = bool(np.array_equal(a, b)) and bool(np.isfinite(b).all()) and float(np.max(np.abs(a))) > 0 + passed &= same + lines.append(f" [{'ok' if same else 'FAIL'}] {name}: bit-identical={same} max|d|={np.max(np.abs(a - b)):.3e}") + print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True) + print("\n".join(lines), flush=True) + return passed + + +def _case(label, overrides, mode="real", gmm_silu=_gmm_silu_stand_in): + """The layer without and with the key on the same parameters; `gmm_silu` is the fused GMM's stand-in.""" + cfg_old = _cfg(**overrides) + cfg_new = _cfg(**overrides, moe_wi_gmm_fused_swiglu=True) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg_old), cfg_old.mesh_axes) + x, token_mask = _inputs(cfg_old, mode) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_old.logical_axis_rules): + _, params, _ = nnx.split(_build(cfg_old, mesh), nnx.Param, ...) + old, calls_old, wi_tilings = _run(cfg_old, mesh, params, x, token_mask, gmm_silu) + new, calls_new, _ = _run(cfg_new, mesh, params, x, token_mask, gmm_silu) + structure = not calls_old and len(calls_new) == cfg_new.num_moe_token_chunks + # gmm_silu gets the tiling the plain path gives the up-projection gmm, and the config's DLHS options. + want = { + "tiling": wi_tilings[0] if len(wi_tilings) == 1 else wi_tilings, + "use_gmm_v2_heuristic_tiling": cfg_new.use_gmm_v2_heuristic_tiling, + "use_dlhs_transpose_rhs": cfg_new.moe_gmm_v2_dlhs_transpose_rhs, + "fuse_dlhs_scale": cfg_new.moe_accumulate_wi_dlhs, + "fuse_3d_dlhs_ps": cfg_new.moe_3d_dlhs_partial_sum, + "preferred_element_type": jnp.bfloat16, + } + got = [{name: tuple(c[name]) if name == "tiling" else c[name] for name in want} for c in calls_new] + bindings = bool(got) and all(g == want for g in got) + print( + f"{label}: [{'ok' if structure else 'FAIL'}] gmm_silu calls traced plain={len(calls_old)}" + f" fused={len(calls_new)}; [{'ok' if bindings else 'FAIL'}] gmm_silu arguments {got[:1]} (want {want})", + flush=True, + ) + return _bit_identical(label, old, new) and structure and bindings + + +def _gmm_offset_ignored_stand_in( + lhs, rhs, group_sizes, preferred_element_type=jnp.float32, tiling=None, group_offset=None, **kwargs +): + """Deliberately wrong: a gmm that ignores the group offset, i.e. computes the first shard's experts everywhere.""" + del group_offset + return _gmm_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, None, **kwargs) + + +def _ragged_dropless_matches_unsorted(label, ragged_gmm=None): + """Key off: the full-buffer (dropless) ring-of-experts ragged sort against the ring of experts without it. + + Both take their GMM routing operands from get_gmm_operands_for_local_experts, but only the ragged path uses a + non-zero group offset on a full buffer. Not a comparison of the key: a reference the key does not touch. + `ragged_gmm` replaces the gmm stand-in in the ragged run only. + """ + cfg_ragged = _cfg(ragged_buffer_factor=-1.0) + cfg_unsorted = _cfg(ragged_buffer_factor=-1.0, use_ragged_sort=False) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg_ragged), cfg_ragged.mesh_axes) + x, token_mask = _inputs(cfg_ragged, "real") + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_ragged.logical_axis_rules): + _, params, _ = nnx.split(_build(cfg_ragged, mesh), nnx.Param, ...) + reference, _, _ = _run(cfg_unsorted, mesh, params, x, token_mask, _gmm_silu_stand_in) + ragged, _, _ = _run(cfg_ragged, mesh, params, x, token_mask, _gmm_silu_stand_in, gmm=ragged_gmm) + (loss_ref, out_ref, grads_ref), (loss_new, out_new, grads_new) = reference, ragged + named = [("output", out_ref, out_new), ("loss", loss_ref, loss_new)] + [ + (jax.tree_util.keystr(path), a, b) + for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_ref)[0], jax.tree_util.tree_leaves(grads_new)) + ] + passed, lines = True, [] + for name, a, b in named: + a, b = np.asarray(a, np.float32), np.asarray(b, np.float32) + rel = float(np.linalg.norm(b - a) / max(np.linalg.norm(a), 1e-30)) + ok = rel < _UNSORTED_REL_L2 and bool(np.isfinite(b).all()) and float(np.max(np.abs(a))) > 0 + passed &= ok + lines.append(f" [{'ok' if ok else 'FAIL'}] {name}: rel_l2 vs use_ragged_sort=False {rel:.2e}") + print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True) + print("\n".join(lines), flush=True) + return passed + + +# The subset of the Qwen3.5-397B RL trainer benchmark's MoE settings that this file needs (the cases above do not +# all set them). +_BENCH_FLAGS = { + "moe_chunk_pipeline": True, + "moe_expert_weight_prefetch": "forward", + "ring_of_experts_row_major_reduce_scatter": True, + "moe_routing_maps": "device", + "ragged_buffer_factor": 2.75, + "fp8_moe_fake_quant": True, + "moe_exclude_padding_from_routing": True, + "moe_gmm_v2_dlhs_transpose_rhs": True, +} + + +def main(): + assert jax.device_count() == _NUM_DEVICES, jax.devices() + pipelined = {"moe_chunk_pipeline": True} + other_keys = { + "moe_unsplit_prefused_wi": True, + "moe_routing_index_math": "count", + "moe_roe_packed_routing_gather": True, + } + results = { + "plain_loop": _case("plain_loop", {}), + "pipelined": _case("pipelined", pipelined), + "dropless": _case("dropless", {"ragged_buffer_factor": -1.0}), + "padding_excluded": _case("padding_excluded", {**pipelined, "moe_exclude_padding_from_routing": True}, "padmask"), + "unsplit_wi": _case("unsplit_wi", {**pipelined, "moe_unsplit_prefused_wi": True}), + "with_other_keys": _case("with_other_keys", {**pipelined, **other_keys}), + "with_other_keys_bench_flags": _case("with_other_keys_bench_flags", {**_BENCH_FLAGS, **other_keys}, "padmask"), + "key_off_ragged_dropless_matches_unsorted": _ragged_dropless_matches_unsorted( + "key_off_ragged_dropless_matches_unsorted" + ), + } + controls = { + "control_halves_swapped": _case("control_halves_swapped", pipelined, gmm_silu=_halves_swapped_stand_in), + "control_offset_ignored": _case( + "control_offset_ignored", {"ragged_buffer_factor": -1.0}, gmm_silu=_offset_ignored_stand_in + ), + "control_key_off_ragged_offset_ignored": _ragged_dropless_matches_unsorted( + "control_key_off_ragged_offset_ignored", ragged_gmm=_gmm_offset_ignored_stand_in + ), + } + print("SUMMARY", {**results, **controls}, flush=True) + failed = [k for k, v in results.items() if not v] + undetected = [k for k, v in controls.items() if v] + if failed or undetected: + print(f"cases that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True) + sys.exit(1) + print(_PASS_TOKEN, flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index 816a1c8b29..51e4ce3cfb 100644 --- a/tests/unit/pyconfig_test.py +++ b/tests/unit/pyconfig_test.py @@ -257,6 +257,32 @@ def test_gdn_packing_allowed_with_kernel_and_pure_jax(self): self.assertTrue(cfg_rl.use_gdn_kernel) self.assertEqual(cfg_rl.max_seq_token_per_tpu, 4096) + def test_gdn_kernel_rejects_packed_rows_without_packing(self): + """With packing=False the kernel skips its segment path, so inputs that still pack rows must say packing=True.""" + + def init_rl(**kwargs): + return pyconfig.initialize( + ["", get_post_train_test_config_path("rl")], + skip_jax_distributed_system=True, + config_class=config_types.RLConfig, + model_name="qwen3-next-80b-a3b", + max_target_length=1024, + use_gdn_kernel=True, + max_seq_token_per_tpu=4096, + **kwargs, + ) + + with self.assertRaisesRegex(ValueError, "max_seq_token_per_tpu > 0 requires packing=True"): + init_rl(packing=False) + self.assertTrue(init_rl(packing=True).packing) + + mmap = {"dataset_type": "grain", "grain_file_type": "mmap_npy", "grain_train_files": "/idx|/data/prefix"} + with self.assertRaisesRegex(ValueError, "grain_file_type='mmap_npy' and reset_attention_mask=True requires packing"): + self._init_qwen3_next_packing(packing=False, use_gdn_kernel=True, **mmap) + self.assertTrue(self._init_qwen3_next_packing(packing=True, use_gdn_kernel=True, **mmap).packing) + self._init_qwen3_next_packing(packing=False, use_gdn_kernel=True, reset_attention_mask=False, **mmap) + self._init_qwen3_next_packing(packing=False, use_gdn_kernel=False, **mmap) + def test_gdn_kernel_fast_fail_constraints(self): """Verifies fast-fail ValueErrors for unsupported GDN kernel configurations.""" # use_gdn_kernel=True with tensor parallelism > 1 raises ValueError on both MaxTextConfig and RLConfig. @@ -308,14 +334,25 @@ def test_gdn_kernel_fast_fail_constraints(self): ) self.assertEqual(cfg_cp4.max_target_length // cfg_cp4.ici_context_parallelism, cfg_cp4.gdn_chunk_size) - # gdn_chunk_size != 64 or head_dim % 128 != 0 raises ValueError when use_gdn_kernel=True. - with self.assertRaisesRegex(ValueError, "requires gdn_chunk_size=64"): - self._init_qwen3_next_packing( - packing=False, - use_gdn_kernel=True, - override_model_config=True, - gdn_chunk_size=32, - ) + # gdn_chunk_size not in (64, 128) or head_dim % 128 != 0 raises ValueError when use_gdn_kernel=True. + for bad_chunk_size in (32, 256): + with self.assertRaisesRegex(ValueError, r"requires gdn_chunk_size in \(64, 128\)"): + self._init_qwen3_next_packing( + packing=False, + use_gdn_kernel=True, + override_model_config=True, + gdn_chunk_size=bad_chunk_size, + ) + cfg_c128 = self._init_qwen3_next_packing( + packing=False, + use_gdn_kernel=True, + override_model_config=True, + gdn_chunk_size=128, + max_target_length=256, + ici_context_parallelism=2, + context_parallel_load_balance=False, + ) + self.assertEqual(cfg_c128.gdn_chunk_size, 128) with self.assertRaisesRegex(ValueError, "multiples of 128"): self._init_qwen3_next_packing( packing=False,