diff --git a/src/maxtext/common/gcloud_stub.py b/src/maxtext/common/gcloud_stub.py index c87ba4123c..044094206a 100644 --- a/src/maxtext/common/gcloud_stub.py +++ b/src/maxtext/common/gcloud_stub.py @@ -330,6 +330,7 @@ def _import(): _goodput_stubs, label="ml_goodput_measurement", stub_if_decoupled=False, + stub_on_error_when_not_decoupled=True, ) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 353656c665..7f88a79991 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2704,6 +2704,10 @@ class VLLM(BaseModel): "the legacy transfer_state_directly / transfer_state_with_mappings paths." ), ) + rollout_backend: Literal["maxtext", "vllm_torchax"] = Field( + "maxtext", + description="Rollout backend for trainer-side weight converter ('maxtext' or 'vllm_torchax').", + ) weight_sync_debug: bool = Field( False, description=( diff --git a/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py b/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py index b0d0deed62..d4f9446475 100644 --- a/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py +++ b/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py @@ -36,35 +36,72 @@ simpler, single-axis case directly instead of adapting that function. """ -from typing import Any +import gc +from typing import Any, Iterator, Tuple, List import jax from flax import nnx from flax.traverse_util import flatten_dict, unflatten_dict -def unscan_layers( +def _unscan_one_key( + key: Tuple[Any, ...], + value: Any, + num_layers: int, + layer_container: str = "layers", + scan_axis: int = 1, +) -> Tuple[List[Tuple[Tuple[Any, ...], Any]], bool]: + """Unscans a single flattened pytree entry. + + Returns: + (list_of_entries, is_scanned), where list_of_entries is a list of (new_key, value_slice) pairs. + """ + if layer_container not in key: + return [(key, value)], False + + idx = key.index(layer_container) + prefix = key[:idx] + suffix = key[idx + 1 :] + arr = getattr(value, "value", value) + + if not hasattr(arr, "shape") or arr.ndim <= scan_axis: + return [(key, value)], False + + if arr.shape[scan_axis] != num_layers: + raise ValueError( + f"unscan_layers: {'.'.join(str(k) for k in key)!r} has shape {arr.shape}, expected axis {scan_axis} to be" + f" num_layers={num_layers}." + ) + + entries = [] + for i in range(num_layers): + sliced = jax.lax.index_in_dim(arr, i, axis=scan_axis, keepdims=False) + new_key = prefix + (f"{layer_container}_{i}",) + suffix + entries.append((new_key, sliced)) + return entries, True + + +def unscan_layers_streaming( state: Any, num_layers: int, layer_container: str = "layers", scan_axis: int = 1, -) -> Any: - """Splits `state`'s scanned `layer_container` axis into per-layer entries. + *, + keys_per_piece: int = 1, +) -> Iterator[Any]: + """Yields unscanned layer pieces incrementally for Raiden weight sync. Args: state: An `nnx.State` (or any pytree exposing `to_pure_dict`/`to_dict`, or a plain nested dict) of MaxText params, scanned along `scan_axis` under a `layer_container` key (MaxText's default scan layout). - num_layers: Number of layers the scanned axis must have. Used both to validate the input and to bound the unscan - loop. - layer_container: The pytree key holding the scanned per-layer params (MaxText's decoder body uses "layers"). + num_layers: Number of layers the scanned axis must have. + layer_container: The pytree key holding the scanned per-layer params. scan_axis: The axis along which layers are scanned (default 1). + keys_per_piece: Number of original flattened keys to batch per yielded piece (default 1). - Returns: - A nested dict with `layer_container` keys replaced by `f"{layer_container}_{i}"` for each layer `i`, each holding - the corresponding `scan_axis` slice of the original array, wrapped in `nnx.Param` -- matching `nnx.state(..., - nnx.Param)`'s leaf type, so downstream consumers (`raiden_synchronizer.flatten_weights`, which unwraps `.value`) - see the same leaf shape whether or not this transform ran. Non-scanned entries (e.g. embeddings, final norm) pass - through unchanged, also rewrapped. + Yields: + Nested dicts with `nnx.Param`-wrapped leaves, each containing `keys_per_piece` keys' worth of + unscanned slices. """ if hasattr(state, "to_pure_dict"): pure = state.to_pure_dict() @@ -73,57 +110,80 @@ def unscan_layers( elif isinstance(state, dict): pure = state else: - return state + yield state + return flat = flatten_dict(pure) - new_flat = {} - unscanned_count = 0 - - # Drain `flat` as we go (pop, not iterate-then-keep) rather than holding - # every original scanned array alive for the whole function: at 30B-A3B - # scale (padded MoE weights are tens of GB each), keeping both the - # original scanned tree and the ~num_layers-times-larger unscanned tree - # alive simultaneously roughly doubles peak host memory during Raiden's - # D2H staging -- confirmed as the direct cause of an OOMKill there. - for key in list(flat.keys()): - value = flat.pop(key) - if layer_container not in key: - new_flat[key] = value - continue - - idx = key.index(layer_container) - prefix = key[:idx] - suffix = key[idx + 1 :] - arr = getattr(value, "value", value) - - if arr is None or not hasattr(arr, "shape") or getattr(arr, "ndim", 0) <= scan_axis: - # Not a per-layer leaf (shouldn't happen for real params under - # `layers`, but don't silently drop anything unexpected). Keep the - # original (possibly already-wrapped) value, matching pre-existing - # behavior -- unlike the actually-scanned case below, there's no - # multi-copy blowup here worth restructuring around. - new_flat[key] = value - continue - del value - - if arr.shape[scan_axis] != num_layers: - raise ValueError( - f"unscan_layers: {'.'.join(key)!r} has shape {arr.shape}, expected axis {scan_axis} to be" - f" num_layers={num_layers}." - ) - - for i in range(num_layers): - sliced = jax.lax.index_in_dim(arr, i, axis=scan_axis, keepdims=False) - new_key = prefix + (f"{layer_container}_{i}",) + suffix - new_flat[new_key] = sliced - del arr - unscanned_count += 1 - if unscanned_count == 0: + has_scanned = any( + layer_container in key + and hasattr(getattr(flat[key], "value", flat[key]), "shape") + and getattr(getattr(flat[key], "value", flat[key]), "ndim", 0) > scan_axis + for key in flat + ) + if not has_scanned: raise ValueError( f"unscan_layers: found no scanned '{layer_container}' entries to unscan " "-- state may already be unscanned, or layer_container is wrong." ) + keys_per_piece = max(1, keys_per_piece) + flat_keys = list(flat.keys()) + for i in range(0, len(flat_keys), keys_per_piece): + chunk_keys = flat_keys[i : i + keys_per_piece] + piece_flat = {} + for key in chunk_keys: + value = flat.pop(key) + outputs, _ = _unscan_one_key( + key, value, num_layers=num_layers, layer_container=layer_container, scan_axis=scan_axis + ) + for new_key, sliced in outputs: + piece_flat[new_key] = sliced + del value, outputs + + nested = unflatten_dict(piece_flat) + del piece_flat + yield jax.tree_util.tree_map( + lambda x: nnx.Param(x) if not isinstance(x, (nnx.Param, nnx.Variable)) else x, + nested, + ) + + del flat + gc.collect() + + +def unscan_layers( + state: Any, + num_layers: int, + layer_container: str = "layers", + scan_axis: int = 1, +) -> Any: + """Splits `state`'s scanned `layer_container` axis into per-layer entries. + + Args: + state: An `nnx.State` (or any pytree exposing `to_pure_dict`/`to_dict`, or a plain nested dict) of MaxText params, + scanned along `scan_axis` under a `layer_container` key (MaxText's default scan layout). + num_layers: Number of layers the scanned axis must have. Used both to validate the input and to bound the unscan + loop. + layer_container: The pytree key holding the scanned per-layer params (MaxText's decoder body uses "layers"). + scan_axis: The axis along which layers are scanned (default 1). + + Returns: + A nested dict with `layer_container` keys replaced by `f"{layer_container}_{i}"` for each layer `i`, each holding + the corresponding `scan_axis` slice of the original array, wrapped in `nnx.Param` -- matching `nnx.state(..., + nnx.Param)`'s leaf type, so downstream consumers (`raiden_synchronizer.flatten_weights`, which unwraps `.value`) + see the same leaf shape whether or not this transform ran. Non-scanned entries (e.g. embeddings, final norm) pass + through unchanged, also rewrapped. + """ + if not hasattr(state, "to_pure_dict") and not hasattr(state, "to_dict") and not isinstance(state, dict): + return state + + new_flat = {} + for piece in unscan_layers_streaming( + state, num_layers=num_layers, layer_container=layer_container, scan_axis=scan_axis + ): + new_flat.update(flatten_dict(piece)) + gc.collect() nested = unflatten_dict(new_flat) - return jax.tree_util.tree_map(nnx.Param, nested) + del new_flat + return nested diff --git a/src/maxtext/integration/vllm/convert_utils.py b/src/maxtext/integration/vllm/convert_utils.py index d349b40792..34fe4832a5 100644 --- a/src/maxtext/integration/vllm/convert_utils.py +++ b/src/maxtext/integration/vllm/convert_utils.py @@ -160,9 +160,40 @@ class ShapeMismatchError(ValueError): """Raised when source and target shapes are incompatible.""" -def _apply_dtype_cast(val: jax.Array | np.ndarray, tgt_dtype: jnp.dtype, src_key: str) -> jax.Array | np.ndarray: +def reclaim_host_memory() -> None: + """Runs garbage collection and triggers libc malloc_trim to return free heap to the OS.""" + import gc # pylint: disable=g-import-not-at-top + gc.collect() + try: + import ctypes # pylint: disable=g-import-not-at-top + ctypes.CDLL("libc.so.6").malloc_trim(0) + except Exception as e: # pylint: disable=broad-exception-caught + logging.debug("reclaim_host_memory: malloc_trim unavailable or failed: %s", e) + + +def normalize_dtype(tgt_dtype: Any) -> Any: + """Normalizes string or dtype representations into a standard jnp.dtype.""" + if tgt_dtype is None: + return None + if isinstance(tgt_dtype, str): + if tgt_dtype in ("bfloat16", "bf16"): + return jnp.bfloat16 + if tgt_dtype in ("float32", "fp32"): + return jnp.float32 + return jnp.dtype(tgt_dtype) + return tgt_dtype + + +def _apply_dtype_cast(val: Any, tgt_dtype: Any, src_key: str) -> Any: """Casts val to target dtype if needed, logging a warning on type mismatch.""" - if val.dtype != tgt_dtype: + tgt_dtype = normalize_dtype(tgt_dtype) + if isinstance(val, jax.ShapeDtypeStruct): + if tgt_dtype is not None and val.dtype != tgt_dtype: + return jax.ShapeDtypeStruct(val.shape, tgt_dtype) + return val + if not hasattr(val, "dtype"): + return val + if tgt_dtype is not None and val.dtype != tgt_dtype: logging.log_first_n( logging.WARNING, "Type mismatch on %s: %s -> %s", @@ -171,7 +202,8 @@ def _apply_dtype_cast(val: jax.Array | np.ndarray, tgt_dtype: jnp.dtype, src_key val.dtype, tgt_dtype, ) - return val.astype(tgt_dtype) + if hasattr(val, "astype"): + return val.astype(tgt_dtype) return val @@ -409,6 +441,8 @@ def _align_per_axis( path here is bulk alignment of scanned MoE weights, where eager dispatch was costing tens of seconds per tensor. """ + if isinstance(arr, jax.ShapeDtypeStruct): + return jax.ShapeDtypeStruct(tgt_shape, arr.dtype) if not hasattr(arr, "shape"): return arr if arr.shape == tgt_shape: @@ -530,7 +564,7 @@ def _fuse_and_unstack_moe( scan_axis: int, n_shards: int, tgt_shape: Tuple[int, ...], - scan_fused_axis: int, + scan_fused_axis: int, # TODO(follow-up): Unused in function body, preserved for caller compatibility. tgt_fused_axis: int, ) -> Tuple[jax.Array | np.ndarray, ...]: """Fuses wi_0/wi_1 per unstacked layer to keep peak intermediate HBM allocation low. @@ -604,6 +638,11 @@ def _bulk_align_and_unstack( A tuple of `num_layers` per-layer arrays at the per-layer target shape. """ per_layer_shape = per_layer_tgt_val.shape + if isinstance(arr, jax.ShapeDtypeStruct) or isinstance(per_layer_tgt_val, jax.ShapeDtypeStruct): + num_layers = arr.shape[scan_axis] + tgt_dtype = getattr(per_layer_tgt_val, "dtype", getattr(arr, "dtype", jnp.float32)) + return tuple(jax.ShapeDtypeStruct(per_layer_shape, tgt_dtype) for _ in range(num_layers)) + scanned_tgt_shape = per_layer_shape[:scan_axis] + (arr.shape[scan_axis],) + per_layer_shape[scan_axis:] scanned_tgt_sharding = _scanned_sharding_from_per_layer(getattr(per_layer_tgt_val, "sharding", None), scan_axis) diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py index 3b0a3c8f05..e633049335 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py @@ -48,19 +48,7 @@ class AttentionMetadata: _HYBRID_LAYER_IMBALANCE_THRESHOLD = 1.5 -def next_power_of_two(x: int) -> int: - """Finds the smallest power of 2 >= x using bit manipulation. - - Args: - x: The input number (should be an integer). - - Returns: - The smallest integer power of 2 that is >= x. - """ - assert x > 0 - if x == 1: - return 1 - return 1 << (x - 1).bit_length() +from maxtext.integration.vllm.moe_padding import compute_padded_moe_mlp_dim, next_power_of_two def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters: @@ -166,11 +154,8 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters # The GMM_v2 kernel requires the MLP dimension per expert to be at least 2x the number of TPU lanes # to ensure efficient execution. See the validate_inputs() method in the following file for more details: # https://github.com/vllm-project/tpu-inference/blob/main/tpu_inference/kernels/megablox/gmm_v2.py - if hidden_size is not None and (hidden_size // moe_mlp_tp_size) % (2 * num_lanes) != 0: - padded_hidden_size = next_power_of_two(hidden_size) - while (padded_hidden_size // moe_mlp_tp_size) < (2 * num_lanes): - padded_hidden_size = next_power_of_two(padded_hidden_size + 1) - + padded_hidden_size = compute_padded_moe_mlp_dim(hidden_size, moe_mlp_tp_size, num_lanes) + if padded_hidden_size is not None and padded_hidden_size != hidden_size: # This inflates every expert weight, so it is a real memory/FLOP cost rather than a # cosmetic reshape: at moe_mlp_tp_size=4 a 512-wide MoE is padded to 1024 (2x the MoE # weights), and at moe_mlp_tp_size=8 to 2048 (4x). Log it at WARNING so it is visible diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py index 1cd7853f96..d8f2ace4a3 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py @@ -464,17 +464,25 @@ def __init__( self, 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.converter = converter self._direct_maxtext_sync = direct_maxtext_sync self._scan_axis = scan_axis self._layer_pattern_length = layer_pattern_length + model_config = getattr(config, "model_config", None) + model_name = getattr(model_config, "model", "") or "" + architectures = getattr(model_config, "architectures", []) or [] + hf_config = getattr(model_config, "hf_config", None) + model_type = getattr(hf_config, "model_type", "") or "" + arch_str = " ".join(str(a) for a in architectures) + self._is_gemma = ( + "gemma" in str(model_name).lower() + or "gemma" in str(model_type).lower() + or "gemma" in str(arch_str).lower() + ) def update_params( self, @@ -494,12 +502,7 @@ def update_params( pass raise if self._converter is None: - if self._direct_maxtext_sync: - updated_weights = unroll_qwen_scanned_weights( - updated_weights, - scan_axis=self._scan_axis, - pattern_length=self._layer_pattern_length, - ) + if self._direct_maxtext_sync and self._is_gemma: updated_weights = unroll_gemma_scanned_weights(updated_weights) try: return super().update_params(updated_weights, filter_types) @@ -742,7 +745,6 @@ def __init__( 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), diff --git a/src/maxtext/integration/vllm/moe_padding.py b/src/maxtext/integration/vllm/moe_padding.py new file mode 100644 index 0000000000..ce5847704f --- /dev/null +++ b/src/maxtext/integration/vllm/moe_padding.py @@ -0,0 +1,61 @@ +# 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 padding utilities for TPU GMM_v2 kernel alignment.""" + +from typing import Optional + + +def next_power_of_two(x: int) -> int: + """Finds the smallest power of 2 >= x using bit manipulation. + + Args: + x: The input number (should be an integer > 0). + + Returns: + The smallest integer power of 2 that is >= x. + """ + assert x > 0 + if x == 1: + return 1 + return 1 << (x - 1).bit_length() + + +def compute_padded_moe_mlp_dim( + hidden_size: Optional[int], + moe_mlp_tp_size: int, + num_lanes: int = 128, +) -> Optional[int]: + """Computes padded MoE intermediate size for GMM_v2 kernel requirements. + + The GMM_v2 kernel requires the MLP dimension per expert to be at least 2x the + number of TPU lanes (e.g. 2 * 128 = 256) to ensure efficient execution. + + Args: + hidden_size: Unpadded MoE intermediate size (e.g. moe_intermediate_size / + base_moe_mlp_dim). + moe_mlp_tp_size: TP size across MLP dimensions (e.g. tp * attn_dp). + num_lanes: Number of TPU lanes (typically 128 for TPU v5p/v6e). + + Returns: + Padded hidden size, or hidden_size if no padding is required / hidden_size is None. + """ + if hidden_size is None or moe_mlp_tp_size <= 0 or num_lanes <= 0: + return hidden_size + + if (hidden_size // moe_mlp_tp_size) % (2 * num_lanes) != 0: + min_required = 2 * num_lanes * moe_mlp_tp_size + return next_power_of_two(max(hidden_size, min_required)) + + return hidden_size diff --git a/src/maxtext/integration/vllm/torchax_converter/__init__.py b/src/maxtext/integration/vllm/torchax_converter/__init__.py index f3582c0090..1ec0362971 100644 --- a/src/maxtext/integration/vllm/torchax_converter/__init__.py +++ b/src/maxtext/integration/vllm/torchax_converter/__init__.py @@ -11,3 +11,15 @@ # 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. + +from maxtext.integration.vllm.torchax_converter.base import BaseMaxTextToVLLMConverter +from maxtext.integration.vllm.torchax_converter.qwen35_moe import Qwen35MaxTextToVLLMConverter +from maxtext.integration.vllm.torchax_converter.qwen3_moe import Qwen3MaxTextToVLLMConverter +from maxtext.integration.vllm.torchax_converter.gemma4_moe import Gemma4MaxTextToVLLMConverter + +__all__ = [ + "BaseMaxTextToVLLMConverter", + "Qwen35MaxTextToVLLMConverter", + "Qwen3MaxTextToVLLMConverter", + "Gemma4MaxTextToVLLMConverter", +] diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py index c8e79ff535..ad6dfb634c 100644 --- a/src/maxtext/integration/vllm/weight_converter.py +++ b/src/maxtext/integration/vllm/weight_converter.py @@ -16,12 +16,14 @@ import abc import dataclasses +import gc import logging +import os import re import jax import jax.numpy as jnp -import gc -from typing import List, Union, Any, Dict, Optional, Mapping, Tuple +import numpy as np +from typing import List, Union, Any, Dict, Optional, Mapping, Tuple, Iterator from flax import traverse_util, nnx from maxtext.integration.vllm.convert_utils import ( _align_per_axis, @@ -33,8 +35,12 @@ _jit_unstack, _scanned_sharding_from_per_layer, _sharding_summary, + normalize_dtype, + reclaim_host_memory, ) +_MOE_MLP_WEIGHTS = frozenset({"wi_0", "wi_1", "wo", "wi"}) + # ========================================== # 1. Operations @@ -276,12 +282,15 @@ def __init__( num_kv_heads: Optional[int] = None, head_dim: Optional[int] = None, config: Any = None, - # Defaults to MoEFusedLayout.PER_SHARD_INTERLEAVE; resolved in the body - # because MoEFusedLayout is defined further down this module. + trainer_config: Any = None, + rollout_backend: str = "maxtext", moe_fused_layout: Optional[str] = None, allow_unused_source_keys: Tuple[str, ...] = (), debug: bool = False, + prefuse_moe_weights: Optional[bool] = None, + target_dtype: Optional[Any] = None, ): + config = trainer_config if config is None else config if rules is not None and not rules: raise ValueError( "WeightConverter(rules=[]) would convert nothing and leave the " @@ -295,9 +304,10 @@ def __init__( # Read by the rollout engine to decide whether to trace the reshard # step that runs after conversion. self.debug = debug + self.rollout_backend = rollout_backend self._direct: Optional["MaxTextToMaxTextConverter"] = None - if rules is None: + if rollout_backend == "maxtext" and rules is None: if config is None: raise ValueError( "WeightConverter(rules=None) needs `config` to derive the " @@ -309,20 +319,29 @@ def __init__( moe_fused_layout=(moe_fused_layout or MoEFusedLayout.PER_SHARD_INTERLEAVE), allow_unused_source_keys=allow_unused_source_keys, debug=debug, + prefuse_moe_weights=prefuse_moe_weights, + target_dtype=target_dtype, ) logging.info("WeightConverter: direct MaxText-to-MaxText mode (debug=%s).", debug) else: + if self.rules is None and config is not None: + model_name = getattr(config, "model_name", "") + if model_name in MODEL_TO_CONVERSION_RULES and MODEL_TO_CONVERSION_RULES[model_name] is not None: + self.rules = MODEL_TO_CONVERSION_RULES[model_name] logging.info( "WeightConverter: torchax rule mode (tp=%d, %d rules).", self.tp, - len(rules), + len(self.rules) if self.rules else 0, ) def convert(self, src_pytree: Any, target_state: Any = None) -> Any: """Converts source weights pytree into target format using rules or direct converter.""" - if self.rules is None: + if self.rollout_backend == "maxtext" and self.rules is None: return self._direct.convert(src_pytree, target_state=target_state) + if self.rules is None: + raise ValueError("WeightConverter in torchax mode requires conversion rules.") + flat_src = traverse_util.flatten_dict(_to_pure_dict(src_pytree), sep=".") gc.collect() @@ -365,6 +384,20 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Any: return _rekey_to_target(result, target_state) + def convert_streaming( + self, + src_pytree: Any, + target_state: Any = None, + *, + groups_per_piece: int = 1, + ) -> Iterator[Dict[str, Any]]: + """Yields converted weight pieces incrementally in direct MaxText-to-MaxText mode.""" + if self.rollout_backend == "maxtext" and self.rules is None: + return self._direct.convert_streaming(src_pytree, target_state=target_state, groups_per_piece=groups_per_piece) + raise NotImplementedError( + "convert_streaming is only supported in direct MaxText-to-MaxText mode (rollout_backend='maxtext' and rules=None)." + ) + # ========================================== # 4. Registries and Builders @@ -550,10 +583,15 @@ def _group_plan(plan: List[_PlanEntry]) -> List[_PlanGroup]: return groups -class ConversionPlanError(ValueError): +class WeightConverterError(Exception): + """Base class for all weight converter exceptions.""" + + +class ConversionPlanError(WeightConverterError, ValueError): """Raised when the source and target trees cannot be fully reconciled.""" + def _is_non_weight_path(key_tuple: Tuple[Any, ...]) -> bool: return any(isinstance(part, str) and part.lstrip("_").startswith(_NON_WEIGHT_PATH_PREFIXES) for part in key_tuple) @@ -657,11 +695,29 @@ def __init__( moe_fused_layout: str = MoEFusedLayout.PER_SHARD_INTERLEAVE, allow_unused_source_keys: Tuple[str, ...] = (), debug: bool = False, + prefuse_moe_weights: Optional[bool] = None, + target_dtype: Optional[Any] = None, + is_pathways: Optional[bool] = None, ): self.config = config self.moe_fused_layout = moe_fused_layout self.allow_unused_source_keys = allow_unused_source_keys self.debug = debug + self.prefuse_moe_weights = ( + prefuse_moe_weights + if prefuse_moe_weights is not None + else getattr(config, "prefuse_moe_weights", False) + ) + self.padded_base_moe_mlp_dim = getattr(config, "padded_base_moe_mlp_dim", None) + self.target_dtype = target_dtype if target_dtype is not None else getattr(config, "weight_dtype", None) + + if is_pathways is not None: + self.is_pathways = is_pathways + else: + backend_platform = getattr(jax.devices()[0], "platform", "").lower() if jax.devices() else "" + self.is_pathways = (backend_platform == "proxy") or ( + "proxy" in os.environ.get("JAX_PLATFORMS", "") and bool(os.environ.get("JAX_BACKEND_TARGET")) + ) self.cycle = int(getattr(config, "inhomogeneous_layer_cycle_interval", 1) or 1) self.num_decoder_layers = int(config.num_decoder_layers) @@ -678,14 +734,20 @@ def __init__( self._groups: Optional[List[_PlanGroup]] = None logging.info( - "MaxTextToMaxTextConverter: %d layers, cycle=%d, %d scanned blocks, " "scan_axis=%d, moe_fused_layout=%s", + "MaxTextToMaxTextConverter: %d layers, cycle=%d, %d scanned blocks, " + "scan_axis=%d, moe_fused_layout=%s, prefuse_moe=%s, padded_moe_dim=%s", self.num_decoder_layers, self.cycle, self.num_blocks, self.scan_axis, self.moe_fused_layout, + self.prefuse_moe_weights, + self.padded_base_moe_mlp_dim, ) + def _resolve_target_dtype(self): + return normalize_dtype(self.target_dtype) + # -------------------------------------------------------------- # # Plan construction # -------------------------------------------------------------- # @@ -704,6 +766,64 @@ def _scanned_candidates( # Homogeneous: a single scanned `layers` container. return [prefix + ("layers",) + suffix] + def _build_target_free_plan( + self, + src_flat: Mapping[Tuple[Any, ...], Any], + ) -> List[_PlanEntry]: + """Builds the conversion plan directly from source keys and config without target state.""" + plan: List[_PlanEntry] = [] + consumed_wi_1 = set() + + for src_key in src_flat: + if _is_non_weight_path(src_key): + continue + if src_key in consumed_wi_1: + continue + + if "layers" not in src_key: + plan.append(_PlanEntry(src_key, (src_key,), None, "identity")) + continue + + idx = src_key.index("layers") + prefix = src_key[:idx] + rest = src_key[idx + 1 :] + + if self.cycle == 1: + slot = 0 + suffix = rest + else: + # Inhomogeneous hybrid cycle: ("decoder", "layers", "layer_0", "input_layernorm", "scale") + slot_token = rest[0] + if isinstance(slot_token, str) and slot_token.startswith("layer_"): + slot = int(slot_token[6:]) + elif isinstance(slot_token, str) and slot_token.isdigit(): + slot = int(slot_token) + elif isinstance(slot_token, int): + slot = slot_token + else: + raise ConversionPlanError(f"Unexpected slot token {slot_token!r} in key {src_key}") + suffix = rest[1:] + + is_wi_0 = bool(suffix and suffix[-1] == "wi_0") + wi_1_key = src_key[:-1] + ("wi_1",) if is_wi_0 else None + fuse_moe = self.prefuse_moe_weights and is_wi_0 and (wi_1_key in src_flat) + + if fuse_moe: + consumed_wi_1.add(wi_1_key) + for b in range(self.num_blocks): + global_idx = b * self.cycle + slot + tgt_key = prefix + (f"layers_{global_idx}",) + suffix[:-1] + ("wi",) + plan.append(_PlanEntry(tgt_key, (src_key, wi_1_key), b, "fuse_moe")) + elif self.prefuse_moe_weights and suffix and suffix[-1] == "wi_1" and (src_key[:-1] + ("wi_0",) in src_flat): + continue + else: + for b in range(self.num_blocks): + global_idx = b * self.cycle + slot + tgt_key = prefix + (f"layers_{global_idx}",) + suffix + plan.append(_PlanEntry(tgt_key, (src_key,), b, "slice")) + + return plan + def _build_plan( self, src_flat: Mapping[Tuple[Any, ...], Any], @@ -764,13 +884,7 @@ def _build_plan( return plan def _validate_plan(self, src_flat, tgt_flat, unmatched, consumed) -> None: - """Fails loudly rather than leaving rollout weights at their dummy values. - - vLLM boots with `load_format="dummy"`, so a target leaf we never write - keeps *random* weights. That produces a quietly wrong reward curve - instead of an error, which is far more expensive to debug than a crash - at startup. - """ + """Fails loudly rather than leaving rollout weights at their dummy values.""" if unmatched: shown = "\n ".join(".".join(map(str, k)) for k in sorted(unmatched)[:40]) raise ConversionPlanError( @@ -804,17 +918,9 @@ def _validate_plan(self, src_flat, tgt_flat, unmatched, consumed) -> None: # Plan execution # -------------------------------------------------------------- # def _fuse_moe_bulk(self, wi_0, wi_1, tgt_val, key_path: str): - """Fuses the *scanned* gate/up kernels, returning one array per block. - - `wi_0`/`wi_1` still carry `num_blocks` at `scan_axis`; `tgt_val` is a - single per-layer target leaf, supplying the fused shape and sharding - that every block in this group shares. - """ + """Fuses the *scanned* gate/up kernels, returning one array per block.""" tgt_shape = tgt_val.shape - # MaxText stores MoE kernels as (experts, in_dim, intermediate); the - # gate/up fusion always doubles the trailing intermediate axis. tgt_fused_axis = len(tgt_shape) - 1 - # Same logical axis, shifted by the scan dim the trainer inserted. scan_fused_axis = tgt_fused_axis if tgt_fused_axis < self.scan_axis else tgt_fused_axis + 1 if self.moe_fused_layout == MoEFusedLayout.PER_SHARD_INTERLEAVE: @@ -839,14 +945,107 @@ def _fuse_moe_bulk(self, wi_0, wi_1, tgt_val, key_path: str): raise ConversionPlanError(f"Unknown moe_fused_layout: {self.moe_fused_layout!r}") - def _execute_group(self, group: _PlanGroup, src_flat, tgt_flat): - """Produces every target leaf in `group`. Returns (target_key, array) pairs. + def _slice_bulk_target_free(self, val: Any, path: str): + last_key = path.split(".")[-1] + if isinstance(val, jax.ShapeDtypeStruct): + unrolled_shape = list(val.shape[: self.scan_axis] + val.shape[self.scan_axis + 1 :]) + if last_key in _MOE_MLP_WEIGHTS and self.padded_base_moe_mlp_dim is not None: + if last_key == "wo": + if self.padded_base_moe_mlp_dim > unrolled_shape[1]: + unrolled_shape[1] = self.padded_base_moe_mlp_dim + elif last_key in ("wi_0", "wi_1", "wi"): + if self.padded_base_moe_mlp_dim > unrolled_shape[-1]: + unrolled_shape[-1] = self.padded_base_moe_mlp_dim + return tuple(jax.ShapeDtypeStruct(tuple(unrolled_shape), val.dtype) for _ in range(val.shape[self.scan_axis])) + + if last_key in _MOE_MLP_WEIGHTS and self.padded_base_moe_mlp_dim is not None: + if last_key == "wo": + intermediate_axis = 2 + if self.padded_base_moe_mlp_dim > val.shape[intermediate_axis]: + pad_amount = self.padded_base_moe_mlp_dim - val.shape[intermediate_axis] + pad_spec = [(0, 0)] * val.ndim + pad_spec[intermediate_axis] = (0, pad_amount) + val = jnp.pad(val, pad_spec) + elif last_key in ("wi_0", "wi_1"): + intermediate_axis = len(val.shape) - 1 + if self.padded_base_moe_mlp_dim > val.shape[intermediate_axis]: + pad_amount = self.padded_base_moe_mlp_dim - val.shape[intermediate_axis] + pad_spec = [(0, 0)] * val.ndim + pad_spec[intermediate_axis] = (0, pad_amount) + val = jnp.pad(val, pad_spec) + + return _jit_unstack(val, self.scan_axis) + + def _fuse_moe_bulk_target_free(self, wi_0: Any, wi_1: Any, path: str): + unpadded_dim = wi_0.shape[-1] + target_intermediate = ( + self.padded_base_moe_mlp_dim + if (self.padded_base_moe_mlp_dim is not None and self.padded_base_moe_mlp_dim > unpadded_dim) + else unpadded_dim + ) + if isinstance(wi_0, jax.ShapeDtypeStruct): + fused_shape = (wi_0.shape[0], wi_0.shape[2], 2 * target_intermediate) + return tuple(jax.ShapeDtypeStruct(fused_shape, wi_0.dtype) for _ in range(wi_0.shape[self.scan_axis])) - The scanned source is cast, aligned and fused *once*; the per-layer - arrays are then read out of a single unstack. Every target in a group - shares a shape and sharding by construction, so the first one is a - sound stand-in for all of them. - """ + tgt_shape = (wi_0.shape[0], wi_0.shape[2], 2 * target_intermediate) + tgt_fused_axis = len(tgt_shape) - 1 + scan_fused_axis = tgt_fused_axis if tgt_fused_axis < self.scan_axis else tgt_fused_axis + 1 + + if self.moe_fused_layout == MoEFusedLayout.PER_SHARD_INTERLEAVE: + n_shards = _get_n_shards(wi_0, scan_fused_axis) + return _fuse_and_unstack_moe( + wi_0, + wi_1, + self.scan_axis, + n_shards, + tgt_shape, + scan_fused_axis, + tgt_fused_axis, + ) + + if self.moe_fused_layout == MoEFusedLayout.CONCAT: + if target_intermediate > unpadded_dim: + pad_spec = [(0, 0)] * wi_0.ndim + pad_spec[-1] = (0, target_intermediate - unpadded_dim) + wi_0 = jnp.pad(wi_0, pad_spec) + wi_1 = jnp.pad(wi_1, pad_spec) + fused = jnp.concatenate([wi_0, wi_1], axis=scan_fused_axis) + return _jit_unstack(fused, self.scan_axis) + + raise ConversionPlanError(f"Unknown moe_fused_layout: {self.moe_fused_layout!r}") + + def _execute_group_target_free(self, group: _PlanGroup, src_flat): + path = group.source_path + target_dtype = self._resolve_target_dtype() + + if group.op == "identity": + raw_val = src_flat[group.source_keys[0]] + val = _apply_dtype_cast(raw_val, target_dtype, path) + return [(tgt_key, val) for _, tgt_key in group.targets] + + if any(idx is None for idx, _ in group.targets): + raise ConversionPlanError( + f"Plan group for {path} has op={group.op!r} but a target with no " + "scan index; only 'identity' targets may omit one." + ) + + if group.op == "fuse_moe": + raw_0 = src_flat[group.source_keys[0]] + raw_1 = src_flat[group.source_keys[1]] + wi_0, wi_1 = (_apply_dtype_cast(raw_0, target_dtype, path), _apply_dtype_cast(raw_1, target_dtype, path)) + self._check_scan_axis(wi_0, path) + per_block = self._fuse_moe_bulk_target_free(wi_0, wi_1, path) + return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets] + + # group.op == "slice" + raw_val = src_flat[group.source_keys[0]] + val = _apply_dtype_cast(raw_val, target_dtype, path) + self._check_scan_axis(val, path) + per_block = self._slice_bulk_target_free(val, path) + return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets] + + def _execute_group(self, group: _PlanGroup, src_flat, tgt_flat): + """Produces every target leaf in `group`. Returns (target_key, array) pairs.""" first_tgt = tgt_flat[group.targets[0][1]] path = group.source_path @@ -886,24 +1085,22 @@ def _check_scan_axis(self, val, path: str) -> None: def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]: """Returns a nested dict of rollout weights, keyed by target paths. - Pure: neither `src_pytree` nor `target_state` is mutated. + Pure: neither `src_pytree` nor `target_state` is mutated. Leaves are wrapped in nnx.Param. """ if target_state is None: - raise ValueError( - "MaxTextToMaxTextConverter requires target_state to resolve the " - "rollout's parameter shapes, shardings and dtypes." - ) + flat_result = {} + for piece in self.convert_streaming(src_pytree, target_state=None): + flat_result.update(traverse_util.flatten_dict(piece)) + reclaim_host_memory() + return traverse_util.unflatten_dict(flat_result) + + src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree)) + src_flat, _ = _strip_root(src_flat, "base") # Read variable types before purifying to plain arrays loses them. skip_paths = _non_param_paths(target_state) - - src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree)) tgt_flat = traverse_util.flatten_dict(_to_pure_dict(target_state)) - # The trainer wraps the model in TunixMaxTextAdapter ("base"); the - # rollout may nest it under one or more "model" levels. Strip both so - # the plan is expressed in a single coordinate system, then re-wrap. - src_flat, _ = _strip_root(src_flat, "base") tgt_flat, tgt_root = _strip_root(tgt_flat, "model") if tgt_root: depth = len(tgt_root) @@ -930,17 +1127,13 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]: if self.debug: for k in group.source_keys: logging.info( - "weight_sync_debug: op=%s source=%s (%d targets) | src %s " "| tgt %s", + "weight_sync_debug: op=%s source=%s (%d targets) | src %s | tgt %s", group.op, ".".join(map(str, k)), len(group.targets), _sharding_summary(src_flat[k]), _sharding_summary(tgt_flat[group.targets[0][1]]), ) - # JAX dispatch is asynchronous, so without a barrier a device - # failure surfaces at an arbitrary later point and the traceback - # names the wrong parameter. Pay the serialization to find out - # which source parameter is actually at fault. try: outs = self._execute_group(group, src_flat, tgt_flat) jax.block_until_ready([out for _, out in outs]) @@ -965,18 +1158,74 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]: else: outs = self._execute_group(group, src_flat, tgt_flat) + for k in group.source_keys: + src_flat.pop(k, None) + for tgt_key, out in outs: tgt_val = tgt_flat[tgt_key] - if out.shape != tgt_val.shape: + if hasattr(out, "shape") and hasattr(tgt_val, "shape") and out.shape != tgt_val.shape: raise ConversionPlanError( f"Shape mismatch after conversion for " f"{'.'.join(map(str, tgt_key))}: produced {out.shape}, " f"rollout expects {tgt_val.shape}." ) result[tgt_root + tgt_key] = out + del outs + del src_flat, tgt_flat gc.collect() - return traverse_util.unflatten_dict(result) + nested = traverse_util.unflatten_dict(result) + del result + reclaim_host_memory() + return jax.tree_util.tree_map( + lambda x: nnx.Param(x) if not isinstance(x, (nnx.Param, nnx.Variable)) else x, + nested, + ) + + def convert_streaming( + self, + src_pytree: Any, + target_state: Any = None, + *, + groups_per_piece: int = 1, + ) -> Iterator[Dict[str, Any]]: + """Yields converted rollout weight pieces incrementally for target-free conversion. + + Pure: `src_pytree` is not mutated. Each yielded piece is a nested dict of `nnx.Param`s + corresponding to `groups_per_piece` plan groups. Memory is freed piece-by-piece as + source keys are consumed. + """ + if target_state is not None: + raise NotImplementedError("convert_streaming only supports target-free conversion (target_state=None).") + + src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree)) + src_flat, src_root = _strip_root(src_flat, "base") + + if self._plan is None: + self._plan = self._build_target_free_plan(src_flat) + self._groups = _group_plan(self._plan) + + groups_per_piece = max(1, groups_per_piece) + for i in range(0, len(self._groups), groups_per_piece): + piece_groups = self._groups[i : i + groups_per_piece] + piece_result: Dict[Tuple[Any, ...], Any] = {} + for group in piece_groups: + outs = self._execute_group_target_free(group, src_flat) + for k in group.source_keys: + src_flat.pop(k, None) + for tgt_key, out in outs: + piece_result[src_root + tgt_key] = out + del outs + + nested = traverse_util.unflatten_dict(piece_result) + del piece_result + yield jax.tree_util.tree_map( + lambda x: nnx.Param(x) if not isinstance(x, (nnx.Param, nnx.Variable)) else x, + nested, + ) + + del src_flat + reclaim_host_memory() def _rekey_to_target(flat_dotted: Dict[str, Any], target_state: Any) -> Dict[str, Any]: diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 356b07eeb1..a03c2639a0 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -23,6 +23,7 @@ from collections.abc import Callable, Mapping import contextlib import dataclasses +import gc import os from typing import Any @@ -37,6 +38,7 @@ from maxtext.common import train_state_nnx from maxtext.configs import pyconfig from maxtext.integration.tunix.weight_mapping import raiden_unscan +from maxtext.integration.vllm.convert_utils import reclaim_host_memory from maxtext.trainers.pre_train import train as maxtext_train from maxtext.training_engine import abstract_engine from maxtext.training_engine import checkpointing @@ -137,6 +139,14 @@ def _batch_signature(dynamic_batch: Any, static_batch: dict[str, Any]) -> Any: return (treedef, shapes, static_batch) +_REPLICATED_BATCH_DIM_WARNING = ( + "Loss input with batch dim %d does not divide mesh axis %r (size %d), so that " + "dimension is replicated instead of sharded: every device along the axis holds and " + "computes the whole micro-batch, %dx the work a sharded one would do there. Results " + "stay correct. If it was not deliberate -- a sequence-packed micro-batch is always " + "size 1 and has no alternative -- make the micro-batch a multiple of the axis size." +) + _UNCOMPARABLE_SIGNATURE_WARNING = ( "Could not compare %s between fwd_bwd calls (%s), so the engine cannot tell whether " "the compiled kernel is still valid and will recompile on EVERY fwd_bwd from now on. %s" @@ -351,6 +361,7 @@ def __init__( self._compile_requested = False self._compiled_signature: Any = None self._signature_compare_warned: bool = False + self._replicated_batch_warned: bool = False if not training_config.model_name: raise ValueError("training_config.model_name must be specified") model_or_model_mesh_pair = model_creation_utils.from_pretrained( @@ -399,6 +410,35 @@ def __init__( self._metrics_recorder = metrics_module.MetricsRecorder() self._throttler = inflight_throttler.InflightThrottler(config=self._config) self._raiden_sync: Any = None + self._last_staged_step: Optional[int] = None + self._staged_metadata: Any = None + vllm_cfg = getattr(self._config, "vllm", {}) + if isinstance(vllm_cfg, dict): + vllm_use_wc = vllm_cfg.get("use_weight_converter", False) + vllm_backend = vllm_cfg.get("rollout_backend", "maxtext") + else: + vllm_use_wc = getattr(vllm_cfg, "use_weight_converter", False) + vllm_backend = getattr(vllm_cfg, "rollout_backend", "maxtext") + + self._use_weight_converter = bool( + getattr(self._config, "use_weight_converter", False) + or vllm_use_wc + or os.environ.get("USE_WEIGHT_CONVERTER", "0").lower() in ("1", "true", "yes") + ) + self._rollout_backend = ( + getattr(self._config, "rollout_backend", None) + or vllm_backend + or os.environ.get("ROLLOUT_BACKEND", "maxtext") + ) + if self._use_weight_converter: + from maxtext.integration.vllm.weight_converter import WeightConverter # pylint: disable=g-import-not-at-top,import-outside-toplevel + self._weight_converter = WeightConverter( + config=self._config, + rollout_backend=self._rollout_backend, + debug=getattr(self._config, "weight_sync_debug", False), + ) + else: + self._weight_converter = None @property def model(self) -> Any: @@ -875,7 +915,8 @@ def _batch_data_shardings(self, dynamic_batch: Any) -> Any: sequence-packed micro-batch, always size 1) replicates that dim instead of sharding it -- every device holds and computes on the same data with no cross-device split, which is correct (there's nothing to reduce back together afterwards) but wastes - compute across the axis for that micro-batch. + compute across the axis for that micro-batch. That is an N-fold cost, so it warns + once per instance rather than living only in this docstring. """ data_sharding = sharding.get_input_data_sharding(self._config, self._mesh) data_spec = tuple(data_sharding.spec) @@ -885,8 +926,16 @@ def leaf_sharding(leaf): return None rank = jnp.ndim(leaf) spec = list(data_spec[:rank]) - if spec and spec[0] is not None and leaf.shape[0] % self._batch_axis_size(spec[0]): - spec[0] = None + if spec and spec[0] is not None: + axis_size = self._batch_axis_size(spec[0]) + if leaf.shape[0] % axis_size: + # Warn once per instance, not per leaf: this runs under a tree_map over every + # loss input, and they normally share a batch dim. Silence here would leave an + # N-fold compute cliff visible only in a docstring. + if not self._replicated_batch_warned: + self._replicated_batch_warned = True + logging.warning(_REPLICATED_BATCH_DIM_WARNING, leaf.shape[0], spec[0], axis_size, axis_size) + spec[0] = None return jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec(*spec)) return jax.tree.map(leaf_sharding, dynamic_batch) @@ -1434,45 +1483,67 @@ def prepare_weight_sync( " tunix build that ships it, or select a different staging_transport." ) from exc + if ( + self._raiden_sync is not None + and self._last_staged_step == self.train_step + and self._staged_metadata is not None + ): + logging.info( + "Trainer re-using staged weight sync for step %d (%d variables)", + self.train_step, + sum(len(m.variables) for m in self._staged_metadata), + ) + return self._staged_metadata + # 1. Drain all in-flight TPU computations to ensure weights are fully updated self._throttler.wait_for_all() + reclaim_host_memory() # 2. Extract clean trainable parameters params_state = self._get_trainable_params_state() - # 2a. The trainer keeps float32 master weights, but the rollout side - # (MaxTextForCausalLM under configs/inference/vllm.yml) loads/serves in - # bfloat16 -- Raiden's manifest preflight rejects a dtype/item_size - # mismatch, and binding mismatched-dtype buffers would be wrong anyway. - # Cast the synced copy down; the trainer's own params_state (used for - # the actual optimizer step) is untouched since this is a fresh tree. - params_state = jax.tree_util.tree_map( - lambda x: x.astype(jnp.bfloat16) if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating) else x, - params_state, - ) - - # 2b. The trainer runs scanned (scan_layers=True) for training speed, but - # the rollout side loads its MaxText model unscanned (MaxTextForCausalLM - # under configs/inference/vllm.yml has scan_layers=False). Raiden matches - # tensors by name, so unscan here -- on the trainer side only -- so the - # names/shapes we bind already match what the sampler reports. - if self._config.scan_layers: - params_state = raiden_unscan.unscan_layers( + if self._use_weight_converter: + if self._weight_converter is None: + from maxtext.integration.vllm.weight_converter import WeightConverter # pylint: disable=g-import-not-at-top,import-outside-toplevel + self._weight_converter = WeightConverter( + config=self._config, + rollout_backend=self._rollout_backend, + debug=getattr(self._config, "weight_sync_debug", False), + ) + converted_state = self._weight_converter.convert(params_state) + else: + # UNCHANGED, deliberately out of scope: this fp32->bf16 cast is an + # on-device (HBM, not host RAM) full materialization -- a different + # memory pool than the host OOM this plan addresses. Candidate + # fast-follow: fold into unscan_layers_streaming's per-piece slicing. + params_state = jax.tree_util.tree_map( + lambda x: x.astype(jnp.bfloat16) if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating) else x, params_state, - num_layers=self._config.num_decoder_layers, - scan_axis=self._config.param_scan_axis, ) + if self._config.scan_layers: + converted_state = raiden_unscan.unscan_layers( + params_state, + num_layers=self._config.num_decoder_layers, + scan_axis=self._config.param_scan_axis, + ) + else: + converted_state = params_state + + del params_state + reclaim_host_memory() # 3. Bind parameters to the Raiden transport. Construct the synchronizer # once, matching the persistent-instance-per-cycle pattern the rebind # optimization depends on. # - # Under Pathways (JAX_PLATFORMS=proxy + JAX_BACKEND_TARGET set, same - # detection tunix's K8sJaxContext.initialize() uses), trainer params + # Under Pathways (JAX_PLATFORMS=proxy + JAX_BACKEND_TARGET set), trainer params # are proxy-backed. Raiden must use FFI (weight_synchronizer_ffi) to bind # directly to device arrays on Pathways TPU workers without host CPU staging, # avoiding client host OOM and multi-minute proxy transfer timeouts. - is_pathways = bool("proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get("JAX_BACKEND_TARGET")) + backend_platform = getattr(jax.devices()[0], "platform", "").lower() if jax.devices() else "" + is_pathways = (backend_platform == "proxy") or ( + "proxy" in os.environ.get("JAX_PLATFORMS", "") and bool(os.environ.get("JAX_BACKEND_TARGET")) + ) if is_pathways and getattr(raiden_synchronizer, "_raiden_ffi", None) is None: raise RuntimeError( "Under Pathways (JAX_PLATFORMS=proxy), Raiden weight synchronization " @@ -1487,11 +1558,13 @@ def prepare_weight_sync( job_name="trainer", worker_index=jax.process_index(), auto_h2d=False, + host_stage=is_pathways, parallelism=4, ) - self._raiden_sync.bind(params_state) - del params_state + self._raiden_sync.bind(converted_state) + del converted_state + reclaim_host_memory() # 4. Initiate Device-to-Host transfer to stage weights for network transfer. if is_pathways or self._raiden_sync.active: @@ -1502,13 +1575,28 @@ def prepare_weight_sync( logging.info("Source weights checksums: %s", self._raiden_sync.checksums()) metadata = self._raiden_sync.work_unit_metadata() + total_variables = len(metadata.variables) + all_metadata = [metadata] + + reclaim_host_memory() + + try: + import resource # pylint: disable=g-import-not-at-top + rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0 + mem_info = f", host memory max RSS: {rss_mb:.1f} MB" + except Exception: # pylint: disable=broad-exception-caught + mem_info = "" + logging.info( - "Trainer prepared weight sync for step %d: registered %d variables on mesh %s", + "Trainer prepared weight sync for step %d: registered %d variables on mesh %s%s", self.train_step, - len(metadata.variables), + total_variables, metadata.mesh_axes, + mem_info, ) - return [metadata] + self._last_staged_step = self.train_step + self._staged_metadata = all_metadata + return all_metadata # Unknown transport: raise rather than return empty metadata. A typo would otherwise # surface only as the coordinator's "empty side" error, with nothing logged anywhere @@ -1517,8 +1605,17 @@ def prepare_weight_sync( def release_weight_sync(self, **kwargs: Any) -> Any: """Releases staged weight buffers after transfer completion.""" + self._last_staged_step = None + self._staged_metadata = None if self._raiden_sync: logging.vlog(1, "Trainer Raiden metrics: %s", self._raiden_sync.metrics()) + reclaim_host_memory() + try: + import resource # pylint: disable=g-import-not-at-top + rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0 + logging.info("Trainer released weight sync: host memory max RSS: %.1f MB", rss_mb) + except Exception: # pylint: disable=broad-exception-caught + pass return True def close(self) -> None: @@ -1527,6 +1624,8 @@ def close(self) -> None: if hasattr(self._raiden_sync, "close"): self._raiden_sync.close() self._raiden_sync = None + self._last_staged_step = None + self._staged_metadata = None self.save_checkpoint(metadata=None, force=True) self._checkpoint_manager.close() diff --git a/tests/post_training/unit/maxtext_engine_e2e_test.py b/tests/post_training/unit/maxtext_engine_e2e_test.py index 1f2bd0e432..e227351a16 100644 --- a/tests/post_training/unit/maxtext_engine_e2e_test.py +++ b/tests/post_training/unit/maxtext_engine_e2e_test.py @@ -16,6 +16,7 @@ from collections.abc import Iterator import dataclasses +import importlib from typing import Any from unittest import mock @@ -33,6 +34,17 @@ import pytest # training_engine imports tunix, so these tests need the post-training dependency bundle. +# The engine's default staging transport is Raiden, whose synchronizer ships with the +# RL tunix build and not with stock tunix. Probe once, the same way the engine does, so +# this loop exercises the real staging path where Raiden exists and the documented +# failure where it does not -- rather than passing or failing on which tunix happens to +# be installed. +try: + importlib.import_module("tunix.experimental.weight_sync.raiden_synchronizer") + _RAIDEN_AVAILABLE = True +except ImportError: + _RAIDEN_AVAILABLE = False + pytestmark = [pytest.mark.post_training] @@ -97,7 +109,13 @@ def run( step_metrics = self.trainer.get_metrics(clear_cache=True) history.append(step_metrics) - _ = self.trainer.prepare_weight_sync() + if _RAIDEN_AVAILABLE: + _ = self.trainer.prepare_weight_sync() + else: + # Without the transport the engine must raise rather than hand back empty + # metadata, which would fail later and far from the cause. + with pytest.raises(RuntimeError, match="raiden_synchronizer"): + self.trainer.prepare_weight_sync() self.trainer.close() return history diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py index 5ee877bbb2..b834c7b459 100644 --- a/tests/post_training/unit/maxtext_engine_test.py +++ b/tests/post_training/unit/maxtext_engine_test.py @@ -16,6 +16,7 @@ # pylint: disable=protected-access import dataclasses +import sys import types from typing import Any from unittest import mock @@ -1331,6 +1332,70 @@ def test_perplexity_is_emitted_alongside_the_loss(self): self.assertIn("perplexity", processed) self.assertAlmostEqual(processed["perplexity"], float(np.exp(6.0)), places=3) + def test_prepare_weight_sync_raises_when_raiden_is_unavailable(self): + """A missing raiden_synchronizer must fail here, not as an empty result downstream. + + Returning empty metadata defers the failure to `WeightSyncCoordinator`, which raises + "metadata collection returned an empty side" -- a count from another process that + never names the missing module. `raiden_synchronizer` ships only on tunix's Raiden + branch, so this is the common case on a released tunix, not a corner. + """ + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + + # Setting the entry to None makes `from ... import raiden_synchronizer` raise + # ImportError, which is what an installed tunix without the module does. + with mock.patch.dict(sys.modules, {"tunix.experimental.weight_sync.raiden_synchronizer": None}): + with self.assertRaisesRegex(RuntimeError, "raiden_synchronizer"): + t.prepare_weight_sync() + + def test_prepare_weight_sync_rejects_an_unknown_transport(self): + """An unrecognised transport must name itself rather than return empty metadata.""" + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + with self.assertRaisesRegex(ValueError, "raidan"): + t.prepare_weight_sync(staging_transport="raidan") + + def _sharded_batch_spec(self, engine, axis="data"): + """A data sharding whose batch dim is actually sharded. + + The single-device test mesh makes `get_input_data_sharding` return a spec with `None` + in the batch position, so the replication branch is unreachable as configured -- an + earlier version of these tests asserted `spec[0] is None` and passed without ever + running the code under test. Stub a spec that shards the batch dim instead. + """ + return jax.sharding.NamedSharding(engine._mesh, jax.sharding.PartitionSpec(axis, None)) # pylint: disable=protected-access + + def test_indivisible_batch_dim_replicates_and_warns_once(self): + """Replicating the batch dim is an N-fold compute cliff, so it must be audible.""" + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + batch = {"a": jnp.zeros((1, 4)), "b": jnp.zeros((1, 4))} + + # Batch dim 1 against a 2-wide axis: indivisible, so the dim must be replicated. + with mock.patch.object(maxtext_engine.sharding, "get_input_data_sharding", return_value=self._sharded_batch_spec(t)): + with mock.patch.object(type(t), "_batch_axis_size", return_value=2): + with self.assertLogs(level="WARNING") as logs: + shardings = t._batch_data_shardings(batch) # pylint: disable=protected-access + t._batch_data_shardings(batch) # pylint: disable=protected-access + + for name, leaf_sharding in shardings.items(): + self.assertIsNone(leaf_sharding.spec[0], f"{name} should have its batch dim replicated") + + # Once per instance, not per leaf and not per call: two leaves over two calls is four + # chances to warn. + warnings = [line for line in logs.output if "does not divide mesh axis" in line] + self.assertLen(warnings, 1) + self.assertIn("2x the work", warnings[0]) + + def test_divisible_batch_dim_stays_sharded_and_is_silent(self): + """The normal case must neither replicate nor warn.""" + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + + with mock.patch.object(maxtext_engine.sharding, "get_input_data_sharding", return_value=self._sharded_batch_spec(t)): + with mock.patch.object(type(t), "_batch_axis_size", return_value=2): + shardings = t._batch_data_shardings({"a": jnp.zeros((4, 4))}) # pylint: disable=protected-access + + self.assertEqual(shardings["a"].spec[0], "data") + self.assertFalse(t._replicated_batch_warned) # pylint: disable=protected-access + if __name__ == "__main__": absltest.main() diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py index 1ecacf3fe0..912d83e900 100644 --- a/tests/post_training/unit/weight_converter_test.py +++ b/tests/post_training/unit/weight_converter_test.py @@ -25,10 +25,13 @@ # Must precede the first JAX import: the cross-mesh tests below need more than # one CPU device, and the backend reads this only at initialization. os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=8") +os.environ.setdefault("JAX_PLATFORMS", "cpu") import types as pytypes # pylint: disable=wrong-import-position import unittest # pylint: disable=wrong-import-position +import logging +from typing import Any import jax # pylint: disable=wrong-import-position import jax.numpy as jnp # pylint: disable=wrong-import-position import numpy as np # pylint: disable=wrong-import-position @@ -439,5 +442,303 @@ def test_padded_moe_fusion_stays_on_the_source_mesh(self): ) +def _profile_conversion_worker(is_streaming: bool, result_queue: Any): + import gc + import resource + import types as pytypes + import jax.numpy as jnp + from maxtext.integration.vllm.weight_converter import MaxTextToMaxTextConverter + + num_layers = 16 + cycle = 2 + scaled_emb = 128 + scaled_experts = 8 + scaled_mlp = 256 + blocks = num_layers // cycle + + cfg = pytypes.SimpleNamespace( + inhomogeneous_layer_cycle_interval=cycle, + num_decoder_layers=num_layers, + param_scan_axis=1, + padded_base_moe_mlp_dim=scaled_mlp, + prefuse_moe_weights=True, + weight_dtype=jnp.float32, + ) + + def _arr(*shape): + return jnp.ones(shape, dtype=jnp.float32) + + layers = {} + for slot in range(cycle): + layers[f"layer_{slot}"] = { + "input_layernorm": {"scale": _arr(scaled_emb, blocks)}, + "post_self_attention_layernorm": {"scale": _arr(scaled_emb, blocks)}, + "self_attention": { + "query": {"kernel": _arr(scaled_emb, blocks, 4, 32)}, + "key": {"kernel": _arr(scaled_emb, blocks, 2, 32)}, + "value": {"kernel": _arr(scaled_emb, blocks, 2, 32)}, + "out": {"kernel": _arr(4, blocks, 32, scaled_emb)}, + }, + "moe_block": { + "gate": {"kernel": _arr(scaled_emb, blocks, scaled_experts)}, + "wi_0": _arr(scaled_experts, blocks, scaled_emb, scaled_mlp), + "wi_1": _arr(scaled_experts, blocks, scaled_emb, scaled_mlp), + "wo": _arr(scaled_experts, blocks, scaled_mlp, scaled_emb), + }, + } + scaled_source = { + "base": { + "token_embedder": {"embedding": _arr(256, scaled_emb)}, + "decoder": {"decoder_norm": {"scale": _arr(scaled_emb)}, "layers": layers}, + } + } + + converter = MaxTextToMaxTextConverter(cfg, prefuse_moe_weights=True) + gc.collect() + before_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + + if is_streaming: + for piece in converter.convert_streaming(scaled_source, target_state=None, groups_per_piece=1): + del piece + gc.collect() + else: + out = converter.convert(scaled_source, target_state=None) + del out + gc.collect() + + after_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + result_queue.put(after_rss - before_rss) + + +class TargetFreeConversionTest(unittest.TestCase): + """Comprehensive test suite for target-free key synthesis and execution.""" + + def test_case_0_raiden_unscan_fails_on_hybrid_cycle(self): + from maxtext.integration.tunix.weight_mapping import raiden_unscan + + source = _source_tree(True) + with self.assertRaises(ValueError) as ctx: + raiden_unscan.unscan_layers(source, num_layers=NUM_LAYERS, scan_axis=SCAN_AXIS) + self.assertIn("expected axis 1 to be num_layers=8", str(ctx.exception)) + + def test_case_1_homogeneous_target_free_unroll(self): + cfg = _config(inhomogeneous_layer_cycle_interval=1, num_decoder_layers=4) + source = { + "base": { + "token_embedder": {"embedding": _arr(16, EMB)}, + "decoder": { + "decoder_norm": {"scale": _arr(EMB)}, + "layers": { + "input_layernorm": {"scale": _arr(EMB, 4)}, + "self_attention": {"query": {"kernel": _arr(EMB, 4, 2, 4)}}, + }, + }, + } + } + converter = WeightConverter(config=cfg, rollout_backend="maxtext") + out = converter.convert(source, target_state=None) + out_root = out["base"] if "base" in out else out + self.assertIn("token_embedder", out_root) + self.assertIn("decoder", out_root) + for i in range(4): + layer_key = f"layers_{i}" + self.assertIn(layer_key, out_root["decoder"]) + scale = getattr(out_root["decoder"][layer_key]["input_layernorm"]["scale"], "value", out_root["decoder"][layer_key]["input_layernorm"]["scale"]) + query = getattr(out_root["decoder"][layer_key]["self_attention"]["query"]["kernel"], "value", out_root["decoder"][layer_key]["self_attention"]["query"]["kernel"]) + self.assertEqual(scale.shape, (EMB,)) + self.assertEqual(query.shape, (EMB, 2, 4)) + + def test_case_2_hybrid_cycle_target_free_unroll(self): + cfg = _config() + source = _source_tree(True) + converter = WeightConverter(config=cfg, rollout_backend="maxtext") + out = converter.convert(source, target_state=None) + out_root = out["base"] if "base" in out else out + src_layers = source["base"]["decoder"]["layers"] + for layer in range(NUM_LAYERS): + slot, block = layer % CYCLE, layer // CYCLE + got = getattr(out_root["decoder"][f"layers_{layer}"]["input_layernorm"]["scale"], "value", out_root["decoder"][f"layers_{layer}"]["input_layernorm"]["scale"]) + want = jnp.take(src_layers[f"layer_{slot}"]["input_layernorm"]["scale"], block, axis=SCAN_AXIS) + np.testing.assert_array_equal(np.asarray(got), np.asarray(want)) + + def test_case_3_prefused_moe_target_free(self): + from maxtext.integration.vllm.moe_padding import compute_padded_moe_mlp_dim + + # Verify helper across topologies + self.assertEqual(compute_padded_moe_mlp_dim(512, 2, 128), 512) + self.assertEqual(compute_padded_moe_mlp_dim(512, 4, 128), 1024) + self.assertEqual(compute_padded_moe_mlp_dim(512, 8, 128), 2048) + + # Verify target-free prefused MoE with padded dim + padded_dim = 16 + cfg = _config(padded_base_moe_mlp_dim=padded_dim, prefuse_moe_weights=True) + source = _source_tree(True) + converter = MaxTextToMaxTextConverter(cfg, prefuse_moe_weights=True) + out = converter.convert(source, target_state=None) + out_root = out["base"] if "base" in out else out + wi = getattr(out_root["decoder"]["layers_0"]["moe_block"]["wi"], "value", out_root["decoder"]["layers_0"]["moe_block"]["wi"]) + wo = getattr(out_root["decoder"]["layers_0"]["moe_block"]["wo"], "value", out_root["decoder"]["layers_0"]["moe_block"]["wo"]) + self.assertEqual(wi.shape, (EXPERTS, EMB, padded_dim * 2)) + self.assertEqual(wo.shape, (EXPERTS, padded_dim, EMB)) + + def test_case_4_abstract_evaluation(self): + cfg = _config(padded_base_moe_mlp_dim=16, prefuse_moe_weights=True) + + def to_struct(x): + arr = getattr(x, "value", x) + return jax.ShapeDtypeStruct(arr.shape, arr.dtype) + + abstract_source = jax.tree_util.tree_map(to_struct, _source_tree(True)) + converter = MaxTextToMaxTextConverter(cfg, prefuse_moe_weights=True) + out = converter.convert(abstract_source, target_state=None) + out_root = out["base"] if "base" in out else out + for leaf in jax.tree_util.tree_leaves(out): + val = getattr(leaf, "value", leaf) + self.assertIsInstance(val, jax.ShapeDtypeStruct) + wi = getattr(out_root["decoder"]["layers_0"]["moe_block"]["wi"], "value", out_root["decoder"]["layers_0"]["moe_block"]["wi"]) + self.assertEqual(wi.shape, (EXPERTS, EMB, 32)) + + def test_case_5_host_memory_profiling(self): + import multiprocessing + ctx = multiprocessing.get_context("spawn") + q_non_stream = ctx.Queue() + p_non_stream = ctx.Process(target=_profile_conversion_worker, args=(False, q_non_stream)) + p_non_stream.start() + delta_non_stream = q_non_stream.get(timeout=60) + p_non_stream.join() + + q_stream = ctx.Queue() + p_stream = ctx.Process(target=_profile_conversion_worker, args=(True, q_stream)) + p_stream.start() + delta_stream = q_stream.get(timeout=60) + p_stream.join() + + logging.info( + "test_case_5_host_memory_profiling: delta_non_stream=%d KB, delta_stream=%d KB", + delta_non_stream, + delta_stream, + ) + self.assertLessEqual(delta_stream, delta_non_stream) + + def test_case_6_parity_vs_raiden_unscan_on_homogeneous(self): + from maxtext.integration.tunix.weight_mapping import raiden_unscan + + cfg = pytypes.SimpleNamespace( + num_decoder_layers=4, + inhomogeneous_layer_cycle_interval=1, + param_scan_axis=1, + weight_dtype=jnp.bfloat16, + prefuse_moe_weights=False, + ) + raw_source = { + "token_embedder": {"embedding": _arr(16, EMB)}, + "decoder": { + "decoder_norm": {"scale": _arr(EMB)}, + "layers": { + "input_layernorm": {"scale": _arr(EMB, 4)}, + "self_attention": {"query": {"kernel": _arr(EMB, 4, 2, 4)}}, + }, + }, + } + bf16_source = jax.tree_util.tree_map( + lambda x: x.astype(jnp.bfloat16) if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating) else x, + raw_source, + ) + baseline_out = raiden_unscan.unscan_layers(bf16_source, num_layers=4, scan_axis=1) + + converter = MaxTextToMaxTextConverter(cfg, prefuse_moe_weights=False, target_dtype=jnp.bfloat16) + converter_out = converter.convert(raw_source, target_state=None) + + base_flat = traverse_util.flatten_dict(baseline_out) + conv_flat = traverse_util.flatten_dict(converter_out) + + self.assertEqual(set(base_flat.keys()), set(conv_flat.keys())) + for k in base_flat: + v_base = getattr(base_flat[k], "value", base_flat[k]) + v_conv = getattr(conv_flat[k], "value", conv_flat[k]) + self.assertEqual(v_base.shape, v_conv.shape, f"Shape mismatch at {k}") + self.assertEqual(v_base.dtype, v_conv.dtype, f"Dtype mismatch at {k}") + np.testing.assert_array_equal(np.asarray(v_base), np.asarray(v_conv), err_msg=f"Value mismatch at {k}") + + def test_case_7_streaming_piece_count_and_parity(self): + cfg = _config( + inhomogeneous_layer_cycle_interval=CYCLE, + num_decoder_layers=NUM_LAYERS, + prefuse_moe_weights=True, + ) + source = _source_tree(True) + converter = WeightConverter(config=cfg, rollout_backend="maxtext") + pieces = list(converter.convert_streaming(source, target_state=None, groups_per_piece=1)) + self.assertEqual(len(pieces), len(converter._direct._groups)) + + # Parity check against fresh non-streaming converter + fresh_converter = WeightConverter(config=cfg, rollout_backend="maxtext") + expected_out = fresh_converter.convert(source, target_state=None) + + merged_flat = {} + for piece in pieces: + piece_flat = traverse_util.flatten_dict(piece) + for k, v in piece_flat.items(): + self.assertNotIn(k, merged_flat, f"Duplicate key across pieces: {k}") + merged_flat[k] = v + + expected_flat = traverse_util.flatten_dict(expected_out) + self.assertEqual(set(merged_flat.keys()), set(expected_flat.keys())) + for k in expected_flat: + v_exp = getattr(expected_flat[k], "value", expected_flat[k]) + v_got = getattr(merged_flat[k], "value", merged_flat[k]) + self.assertEqual(v_exp.shape, v_got.shape, f"Shape mismatch at {k}") + self.assertEqual(v_exp.dtype, v_got.dtype, f"Dtype mismatch at {k}") + np.testing.assert_array_equal(np.asarray(v_exp), np.asarray(v_got), err_msg=f"Value mismatch at {k}") + + def test_case_8_streaming_piece_batching(self): + cfg = _config( + inhomogeneous_layer_cycle_interval=CYCLE, + num_decoder_layers=NUM_LAYERS, + prefuse_moe_weights=True, + ) + source = _source_tree(True) + converter = WeightConverter(config=cfg, rollout_backend="maxtext") + pieces = list(converter.convert_streaming(source, target_state=None, groups_per_piece=2)) + num_groups = len(converter._direct._groups) + expected_piece_count = (num_groups + 1) // 2 + self.assertEqual(len(pieces), expected_piece_count) + + fresh_converter = WeightConverter(config=cfg, rollout_backend="maxtext") + expected_out = fresh_converter.convert(source, target_state=None) + expected_flat = traverse_util.flatten_dict(expected_out) + + merged_flat = {} + for piece in pieces: + piece_flat = traverse_util.flatten_dict(piece) + for k, v in piece_flat.items(): + self.assertNotIn(k, merged_flat, f"Duplicate key across pieces: {k}") + merged_flat[k] = v + + self.assertEqual(set(merged_flat.keys()), set(expected_flat.keys())) + for k in expected_flat: + v_exp = getattr(expected_flat[k], "value", expected_flat[k]) + v_got = getattr(merged_flat[k], "value", merged_flat[k]) + np.testing.assert_array_equal(np.asarray(v_exp), np.asarray(v_got)) + + def test_case_9_weight_converter_convert_streaming_dispatch(self): + cfg = _config( + inhomogeneous_layer_cycle_interval=CYCLE, + num_decoder_layers=NUM_LAYERS, + prefuse_moe_weights=True, + ) + source = _source_tree(True) + # Direct MaxText mode delegates correctly + direct_wc = WeightConverter(config=cfg, rollout_backend="maxtext") + pieces = list(direct_wc.convert_streaming(source, target_state=None)) + self.assertGreater(len(pieces), 0) + + # Torchax rules mode raises NotImplementedError + rule = Rule(source_patterns=["some_pattern"], target_pattern="some_target") + torchax_wc = WeightConverter(rules=[rule], rollout_backend="torchax") + with self.assertRaises(NotImplementedError): + list(torchax_wc.convert_streaming(source)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/prepare_weight_sync_test.py b/tests/unit/prepare_weight_sync_test.py new file mode 100644 index 0000000000..c08a8c7858 --- /dev/null +++ b/tests/unit/prepare_weight_sync_test.py @@ -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. + +"""Unit tests for MaxTextTrainingEngine.prepare_weight_sync single synchronizer logic.""" + +import os +os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=8") +os.environ.setdefault("JAX_PLATFORMS", "cpu") + + +import sys +import types as pytypes +import unittest +from unittest import mock + +# Ensure tunix C-extension / protobuf initializes before transformers/orbax +try: + import tunix.experimental.weight_sync.raiden_synchronizer # pylint: disable=unused-import +except ImportError: + pass + +import jax +import jax.numpy as jnp +from maxtext.training_engine.maxtext_engine import MaxTextTrainingEngine + + +class PrepareWeightSyncTest(unittest.TestCase): + + def setUp(self): + super().setUp() + # Create engine instance without running heavy __init__ + self.engine = MaxTextTrainingEngine.__new__(MaxTextTrainingEngine) + self.engine._raiden_sync = None + self.engine._last_staged_step = None + self.engine._staged_metadata = None + self.engine._train_step = 0 + self.engine._throttler = mock.MagicMock() + self.engine._config = pytypes.SimpleNamespace( + scan_layers=False, + num_decoder_layers=2, + param_scan_axis=1, + weight_sync_debug=False, + ) + self.engine._use_weight_converter = True + self.engine._weight_converter = mock.MagicMock() + self.engine._rollout_backend = "maxtext" + self.engine._get_trainable_params_state = mock.MagicMock(return_value={"layer": jnp.zeros((4, 4))}) + + def _make_dummy_metadata(self, num_vars=2): + meta = mock.MagicMock() + meta.variables = [f"var_{i}" for i in range(num_vars)] + meta.mesh_axes = (1, 1) + return meta + + @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer") + def test_single_synchronizer_creation_and_binding(self, mock_sync_cls): + mock_sync = mock.MagicMock() + mock_sync.active = True + mock_sync.work_unit_metadata.return_value = self._make_dummy_metadata(num_vars=2) + mock_sync.checksums.return_value = {} + mock_sync_cls.return_value = mock_sync + + converted = {"param_0": 0, "param_1": 1} + self.engine._weight_converter.convert.return_value = converted + + metadata = self.engine.prepare_weight_sync() + + self.assertEqual(len(metadata), 1) + self.assertIs(self.engine._raiden_sync, mock_sync) + mock_sync_cls.assert_called_once_with( + job_name="trainer", + worker_index=jax.process_index(), + auto_h2d=False, + host_stage=False, + parallelism=4, + ) + + self.engine._weight_converter.convert.assert_called_once() + mock_sync.bind.assert_called_once_with(converted) + mock_sync.d2h.assert_called_once() + mock_sync.work_unit_metadata.assert_called_once() + + @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer") + def test_rebind_reuses_single_sync_instance(self, mock_sync_cls): + mock_sync = mock.MagicMock() + mock_sync.active = True + mock_sync.work_unit_metadata.return_value = self._make_dummy_metadata(num_vars=2) + mock_sync.checksums.return_value = {} + mock_sync_cls.return_value = mock_sync + + # Round 1 + self.engine._weight_converter.convert.return_value = {"p0": 0} + self.engine.prepare_weight_sync() + self.assertEqual(mock_sync_cls.call_count, 1) + + # Round 2 at step 1 + self.engine._train_step = 1 + self.engine._weight_converter.convert.return_value = {"p0": 0} + self.engine.prepare_weight_sync() + + # Still only 1 synchronizer instance created + self.assertEqual(mock_sync_cls.call_count, 1) + self.assertEqual(mock_sync.bind.call_count, 2) + + def test_release_weight_sync(self): + mock_sync = mock.MagicMock() + self.engine._raiden_sync = mock_sync + self.engine._last_staged_step = 1 + self.engine._staged_metadata = [{"metadata": "dummy"}] + + res = self.engine.release_weight_sync() + + self.assertTrue(res) + self.assertIsNone(self.engine._last_staged_step) + self.assertIsNone(self.engine._staged_metadata) + mock_sync.metrics.assert_called_once() + + def test_release_weight_sync_without_syncs(self): + self.engine._raiden_sync = None + self.engine._last_staged_step = 1 + self.engine._staged_metadata = [{"metadata": "dummy"}] + + res = self.engine.release_weight_sync() + + self.assertTrue(res) + self.assertIsNone(self.engine._last_staged_step) + self.assertIsNone(self.engine._staged_metadata) + + def test_close(self): + mock_sync = mock.MagicMock() + self.engine._raiden_sync = mock_sync + self.engine._last_staged_step = 1 + self.engine._staged_metadata = [{"metadata": "dummy"}] + self.engine.save_checkpoint = mock.MagicMock() + self.engine._checkpoint_manager = mock.MagicMock() + self.engine._throttler = mock.MagicMock() + self.engine._metrics_recorder = mock.MagicMock() + + self.engine.close() + + mock_sync.close.assert_called_once() + self.assertIsNone(self.engine._raiden_sync) + self.assertIsNone(self.engine._last_staged_step) + self.assertIsNone(self.engine._staged_metadata) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/raiden_unscan_test.py b/tests/unit/raiden_unscan_test.py new file mode 100644 index 0000000000..d65d54db9c --- /dev/null +++ b/tests/unit/raiden_unscan_test.py @@ -0,0 +1,238 @@ +# 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 `raiden_unscan.unscan_layers`. + +This transform decides the *names* Raiden binds. The trainer runs scanned +(`scan_layers=True`); the sampler loads its MaxText model unscanned, and Raiden matches +tensors by `jax.tree_util.keystr` path. Nothing downstream cross-checks the two name sets +-- `raiden_handler._validate_metadata` only checks a single manifest's internal +consistency (mesh rank, duplicate variable/layer keys, sharding specs) -- so a naming +error here surfaces as weights that silently never transfer, not as an exception. + +The transform is pure pytree manipulation, so all of this runs on CPU in milliseconds. +""" + +from absl.testing import absltest +from flax import nnx +import jax +import jax.numpy as jnp +from maxtext.integration.tunix.weight_mapping import raiden_unscan +import numpy as np +import pytest + +# This transform exists only for the Tunix/Raiden weight-sync path, so it is graded with +# the rest of that work. `tests/unit` is in cpu-post-training-unit's path list, so the +# marker alone routes it there -- no file move needed (unlike tests/ or tests/integration, +# which are not in that list and would be collected by no job at all). +pytestmark = [pytest.mark.post_training] + + +_NUM_LAYERS = 3 +_IN, _OUT, _VOCAB = 4, 8, 10 + + +def _unwrap(leaf): + """Reads a leaf's array whether or not it is wrapped in an `nnx.Param`.""" + if isinstance(leaf, nnx.Variable): + return leaf[...] + return leaf + + +def _names(tree) -> list[str]: + """The names Raiden binds: exactly what `raiden_synchronizer.flatten_weights` computes.""" + return sorted(jax.tree_util.keystr(p) for p, _ in jax.tree_util.tree_leaves_with_path(tree)) + + +class ScannedInner(nnx.Module): + """One scanned param: the layer axis lives at axis 1 of a single array.""" + + def __init__(self, num_layers: int = _NUM_LAYERS): + self.kernel = nnx.Param(jnp.arange(_IN * num_layers * _OUT, dtype=jnp.float32).reshape(_IN, num_layers, _OUT)) + self.scale = nnx.Param(jnp.arange(_OUT * num_layers, dtype=jnp.float32).reshape(_OUT, num_layers)) + + +class ScannedModel(nnx.Module): + """A trainer-side model: scanned `layers`, plus non-layer params that must pass through.""" + + def __init__(self, num_layers: int = _NUM_LAYERS): + self.layers = ScannedInner(num_layers) + self.embed = nnx.Param(jnp.zeros((_VOCAB, _IN))) + + +class UnscannedInner(nnx.Module): + + def __init__(self): + self.kernel = nnx.Param(jnp.zeros((_IN, _OUT))) + self.scale = nnx.Param(jnp.zeros((_OUT,))) + + +class UnscannedModel(nnx.Module): + """A sampler-side model: one submodule per layer, named `layers_0..N-1`.""" + + def __init__(self, num_layers: int = _NUM_LAYERS): + for i in range(num_layers): + setattr(self, f"layers_{i}", UnscannedInner()) + self.embed = nnx.Param(jnp.zeros((_VOCAB, _IN))) + + +class UnscanLayersTest(absltest.TestCase): + + def _scanned_state(self, num_layers: int = _NUM_LAYERS): + return nnx.state(ScannedModel(num_layers), nnx.Param) + + def test_names_match_an_unscanned_model_exactly(self): + """The point of the transform: trainer names must equal sampler names. + + Raiden binds by `keystr` path on both sides, and nothing validates that the two sets + agree, so this is the assertion that a silent no-transfer would violate. + """ + unscanned = raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS) + sampler_side = nnx.state(UnscannedModel(), nnx.Param) + self.assertEqual(_names(unscanned), _names(sampler_side)) + + def test_plain_dict_and_nnx_state_produce_identical_names(self): + """`unscan_layers` returns a plain nested dict, the sampler binds an `nnx.State`. + + `keystr` renders both identically only because the transform rewraps leaves in + `nnx.Param`. Dropping that rewrap would rename every tensor (`['k']` vs `['k'].value`) + and break every transfer, so pin it. + """ + unscanned = raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS) + self.assertIsInstance(jax.tree_util.tree_leaves(unscanned, is_leaf=lambda x: isinstance(x, nnx.Param))[0], nnx.Param) + self.assertTrue(all(n.endswith(".value") for n in _names(unscanned)), _names(unscanned)) + + def test_slices_carry_the_right_values(self): + """Layer i must receive index i of the scan axis -- not a transpose or an off-by-one.""" + state = self._scanned_state() + original = np.asarray(state.to_pure_dict()["layers"]["kernel"]) + unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS) + + for i in range(_NUM_LAYERS): + got = unscanned[f"layers_{i}"]["kernel"] + got = np.asarray(_unwrap(got)) + self.assertEqual(got.shape, (_IN, _OUT)) + np.testing.assert_array_equal(got, original[:, i, :]) + + def test_rank_two_param_is_also_unscanned(self): + """A rank-2 scanned param (e.g. a norm scale) slices down to rank 1.""" + unscanned = raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS) + for i in range(_NUM_LAYERS): + scale = unscanned[f"layers_{i}"]["scale"] + self.assertEqual(np.asarray(_unwrap(scale)).shape, (_OUT,)) + + def test_non_layer_entries_pass_through_unchanged(self): + """Embeddings and final norms have no layer axis and must survive untouched.""" + state = self._scanned_state() + embed_before = np.asarray(state.to_pure_dict()["embed"]) + unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS) + + self.assertIn("embed", unscanned) + embed_after = unscanned["embed"] + np.testing.assert_array_equal(np.asarray(_unwrap(embed_after)), embed_before) + self.assertNotIn("layers", unscanned) + + def test_layer_count_mismatch_raises(self): + """A wrong num_layers must fail loudly rather than bind truncated weights.""" + with self.assertRaisesRegex(ValueError, "expected axis 1 to be num_layers=99"): + raiden_unscan.unscan_layers(self._scanned_state(), num_layers=99) + + def test_already_unscanned_state_raises(self): + """The anti-silent-no-op guard. + + Without it an already-unscanned (or wrongly-keyed) state would return unchanged and + bind under scanned names, transferring nothing with no error anywhere. + """ + with self.assertRaisesRegex(ValueError, "found no scanned 'layers' entries"): + raiden_unscan.unscan_layers(nnx.state(UnscannedModel(), nnx.Param), num_layers=_NUM_LAYERS) + + def test_wrong_layer_container_raises(self): + with self.assertRaisesRegex(ValueError, "found no scanned 'blocks' entries"): + raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS, layer_container="blocks") + + def test_custom_scan_axis(self): + """`param_scan_axis` is configurable; axis 0 must slice the leading dim.""" + state = {"layers": {"kernel": jnp.arange(_NUM_LAYERS * _OUT, dtype=jnp.float32).reshape(_NUM_LAYERS, _OUT)}} + unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS, scan_axis=0) + for i in range(_NUM_LAYERS): + got = unscanned[f"layers_{i}"]["kernel"] + np.testing.assert_array_equal(np.asarray(_unwrap(got)), np.arange(i * _OUT, (i + 1) * _OUT)) + + def test_plain_dict_input_is_accepted(self): + """The trainer passes an `nnx.State`, but the signature documents plain dicts too.""" + state = { + "layers": {"kernel": jnp.zeros((_IN, _NUM_LAYERS, _OUT))}, + "embed": jnp.zeros((_VOCAB, _IN)), + } + unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS) + self.assertEqual(sorted(unscanned.keys()), ["embed"] + [f"layers_{i}" for i in range(_NUM_LAYERS)]) + + def test_total_element_count_is_preserved(self): + """Unscanning reshapes; it must not drop or duplicate any weight.""" + state = self._scanned_state() + before = sum(int(np.size(x)) for x in jax.tree_util.tree_leaves(state.to_pure_dict())) + unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS) + after = sum(int(np.size(np.asarray(_unwrap(x)))) for x in jax.tree_util.tree_leaves(unscanned)) + self.assertEqual(after, before) + + def test_streaming_piece_count_and_parity(self): + state = self._scanned_state() + pieces = list(raiden_unscan.unscan_layers_streaming(state, num_layers=_NUM_LAYERS, keys_per_piece=1)) + # 3 original flattened keys: embed, layers.kernel, layers.scale + self.assertEqual(len(pieces), 3) + + expected = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS) + merged_flat = {} + for piece in pieces: + piece_flat = raiden_unscan.flatten_dict(piece) + for k, v in piece_flat.items(): + self.assertNotIn(k, merged_flat) + merged_flat[k] = v + + expected_flat = raiden_unscan.flatten_dict(expected) + self.assertEqual(set(merged_flat.keys()), set(expected_flat.keys())) + for k in expected_flat: + v_exp = _unwrap(expected_flat[k]) + v_got = _unwrap(merged_flat[k]) + np.testing.assert_array_equal(np.asarray(v_got), np.asarray(v_exp)) + + def test_streaming_keys_per_piece_batching(self): + state = self._scanned_state() + pieces = list(raiden_unscan.unscan_layers_streaming(state, num_layers=_NUM_LAYERS, keys_per_piece=2)) + # 3 keys with batch size 2 -> ceil(3/2) = 2 pieces + self.assertEqual(len(pieces), 2) + + expected = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS) + expected_flat = raiden_unscan.flatten_dict(expected) + merged_flat = {} + for piece in pieces: + merged_flat.update(raiden_unscan.flatten_dict(piece)) + + self.assertEqual(set(merged_flat.keys()), set(expected_flat.keys())) + for k in expected_flat: + np.testing.assert_array_equal(np.asarray(_unwrap(merged_flat[k])), np.asarray(_unwrap(expected_flat[k]))) + + def test_streaming_leaves_are_nnx_param(self): + state = self._scanned_state() + for piece in raiden_unscan.unscan_layers_streaming(state, num_layers=_NUM_LAYERS): + for leaf in jax.tree_util.tree_leaves(piece, is_leaf=lambda x: isinstance(x, nnx.Param)): + self.assertIsInstance(leaf, nnx.Param) + + def test_streaming_already_unscanned_state_raises(self): + with self.assertRaisesRegex(ValueError, "found no scanned 'layers' entries"): + list(raiden_unscan.unscan_layers_streaming(nnx.state(UnscannedModel(), nnx.Param), num_layers=_NUM_LAYERS)) + + +if __name__ == "__main__": + absltest.main()