From d043c5b8e75197f591f8597ebcf988899de3ace9 Mon Sep 17 00:00:00 2001 From: Yixuan Wang Date: Thu, 10 Sep 2026 12:33:04 -0700 Subject: [PATCH] feat(training_engine): integrate Raiden weight sync and clean up rollout path Combines Raiden FFI weight synchronization integration in MaxTextTrainingEngine and rollout path cleanup. Training Engine & Weight Sync: - Integrate Raiden weight sync in MaxTextTrainingEngine with FFI support and persistent synchronizer lifecycle. - Access config properties directly without getattr. - Generalize WeightConverter and MaxTextToMaxTextConverter for arbitrary kv_tp_size and moe_mlp_tp_size. - Support target-free KV head replication along axis -2 when kv_tp_size > base_num_kv_heads. - Safely default MoE lane_size to 0 for TPU GMM_v2 plain per-shard concatenation. - Add helper utilities is_verify_weights_enabled and resolve_prefuse_moe_weights in convert_utils. - Add unit tests in prepare_weight_sync_test and weight_converter_test. Rollout Path Cleanup: - Remap vLLM hybrid KV cache groups using layer_name_to_kvcache_index before passing to decoder layers. - Strip auxiliary runner kwargs before forwarding to self.model. - Write updated KV caches back to original physical cache slots. - Add rollout_tensor_parallelism field to VLLM config. - Remove redundant rollout weight sync logic and tunix_compat_context in favor of Raiden FFI path. --- .../extra_deps/post_train_github_deps.txt | 2 +- src/maxtext/common/gcloud_stub.py | 1 + src/maxtext/configs/base.yml | 16 + src/maxtext/configs/types.py | 24 +- src/maxtext/integration/vllm/convert_utils.py | 49 ++- .../vllm/maxtext_vllm_adapter/adapter.py | 40 ++- .../integration/vllm/maxtext_vllm_rollout.py | 327 +++++------------- .../integration/vllm/weight_converter.py | 62 +++- .../distillation/scripts/run_distill_xpk.sh | 2 +- .../trainers/post_train/rl/train_rl.py | 77 +---- src/maxtext/training_engine/maxtext_engine.py | 164 ++++++--- .../post_training/unit/convert_utils_test.py | 37 ++ .../unit/maxtext_engine_e2e_test.py | 20 +- .../post_training/unit/maxtext_engine_test.py | 55 ++- .../unit/prepare_weight_sync_test.py | 169 +++++++++ .../unit/qwen35_standalone_converter_test.py | 117 +------ .../unit/router_replay_engine_test.py | 3 +- .../unit/weight_converter_test.py | 51 +++ 18 files changed, 717 insertions(+), 499 deletions(-) create mode 100644 tests/post_training/unit/prepare_weight_sync_test.py diff --git a/src/dependencies/extra_deps/post_train_github_deps.txt b/src/dependencies/extra_deps/post_train_github_deps.txt index b67a0d791c..ddb92f30d7 100644 --- a/src/dependencies/extra_deps/post_train_github_deps.txt +++ b/src/dependencies/extra_deps/post_train_github_deps.txt @@ -1,3 +1,3 @@ -google-tunix @ https://github.com/google/tunix/archive/1b0e3c5e89058d4dddf0ec68ae8be06c127f68ac.zip +google-tunix @ https://github.com/google/tunix/archive/a8d70582f1e2f1fb65973210989e0e148b5ef7ad.zip tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/b67ae5f8f234fd559cf4b376840e7bb9ae6d3275.zip vllm @ git+https://github.com/vllm-project/vllm@d626108b1841888ec90aced33367149a6bbc7e4b 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/base.yml b/src/maxtext/configs/base.yml index 016f31a322..19dd62a6c3 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1374,6 +1374,22 @@ vllm_hf_config_path: "" vllm_hf_overrides: {} # JSON string containing additional configuration for the vLLM model (e.g. '{"maxtext_config": {...}}') vllm_additional_config: {} +# Whether to use the streaming target-free weight converter for trainer-side conversion +use_weight_converter: true +# Tensor parallelism per replica for rollout (-1 for auto-determined) +rollout_tensor_parallelism: -1 +# Rollout backend for trainer-side weight converter ('maxtext' or 'vllm_torchax') +rollout_backend: "maxtext" +# Degree of tensor parallelism for KV cache / attention heads in rollout +kv_tp_size: 1 +# Degree of tensor parallelism for MoE MLP dimension in rollout +moe_mlp_tp_size: 1 +# Enable debug verification and checksum logging during weight sync +weight_sync_debug: false +# Whether to use Raiden FFI transport for weight sync (null derives from platform) +use_raiden_ffi: null +# Weight load format for vLLM in converter validation +vllm_load_format: "dummy" # When use_jax_splash=true, force the layout of the query tensor to be [..., NUM_HEADS, HEAD_DIM, SEQ_LENGTH] force_q_layout: false diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 84acba8936..da1103b6a2 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2741,12 +2741,24 @@ class VLLM(BaseModel): vllm_hf_config_path: str = Field("", description="Path to HuggingFace model config for MaxText model.") use_standalone_converter: bool = Field(False, description="Use the standalone MaxText->torchax vLLM converter") use_weight_converter: bool = Field( - False, + True, description=( "Use an explicit weight converter for trainer->rollout weight sync instead of " "the legacy transfer_state_directly / transfer_state_with_mappings paths." ), ) + use_raiden_ffi: Optional[bool] = Field( + None, + description="Use Raiden FFI transport for weight sync.", + ) + rollout_tensor_parallelism: int = Field( + -1, + description="Tensor parallelism per replica for rollout. If not specified, it will be auto-determined.", + ) + 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=( @@ -2766,6 +2778,16 @@ class VLLM(BaseModel): "rather than data movement. Adds one barrier per sync." ), ) + kv_tp_size: int = Field( + 1, + ge=1, + description="Degree of tensor parallelism for KV cache / attention heads in rollout.", + ) + moe_mlp_tp_size: int = Field( + 1, + ge=1, + description="Degree of tensor parallelism for MoE MLP dimension in rollout.", + ) vllm_load_format: str = Field( "dummy", description="Weight load format for vLLM in converter validation. Options:'auto', 'dummy'.", diff --git a/src/maxtext/integration/vllm/convert_utils.py b/src/maxtext/integration/vllm/convert_utils.py index 1909c6f26b..71e6a7fa2f 100644 --- a/src/maxtext/integration/vllm/convert_utils.py +++ b/src/maxtext/integration/vllm/convert_utils.py @@ -41,8 +41,9 @@ `_bulk_align_and_unstack`; that patch goes away once the port lands. """ -from typing import Mapping, Any, Callable, Dict, Tuple, Optional import functools +import os +from typing import Mapping, Any, Callable, Dict, Tuple, Optional from absl import logging import jax import jax.numpy as jnp @@ -506,7 +507,7 @@ def _interleave_moe_weights( tgt_shape: Tuple[int, ...], n_shards: int, axis: Optional[int] = None, - lane_size: int = DEFAULT_TPU_NUM_LANES, + lane_size: Optional[int] = 0, ) -> jax.Array | np.ndarray: """Interleaves wi_0 and wi_1 per-shard into a single tensor matching TPU GMM layout. @@ -517,9 +518,18 @@ def _interleave_moe_weights( that is the difference between ~3x and ~1x the output size in live transient memory, plus ~5 fewer dispatches per call. + For TPU GMM_v2 kernels (e.g. gmm_v2.py), each TP shard expects plain concatenation + [local_gate, local_up]. The kernel internally performs 128-lane interleaving into + VMEM via interleave_lane(w_gate, w_up). Therefore, lane_size defaults to 0 (plain + concatenation per shard). Pre-interleaving in HBM double-interleaves and corrupts + weights. + `tgt_shape`, `n_shards`, `axis` and `lane_size` are static, so the trace is keyed on them; identical layers share a single compilation. """ + if lane_size is None: + lane_size = 0 + if axis is None: axis = len(tgt_shape) - 1 elif axis < 0: @@ -559,7 +569,7 @@ def _pad_and_chunk(arr): p_wi_1 = p_wi_1.reshape(shape_lanes) combined = jnp.stack([p_wi_0, p_wi_1], axis=axis + 2) else: - # Fallback when dimension is not divisible by lane_size: + # Concatenate wi_0 (gate) and wi_1 (up) per shard: [local_gate, local_up]. combined = jnp.concatenate([p_wi_0, p_wi_1], axis=axis + 1) return combined.reshape(tgt_shape) @@ -692,16 +702,45 @@ def _scanned_sharding_from_per_layer( def resolve_rollout_tp(config: Any, tp: int = 1) -> int: - """Resolves rollout TP from config.""" + """Resolves rollout TP from override, config, or environment.""" if tp > 1: return int(tp) config_tp = 0 if config is not None: - config_tp = int( + raw = ( getattr(config, "rollout_tensor_parallelism", 0) or getattr(getattr(config, "cluster", None), "rollout_tensor_parallelism", 0) or getattr(config, "rollout_mesh_tp", 0) or 0 ) + if raw and int(raw) > 0: + config_tp = int(raw) + + if not config_tp: + env_tp = os.environ.get("ROLLOUT_TENSOR_PARALLEL_SIZE") or os.environ.get("ROLLOUT_MESH_TP") + if env_tp and int(env_tp) > 0: + config_tp = int(env_tp) + return int(config_tp or 1) + + +def resolve_prefuse_moe_weights(config: Any, prefuse_moe_weights: Optional[bool] = None) -> bool: + """Resolves MoE prefuse flag from override, config, or environment.""" + if prefuse_moe_weights is not None: + return bool(prefuse_moe_weights) + if "ROLLOUT_PREFUSE_MOE_WEIGHTS" in os.environ: + return os.environ["ROLLOUT_PREFUSE_MOE_WEIGHTS"].lower() in ("1", "true", "yes") + if config is not None and getattr(config, "rollout_prefuse_moe_weights", None) is not None: + return bool(config.rollout_prefuse_moe_weights) + rollout_backend = getattr(config, "rollout_backend", None) or os.environ.get("ROLLOUT_BACKEND", "maxtext") + if rollout_backend == "maxtext": + return True + if config is not None and getattr(config, "prefuse_moe_weights", None) is not None: + return bool(config.prefuse_moe_weights) + return os.environ.get("PREFUSE_MOE_WEIGHTS", "0").lower() in ("1", "true", "yes") + + +def is_verify_weights_enabled() -> bool: + """Returns whether weight verification / checksum validation is active.""" + return os.environ.get("VERIFY_WEIGHTS", "").lower() == "true" diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py index 6e5454d936..b0a61ae8d7 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.convert_utils import DEFAULT_TPU_NUM_LANES, compute_padded_moe_mlp_dim from maxtext.integration.vllm.hybrid_cache_utils import ( build_qwen_gdn_cache_layout, gather_layer_kv_caches, @@ -124,7 +125,10 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters else vllm_config.model_config.hf_config ) hidden_size = getattr(hf_config, "moe_intermediate_size", None) - num_lanes = pltpu.get_tpu_info().num_lanes + try: + num_lanes = pltpu.get_tpu_info().num_lanes + except Exception: # pylint: disable=broad-exception-caught + num_lanes = DEFAULT_TPU_NUM_LANES num_kv_heads = hf_config.num_key_value_heads # Number of KV heads in global attention layers (None if the field is absent or unset). @@ -172,10 +176,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 @@ -337,6 +339,21 @@ def __call__( positions = _input_positions input_positions = normalize_vllm_input_positions(positions) + # Filter kwargs to only those accepted by self.model. + model_kwargs = dict(kwargs) + for extra_key in ( + "inputs_embeds", + "input_positions", + "layer_name_to_kvcache_index", + "_layer_name_to_kv_cache", + "shared_attention_metadata", + "intermediate_tensors", + "lora_metadata", + "is_first_rank", + "is_last_rank", + ): + model_kwargs.pop(extra_key, None) + with self.mesh, nn.logical_axis_rules(self.maxtext_config.logical_axis_rules): aux_hidden_states = [] expert_indices = None @@ -347,7 +364,7 @@ def __call__( kv_caches=layer_kv_caches, attention_metadata=attention_metadata, model_mode=self.model_mode, - **kwargs, + **model_kwargs, ) if isinstance(res, tuple) and len(res) == 3: @@ -471,6 +488,7 @@ def load_weights(self, rng_key: jax.Array) -> None: if self.maxtext_config.lora.lora_restore_path: lora_utils.restore_lora_from_path(model, self.maxtext_config) self.model = nnx.data(model) + patch_raiden_worker_h2d() def get_mrope_input_positions( self, @@ -614,3 +632,13 @@ def patched_get_kv_cache_spec(self): KVCacheManager.get_kv_cache_spec = patched_get_kv_cache_spec max_logging.log("Successfully applied KVCacheManager patch for hybrid GDN models.") + + +def patch_raiden_worker_h2d(): + """Monkey-patches TPUWorker.raiden_h2d and RaidenWorkerSync to apply Raiden weights to runner.""" + try: + from tunix.experimental.weight_sync.raiden_synchronizer import patch_raiden_worker_sync # pylint: disable=import-outside-toplevel + + patch_raiden_worker_sync() + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.log(f"Skipping raiden worker sync patch: {e}") diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py index 1cd7853f96..033f41b826 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py @@ -32,26 +32,21 @@ import traceback from typing import Any, Optional, Tuple -import jax -import jax.numpy as jnp from flax import nnx from flax.traverse_util import flatten_dict, unflatten_dict - -from tunix.generate import mappings -from tunix.generate import utils as tunix_gen_utils -from tunix.generate.vllm_sampler import VllmConfig, VllmSampler -from tunix.rl import reshard as tunix_reshard -from tunix.rl.rollout import base_rollout, vllm_rollout - +import jax +import jax.numpy as jnp from maxtext.integration.vllm.convert_utils import _sharding_summary -from maxtext.integration.vllm.weight_converter import ( - WeightConverter, - MODEL_TO_CONVERSION_RULES, -) -from maxtext.integration.vllm.torchax_converter.base import BaseMaxTextToVLLMConverter from maxtext.integration.vllm.torchax_converter.gemma4_moe import Gemma4MaxTextToVLLMConverter from maxtext.integration.vllm.torchax_converter.qwen35_moe import Qwen35MaxTextToVLLMConverter from maxtext.integration.vllm.torchax_converter.qwen3_moe import Qwen3MaxTextToVLLMConverter +from maxtext.integration.vllm.weight_converter import ( + MODEL_TO_CONVERSION_RULES, + WeightConverter, +) +from tunix.generate import mappings +from tunix.generate.vllm_sampler import VllmConfig, VllmSampler +from tunix.rl.rollout import base_rollout, vllm_rollout # Sentinel distinguishing "this model has no entry" from "this model has an # entry whose value is None", which means direct-sync-only. @@ -208,6 +203,69 @@ def _find_scanned_layer_idx(key_tuple, container_names=("layers", "scanned_block return -1, None +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 _find_qwen_scanned_layer_idx(key_tuple): """Finds a Qwen scanned block path like `layers.layer_0` or `layers.moe_block`.""" for i in range(len(key_tuple) - 1): @@ -298,69 +356,6 @@ def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Opt 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. @@ -441,7 +436,8 @@ def unroll_gemma_scanned_weights(weights): else: new_flat_w[k] = v - assert unrolled_count > 0, "MaxTextVllmSampler: Detected scanned structure, but failed to unroll any layers!" + if unrolled_count <= 0: + raise ValueError("MaxTextVllmSampler: Detected scanned structure, but failed to unroll any layers!") logging.info( "MaxTextVllmSampler: Successfully unrolled %d scanned tensor components into vLLM-compatible nnx.List format.", @@ -450,31 +446,37 @@ def unroll_gemma_scanned_weights(weights): return unflatten_dict(new_flat_w) +def _log_and_flush_traceback(msg: str) -> None: + """Logs an error with formatted traceback and flushes all logging handlers.""" + logging.error("%s:\n%s", msg, traceback.format_exc()) + for handler in logging.getLogger().handlers: + try: + handler.flush() + except Exception: # pylint: disable=broad-except + pass + + class MaxTextVllmSampler(VllmSampler): - """VllmSampler that hands MaxText weights to a converter before the sync. + """VllmSampler that applies scanned-weight pre-unrolls for direct MaxText sync. The weight-sync implementation itself lives in `VllmSampler.update_params` (Tunix), which owns the KV-cache teardown/rebuild and the `state_leaves` refresh that vLLM needs to actually observe new weights. This subclass only - supplies the converter and applies scanned-weight pre-unrolls for the - legacy / direct-sync paths. + applies scanned-weight pre-unrolls for the legacy / direct-sync paths. """ 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, + model_name: Optional[str] = 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 + engine_kwargs = getattr(config, "engine_kwargs", {}) or {} + self._model_name = model_name or getattr(config, "model", "") or engine_kwargs.get("model", "") or "" + self._is_gemma = "gemma" in str(self._model_name).lower() def update_params( self, @@ -482,25 +484,8 @@ def update_params( filter_types: Optional[Tuple[Any, ...]] = None, ): """Update the vLLM runner weights from a MaxText state tree.""" - if isinstance(self._converter, BaseMaxTextToVLLMConverter): - try: - return self._sync_standalone_converted(updated_weights) - except BaseException: - logging.error("MaxTextVllmSampler standalone sync failed:\n%s", traceback.format_exc()) - for handler in logging.getLogger().handlers: - try: - handler.flush() - except Exception: # pylint: disable=broad-except - 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, - ) - updated_weights = unroll_gemma_scanned_weights(updated_weights) + 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) except BaseException: @@ -508,100 +493,9 @@ def update_params( # down, and the teardown races the normal exception propagation -- the # Python traceback is routinely truncated or lost entirely in the worker # logs. Force it out before re-raising. - logging.error("MaxTextVllmSampler.update_params failed:\n%s", traceback.format_exc()) - for handler in logging.getLogger().handlers: - try: - handler.flush() - except Exception: # pylint: disable=broad-except - pass + _log_and_flush_traceback("MaxTextVllmSampler.update_params failed") raise - def _sync_standalone_converted(self, updated_weights): - """Standalone torchax-converter sync path. - - The converter emits tensors in the tpu-inference runner's *internal* layout, - keyed by its state names, so this bypasses Tunix's mapped/direct transfers: - tear down the KV cache, convert, reshard each tensor onto its existing - sharding (chunked, Pathways-aware) and assign into the runner's flat state - dict in place. - """ - runner = self._model_runner - state = runner.state - if not isinstance(state, dict): - raise TypeError( - "Standalone torchax converters target the vLLM (torchax) model " - "implementation, whose runner state is a flat dict; got " - f"{type(state).__name__}. Remove MaxTextForCausalLM overrides so " - "vLLM runs its native model." - ) - - if self.llm is not None: - self.llm.reset_prefix_cache() - self.llm.collective_rpc("delete_kv_cache") - elif self._driver is not None: - self._driver.llm_engine.reset_prefix_cache() - self._driver.llm_engine.collective_rpc("delete_kv_cache") - jax.effects_barrier() - - start = time.time() - pure = updated_weights.to_pure_dict() if hasattr(updated_weights, "to_pure_dict") else updated_weights - converted = self._converter.convert(pure) - - src = {k: v for k, v in converted.items() if k in state} - version_aliases = sorted(set(converted) - set(src)) - if version_aliases: - logging.info( - "Standalone sync: %d converted tensors have no runner target (vLLM version aliases), e.g. %s", - len(version_aliases), - version_aliases[:3], - ) - uncovered = [k for k in state if k not in src and not k.rsplit(".", 1)[-1].startswith("_") and "rotary_emb" not in k] - if uncovered: - logging.warning( - "Standalone sync: %d runner tensors NOT covered by the converter (stale weights!), e.g. %s", - len(uncovered), - uncovered[:5], - ) - - spec = {k: state[k] for k in src} - expected = {k: (tuple(v.shape), v.dtype) for k, v in spec.items()} - chunk = getattr(self.config, "reshard_chunk_size", None) - delete_dst = getattr(self.config, "delete_dst_buffers", True) - reshard_in_chunks = getattr(tunix_gen_utils, "_reshard_in_chunks", None) - if chunk and reshard_in_chunks is None: - logging.warning("Standalone sync: this Tunix has no _reshard_in_chunks; falling back to one reshard call.") - chunk = None - if chunk: - resharded = reshard_in_chunks( - src_flat=dict(src), - spec_flat=spec, - reshard_fn=tunix_reshard.reshard_pytree, - chunk_size=chunk, - delete_spec_buffers=delete_dst, - ) - else: - shardings = {k: v.sharding for k, v in spec.items()} - if delete_dst: - tunix_gen_utils._delete_target_buffers(spec, src) # pylint: disable=protected-access - resharded = tunix_reshard.reshard_pytree(src, shardings) - - for k in src: - new = resharded[k] - shape, dtype = expected[k] - if tuple(new.shape) != shape or new.dtype != dtype: - raise ValueError( - f"{k}: converter produced {tuple(new.shape)}/{new.dtype}, the runner expects {shape}/{dtype}; " - "the converter's layout is out of date with tpu-inference." - ) - state[k] = new - runner.state_leaves = state - logging.info("Standalone sync: updated %d/%d runner tensors in %.1fs", len(src), len(state), time.time() - start) - - if self.llm is not None: - self.llm.collective_rpc("reinitialize_kv_cache") - elif self._driver is not None: - self._driver.llm_engine.collective_rpc("reinitialize_kv_cache") - class MaxTextVllmRollout(vllm_rollout.VllmRollout): """VllmRollout that uses MaxTextVllmSampler for weight synchronization. @@ -654,40 +548,11 @@ def __init__( # fact indirectly and got it wrong when either field was reformatted. use_hf = "maxtext_config" not in vllm_additional_config and not uses_maxtext_vllm_adapter(maxtext_config) direct_maxtext_sync = not use_hf - use_weight_converter = bool( - getattr(maxtext_config, "use_weight_converter", False) - or vllm_additional_config.get("use_weight_converter", False) - ) - use_standalone_converter = bool( - getattr(maxtext_config, "use_standalone_converter", False) - or vllm_additional_config.get("use_standalone_converter", False) - ) - # Sampler sharding the standalone converter must mirror: attention DP from - # the sharding_strategy blob, expert parallelism from the vLLM engine kwargs. - strategy = {} - sharding_blob = vllm_additional_config.get("sharding") if isinstance(vllm_additional_config, dict) else None - if isinstance(sharding_blob, dict): - strategy = sharding_blob.get("sharding_strategy") or {} - rollout_vllm_kwargs = getattr(rollout_config, "rollout_vllm_kwargs", None) or {} - sharding_hints = { - "attn_dp_size": (int(strategy.get("attn_dp_size") or 1) if strategy.get("enable_dp_attention", False) else 1), - "enable_expert_parallel": bool(rollout_vllm_kwargs.get("enable_expert_parallel", False)), - } - # Accepted from either spelling, matching use_weight_converter above, so a + # Accepted from either spelling, matching MaxTextEngine, so a # debug run can be triggered by editing the same JSON blob. self._weight_sync_debug = bool( getattr(maxtext_config, "weight_sync_debug", False) or vllm_additional_config.get("weight_sync_debug", False) ) - converter = _create_model_converter( - maxtext_config.model_name, - config=maxtext_config, - mesh=mesh, - use_hf_mapping=use_hf, - use_weight_converter=use_weight_converter, - use_standalone_converter=use_standalone_converter, - sharding_hints=sharding_hints, - debug=self._weight_sync_debug, - ) mapping_config = mappings.MappingConfig.build( mapping_obj=rollout_config.rollout_mapping_config, @@ -742,10 +607,8 @@ 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), + model_name=getattr(maxtext_config, "model_name", ""), ) # Counts every weight sync, including the initial one below. See diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py index b927b7173a..26efcc6e5e 100644 --- a/src/maxtext/integration/vllm/weight_converter.py +++ b/src/maxtext/integration/vllm/weight_converter.py @@ -32,6 +32,7 @@ _device_ids, _fuse_and_unstack_moe, _get_n_shards, + _jit_repeat_axes, _jit_unstack, _scanned_sharding_from_per_layer, _sharding_summary, @@ -257,6 +258,8 @@ def __init__( self, rules: Optional[List[Rule]] = None, tp: int = 1, + kv_tp_size: int = 1, + moe_mlp_tp_size: int = 1, num_kv_heads: Optional[int] = None, head_dim: Optional[int] = None, config: Any = None, @@ -279,6 +282,8 @@ def __init__( ) self.rules = rules self.tp = resolve_rollout_tp(config, tp) + self.kv_tp_size = kv_tp_size or getattr(config, "kv_tp_size", 1) or self.tp + self.moe_mlp_tp_size = moe_mlp_tp_size or getattr(config, "moe_mlp_tp_size", 1) or self.tp # Read by the rollout engine to decide whether to trace the reshard # step that runs after conversion. @@ -296,6 +301,8 @@ def __init__( self._direct = MaxTextToMaxTextConverter( config=config, tp=self.tp, + kv_tp_size=self.kv_tp_size, + moe_mlp_tp_size=self.moe_mlp_tp_size, moe_fused_layout=(moe_fused_layout or MoEFusedLayout.PER_SHARD_INTERLEAVE), allow_unused_source_keys=allow_unused_source_keys, debug=debug, @@ -678,9 +685,13 @@ def __init__( debug: bool = False, prefuse_moe_weights: Optional[bool] = None, target_dtype: Optional[Any] = None, + kv_tp_size: int = 1, + moe_mlp_tp_size: int = 1, ): self.config = config self.tp = resolve_rollout_tp(config, tp) + self.kv_tp_size = kv_tp_size or getattr(config, "kv_tp_size", 1) or self.tp + self.moe_mlp_tp_size = moe_mlp_tp_size or getattr(config, "moe_mlp_tp_size", 1) or self.tp self.moe_fused_layout = moe_fused_layout self.allow_unused_source_keys = allow_unused_source_keys self.debug = debug @@ -691,6 +702,15 @@ def __init__( 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) + self.base_num_kv_heads = getattr(config, "base_num_kv_heads", None) or getattr(config, "num_kv_heads", None) + self.kv_replication = 1 + if self.base_num_kv_heads is not None and self.kv_tp_size > self.base_num_kv_heads: + if self.kv_tp_size % self.base_num_kv_heads != 0: + raise ValueError( + f"kv_tp_size ({self.kv_tp_size}) must be divisible by base_num_kv_heads ({self.base_num_kv_heads})." + ) + self.kv_replication = self.kv_tp_size // self.base_num_kv_heads + self.cycle = int(getattr(config, "inhomogeneous_layer_cycle_interval", 1) or 1) self.num_decoder_layers = int(config.num_decoder_layers) if self.num_decoder_layers % self.cycle: @@ -707,7 +727,8 @@ def __init__( logging.info( "MaxTextToMaxTextConverter: %d layers, cycle=%d, %d scanned blocks, " - "scan_axis=%d, moe_fused_layout=%s, prefuse_moe=%s, padded_moe_dim=%s", + "scan_axis=%d, moe_fused_layout=%s, prefuse_moe=%s, padded_moe_dim=%s, " + "kv_tp_size=%d, moe_mlp_tp_size=%d, kv_replication=%d", self.num_decoder_layers, self.cycle, self.num_blocks, @@ -715,6 +736,9 @@ def __init__( self.moe_fused_layout, self.prefuse_moe_weights, self.padded_base_moe_mlp_dim, + self.kv_tp_size, + self.moe_mlp_tp_size, + self.kv_replication, ) def _resolve_target_dtype(self): @@ -906,9 +930,11 @@ def _fuse_moe_bulk(self, wi_0, wi_1, tgt_val, key_path: str): 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(tgt_val, tgt_fused_axis) - if n_shards == 1 and self.tp > 1: - n_shards = self.tp + n_shards = ( + self.moe_mlp_tp_size + if self.moe_mlp_tp_size > 1 + else (self.tp if self.tp > 1 else _get_n_shards(tgt_val, tgt_fused_axis)) + ) return _fuse_and_unstack_moe( wi_0, wi_1, @@ -933,6 +959,7 @@ def _fuse_moe_bulk(self, wi_0, wi_1, tgt_val, key_path: str): def _slice_bulk_target_free(self, val: Any, path: str): """Returns target-free slices for a scanned parameter.""" last_key = path.split(".")[-1] + is_kv = "key.kernel" in path or "value.kernel" in path 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: @@ -942,6 +969,12 @@ def _slice_bulk_target_free(self, val: Any, path: str): 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 + if ( + is_kv + and self.kv_replication > 1 + and (self.base_num_kv_heads is None or unrolled_shape[-2] == self.base_num_kv_heads) + ): + unrolled_shape[-2] = unrolled_shape[-2] * self.kv_replication 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: @@ -960,6 +993,9 @@ def _slice_bulk_target_free(self, val: Any, path: str): pad_spec[intermediate_axis] = (0, pad_amount) val = jnp.pad(val, pad_spec) + if is_kv and self.kv_replication > 1 and (self.base_num_kv_heads is None or val.shape[-2] == self.base_num_kv_heads): + val = _jit_repeat_axes(val, ((-2, self.kv_replication),)) + return _jit_unstack(val, self.scan_axis) def _fuse_moe_bulk_target_free(self, wi_0: Any, wi_1: Any, path: str): @@ -979,7 +1015,11 @@ def _fuse_moe_bulk_target_free(self, wi_0: Any, wi_1: Any, path: str): 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 = self.tp if self.tp > 1 else _get_n_shards(wi_0, scan_fused_axis) + n_shards = ( + self.moe_mlp_tp_size + if self.moe_mlp_tp_size > 1 + else (self.tp if self.tp > 1 else _get_n_shards(wi_0, scan_fused_axis)) + ) return _fuse_and_unstack_moe( wi_0, wi_1, @@ -1010,6 +1050,18 @@ def _execute_group_target_free(self, group: _PlanGroup, src_flat): raw_val = src_flat[group.source_keys[0]] tgt_dt = getattr(raw_val, "dtype", target_dtype) if ("gate" in path or "router" in path) else target_dtype val = _apply_dtype_cast(raw_val, tgt_dt, path) + is_kv = "key.kernel" in path or "value.kernel" in path + if ( + is_kv + and self.kv_replication > 1 + and (self.base_num_kv_heads is None or val.shape[-2] == self.base_num_kv_heads) + ): + if isinstance(val, jax.ShapeDtypeStruct): + new_shape = list(val.shape) + new_shape[-2] = new_shape[-2] * self.kv_replication + val = jax.ShapeDtypeStruct(tuple(new_shape), val.dtype) + else: + val = _jit_repeat_axes(val, ((-2, self.kv_replication),)) return [(tgt_key, val) for _, tgt_key in group.targets] if group.op == "fuse_moe": diff --git a/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh b/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh index 8a6678b34f..1bb3681999 100644 --- a/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh +++ b/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh @@ -164,7 +164,7 @@ require_env() { : "${DISTILL_LAYER_INDICES:=[0,1,2,3,4,5,6,7]}" # Image pinning (used by prep_image). -: "${TUNIX_SOURCE:=git+https://github.com/google/tunix@1b0e3c5e89058d4dddf0ec68ae8be06c127f68ac}" +: "${TUNIX_SOURCE:=git+https://github.com/google/tunix@a8d70582f1e2f1fb65973210989e0e148b5ef7ad}" : "${JAX_PIN:=0.10.0}" : "${JAXLIB_PIN:=0.10.0}" : "${LIBTPU_PIN:=0.0.39}" diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index e6e923ebc5..3bea8289cc 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/src/maxtext/trainers/post_train/rl/train_rl.py @@ -44,7 +44,7 @@ """ from __future__ import annotations -import contextlib +import functools from functools import wraps from typing import Any, Callable, Optional, Sequence @@ -52,7 +52,6 @@ import datasets import grain import jax -import jax.numpy as jnp import json import logging import os @@ -68,79 +67,10 @@ import maxtext.integration.vllm.maxtext_vllm_adapter as adapter adapter.register() -import functools from tunix.rl import rl_cluster as rl_cluster_lib from tunix.rl.rollout import base_rollout from tunix.rl.grpo.grpo_learner import GrpoConfig, GrpoLearner from tunix.sft import metrics_logger, profiler -import tunix.generate.utils as tunix_utils - - -@contextlib.contextmanager -def _tpu_inference_compat_patches(): - """Tactical compat shims for tpu_inference. - - tpu_inference has two call-site assumptions that no longer hold: - 1. jax.lax.with_sharding_constraint: assumes silent reshard on mismatch, - but current jax asserts when all mesh axes are Explicit. Fall back to - jax.sharding.reshard on the AssertionError. - 2. tunix._apply_dtype_cast: tpu_inference JaxEinsum defaults - param_dtype=float32 so its weights initialize as float32, but model - dtype is bfloat16; the cast upgraded synced bfloat16 weights to float32, - which then mismatched in the ragged paged attention kernel. Skip the - bf16->f32 upcast so synced weights stay bfloat16. - - Scoped to rl_train() so the patches don't leak into other importers of this - module. Drop both once tpu_inference is updated upstream. - """ - orig_wsc = jax.lax.with_sharding_constraint - orig_apply_dtype_cast = tunix_utils._apply_dtype_cast # pylint: disable=protected-access - orig_bulk = tunix_utils._bulk_align_and_unstack # pylint: disable=protected-access - orig_unstack = tunix_utils._unstack_scanned_param # pylint: disable=protected-access - - orig_moe_weights = getattr(tunix_utils, "_MOE_MLP_WEIGHTS", None) - - def _compat_wsc(x, shardings): - try: - return orig_wsc(x, shardings) - except AssertionError: - return jax.sharding.reshard(x, shardings) - - def _no_bf16_to_f32_cast(val, tgt_dtype, src_key): - if hasattr(val, "dtype") and val.dtype == jnp.bfloat16 and tgt_dtype == jnp.float32: - return val - return orig_apply_dtype_cast(val, tgt_dtype, src_key) - - def _compat_bulk(arr, scan_axis, per_layer, key_path): - if hasattr(arr, "shape") and len(arr.shape) <= scan_axis: - scan_axis = len(arr.shape) - 1 if len(arr.shape) > 0 else 0 - return orig_bulk(arr, scan_axis, per_layer, key_path) - - def _compat_unstack(src_val, tgt_val, key_path, scan_axis=None): - if scan_axis is not None and hasattr(src_val, "shape") and len(src_val.shape) <= scan_axis: - scan_axis = len(src_val.shape) - 1 if len(src_val.shape) > 0 else 0 - res = orig_unstack(src_val, tgt_val, key_path, scan_axis=scan_axis) - if isinstance(res, tuple) and len(res) == 1 and hasattr(src_val, "shape") and src_val.shape == tgt_val.shape: - return res * 256 - return res - - jax.lax.with_sharding_constraint = _compat_wsc - tunix_utils._apply_dtype_cast = _no_bf16_to_f32_cast # pylint: disable=protected-access - tunix_utils._bulk_align_and_unstack = _compat_bulk # pylint: disable=protected-access - tunix_utils._unstack_scanned_param = _compat_unstack # pylint: disable=protected-access - - if orig_moe_weights is not None: - tunix_utils._MOE_MLP_WEIGHTS = frozenset([*orig_moe_weights, "wo"]) # pylint: disable=protected-access - - try: - yield - finally: - jax.lax.with_sharding_constraint = orig_wsc - tunix_utils._apply_dtype_cast = orig_apply_dtype_cast # pylint: disable=protected-access - tunix_utils._bulk_align_and_unstack = orig_bulk # pylint: disable=protected-access - tunix_utils._unstack_scanned_param = orig_unstack # pylint: disable=protected-access - if orig_moe_weights is not None: - tunix_utils._MOE_MLP_WEIGHTS = orig_moe_weights # pylint: disable=protected-access os.environ["TOKENIZERS_PARALLELISM"] = "0" @@ -649,12 +579,11 @@ def rl_train(argv: Sequence[str], kwargs: dict): trainer_devices: JAX devices for the trainer. sampler_devices: JAX devices for the sampler. """ - with _tpu_inference_compat_patches(): - _rl_train_impl(argv, kwargs) + _rl_train_impl(argv, kwargs) def _rl_train_impl(argv: Sequence[str], kwargs: dict): - """rl_train body — kept separate so _tpu_inference_compat_patches wraps it cleanly.""" + """rl_train execution body.""" trainer_config, sampler_config, trainer_devices, sampler_devices = model_creation_utils.setup_configs_and_devices( argv, kwargs, diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index d572fcacff..648273bdc5 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -23,8 +23,7 @@ from collections.abc import Callable, Mapping import contextlib import dataclasses -import os -from typing import Any +from typing import Any, Optional from absl import logging from flax import nnx @@ -37,6 +36,11 @@ 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 ( + is_verify_weights_enabled, + resolve_prefuse_moe_weights, + resolve_rollout_tp, +) from maxtext.trainers.pre_train import train as maxtext_train from maxtext.training_engine import abstract_engine from maxtext.training_engine import checkpointing @@ -146,6 +150,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" @@ -557,6 +569,7 @@ def __init__( self._compiled_eval: Any = None self._compiled_eval_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") self._model = self._build_model(wrap_with_tunix_adapter, tokenizer_pad_id) @@ -606,6 +619,28 @@ def __init__( self._metrics_logger = metrics_module.MetricsLogger(config=self._config) self._throttler = inflight_throttler.InflightThrottler(config=self._config, metrics_logger=self._metrics_logger) self._raiden_sync: Any = None + self._last_staged_step: Optional[int] = None + self._staged_metadata: Any = None + self._use_weight_converter = bool(self._config.use_weight_converter) + self._rollout_backend = self._config.rollout_backend + rollout_tp = resolve_rollout_tp(self._config) + kv_tp = self._config.kv_tp_size or rollout_tp + moe_tp = self._config.moe_mlp_tp_size or rollout_tp + prefuse_moe = resolve_prefuse_moe_weights(self._config) + 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, + tp=rollout_tp, + kv_tp_size=kv_tp, + moe_mlp_tp_size=moe_tp, + prefuse_moe_weights=prefuse_moe, + rollout_backend=self._rollout_backend, + debug=self._config.weight_sync_debug, + ) + else: + self._weight_converter = None def _build_model(self, wrap_with_tunix_adapter: bool, tokenizer_pad_id: int | None) -> Any: """Returns the model to train, adopting a mesh when this engine was given none.""" @@ -1203,7 +1238,11 @@ def _prepare_batch(self, payload: Any) -> Any: if self._gen_model_input_fn is not None: return self._gen_model_input_fn(payload) if dataclasses.is_dataclass(payload): - return {k: getattr(payload, k) for k in payload.__dataclass_fields__ if getattr(payload, k) is not None} + return { + k: getattr(payload, k) + for k in payload.__dataclass_fields__ + if getattr(payload, k) is not None and k != "metadata" + } return payload def _mesh_sharding(self, leaf: Any) -> jax.sharding.Sharding | None: @@ -1247,7 +1286,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) @@ -1257,8 +1297,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) @@ -1774,6 +1822,10 @@ def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None: metadata: Checkpoint metadata payload from Orchestrator. **kwargs: Additional checkpoint saving options. """ + if not self._config.enable_checkpointing or not self._checkpoint_dir(): + logging.info("Checkpointing is disabled in config; skipping save_checkpoint.") + return + # Drain all inflight computations and log pending metrics before checkpointing. self._throttler.wait_for_all() @@ -2061,54 +2113,50 @@ 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 reusing 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() # 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: + 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, + cycle_interval=self._config.inhomogeneous_layer_cycle_interval, + ) + else: + converted_state = params_state + + del params_state # 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 - # 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")) - if is_pathways and getattr(raiden_synchronizer, "_raiden_ffi", None) is None: - raise RuntimeError( - "Under Pathways (JAX_PLATFORMS=proxy), Raiden weight synchronization " - "requires weight_synchronizer_ffi (from tpu_raiden_jax) to avoid client host OOM " - "and proxy staging timeouts. However, _raiden_ffi is not available in " - "tunix.experimental.weight_sync.raiden_synchronizer. Please ensure a " - "compatible tpu_raiden_jax wheel with FFI support is installed." - ) - if self._raiden_sync is None: self._raiden_sync = raiden_synchronizer.RaidenSynchronizer( job_name="trainer", @@ -2117,25 +2165,30 @@ def prepare_weight_sync( parallelism=4, ) - self._raiden_sync.bind(params_state) - del params_state + self._raiden_sync.bind(converted_state) + del converted_state # 4. Initiate Device-to-Host transfer to stage weights for network transfer. - if is_pathways or self._raiden_sync.active: + if self._raiden_sync.active: self._raiden_sync.d2h() - verify_weights = os.environ.get("VERIFY_WEIGHTS", "").lower() == "true" + verify_weights = is_verify_weights_enabled() if verify_weights: logging.info("Source weights checksums: %s", self._raiden_sync.checksums()) - metadata = self._raiden_sync.work_unit_metadata() + all_metadata = self._raiden_sync.work_unit_metadata_all() + total_variables = sum(len(m.variables) for m in all_metadata) + logging.info( - "Trainer prepared weight sync for step %d: registered %d variables on mesh %s", + "Trainer prepared weight sync for step %d: registered %d work unit(s) with %d variables on mesh %s", self.train_step, - len(metadata.variables), - metadata.mesh_axes, + len(all_metadata), + total_variables, + all_metadata[0].mesh_axes if all_metadata else (), ) - 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 @@ -2144,6 +2197,8 @@ 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()) return True @@ -2154,9 +2209,12 @@ 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() + if self._config.enable_checkpointing and self._checkpoint_dir() and self._checkpoint_manager: + self.save_checkpoint(metadata=None, force=True) + self._checkpoint_manager.close() # Write the metrics and cleanup metrics logger resources self._throttler.cleanup() diff --git a/tests/post_training/unit/convert_utils_test.py b/tests/post_training/unit/convert_utils_test.py index d751e6f2a7..69f059f439 100644 --- a/tests/post_training/unit/convert_utils_test.py +++ b/tests/post_training/unit/convert_utils_test.py @@ -14,6 +14,7 @@ """Tests for MaxText convert_utils functions.""" +import os import unittest import jax.numpy as jnp @@ -24,7 +25,10 @@ _interleave_moe_weights, compute_padded_moe_mlp_dim, DEFAULT_TPU_NUM_LANES, + is_verify_weights_enabled, pad_to_tpu_lanes, + resolve_prefuse_moe_weights, + resolve_rollout_tp, ) pytestmark = [pytest.mark.post_training] @@ -160,6 +164,39 @@ def test_compute_padded_moe_mlp_dim_leaves_aligned_dims_alone(self): self.assertEqual(compute_padded_moe_mlp_dim(2560, 2, lanes), 2560) self.assertIsNone(compute_padded_moe_mlp_dim(None, 4, lanes)) + def test_is_verify_weights_enabled(self): + orig = os.environ.get("VERIFY_WEIGHTS") + try: + os.environ["VERIFY_WEIGHTS"] = "true" + self.assertTrue(is_verify_weights_enabled()) + os.environ["VERIFY_WEIGHTS"] = "false" + self.assertFalse(is_verify_weights_enabled()) + os.environ.pop("VERIFY_WEIGHTS", None) + self.assertFalse(is_verify_weights_enabled()) + finally: + if orig is not None: + os.environ["VERIFY_WEIGHTS"] = orig + else: + os.environ.pop("VERIFY_WEIGHTS", None) + + def test_resolve_prefuse_moe_weights(self): + self.assertTrue(resolve_prefuse_moe_weights(None, prefuse_moe_weights=True)) + self.assertFalse(resolve_prefuse_moe_weights(None, prefuse_moe_weights=False)) + + class DummyConfig: + prefuse_moe_weights = False + rollout_backend = "vllm" + + self.assertFalse(resolve_prefuse_moe_weights(DummyConfig())) + + def test_resolve_rollout_tp(self): + self.assertEqual(resolve_rollout_tp(None, tp=4), 4) + + class DummyConfig: + rollout_tensor_parallelism = 2 + + self.assertEqual(resolve_rollout_tp(DummyConfig()), 2) + if __name__ == "__main__": unittest.main() 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 60d937218a..a683191089 100644 --- a/tests/post_training/unit/maxtext_engine_test.py +++ b/tests/post_training/unit/maxtext_engine_test.py @@ -1014,10 +1014,11 @@ def test_shared_types_are_tunix_classes(self): # What actually arrives at fwd_bwd from GRPOAdapter.create_trainer_payloads. rl_payload = datatypes.RLTrainerPayload( - token_ids=jnp.zeros((1, 4)), - token_mask=jnp.ones((1, 4)), + prompt_ids=jnp.zeros((1, 4)), + prompt_mask=jnp.ones((1, 4)), + completion_ids=jnp.zeros((1, 4)), + completion_mask=jnp.ones((1, 4)), advantages=jnp.zeros((1,)), - loss_mask=jnp.ones((1, 4)), ) self.assertIsInstance(rl_payload, abstract_engine.TrainerPayload) @@ -1349,6 +1350,54 @@ 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_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/prepare_weight_sync_test.py b/tests/post_training/unit/prepare_weight_sync_test.py new file mode 100644 index 0000000000..e38aba8d87 --- /dev/null +++ b/tests/post_training/unit/prepare_weight_sync_test.py @@ -0,0 +1,169 @@ +# 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.""" + +# pylint: disable=protected-access + +import os + +os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=8") +os.environ.setdefault("JAX_PLATFORMS", "cpu") + +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 +import pytest + +from maxtext.configs import pyconfig +from maxtext.training_engine.maxtext_engine import MaxTextTrainingEngine +from tests.utils.test_helpers import get_test_config_path + +pytestmark = [pytest.mark.post_training] + + +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 = pyconfig.initialize( + [None, get_test_config_path()], + run_name="test_prepare_weight_sync", + scan_layers=False, + num_decoder_layers=2, + param_scan_axis=1, + inhomogeneous_layer_cycle_interval=1, + weight_sync_debug=False, + enable_checkpointing=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_all.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, + 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_all.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_all.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._metrics_logger = 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/post_training/unit/qwen35_standalone_converter_test.py b/tests/post_training/unit/qwen35_standalone_converter_test.py index 919679792f..861210eb08 100644 --- a/tests/post_training/unit/qwen35_standalone_converter_test.py +++ b/tests/post_training/unit/qwen35_standalone_converter_test.py @@ -23,18 +23,12 @@ from types import SimpleNamespace import unittest -from unittest import mock import numpy as np import jax.numpy as jnp import pytest -from maxtext.integration.vllm import maxtext_vllm_rollout as rollout_mod -from maxtext.integration.vllm.maxtext_vllm_rollout import ( - MaxTextVllmSampler, - _create_model_converter, -) -from maxtext.integration.vllm.torchax_converter.base import BaseMaxTextToVLLMConverter +from maxtext.integration.vllm.maxtext_vllm_rollout import _create_model_converter from maxtext.integration.vllm.torchax_converter.gemma4_moe import Gemma4MaxTextToVLLMConverter from maxtext.integration.vllm.torchax_converter.qwen35_moe import Qwen35MaxTextToVLLMConverter @@ -386,114 +380,5 @@ def test_standalone_unknown_model_rejected(self): _create_model_converter("llama3.1-8b", config=_make_config(tp=1), mesh=None, use_standalone_converter=True) -class _DummyConverter(BaseMaxTextToVLLMConverter): - """Minimal standalone converter emitting a canned dict.""" - - def __init__(self, out): - super().__init__(_make_config(tp=1), mesh=None) - self._out = out - - def convert(self, model_state, **kwargs): - return dict(self._out) - - def _convert_global(self, params): - pass - - def _convert_attn(self, params): - pass - - def _convert_moe(self, params): - pass - - -class _FakeRunnerSampler(MaxTextVllmSampler): - """Shadows the base class's read-only `_model_runner` property so tests can set it.""" - - _model_runner = None - - -def _make_sampler(state, converter, chunk=None, delete_dst=False, llm=None): - sampler = object.__new__(_FakeRunnerSampler) - sampler._converter = converter # pylint: disable=protected-access - sampler.converter = converter - sampler._model_runner = SimpleNamespace(state=state, state_leaves=None) # pylint: disable=protected-access - sampler.llm = llm - sampler._driver = None # pylint: disable=protected-access - sampler.config = SimpleNamespace(reshard_chunk_size=chunk, delete_dst_buffers=delete_dst) - return sampler - - -def _identity_reshard(src, shardings): - del shardings - return src - - -class StandaloneSyncTest(unittest.TestCase): - - def test_update_params_routes_standalone_converters(self): - sampler = _make_sampler({}, _DummyConverter({})) - sampler._sync_standalone_converted = mock.Mock(return_value="synced") # pylint: disable=protected-access - self.assertEqual(sampler.update_params({"w": 1}), "synced") - sampler._sync_standalone_converted.assert_called_once_with({"w": 1}) # pylint: disable=protected-access - - def test_update_params_reraises_sync_failures(self): - sampler = _make_sampler({}, _DummyConverter({})) - sampler._sync_standalone_converted = mock.Mock(side_effect=RuntimeError("boom")) # pylint: disable=protected-access - with self.assertRaises(RuntimeError): - sampler.update_params({}) - - def test_sync_requires_flat_dict_state(self): - sampler = _make_sampler(object(), _DummyConverter({})) - with self.assertRaisesRegex(TypeError, "flat dict"): - sampler._sync_standalone_converted({}) # pylint: disable=protected-access - - def test_sync_updates_covered_tensors_in_place(self): - state = { - "layers.0.qkv": jnp.zeros((2, 3), jnp.bfloat16), - "layers.0._private": jnp.zeros((1,), jnp.bfloat16), - "layers.0.rotary_emb.cache": jnp.zeros((1,), jnp.bfloat16), - "layers.0.uncovered": jnp.zeros((1,), jnp.bfloat16), - } - new_qkv = jnp.ones((2, 3), jnp.bfloat16) - converter = _DummyConverter({"layers.0.qkv": new_qkv, "layers.0.alias_only": jnp.ones((4,), jnp.bfloat16)}) - llm = mock.Mock() - sampler = _make_sampler(state, converter, llm=llm) - with mock.patch.object(rollout_mod.tunix_reshard, "reshard_pytree", _identity_reshard): - sampler._sync_standalone_converted({}) # pylint: disable=protected-access - self.assertTrue(jnp.array_equal(state["layers.0.qkv"], new_qkv)) - self.assertNotIn("layers.0.alias_only", state) # version alias without a runner target is dropped - self.assertTrue(not state["layers.0.uncovered"].any()) # left untouched (and warned about) - self.assertIs(sampler._model_runner.state_leaves, state) # pylint: disable=protected-access - llm.reset_prefix_cache.assert_called_once_with() - llm.collective_rpc.assert_has_calls([mock.call("delete_kv_cache"), mock.call("reinitialize_kv_cache")]) - - def test_sync_uses_chunked_reshard_when_configured(self): - state = {"w": jnp.zeros((2,), jnp.bfloat16)} - converter = _DummyConverter({"w": jnp.ones((2,), jnp.bfloat16)}) - sampler = _make_sampler(state, converter, chunk=2, delete_dst=True) - chunked = mock.Mock(side_effect=lambda src_flat, **kw: src_flat) - with mock.patch.object(rollout_mod.tunix_gen_utils, "_reshard_in_chunks", chunked, create=True): - sampler._sync_standalone_converted({}) # pylint: disable=protected-access - self.assertEqual(chunked.call_args.kwargs["chunk_size"], 2) - self.assertTrue(state["w"].all()) - - def test_sync_falls_back_when_tunix_lacks_chunked_reshard(self): - state = {"w": jnp.zeros((2,), jnp.bfloat16)} - converter = _DummyConverter({"w": jnp.ones((2,), jnp.bfloat16)}) - sampler = _make_sampler(state, converter, chunk=2, delete_dst=False) - with mock.patch.object(rollout_mod.tunix_gen_utils, "_reshard_in_chunks", None, create=True): - with mock.patch.object(rollout_mod.tunix_reshard, "reshard_pytree", _identity_reshard): - sampler._sync_standalone_converted({}) # pylint: disable=protected-access - self.assertTrue(state["w"].all()) - - def test_sync_rejects_layout_drift(self): - state = {"w": jnp.zeros((2, 2), jnp.bfloat16)} - converter = _DummyConverter({"w": jnp.ones((3, 3), jnp.bfloat16)}) - sampler = _make_sampler(state, converter) - with mock.patch.object(rollout_mod.tunix_reshard, "reshard_pytree", _identity_reshard): - with self.assertRaisesRegex(ValueError, "out of date"): - sampler._sync_standalone_converted({}) # pylint: disable=protected-access - - if __name__ == "__main__": unittest.main() diff --git a/tests/post_training/unit/router_replay_engine_test.py b/tests/post_training/unit/router_replay_engine_test.py index 169e3cf740..9211cf6d61 100644 --- a/tests/post_training/unit/router_replay_engine_test.py +++ b/tests/post_training/unit/router_replay_engine_test.py @@ -73,6 +73,7 @@ def _tiny_qwen35_kwargs(seq_len, batch_size, num_experts, top_k, **overrides): "max_prefill_predict_length": seq_len, "per_device_batch_size": float(batch_size), "weight_dtype": "bfloat16", + "inhomogeneous_layer_cycle_interval": 1, } kwargs.update(overrides) return kwargs @@ -135,7 +136,7 @@ def test_engine_accepts_router_replay_payload_and_loss_is_finite(self): seq_len, batch_size, top_k = 16, 2, 2 cfg = _init_test_cfg( - extra_args=["attention=flash"], + extra_args=["attention=dot_product"], **_tiny_qwen35_kwargs( seq_len, batch_size, diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py index 6c44344fe7..59c2e2a51d 100644 --- a/tests/post_training/unit/weight_converter_test.py +++ b/tests/post_training/unit/weight_converter_test.py @@ -733,6 +733,57 @@ def test_case_8_weight_converter_convert_streaming_dispatch(self): with self.assertRaises(NotImplementedError): list(torchax_wc.convert_streaming(source)) + def test_case_9_target_free_kv_head_replication(self): + cfg = _config( + base_num_kv_heads=2, + inhomogeneous_layer_cycle_interval=1, + num_decoder_layers=2, + ) + key_val = _scanned(EMB, 2, 4) + value_val = _scanned(EMB, 2, 4) + source = { + "base": { + "decoder": { + "layers": { + "self_attention": { + "key": {"kernel": key_val}, + "value": {"kernel": value_val}, + } + } + } + } + } + + # 1. Successful replication: kv_tp_size=4, base_num_kv_heads=2 -> kv_replication=2 + converter = WeightConverter( + config=cfg, + kv_tp_size=4, + rollout_backend="maxtext", + ) + self.assertEqual(converter._direct.kv_replication, 2) # pylint: disable=protected-access + out = converter.convert(source, target_state=None) + out_root = out["base"] if "base" in out else out + for layer_idx in range(2): + layer_key = f"layers_{layer_idx}" + key_kernel = getattr( + out_root["decoder"][layer_key]["self_attention"]["key"]["kernel"], + "value", + out_root["decoder"][layer_key]["self_attention"]["key"]["kernel"], + ) + self.assertEqual(key_kernel.shape, (EMB, 4, 4)) + # Heads 0 and 1 are repeated along axis -2 + np.testing.assert_array_equal(key_kernel[:, 0, :], key_kernel[:, 1, :]) + np.testing.assert_array_equal(key_kernel[:, 2, :], key_kernel[:, 3, :]) + + # 2. Divisibility failure: kv_tp_size=3, base_num_kv_heads=2 + with self.assertRaises(ValueError) as ctx: + WeightConverter( + config=cfg, + kv_tp_size=3, + rollout_backend="maxtext", + ) + self.assertIn("must be divisible by base_num_kv_heads", str(ctx.exception)) + if __name__ == "__main__": unittest.main()