From e80de1cd91a1caace7a178a7779b00d0dcd87a85 Mon Sep 17 00:00:00 2001
From: A9isha <mazumdera@google.com>
Date: Wed, 26 Aug 2026 18:02:37 +0000
Subject: [PATCH 1/7] Make three silent failures in the Raiden training path
 visible

Each change targets a failure that produced no usable signal at the point of
cause: weights that never transfer, metadata that fails in another process, and
a performance cliff recorded only in a docstring.

unscan_layers had no tests, and it is the piece that decides whether trainer and
sampler tensor names agree. Both sides name tensors with jax.tree_util.keystr,
and nothing cross-checks the two sets: raiden_handler._validate_metadata only
validates one manifest's internal consistency (mesh rank, duplicate
variable/layer keys, sharding specs). A naming error therefore surfaces as
weights that silently never transfer. Two cases pin non-obvious invariants:
unscan_layers returns a plain nested dict while the sampler binds an nnx.State,
and keystr renders those identically only because the transform rewraps leaves
in nnx.Param -- dropping that rewrap would rename every tensor ("['k']" vs
"['k'].value"); and an already-unscanned state must raise, since without that
guard it would return unchanged and bind under scanned names.

prepare_weight_sync returned empty metadata on two paths: a missing
raiden_synchronizer (warning-level) and an unrecognised staging_transport (no
log at all). Neither is silent end to end -- WeightSyncCoordinator rejects an
empty side -- but the failure lands far from the cause, surfacing in another
process as "metadata collection returned an empty side", a count that never
names the missing module or the bad transport. The import case is the common one
rather than a corner: raiden_synchronizer ships only on tunix's Raiden branch,
so any released tunix takes it. Both now raise where the cause is known, with
the ImportError chained so the traceback keeps the module name. Because
staging_transport defaults to "raiden", this also reaches callers that never
asked for it, so the engine e2e test now probes for the synchronizer the way the
engine does -- exercising the real staging path where Raiden exists and the
documented failure where it does not.

_batch_data_shardings falls back to replicating the batch dimension when it does
not divide the batch axis's mesh size. That is correct -- every device along the
axis computes the whole micro-batch -- but it costs N times the work a sharded
one would do there. An invisible performance cliff is harder to notice than a
wrong number, because XLA's caching can make it look like nothing worse than a
slow run; the file already warns once per instance when a signature half cannot
be compared, and this extends that treatment. Warned once per instance rather
than per leaf, since the check runs under a tree_map over every loss input and
they normally share a batch dim. A sequence-packed micro-batch is always size 1
and has no alternative, so the message says the fallback may well be deliberate.

Verification. The unscan suite has teeth: renaming the emitted key from
layers_{i} to layer_{i} fails 5 of its 11 tests, including the name-equality
one. Marked post_training and left in tests/unit, which is already in
cpu-post-training-unit's path list, so the marker alone routes it -- tests/ and
tests/integration are not in that list, which is how the engine tests once ended
up collected by no job at all; collection confirms 11 tests in
cpu-post-training-unit and 0 in cpu-unit. The staging and sharding tests fail
without their respective changes. The sharding tests stub both the data spec and
the axis size: a single-device test mesh returns None in the batch position,
making the branch unreachable as configured, and an earlier draft asserted
`spec[0] is None` and passed without running the code under test at all.
---
 src/maxtext/training_engine/maxtext_engine.py |  24 ++-
 .../unit/maxtext_engine_e2e_test.py           |  20 +-
 .../post_training/unit/maxtext_engine_test.py |  65 ++++++
 tests/unit/raiden_unscan_test.py              | 191 ++++++++++++++++++
 4 files changed, 296 insertions(+), 4 deletions(-)
 create mode 100644 tests/unit/raiden_unscan_test.py

diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 356b07eeb1..4cd6c77014 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -137,6 +137,14 @@ def _batch_signature(dynamic_batch: Any, static_batch: dict[str, Any]) -> Any:
   return (treedef, shapes, static_batch)
 
 
+_REPLICATED_BATCH_DIM_WARNING = (
+    "Loss input with batch dim %d does not divide mesh axis %r (size %d), so that "
+    "dimension is replicated instead of sharded: every device along the axis holds and "
+    "computes the whole micro-batch, %dx the work a sharded one would do there. Results "
+    "stay correct. If it was not deliberate -- a sequence-packed micro-batch is always "
+    "size 1 and has no alternative -- make the micro-batch a multiple of the axis size."
+)
+
 _UNCOMPARABLE_SIGNATURE_WARNING = (
     "Could not compare %s between fwd_bwd calls (%s), so the engine cannot tell whether "
     "the compiled kernel is still valid and will recompile on EVERY fwd_bwd from now on. %s"
@@ -351,6 +359,7 @@ def __init__(
     self._compile_requested = False
     self._compiled_signature: Any = None
     self._signature_compare_warned: bool = False
+    self._replicated_batch_warned: bool = False
     if not training_config.model_name:
       raise ValueError("training_config.model_name must be specified")
     model_or_model_mesh_pair = model_creation_utils.from_pretrained(
@@ -875,7 +884,8 @@ def _batch_data_shardings(self, dynamic_batch: Any) -> Any:
     sequence-packed micro-batch, always size 1) replicates that dim instead of sharding
     it -- every device holds and computes on the same data with no cross-device split,
     which is correct (there's nothing to reduce back together afterwards) but wastes
-    compute across the axis for that micro-batch.
+    compute across the axis for that micro-batch. That is an N-fold cost, so it warns
+    once per instance rather than living only in this docstring.
     """
     data_sharding = sharding.get_input_data_sharding(self._config, self._mesh)
     data_spec = tuple(data_sharding.spec)
@@ -885,8 +895,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)
diff --git a/tests/post_training/unit/maxtext_engine_e2e_test.py b/tests/post_training/unit/maxtext_engine_e2e_test.py
index 1f2bd0e432..54896ecc6d 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.worker.raiden_synchronizer")
+  _RAIDEN_AVAILABLE = True
+except ImportError:
+  _RAIDEN_AVAILABLE = False
+
 pytestmark = [pytest.mark.post_training]
 
 
@@ -97,7 +109,13 @@ def run(
       step_metrics = self.trainer.get_metrics(clear_cache=True)
       history.append(step_metrics)
 
-      _ = self.trainer.prepare_weight_sync()
+      if _RAIDEN_AVAILABLE:
+        _ = self.trainer.prepare_weight_sync()
+      else:
+        # Without the transport the engine must raise rather than hand back empty
+        # metadata, which would fail later and far from the cause.
+        with pytest.raises(RuntimeError, match="raiden_synchronizer"):
+          self.trainer.prepare_weight_sync()
 
     self.trainer.close()
     return history
diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py
index 5ee877bbb2..3c28a586fb 100644
--- a/tests/post_training/unit/maxtext_engine_test.py
+++ b/tests/post_training/unit/maxtext_engine_test.py
@@ -16,6 +16,7 @@
 # pylint: disable=protected-access
 
 import dataclasses
+import sys
 import types
 from typing import Any
 from unittest import mock
@@ -1331,6 +1332,70 @@ def test_perplexity_is_emitted_alongside_the_loss(self):
     self.assertIn("perplexity", processed)
     self.assertAlmostEqual(processed["perplexity"], float(np.exp(6.0)), places=3)
 
+  def test_prepare_weight_sync_raises_when_raiden_is_unavailable(self):
+    """A missing raiden_synchronizer must fail here, not as an empty result downstream.
+
+    Returning empty metadata defers the failure to `WeightSyncCoordinator`, which raises
+    "metadata collection returned an empty side" -- a count from another process that
+    never names the missing module. `raiden_synchronizer` ships only on tunix's Raiden
+    branch, so this is the common case on a released tunix, not a corner.
+    """
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+
+    # Setting the entry to None makes `from ... import raiden_synchronizer` raise
+    # ImportError, which is what an installed tunix without the module does.
+    with mock.patch.dict(sys.modules, {"tunix.experimental.worker.raiden_synchronizer": None}):
+      with self.assertRaisesRegex(RuntimeError, "raiden_synchronizer"):
+        t.prepare_weight_sync()
+
+  def test_prepare_weight_sync_rejects_an_unknown_transport(self):
+    """An unrecognised transport must name itself rather than return empty metadata."""
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+    with self.assertRaisesRegex(ValueError, "raidan"):
+      t.prepare_weight_sync(staging_transport="raidan")
+
+  def _sharded_batch_spec(self, engine, axis="data"):
+    """A data sharding whose batch dim is actually sharded.
+
+    The single-device test mesh makes `get_input_data_sharding` return a spec with `None`
+    in the batch position, so the replication branch is unreachable as configured -- an
+    earlier version of these tests asserted `spec[0] is None` and passed without ever
+    running the code under test. Stub a spec that shards the batch dim instead.
+    """
+    return jax.sharding.NamedSharding(engine._mesh, jax.sharding.PartitionSpec(axis, None))  # pylint: disable=protected-access
+
+  def test_indivisible_batch_dim_replicates_and_warns_once(self):
+    """Replicating the batch dim is an N-fold compute cliff, so it must be audible."""
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+    batch = {"a": jnp.zeros((1, 4)), "b": jnp.zeros((1, 4))}
+
+    # Batch dim 1 against a 2-wide axis: indivisible, so the dim must be replicated.
+    with mock.patch.object(maxtext_engine.sharding, "get_input_data_sharding", return_value=self._sharded_batch_spec(t)):
+      with mock.patch.object(type(t), "_batch_axis_size", return_value=2):
+        with self.assertLogs(level="WARNING") as logs:
+          shardings = t._batch_data_shardings(batch)  # pylint: disable=protected-access
+          t._batch_data_shardings(batch)  # pylint: disable=protected-access
+
+    for name, leaf_sharding in shardings.items():
+      self.assertIsNone(leaf_sharding.spec[0], f"{name} should have its batch dim replicated")
+
+    # Once per instance, not per leaf and not per call: two leaves over two calls is four
+    # chances to warn.
+    warnings = [line for line in logs.output if "does not divide mesh axis" in line]
+    self.assertLen(warnings, 1)
+    self.assertIn("2x the work", warnings[0])
+
+  def test_divisible_batch_dim_stays_sharded_and_is_silent(self):
+    """The normal case must neither replicate nor warn."""
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+
+    with mock.patch.object(maxtext_engine.sharding, "get_input_data_sharding", return_value=self._sharded_batch_spec(t)):
+      with mock.patch.object(type(t), "_batch_axis_size", return_value=2):
+        shardings = t._batch_data_shardings({"a": jnp.zeros((4, 4))})  # pylint: disable=protected-access
+
+    self.assertEqual(shardings["a"].spec[0], "data")
+    self.assertFalse(t._replicated_batch_warned)  # pylint: disable=protected-access
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/unit/raiden_unscan_test.py b/tests/unit/raiden_unscan_test.py
new file mode 100644
index 0000000000..694f2449b0
--- /dev/null
+++ b/tests/unit/raiden_unscan_test.py
@@ -0,0 +1,191 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for `raiden_unscan.unscan_layers`.
+
+This transform decides the *names* Raiden binds. The trainer runs scanned
+(`scan_layers=True`); the sampler loads its MaxText model unscanned, and Raiden matches
+tensors by `jax.tree_util.keystr` path. Nothing downstream cross-checks the two name sets
+-- `raiden_handler._validate_metadata` only checks a single manifest's internal
+consistency (mesh rank, duplicate variable/layer keys, sharding specs) -- so a naming
+error here surfaces as weights that silently never transfer, not as an exception.
+
+The transform is pure pytree manipulation, so all of this runs on CPU in milliseconds.
+"""
+
+from absl.testing import absltest
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from maxtext.integration.tunix.weight_mapping import raiden_unscan
+import numpy as np
+import pytest
+
+# This transform exists only for the Tunix/Raiden weight-sync path, so it is graded with
+# the rest of that work. `tests/unit` is in cpu-post-training-unit's path list, so the
+# marker alone routes it there -- no file move needed (unlike tests/ or tests/integration,
+# which are not in that list and would be collected by no job at all).
+pytestmark = [pytest.mark.post_training]
+
+
+_NUM_LAYERS = 3
+_IN, _OUT, _VOCAB = 4, 8, 10
+
+
+def _unwrap(leaf):
+  """Reads a leaf's array whether or not it is wrapped in an `nnx.Param`."""
+  if isinstance(leaf, nnx.Variable):
+    return leaf[...]
+  return leaf
+
+
+def _names(tree) -> list[str]:
+  """The names Raiden binds: exactly what `raiden_synchronizer.flatten_weights` computes."""
+  return sorted(jax.tree_util.keystr(p) for p, _ in jax.tree_util.tree_leaves_with_path(tree))
+
+
+class ScannedInner(nnx.Module):
+  """One scanned param: the layer axis lives at axis 1 of a single array."""
+
+  def __init__(self, num_layers: int = _NUM_LAYERS):
+    self.kernel = nnx.Param(jnp.arange(_IN * num_layers * _OUT, dtype=jnp.float32).reshape(_IN, num_layers, _OUT))
+    self.scale = nnx.Param(jnp.arange(_OUT * num_layers, dtype=jnp.float32).reshape(_OUT, num_layers))
+
+
+class ScannedModel(nnx.Module):
+  """A trainer-side model: scanned `layers`, plus non-layer params that must pass through."""
+
+  def __init__(self, num_layers: int = _NUM_LAYERS):
+    self.layers = ScannedInner(num_layers)
+    self.embed = nnx.Param(jnp.zeros((_VOCAB, _IN)))
+
+
+class UnscannedInner(nnx.Module):
+
+  def __init__(self):
+    self.kernel = nnx.Param(jnp.zeros((_IN, _OUT)))
+    self.scale = nnx.Param(jnp.zeros((_OUT,)))
+
+
+class UnscannedModel(nnx.Module):
+  """A sampler-side model: one submodule per layer, named `layers_0..N-1`."""
+
+  def __init__(self, num_layers: int = _NUM_LAYERS):
+    for i in range(num_layers):
+      setattr(self, f"layers_{i}", UnscannedInner())
+    self.embed = nnx.Param(jnp.zeros((_VOCAB, _IN)))
+
+
+class UnscanLayersTest(absltest.TestCase):
+
+  def _scanned_state(self, num_layers: int = _NUM_LAYERS):
+    return nnx.state(ScannedModel(num_layers), nnx.Param)
+
+  def test_names_match_an_unscanned_model_exactly(self):
+    """The point of the transform: trainer names must equal sampler names.
+
+    Raiden binds by `keystr` path on both sides, and nothing validates that the two sets
+    agree, so this is the assertion that a silent no-transfer would violate.
+    """
+    unscanned = raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS)
+    sampler_side = nnx.state(UnscannedModel(), nnx.Param)
+    self.assertEqual(_names(unscanned), _names(sampler_side))
+
+  def test_plain_dict_and_nnx_state_produce_identical_names(self):
+    """`unscan_layers` returns a plain nested dict, the sampler binds an `nnx.State`.
+
+    `keystr` renders both identically only because the transform rewraps leaves in
+    `nnx.Param`. Dropping that rewrap would rename every tensor (`['k']` vs `['k'].value`)
+    and break every transfer, so pin it.
+    """
+    unscanned = raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS)
+    self.assertIsInstance(jax.tree_util.tree_leaves(unscanned, is_leaf=lambda x: isinstance(x, nnx.Param))[0], nnx.Param)
+    self.assertTrue(all(n.endswith(".value") for n in _names(unscanned)), _names(unscanned))
+
+  def test_slices_carry_the_right_values(self):
+    """Layer i must receive index i of the scan axis -- not a transpose or an off-by-one."""
+    state = self._scanned_state()
+    original = np.asarray(state.to_pure_dict()["layers"]["kernel"])
+    unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS)
+
+    for i in range(_NUM_LAYERS):
+      got = unscanned[f"layers_{i}"]["kernel"]
+      got = np.asarray(_unwrap(got))
+      self.assertEqual(got.shape, (_IN, _OUT))
+      np.testing.assert_array_equal(got, original[:, i, :])
+
+  def test_rank_two_param_is_also_unscanned(self):
+    """A rank-2 scanned param (e.g. a norm scale) slices down to rank 1."""
+    unscanned = raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS)
+    for i in range(_NUM_LAYERS):
+      scale = unscanned[f"layers_{i}"]["scale"]
+      self.assertEqual(np.asarray(_unwrap(scale)).shape, (_OUT,))
+
+  def test_non_layer_entries_pass_through_unchanged(self):
+    """Embeddings and final norms have no layer axis and must survive untouched."""
+    state = self._scanned_state()
+    embed_before = np.asarray(state.to_pure_dict()["embed"])
+    unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS)
+
+    self.assertIn("embed", unscanned)
+    embed_after = unscanned["embed"]
+    np.testing.assert_array_equal(np.asarray(_unwrap(embed_after)), embed_before)
+    self.assertNotIn("layers", unscanned)
+
+  def test_layer_count_mismatch_raises(self):
+    """A wrong num_layers must fail loudly rather than bind truncated weights."""
+    with self.assertRaisesRegex(ValueError, "expected axis 1 to be num_layers=99"):
+      raiden_unscan.unscan_layers(self._scanned_state(), num_layers=99)
+
+  def test_already_unscanned_state_raises(self):
+    """The anti-silent-no-op guard.
+
+    Without it an already-unscanned (or wrongly-keyed) state would return unchanged and
+    bind under scanned names, transferring nothing with no error anywhere.
+    """
+    with self.assertRaisesRegex(ValueError, "found no scanned 'layers' entries"):
+      raiden_unscan.unscan_layers(nnx.state(UnscannedModel(), nnx.Param), num_layers=_NUM_LAYERS)
+
+  def test_wrong_layer_container_raises(self):
+    with self.assertRaisesRegex(ValueError, "found no scanned 'blocks' entries"):
+      raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS, layer_container="blocks")
+
+  def test_custom_scan_axis(self):
+    """`param_scan_axis` is configurable; axis 0 must slice the leading dim."""
+    state = {"layers": {"kernel": jnp.arange(_NUM_LAYERS * _OUT, dtype=jnp.float32).reshape(_NUM_LAYERS, _OUT)}}
+    unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS, scan_axis=0)
+    for i in range(_NUM_LAYERS):
+      got = unscanned[f"layers_{i}"]["kernel"]
+      np.testing.assert_array_equal(np.asarray(_unwrap(got)), np.arange(i * _OUT, (i + 1) * _OUT))
+
+  def test_plain_dict_input_is_accepted(self):
+    """The trainer passes an `nnx.State`, but the signature documents plain dicts too."""
+    state = {
+        "layers": {"kernel": jnp.zeros((_IN, _NUM_LAYERS, _OUT))},
+        "embed": jnp.zeros((_VOCAB, _IN)),
+    }
+    unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS)
+    self.assertEqual(sorted(unscanned.keys()), ["embed"] + [f"layers_{i}" for i in range(_NUM_LAYERS)])
+
+  def test_total_element_count_is_preserved(self):
+    """Unscanning reshapes; it must not drop or duplicate any weight."""
+    state = self._scanned_state()
+    before = sum(int(np.size(x)) for x in jax.tree_util.tree_leaves(state.to_pure_dict()))
+    unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS)
+    after = sum(int(np.size(np.asarray(_unwrap(x)))) for x in jax.tree_util.tree_leaves(unscanned))
+    self.assertEqual(after, before)
+
+
+if __name__ == "__main__":
+  absltest.main()

From 009af47f6e178ac5c6f0b856376edb3759b2223a Mon Sep 17 00:00:00 2001
From: Yixuan Wang <yixuannwang@google.com>
Date: Wed, 2 Sep 2026 04:06:15 +0000
Subject: [PATCH 2/7] Add target-free weight conversion and optimize Raiden
 weight sync memory

- Support target-free key synthesis and unrolling in WeightConverter / MaxTextToMaxTextConverter for hybrid-cycle and MoE layers
- Add MoE padding utility for TPU GMM_v2 kernel alignment
- Cache staged weight sync metadata in MaxTextTrainingEngine and clean up host memory with gc and malloc_trim
- Add comprehensive TargetFreeConversionTest unit test suite
---
 src/maxtext/common/gcloud_stub.py             |   1 +
 src/maxtext/configs/types.py                  |   4 +
 src/maxtext/integration/vllm/convert_utils.py |  27 +-
 .../vllm/maxtext_vllm_adapter/adapter.py      |  21 +-
 .../integration/vllm/maxtext_vllm_rollout.py  |   9 -
 src/maxtext/integration/vllm/moe_padding.py   |  63 +++
 .../vllm/torchax_converter/__init__.py        |  12 +
 .../integration/vllm/weight_converter.py      | 433 ++++++++++++++++--
 src/maxtext/training_engine/maxtext_engine.py | 230 +++++++---
 .../unit/weight_converter_test.py             | 136 ++++++
 10 files changed, 792 insertions(+), 144 deletions(-)
 create mode 100644 src/maxtext/integration/vllm/moe_padding.py

diff --git a/src/maxtext/common/gcloud_stub.py b/src/maxtext/common/gcloud_stub.py
index c87ba4123c..044094206a 100644
--- a/src/maxtext/common/gcloud_stub.py
+++ b/src/maxtext/common/gcloud_stub.py
@@ -330,6 +330,7 @@ def _import():
       _goodput_stubs,
       label="ml_goodput_measurement",
       stub_if_decoupled=False,
+      stub_on_error_when_not_decoupled=True,
   )
 
 
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 353656c665..3eb0c862f8 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -2704,6 +2704,10 @@ class VLLM(BaseModel):
           "the legacy transfer_state_directly / transfer_state_with_mappings paths."
       ),
   )
