[Raiden Weight Sync 6/7] Integrate Raiden weight sync in MaxTextTrainingEngine with FFI support by YixuanWang-99 · Pull Request #5171 · AI-Hypercomputer/maxtext

YixuanWang-99 · 2026-09-09T01:22:48Z

Overview

Part of the stacked Raiden weight-sync enablement PR chain replacing #5089.
Pairs with Tunix [T2a] (google/tunix).

Stack:

[M7] [Raiden Weight Sync 1/7] Add cross-repo drift guard and FFI resolution tests #5166: Cross-repo drift guard
[M1] [Raiden Weight Sync 2/7] Support inhomogeneous layer cycle interval in raiden_unscan #5167: Inhomogeneous layer cycle support in raiden_unscan
[M2] [Raiden Weight Sync 3/7] MoE 128-lane layout and padding for TPU GMM kernels #5168: MoE 128-lane layout and padding (pairs with Tunix T5)
[M3] [Raiden Weight Sync 4/7] Target-free and streaming weight conversion for vLLM #5169: Target-free streaming weight conversion
[M5] [Raiden Weight Sync 5/7] Pass weight_dtype in Qwen3NextSparseMoeBlock shared expert gate #5170: Shared expert gate weight_dtype fix
[M4] AI-Hypercomputer/maxtext (this PR): TrainingEngine Raiden FFI integration (pairs with Tunix T2a)
[M6] AI-Hypercomputer/maxtext: Rollout cleanup

Details

Review thread addressed: Accesses config properties (use_weight_converter, rollout_backend, weight_sync_debug) directly as typed attributes rather than using getattr (addressing @igorts's comment).
Declares use_raiden_ffi: Optional[bool] = Field(None, ...) on config types.
The Step 0 Fix: Replaces conflicting ad-hoc os.environ reads with Tunix's exported resolve_use_ffi and normalize_host_stage.
Rebinds parameters across steps using a single persistent synchronizer instance to eliminate redundant setup overhead.
Adds comprehensive unit testing in tests/unit/prepare_weight_sync_test.py.

Verification

pytest tests/unit/prepare_weight_sync_test.py (5/5 passed).
pytest tests/post_training/unit/maxtext_engine_test.py (54/54 passed).

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

gemini-code-assist

Code Review

This pull request introduces several enhancements to weight synchronization, checkpointing, and configuration in MaxText. Key changes include adding rollout_backend and use_raiden_ffi configurations, integrating a weight converter, implementing a warning for replicated batch dimensions, and refining exception handling during checkpoint restoration. Additionally, new unit tests were added to verify the weight synchronization logic. The review feedback suggests raising a RuntimeError if FFI is explicitly requested but unavailable, logging a warning when running on Pathways without FFI, and adding inhomogeneous_layer_cycle_interval to the test mock configuration to avoid potential attribute errors.

gemini-code-assist · 2026-09-09T01:24:19Z

+        if config_use_ffi is not None:
+          use_raiden_ffi = config_use_ffi and ffi_available


If use_raiden_ffi is explicitly requested via configuration (use_raiden_ffi=True), but the FFI library is not available, the current implementation silently falls back to running without FFI. This can lead to silent performance degradation or out-of-memory (OOM) errors on Pathways. It is safer to raise a RuntimeError when an explicitly requested feature cannot be enabled.

Suggested change

if config_use_ffi is not None:

use_raiden_ffi = config_use_ffi and ffi_available

if config_use_ffi is not None:

if config_use_ffi and not ffi_available:

raise RuntimeError(

"Raiden FFI was explicitly requested (use_raiden_ffi=True), but "

"weight_synchronizer_ffi is not available. Please ensure a compatible "

"tpu_raiden_jax wheel with FFI support is installed."

)

use_raiden_ffi = config_use_ffi

gemini-code-assist · 2026-09-09T01:24:19Z

+      if is_pathways and not use_raiden_ffi:
+        get_ffi = getattr(raiden_synchronizer, "_get_raiden_ffi", None)
+        ffi_available = (get_ffi() is not None) if get_ffi else False
+        if not ffi_available:
+          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."
+          )


When running under Pathways with use_raiden_ffi disabled, but ffi_available is True, the code proceeds without FFI using host staging. Since running without FFI on Pathways is highly discouraged due to client host OOM and proxy staging timeouts, we should log a warning to alert the user.

Suggested change

if is_pathways and not use_raiden_ffi:

