feat(training_engine): support Raiden-FFI weight synchronization and adopt single synchronizer

[igorts-git]

Description

This PR integrates Raiden-FFI device-to-host weight synchronization into MaxTextTrainingEngine and eliminates multi-chunk weight staging in favor of a single persistent synchronizer.
This PR depends on the Raiden-FFI integration in Tunix, which has now been merged into main via google/tunix#2059. Certain APIs have change, thus we update the tunix package version to make sure that CI tests are not broken by this PR.

Additionally, one needs Pathways images that include the required changes.
I used these custom Pathways images to run:
server: us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/shauryag/unsanitized_server:raiden_20260812
proxy_server: us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/shauryag/unsanitized_proxy_server:raiden_20260812

Background & Rationale

Under Pathways (JAX_PLATFORMS=proxy), trainer parameters are proxy-backed. Previously, we synchronized weights by staging them onto client host RAM (host_stage=is_pathways) across multiple sequential chunks (RAIDEN_WEIGHT_SYNC_CHUNKS). At scale, host staging created extreme client RAM pressure, risk of host OOMs, and multi-minute proxy staging timeouts.
With Raiden-FFI (tpu_raiden_jax), RaidenSynchronizer binds directly to TPU device buffers on Pathways TPU workers without client host CPU staging. This eliminates the need for multi-chunk staging and allows all parameters to be bound and staged in a single operation.

Key Changes

1. Enforce Raiden FFI under Pathways: When is_pathways is detected, FFI is enforced and fails fast with an informative error if _raiden_ffi is missing in tunix.experimental.weight_sync.raiden_synchronizer.
2. Single RaidenSynchronizer instance: Removed _split_into_chunks and multi-chunk staging loops. prepare_weight_sync() now uses a single persistent RaidenSynchronizer instance, returning a single [metadata] object.
3. Lifecycle management: Updated release_weight_sync() and close() to manage self._raiden_sync instead of iterating over a chunk list.
4. Staging execution: Explicitly triggers D2H transfer under Pathways (is_pathways or self._raiden_sync.active) to ensure FFI transport initialization and metadata registration, while gracefully bypassing native D2H when running in CPU test environments lacking _ws_lib.
5. E2E Test Configuration: Configured "scan_layers": False in tests/post_training/unit/maxtext_engine_e2e_test.py's setup_config() so that prepare_weight_sync does not attempt to unscan layers on the 1D DummyNNXModel.

Tests

Manual invocation of qwen3.5-0.6B with Pathways trainer using 16 chips and a single rollout worker.

• pytest tests/post_training/unit/maxtext_engine_test.py (46/46 passed)
• pytest tests/post_training/unit/maxtext_engine_constructor_test.py (5/5 passed)
• pytest tests/post_training/unit/maxtext_engine_e2e_test.py (1/1 passed)
• pre-commit run --from-ref origin/main --to-ref HEAD (all checks passed: codespell, pylint, pyink)

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 refactors the Raiden weight synchronization in maxtext_engine.py by removing parameter chunking and updating the import path for raiden_synchronizer. It introduces support for FFI-based synchronization under Pathways to prevent client host OOMs and improves error handling by raising explicit exceptions instead of returning empty metadata. Additionally, a test configuration in maxtext_engine_e2e_test.py is updated. Feedback on the changes points out that an error message still references the old import path instead of the newly updated one.

[github-actions]

🤖 Hi @igorts-git, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This pull request introduces valuable improvements to MaxTextTrainingEngine's Raiden weight synchronization by supporting the modern Raiden-FFI device-to-host path under Pathways and moving from a multi-chunk staging model to a single persistent synchronizer. The overall quality of the changes is high, offering a cleaner implementation that mitigates client host OOM risks and proxy transfer timeouts at scale.

🔍 General Feedback