+  rollout_backend: str = Field(
+      "maxtext",
+      description="Rollout backend for trainer-side weight converter ('maxtext' or 'vllm_torchax').",
+  )
   weight_sync_debug: bool = Field(
       False,
       description=(
diff --git a/src/maxtext/integration/vllm/convert_utils.py b/src/maxtext/integration/vllm/convert_utils.py
index d349b40792..e75d154655 100644
--- a/src/maxtext/integration/vllm/convert_utils.py
+++ b/src/maxtext/integration/vllm/convert_utils.py
@@ -160,9 +160,22 @@ 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 _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:
+  if isinstance(tgt_dtype, str):
+    if tgt_dtype in ("bfloat16", "bf16"):
+      tgt_dtype = jnp.bfloat16
+    elif tgt_dtype in ("float32", "fp32"):
+      tgt_dtype = jnp.float32
+    else:
+      tgt_dtype = jnp.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 +184,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 +423,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:
@@ -604,6 +620,11 @@ def _bulk_align_and_unstack(
     A tuple of `num_layers` per-layer arrays at the per-layer target shape.
   """
   per_layer_shape = per_layer_tgt_val.shape
+  if isinstance(arr, jax.ShapeDtypeStruct) or isinstance(per_layer_tgt_val, jax.ShapeDtypeStruct):
+    num_layers = arr.shape[scan_axis]
+    tgt_dtype = getattr(per_layer_tgt_val, "dtype", getattr(arr, "dtype", jnp.float32))
+    return tuple(jax.ShapeDtypeStruct(per_layer_shape, tgt_dtype) for _ in range(num_layers))
+
   scanned_tgt_shape = per_layer_shape[:scan_axis] + (arr.shape[scan_axis],) + per_layer_shape[scan_axis:]
   scanned_tgt_sharding = _scanned_sharding_from_per_layer(getattr(per_layer_tgt_val, "sharding", None), scan_axis)
 
diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py
index 3b0a3c8f05..e633049335 100644
--- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py
+++ b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py
@@ -48,19 +48,7 @@ class AttentionMetadata:
 _HYBRID_LAYER_IMBALANCE_THRESHOLD = 1.5
 
 
-def next_power_of_two(x: int) -> int:
-  """Finds the smallest power of 2 >= x using bit manipulation.
-
-  Args:
-    x: The input number (should be an integer).
-
-  Returns:
-    The smallest integer power of 2 that is >= x.
-  """
-  assert x > 0
-  if x == 1:
-    return 1
-  return 1 << (x - 1).bit_length()
+from maxtext.integration.vllm.moe_padding import compute_padded_moe_mlp_dim, next_power_of_two
 
 
 def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters:
@@ -166,11 +154,8 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters
   # The GMM_v2 kernel requires the MLP dimension per expert to be at least 2x the number of TPU lanes
   # to ensure efficient execution. See the validate_inputs() method in the following file for more details:
   # https://github.com/vllm-project/tpu-inference/blob/main/tpu_inference/kernels/megablox/gmm_v2.py
-  if hidden_size is not None and (hidden_size // moe_mlp_tp_size) % (2 * num_lanes) != 0:
-    padded_hidden_size = next_power_of_two(hidden_size)
-    while (padded_hidden_size // moe_mlp_tp_size) < (2 * num_lanes):
-      padded_hidden_size = next_power_of_two(padded_hidden_size + 1)
-
+  padded_hidden_size = compute_padded_moe_mlp_dim(hidden_size, moe_mlp_tp_size, num_lanes)
+  if padded_hidden_size is not None and padded_hidden_size != hidden_size:
     # This inflates every expert weight, so it is a real memory/FLOP cost rather than a
     # cosmetic reshape: at moe_mlp_tp_size=4 a 512-wide MoE is padded to 1024 (2x the MoE
     # weights), and at moe_mlp_tp_size=8 to 2048 (4x). Log it at WARNING so it is visible
diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py
index 1cd7853f96..18e8d56f72 100644
--- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py
+++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py
@@ -464,14 +464,11 @@ 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
@@ -495,11 +492,6 @@ def update_params(
         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)
     try:
       return super().update_params(updated_weights, filter_types)
@@ -742,7 +734,6 @@ def __init__(
             additional_config=rollout_additional_config,
             sampling_kwargs=rollout_config.rollout_vllm_sampling_kwargs,
         ),
-        converter=converter,
         direct_maxtext_sync=direct_maxtext_sync,
         scan_axis=getattr(maxtext_config, "param_scan_axis", 1),
         layer_pattern_length=getattr(maxtext_config, "inhomogeneous_layer_cycle_interval", None),
diff --git a/src/maxtext/integration/vllm/moe_padding.py b/src/maxtext/integration/vllm/moe_padding.py
new file mode 100644
index 0000000000..3ad3edc95a
--- /dev/null
+++ b/src/maxtext/integration/vllm/moe_padding.py
@@ -0,0 +1,63 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MoE padding utilities for TPU GMM_v2 kernel alignment."""
+
+from typing import Optional
+
+
+def next_power_of_two(x: int) -> int:
+  """Finds the smallest power of 2 >= x using bit manipulation.
+
+  Args:
+    x: The input number (should be an integer > 0).
+
+  Returns:
+    The smallest integer power of 2 that is >= x.
+  """
+  assert x > 0
+  if x == 1:
+    return 1
+  return 1 << (x - 1).bit_length()
+
+
+def compute_padded_moe_mlp_dim(
+    hidden_size: Optional[int],
+    moe_mlp_tp_size: int,
+    num_lanes: int,
+) -> Optional[int]:
+  """Computes padded MoE intermediate size for GMM_v2 kernel requirements.
+
+  The GMM_v2 kernel requires the MLP dimension per expert to be at least 2x the
+  number of TPU lanes (e.g. 2 * 128 = 256) to ensure efficient execution.
+
+  Args:
+    hidden_size: Unpadded MoE intermediate size (e.g. moe_intermediate_size /
+      base_moe_mlp_dim).
+    moe_mlp_tp_size: TP size across MLP dimensions (e.g. tp * attn_dp).
+    num_lanes: Number of TPU lanes (typically 128 for TPU v5p/v6e).
+
+  Returns:
+    Padded hidden size, or hidden_size if no padding is required / hidden_size is None.
+  """
+  if hidden_size is None or moe_mlp_tp_size <= 0 or num_lanes <= 0:
+    return hidden_size
+
+  if (hidden_size // moe_mlp_tp_size) % (2 * num_lanes) != 0:
+    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)
+    return padded_hidden_size
+
+  return hidden_size
diff --git a/src/maxtext/integration/vllm/torchax_converter/__init__.py b/src/maxtext/integration/vllm/torchax_converter/__init__.py
index f3582c0090..1ec0362971 100644
--- a/src/maxtext/integration/vllm/torchax_converter/__init__.py
+++ b/src/maxtext/integration/vllm/torchax_converter/__init__.py
@@ -11,3 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from maxtext.integration.vllm.torchax_converter.base import BaseMaxTextToVLLMConverter
+from maxtext.integration.vllm.torchax_converter.qwen35_moe import Qwen35MaxTextToVLLMConverter
+from maxtext.integration.vllm.torchax_converter.qwen3_moe import Qwen3MaxTextToVLLMConverter
+from maxtext.integration.vllm.torchax_converter.gemma4_moe import Gemma4MaxTextToVLLMConverter
+
+__all__ = [
+    "BaseMaxTextToVLLMConverter",
+    "Qwen35MaxTextToVLLMConverter",
+    "Qwen3MaxTextToVLLMConverter",
+    "Gemma4MaxTextToVLLMConverter",
+]
diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py
index c8e79ff535..3e828e6cff 100644
--- a/src/maxtext/integration/vllm/weight_converter.py
+++ b/src/maxtext/integration/vllm/weight_converter.py
@@ -16,11 +16,13 @@
 
 import abc
 import dataclasses
+import gc
 import logging
+import os
 import re
 import jax
 import jax.numpy as jnp
-import gc
+import numpy as np
 from typing import List, Union, Any, Dict, Optional, Mapping, Tuple
 from flax import traverse_util, nnx
 from maxtext.integration.vllm.convert_utils import (
@@ -35,6 +37,16 @@
     _sharding_summary,
 )
 
+_MOE_MLP_WEIGHTS = frozenset({"wi_0", "wi_1", "wo", "wi"})
+
+
+def _malloc_trim() -> None:
+  try:
+    import ctypes  # pylint: disable=g-import-not-at-top
+    ctypes.CDLL("libc.so.6").malloc_trim(0)
+  except Exception:
+    pass
+
 
 # ==========================================
 # 1. Operations
@@ -276,12 +288,15 @@ def __init__(
       num_kv_heads: Optional[int] = None,
       head_dim: Optional[int] = None,
       config: Any = None,
-      # Defaults to MoEFusedLayout.PER_SHARD_INTERLEAVE; resolved in the body
-      # because MoEFusedLayout is defined further down this module.
+      trainer_config: Any = None,
+      rollout_backend: str = "maxtext",
       moe_fused_layout: Optional[str] = None,
       allow_unused_source_keys: Tuple[str, ...] = (),
       debug: bool = False,
+      prefuse_moe_weights: Optional[bool] = None,
+      target_dtype: Optional[Any] = None,
   ):
+    config = trainer_config if config is None else config
     if rules is not None and not rules:
       raise ValueError(
           "WeightConverter(rules=[]) would convert nothing and leave the "
@@ -295,9 +310,10 @@ def __init__(
     # Read by the rollout engine to decide whether to trace the reshard
     # step that runs after conversion.
     self.debug = debug
+    self.rollout_backend = rollout_backend
 
     self._direct: Optional["MaxTextToMaxTextConverter"] = None
-    if rules is None:
+    if rollout_backend == "maxtext" and rules is None:
       if config is None:
         raise ValueError(
             "WeightConverter(rules=None) needs `config` to derive the "
@@ -309,20 +325,29 @@ def __init__(
           moe_fused_layout=(moe_fused_layout or MoEFusedLayout.PER_SHARD_INTERLEAVE),
           allow_unused_source_keys=allow_unused_source_keys,
           debug=debug,
+          prefuse_moe_weights=prefuse_moe_weights,
+          target_dtype=target_dtype,
       )
       logging.info("WeightConverter: direct MaxText-to-MaxText mode (debug=%s).", debug)
     else:
+      if self.rules is None and config is not None:
+        model_name = getattr(config, "model_name", "")
+        if model_name in MODEL_TO_CONVERSION_RULES and MODEL_TO_CONVERSION_RULES[model_name] is not None:
+          self.rules = MODEL_TO_CONVERSION_RULES[model_name]
       logging.info(
           "WeightConverter: torchax rule mode (tp=%d, %d rules).",
           self.tp,
-          len(rules),
+          len(self.rules) if self.rules else 0,
       )
 
   def convert(self, src_pytree: Any, target_state: Any = None) -> Any:
     """Converts source weights pytree into target format using rules or direct converter."""
-    if self.rules is None:
+    if self.rollout_backend == "maxtext" and self.rules is None:
       return self._direct.convert(src_pytree, target_state=target_state)
 
+    if self.rules is None:
+      raise ValueError("WeightConverter in torchax mode requires conversion rules.")
+
     flat_src = traverse_util.flatten_dict(_to_pure_dict(src_pytree), sep=".")
     gc.collect()
 
@@ -550,10 +575,22 @@ def _group_plan(plan: List[_PlanEntry]) -> List[_PlanGroup]:
   return groups
 
 
-class ConversionPlanError(ValueError):
+class WeightConverterError(Exception):
+  """Base class for all weight converter exceptions."""
+
+
+class ConversionPlanError(WeightConverterError, ValueError):
   """Raised when the source and target trees cannot be fully reconciled."""
 
 
+class UnscanShapeMismatch(WeightConverterError, ValueError):
+  """Raised when an unscan shape mismatch is detected."""
+
+
+class WeightSyncBindError(WeightConverterError, RuntimeError):
+  """Raised when parameter binding to Raiden synchronizer fails."""
+
+
 def _is_non_weight_path(key_tuple: Tuple[Any, ...]) -> bool:
   return any(isinstance(part, str) and part.lstrip("_").startswith(_NON_WEIGHT_PATH_PREFIXES) for part in key_tuple)
 
@@ -657,11 +694,20 @@ def __init__(
       moe_fused_layout: str = MoEFusedLayout.PER_SHARD_INTERLEAVE,
       allow_unused_source_keys: Tuple[str, ...] = (),
       debug: bool = False,
+      prefuse_moe_weights: Optional[bool] = None,
+      target_dtype: Optional[Any] = None,
   ):
     self.config = config
     self.moe_fused_layout = moe_fused_layout
     self.allow_unused_source_keys = allow_unused_source_keys
     self.debug = debug
+    self.prefuse_moe_weights = (
+        prefuse_moe_weights
+        if prefuse_moe_weights is not None
+        else getattr(config, "prefuse_moe_weights", False)
+    )
+    self.padded_base_moe_mlp_dim = getattr(config, "padded_base_moe_mlp_dim", None)
+    self.target_dtype = target_dtype if target_dtype is not None else getattr(config, "weight_dtype", None)
 
     self.cycle = int(getattr(config, "inhomogeneous_layer_cycle_interval", 1) or 1)
     self.num_decoder_layers = int(config.num_decoder_layers)
@@ -678,14 +724,28 @@ 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):
+    if self.target_dtype is None:
+      return None
+    if isinstance(self.target_dtype, str):
+      if self.target_dtype in ("bfloat16", "bf16"):
+        return jnp.bfloat16
+      if self.target_dtype in ("float32", "fp32"):
+        return jnp.float32
+      return jnp.dtype(self.target_dtype)
+    return self.target_dtype
+
   # -------------------------------------------------------------- #
   # Plan construction
   # -------------------------------------------------------------- #
@@ -704,6 +764,79 @@ 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:
+        # Homogeneous: ("decoder", "layers", "self_attention", "query", "kernel")
+        is_wi_0 = bool(rest and rest[-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 i in range(self.num_decoder_layers):
+            tgt_key = prefix + (f"layers_{i}",) + rest[:-1] + ("wi",)
+            plan.append(_PlanEntry(tgt_key, (src_key, wi_1_key), i, "fuse_moe"))
+        elif self.prefuse_moe_weights and rest and rest[-1] == "wi_1" and (src_key[:-1] + ("wi_0",) in src_flat):
+          continue
+        else:
+          for i in range(self.num_decoder_layers):
+            tgt_key = prefix + (f"layers_{i}",) + rest
+            plan.append(_PlanEntry(tgt_key, (src_key,), i, "slice"))
+
+      else:
+        # Inhomogeneous hybrid cycle: ("decoder", "layers", "layer_0", "input_layernorm", "scale")
+        slot_token = rest[0]
+        if isinstance(slot_token, str) and slot_token.startswith("layer_"):
+          slot = int(slot_token[6:])
+        elif isinstance(slot_token, str) and slot_token.isdigit():
+          slot = int(slot_token)
+        elif isinstance(slot_token, int):
+          slot = slot_token
+        else:
+          raise ConversionPlanError(f"Unexpected slot token {slot_token!r} in key {src_key}")
+
+        suffix = rest[1:]
+        is_wi_0 = bool(suffix and suffix[-1] == "wi_0")
+        wi_1_key = src_key[:-1] + ("wi_1",) if is_wi_0 else None
+        fuse_moe = self.prefuse_moe_weights and is_wi_0 and (wi_1_key in src_flat)
+
+        if fuse_moe:
+          consumed_wi_1.add(wi_1_key)
+          for b in range(self.num_blocks):
+            global_idx = b * self.cycle + slot
+            tgt_key = prefix + (f"layers_{global_idx}",) + suffix[:-1] + ("wi",)
+            plan.append(_PlanEntry(tgt_key, (src_key, wi_1_key), b, "fuse_moe"))
+        elif self.prefuse_moe_weights and suffix and suffix[-1] == "wi_1" and (src_key[:-1] + ("wi_0",) in src_flat):
+          continue
+        else:
+          for b in range(self.num_blocks):
+            global_idx = b * self.cycle + slot
+            tgt_key = prefix + (f"layers_{global_idx}",) + suffix
+            plan.append(_PlanEntry(tgt_key, (src_key,), b, "slice"))
+
+    return plan
+
   def _build_plan(
       self,
       src_flat: Mapping[Tuple[Any, ...], Any],
@@ -764,13 +897,7 @@ def _build_plan(
     return plan
 
   def _validate_plan(self, src_flat, tgt_flat, unmatched, consumed) -> None:
-    """Fails loudly rather than leaving rollout weights at their dummy values.
-
-    vLLM boots with `load_format="dummy"`, so a target leaf we never write
-    keeps *random* weights. That produces a quietly wrong reward curve
-    instead of an error, which is far more expensive to debug than a crash
-    at startup.
-    """
+    """Fails loudly rather than leaving rollout weights at their dummy values."""
     if unmatched:
       shown = "\n  ".join(".".join(map(str, k)) for k in sorted(unmatched)[:40])
       raise ConversionPlanError(
@@ -804,17 +931,9 @@ def _validate_plan(self, src_flat, tgt_flat, unmatched, consumed) -> None:
   # Plan execution
   # -------------------------------------------------------------- #
   def _fuse_moe_bulk(self, wi_0, wi_1, tgt_val, key_path: str):
-    """Fuses the *scanned* gate/up kernels, returning one array per block.
-
-    `wi_0`/`wi_1` still carry `num_blocks` at `scan_axis`; `tgt_val` is a
-    single per-layer target leaf, supplying the fused shape and sharding
-    that every block in this group shares.
-    """
+    """Fuses the *scanned* gate/up kernels, returning one array per block."""
     tgt_shape = tgt_val.shape
-    # MaxText stores MoE kernels as (experts, in_dim, intermediate); the
-    # gate/up fusion always doubles the trailing intermediate axis.
     tgt_fused_axis = len(tgt_shape) - 1
-    # Same logical axis, shifted by the scan dim the trainer inserted.
     scan_fused_axis = tgt_fused_axis if tgt_fused_axis < self.scan_axis else tgt_fused_axis + 1
 
     if self.moe_fused_layout == MoEFusedLayout.PER_SHARD_INTERLEAVE:
@@ -839,14 +958,208 @@ def _fuse_moe_bulk(self, wi_0, wi_1, tgt_val, key_path: str):
 
     raise ConversionPlanError(f"Unknown moe_fused_layout: {self.moe_fused_layout!r}")
 
-  def _execute_group(self, group: _PlanGroup, src_flat, tgt_flat):
-    """Produces every target leaf in `group`. Returns (target_key, array) pairs.
+  def _slice_bulk_target_free(self, val: Any, path: str):
+    last_key = path.split(".")[-1]
+    if isinstance(val, jax.ShapeDtypeStruct):
+      unrolled_shape = list(val.shape[: self.scan_axis] + val.shape[self.scan_axis + 1 :])
+      if last_key in _MOE_MLP_WEIGHTS and self.padded_base_moe_mlp_dim is not None:
+        if last_key == "wo":
+          if self.padded_base_moe_mlp_dim > unrolled_shape[1]:
+            unrolled_shape[1] = self.padded_base_moe_mlp_dim
+        elif last_key in ("wi_0", "wi_1", "wi"):
+          if self.padded_base_moe_mlp_dim > unrolled_shape[-1]:
+            unrolled_shape[-1] = self.padded_base_moe_mlp_dim
+      return tuple(jax.ShapeDtypeStruct(tuple(unrolled_shape), val.dtype) for _ in range(val.shape[self.scan_axis]))
+
+    if last_key in _MOE_MLP_WEIGHTS and self.padded_base_moe_mlp_dim is not None:
+      if last_key == "wo":
+        intermediate_axis = 2
+        if self.padded_base_moe_mlp_dim > val.shape[intermediate_axis]:
+          pad_amount = self.padded_base_moe_mlp_dim - val.shape[intermediate_axis]
+          pad_spec = [(0, 0)] * val.ndim
+          pad_spec[intermediate_axis] = (0, pad_amount)
+          val = jnp.pad(val, pad_spec)
+      elif last_key in ("wi_0", "wi_1"):
+        intermediate_axis = len(val.shape) - 1
+        if self.padded_base_moe_mlp_dim > val.shape[intermediate_axis]:
+          pad_amount = self.padded_base_moe_mlp_dim - val.shape[intermediate_axis]
+          pad_spec = [(0, 0)] * val.ndim
+          pad_spec[intermediate_axis] = (0, pad_amount)
+          val = jnp.pad(val, pad_spec)
+
+    return _jit_unstack(val, self.scan_axis)
+
+  def _fuse_moe_bulk_target_free(self, wi_0: Any, wi_1: Any, path: str):
+    unpadded_dim = wi_0.shape[-1]
+    target_intermediate = (
+        self.padded_base_moe_mlp_dim
+        if (self.padded_base_moe_mlp_dim is not None and self.padded_base_moe_mlp_dim > unpadded_dim)
+        else unpadded_dim
+    )
+    if isinstance(wi_0, jax.ShapeDtypeStruct):
+      fused_shape = (wi_0.shape[0], wi_0.shape[2], 2 * target_intermediate)
+      return tuple(jax.ShapeDtypeStruct(fused_shape, wi_0.dtype) for _ in range(wi_0.shape[self.scan_axis]))
 
-    The scanned source is cast, aligned and fused *once*; the per-layer
-    arrays are then read out of a single unstack. Every target in a group
-    shares a shape and sharding by construction, so the first one is a
-    sound stand-in for all of them.
-    """
+    tgt_shape = (wi_0.shape[0], wi_0.shape[2], 2 * target_intermediate)
+    tgt_fused_axis = len(tgt_shape) - 1
+    scan_fused_axis = tgt_fused_axis if tgt_fused_axis < self.scan_axis else tgt_fused_axis + 1
+
+    if self.moe_fused_layout == MoEFusedLayout.PER_SHARD_INTERLEAVE:
+      n_shards = _get_n_shards(wi_0, scan_fused_axis)
+      return _fuse_and_unstack_moe(
+          wi_0,
+          wi_1,
+          self.scan_axis,
+          n_shards,
+          tgt_shape,
+          scan_fused_axis,
+          tgt_fused_axis,
+      )
+
+    if self.moe_fused_layout == MoEFusedLayout.CONCAT:
+      if target_intermediate > unpadded_dim:
+        pad_spec = [(0, 0)] * wi_0.ndim
+        pad_spec[-1] = (0, target_intermediate - unpadded_dim)
+        wi_0 = jnp.pad(wi_0, pad_spec)
+        wi_1 = jnp.pad(wi_1, pad_spec)
+      fused = jnp.concatenate([wi_0, wi_1], axis=scan_fused_axis)
+      return _jit_unstack(fused, self.scan_axis)
+
+    raise ConversionPlanError(f"Unknown moe_fused_layout: {self.moe_fused_layout!r}")
+
+  def _execute_group_target_free(self, group: _PlanGroup, src_flat):
+    path = group.source_path
+    target_dtype = self._resolve_target_dtype()
+    is_pathways = bool("proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get("JAX_BACKEND_TARGET"))
+
+    if group.op == "identity":
+      raw_val = src_flat[group.source_keys[0]]
+      if isinstance(raw_val, jax.ShapeDtypeStruct):
+        val = _apply_dtype_cast(raw_val, target_dtype, path)
+        return [(tgt_key, val) for _, tgt_key in group.targets]
+      if is_pathways:
+        np_val = np.asarray(jax.device_get(raw_val))
+        if target_dtype is not None:
+          np_val = np_val.astype(target_dtype)
+        cpu = jax.local_devices(backend="cpu")[0]
+        val = jax.device_put(np_val, cpu)
+        del np_val
+        return [(tgt_key, val) for _, tgt_key in group.targets]
+      val = _apply_dtype_cast(raw_val, target_dtype, path)
+      return [(tgt_key, val) for _, tgt_key in group.targets]
+
+    if any(idx is None for idx, _ in group.targets):
+      raise ConversionPlanError(
+          f"Plan group for {path} has op={group.op!r} but a target with no "
+          "scan index; only 'identity' targets may omit one."
+      )
+
+    if group.op == "fuse_moe":
+      raw_0 = src_flat[group.source_keys[0]]
+      raw_1 = src_flat[group.source_keys[1]]
+      if isinstance(raw_0, jax.ShapeDtypeStruct):
+        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]
+
+      if is_pathways:
+        wi_0 = np.asarray(jax.device_get(raw_0))
+        wi_1 = np.asarray(jax.device_get(raw_1))
+        if target_dtype is not None:
+          wi_0 = wi_0.astype(target_dtype)
+          wi_1 = wi_1.astype(target_dtype)
+        self._check_scan_axis(wi_0, path)
+        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
+        )
+        tgt_shape = (wi_0.shape[0], wi_0.shape[2], 2 * target_intermediate)
+        tgt_fused_axis = len(tgt_shape) - 1
+        scan_fused_axis = tgt_fused_axis if tgt_fused_axis < self.scan_axis else tgt_fused_axis + 1
+
+        if self.moe_fused_layout == MoEFusedLayout.PER_SHARD_INTERLEAVE:
+          n_shards = 1
+          target_half_dim = target_intermediate
+          current_total_size = wi_0.shape[scan_fused_axis]
+          chunk_size = current_total_size // n_shards
+          target_chunk_size = target_half_dim // n_shards
+          pad_amount = target_chunk_size - chunk_size
+          if pad_amount > 0:
+            pad_spec = [(0, 0)] * wi_0.ndim
+            pad_spec[scan_fused_axis] = (0, pad_amount)
+            wi_0 = np.pad(wi_0, pad_spec)
+            wi_1 = np.pad(wi_1, pad_spec)
+          fused = np.concatenate([wi_0, wi_1], axis=scan_fused_axis)
+        elif 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 = np.pad(wi_0, pad_spec)
+            wi_1 = np.pad(wi_1, pad_spec)
+          fused = np.concatenate([wi_0, wi_1], axis=scan_fused_axis)
+        else:
+          raise ConversionPlanError(f"Unknown moe_fused_layout: {self.moe_fused_layout!r}")
+
+        cpu = jax.local_devices(backend="cpu")[0]
+        per_block = tuple(
+            jax.device_put(np.ascontiguousarray(fused.take(indices=i, axis=self.scan_axis)), cpu)
+            for i in range(self.num_blocks)
+        )
+        del fused, wi_0, wi_1
+        return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets]
+
+      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]]
+    if isinstance(raw_val, jax.ShapeDtypeStruct):
+      val = _apply_dtype_cast(raw_val, target_dtype, path)
+      self._check_scan_axis(val, path)
+      per_block = self._slice_bulk_target_free(val, path)
+      return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets]
+
+    if is_pathways:
+      np_val = np.asarray(jax.device_get(raw_val))
+      if target_dtype is not None:
+        np_val = np_val.astype(target_dtype)
+      self._check_scan_axis(np_val, path)
+      last_key = path.split(".")[-1]
+      if last_key in _MOE_MLP_WEIGHTS and self.padded_base_moe_mlp_dim is not None:
+        if last_key == "wo":
+          intermediate_axis = 2
+          if self.padded_base_moe_mlp_dim > np_val.shape[intermediate_axis]:
+            pad_amount = self.padded_base_moe_mlp_dim - np_val.shape[intermediate_axis]
+            pad_spec = [(0, 0)] * np_val.ndim
+            pad_spec[intermediate_axis] = (0, pad_amount)
+            np_val = np.pad(np_val, pad_spec)
+        elif last_key in ("wi_0", "wi_1", "wi"):
+          intermediate_axis = len(np_val.shape) - 1
+          if self.padded_base_moe_mlp_dim > np_val.shape[intermediate_axis]:
+            pad_amount = self.padded_base_moe_mlp_dim - np_val.shape[intermediate_axis]
+            pad_spec = [(0, 0)] * np_val.ndim
+            pad_spec[intermediate_axis] = (0, pad_amount)
+            np_val = np.pad(np_val, pad_spec)
+      cpu = jax.local_devices(backend="cpu")[0]
+      per_block = tuple(
+          jax.device_put(np.ascontiguousarray(np_val.take(indices=i, axis=self.scan_axis)), cpu)
+          for i in range(self.num_blocks)
+      )
+      del np_val
+      return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets]
+
+    val = _apply_dtype_cast(raw_val, target_dtype, path)
+    self._check_scan_axis(val, path)
+    per_block = self._slice_bulk_target_free(val, path)
+    return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets]
+
+  def _execute_group(self, group: _PlanGroup, src_flat, tgt_flat):
+    """Produces every target leaf in `group`. Returns (target_key, array) pairs."""
     first_tgt = tgt_flat[group.targets[0][1]]
     path = group.source_path
 
@@ -886,24 +1199,40 @@ 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.
     """
+    src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree))
+    src_flat, _ = _strip_root(src_flat, "base")
+
     if target_state is None:
-      raise ValueError(
-          "MaxTextToMaxTextConverter requires target_state to resolve the "
-          "rollout's parameter shapes, shardings and dtypes."
+      if self._plan is None:
+        self._plan = self._build_target_free_plan(src_flat)
+        self._groups = _group_plan(self._plan)
+
+      result: Dict[Tuple[Any, ...], Any] = {}
+      for group in self._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:
+          result[tgt_key] = out
+        del outs
+
+      del src_flat
+      gc.collect()
+      nested = traverse_util.unflatten_dict(result)
+      del result
+      gc.collect()
+      _malloc_trim()
+      return jax.tree_util.tree_map(
+          lambda x: nnx.Param(x) if not isinstance(x, (nnx.Param, nnx.Variable)) else x,
+          nested,
       )
 
     # Read variable types before purifying to plain arrays loses them.
     skip_paths = _non_param_paths(target_state)
-
-    src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree))
     tgt_flat = traverse_util.flatten_dict(_to_pure_dict(target_state))
 
-    # The trainer wraps the model in TunixMaxTextAdapter ("base"); the
-    # rollout may nest it under one or more "model" levels. Strip both so
-    # the plan is expressed in a single coordinate system, then re-wrap.
-    src_flat, _ = _strip_root(src_flat, "base")
     tgt_flat, tgt_root = _strip_root(tgt_flat, "model")
     if tgt_root:
       depth = len(tgt_root)
@@ -930,17 +1259,13 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]:
       if self.debug:
         for k in group.source_keys:
           logging.info(
-              "weight_sync_debug: op=%s source=%s (%d targets) | src %s " "| tgt %s",
+              "weight_sync_debug: op=%s source=%s (%d targets) | src %s | tgt %s",
               group.op,
               ".".join(map(str, k)),
               len(group.targets),
               _sharding_summary(src_flat[k]),
               _sharding_summary(tgt_flat[group.targets[0][1]]),
           )
-        # JAX dispatch is asynchronous, so without a barrier a device
-        # failure surfaces at an arbitrary later point and the traceback
-        # names the wrong parameter. Pay the serialization to find out
-        # which source parameter is actually at fault.
         try:
           outs = self._execute_group(group, src_flat, tgt_flat)
           jax.block_until_ready([out for _, out in outs])
@@ -965,18 +1290,30 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]:
       else:
         outs = self._execute_group(group, src_flat, tgt_flat)
 
+      for k in group.source_keys:
+        src_flat.pop(k, None)
+
       for tgt_key, out in outs:
         tgt_val = tgt_flat[tgt_key]
-        if out.shape != tgt_val.shape:
+        if hasattr(out, "shape") and hasattr(tgt_val, "shape") and out.shape != tgt_val.shape:
           raise ConversionPlanError(
               f"Shape mismatch after conversion for "
               f"{'.'.join(map(str, tgt_key))}: produced {out.shape}, "
               f"rollout expects {tgt_val.shape}."
           )
         result[tgt_root + tgt_key] = out
+      del outs
 
+    del src_flat, tgt_flat
+    gc.collect()
+    nested = traverse_util.unflatten_dict(result)
+    del result
     gc.collect()
-    return traverse_util.unflatten_dict(result)
+    _malloc_trim()
+    return jax.tree_util.tree_map(
+        lambda x: nnx.Param(x) if not isinstance(x, (nnx.Param, nnx.Variable)) else x,
+        nested,
+    )
 
 
 def _rekey_to_target(flat_dotted: Dict[str, Any], target_state: Any) -> Dict[str, Any]:
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 4cd6c77014..9b6aefd316 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -23,6 +23,7 @@
 from collections.abc import Callable, Mapping
 import contextlib
 import dataclasses
+import gc
 import os
 from typing import Any
 
@@ -30,6 +31,8 @@
 from flax import nnx
 from flax import struct
 from flax.linen import partitioning as nn_partitioning
+from flax.traverse_util import flatten_dict
+from flax.traverse_util import unflatten_dict
 import jax
 import jax.numpy as jnp
 from jax.typing import ArrayLike  # pylint: disable=g-importing-member
@@ -66,6 +69,14 @@
 )
 
 
+def _malloc_trim() -> None:
+  try:
+    import ctypes  # pylint: disable=g-import-not-at-top
+    ctypes.CDLL("libc.so.6").malloc_trim(0)
+  except Exception:
+    pass
+
+
 def _is_jax_dynamic(value: Any) -> bool:
   """Returns True if `value` can cross a `jax.jit` boundary as a traced argument.
 
@@ -405,9 +416,26 @@ def __init__(
         checkpoint_dir=self._config.checkpoint_dir,
         config=self._config,
     )
-    self._metrics_recorder = metrics_module.MetricsRecorder()
     self._throttler = inflight_throttler.InflightThrottler(config=self._config)
-    self._raiden_sync: Any = None
+    self._raiden_syncs: Any = None
+    self._last_staged_step: Optional[int] = None
+    self._staged_metadata: Any = None
+    self._use_weight_converter = bool(
+        getattr(self._config, "use_weight_converter", False)
+        or os.environ.get("USE_WEIGHT_CONVERTER", "0").lower() in ("1", "true", "yes")
+    )
+    self._rollout_backend = (
+        getattr(self._config, "rollout_backend", "maxtext") or os.environ.get("ROLLOUT_BACKEND", "maxtext")
+    )
+    if self._use_weight_converter:
+      from maxtext.integration.vllm.weight_converter import WeightConverter  # pylint: disable=g-import-not-at-top,import-outside-toplevel
+      self._weight_converter = WeightConverter(
+          config=self._config,
+          rollout_backend=self._rollout_backend,
+          debug=getattr(self._config, "weight_sync_debug", False),
+      )
+    else:
+      self._weight_converter = None
 
   @property
   def model(self) -> Any:
@@ -1423,6 +1451,20 @@ def _get_trainable_params_state(self) -> Any:
       return nnx.state(model, nnx.Param)
     return self.model
 
+  def _split_into_chunks(self, nested_state: Any, num_chunks: int) -> list[Any]:
+    """Splits a nested param dict into `num_chunks` nested dicts of near-equal leaf count."""
+    if hasattr(nested_state, "to_pure_dict"):
+      pure_state = nested_state.to_pure_dict()
+    elif hasattr(nested_state, "to_dict"):
+      pure_state = nested_state.to_dict()
+    else:
+      pure_state = nested_state
+    flat = flatten_dict(pure_state)
+    chunk_flats = [{} for _ in range(num_chunks)]
+    for i, key in enumerate(flat):
+      chunk_flats[i % num_chunks][key] = flat[key]
+    return [unflatten_dict(cf) for cf in chunk_flats]
+
   def prepare_weight_sync(
       self,
       staging_transport: str = "raiden",
@@ -1452,81 +1494,130 @@ def prepare_weight_sync(
             " tunix build that ships it, or select a different staging_transport."
         ) from exc
 
+      if (
+          self._raiden_syncs is not None
+          and self._last_staged_step == self.train_step
+          and self._staged_metadata is not None
+      ):
+        logging.info(
+            "Trainer re-using staged weight sync for step %d (%d variables)",
+            self.train_step,
+            sum(len(m.variables) for m in self._staged_metadata),
+        )
+        return self._staged_metadata
+
+      if self._raiden_syncs is not None:
+        for sync in self._raiden_syncs:
+          sync.release_host_arrays()
+        gc.collect()
+        _malloc_trim()
+
       # 1. Drain all in-flight TPU computations to ensure weights are fully updated
       self._throttler.wait_for_all()
+      gc.collect()
 
       # 2. Extract clean trainable parameters
       params_state = self._get_trainable_params_state()
 
-      # 2a. The trainer keeps float32 master weights, but the rollout side
-      # (MaxTextForCausalLM under configs/inference/vllm.yml) loads/serves in
-      # bfloat16 -- Raiden's manifest preflight rejects a dtype/item_size
-      # mismatch, and binding mismatched-dtype buffers would be wrong anyway.
-      # Cast the synced copy down; the trainer's own params_state (used for
-      # the actual optimizer step) is untouched since this is a fresh tree.
-      params_state = jax.tree_util.tree_map(
-          lambda x: x.astype(jnp.bfloat16) if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating) else x,
-          params_state,
-      )
-
-      # 2b. The trainer runs scanned (scan_layers=True) for training speed, but
-      # the rollout side loads its MaxText model unscanned (MaxTextForCausalLM
-      # under configs/inference/vllm.yml has scan_layers=False). Raiden matches
-      # tensors by name, so unscan here -- on the trainer side only -- so the
-      # names/shapes we bind already match what the sampler reports.
-      if self._config.scan_layers:
-        params_state = raiden_unscan.unscan_layers(
+      if self._use_weight_converter:
+        if self._weight_converter is None:
+          from maxtext.integration.vllm.weight_converter import WeightConverter  # pylint: disable=g-import-not-at-top,import-outside-toplevel
+          self._weight_converter = WeightConverter(
+              config=self._config,
+              rollout_backend=self._rollout_backend,
+              debug=getattr(self._config, "weight_sync_debug", False),
+          )
+        params_state = self._weight_converter.convert(params_state)
+        gc.collect()
+      else:
+        # 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,
-            num_layers=self._config.num_decoder_layers,
-            scan_axis=self._config.param_scan_axis,
-        )
-
-      # 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",
-            worker_index=jax.process_index(),
-            auto_h2d=False,
-            parallelism=4,
-        )
+        # 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(
+              params_state,
+              num_layers=self._config.num_decoder_layers,
+              scan_axis=self._config.param_scan_axis,
+          )
 
-      self._raiden_sync.bind(params_state)
+      # 3. Bind parameters to the Raiden transport, one chunk at a time (see
+      # _split_into_chunks) -- construct the per-chunk synchronizers once,
+      # matching the persistent-instance-per-cycle pattern the rebind
+      # optimization (fewer stale holds) depends on.
+      num_chunks = max(1, int(os.environ.get("RAIDEN_WEIGHT_SYNC_CHUNKS", "1")))
+      if self._raiden_syncs is None:
+        # Under Pathways (JAX_PLATFORMS=proxy + JAX_BACKEND_TARGET set, same
+        # detection tunix's K8sJaxContext.initialize() uses), trainer params
+        # are proxy-backed and Raiden can't bind them in place -- host_stage
+        # pulls them to client host memory first. Direct-TPU trainers skip
+        # that extra copy since their params already live on TPU.
+        is_pathways = bool("proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get("JAX_BACKEND_TARGET"))
+        # worker_index must be unique per chunk (it seeds WorkUnitId's
+        # job_replica_id) -- otherwise every chunk's work unit collides under
+        # the same id in the handler's registry and only one survives
+        # registration.
+        self._raiden_syncs = [
+            raiden_synchronizer.RaidenSynchronizer(
+                job_name="trainer",
+                worker_index=jax.process_index() if num_chunks == 1 else (jax.process_index() * num_chunks + i + 1),
+                auto_h2d=False,
+                host_stage=is_pathways,
+                parallelism=4,
+            )
+            for i in range(num_chunks)
+        ]
+
+      chunks = self._split_into_chunks(params_state, num_chunks) if num_chunks > 1 else [params_state]
       del params_state
-
-      # 4. Initiate Device-to-Host transfer to stage weights for network transfer.
-      if is_pathways or self._raiden_sync.active:
-        self._raiden_sync.d2h()
+      gc.collect()
 
       verify_weights = os.environ.get("VERIFY_WEIGHTS", "").lower() == "true"
-      if verify_weights:
-        logging.info("Source weights checksums: %s", self._raiden_sync.checksums())
+      all_metadata = []
+      total_variables = 0
+      for chunk_idx, (sync, chunk_state) in enumerate(zip(self._raiden_syncs, chunks)):
+        sync.bind(chunk_state)
+        del chunk_state
+        gc.collect()
+
+        # 4. Initiate Device-to-Host transfer to stage this chunk for network
+        # transfer before moving on to the next chunk.
+        if sync.active:
+          sync.d2h()
+
+        if verify_weights:
+          logging.info("Source weights checksums (chunk %d): %s", chunk_idx, sync.checksums())
+
+        metadata = sync.work_unit_metadata()
+        total_variables += len(metadata.variables)
+        all_metadata.append(metadata)
+        sync.release_host_arrays()
+
+      gc.collect()
+      _malloc_trim()
 
-      metadata = self._raiden_sync.work_unit_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 variables across %d chunk(s) on mesh %s",
           self.train_step,
-          len(metadata.variables),
-          metadata.mesh_axes,
+          total_variables,
+          num_chunks,
+          all_metadata[0].mesh_axes if all_metadata else None,
       )
-      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
@@ -1535,16 +1626,23 @@ def prepare_weight_sync(
 
   def release_weight_sync(self, **kwargs: Any) -> Any:
     """Releases staged weight buffers after transfer completion."""
-    if self._raiden_sync:
-      logging.vlog(1, "Trainer Raiden metrics: %s", self._raiden_sync.metrics())
+    if self._raiden_syncs:
+      for sync in self._raiden_syncs:
+        logging.vlog(1, "Trainer Raiden metrics: %s", sync.metrics())
+        sync.release_host_arrays()
+    gc.collect()
+    _malloc_trim()
     return True
 
   def close(self) -> None:
     """Closes the trainer, writes buffered metrics and final checkpoint."""
-    if self._raiden_sync:
-      if hasattr(self._raiden_sync, "close"):
-        self._raiden_sync.close()
-      self._raiden_sync = None
+    if self._raiden_syncs:
+      for sync in self._raiden_syncs:
+        if hasattr(sync, "close"):
+          sync.close()
+      self._raiden_syncs = None
+    self._last_staged_step = None
+    self._staged_metadata = None
 
     self.save_checkpoint(metadata=None, force=True)
     self._checkpoint_manager.close()
diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py
index 1ecacf3fe0..6652f65790 100644
--- a/tests/post_training/unit/weight_converter_test.py
+++ b/tests/post_training/unit/weight_converter_test.py
@@ -439,5 +439,141 @@ def test_padded_moe_fusion_stays_on_the_source_mesh(self):
       )
 
 
+class TargetFreeConversionTest(unittest.TestCase):
+  """Comprehensive test suite for target-free key synthesis and execution."""
+
+  def test_case_0_raiden_unscan_fails_on_hybrid_cycle(self):
+    from maxtext.integration.tunix.weight_mapping import raiden_unscan
+
+    source = _source_tree(True)
+    with self.assertRaises(ValueError) as ctx:
+      raiden_unscan.unscan_layers(source, num_layers=NUM_LAYERS, scan_axis=SCAN_AXIS)
+    self.assertIn("expected axis 1 to be num_layers=8", str(ctx.exception))
+
+  def test_case_1_homogeneous_target_free_unroll(self):
+    cfg = _config(inhomogeneous_layer_cycle_interval=1, num_decoder_layers=4)
+    source = {
+        "base": {
+            "token_embedder": {"embedding": _arr(16, EMB)},
+            "decoder": {
+                "decoder_norm": {"scale": _arr(EMB)},
+                "layers": {
+                    "input_layernorm": {"scale": _arr(EMB, 4)},
+                    "self_attention": {"query": {"kernel": _arr(EMB, 4, 2, 4)}},
+                },
+            },
+        }
+    }
+    converter = WeightConverter(config=cfg, rollout_backend="maxtext")
+    out = converter.convert(source, target_state=None)
+    self.assertIn("token_embedder", out)
+    self.assertIn("decoder", out)
+    for i in range(4):
+      layer_key = f"layers_{i}"
+      self.assertIn(layer_key, out["decoder"])
+      scale = getattr(out["decoder"][layer_key]["input_layernorm"]["scale"], "value", out["decoder"][layer_key]["input_layernorm"]["scale"])
+      query = getattr(out["decoder"][layer_key]["self_attention"]["query"]["kernel"], "value", out["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)
+    src_layers = source["base"]["decoder"]["layers"]
+    for layer in range(NUM_LAYERS):
+      slot, block = layer % CYCLE, layer // CYCLE
+      got = getattr(out["decoder"][f"layers_{layer}"]["input_layernorm"]["scale"], "value", out["decoder"][f"layers_{layer}"]["input_layernorm"]["scale"])
+      want = jnp.take(src_layers[f"layer_{slot}"]["input_layernorm"]["scale"], block, axis=SCAN_AXIS)
+      np.testing.assert_array_equal(np.asarray(got), np.asarray(want))
+
+  def test_case_3_prefused_moe_target_free(self):
+    from maxtext.integration.vllm.moe_padding import compute_padded_moe_mlp_dim
+
+    # Verify helper across topologies
+    self.assertEqual(compute_padded_moe_mlp_dim(512, 2, 128), 512)
+    self.assertEqual(compute_padded_moe_mlp_dim(512, 4, 128), 1024)
+    self.assertEqual(compute_padded_moe_mlp_dim(512, 8, 128), 2048)
+
+    # Verify target-free prefused MoE with padded dim
+    padded_dim = 16
+    cfg = _config(padded_base_moe_mlp_dim=padded_dim, prefuse_moe_weights=True)
+    source = _source_tree(True)
+    converter = MaxTextToMaxTextConverter(cfg, prefuse_moe_weights=True)
+    out = converter.convert(source, target_state=None)
+    wi = getattr(out["decoder"]["layers_0"]["moe_block"]["wi"], "value", out["decoder"]["layers_0"]["moe_block"]["wi"])
+    wo = getattr(out["decoder"]["layers_0"]["moe_block"]["wo"], "value", out["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)
+    for leaf in jax.tree_util.tree_leaves(out):
+      val = getattr(leaf, "value", leaf)
+      self.assertIsInstance(val, jax.ShapeDtypeStruct)
+    wi = getattr(out["decoder"]["layers_0"]["moe_block"]["wi"], "value", out["decoder"]["layers_0"]["moe_block"]["wi"])
+    self.assertEqual(wi.shape, (EXPERTS, EMB, 32))
+
+  def test_case_5_host_memory_profiling(self):
+    import resource
+
+    cfg = _config()
+    source = _source_tree(True)
+    converter = WeightConverter(config=cfg, rollout_backend="maxtext")
+    before_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+    out = converter.convert(source, target_state=None)
+    after_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+    self.assertIsNotNone(out)
+
+  def test_case_6_parity_vs_raiden_unscan_on_homogeneous(self):
+    from maxtext.integration.tunix.weight_mapping import raiden_unscan
+
+    cfg = pytypes.SimpleNamespace(
+        num_decoder_layers=4,
+        inhomogeneous_layer_cycle_interval=1,
+        param_scan_axis=1,
+        weight_dtype=jnp.bfloat16,
+        prefuse_moe_weights=False,
+    )
+    raw_source = {
+        "token_embedder": {"embedding": _arr(16, EMB)},
+        "decoder": {
+            "decoder_norm": {"scale": _arr(EMB)},
+            "layers": {
+                "input_layernorm": {"scale": _arr(EMB, 4)},
+                "self_attention": {"query": {"kernel": _arr(EMB, 4, 2, 4)}},
+            },
+        },
+    }
+    bf16_source = jax.tree_util.tree_map(
+        lambda x: x.astype(jnp.bfloat16) if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating) else x,
+        raw_source,
+    )
+    baseline_out = raiden_unscan.unscan_layers(bf16_source, num_layers=4, scan_axis=1)
+
+    converter = MaxTextToMaxTextConverter(cfg, prefuse_moe_weights=False, target_dtype=jnp.bfloat16)
+    converter_out = converter.convert(raw_source, target_state=None)
+
+    base_flat = traverse_util.flatten_dict(baseline_out)
+    conv_flat = traverse_util.flatten_dict(converter_out)
+
+    self.assertEqual(set(base_flat.keys()), set(conv_flat.keys()))
+    for k in base_flat:
+      v_base = getattr(base_flat[k], "value", base_flat[k])
+      v_conv = getattr(conv_flat[k], "value", conv_flat[k])
+      self.assertEqual(v_base.shape, v_conv.shape, f"Shape mismatch at {k}")
+      self.assertEqual(v_base.dtype, v_conv.dtype, f"Dtype mismatch at {k}")
+      np.testing.assert_array_equal(np.asarray(v_base), np.asarray(v_conv), err_msg=f"Value mismatch at {k}")
+
+
 if __name__ == "__main__":
   unittest.main()

From 4f7ea64c353e61c003716f0cafcd133f448bd1dc Mon Sep 17 00:00:00 2001
From: Yixuan Wang <yixuannwang@google.com>
Date: Wed, 2 Sep 2026 04:34:54 +0000
Subject: [PATCH 3/7] Address review feedback: fix converter config resolution,
 restore metrics recorder, and add cache invalidation

- Handle nested vllm dict/object in HyperParameters for use_weight_converter and rollout_backend
- Restore self._metrics_recorder = metrics_module.MetricsRecorder() in MaxTextTrainingEngine
- Invalidate staged metadata cache in release_weight_sync()
- Gate unroll_gemma_scanned_weights by Gemma model identity in MaxTextVllmSampler
- Set default num_lanes=128 in compute_padded_moe_mlp_dim
- Clarify memory lifecycle in WeightConverter convert docstrings and enhance test_case_5 memory profiling
---
 .../integration/vllm/maxtext_vllm_rollout.py  | 13 ++++-
 src/maxtext/integration/vllm/moe_padding.py   |  2 +-
 .../integration/vllm/weight_converter.py      |  5 ++
 src/maxtext/training_engine/maxtext_engine.py | 16 +++++-
 .../unit/weight_converter_test.py             | 51 +++++++++++++++++--
 5 files changed, 81 insertions(+), 6 deletions(-)

diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py
index 18e8d56f72..d8f2ace4a3 100644
--- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py
+++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py
@@ -472,6 +472,17 @@ def __init__(
     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,
@@ -491,7 +502,7 @@ def update_params(
             pass
         raise
     if self._converter is None:
-      if self._direct_maxtext_sync:
+      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)
diff --git a/src/maxtext/integration/vllm/moe_padding.py b/src/maxtext/integration/vllm/moe_padding.py
index 3ad3edc95a..fa11df1ed2 100644
--- a/src/maxtext/integration/vllm/moe_padding.py
+++ b/src/maxtext/integration/vllm/moe_padding.py
@@ -35,7 +35,7 @@ def next_power_of_two(x: int) -> int:
 def compute_padded_moe_mlp_dim(
     hidden_size: Optional[int],
     moe_mlp_tp_size: int,
-    num_lanes: int,
+    num_lanes: int = 128,
 ) -> Optional[int]:
   """Computes padded MoE intermediate size for GMM_v2 kernel requirements.
 
diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py
index 3e828e6cff..45be490e7d 100644
--- a/src/maxtext/integration/vllm/weight_converter.py
+++ b/src/maxtext/integration/vllm/weight_converter.py
@@ -1200,6 +1200,10 @@ 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. Leaves are wrapped in nnx.Param.
+    Note on memory lifecycle: `src_flat.pop()` frees internal flattened dict references
+    as `result` is constructed to prevent dictionary growth overhead, but the caller's
+    `src_pytree` retains full-tree references until `convert()` returns and the caller
+    rebinds or drops its handle.
     """
     src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree))
     src_flat, _ = _strip_root(src_flat, "base")
@@ -1212,6 +1216,7 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]:
       result: Dict[Tuple[Any, ...], Any] = {}
       for group in self._groups:
         outs = self._execute_group_target_free(group, src_flat)
+        # Drop processed source entries from src_flat to reduce dict overhead
         for k in group.source_keys:
           src_flat.pop(k, None)
         for tgt_key, out in outs:
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 9b6aefd316..57ebc94e5f 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -416,16 +416,28 @@ def __init__(
         checkpoint_dir=self._config.checkpoint_dir,
         config=self._config,
     )
+    self._metrics_recorder = metrics_module.MetricsRecorder()
     self._throttler = inflight_throttler.InflightThrottler(config=self._config)
     self._raiden_syncs: Any = None
     self._last_staged_step: Optional[int] = None
     self._staged_metadata: Any = None
+    vllm_cfg = getattr(self._config, "vllm", {})
+    if isinstance(vllm_cfg, dict):
+      vllm_use_wc = vllm_cfg.get("use_weight_converter", False)
+      vllm_backend = vllm_cfg.get("rollout_backend", "maxtext")
+    else:
+      vllm_use_wc = getattr(vllm_cfg, "use_weight_converter", False)
+      vllm_backend = getattr(vllm_cfg, "rollout_backend", "maxtext")
+
     self._use_weight_converter = bool(
         getattr(self._config, "use_weight_converter", False)
+        or vllm_use_wc
         or os.environ.get("USE_WEIGHT_CONVERTER", "0").lower() in ("1", "true", "yes")
     )
     self._rollout_backend = (
-        getattr(self._config, "rollout_backend", "maxtext") or os.environ.get("ROLLOUT_BACKEND", "maxtext")
+        getattr(self._config, "rollout_backend", None)
+        or vllm_backend
+        or os.environ.get("ROLLOUT_BACKEND", "maxtext")
     )
     if self._use_weight_converter:
       from maxtext.integration.vllm.weight_converter import WeightConverter  # pylint: disable=g-import-not-at-top,import-outside-toplevel
@@ -1626,6 +1638,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_syncs:
       for sync in self._raiden_syncs:
         logging.vlog(1, "Trainer Raiden metrics: %s", sync.metrics())
diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py
index 6652f65790..30814f584b 100644
--- a/tests/post_training/unit/weight_converter_test.py
+++ b/tests/post_training/unit/weight_converter_test.py
@@ -526,13 +526,58 @@ def to_struct(x):
   def test_case_5_host_memory_profiling(self):
     import resource
 
-    cfg = _config()
-    source = _source_tree(True)
+    # Test with realistic scaled dimensions
+    scaled_emb = 128
+    scaled_experts = 8
+    scaled_mlp = 256
+    cfg = _config(
+        inhomogeneous_layer_cycle_interval=CYCLE,
+        num_decoder_layers=NUM_LAYERS,
+        padded_base_moe_mlp_dim=scaled_mlp,
+        prefuse_moe_weights=True,
+    )
+    # Build scaled source tree
+    blocks = NUM_LAYERS // CYCLE
+    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(blocks, 4, 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 = WeightConverter(config=cfg, rollout_backend="maxtext")
     before_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
-    out = converter.convert(source, target_state=None)
+    out = converter.convert(scaled_source, target_state=None)
     after_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+    logging.info(
+        "test_case_5_host_memory_profiling: before_rss=%d KB, after_rss=%d KB, delta=%d KB",
+        before_rss,
+        after_rss,
+        after_rss - before_rss,
+    )
     self.assertIsNotNone(out)
+    self.assertIn("decoder", out)
+    self.assertIn(f"layers_{NUM_LAYERS - 1}", out["decoder"])
+    wi = getattr(out["decoder"]["layers_0"]["moe_block"]["wi"], "value", out["decoder"]["layers_0"]["moe_block"]["wi"])
+    self.assertEqual(wi.shape, (scaled_experts, scaled_emb, scaled_mlp * 2))
 
   def test_case_6_parity_vs_raiden_unscan_on_homogeneous(self):
     from maxtext.integration.tunix.weight_mapping import raiden_unscan

From 8f75302020f9438e941fceb1d4af07871ea1abc5 Mon Sep 17 00:00:00 2001
From: Yixuan Wang <yixuannwang@google.com>
Date: Wed, 2 Sep 2026 21:24:49 +0000
Subject: [PATCH 4/7] Implement streaming trainer-side weight conversion for
 Raiden weight sync

- Add convert_streaming() to WeightConverter and MaxTextToMaxTextConverter for incremental transformation and eager memory release per group
- Add unscan_layers_streaming() to raiden_unscan with shared _unscan_one_key() helper
- Refactor MaxTextTrainingEngine.prepare_weight_sync() to stream piece-by-piece with unique strided worker indices
- Support RAIDEN_STREAM_PIECE_BATCH env var and deprecate RAIDEN_WEIGHT_SYNC_CHUNKS
- Add unit test coverage across weight converter, raiden unscan, and prepare weight sync suites
---
 .../tunix/weight_mapping/raiden_unscan.py     | 176 ++++++++-----
 .../integration/vllm/weight_converter.py      | 137 +++++++---
 src/maxtext/training_engine/maxtext_engine.py | 131 +++++-----
 .../unit/weight_converter_test.py             | 240 +++++++++++++-----
 tests/unit/prepare_weight_sync_test.py        | 165 ++++++++++++
 tests/unit/raiden_unscan_test.py              |  47 ++++
 6 files changed, 682 insertions(+), 214 deletions(-)
 create mode 100644 tests/unit/prepare_weight_sync_test.py

diff --git a/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py b/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py
index b0d0deed62..e372d014f5 100644
--- a/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py
+++ b/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py
@@ -36,35 +36,72 @@
 simpler, single-axis case directly instead of adapting that function.
 """
 
-from typing import Any
+import gc
+from typing import Any, Iterator, Tuple, List, Dict
 
 import jax
 from flax import nnx
 from flax.traverse_util import flatten_dict, unflatten_dict
 
 
-def unscan_layers(
+def _unscan_one_key(
+    key: Tuple[Any, ...],
+    value: Any,
+    num_layers: int,
+    layer_container: str = "layers",
+    scan_axis: int = 1,
+) -> Tuple[List[Tuple[Tuple[Any, ...], Any]], bool]:
+  """Unscans a single flattened pytree entry.
+
+  Returns:
+    (list_of_entries, is_scanned), where list_of_entries is a list of (new_key, value_slice) pairs.
+  """
+  if layer_container not in key:
+    return [(key, value)], False
+
+  idx = key.index(layer_container)
+  prefix = key[:idx]
+  suffix = key[idx + 1 :]
+  arr = getattr(value, "value", value)
+
+  if not hasattr(arr, "shape") or arr.ndim <= scan_axis:
+    return [(key, value)], False
+
+  if arr.shape[scan_axis] != num_layers:
+    raise ValueError(
+        f"unscan_layers: {'.'.join(str(k) for k in key)!r} has shape {arr.shape}, expected axis {scan_axis} to be"
+        f" num_layers={num_layers}."
+    )
+
+  entries = []
+  for i in range(num_layers):
+    sliced = jax.lax.index_in_dim(arr, i, axis=scan_axis, keepdims=False)
+    new_key = prefix + (f"{layer_container}_{i}",) + suffix
+    entries.append((new_key, sliced))
+  return entries, True
+
+
+def unscan_layers_streaming(
     state: Any,
     num_layers: int,
     layer_container: str = "layers",
     scan_axis: int = 1,
-) -> Any:
-  """Splits `state`'s scanned `layer_container` axis into per-layer entries.
+    *,
+    keys_per_piece: int = 1,
+) -> Iterator[Any]:
+  """Yields unscanned layer pieces incrementally for Raiden weight sync.
 
   Args:
     state: An `nnx.State` (or any pytree exposing `to_pure_dict`/`to_dict`, or a plain nested dict) of MaxText params,
       scanned along `scan_axis` under a `layer_container` key (MaxText's default scan layout).
-    num_layers: Number of layers the scanned axis must have. Used both to validate the input and to bound the unscan
-      loop.
-    layer_container: The pytree key holding the scanned per-layer params (MaxText's decoder body uses "layers").
+    num_layers: Number of layers the scanned axis must have.
+    layer_container: The pytree key holding the scanned per-layer params.
     scan_axis: The axis along which layers are scanned (default 1).
+    keys_per_piece: Number of original flattened keys to batch per yielded piece (default 1).
 
-  Returns:
-    A nested dict with `layer_container` keys replaced by `f"{layer_container}_{i}"` for each layer `i`, each holding
-    the corresponding `scan_axis` slice of the original array, wrapped in `nnx.Param` -- matching `nnx.state(...,
-    nnx.Param)`'s leaf type, so downstream consumers (`raiden_synchronizer.flatten_weights`, which unwraps `.value`)
-    see the same leaf shape whether or not this transform ran. Non-scanned entries (e.g. embeddings, final norm) pass
-    through unchanged, also rewrapped.
+  Yields:
+    Nested dicts with `nnx.Param`-wrapped leaves, each containing `keys_per_piece` keys' worth of
+    unscanned slices.
   """
   if hasattr(state, "to_pure_dict"):
     pure = state.to_pure_dict()
@@ -73,57 +110,80 @@ def unscan_layers(
   elif isinstance(state, dict):
     pure = state
   else:
-    return state
+    yield state
+    return
 
   flat = flatten_dict(pure)
-  new_flat = {}
-  unscanned_count = 0
-
-  # Drain `flat` as we go (pop, not iterate-then-keep) rather than holding
-  # every original scanned array alive for the whole function: at 30B-A3B
-  # scale (padded MoE weights are tens of GB each), keeping both the
-  # original scanned tree and the ~num_layers-times-larger unscanned tree
-  # alive simultaneously roughly doubles peak host memory during Raiden's
-  # D2H staging -- confirmed as the direct cause of an OOMKill there.
-  for key in list(flat.keys()):
-    value = flat.pop(key)
-    if layer_container not in key:
-      new_flat[key] = value
-      continue
-
-    idx = key.index(layer_container)
-    prefix = key[:idx]
-    suffix = key[idx + 1 :]
-    arr = getattr(value, "value", value)
-
-    if arr is None or not hasattr(arr, "shape") or getattr(arr, "ndim", 0) <= scan_axis:
-      # Not a per-layer leaf (shouldn't happen for real params under
-      # `layers`, but don't silently drop anything unexpected). Keep the
-      # original (possibly already-wrapped) value, matching pre-existing
-      # behavior -- unlike the actually-scanned case below, there's no
-      # multi-copy blowup here worth restructuring around.
-      new_flat[key] = value
-      continue
-    del value
-
-    if arr.shape[scan_axis] != num_layers:
-      raise ValueError(
-          f"unscan_layers: {'.'.join(key)!r} has shape {arr.shape}, expected axis {scan_axis} to be"
-          f" num_layers={num_layers}."
-      )
-
-    for i in range(num_layers):
-      sliced = jax.lax.index_in_dim(arr, i, axis=scan_axis, keepdims=False)
-      new_key = prefix + (f"{layer_container}_{i}",) + suffix
-      new_flat[new_key] = sliced
-    del arr
-    unscanned_count += 1
 
-  if unscanned_count == 0:
+  has_scanned = any(
+      layer_container in key
+      and hasattr(getattr(flat[key], "value", flat[key]), "shape")
+      and getattr(getattr(flat[key], "value", flat[key]), "ndim", 0) > scan_axis
+      for key in flat
+  )
+  if not has_scanned:
     raise ValueError(
         f"unscan_layers: found no scanned '{layer_container}' entries to unscan "
         "-- state may already be unscanned, or layer_container is wrong."
     )
 
+  keys_per_piece = max(1, keys_per_piece)
+  flat_keys = list(flat.keys())
+  for i in range(0, len(flat_keys), keys_per_piece):
+    chunk_keys = flat_keys[i : i + keys_per_piece]
+    piece_flat = {}
+    for key in chunk_keys:
+      value = flat.pop(key)
+      outputs, _ = _unscan_one_key(
+          key, value, num_layers=num_layers, layer_container=layer_container, scan_axis=scan_axis
+      )
+      for new_key, sliced in outputs:
+        piece_flat[new_key] = sliced
+      del value, outputs
+
+    nested = unflatten_dict(piece_flat)
+    del piece_flat
+    yield jax.tree_util.tree_map(
+        lambda x: nnx.Param(x) if not isinstance(x, (nnx.Param, nnx.Variable)) else x,
+        nested,
+    )
+
+  del flat
+  gc.collect()
+
+
+def unscan_layers(
+    state: Any,
+    num_layers: int,
+    layer_container: str = "layers",
+    scan_axis: int = 1,
+) -> Any:
+  """Splits `state`'s scanned `layer_container` axis into per-layer entries.
+
+  Args:
+    state: An `nnx.State` (or any pytree exposing `to_pure_dict`/`to_dict`, or a plain nested dict) of MaxText params,
+      scanned along `scan_axis` under a `layer_container` key (MaxText's default scan layout).
+    num_layers: Number of layers the scanned axis must have. Used both to validate the input and to bound the unscan
+      loop.
+    layer_container: The pytree key holding the scanned per-layer params (MaxText's decoder body uses "layers").
+    scan_axis: The axis along which layers are scanned (default 1).
+
+  Returns:
+    A nested dict with `layer_container` keys replaced by `f"{layer_container}_{i}"` for each layer `i`, each holding
+    the corresponding `scan_axis` slice of the original array, wrapped in `nnx.Param` -- matching `nnx.state(...,
+    nnx.Param)`'s leaf type, so downstream consumers (`raiden_synchronizer.flatten_weights`, which unwraps `.value`)
+    see the same leaf shape whether or not this transform ran. Non-scanned entries (e.g. embeddings, final norm) pass
+    through unchanged, also rewrapped.
+  """
+  if not hasattr(state, "to_pure_dict") and not hasattr(state, "to_dict") and not isinstance(state, dict):
+    return state
+
+  new_flat = {}
+  for piece in unscan_layers_streaming(
+      state, num_layers=num_layers, layer_container=layer_container, scan_axis=scan_axis
+  ):
+    new_flat.update(flatten_dict(piece))
+  gc.collect()
   nested = unflatten_dict(new_flat)
-  return jax.tree_util.tree_map(nnx.Param, nested)
+  del new_flat
+  return nested
diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py
index 45be490e7d..7c060410d2 100644
--- a/src/maxtext/integration/vllm/weight_converter.py
+++ b/src/maxtext/integration/vllm/weight_converter.py
@@ -23,7 +23,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from typing import List, Union, Any, Dict, Optional, Mapping, Tuple
+from typing import List, Union, Any, Dict, Optional, Mapping, Tuple, Iterator
 from flax import traverse_util, nnx
 from maxtext.integration.vllm.convert_utils import (
     _align_per_axis,
@@ -390,6 +390,59 @@ 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)."
+    )
+
+    # HuggingFace-shaped target: rename and restructure per the rule table.
+    result = {}
+    unfired = []
+    for rule in self.rules:
+      tensors = [flat_src.pop(src_pat) for src_pat in rule.source_patterns if src_pat in flat_src]
+      if not tensors:
+        # A rule can legitimately not fire (e.g. the `wi` rule when the
+        # trainer stores split `wi_0`/`wi_1`), but a table where *many*
+        # rules miss means the source layout has drifted.
+        unfired.append(rule.target_pattern)
+        continue
+
+      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):
+          result[rule.target_pattern.format(i)] = tensor
+      elif isinstance(out, list) and len(out) == 1:
+        result[rule.target_pattern] = out[0]
+      else:
+        result[rule.target_pattern] = out
+
+      del out, tensors
+      gc.collect()
+
+    if not result:
+      raise ValueError(
+          "No conversion rule matched the trainer state; the rollout would "
+          f"keep its dummy weights. Rule targets: {unfired}"
+      )
+    if unfired:
+      logging.info("Conversion rules that did not fire: %s", unfired)
+
+    return _rekey_to_target(result, target_state)
+
 
 # ==========================================
 # 4. Registries and Builders
@@ -1200,39 +1253,17 @@ 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. Leaves are wrapped in nnx.Param.
-    Note on memory lifecycle: `src_flat.pop()` frees internal flattened dict references
-    as `result` is constructed to prevent dictionary growth overhead, but the caller's
-    `src_pytree` retains full-tree references until `convert()` returns and the caller
-    rebinds or drops its handle.
     """
-    src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree))
-    src_flat, _ = _strip_root(src_flat, "base")
-
     if target_state is None:
-      if self._plan is None:
-        self._plan = self._build_target_free_plan(src_flat)
-        self._groups = _group_plan(self._plan)
-
-      result: Dict[Tuple[Any, ...], Any] = {}
-      for group in self._groups:
-        outs = self._execute_group_target_free(group, src_flat)
-        # Drop processed source entries from src_flat to reduce dict overhead
-        for k in group.source_keys:
-          src_flat.pop(k, None)
-        for tgt_key, out in outs:
-          result[tgt_key] = out
-        del outs
-
-      del src_flat
-      gc.collect()
-      nested = traverse_util.unflatten_dict(result)
-      del result
+      flat_result = {}
+      for piece in self.convert_streaming(src_pytree, target_state=None):
+        flat_result.update(traverse_util.flatten_dict(piece))
       gc.collect()
       _malloc_trim()
-      return jax.tree_util.tree_map(
-          lambda x: nnx.Param(x) if not isinstance(x, (nnx.Param, nnx.Variable)) else x,
-          nested,
-      )
+      return traverse_util.unflatten_dict(flat_result)
+
+    src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree))
+    src_flat, src_root = _strip_root(src_flat, "base")
 
     # Read variable types before purifying to plain arrays loses them.
     skip_paths = _non_param_paths(target_state)
@@ -1320,6 +1351,52 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]:
         nested,
     )
 
+  def convert_streaming(
+      self,
+      src_pytree: Any,
+      target_state: Any = None,
+      *,
+      groups_per_piece: int = 1,
+  ) -> Iterator[Dict[str, Any]]:
+    """Yields converted rollout weight pieces incrementally for target-free conversion.
+
+    Pure: `src_pytree` is not mutated. Each yielded piece is a nested dict of `nnx.Param`s
+    corresponding to `groups_per_piece` plan groups. Memory is freed piece-by-piece as
+    source keys are consumed.
+    """
+    if target_state is not None:
+      raise NotImplementedError("convert_streaming only supports target-free conversion (target_state=None).")
+
+    src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree))
+    src_flat, src_root = _strip_root(src_flat, "base")
+
+    if self._plan is None:
+      self._plan = self._build_target_free_plan(src_flat)
+      self._groups = _group_plan(self._plan)
+
+    groups_per_piece = max(1, groups_per_piece)
+    for i in range(0, len(self._groups), groups_per_piece):
+      piece_groups = self._groups[i : i + groups_per_piece]
+      piece_result: Dict[Tuple[Any, ...], Any] = {}
+      for group in piece_groups:
+        outs = self._execute_group_target_free(group, src_flat)
+        for k in group.source_keys:
+          src_flat.pop(k, None)
+        for tgt_key, out in outs:
+          piece_result[src_root + tgt_key] = out
+        del outs
+
+      nested = traverse_util.unflatten_dict(piece_result)
+      del piece_result
+      yield jax.tree_util.tree_map(
+          lambda x: nnx.Param(x) if not isinstance(x, (nnx.Param, nnx.Variable)) else x,
+          nested,
+      )
+
+    del src_flat
+    gc.collect()
+    _malloc_trim()
+
 
 def _rekey_to_target(flat_dotted: Dict[str, Any], target_state: Any) -> Dict[str, Any]:
   """Re-keys dotted converter output onto the target state's own key form.
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 57ebc94e5f..0ebcd9b0c3 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -31,8 +31,6 @@
 from flax import nnx
 from flax import struct
 from flax.linen import partitioning as nn_partitioning
-from flax.traverse_util import flatten_dict
-from flax.traverse_util import unflatten_dict
 import jax
 import jax.numpy as jnp
 from jax.typing import ArrayLike  # pylint: disable=g-importing-member
@@ -68,6 +66,8 @@
     "~2 ms for the pure-state equivalent. Logged once per engine instance."
 )
 
