diff --git a/src/maxtext/configs/post_train/rl.yml b/src/maxtext/configs/post_train/rl.yml index f1f245bc58..d372e6978c 100644 --- a/src/maxtext/configs/post_train/rl.yml +++ b/src/maxtext/configs/post_train/rl.yml @@ -213,6 +213,7 @@ math_verify_num_procs: null # ====== Special tokens/templates for GSM8K reasoning ====== reasoning_start_token: '' reasoning_end_token: '' +reasoning_start_token_in_prompt: false solution_start_token: '' solution_end_token: '' data_template_path: 'maxtext/examples/chat_templates/gsm8k_rl.json' diff --git a/src/maxtext/configs/post_train/rl_gsm8k_qwen35_35b_v5p64.yml b/src/maxtext/configs/post_train/rl_gsm8k_qwen35_35b_v5p64.yml new file mode 100644 index 0000000000..8b5713d603 --- /dev/null +++ b/src/maxtext/configs/post_train/rl_gsm8k_qwen35_35b_v5p64.yml @@ -0,0 +1,159 @@ +# 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. + +base_config: "rl.yml" + +# ====== Model ====== +model_name: qwen3.5-35b-a3b +tokenizer_path: Qwen/Qwen3.5-35B-A3B +tokenizer_type: huggingface +scan_layers: true + +# ====== Hardware ====== +# v5p-64 = v5p-8 (4 chips) * 8 nodes +# Inference: 4 nodes +# Training: 4 nodes +trainer_devices_fraction: 0.5 +sampler_devices_fraction: 0.5 +chips_per_vm: 4 +use_pathways: true +allow_split_physical_axes: true + +# Rollout: TP = 4 (across 4 chips). +use_standalone_converter: false +rollout_data_parallelism: -1 +rollout_tensor_parallelism: 4 +rollout_expert_parallelism: 1 + +# Use the MaxText vLLM adapter so actor and rollout share the same Qwen3.5 +# parameter layout. These overrides were previously supplied only by the +# experiment JobSet. +vllm_hf_overrides: + architectures: ["MaxTextForCausalLM"] +vllm_additional_config: + maxtext_config: + model_name: qwen3.5-35b-a3b + model_call_mode: inference + attention: vllm_rpa + allow_split_physical_axes: true + log_config: false + weight_dtype: bfloat16 + prefuse_moe_weights: true + +# One checkpoint restore plus an in-memory clone avoids reading/converting the +# same model twice during startup. +load_checkpoint_only_once: true + +# ====== GRPO ====== +rl: + num_generations: 4 + num_iterations: 1 + grpo_beta: 0.08 + grpo_epsilon: 0.2 + loss_algo: "grpo" + loss_agg_mode: "sequence-mean-token-mean" + use_agentic_rollout: false + +# ====== Training Schedule ====== +# GSM8K train yields 934 full batches at batch_size=8 (7,472 examples, with +# drop_remainder=True). One epoch covers the full training split once. +batch_size: 8 +num_batches: 934 +num_epoch: 1 +train_fraction: 1.0 +learning_rate_schedule_steps: 934 + +train_micro_batch_size: 1 +rollout_micro_batch_size: 8 + +learning_rate: 3e-6 +warmup_steps_fraction: 0.1 +adam_b1: 0.9 +adam_b2: 0.99 +adam_weight_decay: 0.1 +gradient_clipping_threshold: 0.1 + +log_period: 20 +eval_interval: 100 + +# ====== Evaluation ====== +# 20 * 32 = 640 held-out GSM8K prompts per eval pass. +num_test_batches: 20 +eval_batch_size: 32 +# Qwen3.5 is a thinking model; greedy decoding can fall into endless token +# loops. Match the model's published sampling defaults. +eval_sampling_strategy: "standard" +generation_configs: + standard: + eval_temperature: 1.0 + eval_top_k: 20 + eval_top_p: 0.95 +num_eval_passes: 1 +eval_mode: "pass_at_1" + +# ====== Rollout / Generation ====== +max_prefill_predict_length: 256 +max_target_length: 1024 +kv_cache_buffer: 256 + +decode_sampling_temperature: 1.0 +decode_sampling_top_k: 20 +decode_sampling_nucleus_p: 0.95 + +hbm_utilization_vllm: 0.60 +swap_space_vllm_gb: 2 +max_num_seqs: 32 +max_num_batched_tokens: 16384 +async_scheduling: false +enable_dp_attention: false +# The MaxText Qwen3.5 adapter does not save and restore GDN recurrent state in +# vLLM's block-addressed prefix cache. Reusing only the attention prefix would +# pair it with unrelated recurrent state and corrupt generation. +enable_prefix_caching: false + +# ====== Checkpointing ====== +enable_checkpointing: true +async_checkpointing: false +checkpoint_period: 250 +max_num_checkpoints_to_keep: 4 + +# ====== Dataset ====== +dataset_name: "openai/gsm8k" +eval_dataset_name: "openai/gsm8k" +train_split: "train" +eval_split: "test" +hf_subset: "main" +data_template_path: "maxtext/examples/chat_templates/qwen35_math_rl.json" +reasoning_start_token: "" +reasoning_end_token: "" +reasoning_start_token_in_prompt: true + +# Qwen3.5's chat template opens its native block in the prompt. The +# model closes it, then emits the task-specific answer block requested above. +stop_strings: [""] + +# ====== Reward ====== +reward_exact_answer: 1.0 +reward_white_space_format_match: 1.0 +reward_exact_format_match: 0.1 +reward_partial_format_match: 0.0 +reward_ratio_guess_to_answer_high: 0.0 +reward_ratio_guess_to_answer_low: 0.0 +penalty_incorrect_format: 0.0 +penalty_incorrect_answer: 0.0 + +math_verify_timeout: 120 +math_verify_num_procs: null + +debug: false diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index c4a270a567..b5235a7276 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2439,6 +2439,10 @@ class RLSpecialTokens(BaseModel): reasoning_start_token: str = Field("", description="Token to mark the beginning of a reasoning section.") reasoning_end_token: str = Field("", description="Token to mark the end of a reasoning section.") + reasoning_start_token_in_prompt: bool = Field( + False, + description="Whether the chat template prefilled the reasoning start token, so it is absent from the completion.", + ) solution_start_token: str = Field("", description="Token to mark the beginning of a solution section.") solution_end_token: str = Field("", description="Token to mark the end of a solution section.") @@ -3942,12 +3946,14 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de class RLConfig( LogitsAndLoss, Engram, + ManifoldConstrainedHyperConnections, RematAndOffload, Attention, Llama4Attention, LayoutAndSharding, InferenceLayout, InferenceGeneral, + PrefixCaching, Decoding, IciParallelism, DcnParallelism, diff --git a/src/maxtext/examples/chat_templates/qwen35_math_rl.json b/src/maxtext/examples/chat_templates/qwen35_math_rl.json new file mode 100644 index 0000000000..280ced4beb --- /dev/null +++ b/src/maxtext/examples/chat_templates/qwen35_math_rl.json @@ -0,0 +1,4 @@ +{ + "SYSTEM_PROMPT": "Solve the problem step by step. After your reasoning, place only the final numerical answer between {solution_start_token} and {solution_end_token}.", + "TEMPLATE": "{system_prompt}\n\n{question}" +} diff --git a/src/maxtext/integration/vllm/_hybrid_cache.py b/src/maxtext/integration/vllm/_hybrid_cache.py new file mode 100644 index 0000000000..c2edc99d21 --- /dev/null +++ b/src/maxtext/integration/vllm/_hybrid_cache.py @@ -0,0 +1,38 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hybrid cache-layout helpers for MaxText's vLLM adapter.""" + +import math +from typing import Any + + +def build_qwen_gdn_cache_layout(cfg: Any, torch_module: Any): + """Returns the shapes, dtypes, and unpadded bytes for a Qwen GDN cache.""" + key_dim = cfg.gdn_key_head_dim * cfg.gdn_num_key_heads + value_dim = cfg.gdn_value_head_dim * cfg.gdn_num_value_heads + conv_dim = key_dim * 2 + value_dim + + shapes = ( + (cfg.gdn_conv_kernel_dim - 1, conv_dim), + (cfg.gdn_num_value_heads, cfg.gdn_key_head_dim, cfg.gdn_value_head_dim), + ) + # This is the TPU Inference / upstream vLLM contract regardless of model + # weight or attention-KV dtype. + dtypes = (torch_module.bfloat16, torch_module.float32) + page_size_bytes = sum( + math.prod(shape) * torch_module.empty((), dtype=dtype).element_size() + for shape, dtype in zip(shapes, dtypes, strict=True) + ) + return shapes, dtypes, page_size_bytes diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py index bacd28e551..229f460577 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py @@ -23,6 +23,7 @@ from jax.sharding import Mesh from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE from maxtext.configs import pyconfig +from maxtext.integration.vllm._hybrid_cache import build_qwen_gdn_cache_layout from maxtext.utils import lora_utils from maxtext.utils import max_logging from maxtext.utils import model_creation_utils @@ -378,7 +379,6 @@ def patch_kv_cache_manager(): from tpu_inference.runner.kv_cache_manager import KVCacheManager from vllm.v1.kv_cache_interface import MambaSpec import torch - import numpy as np except ImportError as e: # Gracefully handle missing imports in standard JAX environments (e.g. unit tests on CPU) max_logging.log(f"Skipping KVCacheManager patch (tpu_inference or dependencies not installed): {e}") @@ -414,31 +414,12 @@ def patched_get_kv_cache_spec(self): if decoder_block_str in ("qwen3_next", "qwen3_5"): interval = cfg.inhomogeneous_layer_cycle_interval - num_v_heads = cfg.gdn_num_value_heads - num_k_heads = cfg.gdn_num_key_heads - head_k_dim = cfg.gdn_key_head_dim - head_v_dim = cfg.gdn_value_head_dim - conv_kernel_size = cfg.gdn_conv_kernel_dim - - key_dim = head_k_dim * num_k_heads - value_dim = head_v_dim * num_v_heads - conv_dim = key_dim * 2 + value_dim - - conv_state_shape = (conv_kernel_size - 1, conv_dim) - recurrent_state_shape = (num_v_heads, head_k_dim, head_v_dim) - - mamba_shapes = (conv_state_shape, recurrent_state_shape) - - torch_dtype = torch.bfloat16 - if str(cfg.dtype) == "float32": - torch_dtype = torch.float32 - elif str(cfg.dtype) == "float16": - torch_dtype = torch.float16 - mamba_dtypes = (torch_dtype, torch_dtype) - - # Calculate unpadded mamba page size - dtype_size = 4 if torch_dtype == torch.float32 else 2 - unpadded_mamba_page_size = sum(int(np.prod(shape)) * dtype_size for shape in mamba_shapes) + # Qwen GDN keeps its short convolution history in BF16, but recurrence is + # accumulated and persisted in FP32. Declaring both caches as the model + # dtype silently quantizes the recurrent state after every generated token. + mamba_shapes, mamba_dtypes, unpadded_mamba_page_size = build_qwen_gdn_cache_layout( + cfg, torch + ) # Calculate attn_page_size_bytes from tpu_inference.layers.common.sharding import ShardingAxisName diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py index 636eee4d56..fc1e7e7f95 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py @@ -14,19 +14,21 @@ """MaxText-specific VllmSampler and VllmRollout subclasses. -These replace the Tunix built-in key-mapping path with model-specific -MaxText to vLLM converters, which handle: +These extend Tunix weight synchronization with both model-specific native-vLLM +converters and scanned MaxText-to-MaxText state unrolling. The converters handle: - QKV fusion with GQA interleaving (attention) - MoE expert gate+up fusion (w13_weight chunk-interleaved for TP) - MoE gate / down transpose - Layer-norm and LM-head transposes """ -from typing import Any, Optional, Tuple - +import copy import gc import logging +import re import time +from typing import Any, Optional, Tuple + import jax import jax.numpy as jnp from flax import nnx @@ -55,6 +57,57 @@ def _create_model_converter(model_name: str, config: Any, mesh: jax.sharding.Mes return None +def uses_maxtext_vllm_adapter(config: Any) -> bool: + """Returns whether vLLM is configured to instantiate MaxTextForCausalLM.""" + overrides = getattr(config, "vllm_hf_overrides", None) + if isinstance(overrides, str): + return "MaxTextForCausalLM" in overrides + if isinstance(overrides, dict): + architectures = overrides.get("architectures", ()) + if isinstance(architectures, str): + architectures = (architectures,) + return "MaxTextForCausalLM" in architectures + return False + + +def requires_maxtext_scanned_weight_unroll(config: Any) -> bool: + """Returns whether direct MaxText-to-MaxText sync needs a custom unroll.""" + return bool(getattr(config, "scan_layers", False) and uses_maxtext_vllm_adapter(config)) + + +def prepare_direct_sync_additional_config( + additional_config: Optional[dict[str, Any]], + *, + direct_maxtext_sync: bool, + num_experts: int, + tensor_parallel_size: int, +) -> Optional[dict[str, Any]]: + """Makes the direct MaxText MoE target use TPU-safe prefused weights. + + TPU inference shards the fused gate/up dimension across tensor-parallel + devices. Each shard must therefore contain its local gate chunk followed by + its local up chunk. The unfused MaxText inference path concatenates the two + complete tensors globally, which gives incorrect local shards when TP > 1. + Tunix's direct-sync MoE fusion builds the required per-shard layout once at + weight-load time when the target exposes a prefused ``wi`` parameter. + """ + if not direct_maxtext_sync or num_experts <= 1 or tensor_parallel_size <= 1: + return additional_config + + prepared = copy.deepcopy(additional_config) if additional_config is not None else {} + maxtext_overrides = prepared.setdefault("maxtext_config", {}) + if not isinstance(maxtext_overrides, dict): + raise ValueError("vLLM additional_config.maxtext_config must be a dictionary for direct MaxText sync.") + + if not maxtext_overrides.get("prefuse_moe_weights", False): + logging.info( + "MaxTextVllmRollout: enabling prefuse_moe_weights for correct MoE gate/up layout with TP=%d.", + tensor_parallel_size, + ) + maxtext_overrides["prefuse_moe_weights"] = True + return prepared + + def _find_scanned_layer_idx(key_tuple, container_names=("layers", "scanned_blocks", "layers_remainder")): """Returns (container_idx, container_name) if a scanned layer structure is found, else (-1, None).""" for name in container_names: @@ -64,6 +117,155 @@ def _find_scanned_layer_idx(key_tuple, container_names=("layers", "scanned_block return -1, None +def _find_qwen_scanned_layer_idx(key_tuple): + """Finds a Qwen heterogeneous scanned block path like `layers.layer_0`.""" + for i in range(len(key_tuple) - 1): + if key_tuple[i] != "layers" or not isinstance(key_tuple[i + 1], str): + continue + match = re.fullmatch(r"layer_(\d+)", key_tuple[i + 1]) + if match: + return i, int(match.group(1)) + return -1, -1 + + +def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Optional[int] = None): + """Unroll Qwen's heterogeneous scanned blocks for an unscanned MaxText target. + + Qwen 3 Next/3.5 training stores a repeating layer cycle as + `decoder.layers.layer_{slot}`, with repetitions stacked on `scan_axis`. + The inference model stores every layer as a direct decoder attribute named + `layers_{global_index}`. Tunix's generic direct-sync mapper cannot bridge + these two structures and otherwise silently leaves all destination layers at + their random initialization. + """ + if hasattr(weights, "filter") and hasattr(weights, "to_pure_dict"): + # NNX stacks non-parameter state (notably RNG state) on axis 0 even when + # parameters use param_scan_axis=1. Only parameters belong in weight sync. + pure_dict = weights.filter(nnx.Param).to_pure_dict() + elif hasattr(weights, "to_pure_dict"): + pure_dict = weights.to_pure_dict() + elif hasattr(weights, "to_dict"): + pure_dict = weights.to_dict() + elif isinstance(weights, dict): + pure_dict = weights + else: + return weights + + flat_w = flatten_dict(pure_dict) + scanned_keys = [] + slot_indices = set() + scan_lengths = set() + for key, value in flat_w.items(): + container_idx, slot_idx = _find_qwen_scanned_layer_idx(key) + if container_idx == -1 or "dropout" in key or "rngs" in key: + continue + if not hasattr(value, "shape") or len(value.shape) <= scan_axis: + raise ValueError(f"Qwen scanned parameter {'.'.join(key)} has no scan axis {scan_axis}: {value!r}") + scanned_keys.append((key, value, container_idx, slot_idx)) + slot_indices.add(slot_idx) + scan_lengths.add(value.shape[scan_axis]) + + if not scanned_keys: + return weights + + if pattern_length is None: + expected_slots = set(range(max(slot_indices) + 1)) + if slot_indices != expected_slots: + raise ValueError( + "Qwen scanned layer slots must be contiguous when pattern_length is omitted; " + f"found {sorted(slot_indices)}" + ) + pattern_length = len(slot_indices) + elif pattern_length <= max(slot_indices): + raise ValueError( + f"Qwen scanned layer slot {max(slot_indices)} is outside configured pattern length {pattern_length}" + ) + if len(scan_lengths) != 1: + raise ValueError(f"Qwen scanned parameters disagree on scan length: {sorted(scan_lengths)}") + + scan_length = scan_lengths.pop() + scanned_key_paths = {key for key, _, _, _ in scanned_keys} + new_flat_w = {key: value for key, value in flat_w.items() if key not in scanned_key_paths} + + for key, value, container_idx, slot_idx in scanned_keys: + prefix = key[:container_idx] + suffix = key[container_idx + 2 :] + for repetition in range(scan_length): + global_idx = repetition * pattern_length + slot_idx + new_key = prefix + (f"layers_{global_idx}",) + suffix + new_flat_w[new_key] = jnp.take(value, repetition, axis=scan_axis) + + logging.info( + "MaxTextVllmSampler: unrolled %d Qwen tensor components across %d layers for direct MaxText weight sync.", + len(scanned_keys), + scan_length * pattern_length, + ) + return unflatten_dict(new_flat_w) + + +def validate_direct_sync_layer_coverage(source, target) -> int: + """Fail if an unrolled source would leave MaxText target layers untouched. + + Tunix intentionally intersects direct-sync trees. For heterogeneous Qwen + scans, a schema error can therefore skip every transformer layer without an + exception. This check runs on the initial full-parameter load and requires + every unscanned target-layer parameter path to exist in the source. + """ + + def to_pure_params(state): + if hasattr(state, "filter") and hasattr(state, "to_pure_dict"): + return state.filter(nnx.Param).to_pure_dict() + if hasattr(state, "to_pure_dict"): + return state.to_pure_dict() + if hasattr(state, "to_dict"): + return state.to_dict() + return state + + def unwrap(state, wrapper): + while isinstance(state, dict) and wrapper in state: + state = state[wrapper] + return state + + source = unwrap(to_pure_params(source), "base") + target = unwrap(to_pure_params(target), "model") + if not isinstance(source, dict) or not isinstance(target, dict): + return 0 + + source_flat = flatten_dict(source) + target_flat = flatten_dict(target) + + def is_unscanned_layer_path(path): + return any(isinstance(part, str) and re.fullmatch(r"layers_\d+", part) for part in path) + + source_layer_keys = {key for key in source_flat if is_unscanned_layer_path(key)} + target_layer_keys = {key for key in target_flat if is_unscanned_layer_path(key)} + + def source_covers(target_key): + if target_key in source_layer_keys: + return True + # Tunix fuses split training weights into the inference-only prefused + # parameter before transfer. Treat the pair as coverage for target `wi`. + if target_key and target_key[-1] == "wi": + prefix = target_key[:-1] + return prefix + ("wi_0",) in source_layer_keys and prefix + ("wi_1",) in source_layer_keys + return False + + missing = {key for key in target_layer_keys if not source_covers(key)} + if not target_layer_keys or missing: + examples = [".".join(map(str, key)) for key in sorted(missing)[:5]] + raise ValueError( + "Direct MaxText weight sync would leave rollout transformer parameters at random initialization: " + f"matched {len(target_layer_keys) - len(missing)}/{len(target_layer_keys)} target layer parameters; " + f"missing examples: {examples}" + ) + + logging.info( + "MaxTextVllmSampler: verified direct-sync coverage for all %d rollout layer parameters.", + len(target_layer_keys), + ) + return len(target_layer_keys) + + def unroll_gemma_scanned_weights(weights): """Workaround for tunix unstacking bug with Gemma 3/4 scanned blocks. @@ -158,8 +360,8 @@ class MaxTextVllmSampler(VllmSampler): When a converter is supplied, update_params bypasses transfer_state_with_mappings entirely and instead runs converter.convert() followed by a direct device_put - into the vLLM model-runner state dict. If no converter is supplied the base-class - behaviour is preserved, so this class is safe to use as a drop-in replacement. + into the vLLM model-runner state dict. If no converter is supplied the base-class + path is used after MaxText's heterogeneous scanned states are unrolled. """ def __init__( @@ -167,9 +369,15 @@ def __init__( tokenizer: Any, config: VllmConfig, converter: Any = None, + direct_maxtext_sync: bool = False, + scan_axis: int = 1, + layer_pattern_length: Optional[int] = None, ): super().__init__(tokenizer=tokenizer, config=config) self._converter = converter + self._direct_maxtext_sync = direct_maxtext_sync + self._scan_axis = scan_axis + self._layer_pattern_length = layer_pattern_length def update_params( self, @@ -178,7 +386,15 @@ def update_params( ): """Update the vLLM runner weights from a MaxText state tree.""" if self._converter is None: - updated_weights = unroll_gemma_scanned_weights(updated_weights) + if self._direct_maxtext_sync: + updated_weights = unroll_qwen_scanned_weights( + updated_weights, + scan_axis=self._scan_axis, + pattern_length=self._layer_pattern_length, + ) + updated_weights = unroll_gemma_scanned_weights(updated_weights) + if filter_types is None: + validate_direct_sync_layer_coverage(updated_weights, self.transformer_state) super().update_params(updated_weights, filter_types) return None @@ -245,10 +461,10 @@ def update_params( class MaxTextVllmRollout(vllm_rollout.VllmRollout): - """VllmRollout that uses MaxTextVllmSampler for weight synchronisation. + """VllmRollout that uses MaxTextVllmSampler for weight synchronization. - The extra `maxtext_config` argument is forwarded to the model-specific converter - together with `mesh`. All other arguments mirror VllmRollout.__init__. + The extra `maxtext_config` selects either a native-vLLM converter or direct + MaxText adapter synchronization. All other arguments mirror VllmRollout.__init__. Usage (direct): rollout = MaxTextVllmRollout( @@ -281,7 +497,15 @@ def __init__( if cache_config_or_size is None: cache_config_or_size = rollout_config.kv_cache_size - converter = _create_model_converter(maxtext_config.model_name, config=maxtext_config, mesh=mesh) + # Native vLLM models need explicit MaxText-to-HF conversion. The MaxText + # adapter instead has the same tensor layouts as the actor and uses direct + # structural sync after scanned layers are unrolled above. + direct_maxtext_sync = uses_maxtext_vllm_adapter(maxtext_config) + converter = ( + None + if direct_maxtext_sync + else _create_model_converter(maxtext_config.model_name, config=maxtext_config, mesh=mesh) + ) mapping_config = mappings.MappingConfig.build( mapping_obj=rollout_config.rollout_mapping_config, @@ -291,16 +515,27 @@ def __init__( engine_kwargs = { "max_model_len": cache_config_or_size, "model": rollout_config.rollout_vllm_model_version, - "swap_space": getattr(rollout_config, "rollout_vllm_swap_space_size_gb", maxtext_config.swap_space_vllm_gb), # Async scheduling causes KeyError in dp_scheduler on slow models # (30B+) where inference latency exceeds the scheduler's window. "async_scheduling": rollout_config.rollout_vllm_async_scheduling, + "max_num_batched_tokens": rollout_config.rollout_vllm_max_num_batched_tokens, + "max_num_seqs": rollout_config.rollout_vllm_max_num_seqs, + "hf_config_path": rollout_config.rollout_vllm_hf_config_path, + "max_logprobs": 1, + "logprobs_mode": rollout_config.rollout_vllm_logprobs_mode, } # Merge additional kwargs like dtype and hf_overrides provided by train_rl.py if hasattr(rollout_config, "rollout_vllm_kwargs") and rollout_config.rollout_vllm_kwargs: engine_kwargs.update(rollout_config.rollout_vllm_kwargs) + rollout_additional_config = prepare_direct_sync_additional_config( + getattr(rollout_config, "rollout_vllm_additional_config", None), + direct_maxtext_sync=direct_maxtext_sync, + num_experts=getattr(maxtext_config, "num_experts", 1), + tensor_parallel_size=rollout_config.tensor_parallel_size, + ) + self._sampler = MaxTextVllmSampler( tokenizer=tokenizer, config=VllmConfig( # pylint: disable=unexpected-keyword-arg,no-value-for-parameter @@ -311,15 +546,25 @@ def __init__( mapping_config=mapping_config, lora_config=rollout_config.rollout_vllm_lora_config, server_mode=rollout_config.rollout_vllm_server_mode, + server_mode_submission_threshold=rollout_config.rollout_vllm_server_mode_submission_threshold, + server_mode_submission_timeout_s=rollout_config.rollout_vllm_server_mode_submission_timeout_s, + return_logprobs=rollout_config.return_logprobs, tensor_parallel_size=rollout_config.tensor_parallel_size, data_parallel_size=rollout_config.data_parallel_size, + expert_parallel_size=rollout_config.expert_parallel_size, enable_dp_attention=rollout_config.rollout_vllm_enable_dp_attention, + delete_dst_buffers=rollout_config.rollout_vllm_delete_dst_buffers, + reshard_chunk_size=rollout_config.rollout_vllm_reshard_chunk_size, engine_kwargs=engine_kwargs, - additional_config=getattr(rollout_config, "rollout_vllm_additional_config", None), + additional_config=rollout_additional_config, + sampling_kwargs=rollout_config.rollout_vllm_sampling_kwargs, ), converter=converter, + direct_maxtext_sync=direct_maxtext_sync, + scan_axis=getattr(maxtext_config, "param_scan_axis", 1), + layer_pattern_length=getattr(maxtext_config, "inhomogeneous_layer_cycle_interval", None), ) # Initial weight sync: run the converter so vLLM starts with real weights. - state = nnx.state(rollout_actor) + state = nnx.state(rollout_actor, nnx.Param) self._sampler.load_checkpoint(state) diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index 5a65b35ba4..7e0aaba831 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -851,6 +851,11 @@ def init_rotary_embedding(self): cast_as_fprop_dtype=True, fprop_dtype=self.dtype, mrope_section=self.mrope_section, + partial_rotary_factor=( + self.partial_rotary_factor + if self.partial_rotary_factor is not None + else self.config.partial_rotary_factor + ), rngs=self.rngs, ) diff --git a/src/maxtext/layers/embeddings.py b/src/maxtext/layers/embeddings.py index 05bc1fa193..8376d78b52 100644 --- a/src/maxtext/layers/embeddings.py +++ b/src/maxtext/layers/embeddings.py @@ -1784,6 +1784,7 @@ def __init__( cast_as_fprop_dtype: bool = True, fprop_dtype: DType = jnp.bfloat16, mrope_section: tuple[int, int, int] | None = None, + partial_rotary_factor: float = 1.0, attention_scaling: float = 1.0, rngs: nnx.Rngs = None, ): @@ -1792,19 +1793,30 @@ def __init__( Args: min_timescale: Start of the geometric index (typically 1). max_timescale: End of the geometric index (rope_theta, e.g., 1000000). - embedding_dims: Dimension of the embedding (head_dim). + embedding_dims: Dimension of the attention head. cast_as_fprop_dtype: Whether to cast output to fprop dtype. fprop_dtype: The dtype of the output. mrope_section: Tuple of (temporal_dim, height_dim, width_dim) for MRoPE. Defaults to [24, 20, 20] if None. + partial_rotary_factor: Fraction of the head dimensions to rotate. The + remaining suffix is passed through unchanged. attention_scaling: Scaling factor applied to cos/sin embeddings. Defaults to 1.0. rngs: rng keys passed in by nnx.bridge.to_linen. """ + if partial_rotary_factor is None or not 0.0 < partial_rotary_factor <= 1.0: + raise ValueError(f"partial_rotary_factor must be in (0, 1], got {partial_rotary_factor}.") + + self.head_dim = embedding_dims + self.partial_rotary_factor = partial_rotary_factor + self.rotary_dim = int(self.head_dim * self.partial_rotary_factor) + if self.rotary_dim <= 0 or self.rotary_dim % 2: + raise ValueError("Rotary dim for rotary position embedding must be a positive multiple of 2.") + super().__init__( min_timescale=min_timescale, max_timescale=max_timescale, mesh=None, - embedding_dims=embedding_dims, + embedding_dims=self.rotary_dim, cast_as_fprop_dtype=cast_as_fprop_dtype, fprop_dtype=fprop_dtype, rngs=rngs, @@ -1812,8 +1824,11 @@ def __init__( self.mrope_section = mrope_section if mrope_section is not None else (24, 20, 20) self.attention_scaling = attention_scaling - if self.embedding_dims % 2: - raise ValueError("Embedding dim for rotary position embedding must be a multiple of 2.") + if sum(self.mrope_section) != self.rotary_dim // 2: + raise ValueError( + f"mrope_section must describe rotary_dim / 2 frequencies; got {self.mrope_section} " + f"for rotary_dim={self.rotary_dim}." + ) def _apply_interleaved_mrope(self, freqs: jax.Array) -> jax.Array: """Apply interleaved MRoPE pattern to 3D rotary embeddings. @@ -1822,16 +1837,16 @@ def _apply_interleaved_mrope(self, freqs: jax.Array) -> jax.Array: interleaved [THTHWHTHW...], preserving frequency continuity. Args: - freqs: Shape (3, batch, seq_len, head_dim // 2) + freqs: Shape (3, batch, seq_len, rotary_dim // 2) Dimension 0: temporal frequencies Dimension 1: height frequencies Dimension 2: width frequencies Returns: - freqs_t: Shape (batch, seq_len, head_dim // 2) with interleaved pattern + freqs_t: Shape (batch, seq_len, rotary_dim // 2) with interleaved pattern """ # Start with temporal frequencies (dimension 0) - freqs_t = freqs[0] # (batch, seq_len, head_dim // 2) + freqs_t = freqs[0] # (batch, seq_len, rotary_dim // 2) # Create interleaved pattern # For each spatial dimension (H, W), place frequencies at positions: @@ -1854,7 +1869,9 @@ def __call__( """Generates rotary position embeddings for multimodal sequences. Args: - inputs: Input tensor of shape [batch, sequence, heads, head_dim]. + inputs: Input tensor of shape [batch, sequence, heads, head_dim]. MRoPE + is applied to the first ``rotary_dim`` features and the rest are + returned unchanged. position: Position IDs with shape: - [batch, sequence] for text-only (2D) - [3, batch, sequence] for multimodal with vision (3D) @@ -1865,9 +1882,9 @@ def __call__( """ if len(inputs.shape) != 4: raise ValueError("Input is assumed to be a rank 4 tensor of shape [batch, sequence, heads, head_dim].") - if self.embedding_dims != inputs.shape[3]: + if self.head_dim != inputs.shape[3]: raise ValueError( - "The embedding dims of the rotary position embedding must match the hidden dimension of the inputs." + "The head dim of the rotary position embedding must match the hidden dimension of the inputs." ) # Handle both 2D (text-only) and 3D (multimodal) position IDs @@ -1877,25 +1894,27 @@ def __call__( elif position.ndim != 3 or position.shape[0] != 3: raise ValueError(f"Position IDs must be 2D (batch, seq) or 3D (3, batch, seq), got shape {position.shape}") - # Compute frequencies: (3, batch, seq, 1) @ (head_dim // 2, 1) -> (3, batch, seq, head_dim // 2) - inv_freq_expanded = (1.0 / self.timescale)[jnp.newaxis, jnp.newaxis, jnp.newaxis, :] # (1, 1, 1, head_dim//2) + # Compute frequencies over the rotated prefix only. + inv_freq_expanded = (1.0 / self.timescale)[jnp.newaxis, jnp.newaxis, jnp.newaxis, :] position_expanded = position[..., jnp.newaxis] # (3, batch, seq, 1) - freqs = position_expanded * inv_freq_expanded # (3, batch, seq, head_dim//2) + freqs = position_expanded * inv_freq_expanded # (3, batch, seq, rotary_dim // 2) # Apply interleaved MRoPE pattern for 3D positions - freqs = self._apply_interleaved_mrope(freqs) # (batch, seq, head_dim//2) + freqs = self._apply_interleaved_mrope(freqs) # (batch, seq, rotary_dim // 2) # Compute sin and cos - # Concatenate to get full head_dim: (batch, seq, head_dim//2) -> (batch, seq, head_dim) - emb = jnp.concatenate([freqs, freqs], axis=-1) # Duplicate for both halves - cos_emb = jnp.cos(emb) * self.attention_scaling # (batch, seq, head_dim) - sin_emb = jnp.sin(emb) * self.attention_scaling # (batch, seq, head_dim) + # Duplicate frequencies for the two halves of the rotated prefix. + emb = jnp.concatenate([freqs, freqs], axis=-1) + cos_emb = jnp.cos(emb) * self.attention_scaling + sin_emb = jnp.sin(emb) * self.attention_scaling - # Expand for heads dimension: (batch, seq, head_dim) -> (batch, seq, 1, head_dim) + # Expand for the heads dimension. cos_emb = cos_emb[:, :, jnp.newaxis, :] sin_emb = sin_emb[:, :, jnp.newaxis, :] - x_out = self.apply_rotary(inputs, cos_emb, sin_emb) + inputs_rotary, inputs_pass = jnp.split(inputs, [self.rotary_dim], axis=-1) + rotated = self.apply_rotary(inputs_rotary, cos_emb, sin_emb) + x_out = jnp.concatenate([rotated, inputs_pass], axis=-1) if self.cast_as_fprop_dtype: x_out = x_out.astype(self.fprop_dtype) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index b310d97fb5..6ffef35f8c 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -42,8 +42,15 @@ from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils import maxtext_utils -from maxtext.utils.sharding import create_sharding, maybe_shard_with_logical, maybe_shard_with_pspec -from maxtext.utils.sharding import logical_to_mesh_axes, remove_expert_from_partition_spec, get_logical_axis_rules +from maxtext.utils.sharding import ( + create_sharding, + get_logical_axis_rules, + logical_to_mesh_axes, + maybe_shard_with_logical, + maybe_shard_with_pspec, + remove_expert_from_partition_spec, + remove_incompatible_mesh_axes_from_partition_spec, +) import numpy as np import qwix from qwix.contrib.sparsity import sparsity_module @@ -1647,6 +1654,48 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): decoder_tokens_pspec, ) = get_routed_moe_shardings(is_batch_sharded_by_expert, input_ids is not None) w0_pspec, w1_pspec, wo_pspec = maybe_aqt_partition(w0_kernel, w0_pspec, w1_kernel, w1_pspec, wo_kernel, wo_pspec) + output_pspec = self._logical_to_mesh_axes( + ( + batch_logical_axis, + "activation_norm_length", + "activation_embed", + ) + ) + # Replicating the batch is only safe when expert routing does not depend on + # batch shards. Keep the existing strict behavior for expert-parallel meshes. + if self.get_expert_parallelism_size() == 1: + input_partition_pspec = remove_incompatible_mesh_axes_from_partition_spec( + input_partition_pspec, + inputs.shape, + self.mesh, + dims=(0,), + ) + gate_logits_pspec = remove_incompatible_mesh_axes_from_partition_spec( + gate_logits_pspec, + gate_logits.shape, + self.mesh, + dims=(0,), + ) + if pre_bias_logits_pspec is not None: + pre_bias_logits_pspec = remove_incompatible_mesh_axes_from_partition_spec( + pre_bias_logits_pspec, + pre_bias_logits.shape, + self.mesh, + dims=(0,), + ) + if decoder_tokens_pspec is not None: + decoder_tokens_pspec = remove_incompatible_mesh_axes_from_partition_spec( + decoder_tokens_pspec, + input_ids.shape, + self.mesh, + dims=(0,), + ) + output_pspec = remove_incompatible_mesh_axes_from_partition_spec( + output_pspec, + inputs.shape, + self.mesh, + dims=(0,), + ) def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None): # The ring-of-experts strategy first duplicates the inputs to all @@ -2234,13 +2283,7 @@ def _moe_body( P(), # Replicate the input key ), out_specs=( - self._logical_to_mesh_axes( - ( - batch_logical_axis, - "activation_norm_length", - "activation_embed", - ) - ), + output_pspec, P(), # Handle None or replicate the output P(), # Handle None or replicate the output ), @@ -2341,18 +2384,9 @@ def sparse_matmul_route_and_compute( w1_kernel = self._maybe_shard_with_logical(w1_kernel, ("exp_with_fsdp", None, "mlp_no_fsdp")) wo_kernel = self._maybe_shard_with_logical(wo_kernel, ("exp_with_fsdp", "mlp_no_fsdp", None)) - input_axes = (batch_logical_axis, "activation_norm_length", None) - - gate_logits_axes = (batch_logical_axis, "activation_norm_length", None) - # NOTE: deepseek2 has a different pattern - if self.config.model_name.startswith(("deepseek3", "deepseek4")): - pre_bias_logits_axes = (batch_logical_axis, "activation_norm_length", None) - else: - pre_bias_logits_axes = None - - inputs = self._maybe_shard_with_logical(inputs, input_axes) - gate_logits = self._maybe_shard_with_logical(gate_logits, gate_logits_axes) - pre_bias_logits = self._maybe_shard_with_logical(pre_bias_logits, pre_bias_logits_axes) + inputs = self._maybe_shard_with_pspec(inputs, input_partition_pspec) + gate_logits = self._maybe_shard_with_pspec(gate_logits, gate_logits_pspec) + pre_bias_logits = self._maybe_shard_with_pspec(pre_bias_logits, pre_bias_logits_pspec) w0_kernel = self._maybe_shard_with_pspec(w0_kernel, w0_pspec) w1_kernel = self._maybe_shard_with_pspec(w1_kernel, w1_pspec) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 7cb710bf1e..76ce8bf0cd 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -33,7 +33,11 @@ from maxtext.common.common_types import AttentionType, Config, DType, Array, BATCH, EMBED, MODEL_MODE_TRAIN, LENGTH, MODEL_MODE_AUTOREGRESSIVE from maxtext.common.common_types import KV_BATCH, KV_HEAD -from maxtext.utils.sharding import logical_to_mesh_axes, get_logical_axis_rules +from maxtext.utils.sharding import ( + get_logical_axis_rules, + logical_to_mesh_axes, + remove_incompatible_mesh_axes_from_partition_spec, +) from maxtext.layers import attentions from maxtext.layers import initializers as max_initializers from maxtext.layers import moe @@ -614,6 +618,13 @@ def __call__( if self.mesh is not None: logical_rules = get_logical_axis_rules() qkvz_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) + # Training microbatches can be smaller than the physical KV_BATCH mesh partition. + qkvz_pspec = remove_incompatible_mesh_axes_from_partition_spec( + qkvz_pspec, + mixed_qkvz.shape, + self.mesh, + dims=(0,), + ) qkvz_sharding = jax.sharding.NamedSharding(self.mesh, qkvz_pspec) mixed_qkvz = jax.lax.with_sharding_constraint(mixed_qkvz, qkvz_sharding) @@ -660,7 +671,10 @@ def __call__( try: from tpu_inference.layers.common.gdn_attention import run_jax_gdn_attention # pylint: disable=import-outside-toplevel # pytype: disable=import-error from tpu_inference.layers.common.sharding import ShardingAxisName # pylint: disable=import-outside-toplevel # pytype: disable=import-error - from tpu_inference.layers.common.utils import reorder_concatenated_tensor_for_sharding # pylint: disable=import-outside-toplevel # pytype: disable=import-error + from tpu_inference.layers.common.utils import ( # pylint: disable=import-outside-toplevel # pytype: disable=import-error + reorder_concatenated_tensor_for_sharding, + truncate_sharded_tensor, + ) from tpu_inference.utils import get_mesh_shape_product # pylint: disable=import-outside-toplevel # pytype: disable=import-error from jax.sharding import PartitionSpec as P_spec # pylint: disable=import-outside-toplevel # pytype: disable=import-error except ImportError as e: @@ -702,6 +716,26 @@ def __call__( conv_state_paged, recurrent_state_paged = kv_cache + # Compile against the active request bucket rather than the runner's + # maximum-size metadata buffers. + dp_size = get_mesh_shape_product(self.mesh, attn_data) + padded_num_reqs_per_dp = attention_metadata.padded_num_reqs // dp_size # pyrefly: ignore[missing-attribute] + state_indices = truncate_sharded_tensor( + attention_metadata.mamba_state_indices.astype(jnp.int32), # pyrefly: ignore[missing-attribute] + padded_num_reqs_per_dp, + dp_size, + ) + query_start_loc = truncate_sharded_tensor( + attention_metadata.query_start_loc, # pyrefly: ignore[missing-attribute] + padded_num_reqs_per_dp + 1, + dp_size, + ) + seq_lens = truncate_sharded_tensor( + attention_metadata.seq_lens, # pyrefly: ignore[missing-attribute] + padded_num_reqs_per_dp, + dp_size, + ) + (new_conv_state_paged, new_recurrent_state_paged), gdn_output = run_jax_gdn_attention( mixed_qkv, b_flat, @@ -712,10 +746,10 @@ def __call__( None, # conv_bias: MaxText conv1d uses use_bias=False. jnp.asarray(self.A_log[...], dtype=cfg.dtype), jnp.asarray(self.dt_bias[...], dtype=cfg.dtype), - attention_metadata.mamba_state_indices.astype(jnp.int32), # pyrefly: ignore[missing-attribute] - attention_metadata.query_start_loc, # pyrefly: ignore[missing-attribute] + state_indices, + query_start_loc, attention_metadata.request_distribution, # pyrefly: ignore[missing-attribute] - attention_metadata.seq_lens, # pyrefly: ignore[missing-attribute] + seq_lens, self.num_k_heads, self.num_v_heads, self.head_k_dim, @@ -849,6 +883,25 @@ def extract_state(c_in, v_len): qkv_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) g_beta_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD), mesh=self.mesh, rules=logical_rules) state_pspec = logical_to_mesh_axes((KV_BATCH, KV_HEAD, None, None), mesh=self.mesh, rules=logical_rules) + # Keep every shard_map input/output batch spec consistent when replication is required. + qkv_pspec = remove_incompatible_mesh_axes_from_partition_spec( + qkv_pspec, + query.shape, + self.mesh, + dims=(0,), + ) + g_beta_pspec = remove_incompatible_mesh_axes_from_partition_spec( + g_beta_pspec, + g.shape, + self.mesh, + dims=(0,), + ) + state_pspec = remove_incompatible_mesh_axes_from_partition_spec( + state_pspec, + recurrent_state_arg.shape, + self.mesh, + dims=(0,), + ) @functools.partial( jax.shard_map, diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index b67420313a..26556d3ab7 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/src/maxtext/trainers/post_train/rl/train_rl.py @@ -136,15 +136,41 @@ def _compat_unstack(src_val, tgt_val, key_path, scan_axis=None): os.environ["TOKENIZERS_PARALLELISM"] = "0" +from maxtext.common.common_types import DecoderBlockType from maxtext.configs import pyconfig, types from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR -from maxtext.integration.vllm.maxtext_vllm_rollout import MaxTextVllmRollout +from maxtext.integration.vllm.maxtext_vllm_rollout import MaxTextVllmRollout, requires_maxtext_scanned_weight_unroll from maxtext.trainers.post_train.rl.evaluate_rl import evaluate from maxtext.trainers.post_train.rl import utils_rl from maxtext.input_pipeline.instruction_data_processing import load_data_template_from_file from maxtext.utils import max_logging, max_utils, model_creation_utils +_RECURRENT_ROLLOUT_DECODER_BLOCKS = frozenset( + (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5) +) + + +def rollout_prefix_caching_enabled(trainer_config: Any) -> bool: + """Returns whether vLLM prefix caching is safe for this rollout model. + + The MaxText vLLM adapter stores Qwen3-Next/Qwen3.5 GDN state in a + per-request slot, outside vLLM's block-addressed prefix cache. Reusing only + the attention KV blocks would therefore skip the cached prefix tokens while + starting the recurrent layers from a fresh (potentially stale) state slot. + Keep prefix caching disabled for these hybrid recurrent models until the + adapter implements recurrent-state prefix save/restore. + """ + enabled = bool(getattr(trainer_config, "enable_prefix_caching", False)) + if enabled and getattr(trainer_config, "decoder_block", None) in _RECURRENT_ROLLOUT_DECODER_BLOCKS: + max_logging.log( + "Disabling vLLM prefix caching for the hybrid recurrent rollout model: " + "GDN state is not part of the MaxText adapter's prefix cache." + ) + return False + return enabled + + def get_dataset( tmvp_config: Any, split: str = "train", @@ -463,7 +489,7 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments rl_rollout_engine = ( functools.partial(MaxTextVllmRollout, maxtext_config=trainer_config) - if trainer_config.use_standalone_converter + if trainer_config.use_standalone_converter or requires_maxtext_scanned_weight_unroll(trainer_config) else "vllm" ) @@ -514,7 +540,7 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments rollout_vllm_kwargs={ "hf_overrides": trainer_config.vllm_hf_overrides, "enable_expert_parallel": sampler_config.enable_expert_parallel, - "enable_prefix_caching": True, # Enable prefix caching to speed up generation for long prompts + "enable_prefix_caching": rollout_prefix_caching_enabled(trainer_config), # Ensures vLLM model initializes with correct dtype (not float32 default) "dtype": trainer_config.weight_dtype.value, }, diff --git a/src/maxtext/trainers/post_train/rl/utils_rl.py b/src/maxtext/trainers/post_train/rl/utils_rl.py index 15028a783f..f6b891e3c1 100644 --- a/src/maxtext/trainers/post_train/rl/utils_rl.py +++ b/src/maxtext/trainers/post_train/rl/utils_rl.py @@ -164,10 +164,18 @@ def boxed(x: str) -> str: def get_match_format_regex(tmvp_config: Any) -> re.Pattern[str]: """Returns a compiled regex to extract the answer from a completion.""" + # Some thinking-model chat templates (Qwen3.5) prefill the opening + # reasoning marker in the prompt. The generated completion then begins with + # the reasoning body and contains only the closing marker. + reasoning_start = ( + "" + if getattr(tmvp_config, "reasoning_start_token_in_prompt", False) + else re.escape(tmvp_config.reasoning_start_token) + ) match_format = re.compile( ( - rf"{tmvp_config.reasoning_start_token}.+{tmvp_config.reasoning_end_token}.*?" - rf"{tmvp_config.solution_start_token}(.+?){tmvp_config.solution_end_token}" + rf"{reasoning_start}.+{re.escape(tmvp_config.reasoning_end_token)}.*?" + rf"{re.escape(tmvp_config.solution_start_token)}(.+?){re.escape(tmvp_config.solution_end_token)}" ), flags=re.MULTILINE | re.DOTALL, ) @@ -220,7 +228,10 @@ def match_format_approximately(prompts: list[str], completions: list[str], tmvp_ # If we see 1, then plus some points! score += ( tmvp_config.reward_partial_format_match - if completion.count(tmvp_config.reasoning_start_token) == 1 + if ( + getattr(tmvp_config, "reasoning_start_token_in_prompt", False) + or completion.count(tmvp_config.reasoning_start_token) == 1 + ) else tmvp_config.penalty_incorrect_format ) score += ( @@ -489,13 +500,19 @@ def extract_answer(response: str, tmvp_config: Any) -> str: Strategy (priority order): 1. Narrow the search scope to the LAST `{solution_start_token}...{solution_end_token}` block (default - `...`) if present; otherwise use the full response. + `...`) if present. For a native thinking response, use + only the section after its reasoning close; otherwise use the full + response. 2. Inside the search scope, find the last `\\boxed{N}` via a brace- balanced scan (handles nested braces in LaTeX). Fall back to a permissive `\\boxed{N}` regex if no balanced match is found. 3. If no boxed expression is found, fall back to the same configured solution-tag regex over the full response, for recipes that emit the answer as plain text rather than `\\boxed{N}`. + 4. For a thinking template that prefilled its opening marker, accept an + explicitly stated numeric conclusion in the final-answer section after + the closing reasoning marker. This keeps a valid native Qwen response + scoreable even if it uses prose instead of the requested answer tags. Step 1 + 2 are required for modern reasoning models (Qwen3, DeepSeek-R1, etc.) that emit `...\\boxed{N}` or `\\boxed{N}` @@ -506,11 +523,23 @@ def extract_answer(response: str, tmvp_config: Any) -> str: The solution tags have a single source of truth: both the scoping (step 1) and the plain-text fallback (step 3) reuse `get_answer_fallback_regex`, - built from `tmvp_config.solution_start_token` / `solution_end_token`. + built from `tmvp_config.solution_start_token` / `solution_end_token`. Step 4 + is enabled only for recipes that explicitly mark the reasoning opener as + part of the prompt. """ answer_tag_regex = get_answer_fallback_regex(tmvp_config) answer_matches = answer_tag_regex.findall(response) - content = answer_matches[-1] if answer_matches else response + native_final_section = None + if getattr(tmvp_config, "reasoning_start_token_in_prompt", False): + reasoning_end = tmvp_config.reasoning_end_token + if reasoning_end in response: + native_final_section = response.rsplit(reasoning_end, maxsplit=1)[-1] + if answer_matches: + content = answer_matches[-1] + elif native_final_section is not None: + content = native_final_section + else: + content = response boxed_matches: list[str] = [] stack: list[int] = [] for i, ch in enumerate(content): @@ -530,6 +559,29 @@ def extract_answer(response: str, tmvp_config: Any) -> str: fallback_matches = answer_tag_regex.findall(response) if fallback_matches: return fallback_matches[-1].strip() + + if native_final_section is not None: + # GSM8K gold answers are numeric. Require either an explicit conclusion + # phrase or a numeric-only final section; do not reward an arbitrary + # trailing number in prose. + numeric_pattern = ( + r"[-+]?(?:\d{1,3}(?:,\d{3})+|\d+|\.\d+)(?:\.\d+)?" + r"(?:\s*/\s*[-+]?\d+(?:\.\d+)?)?" + ) + conclusion_matches = re.findall( + rf"(?:\b(?:the\s+)?(?:final\s+)?answer\s*(?:is|=|:)?|\b(?:therefore|thus|hence)\s*[,:]?)" + rf"\s*\$?({numeric_pattern})\$?", + native_final_section, + flags=re.IGNORECASE, + ) + if conclusion_matches: + return conclusion_matches[-1].strip() + numeric_only = re.fullmatch( + rf"\s*\$?({numeric_pattern})\$?\s*[.!]?\s*", + native_final_section, + ) + if numeric_only: + return numeric_only.group(1).strip() return FALLBACK_ANSWER diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index e02e598a12..1de76fdb7e 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -910,6 +910,52 @@ def remove_mesh_axes_from_partition_spec(pspec, axes_to_remove, dims=None): return jax.sharding.PartitionSpec(*new_spec) +def remove_incompatible_mesh_axes_from_partition_spec(pspec, shape, mesh, dims=None): + """Replicate tensor dimensions that cannot be evenly sharded by their mesh axes. + + `shard_map` requires every tensor dimension to be evenly divisible by the + product of the mesh axes assigned to that dimension. When that requirement is + not met, remove the assigned axes from the dimension so JAX replicates it + instead. + + Args: + pspec: Physical PartitionSpec to make compatible with `shape`. + shape: Global tensor shape described by `pspec`. + mesh: Device mesh containing the physical axes referenced by `pspec`. + dims: Dim indices to check; `None` (the default) checks every dim. + + Returns: + A PartitionSpec whose checked dimensions are evenly shardable. + """ + if len(pspec) > len(shape): + raise ValueError(f"PartitionSpec rank {len(pspec)} exceeds tensor rank {len(shape)}") + + dims_to_check = None if dims is None else set(dims) + compatible_pspec = pspec + for dim, (dim_size, partition) in enumerate(zip(shape, pspec)): + if (dims_to_check is not None and dim not in dims_to_check) or partition is None or partition == P.UNCONSTRAINED: + continue + + if isinstance(partition, str): + mesh_axes = (partition,) + elif isinstance(partition, (list, tuple)): + mesh_axes = tuple(partition) + else: + raise ValueError(f"Unsupported axis type: {type(partition)}") + + shard_count = 1 + for mesh_axis in mesh_axes: + shard_count *= mesh.shape[mesh_axis] + if dim_size % shard_count: + compatible_pspec = remove_mesh_axes_from_partition_spec( + compatible_pspec, + mesh_axes, + dims=(dim,), + ) + + return compatible_pspec + + def remove_mesh_axes_from_sharding(sharding_tree, axes_to_remove): """Recursively traverses a sharding tree removing `axes_to_remove` from each spec.""" diff --git a/tests/post_training/unit/extract_answer_test.py b/tests/post_training/unit/extract_answer_test.py index dad5d10dd2..d27891d96c 100644 --- a/tests/post_training/unit/extract_answer_test.py +++ b/tests/post_training/unit/extract_answer_test.py @@ -14,11 +14,12 @@ """Unit tests for utils_rl.extract_answer (CPU-only). -Covers the two-part contract of the boxed-extraction change: +Covers the answer-extraction contract: 1. `\\boxed{N}` is extracted (with/without tags, nested LaTeX, multiple boxed, whitespace, negatives, and answer-tag scoping). 2. Legacy plain-text answers inside the solution tags still work, so existing recipes that do not emit `\\boxed` are unaffected. + 3. Native thinking-model final prose is scoped after ``. """ import unittest @@ -102,6 +103,67 @@ def test_legacy_last_answer_wins(self): got = utils_rl.extract_answer("1 ... 5", self.config) self.assertEqual(got, "5") + @pytest.mark.cpu_only + def test_qwen_native_final_answer_without_custom_tags(self): + """Qwen's native post-thinking prose remains scoreable.""" + config = SimpleNamespace( + reasoning_start_token="", + reasoning_end_token="", + reasoning_start_token_in_prompt=True, + solution_start_token="", + solution_end_token="", + ) + response = "We calculate 6 * 7.\nThe final answer is 42. do do do" + + self.assertEqual(utils_rl.extract_answer(response, config), "42") + + @pytest.mark.cpu_only + def test_qwen_native_final_section_ignores_reasoning_box(self): + """An intermediate boxed value must not override the native final answer.""" + config = SimpleNamespace( + reasoning_start_token="", + reasoning_end_token="", + reasoning_start_token_in_prompt=True, + solution_start_token="", + solution_end_token="", + ) + response = "Maybe \\boxed{41}, but recalculate.\nThe answer is 42." + + self.assertEqual(utils_rl.extract_answer(response, config), "42") + + @pytest.mark.cpu_only + def test_native_final_fallback_does_not_read_reasoning_numbers(self): + """An unfinished thinking trace must not leak an intermediate value.""" + config = SimpleNamespace( + reasoning_start_token="", + reasoning_end_token="", + reasoning_start_token_in_prompt=True, + solution_start_token="", + solution_end_token="", + ) + + self.assertEqual( + utils_rl.extract_answer("Try 41, then 42", config), + utils_rl.FALLBACK_ANSWER, + ) + + @pytest.mark.cpu_only + def test_native_final_fallback_does_not_take_unrelated_trailing_number(self): + """Only an explicit conclusion is accepted from native final prose.""" + config = SimpleNamespace( + reasoning_start_token="", + reasoning_end_token="", + reasoning_start_token_in_prompt=True, + solution_start_token="", + solution_end_token="", + ) + response = "Work\nI considered 41 and 42, then wrote this in 2026." + + self.assertEqual( + utils_rl.extract_answer(response, config), + utils_rl.FALLBACK_ANSWER, + ) + # ---- no answer ---- def test_no_answer_returns_fallback_constant(self): diff --git a/tests/post_training/unit/qwen35_recipe_test.py b/tests/post_training/unit/qwen35_recipe_test.py new file mode 100644 index 0000000000..5d33eda126 --- /dev/null +++ b/tests/post_training/unit/qwen35_recipe_test.py @@ -0,0 +1,67 @@ +# 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. + +"""Static checks for the Qwen3.5 35B GSM8K post-training recipe.""" + +import json +from pathlib import Path + +import pytest +import yaml + + +pytestmark = [pytest.mark.post_training, pytest.mark.cpu_only] +_REPO_ROOT = Path(__file__).resolve().parents[3] +_RECIPE_PATH = _REPO_ROOT / "src/maxtext/configs/post_train/rl_gsm8k_qwen35_35b_v5p64.yml" +_TEMPLATE_PATH = _REPO_ROOT / "src/maxtext/examples/chat_templates/qwen35_math_rl.json" + + +def test_qwen35_recipe_keeps_only_reproducible_runtime_configuration(): + recipe = yaml.safe_load(_RECIPE_PATH.read_text(encoding="utf-8")) + + assert recipe["model_name"] == "qwen3.5-35b-a3b" + assert recipe["scan_layers"] + assert recipe["vllm_hf_overrides"] == {"architectures": ["MaxTextForCausalLM"]} + assert recipe["vllm_additional_config"]["maxtext_config"] == { + "model_name": "qwen3.5-35b-a3b", + "model_call_mode": "inference", + "attention": "vllm_rpa", + "allow_split_physical_axes": True, + "log_config": False, + "weight_dtype": "bfloat16", + "prefuse_moe_weights": True, + } + assert not recipe["enable_prefix_caching"] + assert recipe["reasoning_start_token_in_prompt"] + assert recipe["stop_strings"] == [""] + assert not recipe["debug"] + + experiment_only_keys = { + "base_output_directory", + "enable_tunix_perf_metrics", + "load_parameters_path", + "run_name", + } + assert experiment_only_keys.isdisjoint(recipe) + + +def test_qwen35_template_does_not_nest_model_specific_chat_or_reasoning_tags(): + template = json.loads(_TEMPLATE_PATH.read_text(encoding="utf-8")) + + assert "{solution_start_token}" in template["SYSTEM_PROMPT"] + assert "{solution_end_token}" in template["SYSTEM_PROMPT"] + assert "{reasoning_start_token}" not in template["SYSTEM_PROMPT"] + assert "{reasoning_end_token}" not in template["SYSTEM_PROMPT"] + assert "<|im_start|>" not in template["TEMPLATE"] + assert "" not in template["TEMPLATE"] diff --git a/tests/post_training/unit/rl_utils_test.py b/tests/post_training/unit/rl_utils_test.py index 93c3e09b8c..e84a189ec3 100644 --- a/tests/post_training/unit/rl_utils_test.py +++ b/tests/post_training/unit/rl_utils_test.py @@ -128,6 +128,19 @@ def test_score_multiple_completions(self): self.assertEqual(scores[0], 2.0) self.assertEqual(scores[1], -2.0) + @pytest.mark.cpu_only + def test_prefilled_reasoning_start_scores_native_completion(self): + self.config.reasoning_start_token = "" + self.config.reasoning_end_token = "" + self.config.reasoning_start_token_in_prompt = True + completion = "40 + 2 = 42\n\n42" + + self.assertEqual(self._score([completion])[0], 2.0) + self.assertEqual( + utils_rl.match_format_exactly(None, [completion], self.config)[0], + self.config.reward_exact_format_match, + ) + class TestCheckNumbers(unittest.TestCase): """Tests for utils_rl.check_numbers. diff --git a/tests/post_training/unit/train_rl_test.py b/tests/post_training/unit/train_rl_test.py index 48c2d6b3c0..41daec72e9 100644 --- a/tests/post_training/unit/train_rl_test.py +++ b/tests/post_training/unit/train_rl_test.py @@ -42,12 +42,45 @@ def _get_mock_devices(devices_per_slice, num_slices=1): class TrainRLTest(unittest.TestCase): """Tests for train_rl.py.""" - def test_rl_config_includes_decoder_engram_defaults(self): - """RL models must expose Engram fields consumed by the shared decoder.""" + def test_rl_config_includes_shared_decoder_defaults(self): + """RL models must expose fields consumed by the shared decoder.""" config = types.RLConfig(model_name="gemma4-26b") self.assertEqual(config.engram_layers, []) self.assertEqual(config.engram_max_ngram_size, 3) + self.assertEqual(config.mhc_expansion_rate, 1) + self.assertEqual(config.sinkhorn_iterations, 20) + self.assertFalse(config.enable_mhc_lite) + self.assertFalse(config.enable_prefix_caching) + + @pytest.mark.cpu_only + def test_rollout_prefix_caching_respects_config_for_attention_model(self): + config = SimpleNamespace( + enable_prefix_caching=True, + decoder_block=train_rl.DecoderBlockType.QWEN3, + ) + + self.assertTrue(train_rl.rollout_prefix_caching_enabled(config)) + + @pytest.mark.cpu_only + def test_rollout_prefix_caching_disabled_for_recurrent_model(self): + for decoder_block in (train_rl.DecoderBlockType.QWEN3_NEXT, train_rl.DecoderBlockType.QWEN3_5): + with self.subTest(decoder_block=decoder_block): + config = SimpleNamespace( + enable_prefix_caching=True, + decoder_block=decoder_block, + ) + + with mock.patch.object(train_rl.max_logging, "log") as mock_log: + self.assertFalse(train_rl.rollout_prefix_caching_enabled(config)) + + mock_log.assert_called_once() + + @pytest.mark.cpu_only + def test_rollout_prefix_caching_defaults_to_disabled(self): + config = SimpleNamespace(decoder_block=train_rl.DecoderBlockType.QWEN3) + + self.assertFalse(train_rl.rollout_prefix_caching_enabled(config)) def test_setup_configs_and_devices_pathways_split(self): """Test setup_configs_and_devices with multiple VMs and Pathways.""" diff --git a/tests/post_training/unit/vllm_hybrid_cache_test.py b/tests/post_training/unit/vllm_hybrid_cache_test.py new file mode 100644 index 0000000000..f03ef1dcf0 --- /dev/null +++ b/tests/post_training/unit/vllm_hybrid_cache_test.py @@ -0,0 +1,50 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for MaxText vLLM hybrid-cache layout helpers.""" + +from types import SimpleNamespace +import unittest + +import pytest +import torch + +from maxtext.integration.vllm._hybrid_cache import build_qwen_gdn_cache_layout + + +pytestmark = [pytest.mark.post_training] + + +class QwenGdnCacheLayoutTest(unittest.TestCase): + """Verify the mixed-precision recurrent-cache contract.""" + + @pytest.mark.cpu_only + def test_recurrent_state_is_float32_and_page_size_uses_each_dtype(self): + cfg = SimpleNamespace( + gdn_num_value_heads=32, + gdn_num_key_heads=16, + gdn_key_head_dim=128, + gdn_value_head_dim=128, + gdn_conv_kernel_dim=4, + ) + + shapes, dtypes, page_size_bytes = build_qwen_gdn_cache_layout(cfg, torch) + + self.assertEqual(shapes, ((3, 8192), (32, 128, 128))) + self.assertEqual(dtypes, (torch.bfloat16, torch.float32)) + self.assertEqual(page_size_bytes, 2_146_304) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/post_training/unit/vllm_rollout_unroll_test.py b/tests/post_training/unit/vllm_rollout_unroll_test.py index 7028873df4..99e657cb95 100644 --- a/tests/post_training/unit/vllm_rollout_unroll_test.py +++ b/tests/post_training/unit/vllm_rollout_unroll_test.py @@ -12,13 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the Gemma scanned weights unrolling workaround.""" +"""Unit tests for MaxText scanned-weight unrolling workarounds.""" +from types import SimpleNamespace import unittest +from unittest import mock import numpy as np import pytest -from maxtext.integration.vllm.maxtext_vllm_rollout import unroll_gemma_scanned_weights +from maxtext.integration.vllm.maxtext_vllm_rollout import ( + MaxTextVllmRollout, + prepare_direct_sync_additional_config, + requires_maxtext_scanned_weight_unroll, + unroll_gemma_scanned_weights, + unroll_qwen_scanned_weights, + uses_maxtext_vllm_adapter, + validate_direct_sync_layer_coverage, +) pytestmark = pytest.mark.post_training @@ -168,3 +178,280 @@ def test_correctly_unrolls_gemma3_gemma4_scanned_blocks(self): self.assertIsInstance(list(decoder_dict.keys())[0], str) np.testing.assert_array_equal(decoder_dict["layers_0"]["mlp"]["wi_0"], np.array([[0], [0]])) np.testing.assert_array_equal(decoder_dict["layers_6"]["mlp"]["wi_0"], np.array([[6], [6]])) + + +class QwenScannedWeightsUnrollTest(unittest.TestCase): + """Verify heterogeneous Qwen blocks map to unscanned decoder attributes.""" + + @pytest.mark.cpu_only + def test_interleaves_slots_and_repetitions(self): + slot_0 = np.zeros((2, 2, 1), dtype=np.float32) + slot_0[:, 0, :] = 0 + slot_0[:, 1, :] = 2 + slot_1 = np.zeros((2, 2, 1), dtype=np.float32) + slot_1[:, 0, :] = 1 + slot_1[:, 1, :] = 3 + weights = MockWeights( + { + "base": { + "decoder": { + "layers": { + "layer_0": {"probe": slot_0}, + "layer_1": {"probe": slot_1, "rngs": {"key": np.ones(2, dtype=np.uint32)}}, + }, + "decoder_norm": {"scale": np.ones(2)}, + } + } + } + ) + + unrolled = unroll_qwen_scanned_weights(weights) + + decoder = unrolled["base"]["decoder"] + for layer_idx in range(4): + self.assertIn(f"layers_{layer_idx}", decoder) + np.testing.assert_array_equal( + decoder[f"layers_{layer_idx}"]["probe"], + np.full((2, 1), layer_idx, dtype=np.float32), + ) + np.testing.assert_array_equal(decoder["decoder_norm"]["scale"], np.ones(2)) + self.assertNotIn("probe", decoder["layers"]["layer_1"]) + np.testing.assert_array_equal(decoder["layers"]["layer_1"]["rngs"]["key"], np.ones(2, dtype=np.uint32)) + target = { + "model": { + "decoder": { + f"layers_{layer_idx}": {"probe": np.zeros((2, 1), dtype=np.float32)} for layer_idx in range(4) + } + } + } + self.assertEqual(validate_direct_sync_layer_coverage(unrolled, target), 4) + + @pytest.mark.cpu_only + def test_rejects_inconsistent_scan_lengths(self): + weights = MockWeights( + { + "decoder": { + "layers": { + "layer_0": {"probe": np.ones((2, 2, 1))}, + "layer_1": {"probe": np.ones((2, 3, 1))}, + } + } + } + ) + + with self.assertRaisesRegex(ValueError, "disagree on scan length"): + unroll_qwen_scanned_weights(weights) + + @pytest.mark.cpu_only + def test_supports_nondefault_axis_and_sparse_slots(self): + slot_1 = np.stack( + [np.full((2, 1), 1, dtype=np.float32), np.full((2, 1), 3, dtype=np.float32)], + axis=0, + ) + weights = MockWeights({"decoder": {"layers": {"layer_1": {"probe": slot_1}}}}) + + unrolled = unroll_qwen_scanned_weights(weights, scan_axis=0, pattern_length=2) + + self.assertNotIn("layers_0", unrolled["decoder"]) + np.testing.assert_array_equal(unrolled["decoder"]["layers_1"]["probe"], np.full((2, 1), 1)) + np.testing.assert_array_equal(unrolled["decoder"]["layers_3"]["probe"], np.full((2, 1), 3)) + + @pytest.mark.cpu_only + def test_rejects_missing_target_layer_parameters(self): + source = {"base": {"decoder": {"layers_0": {"probe": np.ones((2, 1))}}}} + target = { + "model": { + "decoder": { + "layers_0": {"probe": np.zeros((2, 1))}, + "layers_1": {"probe": np.zeros((2, 1))}, + } + } + } + + with self.assertRaisesRegex(ValueError, "leave rollout transformer parameters at random initialization"): + validate_direct_sync_layer_coverage(source, target) + + @pytest.mark.cpu_only + def test_rejects_source_without_unrolled_layers(self): + source = {"base": {"decoder": {"layers": {"layer_0": {"probe": np.ones((2, 1))}}}}} + target = {"model": {"decoder": {"layers_0": {"probe": np.zeros((2, 1))}}}} + + with self.assertRaisesRegex(ValueError, "matched 0/1 target layer parameters"): + validate_direct_sync_layer_coverage(source, target) + + @pytest.mark.cpu_only + def test_accepts_split_moe_weights_for_prefused_target(self): + source = { + "base": { + "decoder": { + "layers_0": { + "mlp": { + "routed_experts": { + "wi_0": np.ones((2, 3, 4)), + "wi_1": np.ones((2, 3, 4)), + } + } + } + } + } + } + target = { + "model": { + "decoder": { + "layers_0": { + "mlp": {"routed_experts": {"wi": np.zeros((2, 3, 8))}} + } + } + } + } + + self.assertEqual(validate_direct_sync_layer_coverage(source, target), 1) + + +class DirectSyncRolloutConfigTest(unittest.TestCase): + """Verify TP-sharded MoE rollout targets request the safe fused layout.""" + + @pytest.mark.cpu_only + def test_enables_prefusion_for_direct_moe_tp(self): + original = {"maxtext_config": {"model_name": "qwen3.5-35b-a3b"}} + + prepared = prepare_direct_sync_additional_config( + original, + direct_maxtext_sync=True, + num_experts=256, + tensor_parallel_size=4, + ) + + self.assertTrue(prepared["maxtext_config"]["prefuse_moe_weights"]) + self.assertNotIn("prefuse_moe_weights", original["maxtext_config"]) + + @pytest.mark.cpu_only + def test_leaves_dense_or_single_tp_config_unchanged(self): + original = {"maxtext_config": {"model_name": "qwen3-0.6b"}} + + self.assertIs( + prepare_direct_sync_additional_config( + original, + direct_maxtext_sync=True, + num_experts=1, + tensor_parallel_size=4, + ), + original, + ) + self.assertIs( + prepare_direct_sync_additional_config( + original, + direct_maxtext_sync=True, + num_experts=256, + tensor_parallel_size=1, + ), + original, + ) + + +class MaxTextVllmRolloutConfigForwardingTest(unittest.TestCase): + """Verify the custom rollout preserves Tunix rollout options.""" + + @pytest.mark.cpu_only + def test_forwards_sampling_parallelism_and_capacity_options(self): + sampling_kwargs = { + "stop": [""], + "detokenize": True, + "include_stop_str_in_output": True, + } + rollout_config = SimpleNamespace( + kv_cache_size=1280, + rollout_mapping_config=None, + rollout_vllm_model_version="Qwen/Qwen3.5-35B-A3B", + rollout_vllm_swap_space_size_gb=2, + rollout_vllm_async_scheduling=False, + rollout_vllm_max_num_batched_tokens=16384, + rollout_vllm_max_num_seqs=32, + rollout_vllm_hf_config_path=None, + rollout_vllm_logprobs_mode="raw_logprobs", + rollout_vllm_kwargs={"dtype": "bfloat16"}, + rollout_vllm_additional_config={"maxtext_config": {"model_name": "qwen3.5-35b-a3b"}}, + rollout_vllm_hbm_utilization=0.6, + rollout_vllm_init_with_random_weights=True, + rollout_vllm_tpu_backend_type="jax", + rollout_vllm_lora_config=None, + rollout_vllm_server_mode=False, + rollout_vllm_server_mode_submission_threshold=7, + rollout_vllm_server_mode_submission_timeout_s=3.0, + return_logprobs=True, + tensor_parallel_size=4, + data_parallel_size=2, + expert_parallel_size=1, + rollout_vllm_enable_dp_attention=False, + rollout_vllm_delete_dst_buffers=True, + rollout_vllm_reshard_chunk_size=8, + rollout_vllm_sampling_kwargs=sampling_kwargs, + ) + maxtext_config = SimpleNamespace( + model_name="qwen3.5-35b-a3b", + num_experts=256, + param_scan_axis=1, + inhomogeneous_layer_cycle_interval=4, + swap_space_vllm_gb=2, + vllm_hf_overrides={"architectures": ["MaxTextForCausalLM"]}, + ) + fake_sampler = mock.MagicMock() + + with ( + mock.patch( + "maxtext.integration.vllm.maxtext_vllm_rollout.mappings.MappingConfig.build", + return_value=object(), + ), + mock.patch( + "maxtext.integration.vllm.maxtext_vllm_rollout.VllmConfig", + side_effect=lambda **kwargs: SimpleNamespace(**kwargs), + ), + mock.patch( + "maxtext.integration.vllm.maxtext_vllm_rollout.MaxTextVllmSampler", + return_value=fake_sampler, + ) as sampler_cls, + mock.patch("maxtext.integration.vllm.maxtext_vllm_rollout.nnx.state", return_value={"base": {}}), + ): + MaxTextVllmRollout( + rollout_actor=object(), + tokenizer=object(), + mesh=object(), + rollout_config=rollout_config, + maxtext_config=maxtext_config, + ) + + config = sampler_cls.call_args.kwargs["config"] + self.assertEqual(config.sampling_kwargs, sampling_kwargs) + self.assertEqual(config.expert_parallel_size, 1) + self.assertEqual(config.return_logprobs, True) + self.assertEqual(config.reshard_chunk_size, 8) + self.assertEqual(config.server_mode_submission_threshold, 7) + self.assertEqual(config.server_mode_submission_timeout_s, 3.0) + self.assertEqual(config.engine_kwargs["max_num_batched_tokens"], 16384) + self.assertEqual(config.engine_kwargs["max_num_seqs"], 32) + self.assertEqual(config.engine_kwargs["logprobs_mode"], "raw_logprobs") + self.assertNotIn("swap_space", config.engine_kwargs) + fake_sampler.load_checkpoint.assert_called_once_with({"base": {}}) + + +class MaxTextAdapterSelectionTest(unittest.TestCase): + """Verify only scanned MaxText-adapter rollouts take the custom sync path.""" + + @pytest.mark.cpu_only + def test_detects_dict_and_string_overrides(self): + dict_config = SimpleNamespace(vllm_hf_overrides={"architectures": ["MaxTextForCausalLM"]}, scan_layers=True) + string_config = SimpleNamespace(vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}', scan_layers=True) + + self.assertTrue(uses_maxtext_vllm_adapter(dict_config)) + self.assertTrue(uses_maxtext_vllm_adapter(string_config)) + self.assertTrue(requires_maxtext_scanned_weight_unroll(dict_config)) + + @pytest.mark.cpu_only + def test_bypasses_native_or_unscanned_rollouts(self): + native_config = SimpleNamespace( + vllm_hf_overrides={"architectures": ["Qwen3_5MoeForConditionalGeneration"]}, scan_layers=True + ) + unscanned_config = SimpleNamespace(vllm_hf_overrides={"architectures": ["MaxTextForCausalLM"]}, scan_layers=False) + + self.assertFalse(requires_maxtext_scanned_weight_unroll(native_config)) + self.assertFalse(requires_maxtext_scanned_weight_unroll(unscanned_config)) diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index de91cb647c..3bef7a4b30 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -31,6 +31,7 @@ from jax.sharding import AxisType, Mesh from maxtext.utils import max_utils from maxtext.utils import maxtext_utils +from maxtext.utils import sharding from maxtext.common.gcloud_stub import is_decoupled from maxtext.common.common_types import ( @@ -3029,6 +3030,106 @@ def get_structured_data(self, dtype): ) return lnx + @pytest.mark.cpu_only + def test_train_path_checks_all_batch_sharding_specs(self): + """The non-paged GDN path makes every batch-sharded spec shape-compatible.""" + lnx = self.get_structured_data(self.cfg.dtype) + gdn = Qwen3NextGatedDeltaNet( + config=self.cfg, + inputs_shape=lnx.shape, + mesh=self.mesh, + dtype=self.cfg.dtype, + model_mode=MODEL_MODE_TRAIN, + rngs=self.nnx_rng, + ) + + with mock.patch( + "maxtext.models.qwen3.remove_incompatible_mesh_axes_from_partition_spec", + wraps=sharding.remove_incompatible_mesh_axes_from_partition_spec, + ) as make_compatible: + output, _ = gdn(lnx, model_mode=MODEL_MODE_TRAIN) + + self.assertEqual(output.shape, lnx.shape) + self.assertEqual(make_compatible.call_count, 4) + self.assertEqual([len(call.args[1]) for call in make_compatible.call_args_list], [4, 4, 3, 4]) + self.assertTrue(all(call.kwargs["dims"] == (0,) for call in make_compatible.call_args_list)) + + @pytest.mark.cpu_only + @pytest.mark.post_training + def test_paged_state_truncates_metadata_to_active_requests(self): + """The paged-state bridge trims maximum-size metadata buffers.""" + gdn_attention = pytest.importorskip("tpu_inference.layers.common.gdn_attention") + + cfg = pyconfig.initialize( + [sys.argv[0], get_test_config_path("inference/vllm.yml")], + run_name="paged_gdn_metadata_test", + enable_checkpointing=False, + log_config=False, + base_emb_dim=16, + gdn_num_value_heads=2, + gdn_num_key_heads=2, + gdn_key_head_dim=4, + gdn_value_head_dim=4, + gdn_conv_kernel_dim=4, + gdn_chunk_size=4, + dtype="float32", + weight_dtype="float32", + max_prefill_predict_length=2, + max_target_length=4, + per_device_batch_size=1.0, + ) + devices_array = maxtext_utils.create_device_mesh(cfg) + mesh = Mesh(devices_array, cfg.mesh_axes) + hidden_states = jnp.ones((1, 1, cfg.emb_dim), dtype=cfg.dtype) + gdn = Qwen3NextGatedDeltaNet( + config=cfg, + inputs_shape=hidden_states.shape, + mesh=mesh, + dtype=cfg.dtype, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + rngs=nnx.Rngs(params=0, dropout=1), + ) + + num_blocks = 2 + key_dim = cfg.gdn_num_key_heads * cfg.gdn_key_head_dim + value_dim = cfg.gdn_num_value_heads * cfg.gdn_value_head_dim + conv_dim = 2 * key_dim + value_dim + conv_state = jnp.zeros((num_blocks, cfg.gdn_conv_kernel_dim - 1, conv_dim), dtype=cfg.dtype) + recurrent_state = jnp.zeros( + (num_blocks, cfg.gdn_num_value_heads, cfg.gdn_key_head_dim, cfg.gdn_value_head_dim), + dtype=cfg.dtype, + ) + attention_metadata = types.SimpleNamespace( + padded_num_reqs=1, + mamba_state_indices=jnp.array([1, 101, 102], dtype=jnp.int32), + query_start_loc=jnp.array([0, 1, 101, 201], dtype=jnp.int32), + request_distribution=jnp.array([0, 0, 1], dtype=jnp.int32), + seq_lens=jnp.array([1, 101, 102], dtype=jnp.int32), + ) + + with mock.patch.object(gdn_attention, "run_jax_gdn_attention", autospec=True) as mock_run_gdn: + mock_run_gdn.return_value = ( + (conv_state, recurrent_state), + jnp.zeros((hidden_states.shape[1], value_dim), dtype=cfg.dtype), + ) + output, new_cache = gdn( + hidden_states, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + kv_cache=(conv_state, recurrent_state), + attention_metadata=attention_metadata, + ) + + mock_run_gdn.assert_called_once() + self.assertEqual(len(mock_run_gdn.call_args.args), 18) + self.assertEqual(set(mock_run_gdn.call_args.kwargs), {"mesh"}) + self.assertIs(mock_run_gdn.call_args.kwargs["mesh"], mesh) + np.testing.assert_array_equal(mock_run_gdn.call_args.args[9], jnp.array([1], dtype=jnp.int32)) + np.testing.assert_array_equal(mock_run_gdn.call_args.args[10], jnp.array([0, 1], dtype=jnp.int32)) + np.testing.assert_array_equal(mock_run_gdn.call_args.args[12], jnp.array([1], dtype=jnp.int32)) + self.assertEqual(output.shape, hidden_states.shape) + self.assertEqual(new_cache[0].shape, conv_state.shape) + self.assertEqual(new_cache[1].shape, recurrent_state.shape) + @pytest.mark.tpu_only def test_autoregression(self): cfg = self.cfg diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 7781f4b797..02202f8684 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -14,6 +14,8 @@ """Mixture of Experts (MoE) tests.""" import unittest +from types import SimpleNamespace +from unittest import mock from absl.testing import parameterized import pytest @@ -24,7 +26,7 @@ import jax.numpy as jnp import numpy as np import qwix -from jax.sharding import Mesh +from jax.sharding import Mesh, PartitionSpec as P from maxtext.configs import pyconfig from maxtext.common.common_types import Config, DType from maxtext.layers import linears @@ -437,6 +439,68 @@ def get_moe_loop( return module +@pytest.mark.parametrize( + ("expert_parallelism", "batch_partition"), + ( + (1, None), + (2, ("fsdp", "expert")), + ), +) +def test_sparse_matmul_repairs_batch_specs_only_without_expert_parallelism(expert_parallelism, batch_partition): + """Sparse MoE only replicates batches when expert routing remains local.""" + fake_moe = SimpleNamespace( + config=SimpleNamespace( + shard_exp_on_fsdp=False, + use_2d_fsdp_sharding=False, + model_name="qwen3.5-35b-a3b", + check_vma=False, + moe_fsdp_use_two_stage_all_gather=False, + ), + mesh=SimpleNamespace(shape={"fsdp": 32, "expert": expert_parallelism}), + rngs=object(), + get_expert_parallelism_size=lambda: expert_parallelism, + ) + original_batch_partition = "fsdp" if expert_parallelism == 1 else ("fsdp", "expert") + fake_moe._logical_to_mesh_axes = lambda logical_axes: P( + *(original_batch_partition if axis == "activation_batch" else None for axis in logical_axes) + ) + fake_moe._maybe_shard_with_pspec = lambda value, _pspec: value + + inputs = SimpleNamespace(shape=(4, 1024, 2048)) + gate_logits = SimpleNamespace(shape=(4, 1024, 256)) + w0 = SimpleNamespace(shape=(256, 2048, 512)) + w1 = SimpleNamespace(shape=(256, 2048, 512)) + wo = SimpleNamespace(shape=(256, 512, 2048)) + captured = {} + + def fake_shard_map(function, *, mesh, in_specs, out_specs, check_vma): + del function, mesh, check_vma + captured["in_specs"] = in_specs + captured["out_specs"] = out_specs + return lambda x, *_args: (x, None, None) + + with mock.patch.object(jax, "shard_map", side_effect=fake_shard_map): + output, _, _ = moe.RoutedMoE.sparse_matmul( + fake_moe, + inputs, + gate_logits, + None, + w0, + w1, + wo, + None, + None, + None, + ) + + assert output is inputs + assert captured["in_specs"][0] == P(batch_partition, None, None) + assert captured["in_specs"][1] == P(batch_partition, None, None) + assert captured["in_specs"][2] is None + assert captured["in_specs"][9] is None + assert captured["out_specs"][0] == P(batch_partition, None, None) + + class RoutedMoeTest(parameterized.TestCase): """Routed Mixture of Experts test.""" diff --git a/tests/unit/qwen35_partial_mrope_test.py b/tests/unit/qwen35_partial_mrope_test.py new file mode 100644 index 0000000000..3f1cba3adb --- /dev/null +++ b/tests/unit/qwen35_partial_mrope_test.py @@ -0,0 +1,138 @@ +# 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 Qwen3.5 partial multi-dimensional rotary embeddings.""" + +import unittest +from types import SimpleNamespace + +from flax import nnx +import jax.numpy as jnp +import numpy as np + +from maxtext.layers.attentions import Attention +from maxtext.layers.embeddings import Qwen3OmniMoeThinkerTextRotaryEmbedding + + +_HEAD_DIM = 256 +_PARTIAL_ROTARY_FACTOR = 0.25 +_ROTARY_DIM = 64 +_MROPE_SECTION = (11, 11, 10) +_ROPE_THETA = 10_000_000 + + +def _reference_partial_mrope(inputs: np.ndarray, positions: np.ndarray) -> np.ndarray: + """Independent NumPy implementation of the Qwen3.5 partial-MRoPE rule.""" + if positions.ndim == 2: + positions = np.broadcast_to(positions[np.newaxis, ...], (3,) + positions.shape) + + inv_freq = 1.0 / ( + _ROPE_THETA ** (np.arange(0, _ROTARY_DIM, 2, dtype=np.float32) / _ROTARY_DIM) + ) + freqs = positions[..., np.newaxis].astype(np.float32) * inv_freq + interleaved = np.array(freqs[0], copy=True) + for dim, offset in enumerate((1, 2), start=1): + idx = slice(offset, _MROPE_SECTION[dim] * 3, 3) + interleaved[..., idx] = freqs[dim, ..., idx] + + angles = np.concatenate([interleaved, interleaved], axis=-1) + cos = np.cos(angles)[:, :, np.newaxis, :] + sin = np.sin(angles)[:, :, np.newaxis, :] + + rotary, passthrough = np.split(inputs, [_ROTARY_DIM], axis=-1) + first_half, second_half = np.split(rotary, 2, axis=-1) + rotate_half = np.concatenate([-second_half, first_half], axis=-1) + rotated = rotary * cos + rotate_half * sin + return np.concatenate([rotated, passthrough], axis=-1) + + +class Qwen35PartialMropeTest(unittest.TestCase): + + def test_matches_reference_and_preserves_unrotated_suffix(self): + inputs = np.linspace(-1.0, 1.0, num=2 * 4 * 3 * _HEAD_DIM, dtype=np.float32).reshape( + 2, 4, 3, _HEAD_DIM + ) + text_positions = np.broadcast_to(np.arange(4, dtype=np.int32), (2, 4)) + multimodal_positions = np.stack( + [text_positions, text_positions + 3, text_positions + 7], axis=0 + ) + embedding = Qwen3OmniMoeThinkerTextRotaryEmbedding( + min_timescale=1, + max_timescale=_ROPE_THETA, + embedding_dims=_HEAD_DIM, + partial_rotary_factor=_PARTIAL_ROTARY_FACTOR, + cast_as_fprop_dtype=False, + fprop_dtype=jnp.float32, + mrope_section=_MROPE_SECTION, + rngs=nnx.Rngs(0), + ) + + self.assertEqual(embedding.rotary_dim, _ROTARY_DIM) + self.assertEqual(embedding.timescale.shape, (_ROTARY_DIM // 2,)) + for positions in (text_positions, multimodal_positions): + with self.subTest(position_rank=positions.ndim): + actual = np.asarray(embedding(jnp.asarray(inputs), jnp.asarray(positions))) + expected = _reference_partial_mrope(inputs, positions) + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + np.testing.assert_array_equal(actual[..., _ROTARY_DIM:], inputs[..., _ROTARY_DIM:]) + + def test_factor_one_preserves_full_width_mrope_behavior(self): + inputs = np.linspace(-1.0, 1.0, num=2 * 4 * 3 * _ROTARY_DIM, dtype=np.float32).reshape( + 2, 4, 3, _ROTARY_DIM + ) + positions = np.broadcast_to(np.arange(4, dtype=np.int32), (2, 4)) + embedding = Qwen3OmniMoeThinkerTextRotaryEmbedding( + min_timescale=1, + max_timescale=_ROPE_THETA, + embedding_dims=_ROTARY_DIM, + partial_rotary_factor=1.0, + cast_as_fprop_dtype=False, + fprop_dtype=jnp.float32, + mrope_section=_MROPE_SECTION, + rngs=nnx.Rngs(0), + ) + + actual = np.asarray(embedding(jnp.asarray(inputs), jnp.asarray(positions))) + expected = _reference_partial_mrope(inputs, positions) + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + + def test_attention_wires_qwen35_partial_factor_into_mrope(self): + attention = SimpleNamespace( + config=SimpleNamespace( + attention_type="global", + rope_use_scale=False, + rope_min_timescale=1, + partial_rotary_factor=_PARTIAL_ROTARY_FACTOR, + ), + qk_rope_head_dim=0, + head_dim=_HEAD_DIM, + rope_type="default", + is_vision=False, + use_mrope=True, + rope_max_timescale=_ROPE_THETA, + dtype=jnp.float32, + mrope_section=_MROPE_SECTION, + partial_rotary_factor=None, + rngs=nnx.Rngs(0), + ) + + embedding = Attention.init_rotary_embedding(attention) + + self.assertIsInstance(embedding, Qwen3OmniMoeThinkerTextRotaryEmbedding) + self.assertEqual(embedding.head_dim, _HEAD_DIM) + self.assertEqual(embedding.rotary_dim, _ROTARY_DIM) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sharding_test.py b/tests/unit/sharding_test.py index 390c865b73..b8c3baf640 100644 --- a/tests/unit/sharding_test.py +++ b/tests/unit/sharding_test.py @@ -16,6 +16,7 @@ import datetime +from types import SimpleNamespace import numpy as np import jax @@ -25,6 +26,7 @@ from jax.sharding import Mesh from jax.experimental import mesh_utils from jax.lax import with_sharding_constraint +from maxtext.utils import sharding # Global model and data constants PER_DEVICE_BATCH_SIZE = 131072 @@ -39,6 +41,50 @@ multiply_layers_and_grad = None +@pytest.mark.parametrize( + ("batch_size", "expected"), + ( + (4, PartitionSpec(None, None, "tensor", None)), + (32, PartitionSpec("fsdp", None, "tensor", None)), + ), +) +def test_remove_incompatible_batch_mesh_axes(batch_size, expected): + """Only a non-divisible batch axis is replicated; valid head sharding remains.""" + mesh = SimpleNamespace(shape={"fsdp": 32, "tensor": 2}) + pspec = PartitionSpec("fsdp", None, "tensor", None) + + actual = sharding.remove_incompatible_mesh_axes_from_partition_spec( + pspec, + (batch_size, 1024, 32, 128), + mesh, + dims=(0,), + ) + + assert actual == expected + + +@pytest.mark.parametrize( + ("batch_size", "expected"), + ( + (4, PartitionSpec(None, None, "tensor")), + (8, PartitionSpec(("data", "fsdp"), None, "tensor")), + ), +) +def test_remove_incompatible_composite_batch_mesh_axes(batch_size, expected): + """Compatibility uses the product of all mesh axes assigned to a dimension.""" + mesh = SimpleNamespace(shape={"data": 2, "fsdp": 4, "tensor": 2}) + pspec = PartitionSpec(("data", "fsdp"), None, "tensor") + + actual = sharding.remove_incompatible_mesh_axes_from_partition_spec( + pspec, + (batch_size, 16, 8), + mesh, + dims=(0,), + ) + + assert actual == expected + + def simple_timeit(f, tries=5, verbose=True): """Simple utility to time a function for multiple runs""" outcomes = []