Repository navigation
Add general scan support for M3 - #5875
hengtaoguo wants to merge 1 commit 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 introduces post-construction scanning of repeated decoder layers or blocks for Qwen3 models, implementing the ScannedLayers module and integration utilities in scan.py along with corresponding unit tests. The review feedback highlights several compatibility issues with standard Flax NNX APIs in the new scan.py file, specifically pointing out that nnx.Variable lacks get_metadata() and get_value() methods, and that nnx.graphdef is not a valid function.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
🤖 Hi @hengtaoguo, 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[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
📋 Review Summary
Verdict: Core scan transformation implemented and verified for dense weights, but contains scalar parameter stacking/transposition bugs and an unedited boilerplate description.
Top issues:
⚠️ [Description] The PR description contains only unedited template boilerplate with unchecked boxes and no problem statement, architecture context, or reproduction instructions. → Replace the boilerplate with the problem, design of the M3 post-construction scan transform, limitations (no remat yet, DeepStack incompatible), and reproduction steps.- 🔴 [Bug]
_stack_statesandScannedLayers.__call__assume all parameters have rank >param_scan_axis, which will crash on 0D (scalar) parameters withAxisErrorwhenparam_scan_axis=1. → Guard parameter stacking and axis transpositions withndim > param_scan_axis, defaulting to axis 0 (see inline comments onsrc/maxtext/m3/infra/scan.py). - [Bug]
_remove_axis_from_speconly deletes iflen(spec) > axis, which fails to invert_insert_axis_in_specwhen the original spec had length <=axis, leaving mismatched PartitionSpec ranks. → Delete indexmin(axis, len(spec) - 1)so the inserted axis is consistently removed (see inline comment onsrc/maxtext/m3/infra/scan.py). - [Perf] Mapping sharding metadata updates 3 times inside the inner scan
bodyadds tracing overhead while serving no functional purpose, as JAX/XLA ignores variable metadata during forward execution. → Remove_map_sharding_metadatacalls from the inner scanbody(see inline comment onsrc/maxtext/m3/infra/scan.py). - [Tests] Core configuration variants like
param_scan_axis=0and sharded execution with an activeMeshare untested intests/unit/m3/scan_test.py. → Parameterize_configand add tests validatingparam_scan_axis=0and sharding specs under a mesh (see inline comment ontests/unit/m3/scan_test.py).
Inferred spec & validation
- Sources: PR title; code before and after;
m3/README.mdand legacylayers/nnx_scan.py. - Problem: M3 lacked layer scanning support for decoder stacks, forcing full unrolling during compilation and execution.
| Spec item | Implementation | Validating test |
|---|---|---|
R1 (inferred): Post-construction scan transform (ScannedLayers, maybe_apply_scan) stacks layers and runs them via jax.lax.scan. |
✅ | tests/unit/m3/scan_test.py::test_qwen3_scanned_matches_unscanned ✅ |
| R2 (inferred): Scanned model preserves forward logits and backward gradients compared to unscanned model. | ✅ | tests/unit/m3/scan_test.py::test_qwen3_scanned_matches_unscanned, test_qwen3_scanned_gradients_match ✅ |
R3 (inferred): Parameter stacking and scanning supports arbitrary param_scan_axis (including 0 and 1) and handles arbitrary parameter ranks. |
param_scan_axis=1) |
tests/unit/m3/scan_test.py ❌ none (only default param_scan_axis=1 tested on dense 2D/1D weights) |
R4 (inferred): Guard unsupported models (e.g. Qwen3VLDecoder with DeepStack injections) by rejecting scanning before mutating model. |
✅ | tests/unit/m3/scan_test.py::test_scan_rejects_unsupported_model_before_mutation ✅ |
R5 (inferred): Integrate scan into M3 training loop (train/loop.py) under scan_layers=true. |
✅ | tests/unit/m3/train_loop_test.py::test_train_loop_losses_match_with_and_without_scan ✅ |
C1 (inferred): scan_layers=false remains an exact no-op with unchanged model structure. |
✅ preserved | tests/unit/m3/scan_test.py::test_qwen3_scan_is_noop_when_disabled ✅ |
| C2 (inferred): Maintain transparent parameter paths for checkpointing, optimizer, and weight inspectability. | ✅ preserved | tests/unit/m3/scan_test.py::test_qwen3_scanned_gradients_match ✅ |
d160cdd to
2019109
Compare
0e92ea8 to
c8f7b0e
Compare
87cc16f to
53a1aa7
Compare
726e0f2 to
784b258
Compare
There was a problem hiding this comment.
LGTM. Just two callouts:
- (Also discussed offline) Is there a way we can enforce model files follow the patterns for scan to get picked up properly?
- Does merged code from @fenghuizhang's GLM-5.3-Flash or @RissyRan's Qwen3.8-Flash-Next need updates with this?
RissyRan
left a comment
RissyRan
left a comment
There was a problem hiding this comment.
Thanks for adding general scan support for m3! I left comments on two blocking issues (GLM-5.3-Flash failing under the default scan_layers=true, and loading m3 checkpoints saved before this PR) and six correctness issues in scan.py. These come from reading the code at 784b2587; I haven't run these cases, so please let me know if I've misread anything.
| """ | ||
| if not config.scan_layers: | ||
| return model | ||
| decoder = model.decoder |
There was a problem hiding this comment.
Blocking: model.decoder is read unconditionally, so any m3 model without a decoder attribute now fails whenever scan_layers=true, which is the base.yml default. Glm5Model (registered as glm5.3-flash, with use_m3_model: true in its yml) has embed_tokens / prefix_layers / super_block / suffix_layers and no decoder. As a result, create_model, from_pretrained and the m3 loop all raise AttributeError: 'Glm5Model' object has no attribute 'decoder'. Setting scan_layers=false avoids it, but it also disables GLM's own scan (glm5_scan reads config.scan_layers).
The same default also turns the layer_names is None case (e.g. Qwen3-VL) into a hard ValueError for runs that previously ignored scan_layers.
Suggestion: treat a missing decoder / layer_names, or a model that scans internally, as a no-op rather than an error (e.g. decoder = getattr(model, "decoder", None) and return early). Models would then opt in, instead of every m3 model going through this path.
|
|
||
| return m3_models.create_model(config, mesh, rngs=rngs, quant_mode_str=quant_mode_str) | ||
| model = m3_models.create_model(config, mesh, rngs=rngs, quant_mode_str=quant_mode_str) | ||
| model = m3_scan.maybe_apply_scan(model, config, mesh=mesh) |
There was a problem hiding this comment.
Blocking (checkpoint compatibility): with this change, m3 models start honoring scan_layers (default true). However, m3 checkpoints saved before this PR wrote custom_metadata["scan_layers"] = config.scan_layers (checkpointing.py:1613), i.e. True, while their trees hold unscanned decoder/layers_N.
In from_pretrained, verify_and_sync_scan_layers then either:
- keeps
True(default), so the model is built with a stackeddecoder/layersand restore fails on the tree mismatch; or - raises "Configuration mismatch" when the user passes
scan_layers=falseexplicitly.
So those checkpoints can't be loaded with any setting. Could we add a migration path, e.g. an m3-specific metadata key, or accepting unscanned trees for checkpoints that predate this change?
|
|
||
| def stack(*variables): | ||
| first = variables[0] | ||
| if not _is_variable(first): |
There was a problem hiding this comment.
stack returns layer 0's value for any non-Variable leaf. NNX graph flattening keeps raw jax.Array / np.ndarray attributes as bare state leaves, so a unit that stores a constant buffer (e.g. self.inv_freq = jnp.arange(64)) ends up with a single unstacked copy. That leaf then lands in rest, and lax.scan slices its leading axis as if it were the layer axis: a length mismatch if dim 0 != L, or silently one row per layer if it happens to equal L. Could we stack these leaves as well, or reject them explicitly?
| layer = nnx.merge(self._layer_graphdef, layer_params, layer_rest) | ||
| out = layer(carry, *args, **kwargs) | ||
| # Params come back unchanged; only emit the (possibly mutated) non-Param state. | ||
| _, _, new_rest = nnx.split(layer, nnx.Param, ...) |
There was a problem hiding this comment.
Only non-Param state is emitted from the body, so any nnx.Param created during the call is silently dropped. With the documented order construct → scan → quantize, Qwix LoRA / quantized layers create their params lazily inside this call (LoRA adapters are Param subclasses), so they would be discarded on every step and never trained.
Conversely, a new non-Param variable created in the call (sow, quantization stats) is stacked and added by nnx.update. That changes the model graphdef after the first step, which makes nnx.cached_partial raise "graph structure ... was mutated".
maxtext.layers.nnx_scan.apply_scanned_layers already handles this case. Could we reuse it, or mirror its handling here?
| _stack_states([nnx.state(l) for l in layers], param_scan_axis, mesh), | ||
| ) | ||
| # Adopt the stacked layer's attributes so parameter paths match a single layer's. | ||
| for name, value in vars(stacked).items(): |
There was a problem hiding this comment.
Copying attributes with plain setattr skips names starting with _ and loses explicit nnx.data / nnx.static status:
- A unit with a private submodule or Variable (e.g.
self._norm = nnx.RMSNorm(...)) loses it fromScannedLayers' state, andnnx.merge(self._layer_graphdef, ...)in the body then fails on the missing leaves. - A unit holding
nnx.data([...])can hit flax's "Found data on value ... assigned to static attribute" error here.
Could we copy every attribute, including _-prefixed ones, and preserve its data/static status?
|
|
||
| def _insert_axis_in_spec(spec: Any, axis: int) -> Any: | ||
| """Inserts a replicated (None) entry for the new layer axis into a sharding spec.""" | ||
| if isinstance(spec, (P, tuple)): |
There was a problem hiding this comment.
_insert_axis_in_spec and _remove_axis_from_spec only handle P and tuple. If a Variable's out_sharding is a NamedSharding or a list, it passes through unchanged. The stacked rank-(d+1) array then keeps a rank-d spec, and _stack_states also skips the device_put (it checks isinstance(spec, P)). Downstream sharded creation (get_nnx_var_named_sharding_with_scan_axis / create_nnx_sharded_model) would then fail on the rank mismatch. nnx.spmd.add_axis / remove_axis may already cover these cases.
|
|
||
| stacked = nnx.merge( | ||
| graphdef, | ||
| _stack_states([nnx.state(l) for l in layers], param_scan_axis, mesh), |
There was a problem hiding this comment.
Stacking nnx.state(l) per layer turns a Variable or submodule that is shared across layers, or shared with a module outside the stack, into L independent copies. Weight tying is then silently broken: gradients no longer accumulate into the shared weight, the outside copy stops seeing in-layer updates, and the parameter count grows. Could ScannedLayers.__init__ detect shared references across layers (e.g. by object identity) and raise?
| f"{type(decoder).__name__} does not expose scan units through " | ||
| "`layer_names`; m3 layer scanning is currently unsupported for this model." | ||
| ) | ||
| if len(layer_names) < 2: |
There was a problem hiding this comment.
With scan_layers=true, a decoder with fewer than two layers is left unscanned (decoder/layers_0), while the checkpoint metadata still records scan_layers=True. The legacy path produces a stacked layers tree with L=1 in that case, so a 1-layer scanned checkpoint can't move between the two paths. lax.scan handles length 1 fine, so the < 2 guard (here and in ScannedLayers.__init__) could be dropped.
Description
Add post-construction scanning for m3 Qwen3 through
scan_layers=trueinm3.train.loop. The transform stacks independent layer weights and executes them withjax.lax.scan, preserving existing unscanned parameter paths.New models opt into m3 layer scanning by exposing
decoder.layer_namesas an ordered sequence of structurally identical layers or model-specific blocks, keeping scan.py free of model-name branches.Fix: Adding a small patch to reduce Python overhead in the training loop (original #5900, before vs after: Cache the model/optimizer binding and avoid extracting model state for profiling.
Tests
New
scan_tests.pyguards:Profiling comparison:
Scanned ckpt logits check
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.