+_RAIDEN_WORKER_INDEX_STRIDE = 10_000  # >> any plausible piece count (dozens-to-low-hundreds of groups)
+
 
 def _malloc_trim() -> None:
   try:
@@ -421,6 +421,7 @@ def __init__(
     self._raiden_syncs: Any = None
     self._last_staged_step: Optional[int] = None
     self._staged_metadata: Any = None
+    self._warned_raiden_sync_chunks: bool = False
     vllm_cfg = getattr(self._config, "vllm", {})
     if isinstance(vllm_cfg, dict):
       vllm_use_wc = vllm_cfg.get("use_weight_converter", False)
@@ -1463,19 +1464,8 @@ def _get_trainable_params_state(self) -> Any:
       return nnx.state(model, nnx.Param)
     return self.model
 
-  def _split_into_chunks(self, nested_state: Any, num_chunks: int) -> list[Any]:
-    """Splits a nested param dict into `num_chunks` nested dicts of near-equal leaf count."""
-    if hasattr(nested_state, "to_pure_dict"):
-      pure_state = nested_state.to_pure_dict()
-    elif hasattr(nested_state, "to_dict"):
-      pure_state = nested_state.to_dict()
-    else:
-      pure_state = nested_state
-    flat = flatten_dict(pure_state)
-    chunk_flats = [{} for _ in range(num_chunks)]
-    for i, key in enumerate(flat):
-      chunk_flats[i % num_chunks][key] = flat[key]
-    return [unflatten_dict(cf) for cf in chunk_flats]
+  def _raiden_worker_index(self, piece_idx: int) -> int:
+    return jax.process_index() * _RAIDEN_WORKER_INDEX_STRIDE + piece_idx + 1
 
   def prepare_weight_sync(
       self,
@@ -1530,6 +1520,13 @@ def prepare_weight_sync(
 
       # 2. Extract clean trainable parameters
       params_state = self._get_trainable_params_state()
+      piece_batch = max(1, int(os.environ.get("RAIDEN_STREAM_PIECE_BATCH", "1")))
+      if "RAIDEN_WEIGHT_SYNC_CHUNKS" in os.environ and not getattr(self, "_warned_raiden_sync_chunks", False):
+        logging.warning(
+            "RAIDEN_WEIGHT_SYNC_CHUNKS is deprecated and no longer affects Raiden staging; "
+            "use RAIDEN_STREAM_PIECE_BATCH instead."
+        )
+        self._warned_raiden_sync_chunks = True
 
       if self._use_weight_converter:
         if self._weight_converter is None:
@@ -1539,92 +1536,94 @@ def prepare_weight_sync(
               rollout_backend=self._rollout_backend,
               debug=getattr(self._config, "weight_sync_debug", False),
           )
-        params_state = self._weight_converter.convert(params_state)
-        gc.collect()
+        piece_iter = self._weight_converter.convert_streaming(params_state, groups_per_piece=piece_batch)
       else:
-        # 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.
+        # 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,
         )
-
-        # 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(
+          piece_iter = raiden_unscan.unscan_layers_streaming(
               params_state,
               num_layers=self._config.num_decoder_layers,
               scan_axis=self._config.param_scan_axis,
+              keys_per_piece=piece_batch,
           )
+        else:
+          piece_iter = iter([params_state])
 
-      # 3. Bind parameters to the Raiden transport, one chunk at a time (see
-      # _split_into_chunks) -- construct the per-chunk synchronizers once,
-      # matching the persistent-instance-per-cycle pattern the rebind
-      # optimization (fewer stale holds) depends on.
-      num_chunks = max(1, int(os.environ.get("RAIDEN_WEIGHT_SYNC_CHUNKS", "1")))
-      if self._raiden_syncs is None:
-        # Under Pathways (JAX_PLATFORMS=proxy + JAX_BACKEND_TARGET set, same
-        # detection tunix's K8sJaxContext.initialize() uses), trainer params
-        # are proxy-backed and Raiden can't bind them in place -- host_stage
-        # pulls them to client host memory first. Direct-TPU trainers skip
-        # that extra copy since their params already live on TPU.
-        is_pathways = bool("proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get("JAX_BACKEND_TARGET"))
-        # worker_index must be unique per chunk (it seeds WorkUnitId's
-        # job_replica_id) -- otherwise every chunk's work unit collides under
-        # the same id in the handler's registry and only one survives
-        # registration.
-        self._raiden_syncs = [
-            raiden_synchronizer.RaidenSynchronizer(
-                job_name="trainer",
-                worker_index=jax.process_index() if num_chunks == 1 else (jax.process_index() * num_chunks + i + 1),
-                auto_h2d=False,
-                host_stage=is_pathways,
-                parallelism=4,
-            )
-            for i in range(num_chunks)
-        ]
-
-      chunks = self._split_into_chunks(params_state, num_chunks) if num_chunks > 1 else [params_state]
       del params_state
       gc.collect()
 
+      if self._raiden_syncs is None:
+        self._raiden_syncs = []
+
+      expected_num_pieces = len(self._raiden_syncs) if self._raiden_syncs else None
+      is_pathways = bool("proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get("JAX_BACKEND_TARGET"))
       verify_weights = os.environ.get("VERIFY_WEIGHTS", "").lower() == "true"
       all_metadata = []
       total_variables = 0
-      for chunk_idx, (sync, chunk_state) in enumerate(zip(self._raiden_syncs, chunks)):
-        sync.bind(chunk_state)
-        del chunk_state
+      piece_idx = -1
+
+      for piece_idx, piece in enumerate(piece_iter):
+        if expected_num_pieces is None and piece_idx >= len(self._raiden_syncs):
+          self._raiden_syncs.append(
+              raiden_synchronizer.RaidenSynchronizer(
+                  job_name="trainer",
+                  worker_index=self._raiden_worker_index(piece_idx),
+                  auto_h2d=False,
+                  host_stage=is_pathways,
+                  parallelism=4,
+              )
+          )
+        elif piece_idx >= len(self._raiden_syncs):
+          del piece
+          break
+
+        sync = self._raiden_syncs[piece_idx]
+        sync.bind(piece)
+        del piece
         gc.collect()
 
-        # 4. Initiate Device-to-Host transfer to stage this chunk for network
-        # transfer before moving on to the next chunk.
+        # 4. Initiate Device-to-Host transfer to stage this piece for network
+        # transfer before moving on to the next piece.
         if sync.active:
           sync.d2h()
 
         if verify_weights:
-          logging.info("Source weights checksums (chunk %d): %s", chunk_idx, sync.checksums())
+          logging.info("Source weights checksums (piece %d): %s", piece_idx, sync.checksums())
 
         metadata = sync.work_unit_metadata()
         total_variables += len(metadata.variables)
         all_metadata.append(metadata)
         sync.release_host_arrays()
 
+      if expected_num_pieces is not None:
+        remaining = 0
+        for _ in piece_iter:
+          remaining += 1
+        num_pieces = (piece_idx + 1) + remaining
+        if num_pieces != expected_num_pieces:
+          raise RuntimeError(
+              f"weight-sync piece count changed from {expected_num_pieces} to {num_pieces} "
+              "between rounds; the cached conversion plan should make this impossible "
+              "unless the model/config changed mid-run."
+          )
+      else:
+        num_pieces = piece_idx + 1
+
       gc.collect()
       _malloc_trim()
 
       logging.info(
-          "Trainer prepared weight sync for step %d: registered %d variables across %d chunk(s) on mesh %s",
+          "Trainer prepared weight sync for step %d: registered %d variables across %d piece(s) on mesh %s",
           self.train_step,
           total_variables,
-          num_chunks,
+          num_pieces,
           all_metadata[0].mesh_axes if all_metadata else None,
       )
       self._last_staged_step = self.train_step
diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py
index 30814f584b..912d83e900 100644
--- a/tests/post_training/unit/weight_converter_test.py
+++ b/tests/post_training/unit/weight_converter_test.py
@@ -25,10 +25,13 @@
 # Must precede the first JAX import: the cross-mesh tests below need more than
 # one CPU device, and the backend reads this only at initialization.
 os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=8")
+os.environ.setdefault("JAX_PLATFORMS", "cpu")
 
 import types as pytypes  # pylint: disable=wrong-import-position
 import unittest  # pylint: disable=wrong-import-position
 
+import logging
+from typing import Any
 import jax  # pylint: disable=wrong-import-position
 import jax.numpy as jnp  # pylint: disable=wrong-import-position
 import numpy as np  # pylint: disable=wrong-import-position
@@ -439,6 +442,74 @@ def test_padded_moe_fusion_stays_on_the_source_mesh(self):
       )
 
 
+def _profile_conversion_worker(is_streaming: bool, result_queue: Any):
+  import gc
+  import resource
+  import types as pytypes
+  import jax.numpy as jnp
+  from maxtext.integration.vllm.weight_converter import MaxTextToMaxTextConverter
+
+  num_layers = 16
+  cycle = 2
+  scaled_emb = 128
+  scaled_experts = 8
+  scaled_mlp = 256
+  blocks = num_layers // cycle
+
+  cfg = pytypes.SimpleNamespace(
+      inhomogeneous_layer_cycle_interval=cycle,
+      num_decoder_layers=num_layers,
+      param_scan_axis=1,
+      padded_base_moe_mlp_dim=scaled_mlp,
+      prefuse_moe_weights=True,
+      weight_dtype=jnp.float32,
+  )
+
+  def _arr(*shape):
+    return jnp.ones(shape, dtype=jnp.float32)
+
+  layers = {}
+  for slot in range(cycle):
+    layers[f"layer_{slot}"] = {
+        "input_layernorm": {"scale": _arr(scaled_emb, blocks)},
+        "post_self_attention_layernorm": {"scale": _arr(scaled_emb, blocks)},
+        "self_attention": {
+            "query": {"kernel": _arr(scaled_emb, blocks, 4, 32)},
+            "key": {"kernel": _arr(scaled_emb, blocks, 2, 32)},
+            "value": {"kernel": _arr(scaled_emb, blocks, 2, 32)},
+            "out": {"kernel": _arr(4, blocks, 32, scaled_emb)},
+        },
+        "moe_block": {
+            "gate": {"kernel": _arr(scaled_emb, blocks, scaled_experts)},
+            "wi_0": _arr(scaled_experts, blocks, scaled_emb, scaled_mlp),
+            "wi_1": _arr(scaled_experts, blocks, scaled_emb, scaled_mlp),
+            "wo": _arr(scaled_experts, blocks, scaled_mlp, scaled_emb),
+        },
+    }
+  scaled_source = {
+      "base": {
+          "token_embedder": {"embedding": _arr(256, scaled_emb)},
+          "decoder": {"decoder_norm": {"scale": _arr(scaled_emb)}, "layers": layers},
+      }
+  }
+
+  converter = MaxTextToMaxTextConverter(cfg, prefuse_moe_weights=True)
+  gc.collect()
+  before_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+
+  if is_streaming:
+    for piece in converter.convert_streaming(scaled_source, target_state=None, groups_per_piece=1):
+      del piece
+      gc.collect()
+  else:
+    out = converter.convert(scaled_source, target_state=None)
+    del out
+    gc.collect()
+
+  after_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+  result_queue.put(after_rss - before_rss)
+
+
 class TargetFreeConversionTest(unittest.TestCase):
   """Comprehensive test suite for target-free key synthesis and execution."""
 
@@ -466,13 +537,14 @@ def test_case_1_homogeneous_target_free_unroll(self):
     }
     converter = WeightConverter(config=cfg, rollout_backend="maxtext")
     out = converter.convert(source, target_state=None)
-    self.assertIn("token_embedder", out)
-    self.assertIn("decoder", out)
+    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["decoder"])
-      scale = getattr(out["decoder"][layer_key]["input_layernorm"]["scale"], "value", out["decoder"][layer_key]["input_layernorm"]["scale"])
-      query = getattr(out["decoder"][layer_key]["self_attention"]["query"]["kernel"], "value", out["decoder"][layer_key]["self_attention"]["query"]["kernel"])
+      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))
 
