[Raiden Weight Sync 6/7] Integrate Raiden weight sync in MaxTextTrainingEngine with FFI support

src/maxtext/configs/base.yml

@@ -1374,6 +1374,18 @@ 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
+# 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
+# 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
 

src/maxtext/configs/types.py

@@ -2741,12 +2741,20 @@ 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_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 +2774,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'.",

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
@@ -93,7 +94,9 @@ def compute_padded_moe_mlp_dim(
 
   min_required = 2 * num_lanes * moe_mlp_tp_size
   if (hidden_size // moe_mlp_tp_size) % (2 * num_lanes) != 0:
-    return ((max(hidden_size, min_required) + min_required - 1) // min_required) * min_required
+    return (
+        (max(hidden_size, min_required) + min_required - 1) // min_required
+    ) * min_required
 
   return hidden_size
 
@@ -178,7 +181,9 @@ def _fuse_moe_weights(
         fused_shape_tuple,
         axis,
     )
-    new_src_flat[wi_target_key] = _interleave_moe_weights(wi_0, wi_1, fused_shape_tuple, n_shards, axis=axis)
+    new_src_flat[wi_target_key] = _interleave_moe_weights(
+        wi_0, wi_1, fused_shape_tuple, n_shards, axis=axis
+    )
     del wi_0, wi_1
 
   return new_src_flat
@@ -292,7 +297,8 @@ def _unstack_scanned_param(
       return tuple(jnp.unstack(src_val))
     else:
       logging.warning(
-          "Shape mismatch in scanned param '%s'. Src: %s, Tgt: %s. Cannot" " determine scan axis.",
+          "Shape mismatch in scanned param '%s'. Src: %s, Tgt: %s. Cannot"
+          " determine scan axis.",
           key_path,
           src_shape,
           tgt_shape,
@@ -330,7 +336,9 @@ def _get_n_shards(arr: jax.Array | np.ndarray, axis: int) -> int:
   """Returns the number of shards for a given axis of an array."""
   sharding = getattr(arr, "sharding", None)
   if isinstance(sharding, jax.sharding.NamedSharding):
-    return _partition_size(_spec_at_axis(sharding, axis), sharding.mesh)  # pyrefly: ignore[bad-argument-type]
+    return _partition_size(
+        _spec_at_axis(sharding, axis), sharding.mesh
+    )  # pyrefly: ignore[bad-argument-type]
   return 1
 
 
@@ -450,14 +458,18 @@ def _align_per_axis(
   if arr.shape == tgt_shape:
     return arr
   if len(arr.shape) != len(tgt_shape):
-    raise ShapeMismatchError(f"Rank mismatch for {key_path}: src={arr.shape} vs tgt={tgt_shape}")
+    raise ShapeMismatchError(
+        f"Rank mismatch for {key_path}: src={arr.shape} vs tgt={tgt_shape}"
+    )
 
   mismatches = []
   for axis, (s, t) in enumerate(zip(arr.shape, tgt_shape)):
     if s == t:
       continue
     if t < s:
-      raise ShapeMismatchError(f"Cannot shrink axis {axis} for {key_path}: src={s} -> tgt={t}")
+      raise ShapeMismatchError(
+          f"Cannot shrink axis {axis} for {key_path}: src={s} -> tgt={t}"
+      )
     mismatches.append((axis, s, t))
   if not mismatches:
     return arr
@@ -468,7 +480,9 @@ def _align_per_axis(
       mesh = tgt_sharding.mesh
       pad_specs = []
       for axis, s, t in mismatches:
-        n_shards = _partition_size(_spec_at_axis(tgt_sharding, axis), mesh)  # pyrefly: ignore[bad-argument-type]
+        n_shards = _partition_size(
+            _spec_at_axis(tgt_sharding, axis), mesh
+        )  # pyrefly: ignore[bad-argument-type]
         if t % n_shards != 0:
           raise ValueError(
               f"Target dimension {t} on axis {axis} for {key_path} is not "
@@ -499,14 +513,16 @@ def _align_per_axis(
   return _jit_repeat_axes(arr, tuple(repeats))
 
 
-@functools.partial(jax.jit, static_argnames=("tgt_shape", "n_shards", "axis", "lane_size"))
+@functools.partial(
+    jax.jit, static_argnames=("tgt_shape", "n_shards", "axis", "lane_size")
+)
 def _interleave_moe_weights(
     wi_0: jax.Array | np.ndarray,
     wi_1: jax.Array | np.ndarray,
     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 +533,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 +584,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)
@@ -661,10 +686,18 @@ def _bulk_align_and_unstack(
   if isinstance(arr, 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))
+    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)
+  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
+  )
 
   if arr.shape == scanned_tgt_shape:
     return _jit_unstack(arr, scan_axis)
@@ -692,16 +725,54 @@ 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"

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:
+    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
@@ -471,6 +473,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 +617,14 @@ 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
+
+    patch_raiden_worker_sync()
+  except Exception as e:
+    max_logging.log(f"Skipping raiden worker sync patch: {e}")
+