get_ffi = getattr(raiden_synchronizer, "_get_raiden_ffi", None)

ffi_available = (get_ffi() is not None) if get_ffi else False

if not ffi_available:

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 is_pathways and not use_raiden_ffi:

get_ffi = getattr(raiden_synchronizer, "_get_raiden_ffi", None)

ffi_available = (get_ffi() is not None) if get_ffi else False

if not ffi_available:

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."

)

else:

logging.warning(

"Under Pathways (JAX_PLATFORMS=proxy), running Raiden weight synchronization "

"without FFI is highly discouraged as it can cause client host OOM and proxy "

"staging timeouts. Consider enabling use_raiden_ffi."

)

gemini-code-assist · 2026-09-09T01:24:19Z

+    self.engine._config = pytypes.SimpleNamespace(
+        scan_layers=False,
+        num_decoder_layers=2,
+        param_scan_axis=1,
+        weight_sync_debug=False,
+        use_raiden_ffi=None,
+    )


The mock config in the unit test is missing inhomogeneous_layer_cycle_interval. If scan_layers is ever set to True in future test modifications, this will cause an AttributeError when accessing self._config.inhomogeneous_layer_cycle_interval in prepare_weight_sync. Adding it makes the test setup more robust.

Suggested change

self.engine._config = pytypes.SimpleNamespace(

scan_layers=False,

num_decoder_layers=2,

param_scan_axis=1,

weight_sync_debug=False,

use_raiden_ffi=None,

)

self.engine._config = pytypes.SimpleNamespace(

scan_layers=False,

num_decoder_layers=2,

param_scan_axis=1,

inhomogeneous_layer_cycle_interval=1,

weight_sync_debug=False,

use_raiden_ffi=None,

)

SurbhiJainUSC · 2026-09-09T16:58:51Z

          "rather than data movement. Adds one barrier per sync."
      ),
  )
+  use_raiden_ffi: Optional[bool] = Field(


When should we set this flag to False?

SurbhiJainUSC · 2026-09-09T17:00:02Z

+    max_logging.log(f"Skipping raiden worker sync patch: {e}")
+
+
+patch_raiden_worker_h2d()


Is this required to call here?

SurbhiJainUSC · 2026-09-09T17:01:11Z

      return 0
    try:
      metadata = self._checkpoint_manager.metadata(step)
-    except Exception as e:  # pylint: disable=broad-except


Why this needs to be changed?

SurbhiJainUSC · 2026-09-09T17:02:20Z

      )
    except Exception as e:  # pylint: disable=broad-except
      logging.exception("Failed to restore checkpoint: %s", e)
-      return None, None, None


Why do we need to change this?