@@ -481,10 +553,11 @@ def test_case_2_hybrid_cycle_target_free_unroll(self):
     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["decoder"][f"layers_{layer}"]["input_layernorm"]["scale"], "value", out["decoder"][f"layers_{layer}"]["input_layernorm"]["scale"])
+      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))
 
@@ -502,8 +575,9 @@ def test_case_3_prefused_moe_target_free(self):
     source = _source_tree(True)
     converter = MaxTextToMaxTextConverter(cfg, prefuse_moe_weights=True)
     out = converter.convert(source, target_state=None)
-    wi = getattr(out["decoder"]["layers_0"]["moe_block"]["wi"], "value", out["decoder"]["layers_0"]["moe_block"]["wi"])
-    wo = getattr(out["decoder"]["layers_0"]["moe_block"]["wo"], "value", out["decoder"]["layers_0"]["moe_block"]["wo"])
+    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))
 
@@ -517,67 +591,34 @@ def to_struct(x):
     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["decoder"]["layers_0"]["moe_block"]["wi"], "value", out["decoder"]["layers_0"]["moe_block"]["wi"])
+    wi = getattr(out_root["decoder"]["layers_0"]["moe_block"]["wi"], "value", out_root["decoder"]["layers_0"]["moe_block"]["wi"])
     self.assertEqual(wi.shape, (EXPERTS, EMB, 32))
 
   def test_case_5_host_memory_profiling(self):
