[Trellis] Weight Conversion that is compatible with Raiden

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,
   )
 
 

src/maxtext/configs/types.py

@@ -2704,6 +2704,10 @@ class VLLM(BaseModel):
           "the legacy transfer_state_directly / transfer_state_with_mappings paths."
       ),
   )
+  rollout_backend: Literal["maxtext", "vllm_torchax"] = Field(
+      "maxtext",
+      description="Rollout backend for trainer-side weight converter ('maxtext' or 'vllm_torchax').",
+  )
   weight_sync_debug: bool = Field(
       False,
       description=(

src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py

@@ -36,35 +36,72 @@
 simpler, single-axis case directly instead of adapting that function.
 """
 
-from typing import Any
+import gc
+from typing import Any, Iterator, Tuple, List
 
 import jax
 from flax import nnx
 from flax.traverse_util import flatten_dict, unflatten_dict
 
 
-def unscan_layers(
+def _unscan_one_key(
+    key: Tuple[Any, ...],
+    value: Any,
+    num_layers: int,
+    layer_container: str = "layers",
+    scan_axis: int = 1,
+) -> Tuple[List[Tuple[Tuple[Any, ...], Any]], bool]:
+  """Unscans a single flattened pytree entry.
+
+  Returns:
+    (list_of_entries, is_scanned), where list_of_entries is a list of (new_key, value_slice) pairs.
+  """
+  if layer_container not in key:
+    return [(key, value)], False
+
+  idx = key.index(layer_container)
+  prefix = key[:idx]
+  suffix = key[idx + 1 :]
+  arr = getattr(value, "value", value)
+
+  if not hasattr(arr, "shape") or arr.ndim <= scan_axis:
+    return [(key, value)], False
+
+  if arr.shape[scan_axis] != num_layers:
+    raise ValueError(
+        f"unscan_layers: {'.'.join(str(k) for k in key)!r} has shape {arr.shape}, expected axis {scan_axis} to be"
+        f" num_layers={num_layers}."
+    )
+
+  entries = []
+  for i in range(num_layers):
+    sliced = jax.lax.index_in_dim(arr, i, axis=scan_axis, keepdims=False)
+    new_key = prefix + (f"{layer_container}_{i}",) + suffix
+    entries.append((new_key, sliced))
+  return entries, True
+
+
+def unscan_layers_streaming(
     state: Any,
     num_layers: int,
     layer_container: str = "layers",
     scan_axis: int = 1,
-) -> Any:
-  """Splits `state`'s scanned `layer_container` axis into per-layer entries.
+    *,
+    keys_per_piece: int = 1,
+) -> Iterator[Any]:
+  """Yields unscanned layer pieces incrementally for Raiden weight sync.
 
   Args:
     state: An `nnx.State` (or any pytree exposing `to_pure_dict`/`to_dict`, or a plain nested dict) of MaxText params,
       scanned along `scan_axis` under a `layer_container` key (MaxText's default scan layout).
-    num_layers: Number of layers the scanned axis must have. Used both to validate the input and to bound the unscan
-      loop.
-    layer_container: The pytree key holding the scanned per-layer params (MaxText's decoder body uses "layers").
+    num_layers: Number of layers the scanned axis must have.
+    layer_container: The pytree key holding the scanned per-layer params.
     scan_axis: The axis along which layers are scanned (default 1).
+    keys_per_piece: Number of original flattened keys to batch per yielded piece (default 1).
 
-  Returns:
-    A nested dict with `layer_container` keys replaced by `f"{layer_container}_{i}"` for each layer `i`, each holding
-    the corresponding `scan_axis` slice of the original array, wrapped in `nnx.Param` -- matching `nnx.state(...,
-    nnx.Param)`'s leaf type, so downstream consumers (`raiden_synchronizer.flatten_weights`, which unwraps `.value`)
-    see the same leaf shape whether or not this transform ran. Non-scanned entries (e.g. embeddings, final norm) pass
-    through unchanged, also rewrapped.
+  Yields:
+    Nested dicts with `nnx.Param`-wrapped leaves, each containing `keys_per_piece` keys' worth of
+    unscanned slices.
   """
   if hasattr(state, "to_pure_dict"):
     pure = state.to_pure_dict()
@@ -73,57 +110,80 @@ def unscan_layers(
   elif isinstance(state, dict):
     pure = state
   else:
-    return state
+    yield state
+    return
 
   flat = flatten_dict(pure)
-  new_flat = {}
-  unscanned_count = 0
-
-  # Drain `flat` as we go (pop, not iterate-then-keep) rather than holding
-  # every original scanned array alive for the whole function: at 30B-A3B
-  # scale (padded MoE weights are tens of GB each), keeping both the
-  # original scanned tree and the ~num_layers-times-larger unscanned tree
-  # alive simultaneously roughly doubles peak host memory during Raiden's
-  # D2H staging -- confirmed as the direct cause of an OOMKill there.
-  for key in list(flat.keys()):
-    value = flat.pop(key)
-    if layer_container not in key:
-      new_flat[key] = value
-      continue
-
-    idx = key.index(layer_container)
-    prefix = key[:idx]
-    suffix = key[idx + 1 :]
-    arr = getattr(value, "value", value)
-
-    if arr is None or not hasattr(arr, "shape") or getattr(arr, "ndim", 0) <= scan_axis:
-      # Not a per-layer leaf (shouldn't happen for real params under
-      # `layers`, but don't silently drop anything unexpected). Keep the
-      # original (possibly already-wrapped) value, matching pre-existing
-      # behavior -- unlike the actually-scanned case below, there's no
-      # multi-copy blowup here worth restructuring around.
-      new_flat[key] = value
-      continue
-    del value
-
-    if arr.shape[scan_axis] != num_layers:
-      raise ValueError(
-          f"unscan_layers: {'.'.join(key)!r} has shape {arr.shape}, expected axis {scan_axis} to be"
-          f" num_layers={num_layers}."
-      )
-
-    for i in range(num_layers):
-      sliced = jax.lax.index_in_dim(arr, i, axis=scan_axis, keepdims=False)
-      new_key = prefix + (f"{layer_container}_{i}",) + suffix
-      new_flat[new_key] = sliced
-    del arr
-    unscanned_count += 1
 
-  if unscanned_count == 0:
+  has_scanned = any(
+      layer_container in key
+      and hasattr(getattr(flat[key], "value", flat[key]), "shape")
+      and getattr(getattr(flat[key], "value", flat[key]), "ndim", 0) > scan_axis
+      for key in flat
+  )
+  if not has_scanned:
     raise ValueError(
         f"unscan_layers: found no scanned '{layer_container}' entries to unscan "
         "-- state may already be unscanned, or layer_container is wrong."
     )
 
+  keys_per_piece = max(1, keys_per_piece)
+  flat_keys = list(flat.keys())
+  for i in range(0, len(flat_keys), keys_per_piece):
+    chunk_keys = flat_keys[i : i + keys_per_piece]
+    piece_flat = {}
+    for key in chunk_keys:
+      value = flat.pop(key)
+      outputs, _ = _unscan_one_key(
+          key, value, num_layers=num_layers, layer_container=layer_container, scan_axis=scan_axis
+      )
+      for new_key, sliced in outputs:
+        piece_flat[new_key] = sliced
+      del value, outputs
+
+    nested = unflatten_dict(piece_flat)
+    del piece_flat
+    yield jax.tree_util.tree_map(
+        lambda x: nnx.Param(x) if not isinstance(x, (nnx.Param, nnx.Variable)) else x,
+        nested,
+    )
+
+  del flat
+  gc.collect()
+
+
+def unscan_layers(
+    state: Any,
+    num_layers: int,
+    layer_container: str = "layers",
+    scan_axis: int = 1,
+) -> Any:
+  """Splits `state`'s scanned `layer_container` axis into per-layer entries.
+
+  Args:
+    state: An `nnx.State` (or any pytree exposing `to_pure_dict`/`to_dict`, or a plain nested dict) of MaxText params,
+      scanned along `scan_axis` under a `layer_container` key (MaxText's default scan layout).
+    num_layers: Number of layers the scanned axis must have. Used both to validate the input and to bound the unscan
+      loop.
+    layer_container: The pytree key holding the scanned per-layer params (MaxText's decoder body uses "layers").
+    scan_axis: The axis along which layers are scanned (default 1).
+
+  Returns:
+    A nested dict with `layer_container` keys replaced by `f"{layer_container}_{i}"` for each layer `i`, each holding
+    the corresponding `scan_axis` slice of the original array, wrapped in `nnx.Param` -- matching `nnx.state(...,
+    nnx.Param)`'s leaf type, so downstream consumers (`raiden_synchronizer.flatten_weights`, which unwraps `.value`)
+    see the same leaf shape whether or not this transform ran. Non-scanned entries (e.g. embeddings, final norm) pass
+    through unchanged, also rewrapped.
+  """
+  if not hasattr(state, "to_pure_dict") and not hasattr(state, "to_dict") and not isinstance(state, dict):
+    return state
+
+  new_flat = {}
+  for piece in unscan_layers_streaming(
+      state, num_layers=num_layers, layer_container=layer_container, scan_axis=scan_axis
+  ):
+    new_flat.update(flatten_dict(piece))
+  gc.collect()
   nested = unflatten_dict(new_flat)
-  return jax.tree_util.tree_map(nnx.Param, nested)
+  del new_flat
+  return nested

src/maxtext/integration/vllm/convert_utils.py

@@ -160,9 +160,40 @@ class ShapeMismatchError(ValueError):
   """Raised when source and target shapes are incompatible."""
 
 
-def _apply_dtype_cast(val: jax.Array | np.ndarray, tgt_dtype: jnp.dtype, src_key: str) -> jax.Array | np.ndarray:
+def reclaim_host_memory() -> None:
+  """Runs garbage collection and triggers libc malloc_trim to return free heap to the OS."""
+  import gc  # pylint: disable=g-import-not-at-top
+  gc.collect()
+  try:
+    import ctypes  # pylint: disable=g-import-not-at-top
+    ctypes.CDLL("libc.so.6").malloc_trim(0)
+  except Exception as e:  # pylint: disable=broad-exception-caught
+    logging.debug("reclaim_host_memory: malloc_trim unavailable or failed: %s", e)
+
+
+def normalize_dtype(tgt_dtype: Any) -> Any:
+  """Normalizes string or dtype representations into a standard jnp.dtype."""
+  if tgt_dtype is None:
+    return None
+  if isinstance(tgt_dtype, str):
+    if tgt_dtype in ("bfloat16", "bf16"):
+      return jnp.bfloat16
+    if tgt_dtype in ("float32", "fp32"):
+      return jnp.float32
+    return jnp.dtype(tgt_dtype)
+  return tgt_dtype
+
+
+def _apply_dtype_cast(val: Any, tgt_dtype: Any, src_key: str) -> Any:
   """Casts val to target dtype if needed, logging a warning on type mismatch."""
-  if val.dtype != tgt_dtype:
+  tgt_dtype = normalize_dtype(tgt_dtype)
+  if isinstance(val, jax.ShapeDtypeStruct):
+    if tgt_dtype is not None and val.dtype != tgt_dtype:
+      return jax.ShapeDtypeStruct(val.shape, tgt_dtype)
+    return val
+  if not hasattr(val, "dtype"):
+    return val
+  if tgt_dtype is not None and val.dtype != tgt_dtype:
     logging.log_first_n(
         logging.WARNING,
         "Type mismatch on %s: %s -> %s",
@@ -171,7 +202,8 @@ def _apply_dtype_cast(val: jax.Array | np.ndarray, tgt_dtype: jnp.dtype, src_key
         val.dtype,
         tgt_dtype,
     )
-    return val.astype(tgt_dtype)
+    if hasattr(val, "astype"):
+      return val.astype(tgt_dtype)
   return val
 
 
@@ -409,6 +441,8 @@ def _align_per_axis(
   path here is bulk alignment of scanned MoE weights, where eager
   dispatch was costing tens of seconds per tensor.
   """
+  if isinstance(arr, jax.ShapeDtypeStruct):
+    return jax.ShapeDtypeStruct(tgt_shape, arr.dtype)
   if not hasattr(arr, "shape"):
     return arr
   if arr.shape == tgt_shape:
@@ -530,7 +564,7 @@ def _fuse_and_unstack_moe(
     scan_axis: int,
     n_shards: int,
     tgt_shape: Tuple[int, ...],
-    scan_fused_axis: int,
+    scan_fused_axis: int,  # TODO(follow-up): Unused in function body, preserved for caller compatibility.
     tgt_fused_axis: int,
 ) -> Tuple[jax.Array | np.ndarray, ...]:
   """Fuses wi_0/wi_1 per unstacked layer to keep peak intermediate HBM allocation low.
@@ -604,6 +638,11 @@ def _bulk_align_and_unstack(
     A tuple of `num_layers` per-layer arrays at the per-layer target shape.
   """
   per_layer_shape = per_layer_tgt_val.shape
+  if isinstance(arr, jax.ShapeDtypeStruct) or isinstance(per_layer_tgt_val, jax.ShapeDtypeStruct):
+    num_layers = arr.shape[scan_axis]
+    tgt_dtype = getattr(per_layer_tgt_val, "dtype", getattr(arr, "dtype", jnp.float32))
+    return tuple(jax.ShapeDtypeStruct(per_layer_shape, tgt_dtype) for _ in range(num_layers))
+
   scanned_tgt_shape = per_layer_shape[:scan_axis] + (arr.shape[scan_axis],) + per_layer_shape[scan_axis:]
   scanned_tgt_sharding = _scanned_sharding_from_per_layer(getattr(per_layer_tgt_val, "sharding", None), scan_axis)
 

src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py

@@ -48,19 +48,7 @@ class AttentionMetadata:
 _HYBRID_LAYER_IMBALANCE_THRESHOLD = 1.5
 
 
-def next_power_of_two(x: int) -> int:
-  """Finds the smallest power of 2 >= x using bit manipulation.
-
-  Args:
-    x: The input number (should be an integer).
-
-  Returns:
-    The smallest integer power of 2 that is >= x.
-  """
-  assert x > 0
-  if x == 1:
-    return 1
-  return 1 << (x - 1).bit_length()
+from maxtext.integration.vllm.moe_padding import compute_padded_moe_mlp_dim, next_power_of_two
 
 
 def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters:
@@ -166,11 +154,8 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters
   # The GMM_v2 kernel requires the MLP dimension per expert to be at least 2x the number of TPU lanes
   # to ensure efficient execution. See the validate_inputs() method in the following file for more details:
   # https://github.com/vllm-project/tpu-inference/blob/main/tpu_inference/kernels/megablox/gmm_v2.py
-  if hidden_size is not None and (hidden_size // moe_mlp_tp_size) % (2 * num_lanes) != 0:
-    padded_hidden_size = next_power_of_two(hidden_size)
-    while (padded_hidden_size // moe_mlp_tp_size) < (2 * num_lanes):
-      padded_hidden_size = next_power_of_two(padded_hidden_size + 1)
-
+  padded_hidden_size = compute_padded_moe_mlp_dim(hidden_size, moe_mlp_tp_size, num_lanes)
+  if padded_hidden_size is not None and padded_hidden_size != hidden_size:
     # This inflates every expert weight, so it is a real memory/FLOP cost rather than a
     # cosmetic reshape: at moe_mlp_tp_size=4 a 512-wide MoE is padded to 1024 (2x the MoE
     # weights), and at moe_mlp_tp_size=8 to 2048 (4x). Log it at WARNING so it is visible

src/maxtext/integration/vllm/maxtext_vllm_rollout.py

@@ -464,17 +464,25 @@ def __init__(
       self,
       tokenizer: Any,
       config: VllmConfig,
-      converter: Any = None,
       direct_maxtext_sync: bool = False,
       scan_axis: int = 1,
       layer_pattern_length: Optional[int] = None,
   ):
     super().__init__(tokenizer=tokenizer, config=config)
-    self._converter = converter
-    self.converter = converter
     self._direct_maxtext_sync = direct_maxtext_sync
     self._scan_axis = scan_axis
     self._layer_pattern_length = layer_pattern_length
+    model_config = getattr(config, "model_config", None)
+    model_name = getattr(model_config, "model", "") or ""
+    architectures = getattr(model_config, "architectures", []) or []
+    hf_config = getattr(model_config, "hf_config", None)
+    model_type = getattr(hf_config, "model_type", "") or ""
+    arch_str = " ".join(str(a) for a in architectures)
+    self._is_gemma = (
+        "gemma" in str(model_name).lower()
+        or "gemma" in str(model_type).lower()
+        or "gemma" in str(arch_str).lower()
+    )
 
   def update_params(
       self,
@@ -494,12 +502,7 @@ def update_params(
             pass
         raise
     if self._converter is None:
-      if self._direct_maxtext_sync:
-        updated_weights = unroll_qwen_scanned_weights(
-            updated_weights,
-            scan_axis=self._scan_axis,
-            pattern_length=self._layer_pattern_length,
-        )
+      if self._direct_maxtext_sync and self._is_gemma:
         updated_weights = unroll_gemma_scanned_weights(updated_weights)
     try:
       return super().update_params(updated_weights, filter_types)
@@ -742,7 +745,6 @@ def __init__(
             additional_config=rollout_additional_config,
             sampling_kwargs=rollout_config.rollout_vllm_sampling_kwargs,
         ),
-        converter=converter,
         direct_maxtext_sync=direct_maxtext_sync,
         scan_axis=getattr(maxtext_config, "param_scan_axis", 1),
         layer_pattern_length=getattr(maxtext_config, "inhomogeneous_layer_cycle_interval", None),