SurbhiJainUSC · 2026-09-09T17:06:10Z

      metadata: Checkpoint metadata payload from Orchestrator.
      **kwargs: Additional checkpoint saving options.
    """
+    if os.environ.get("DISABLE_CHECKPOINTING", "false").lower() in ("true", "1"):


This is controlled by orchestrator, I think we should not gate this API behind any flag.

…g and assertion guard - Access config fields directly without getattr - Add use_raiden_ffi assertion guard in configs/types.py and prepare_weight_sync - Remove deprecated resolve_use_ffi and normalize_host_stage calls - Simplify RaidenSynchronizer instantiation and update active D2H check - Set default use_weight_converter to true in base.yml - Support single synchronizer lifecycle and rebind - Add comprehensive prepare_weight_sync unit tests

…scan - Support cycle_interval in unscan_layers - Raise descriptive ValueError naming the key and missing cycle-slot prefix when slot is None and cycle_interval > 1 - Add unit tests for unscan_layers and inhomogeneous layer cycles

- Add 128-lane interleaving and padding in convert_utils - Remove dead scan_fused_axis parameter - Add compute_padded_moe_mlp_dim in moe_padding - Add unit tests for convert_utils padding and fusion

…dels - Support target-free weight conversion directly to vLLM format - Enable streaming parameter conversion to minimize host peak RSS - Add comprehensive test suite in weight_converter_test

…g and FFI support - Access config fields directly without getattr - Integrate Tunix FFI resolution and host staging normalization - Add use_raiden_ffi field to config types - Support single synchronizer lifecycle and rebind - Add comprehensive prepare_weight_sync tests

This was referenced Sep 9, 2026

[Raiden Weight Sync 7/7] Clean up rollout path and remove redundant weight sync logic #5172

Open

[Trellis]Qwen3.5-35B; Weight Conversion and Raiden Weight Sync Integration #5089

Open

gemini-code-assist Bot reviewed Sep 9, 2026

View reviewed changes

YixuanWang-99 force-pushed the yixuann-m4-engine-raiden branch from 9844f2e to 6786aae Compare September 9, 2026 02:49

SurbhiJainUSC reviewed Sep 9, 2026

View reviewed changes

YixuanWang-99 force-pushed the yixuann-m4-engine-raiden branch from 6786aae to a6ade7c Compare September 9, 2026 17:56

YixuanWang-99 force-pushed the yixuann-m5-decoder-unroll branch from 890ab06 to 29e3d54 Compare September 9, 2026 17:56

SurbhiJainUSC force-pushed the yixuann-m5-decoder-unroll branch from 29e3d54 to b033d1a Compare September 9, 2026 18:44

SurbhiJainUSC requested a review from parambole as a code owner September 9, 2026 18:44

YixuanWang-99 force-pushed the yixuann-m4-engine-raiden branch from a6ade7c to ff6a0eb Compare September 9, 2026 19:05

YixuanWang-99 force-pushed the yixuann-m5-decoder-unroll branch from b033d1a to 95adff1 Compare September 9, 2026 19:05

YixuanWang-99 force-pushed the yixuann-m4-engine-raiden branch from ff6a0eb to 75ae687 Compare September 9, 2026 19:13

YixuanWang-99 force-pushed the yixuann-m5-decoder-unroll branch from 95adff1 to 20f7c50 Compare September 9, 2026 19:13

SurbhiJainUSC force-pushed the yixuann-m5-decoder-unroll branch from 20f7c50 to 212f591 Compare September 9, 2026 19:37

Base automatically changed from yixuann-m5-decoder-unroll to main September 9, 2026 20:10

YixuanWang-99 force-pushed the yixuann-m4-engine-raiden branch 2 times, most recently from 7b3bf86 to 870c66b Compare September 9, 2026 20:53

YixuanWang-99 changed the base branch from main to yixuann-m3-weight-converter September 9, 2026 20:53

A9isha force-pushed the yixuann-m3-weight-converter branch from db7d733 to b948c90 Compare September 9, 2026 20:57

A9isha requested a review from xuefgu as a code owner September 9, 2026 20:57

A9isha force-pushed the yixuann-m3-weight-converter branch 2 times, most recently from e97ce8f to 9e415a2 Compare September 9, 2026 21:24

SurbhiJainUSC changed the base branch from yixuann-m3-weight-converter to main September 9, 2026 22:00

YixuanWang-99 force-pushed the yixuann-m4-engine-raiden branch from 870c66b to 1b747da Compare September 10, 2026 04:07

YixuanWang-99 added 4 commits September 10, 2026 05:27

feat(vllm): implement MoE 128-lane chunking and padding for GMM layout

67cdbbc

- Add 128-lane interleaving and padding in convert_utils - Remove dead scan_fused_axis parameter - Add compute_padded_moe_mlp_dim in moe_padding - Add unit tests for convert_utils padding and fusion

feat(vllm): add target-free streaming weight converter for MaxText mo…

e84826e

…dels - Support target-free weight conversion directly to vLLM format - Enable streaming parameter conversion to minimize host peak RSS - Add comprehensive test suite in weight_converter_test

Navigation Menu

Uh oh!

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

[Raiden Weight Sync 6/7] Integrate Raiden weight sync in MaxTextTrainingEngine with FFI support#5171
YixuanWang-99 wants to merge 5 commits into
mainfrom
yixuann-m4-engine-raiden

YixuanWang-99 commented Sep 9, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

SurbhiJainUSC Sep 9, 2026

Uh oh!

SurbhiJainUSC Sep 9, 2026

Uh oh!

SurbhiJainUSC Sep 9, 2026

Uh oh!

SurbhiJainUSC Sep 9, 2026

Uh oh!

SurbhiJainUSC Sep 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		if config_use_ffi is not None:
		use_raiden_ffi = config_use_ffi and ffi_available

		max_logging.log(f"Skipping raiden worker sync patch: {e}")


		patch_raiden_worker_h2d()

Uh oh!

Conversation

YixuanWang-99 commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Details

Verification

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

SurbhiJainUSC Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

SurbhiJainUSC Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

SurbhiJainUSC Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

SurbhiJainUSC Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

SurbhiJainUSC Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

YixuanWang-99 commented Sep 9, 2026 •

edited

Loading