-    import resource
+    import multiprocessing
+    ctx = multiprocessing.get_context("spawn")
+    q_non_stream = ctx.Queue()
+    p_non_stream = ctx.Process(target=_profile_conversion_worker, args=(False, q_non_stream))
+    p_non_stream.start()
+    delta_non_stream = q_non_stream.get(timeout=60)
+    p_non_stream.join()
+
+    q_stream = ctx.Queue()
+    p_stream = ctx.Process(target=_profile_conversion_worker, args=(True, q_stream))
+    p_stream.start()
+    delta_stream = q_stream.get(timeout=60)
+    p_stream.join()
 
-    # Test with realistic scaled dimensions
-    scaled_emb = 128
-    scaled_experts = 8
-    scaled_mlp = 256
-    cfg = _config(
-        inhomogeneous_layer_cycle_interval=CYCLE,
-        num_decoder_layers=NUM_LAYERS,
-        padded_base_moe_mlp_dim=scaled_mlp,
-        prefuse_moe_weights=True,
-    )
-    # Build scaled source tree
-    blocks = NUM_LAYERS // CYCLE
-    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(blocks, 4, 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 = WeightConverter(config=cfg, rollout_backend="maxtext")
-    before_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
-    out = converter.convert(scaled_source, target_state=None)
-    after_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
     logging.info(
-        "test_case_5_host_memory_profiling: before_rss=%d KB, after_rss=%d KB, delta=%d KB",
-        before_rss,
-        after_rss,
-        after_rss - before_rss,
+        "test_case_5_host_memory_profiling: delta_non_stream=%d KB, delta_stream=%d KB",
+        delta_non_stream,
+        delta_stream,
     )
-    self.assertIsNotNone(out)
-    self.assertIn("decoder", out)
-    self.assertIn(f"layers_{NUM_LAYERS - 1}", out["decoder"])
-    wi = getattr(out["decoder"]["layers_0"]["moe_block"]["wi"], "value", out["decoder"]["layers_0"]["moe_block"]["wi"])
-    self.assertEqual(wi.shape, (scaled_experts, scaled_emb, scaled_mlp * 2))
+    self.assertLessEqual(delta_stream, delta_non_stream)
 
   def test_case_6_parity_vs_raiden_unscan_on_homogeneous(self):
     from maxtext.integration.tunix.weight_mapping import raiden_unscan
@@ -619,6 +660,85 @@ def test_case_6_parity_vs_raiden_unscan_on_homogeneous(self):
       self.assertEqual(v_base.dtype, v_conv.dtype, f"Dtype mismatch at {k}")
       np.testing.assert_array_equal(np.asarray(v_base), np.asarray(v_conv), err_msg=f"Value mismatch at {k}")
 
+  def test_case_7_streaming_piece_count_and_parity(self):
+    cfg = _config(
+        inhomogeneous_layer_cycle_interval=CYCLE,
+        num_decoder_layers=NUM_LAYERS,
+        prefuse_moe_weights=True,
+    )
+    source = _source_tree(True)
+    converter = WeightConverter(config=cfg, rollout_backend="maxtext")
+    pieces = list(converter.convert_streaming(source, target_state=None, groups_per_piece=1))
+    self.assertEqual(len(pieces), len(converter._direct._groups))
+
+    # Parity check against fresh non-streaming converter
+    fresh_converter = WeightConverter(config=cfg, rollout_backend="maxtext")
+    expected_out = fresh_converter.convert(source, target_state=None)
+
+    merged_flat = {}
+    for piece in pieces:
+      piece_flat = traverse_util.flatten_dict(piece)
+      for k, v in piece_flat.items():
+        self.assertNotIn(k, merged_flat, f"Duplicate key across pieces: {k}")
+        merged_flat[k] = v
+
+    expected_flat = traverse_util.flatten_dict(expected_out)
+    self.assertEqual(set(merged_flat.keys()), set(expected_flat.keys()))
+    for k in expected_flat:
+      v_exp = getattr(expected_flat[k], "value", expected_flat[k])
+      v_got = getattr(merged_flat[k], "value", merged_flat[k])
+      self.assertEqual(v_exp.shape, v_got.shape, f"Shape mismatch at {k}")
+      self.assertEqual(v_exp.dtype, v_got.dtype, f"Dtype mismatch at {k}")
+      np.testing.assert_array_equal(np.asarray(v_exp), np.asarray(v_got), err_msg=f"Value mismatch at {k}")
+
+  def test_case_8_streaming_piece_batching(self):
+    cfg = _config(
+        inhomogeneous_layer_cycle_interval=CYCLE,
+        num_decoder_layers=NUM_LAYERS,
+        prefuse_moe_weights=True,
+    )
+    source = _source_tree(True)
+    converter = WeightConverter(config=cfg, rollout_backend="maxtext")
+    pieces = list(converter.convert_streaming(source, target_state=None, groups_per_piece=2))
+    num_groups = len(converter._direct._groups)
+    expected_piece_count = (num_groups + 1) // 2
+    self.assertEqual(len(pieces), expected_piece_count)
+
+    fresh_converter = WeightConverter(config=cfg, rollout_backend="maxtext")
+    expected_out = fresh_converter.convert(source, target_state=None)
+    expected_flat = traverse_util.flatten_dict(expected_out)
+
+    merged_flat = {}
+    for piece in pieces:
+      piece_flat = traverse_util.flatten_dict(piece)
+      for k, v in piece_flat.items():
+        self.assertNotIn(k, merged_flat, f"Duplicate key across pieces: {k}")
+        merged_flat[k] = v
+
+    self.assertEqual(set(merged_flat.keys()), set(expected_flat.keys()))
+    for k in expected_flat:
+      v_exp = getattr(expected_flat[k], "value", expected_flat[k])
+      v_got = getattr(merged_flat[k], "value", merged_flat[k])
+      np.testing.assert_array_equal(np.asarray(v_exp), np.asarray(v_got))
+
+  def test_case_9_weight_converter_convert_streaming_dispatch(self):
+    cfg = _config(
+        inhomogeneous_layer_cycle_interval=CYCLE,
+        num_decoder_layers=NUM_LAYERS,
+        prefuse_moe_weights=True,
+    )
+    source = _source_tree(True)
+    # Direct MaxText mode delegates correctly
+    direct_wc = WeightConverter(config=cfg, rollout_backend="maxtext")
+    pieces = list(direct_wc.convert_streaming(source, target_state=None))
+    self.assertGreater(len(pieces), 0)
+
+    # Torchax rules mode raises NotImplementedError
+    rule = Rule(source_patterns=["some_pattern"], target_pattern="some_target")
+    torchax_wc = WeightConverter(rules=[rule], rollout_backend="torchax")
+    with self.assertRaises(NotImplementedError):
+      list(torchax_wc.convert_streaming(source))
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/prepare_weight_sync_test.py b/tests/unit/prepare_weight_sync_test.py
new file mode 100644
index 0000000000..0976a6efc3
--- /dev/null
+++ b/tests/unit/prepare_weight_sync_test.py
@@ -0,0 +1,165 @@
+# 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 streaming logic."""
+
+import os
+os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=8")
+os.environ.setdefault("JAX_PLATFORMS", "cpu")
+
+
+import sys
+import types as pytypes
+import unittest
+from unittest import mock
+
+# Ensure tunix C-extension / protobuf initializes before transformers/orbax
+try:
+  import tunix.experimental.weight_sync.raiden_synchronizer  # pylint: disable=unused-import
+except ImportError:
+  pass
+
+import jax
+import jax.numpy as jnp
+from maxtext.training_engine.maxtext_engine import MaxTextTrainingEngine, _RAIDEN_WORKER_INDEX_STRIDE
+
+
+class PrepareWeightSyncTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    # Create engine instance without running heavy __init__
+    self.engine = MaxTextTrainingEngine.__new__(MaxTextTrainingEngine)
+    self.engine._raiden_syncs = None
+    self.engine._last_staged_step = None
+    self.engine._staged_metadata = None
+    self.engine._train_step = 0
+    self.engine._throttler = mock.MagicMock()
+    self.engine._config = pytypes.SimpleNamespace(
+        scan_layers=False,
+        num_decoder_layers=2,
+        param_scan_axis=1,
+        weight_sync_debug=False,
+    )
+    self.engine._use_weight_converter = True
+    self.engine._weight_converter = mock.MagicMock()
+    self.engine._rollout_backend = "maxtext"
+    self.engine._warned_raiden_sync_chunks = False
+    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_streaming_grows_syncs_and_accumulates_metadata(self, mock_sync_cls):
+    created_syncs = []
+
+    def make_sync(*args, **kwargs):
+      s = mock.MagicMock()
+      s.active = True
+      s.worker_index = kwargs.get("worker_index")
+      s.work_unit_metadata.return_value = self._make_dummy_metadata(num_vars=2)
+      s.checksums.return_value = {}
+      created_syncs.append(s)
+      return s
+
+    mock_sync_cls.side_effect = make_sync
+
+    pieces = [{"piece_0": 0}, {"piece_1": 1}, {"piece_2": 2}]
+    self.engine._weight_converter.convert_streaming.return_value = iter(pieces)
+
+    metadata = self.engine.prepare_weight_sync()
+
+    self.assertEqual(len(metadata), 3)
+    self.assertEqual(len(self.engine._raiden_syncs), 3)
+    self.assertEqual(len(created_syncs), 3)
+
+    # Worker indices must be unique and properly strided
+    expected_indices = [
+        jax.process_index() * _RAIDEN_WORKER_INDEX_STRIDE + i + 1 for i in range(3)
+    ]
+    actual_indices = [s.worker_index for s in created_syncs]
+    self.assertEqual(actual_indices, expected_indices)
+
+    # Check that bind, d2h, metadata, release were called on each sync
+    for s, p in zip(created_syncs, pieces):
+      s.bind.assert_called_once_with(p)
+      s.d2h.assert_called_once()
+      s.work_unit_metadata.assert_called_once()
+      s.release_host_arrays.assert_called_once()
+
+  @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
+  def test_rebind_reuses_sync_instances(self, mock_sync_cls):
+    mock_syncs = []
+
+    def make_sync(*args, **kwargs):
+      s = mock.MagicMock()
+      s.active = True
+      s.worker_index = kwargs.get("worker_index")
+      s.work_unit_metadata.return_value = self._make_dummy_metadata(num_vars=2)
+      mock_syncs.append(s)
+      return s
+
+    mock_sync_cls.side_effect = make_sync
+
+    # Round 1
+    self.engine._weight_converter.convert_streaming.return_value = iter([{"p0": 0}, {"p1": 1}])
+    self.engine.prepare_weight_sync()
+    self.assertEqual(len(mock_syncs), 2)
+    first_round_syncs = list(self.engine._raiden_syncs)
+
+    # Round 2 at step 1
+    self.engine._train_step = 1
+    self.engine._weight_converter.convert_streaming.return_value = iter([{"p0": 0}, {"p1": 1}])
+    self.engine.prepare_weight_sync()
+
+    # No new instances created
+    self.assertEqual(len(mock_syncs), 2)
+    self.assertEqual(self.engine._raiden_syncs, first_round_syncs)
+
+  @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
+  def test_piece_count_mismatch_between_rounds_raises(self, mock_sync_cls):
+    mock_sync_cls.side_effect = lambda *a, **kw: mock.MagicMock(
+        active=True, work_unit_metadata=mock.MagicMock(return_value=self._make_dummy_metadata())
+    )
+
+    # Round 1 has 2 pieces
+    self.engine._weight_converter.convert_streaming.return_value = iter([{"p0": 0}, {"p1": 1}])
+    self.engine.prepare_weight_sync()
+
+    # Round 2 has 3 pieces
+    self.engine._train_step = 1
+    self.engine._weight_converter.convert_streaming.return_value = iter([{"p0": 0}, {"p1": 1}, {"p2": 2}])
+    with self.assertRaisesRegex(RuntimeError, "weight-sync piece count changed from 2 to 3"):
+      self.engine.prepare_weight_sync()
+
+  @mock.patch.dict(os.environ, {"RAIDEN_WEIGHT_SYNC_CHUNKS": "4"})
+  @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
+  def test_deprecated_chunks_env_var_warning(self, mock_sync_cls):
+    mock_sync_cls.side_effect = lambda *a, **kw: mock.MagicMock(
+        active=True, work_unit_metadata=mock.MagicMock(return_value=self._make_dummy_metadata())
+    )
+    self.engine._weight_converter.convert_streaming.return_value = iter([{"p0": 0}])
+    with mock.patch("absl.logging.warning") as mock_warn:
+      self.engine.prepare_weight_sync()
+      mock_warn.assert_called()
+      self.assertTrue(any("RAIDEN_WEIGHT_SYNC_CHUNKS is deprecated" in str(call) for call in mock_warn.call_args_list))
+    self.assertTrue(self.engine._warned_raiden_sync_chunks)
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/tests/unit/raiden_unscan_test.py b/tests/unit/raiden_unscan_test.py
index 694f2449b0..d65d54db9c 100644
--- a/tests/unit/raiden_unscan_test.py
+++ b/tests/unit/raiden_unscan_test.py
@@ -186,6 +186,53 @@ def test_total_element_count_is_preserved(self):
     after = sum(int(np.size(np.asarray(_unwrap(x)))) for x in jax.tree_util.tree_leaves(unscanned))
     self.assertEqual(after, before)
 
+  def test_streaming_piece_count_and_parity(self):
+    state = self._scanned_state()
+    pieces = list(raiden_unscan.unscan_layers_streaming(state, num_layers=_NUM_LAYERS, keys_per_piece=1))
+    # 3 original flattened keys: embed, layers.kernel, layers.scale
+    self.assertEqual(len(pieces), 3)
+
+    expected = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS)
+    merged_flat = {}
+    for piece in pieces:
+      piece_flat = raiden_unscan.flatten_dict(piece)
+      for k, v in piece_flat.items():
+        self.assertNotIn(k, merged_flat)
+        merged_flat[k] = v
+
+    expected_flat = raiden_unscan.flatten_dict(expected)
+    self.assertEqual(set(merged_flat.keys()), set(expected_flat.keys()))
+    for k in expected_flat:
+      v_exp = _unwrap(expected_flat[k])
+      v_got = _unwrap(merged_flat[k])
+      np.testing.assert_array_equal(np.asarray(v_got), np.asarray(v_exp))
+
+  def test_streaming_keys_per_piece_batching(self):
+    state = self._scanned_state()
+    pieces = list(raiden_unscan.unscan_layers_streaming(state, num_layers=_NUM_LAYERS, keys_per_piece=2))
+    # 3 keys with batch size 2 -> ceil(3/2) = 2 pieces
+    self.assertEqual(len(pieces), 2)
+
+    expected = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS)
+    expected_flat = raiden_unscan.flatten_dict(expected)
+    merged_flat = {}
+    for piece in pieces:
+      merged_flat.update(raiden_unscan.flatten_dict(piece))
+
+    self.assertEqual(set(merged_flat.keys()), set(expected_flat.keys()))
+    for k in expected_flat:
+      np.testing.assert_array_equal(np.asarray(_unwrap(merged_flat[k])), np.asarray(_unwrap(expected_flat[k])))
+
+  def test_streaming_leaves_are_nnx_param(self):
+    state = self._scanned_state()
+    for piece in raiden_unscan.unscan_layers_streaming(state, num_layers=_NUM_LAYERS):
+      for leaf in jax.tree_util.tree_leaves(piece, is_leaf=lambda x: isinstance(x, nnx.Param)):
+        self.assertIsInstance(leaf, nnx.Param)
+
+  def test_streaming_already_unscanned_state_raises(self):
+    with self.assertRaisesRegex(ValueError, "found no scanned 'layers' entries"):
+      list(raiden_unscan.unscan_layers_streaming(nnx.state(UnscannedModel(), nnx.Param), num_layers=_NUM_LAYERS))
+
 
 if __name__ == "__main__":
   absltest.main()

