From 7754d49042e624aa2e57b93912a175b5df896ba8 Mon Sep 17 00:00:00 2001 From: A9isha Date: Thu, 10 Sep 2026 01:03:45 +0000 Subject: [PATCH] [Raiden Weight Sync 4/7] Target-free and streaming weight conversion for vLLM --- src/maxtext/integration/vllm/convert_utils.py | 16 + .../integration/vllm/weight_converter.py | 365 ++++++++++++++---- .../unit/weight_converter_test.py | 295 ++++++++++++++ 3 files changed, 609 insertions(+), 67 deletions(-) diff --git a/src/maxtext/integration/vllm/convert_utils.py b/src/maxtext/integration/vllm/convert_utils.py index 893e4e2af9..1909c6f26b 100644 --- a/src/maxtext/integration/vllm/convert_utils.py +++ b/src/maxtext/integration/vllm/convert_utils.py @@ -689,3 +689,19 @@ def _scanned_sharding_from_per_layer( jax.sharding.PartitionSpec(*spec), memory_kind=per_layer_sharding.memory_kind, ) + + +def resolve_rollout_tp(config: Any, tp: int = 1) -> int: + """Resolves rollout TP from config.""" + if tp > 1: + return int(tp) + + config_tp = 0 + if config is not None: + config_tp = int( + 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 + ) + return int(config_tp or 1) diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py index c8e79ff535..b927b7173a 100644 --- a/src/maxtext/integration/vllm/weight_converter.py +++ b/src/maxtext/integration/vllm/weight_converter.py @@ -16,14 +16,16 @@ import abc import dataclasses +import gc import logging import re +from typing import Any, Dict, Iterator, List, Mapping, Optional, Tuple, Union + +from flax import nnx, traverse_util import jax import jax.numpy as jnp -import gc -from typing import List, Union, Any, Dict, Optional, Mapping, Tuple -from flax import traverse_util, nnx from maxtext.integration.vllm.convert_utils import ( + MOE_MLP_WEIGHTS, _align_per_axis, _apply_dtype_cast, _bulk_align_and_unstack, @@ -33,6 +35,9 @@ _jit_unstack, _scanned_sharding_from_per_layer, _sharding_summary, + normalize_dtype, + pad_to_tpu_lanes, + resolve_rollout_tp, ) @@ -47,20 +52,6 @@ def __call__(self, tensors: List[Any], **kwargs) -> Any: pass -class Concatenate(Operation): - """Concatenates input tensors along a given dimension.""" - - def __init__(self, dim: int): - self.dim = dim - - def __call__(self, tensors, **kwargs): - @jax.jit - def _f(*ts): - return jnp.concatenate(ts, axis=self.dim) - - return _f(*tensors) - - class Transpose(Operation): """Transposes the input tensor along specified axes.""" @@ -164,7 +155,7 @@ def _fuse_single(w0, w1): w1 = jnp.transpose(w1, (0, 2, 1)) num_experts, d_inner, d_model = w0.shape chunk_size = d_inner // tp - padded_chunk_size = ((chunk_size + 127) // 128) * 128 + padded_chunk_size = pad_to_tpu_lanes(chunk_size) pad_amount = padded_chunk_size - chunk_size gate_chunks = w0.reshape(num_experts, tp, chunk_size, d_model) up_chunks = w1.reshape(num_experts, tp, chunk_size, d_model) @@ -198,7 +189,7 @@ def _fuse_single(w): w0 = w[:, :d_inner, :] w1 = w[:, d_inner:, :] chunk_size = d_inner // tp - padded_chunk_size = ((chunk_size + 127) // 128) * 128 + padded_chunk_size = pad_to_tpu_lanes(chunk_size) pad_amount = padded_chunk_size - chunk_size gate_chunks = w0.reshape(num_experts, tp, chunk_size, d_model) up_chunks = w1.reshape(num_experts, tp, chunk_size, d_model) @@ -215,13 +206,6 @@ def _fuse_single(w): return list(jnp.unstack(fused, axis=0)) -class Identity(Operation): - """Returns the input tensor unmodified.""" - - def __call__(self, tensors, **kwargs): - return tensors[0] - - # ========================================== # 2. Rule # ========================================== @@ -276,12 +260,17 @@ def __init__( num_kv_heads: Optional[int] = None, head_dim: Optional[int] = None, config: Any = None, + trainer_config: Any = None, + rollout_backend: str = "maxtext", # Defaults to MoEFusedLayout.PER_SHARD_INTERLEAVE; resolved in the body # because MoEFusedLayout is defined further down this module. 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 " @@ -289,15 +278,15 @@ def __init__( "MaxText-to-MaxText path, or a non-empty rule list." ) self.rules = rules - self.tp = tp - self.num_kv_heads = num_kv_heads - self.head_dim = head_dim + self.tp = resolve_rollout_tp(config, tp) + # 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 " @@ -306,23 +295,42 @@ def __init__( ) self._direct = MaxTextToMaxTextConverter( config=config, + tp=self.tp, 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: + keys = [ + config.model_name, + config.decoder_block.value if hasattr(config.decoder_block, "value") else config.decoder_block, + ] + for candidate_key in keys: + if ( + candidate_key + and candidate_key in MODEL_TO_CONVERSION_RULES + and MODEL_TO_CONVERSION_RULES[candidate_key] is not None + ): + self.rules = MODEL_TO_CONVERSION_RULES[candidate_key] + break 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() @@ -341,8 +349,6 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Any: out = tensors for op in rule.operations: out = op(out, tp=self.tp) - if not isinstance(out, list) and op != rule.operations[-1]: - out = [out] if isinstance(out, list) and len(out) > 1 and "{}" in rule.target_pattern: for i, tensor in enumerate(out): @@ -365,6 +371,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 @@ -431,8 +451,6 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Any: ), ], } -# Backward compatibility alias -_MODEL_TO_CONVERSION_RULES = MODEL_TO_CONVERSION_RULES # ========================================== @@ -654,14 +672,24 @@ class MaxTextToMaxTextConverter: def __init__( self, config: Any, + tp: int = 1, 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, ): self.config = config + self.tp = resolve_rollout_tp(config, tp) self.moe_fused_layout = moe_fused_layout self.allow_unused_source_keys = allow_unused_source_keys self.debug = debug + if prefuse_moe_weights is not None: + self.prefuse_moe_weights = prefuse_moe_weights + else: + self.prefuse_moe_weights = 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) self.cycle = int(getattr(config, "inhomogeneous_layer_cycle_interval", 1) or 1) self.num_decoder_layers = int(config.num_decoder_layers) @@ -678,14 +706,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 +738,65 @@ 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] + match = re.fullmatch(r"layer_(\d+)", slot_token) if isinstance(slot_token, str) else None + if match: + slot = int(match.group(1)) + 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], @@ -804,12 +897,7 @@ 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. @@ -818,11 +906,14 @@ 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 return _fuse_and_unstack_moe( wi_0, wi_1, self.scan_axis, - _get_n_shards(tgt_val, tgt_fused_axis), + n_shards, tgt_shape, scan_fused_axis, tgt_fused_axis, @@ -839,14 +930,106 @@ 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): + """Returns target-free slices for a scanned parameter.""" + 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.scan_axis <= 1 else 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) + 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): + """Fuses gate/up MoE kernels without a target state to derive shapes from.""" + 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 + ) + unrolled_shape = list(wi_0.shape[: self.scan_axis] + wi_0.shape[self.scan_axis + 1 :]) + tgt_shape = tuple(unrolled_shape[:-1] + [2 * target_intermediate]) + if isinstance(wi_0, jax.ShapeDtypeStruct): + return tuple(jax.ShapeDtypeStruct(tgt_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_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 = self.tp if self.tp > 1 else _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): + """Executes a conversion plan group without a target state.""" + path = group.source_path + target_dtype = self._resolve_target_dtype() + + if group.op == "identity": + 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) + return [(tgt_key, val) for _, tgt_key in group.targets] + + 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]] + 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) + 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 @@ -855,12 +1038,6 @@ def _execute_group(self, group: _PlanGroup, src_flat, tgt_flat): out = _align_per_axis(val, first_tgt.shape, getattr(first_tgt, "sharding", None), path) return [(tgt_key, out) 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": wi_0, wi_1 = (_apply_dtype_cast(src_flat[k], first_tgt.dtype, path) for k in group.source_keys) self._check_scan_axis(wi_0, path) @@ -886,24 +1063,24 @@ 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)) + 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,7 +1107,7 @@ 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), @@ -965,6 +1142,9 @@ 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: @@ -974,9 +1154,60 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]: 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 + return jax.tree_util.tree_map( + nnx.Param, + 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( + nnx.Param, + nested, + ) + + del src_flat def _rekey_to_target(flat_dotted: Dict[str, Any], target_state: Any) -> Dict[str, Any]: diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py index 1ecacf3fe0..6c44344fe7 100644 --- a/tests/post_training/unit/weight_converter_test.py +++ b/tests/post_training/unit/weight_converter_test.py @@ -21,10 +21,13 @@ """ import os +import gc +import resource # 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 @@ -58,10 +61,13 @@ def _config(**overrides): + """Returns a mock configuration object for tests.""" cfg = pytypes.SimpleNamespace( num_decoder_layers=NUM_LAYERS, inhomogeneous_layer_cycle_interval=CYCLE, param_scan_axis=SCAN_AXIS, + model_name="test_model", + decoder_block="default", ) for key, value in overrides.items(): setattr(cfg, key, value) @@ -439,5 +445,294 @@ def test_padded_moe_fusion_stays_on_the_source_mesh(self): ) +def _profile_conversion_worker(is_streaming, result_queue): + """Worker process to profile memory usage during conversion.""" + 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): # pylint: disable=redefined-outer-name + 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 # pylint: disable=import-outside-toplevel + + 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 8 (num_layers=8, cycle_interval=1)", 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.convert_utils import compute_padded_moe_mlp_dim # pylint: disable=import-outside-toplevel + + # 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_parity_vs_raiden_unscan_on_homogeneous(self): + from maxtext.integration.tunix.weight_mapping import raiden_unscan # pylint: disable=import-outside-toplevel + + 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_6_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)) # pylint: disable=protected-access + + # 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_7_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) # pylint: disable=protected-access + 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_8_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()