• Simplification: Eliminating multi-chunk splitting and staging significantly simplifies parameter binding and memory management.
• Robustness: Raising explicit RuntimeError and ValueError exceptions for missing/misconfigured transports or missing _raiden_ffi modules helps fail-fast and provides better diagnostics.
• Code Cleanliness: The redundant duplicate initialization of self._raiden_syncs was properly resolved.

[codecov]

Codecov Report

❌ Patch coverage is 51.85185% with 13 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/training_engine/maxtext_engine.py	51.85%	9 Missing and 4 partials ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @igorts-git, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This pull request introduces a major optimization to the training engine by integrating Raiden-FFI device-to-host weight synchronization and adopting a single persistent synchronizer. This successfully resolves extreme host RAM pressure, potential OOMs, and transfer timeouts under Pathways scaling. Overall, the implementation is highly focused and architecturally sound.

🔍 General Feedback

• Single Synchronizer Simplification: Eliminating the complex multi-chunking and splitting logic in favor of a single persistent synchronizer is a fantastic simplification that significantly reduces code complexity.
• Fail-fast FFI Enforcement: The explicit validation of _raiden_ffi presence under Pathways provides highly informative error handling that will save debugging time in incompatible environments.
• Test Integrity: The integration of "scan_layers": False into the post-training E2E unit tests correctly aligns the dummy NNX model's structure with its configuration, preventing unexpected failures.

[github-actions]

🤖 Hi @igorts-git, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This PR introduces a significant architectural simplification and memory optimization by transitioning from multi-chunk host CPU weight staging to a single persistent RaidenSynchronizer using Raiden-FFI under Pathways. The overall code changes are highly focused, clean, and greatly improve performance while reducing client RAM pressure under Pathways.

🔍 General Feedback

• Simplification: Removing multi-chunk weight staging and _split_into_chunks() greatly reduces complexity and code size.
• Fail-Fast Error Handling: Emphasizing immediate error reporting on missing imports/FFI is a great developer experience enhancement.
• Resource Cleanup: We provided an inline comment to ensure release_host_arrays is called within release_weight_sync() for optimal resource/staged-buffer cleanup.

[A9isha]

Verified this branch end to end on a v5p-8 host (4 chips) running tunix's math_gsm8k_dist GRPO example — Qwen3-0.6B, MaxText trainer, vLLM rollout, TP=2, Raiden weight sync. Trainer source checksums and vLLM destination checksums match exactly:

__grand_total__   13217973.204956055
__tensor_count__  310
__element_count__ 596049920

Identical on both sides, and identical to what the pre-collapse multi-chunk path produced, so the single-synchronizer change is a clean swap on the num_chunks == 1 path. I'd arrived at the same host_stage= and release_host_arrays() removals independently against current tunix — this PR covers both, so I'm dropping my version in favour of it.

One note on the payload freeze:

@dataclasses.dataclass(frozen=True, kw_only=True)
class RouterReplayTrainerPayload(abstract_engine.TrainerPayload):

TrainerPayload is declared @flax.struct.dataclass(frozen=True, kw_only=True) in tunix/experimental/common/datatypes.py. Freezing does clear the TypeError: cannot inherit non-frozen dataclass from a frozen one, but flax.struct.dataclass registers each class as a pytree node individually, and a plain dataclasses.dataclass subclass doesn't inherit that registration. On this branch:

TrainerPayload (base):             PyTreeDef(CustomNode(TrainerPayload[()], [None, None, None, None]))
RouterReplayTrainerPayload (here): PyTreeDef(*)

The subclass flattens to a single opaque leaf rather than into its fields. It doesn't affect the run above — Qwen3-0.6B is dense, so no router-replay payload is ever constructed — but on an MoE path that passes one across a jit boundary, the arrays inside would be treated as static rather than traced.

@struct.dataclass(frozen=True, kw_only=True) (with from flax import struct) keeps the registration and is otherwise equivalent here. Happy to send it as a follow-up if you'd rather not touch this PR.

[A9isha]

Approved with one comment