[WIP] [Trellis] Weight Conversion that is compatible with Raiden - #5089
[WIP] [Trellis] Weight Conversion that is compatible with Raiden#5089YixuanWang-99 wants to merge 4 commits into
Conversation
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.
raiden_synchronizer moved from tunix.experimental.worker to tunix.experimental.weight_sync, so the staging_transport == "raiden" branch raised the ImportError it treats as fatal. Also treat an empty accumulated_metrics like a missing one: the is-not-None check passed for an empty dict and handed Orbax an empty PyTree to save, which fails rather than no-opping.
![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! |
- 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
YixuanWang-99
changed the title
…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
2 participants
Description
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456
You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.