From adcd343476d0e2d1b21230be3ceee3ef85effe46 Mon Sep 17 00:00:00 2001
From: Yixuan Wang <yixuannwang@google.com>
Date: Wed, 2 Sep 2026 21:52:00 +0000
Subject: [PATCH 5/7] Fix target-free weight keys, clean up dead code and piece
 count checks

- Drop src_root ('base') prefix from target-free piece outputs in MaxTextToMaxTextConverter.convert_streaming()
- Remove out_root workaround in weight_converter_test.py test cases 1-4
- Remove dead code in WeightConverter.convert_streaming()
- Clean up _warned_raiden_sync_chunks check and simplify piece count mismatch validation in prepare_weight_sync()
---
 .../integration/vllm/weight_converter.py      | 45 ++-----------------
 src/maxtext/training_engine/maxtext_engine.py | 28 ++++--------
 .../unit/weight_converter_test.py             | 22 ++++-----
 3 files changed, 21 insertions(+), 74 deletions(-)

diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py
index 7c060410d2..2c78fc5988 100644
--- a/src/maxtext/integration/vllm/weight_converter.py
+++ b/src/maxtext/integration/vllm/weight_converter.py
@@ -404,45 +404,6 @@ def convert_streaming(
         "convert_streaming is only supported in direct MaxText-to-MaxText mode (rollout_backend='maxtext' and rules=None)."
     )
 
-    # HuggingFace-shaped target: rename and restructure per the rule table.
-    result = {}
-    unfired = []
-    for rule in self.rules:
-      tensors = [flat_src.pop(src_pat) for src_pat in rule.source_patterns if src_pat in flat_src]
-      if not tensors:
-        # A rule can legitimately not fire (e.g. the `wi` rule when the
-        # trainer stores split `wi_0`/`wi_1`), but a table where *many*
-        # rules miss means the source layout has drifted.
-        unfired.append(rule.target_pattern)
-        continue
-
-      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):
-          result[rule.target_pattern.format(i)] = tensor
-      elif isinstance(out, list) and len(out) == 1:
-        result[rule.target_pattern] = out[0]
-      else:
-        result[rule.target_pattern] = out
-
-      del out, tensors
-      gc.collect()
-
-    if not result:
-      raise ValueError(
-          "No conversion rule matched the trainer state; the rollout would "
-          f"keep its dummy weights. Rule targets: {unfired}"
-      )
-    if unfired:
-      logging.info("Conversion rules that did not fire: %s", unfired)
-
-    return _rekey_to_target(result, target_state)
-
 
 # ==========================================
 # 4. Registries and Builders
@@ -1263,7 +1224,7 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]:
       return traverse_util.unflatten_dict(flat_result)
 
     src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree))
-    src_flat, src_root = _strip_root(src_flat, "base")
+    src_flat, _ = _strip_root(src_flat, "base")
 
     # Read variable types before purifying to plain arrays loses them.
     skip_paths = _non_param_paths(target_state)
