[Trellis]Qwen3.5-35B; Weight Conversion and Raiden Weight Sync Integration - #5089
[Trellis]Qwen3.5-35B; Weight Conversion and Raiden Weight Sync Integration#5089YixuanWang-99 wants to merge 12 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for Qwen 3.5 hybrid cycle layers in the weight unscanning and synchronization pipeline, updates Raiden synchronizer import paths, refactors vLLM integration to use MaxTextVllmSampler, and introduces warnings for replicated batch dimensions. The code review feedback highlights a critical bug where self._raiden_syncs was accidentally removed from maxtext_engine.py's initialization, causing an AttributeError. Additionally, the reviewer pointed out outdated import paths in error messages and test probes, a risky string-stripping operation (rstrip('s')) in raiden_unscan.py, a potential unhandled case for cycle-slot matching, and an unused parameter in _fuse_and_unstack_moe.
| self._signature_compare_warned: bool = False | ||
| self._raiden_syncs: Any = None | ||
| self._replicated_batch_warned: bool = False |
There was a problem hiding this comment.
The initialization of self._raiden_syncs was accidentally removed from __init__ when adding self._replicated_batch_warned. This will cause an immediate AttributeError when prepare_weight_sync, release_weight_sync, or close is called. Please restore self._raiden_syncs: Any = None in __init__.
self._signature_compare_warned: bool = False
self._replicated_batch_warned: bool = False
self._raiden_syncs: Any = None| raise RuntimeError( | ||
| "staging_transport='raiden' requires tunix.experimental.worker." | ||
| "raiden_synchronizer, which the installed tunix does not provide. Install a" |
There was a problem hiding this comment.
The error message still refers to the old module path tunix.experimental.worker.raiden_synchronizer. Since the import path was updated to tunix.experimental.weight_sync.raiden_synchronizer, please update the error message to match the new path to avoid confusion during debugging.
| raise RuntimeError( | |
| "staging_transport='raiden' requires tunix.experimental.worker." | |
| "raiden_synchronizer, which the installed tunix does not provide. Install a" | |
| raise RuntimeError( | |
| "staging_transport='raiden' requires tunix.experimental.weight_sync." | |
| "raiden_synchronizer, which the installed tunix does not provide. Install a" |
| try: | ||
| importlib.import_module("tunix.experimental.worker.raiden_synchronizer") | ||
| _RAIDEN_AVAILABLE = True |
There was a problem hiding this comment.
The test still probes the old module path tunix.experimental.worker.raiden_synchronizer to determine _RAIDEN_AVAILABLE. Since the engine now imports from tunix.experimental.weight_sync.raiden_synchronizer, this probe will evaluate to False even if the synchronizer is available at the new path, causing the test to bypass the real staging path. Please update the probe to use the new module path.
| try: | |
| importlib.import_module("tunix.experimental.worker.raiden_synchronizer") | |
| _RAIDEN_AVAILABLE = True | |
| try: | |
| importlib.import_module("tunix.experimental.weight_sync.raiden_synchronizer") | |
| _RAIDEN_AVAILABLE = True |
| # suffix and fold it into the global layer index below. | ||
| slot = None | ||
| if cycle_interval > 1 and suffix and isinstance(suffix[0], str): | ||
| match = re.fullmatch(rf"{re.escape(layer_container.rstrip('s'))}_(\d+)", suffix[0]) |
There was a problem hiding this comment.
Using rstrip('s') to strip the trailing 's' from layer_container is risky because it removes all trailing 's' characters (e.g., "class" would become "cla", "loss" would become "lo"). A safer and more robust approach is to only strip a single trailing 's' if it exists.
| match = re.fullmatch(rf"{re.escape(layer_container.rstrip('s'))}_(\d+)", suffix[0]) | |
| match = re.fullmatch(rf"{re.escape(layer_container[:-1] if layer_container.endswith('s') else layer_container)}_(\d+)", suffix[0]) |
| expected_axis_len = num_reps if slot is not None else num_layers | ||
| if arr.shape[scan_axis] != expected_axis_len: |
There was a problem hiding this comment.
If cycle_interval > 1 but slot is None (e.g., due to a regex mismatch or a homogeneous parameter), expected_axis_len defaults to num_layers. However, the actual scanned axis length is num_reps. This mismatch will cause a confusing ValueError claiming a shape mismatch (expecting num_layers instead of num_reps), or an out-of-bounds error during slicing. Consider raising a descriptive error directly if slot is None when cycle_interval > 1.
| expected_axis_len = num_reps if slot is not None else num_layers | |
| if arr.shape[scan_axis] != expected_axis_len: | |
| if cycle_interval > 1 and slot is None: | |
| raise ValueError( | |
| f"unscan_layers: {'.'.join(str(k) for k in key)!r} is missing the expected cycle-slot prefix " | |
| f"under {layer_container!r}." | |
| ) | |
| expected_axis_len = num_reps if slot is not None else num_layers | |
| if arr.shape[scan_axis] != expected_axis_len: |
| scan_fused_axis: int, | ||
| tgt_fused_axis: int, |
There was a problem hiding this comment.
The scan_fused_axis argument is no longer used in the new implementation of _fuse_and_unstack_moe. Since weight_converter.py is not part of this PR's diff, we cannot safely remove it from the signature without breaking the caller. However, please consider cleaning this up in a future refactoring of both files to remove the dead code and unused argument.
5fb2ac6 to
e80a00c
Compare
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
YixuanWang-99
changed the title
0e43634 to
2e967ce
Compare
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.
- 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
…rics 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
…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
- 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()
2e967ce to
adcd343
Compare
…ase 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
YixuanWang-99
changed the title
… 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
…ight sync - Map layer names to KV cache indices for hybrid attention/GDN architectures in decoders and adapter - Support cycle_interval in raiden_unscan and optimize memory during unflattening to prevent host OOM - Honor DISABLE_CHECKPOINTING and USE_RAIDEN_FFI in MaxTextTrainingEngine - Preserve gate and router parameter dtype during weight conversion - Add unit tests for map_layer_names_to_indices
… sync; filter metadata in training engine
… engine integration
- Centralize tunix_compat_context and resolve_rollout_tp in convert_utils - Streamline MaxTextVllmSampler and MaxTextVllmRollout by delegating sync lifecycle - Support RAIDEN_USE_FFI env var and host array release in MaxTextTrainingEngine - Enforce strict layer mapping validation in decoders and nnx_decoders - Add convert_utils_test and cross_repo_drift_test unit tests
YixuanWang-99
changed the title
| self._last_staged_step: Optional[int] = None | ||
| self._staged_metadata: Any = None | ||
| self._use_weight_converter = bool( | ||
| getattr(self._config, "use_weight_converter", False) |
There was a problem hiding this comment.
I am always very suspicious when AI Agents add "getattr" or "hasattr", it is usually an indications that they are referencing an old implementation or hallucinating APIs that don't exist. Our config either has a use_weight_converter flag or it does not. Using getattr is error-prone in case of misspelling or flag name change.
AI agents also tend to add those statements because they are lazy to update the configs in the unit tests.
| "compatible tpu_raiden_jax wheel with FFI support is installed." | ||
| ) | ||
|
|
||
| raiden_ffi_env = os.environ.get("RAIDEN_USE_FFI") |
There was a problem hiding this comment.
Why do we ever want to override this? IIUC, for any reasonable large model we must use FFI or we will OOM.
| job_name="trainer", | ||
| worker_index=jax.process_index(), | ||
| auto_h2d=False, | ||
| host_stage=host_stage, |
There was a problem hiding this comment.
IIRC, "host_stage" was removed from the API in the last week. Do we have to add it back or is it a hallunication?
|
Thanks everyone for the review feedback! Per the feedback on PR size and scope, this PR is being split into a stacked chain of focused, reviewable PRs targeting Stacked PR Chain:
Notes on review comments & dropped duplicates:
|
3 participants
Description
Integrates Raiden weight synchronization into
MaxTextTrainingEnginefor distributed RL post-training (Trellis / GRPO). Introduces trainer-side target-free weight conversion, MoE kernel layout alignment, hybrid KV cache mapping, and aggressive host memory reclamation during device-to-host staging.Companion Tunix PR: google/tunix#2083
Key Changes
weight_converter.py,convert_utils.py): Synthesizes rollout target shapes and keys directly from trainer parameter state—unrolling scanned layers and prefusing MoE gate/up weights (wi_0+wi_1wi) without requiring rollout enginetarget_state.convert_utils.py,moe_padding.py): Implements 128-element lane interleaving (TPU_V5P_SUBCORE_LANE_SIZE = 128) for fused MoE weights to satisfy TPU GMM v2 kernel requirements, and dynamically pads intermediate dimensions.hybrid_cache_utils.py,decoders.py,nnx_decoders.py): Maps physical layer names to KV cache slot indices for hybrid architectures (e.g., Qwen 3.5 GDN recurrent + attention layers) and supports inhomogeneous layer unrolling.maxtext_engine.py): Addsreclaim_host_memory()(gc.collect()+malloc_trim(0)) andrelease_host_arrays()to eliminate host OOMs during D2H transfer; adds Pathways FFI support (RAIDEN_USE_FFI).cross_repo_drift_test.pyto guard against parameter/contract drift with Tunix, and replaces silent failure fallbacks with fail-fast exceptions.Testing
cross_repo_drift_test.py,convert_utils_test.py,weight_converter_test.py,raiden_unscan_test.py,vllm_hybrid_cache_test.py, andprepare_weight_sync_test.py.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.