@@ -1368,7 +1329,7 @@ def convert_streaming(
       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")
+    src_flat, _ = _strip_root(src_flat, "base")
 
     if self._plan is None:
       self._plan = self._build_target_free_plan(src_flat)
@@ -1383,7 +1344,7 @@ def convert_streaming(
         for k in group.source_keys:
           src_flat.pop(k, None)
         for tgt_key, out in outs:
-          piece_result[src_root + tgt_key] = out
+          piece_result[tgt_key] = out
         del outs
 
       nested = traverse_util.unflatten_dict(piece_result)
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 0ebcd9b0c3..7d3937c46f 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -1521,7 +1521,7 @@ def prepare_weight_sync(
       # 2. Extract clean trainable parameters
       params_state = self._get_trainable_params_state()
       piece_batch = max(1, int(os.environ.get("RAIDEN_STREAM_PIECE_BATCH", "1")))
-      if "RAIDEN_WEIGHT_SYNC_CHUNKS" in os.environ and not getattr(self, "_warned_raiden_sync_chunks", False):
+      if "RAIDEN_WEIGHT_SYNC_CHUNKS" in os.environ and not self._warned_raiden_sync_chunks:
         logging.warning(
             "RAIDEN_WEIGHT_SYNC_CHUNKS is deprecated and no longer affects Raiden staging; "
             "use RAIDEN_STREAM_PIECE_BATCH instead."
@@ -1567,10 +1567,9 @@ def prepare_weight_sync(
       verify_weights = os.environ.get("VERIFY_WEIGHTS", "").lower() == "true"
       all_metadata = []
       total_variables = 0
-      piece_idx = -1
 
       for piece_idx, piece in enumerate(piece_iter):
-        if expected_num_pieces is None and piece_idx >= len(self._raiden_syncs):
+        if piece_idx >= len(self._raiden_syncs):
           self._raiden_syncs.append(
               raiden_synchronizer.RaidenSynchronizer(
                   job_name="trainer",
@@ -1580,9 +1579,6 @@ def prepare_weight_sync(
                   parallelism=4,
               )
           )
-        elif piece_idx >= len(self._raiden_syncs):
-          del piece
-          break
 
         sync = self._raiden_syncs[piece_idx]
         sync.bind(piece)
@@ -1602,19 +1598,13 @@ def prepare_weight_sync(
         all_metadata.append(metadata)
         sync.release_host_arrays()
 
-      if expected_num_pieces is not None:
-        remaining = 0
-        for _ in piece_iter:
-          remaining += 1
-        num_pieces = (piece_idx + 1) + remaining
-        if num_pieces != expected_num_pieces:
-          raise RuntimeError(
-              f"weight-sync piece count changed from {expected_num_pieces} to {num_pieces} "
-              "between rounds; the cached conversion plan should make this impossible "
-              "unless the model/config changed mid-run."
-          )
-      else:
-        num_pieces = piece_idx + 1
+      num_pieces = len(all_metadata)
+      if expected_num_pieces is not None and num_pieces != expected_num_pieces:
+        raise RuntimeError(
+            f"weight-sync piece count changed from {expected_num_pieces} to {num_pieces} "
+            "between rounds; the cached conversion plan should make this impossible "
+            "unless the model/config changed mid-run."
+        )
 
       gc.collect()
       _malloc_trim()
diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py
index 912d83e900..9c58974041 100644
--- a/tests/post_training/unit/weight_converter_test.py
+++ b/tests/post_training/unit/weight_converter_test.py
@@ -537,14 +537,13 @@ def test_case_1_homogeneous_target_free_unroll(self):
     }
     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)
+    self.assertIn("token_embedder", out)
+    self.assertIn("decoder", out)
     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.assertIn(layer_key, out["decoder"])
+      scale = getattr(out["decoder"][layer_key]["input_layernorm"]["scale"], "value", out["decoder"][layer_key]["input_layernorm"]["scale"])
+      query = getattr(out["decoder"][layer_key]["self_attention"]["query"]["kernel"], "value", out["decoder"][layer_key]["self_attention"]["query"]["kernel"])
       self.assertEqual(scale.shape, (EMB,))
       self.assertEqual(query.shape, (EMB, 2, 4))
 
@@ -553,11 +552,10 @@ def test_case_2_hybrid_cycle_target_free_unroll(self):
     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"])
+      got = getattr(out["decoder"][f"layers_{layer}"]["input_layernorm"]["scale"], "value", out["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))
 
@@ -575,9 +573,8 @@ def test_case_3_prefused_moe_target_free(self):
     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"])
+    wi = getattr(out["decoder"]["layers_0"]["moe_block"]["wi"], "value", out["decoder"]["layers_0"]["moe_block"]["wi"])
+    wo = getattr(out["decoder"]["layers_0"]["moe_block"]["wo"], "value", out["decoder"]["layers_0"]["moe_block"]["wo"])
     self.assertEqual(wi.shape, (EXPERTS, EMB, padded_dim * 2))
     self.assertEqual(wo.shape, (EXPERTS, padded_dim, EMB))
 
@@ -591,11 +588,10 @@ def to_struct(x):
     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"])
+    wi = getattr(out["decoder"]["layers_0"]["moe_block"]["wi"], "value", out["decoder"]["layers_0"]["moe_block"]["wi"])
     self.assertEqual(wi.shape, (EXPERTS, EMB, 32))
 
   def test_case_5_host_memory_profiling(self):

From 72bd221cfd4fc1f3ddfe7e44a5baa2f24065e8ba Mon Sep 17 00:00:00 2001
From: Yixuan Wang <yixuannwang@google.com>
Date: Thu, 3 Sep 2026 22:17:22 +0000
Subject: [PATCH 6/7] Simplify Raiden weight sync to single-piece conversion
 and preserve base root in streaming converter

- Revert streaming piece-by-piece conversion in MaxTextTrainingEngine to single-piece convert and bind
- Restore base root prefix in MaxTextToMaxTextConverter.convert_streaming()
- Update prepare_weight_sync_test suite to reflect single sync instance
---
 .../integration/vllm/weight_converter.py      |  6 +-
 src/maxtext/training_engine/maxtext_engine.py | 77 +++++++------------
 tests/unit/prepare_weight_sync_test.py        | 54 ++++---------
 3 files changed, 48 insertions(+), 89 deletions(-)

diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py
index 2c78fc5988..9d196094c6 100644
--- a/src/maxtext/integration/vllm/weight_converter.py
+++ b/src/maxtext/integration/vllm/weight_converter.py
@@ -1329,7 +1329,9 @@ def convert_streaming(
       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, _ = _strip_root(src_flat, "base")
+    src_flat, src_root = _strip_root(src_flat, "base")
+    if not src_root:
+      src_root = ("base",)
 
     if self._plan is None:
       self._plan = self._build_target_free_plan(src_flat)
@@ -1344,7 +1346,7 @@ def convert_streaming(
         for k in group.source_keys:
           src_flat.pop(k, None)
         for tgt_key, out in outs:
-          piece_result[tgt_key] = out
+          piece_result[src_root + tgt_key] = out
         del outs
 
       nested = traverse_util.unflatten_dict(piece_result)
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 7d3937c46f..c9b6fdae7c 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -1536,7 +1536,7 @@ def prepare_weight_sync(
               rollout_backend=self._rollout_backend,
               debug=getattr(self._config, "weight_sync_debug", False),
           )
-        piece_iter = self._weight_converter.convert_streaming(params_state, groups_per_piece=piece_batch)
+        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
@@ -1547,73 +1547,54 @@ def prepare_weight_sync(
             params_state,
         )
         if self._config.scan_layers:
-          piece_iter = raiden_unscan.unscan_layers_streaming(
+          converted_state = raiden_unscan.unscan_layers(
               params_state,
               num_layers=self._config.num_decoder_layers,
               scan_axis=self._config.param_scan_axis,
-              keys_per_piece=piece_batch,
           )
         else:
-          piece_iter = iter([params_state])
+          converted_state = params_state
 
       del params_state
       gc.collect()
 
       if self._raiden_syncs is None:
-        self._raiden_syncs = []
-
-      expected_num_pieces = len(self._raiden_syncs) if self._raiden_syncs else None
-      is_pathways = bool("proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get("JAX_BACKEND_TARGET"))
-      verify_weights = os.environ.get("VERIFY_WEIGHTS", "").lower() == "true"
-      all_metadata = []
-      total_variables = 0
-
-      for piece_idx, piece in enumerate(piece_iter):
-        if piece_idx >= len(self._raiden_syncs):
-          self._raiden_syncs.append(
-              raiden_synchronizer.RaidenSynchronizer(
-                  job_name="trainer",
-                  worker_index=self._raiden_worker_index(piece_idx),
-                  auto_h2d=False,
-                  host_stage=is_pathways,
-                  parallelism=4,
-              )
-          )
-
-        sync = self._raiden_syncs[piece_idx]
-        sync.bind(piece)
-        del piece
-        gc.collect()
-
-        # 4. Initiate Device-to-Host transfer to stage this piece for network
-        # transfer before moving on to the next piece.
-        if sync.active:
-          sync.d2h()
+        is_pathways = bool("proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get("JAX_BACKEND_TARGET"))
+        self._raiden_syncs = [
+            raiden_synchronizer.RaidenSynchronizer(
+                job_name="trainer",
+                worker_index=jax.process_index(),
+                auto_h2d=False,
+                host_stage=is_pathways,
+                parallelism=4,
+            )
+        ]
+
+      sync = self._raiden_syncs[0]
+      sync.bind(converted_state)
+      del converted_state
+      gc.collect()
+      _malloc_trim()
 
-        if verify_weights:
-          logging.info("Source weights checksums (piece %d): %s", piece_idx, sync.checksums())
+      # 4. Initiate Device-to-Host transfer to stage for network transfer.
+      if sync.active:
+        sync.d2h()
 
-        metadata = sync.work_unit_metadata()
-        total_variables += len(metadata.variables)
-        all_metadata.append(metadata)
-        sync.release_host_arrays()
+      verify_weights = os.environ.get("VERIFY_WEIGHTS", "").lower() == "true"
+      if verify_weights:
+        logging.info("Source weights checksums: %s", sync.checksums())
 
-      num_pieces = len(all_metadata)
-      if expected_num_pieces is not None and num_pieces != expected_num_pieces:
-        raise RuntimeError(
-            f"weight-sync piece count changed from {expected_num_pieces} to {num_pieces} "
-            "between rounds; the cached conversion plan should make this impossible "
-            "unless the model/config changed mid-run."
-        )
+      metadata = sync.work_unit_metadata()
+      total_variables = len(metadata.variables)
+      all_metadata = [metadata]
 
       gc.collect()
       _malloc_trim()
 
       logging.info(
-          "Trainer prepared weight sync for step %d: registered %d variables across %d piece(s) on mesh %s",
+          "Trainer prepared weight sync for step %d: registered %d variables across 1 piece(s) on mesh %s",
           self.train_step,
           total_variables,
-          num_pieces,
           all_metadata[0].mesh_axes if all_metadata else None,
       )
       self._last_staged_step = self.train_step
diff --git a/tests/unit/prepare_weight_sync_test.py b/tests/unit/prepare_weight_sync_test.py
index 0976a6efc3..382e28b5df 100644
--- a/tests/unit/prepare_weight_sync_test.py
+++ b/tests/unit/prepare_weight_sync_test.py
@@ -65,7 +65,7 @@ def _make_dummy_metadata(self, num_vars=2):
     return meta
 
   @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
-  def test_streaming_grows_syncs_and_accumulates_metadata(self, mock_sync_cls):
+  def test_sync_binds_converted_state_and_accumulates_metadata(self, mock_sync_cls):
     created_syncs = []
 
     def make_sync(*args, **kwargs):
@@ -79,28 +79,20 @@ def make_sync(*args, **kwargs):
 
     mock_sync_cls.side_effect = make_sync
 
-    pieces = [{"piece_0": 0}, {"piece_1": 1}, {"piece_2": 2}]
-    self.engine._weight_converter.convert_streaming.return_value = iter(pieces)
+    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), 3)
-    self.assertEqual(len(self.engine._raiden_syncs), 3)
-    self.assertEqual(len(created_syncs), 3)
+    self.assertEqual(len(metadata), 1)
+    self.assertEqual(len(self.engine._raiden_syncs), 1)
+    self.assertEqual(len(created_syncs), 1)
 
-    # Worker indices must be unique and properly strided
-    expected_indices = [
-        jax.process_index() * _RAIDEN_WORKER_INDEX_STRIDE + i + 1 for i in range(3)
-    ]
-    actual_indices = [s.worker_index for s in created_syncs]
-    self.assertEqual(actual_indices, expected_indices)
-
-    # Check that bind, d2h, metadata, release were called on each sync
-    for s, p in zip(created_syncs, pieces):
-      s.bind.assert_called_once_with(p)
-      s.d2h.assert_called_once()
-      s.work_unit_metadata.assert_called_once()
-      s.release_host_arrays.assert_called_once()
+    s = created_syncs[0]
+    self.assertEqual(s.worker_index, jax.process_index())
+    s.bind.assert_called_once_with(converted)
+    s.d2h.assert_called_once()
+    s.work_unit_metadata.assert_called_once()
 
   @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
   def test_rebind_reuses_sync_instances(self, mock_sync_cls):
@@ -117,36 +109,20 @@ def make_sync(*args, **kwargs):
     mock_sync_cls.side_effect = make_sync
 
     # Round 1
-    self.engine._weight_converter.convert_streaming.return_value = iter([{"p0": 0}, {"p1": 1}])
+    self.engine._weight_converter.convert.return_value = {"p0": 0}
     self.engine.prepare_weight_sync()
-    self.assertEqual(len(mock_syncs), 2)
+    self.assertEqual(len(mock_syncs), 1)
     first_round_syncs = list(self.engine._raiden_syncs)
 
     # Round 2 at step 1
     self.engine._train_step = 1
-    self.engine._weight_converter.convert_streaming.return_value = iter([{"p0": 0}, {"p1": 1}])
+    self.engine._weight_converter.convert.return_value = {"p0": 0}
     self.engine.prepare_weight_sync()
 
     # No new instances created
-    self.assertEqual(len(mock_syncs), 2)
+    self.assertEqual(len(mock_syncs), 1)
     self.assertEqual(self.engine._raiden_syncs, first_round_syncs)
 
-  @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
-  def test_piece_count_mismatch_between_rounds_raises(self, mock_sync_cls):
-    mock_sync_cls.side_effect = lambda *a, **kw: mock.MagicMock(
-        active=True, work_unit_metadata=mock.MagicMock(return_value=self._make_dummy_metadata())
-    )
-
-    # Round 1 has 2 pieces
-    self.engine._weight_converter.convert_streaming.return_value = iter([{"p0": 0}, {"p1": 1}])
-    self.engine.prepare_weight_sync()
-
-    # Round 2 has 3 pieces
-    self.engine._train_step = 1
-    self.engine._weight_converter.convert_streaming.return_value = iter([{"p0": 0}, {"p1": 1}, {"p2": 2}])
-    with self.assertRaisesRegex(RuntimeError, "weight-sync piece count changed from 2 to 3"):
-      self.engine.prepare_weight_sync()
-
   @mock.patch.dict(os.environ, {"RAIDEN_WEIGHT_SYNC_CHUNKS": "4"})
   @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
   def test_deprecated_chunks_env_var_warning(self, mock_sync_cls):

From aba90c78a99d31ac5c8775a7b755a3d18c8cc2fd Mon Sep 17 00:00:00 2001
From: Yixuan Wang <yixuannwang@google.com>
Date: Fri, 4 Sep 2026 07:26:17 +0000
Subject: [PATCH 7/7] Support Pathways FFI Raiden weight sync, add host memory
 reclamation, and consolidate sync instances

- Under Pathways (JAX_PLATFORMS=proxy), require weight_synchronizer_ffi to avoid client host OOM
- Consolidate to single RaidenSynchronizer instance in MaxTextTrainingEngine
- Add reclaim_host_memory() utility invoking gc.collect() and malloc_trim(0)
- Add weight_sync_debug flag to HyperParameters config
- Update unit tests across maxtext_engine, prepare_weight_sync, and weight_converter
---
 src/maxtext/configs/types.py                  |   2 +-
 .../tunix/weight_mapping/raiden_unscan.py     |   2 +-
 src/maxtext/integration/vllm/convert_utils.py |  34 ++-
 src/maxtext/integration/vllm/moe_padding.py   |   6 +-
 .../integration/vllm/weight_converter.py      | 203 +++---------------
 src/maxtext/training_engine/maxtext_engine.py | 135 ++++++------
 .../unit/maxtext_engine_e2e_test.py           |   2 +-
 .../post_training/unit/maxtext_engine_test.py |   2 +-
 .../unit/weight_converter_test.py             |  22 +-
 tests/unit/prepare_weight_sync_test.py        | 124 ++++++-----
 10 files changed, 218 insertions(+), 314 deletions(-)

diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 3eb0c862f8..7f88a79991 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -2704,7 +2704,7 @@ class VLLM(BaseModel):
           "the legacy transfer_state_directly / transfer_state_with_mappings paths."
       ),
   )
-  rollout_backend: str = Field(
+  rollout_backend: Literal["maxtext", "vllm_torchax"] = Field(
       "maxtext",
       description="Rollout backend for trainer-side weight converter ('maxtext' or 'vllm_torchax').",
   )
diff --git a/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py b/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py
index e372d014f5..d4f9446475 100644
--- a/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py
+++ b/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py
@@ -37,7 +37,7 @@
 """
 
 import gc
-from typing import Any, Iterator, Tuple, List, Dict
+from typing import Any, Iterator, Tuple, List
 
 import jax
 from flax import nnx
diff --git a/src/maxtext/integration/vllm/convert_utils.py b/src/maxtext/integration/vllm/convert_utils.py
index e75d154655..34fe4832a5 100644
--- a/src/maxtext/integration/vllm/convert_utils.py
+++ b/src/maxtext/integration/vllm/convert_utils.py
@@ -160,15 +160,33 @@ class ShapeMismatchError(ValueError):
   """Raised when source and target shapes are incompatible."""
 
 
-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."""
+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"):
-      tgt_dtype = jnp.bfloat16
-    elif tgt_dtype in ("float32", "fp32"):
-      tgt_dtype = jnp.float32
-    else:
-      tgt_dtype = jnp.dtype(tgt_dtype)
+      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."""
+  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)
@@ -546,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.
diff --git a/src/maxtext/integration/vllm/moe_padding.py b/src/maxtext/integration/vllm/moe_padding.py
index fa11df1ed2..ce5847704f 100644
--- a/src/maxtext/integration/vllm/moe_padding.py
+++ b/src/maxtext/integration/vllm/moe_padding.py
@@ -55,9 +55,7 @@ def compute_padded_moe_mlp_dim(
     return hidden_size
 
   if (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)
-    return padded_hidden_size
+    min_required = 2 * num_lanes * moe_mlp_tp_size
+    return next_power_of_two(max(hidden_size, min_required))
 
   return hidden_size
diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py
index 9d196094c6..ad6dfb634c 100644
--- a/src/maxtext/integration/vllm/weight_converter.py
+++ b/src/maxtext/integration/vllm/weight_converter.py
@@ -35,19 +35,13 @@
     _jit_unstack,
     _scanned_sharding_from_per_layer,
     _sharding_summary,
+    normalize_dtype,
+    reclaim_host_memory,
 )
 
 _MOE_MLP_WEIGHTS = frozenset({"wi_0", "wi_1", "wo", "wi"})
 
 
-def _malloc_trim() -> None:
-  try:
-    import ctypes  # pylint: disable=g-import-not-at-top
-    ctypes.CDLL("libc.so.6").malloc_trim(0)
-  except Exception:
-    pass
-
-
 # ==========================================
 # 1. Operations
 # ==========================================
@@ -597,13 +591,6 @@ class ConversionPlanError(WeightConverterError, ValueError):
   """Raised when the source and target trees cannot be fully reconciled."""
 
 
-class UnscanShapeMismatch(WeightConverterError, ValueError):
-  """Raised when an unscan shape mismatch is detected."""
-
-
-class WeightSyncBindError(WeightConverterError, RuntimeError):
-  """Raised when parameter binding to Raiden synchronizer fails."""
-
 
 def _is_non_weight_path(key_tuple: Tuple[Any, ...]) -> bool:
   return any(isinstance(part, str) and part.lstrip("_").startswith(_NON_WEIGHT_PATH_PREFIXES) for part in key_tuple)
@@ -710,6 +697,7 @@ def __init__(
       debug: bool = False,
       prefuse_moe_weights: Optional[bool] = None,
       target_dtype: Optional[Any] = None,
+      is_pathways: Optional[bool] = None,
   ):
     self.config = config
     self.moe_fused_layout = moe_fused_layout
@@ -723,6 +711,14 @@ 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)
 
+    if is_pathways is not None:
+      self.is_pathways = is_pathways
+    else:
+      backend_platform = getattr(jax.devices()[0], "platform", "").lower() if jax.devices() else ""
+      self.is_pathways = (backend_platform == "proxy") or (
+          "proxy" in os.environ.get("JAX_PLATFORMS", "") and bool(os.environ.get("JAX_BACKEND_TARGET"))
+      )
+
     self.cycle = int(getattr(config, "inhomogeneous_layer_cycle_interval", 1) or 1)
     self.num_decoder_layers = int(config.num_decoder_layers)
     if self.num_decoder_layers % self.cycle:
@@ -750,15 +746,7 @@ def __init__(
     )
 
   def _resolve_target_dtype(self):
-    if self.target_dtype is None:
-      return None
-    if isinstance(self.target_dtype, str):
-      if self.target_dtype in ("bfloat16", "bf16"):
-        return jnp.bfloat16
-      if self.target_dtype in ("float32", "fp32"):
-        return jnp.float32
-      return jnp.dtype(self.target_dtype)
-    return self.target_dtype
+    return normalize_dtype(self.target_dtype)
 
   # -------------------------------------------------------------- #
   # Plan construction
@@ -801,23 +789,8 @@ def _build_target_free_plan(
       rest = src_key[idx + 1 :]
 
       if self.cycle == 1:
-        # Homogeneous: ("decoder", "layers", "self_attention", "query", "kernel")
-        is_wi_0 = bool(rest and rest[-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 i in range(self.num_decoder_layers):
-            tgt_key = prefix + (f"layers_{i}",) + rest[:-1] + ("wi",)
-            plan.append(_PlanEntry(tgt_key, (src_key, wi_1_key), i, "fuse_moe"))
-        elif self.prefuse_moe_weights and rest and rest[-1] == "wi_1" and (src_key[:-1] + ("wi_0",) in src_flat):
-          continue
-        else:
-          for i in range(self.num_decoder_layers):
-            tgt_key = prefix + (f"layers_{i}",) + rest
-            plan.append(_PlanEntry(tgt_key, (src_key,), i, "slice"))
-
+        slot = 0
+        suffix = rest
       else:
         # Inhomogeneous hybrid cycle: ("decoder", "layers", "layer_0", "input_layernorm", "scale")
         slot_token = rest[0]
@@ -829,25 +802,25 @@ def _build_target_free_plan(
           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"))
+
+      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
 
@@ -1044,21 +1017,9 @@ def _fuse_moe_bulk_target_free(self, wi_0: Any, wi_1: Any, path: str):
   def _execute_group_target_free(self, group: _PlanGroup, src_flat):
     path = group.source_path
     target_dtype = self._resolve_target_dtype()
-    is_pathways = bool("proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get("JAX_BACKEND_TARGET"))
 
     if group.op == "identity":
       raw_val = src_flat[group.source_keys[0]]
-      if isinstance(raw_val, jax.ShapeDtypeStruct):
-        val = _apply_dtype_cast(raw_val, target_dtype, path)
-        return [(tgt_key, val) for _, tgt_key in group.targets]
-      if is_pathways:
-        np_val = np.asarray(jax.device_get(raw_val))
-        if target_dtype is not None:
-          np_val = np_val.astype(target_dtype)
-        cpu = jax.local_devices(backend="cpu")[0]
-        val = jax.device_put(np_val, cpu)
-        del np_val
-        return [(tgt_key, val) for _, tgt_key in group.targets]
       val = _apply_dtype_cast(raw_val, target_dtype, path)
       return [(tgt_key, val) for _, tgt_key in group.targets]
 
@@ -1071,60 +1032,6 @@ def _execute_group_target_free(self, group: _PlanGroup, src_flat):
     if group.op == "fuse_moe":
       raw_0 = src_flat[group.source_keys[0]]
       raw_1 = src_flat[group.source_keys[1]]
-      if isinstance(raw_0, jax.ShapeDtypeStruct):
-        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]
-
-      if is_pathways:
-        wi_0 = np.asarray(jax.device_get(raw_0))
-        wi_1 = np.asarray(jax.device_get(raw_1))
-        if target_dtype is not None:
-          wi_0 = wi_0.astype(target_dtype)
-          wi_1 = wi_1.astype(target_dtype)
-        self._check_scan_axis(wi_0, path)
-        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
-        )
-        tgt_shape = (wi_0.shape[0], wi_0.shape[2], 2 * target_intermediate)
-        tgt_fused_axis = len(tgt_shape) - 1
-        scan_fused_axis = tgt_fused_axis if tgt_fused_axis < self.scan_axis else tgt_fused_axis + 1
-
-        if self.moe_fused_layout == MoEFusedLayout.PER_SHARD_INTERLEAVE:
-          n_shards = 1
-          target_half_dim = target_intermediate
-          current_total_size = wi_0.shape[scan_fused_axis]
-          chunk_size = current_total_size // n_shards
-          target_chunk_size = target_half_dim // n_shards
-          pad_amount = target_chunk_size - chunk_size
-          if pad_amount > 0:
-            pad_spec = [(0, 0)] * wi_0.ndim
-            pad_spec[scan_fused_axis] = (0, pad_amount)
-            wi_0 = np.pad(wi_0, pad_spec)
-            wi_1 = np.pad(wi_1, pad_spec)
-          fused = np.concatenate([wi_0, wi_1], axis=scan_fused_axis)
-        elif 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 = np.pad(wi_0, pad_spec)
-            wi_1 = np.pad(wi_1, pad_spec)
-          fused = np.concatenate([wi_0, wi_1], axis=scan_fused_axis)
-        else:
-          raise ConversionPlanError(f"Unknown moe_fused_layout: {self.moe_fused_layout!r}")
-
-        cpu = jax.local_devices(backend="cpu")[0]
-        per_block = tuple(
-            jax.device_put(np.ascontiguousarray(fused.take(indices=i, axis=self.scan_axis)), cpu)
-            for i in range(self.num_blocks)
-        )
-        del fused, wi_0, wi_1
-        return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets]
-
       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)
@@ -1132,41 +1039,6 @@ def _execute_group_target_free(self, group: _PlanGroup, src_flat):
 
     # group.op == "slice"
     raw_val = src_flat[group.source_keys[0]]
-    if isinstance(raw_val, jax.ShapeDtypeStruct):
-      val = _apply_dtype_cast(raw_val, target_dtype, path)
-      self._check_scan_axis(val, path)
-      per_block = self._slice_bulk_target_free(val, path)
-      return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets]
-
-    if is_pathways:
-      np_val = np.asarray(jax.device_get(raw_val))
-      if target_dtype is not None:
-        np_val = np_val.astype(target_dtype)
-      self._check_scan_axis(np_val, path)
-      last_key = path.split(".")[-1]
-      if last_key in _MOE_MLP_WEIGHTS and self.padded_base_moe_mlp_dim is not None:
-        if last_key == "wo":
-          intermediate_axis = 2
-          if self.padded_base_moe_mlp_dim > np_val.shape[intermediate_axis]:
-            pad_amount = self.padded_base_moe_mlp_dim - np_val.shape[intermediate_axis]
-            pad_spec = [(0, 0)] * np_val.ndim
-            pad_spec[intermediate_axis] = (0, pad_amount)
-            np_val = np.pad(np_val, pad_spec)
-        elif last_key in ("wi_0", "wi_1", "wi"):
-          intermediate_axis = len(np_val.shape) - 1
-          if self.padded_base_moe_mlp_dim > np_val.shape[intermediate_axis]:
-            pad_amount = self.padded_base_moe_mlp_dim - np_val.shape[intermediate_axis]
-            pad_spec = [(0, 0)] * np_val.ndim
-            pad_spec[intermediate_axis] = (0, pad_amount)
-            np_val = np.pad(np_val, pad_spec)
-      cpu = jax.local_devices(backend="cpu")[0]
-      per_block = tuple(
-          jax.device_put(np.ascontiguousarray(np_val.take(indices=i, axis=self.scan_axis)), cpu)
-          for i in range(self.num_blocks)
-      )
-      del np_val
-      return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets]
-
     val = _apply_dtype_cast(raw_val, target_dtype, path)
     self._check_scan_axis(val, path)
     per_block = self._slice_bulk_target_free(val, path)
@@ -1219,8 +1091,7 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]:
       flat_result = {}
       for piece in self.convert_streaming(src_pytree, target_state=None):
         flat_result.update(traverse_util.flatten_dict(piece))
-      gc.collect()
-      _malloc_trim()
+      reclaim_host_memory()
       return traverse_util.unflatten_dict(flat_result)
 
     src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree))
@@ -1305,8 +1176,7 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]:
     gc.collect()
     nested = traverse_util.unflatten_dict(result)
     del result
-    gc.collect()
-    _malloc_trim()
+    reclaim_host_memory()
     return jax.tree_util.tree_map(
         lambda x: nnx.Param(x) if not isinstance(x, (nnx.Param, nnx.Variable)) else x,
         nested,
@@ -1330,8 +1200,6 @@ def convert_streaming(
 
     src_flat = traverse_util.flatten_dict(_to_pure_dict(src_pytree))
     src_flat, src_root = _strip_root(src_flat, "base")
-    if not src_root:
-      src_root = ("base",)
 
     if self._plan is None:
       self._plan = self._build_target_free_plan(src_flat)
@@ -1357,8 +1225,7 @@ def convert_streaming(
       )
 
     del src_flat
-    gc.collect()
-    _malloc_trim()
+    reclaim_host_memory()
 
 
 def _rekey_to_target(flat_dotted: Dict[str, Any], target_state: Any) -> Dict[str, Any]:
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index c9b6fdae7c..a03c2639a0 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -38,6 +38,7 @@
 from maxtext.common import train_state_nnx
 from maxtext.configs import pyconfig
 from maxtext.integration.tunix.weight_mapping import raiden_unscan
+from maxtext.integration.vllm.convert_utils import reclaim_host_memory
 from maxtext.trainers.pre_train import train as maxtext_train
 from maxtext.training_engine import abstract_engine
 from maxtext.training_engine import checkpointing
@@ -66,16 +67,6 @@
     "~2 ms for the pure-state equivalent. Logged once per engine instance."
 )
 
-_RAIDEN_WORKER_INDEX_STRIDE = 10_000  # >> any plausible piece count (dozens-to-low-hundreds of groups)
-
-
-def _malloc_trim() -> None:
-  try:
-    import ctypes  # pylint: disable=g-import-not-at-top
-    ctypes.CDLL("libc.so.6").malloc_trim(0)
-  except Exception:
-    pass
-
 
 def _is_jax_dynamic(value: Any) -> bool:
   """Returns True if `value` can cross a `jax.jit` boundary as a traced argument.
@@ -418,10 +409,9 @@ def __init__(
     )
     self._metrics_recorder = metrics_module.MetricsRecorder()
     self._throttler = inflight_throttler.InflightThrottler(config=self._config)
-    self._raiden_syncs: Any = None
+    self._raiden_sync: Any = None
     self._last_staged_step: Optional[int] = None
     self._staged_metadata: Any = None
-    self._warned_raiden_sync_chunks: bool = False
     vllm_cfg = getattr(self._config, "vllm", {})
     if isinstance(vllm_cfg, dict):
       vllm_use_wc = vllm_cfg.get("use_weight_converter", False)
@@ -1464,9 +1454,6 @@ def _get_trainable_params_state(self) -> Any:
       return nnx.state(model, nnx.Param)
     return self.model
 
-  def _raiden_worker_index(self, piece_idx: int) -> int:
-    return jax.process_index() * _RAIDEN_WORKER_INDEX_STRIDE + piece_idx + 1
-
   def prepare_weight_sync(
       self,
       staging_transport: str = "raiden",
@@ -1497,7 +1484,7 @@ def prepare_weight_sync(
         ) from exc
 
       if (
-          self._raiden_syncs is not None
+          self._raiden_sync is not None
           and self._last_staged_step == self.train_step
           and self._staged_metadata is not None
       ):
@@ -1508,25 +1495,12 @@ def prepare_weight_sync(
         )
         return self._staged_metadata
 
-      if self._raiden_syncs is not None:
-        for sync in self._raiden_syncs:
-          sync.release_host_arrays()
-        gc.collect()
-        _malloc_trim()
-
       # 1. Drain all in-flight TPU computations to ensure weights are fully updated
       self._throttler.wait_for_all()
-      gc.collect()
+      reclaim_host_memory()
 
       # 2. Extract clean trainable parameters
       params_state = self._get_trainable_params_state()
-      piece_batch = max(1, int(os.environ.get("RAIDEN_STREAM_PIECE_BATCH", "1")))
-      if "RAIDEN_WEIGHT_SYNC_CHUNKS" in os.environ and not self._warned_raiden_sync_chunks:
-        logging.warning(
-            "RAIDEN_WEIGHT_SYNC_CHUNKS is deprecated and no longer affects Raiden staging; "
-            "use RAIDEN_STREAM_PIECE_BATCH instead."
-        )
-        self._warned_raiden_sync_chunks = True
 
       if self._use_weight_converter:
         if self._weight_converter is None:
@@ -1556,46 +1530,69 @@ def prepare_weight_sync(
           converted_state = params_state
 
       del params_state
-      gc.collect()
-
-      if self._raiden_syncs is None:
-        is_pathways = bool("proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get("JAX_BACKEND_TARGET"))
-        self._raiden_syncs = [
-            raiden_synchronizer.RaidenSynchronizer(
-                job_name="trainer",
-                worker_index=jax.process_index(),
-                auto_h2d=False,
-                host_stage=is_pathways,
-                parallelism=4,
-            )
-        ]
-
-      sync = self._raiden_syncs[0]
-      sync.bind(converted_state)
+      reclaim_host_memory()
+
+      # 3. Bind parameters to the Raiden transport. Construct the synchronizer
+      # once, matching the persistent-instance-per-cycle pattern the rebind
+      # optimization depends on.
+      #
+      # Under Pathways (JAX_PLATFORMS=proxy + JAX_BACKEND_TARGET set), 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.
+      backend_platform = getattr(jax.devices()[0], "platform", "").lower() if jax.devices() else ""
+      is_pathways = (backend_platform == "proxy") or (
+          "proxy" in os.environ.get("JAX_PLATFORMS", "") and bool(os.environ.get("JAX_BACKEND_TARGET"))
+      )
+      if is_pathways and getattr(raiden_synchronizer, "_raiden_ffi", None) is None:
+        raise RuntimeError(
+            "Under Pathways (JAX_PLATFORMS=proxy), Raiden weight synchronization "
+            "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",
+            worker_index=jax.process_index(),
+            auto_h2d=False,
+            host_stage=is_pathways,
+            parallelism=4,
+        )
+
+      self._raiden_sync.bind(converted_state)
       del converted_state
-      gc.collect()
-      _malloc_trim()
+      reclaim_host_memory()
 
-      # 4. Initiate Device-to-Host transfer to stage for network transfer.
-      if sync.active:
-        sync.d2h()
+      # 4. Initiate Device-to-Host transfer to stage weights for network transfer.
+      if is_pathways or self._raiden_sync.active:
+        self._raiden_sync.d2h()
 
       verify_weights = os.environ.get("VERIFY_WEIGHTS", "").lower() == "true"
       if verify_weights:
-        logging.info("Source weights checksums: %s", sync.checksums())
+        logging.info("Source weights checksums: %s", self._raiden_sync.checksums())
 
-      metadata = sync.work_unit_metadata()
+      metadata = self._raiden_sync.work_unit_metadata()
       total_variables = len(metadata.variables)
       all_metadata = [metadata]
 
-      gc.collect()
-      _malloc_trim()
+      reclaim_host_memory()
+
+      try:
+        import resource  # pylint: disable=g-import-not-at-top
+        rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0
+        mem_info = f", host memory max RSS: {rss_mb:.1f} MB"
+      except Exception:  # pylint: disable=broad-exception-caught
+        mem_info = ""
 
       logging.info(
-          "Trainer prepared weight sync for step %d: registered %d variables across 1 piece(s) on mesh %s",
+          "Trainer prepared weight sync for step %d: registered %d variables on mesh %s%s",
           self.train_step,
           total_variables,
-          all_metadata[0].mesh_axes if all_metadata else None,
+          metadata.mesh_axes,
+          mem_info,
       )
       self._last_staged_step = self.train_step
       self._staged_metadata = all_metadata
@@ -1610,21 +1607,23 @@ 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_syncs:
-      for sync in self._raiden_syncs:
-        logging.vlog(1, "Trainer Raiden metrics: %s", sync.metrics())
-        sync.release_host_arrays()
-    gc.collect()
-    _malloc_trim()
+    if self._raiden_sync:
+      logging.vlog(1, "Trainer Raiden metrics: %s", self._raiden_sync.metrics())
+    reclaim_host_memory()
+    try:
+      import resource  # pylint: disable=g-import-not-at-top
+      rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0
+      logging.info("Trainer released weight sync: host memory max RSS: %.1f MB", rss_mb)
+    except Exception:  # pylint: disable=broad-exception-caught
+      pass
     return True
 
   def close(self) -> None:
     """Closes the trainer, writes buffered metrics and final checkpoint."""
-    if self._raiden_syncs:
-      for sync in self._raiden_syncs:
-        if hasattr(sync, "close"):
-          sync.close()
-      self._raiden_syncs = None
+    if self._raiden_sync:
+      if hasattr(self._raiden_sync, "close"):
+        self._raiden_sync.close()
+      self._raiden_sync = None
     self._last_staged_step = None
     self._staged_metadata = None
 
diff --git a/tests/post_training/unit/maxtext_engine_e2e_test.py b/tests/post_training/unit/maxtext_engine_e2e_test.py
index 54896ecc6d..e227351a16 100644
--- a/tests/post_training/unit/maxtext_engine_e2e_test.py
+++ b/tests/post_training/unit/maxtext_engine_e2e_test.py
@@ -40,7 +40,7 @@
 # failure where it does not -- rather than passing or failing on which tunix happens to
 # be installed.
 try:
-  importlib.import_module("tunix.experimental.worker.raiden_synchronizer")
+  importlib.import_module("tunix.experimental.weight_sync.raiden_synchronizer")
   _RAIDEN_AVAILABLE = True
 except ImportError:
   _RAIDEN_AVAILABLE = False
diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py
index 3c28a586fb..b834c7b459 100644
--- a/tests/post_training/unit/maxtext_engine_test.py
+++ b/tests/post_training/unit/maxtext_engine_test.py
@@ -1344,7 +1344,7 @@ def test_prepare_weight_sync_raises_when_raiden_is_unavailable(self):
 
     # Setting the entry to None makes `from ... import raiden_synchronizer` raise
     # ImportError, which is what an installed tunix without the module does.
-    with mock.patch.dict(sys.modules, {"tunix.experimental.worker.raiden_synchronizer": None}):
+    with mock.patch.dict(sys.modules, {"tunix.experimental.weight_sync.raiden_synchronizer": None}):
       with self.assertRaisesRegex(RuntimeError, "raiden_synchronizer"):
         t.prepare_weight_sync()
 
diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py
index 9c58974041..912d83e900 100644
--- a/tests/post_training/unit/weight_converter_test.py
+++ b/tests/post_training/unit/weight_converter_test.py
@@ -537,13 +537,14 @@ def test_case_1_homogeneous_target_free_unroll(self):
     }
     converter = WeightConverter(config=cfg, rollout_backend="maxtext")
     out = converter.convert(source, target_state=None)
-    self.assertIn("token_embedder", out)
-    self.assertIn("decoder", out)
+    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["decoder"])
-      scale = getattr(out["decoder"][layer_key]["input_layernorm"]["scale"], "value", out["decoder"][layer_key]["input_layernorm"]["scale"])
-      query = getattr(out["decoder"][layer_key]["self_attention"]["query"]["kernel"], "value", out["decoder"][layer_key]["self_attention"]["query"]["kernel"])
+      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))
 
@@ -552,10 +553,11 @@ def test_case_2_hybrid_cycle_target_free_unroll(self):
     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["decoder"][f"layers_{layer}"]["input_layernorm"]["scale"], "value", out["decoder"][f"layers_{layer}"]["input_layernorm"]["scale"])
+      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))
 
@@ -573,8 +575,9 @@ def test_case_3_prefused_moe_target_free(self):
     source = _source_tree(True)
     converter = MaxTextToMaxTextConverter(cfg, prefuse_moe_weights=True)
     out = converter.convert(source, target_state=None)
-    wi = getattr(out["decoder"]["layers_0"]["moe_block"]["wi"], "value", out["decoder"]["layers_0"]["moe_block"]["wi"])
-    wo = getattr(out["decoder"]["layers_0"]["moe_block"]["wo"], "value", out["decoder"]["layers_0"]["moe_block"]["wo"])
+    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))
 
@@ -588,10 +591,11 @@ def to_struct(x):
     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["decoder"]["layers_0"]["moe_block"]["wi"], "value", out["decoder"]["layers_0"]["moe_block"]["wi"])
+    wi = getattr(out_root["decoder"]["layers_0"]["moe_block"]["wi"], "value", out_root["decoder"]["layers_0"]["moe_block"]["wi"])
     self.assertEqual(wi.shape, (EXPERTS, EMB, 32))
 
   def test_case_5_host_memory_profiling(self):
diff --git a/tests/unit/prepare_weight_sync_test.py b/tests/unit/prepare_weight_sync_test.py
index 382e28b5df..c08a8c7858 100644
--- a/tests/unit/prepare_weight_sync_test.py
+++ b/tests/unit/prepare_weight_sync_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Unit tests for MaxTextTrainingEngine.prepare_weight_sync streaming logic."""
+"""Unit tests for MaxTextTrainingEngine.prepare_weight_sync single synchronizer logic."""
 
 import os
 os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=8")
@@ -32,7 +32,7 @@
 
 import jax
 import jax.numpy as jnp
-from maxtext.training_engine.maxtext_engine import MaxTextTrainingEngine, _RAIDEN_WORKER_INDEX_STRIDE
+from maxtext.training_engine.maxtext_engine import MaxTextTrainingEngine
 
 
 class PrepareWeightSyncTest(unittest.TestCase):
@@ -41,7 +41,7 @@ def setUp(self):
     super().setUp()
     # Create engine instance without running heavy __init__
     self.engine = MaxTextTrainingEngine.__new__(MaxTextTrainingEngine)
-    self.engine._raiden_syncs = None
+    self.engine._raiden_sync = None
     self.engine._last_staged_step = None
     self.engine._staged_metadata = None
     self.engine._train_step = 0
@@ -55,7 +55,6 @@ def setUp(self):
     self.engine._use_weight_converter = True
     self.engine._weight_converter = mock.MagicMock()
     self.engine._rollout_backend = "maxtext"
-    self.engine._warned_raiden_sync_chunks = False
     self.engine._get_trainable_params_state = mock.MagicMock(return_value={"layer": jnp.zeros((4, 4))})
 
   def _make_dummy_metadata(self, num_vars=2):
@@ -65,19 +64,12 @@ def _make_dummy_metadata(self, num_vars=2):
     return meta
 
   @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
-  def test_sync_binds_converted_state_and_accumulates_metadata(self, mock_sync_cls):
-    created_syncs = []
-
-    def make_sync(*args, **kwargs):
-      s = mock.MagicMock()
-      s.active = True
-      s.worker_index = kwargs.get("worker_index")
-      s.work_unit_metadata.return_value = self._make_dummy_metadata(num_vars=2)
-      s.checksums.return_value = {}
-      created_syncs.append(s)
-      return s
-
-    mock_sync_cls.side_effect = make_sync
+  def test_single_synchronizer_creation_and_binding(self, mock_sync_cls):
+    mock_sync = mock.MagicMock()
+    mock_sync.active = True
+    mock_sync.work_unit_metadata.return_value = self._make_dummy_metadata(num_vars=2)
+    mock_sync.checksums.return_value = {}
+    mock_sync_cls.return_value = mock_sync
 
     converted = {"param_0": 0, "param_1": 1}
     self.engine._weight_converter.convert.return_value = converted
@@ -85,56 +77,82 @@ def make_sync(*args, **kwargs):
     metadata = self.engine.prepare_weight_sync()
 
     self.assertEqual(len(metadata), 1)
-    self.assertEqual(len(self.engine._raiden_syncs), 1)
-    self.assertEqual(len(created_syncs), 1)
+    self.assertIs(self.engine._raiden_sync, mock_sync)
+    mock_sync_cls.assert_called_once_with(
+        job_name="trainer",
+        worker_index=jax.process_index(),
+        auto_h2d=False,
+        host_stage=False,
+        parallelism=4,
+    )
 
-    s = created_syncs[0]
-    self.assertEqual(s.worker_index, jax.process_index())
-    s.bind.assert_called_once_with(converted)
-    s.d2h.assert_called_once()
-    s.work_unit_metadata.assert_called_once()
+    self.engine._weight_converter.convert.assert_called_once()
+    mock_sync.bind.assert_called_once_with(converted)
+    mock_sync.d2h.assert_called_once()
+    mock_sync.work_unit_metadata.assert_called_once()
 
   @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
-  def test_rebind_reuses_sync_instances(self, mock_sync_cls):
-    mock_syncs = []
-
-    def make_sync(*args, **kwargs):
-      s = mock.MagicMock()
-      s.active = True
-      s.worker_index = kwargs.get("worker_index")
-      s.work_unit_metadata.return_value = self._make_dummy_metadata(num_vars=2)
-      mock_syncs.append(s)
-      return s
-
-    mock_sync_cls.side_effect = make_sync
+  def test_rebind_reuses_single_sync_instance(self, mock_sync_cls):
+    mock_sync = mock.MagicMock()
+    mock_sync.active = True
+    mock_sync.work_unit_metadata.return_value = self._make_dummy_metadata(num_vars=2)
+    mock_sync.checksums.return_value = {}
+    mock_sync_cls.return_value = mock_sync
 
     # Round 1
     self.engine._weight_converter.convert.return_value = {"p0": 0}
     self.engine.prepare_weight_sync()
-    self.assertEqual(len(mock_syncs), 1)
-    first_round_syncs = list(self.engine._raiden_syncs)
+    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()
 
-    # No new instances created
-    self.assertEqual(len(mock_syncs), 1)
-    self.assertEqual(self.engine._raiden_syncs, first_round_syncs)
+    # 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()
 
-  @mock.patch.dict(os.environ, {"RAIDEN_WEIGHT_SYNC_CHUNKS": "4"})
-  @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
-  def test_deprecated_chunks_env_var_warning(self, mock_sync_cls):
-    mock_sync_cls.side_effect = lambda *a, **kw: mock.MagicMock(
-        active=True, work_unit_metadata=mock.MagicMock(return_value=self._make_dummy_metadata())
-    )
-    self.engine._weight_converter.convert_streaming.return_value = iter([{"p0": 0}])
-    with mock.patch("absl.logging.warning") as mock_warn:
-      self.engine.prepare_weight_sync()
-      mock_warn.assert_called()
-      self.assertTrue(any("RAIDEN_WEIGHT_SYNC_CHUNKS is deprecated" in str(call) for call in mock_warn.call_args_list))
-    self.assertTrue(self.engine._warned_raiden_sync_chunks)
+    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__":