From 784b25870c11ee3b8c6fefcc25ade5d58f2f25cf Mon Sep 17 00:00:00 2001 From: hengtaoguo Date: Tue, 20 Oct 2026 23:37:18 +0000 Subject: [PATCH] Add general scan support for M3 --- src/maxtext/m3/infra/__init__.py | 4 +- src/maxtext/m3/infra/scan.py | 169 ++++++++++++++++++ src/maxtext/m3/models/qwen3/modeling_qwen3.py | 6 +- .../m3/models/qwen3_vl/modeling_qwen3_vl.py | 2 + src/maxtext/m3/train/loop.py | 11 +- src/maxtext/utils/model_creation_utils.py | 5 +- tests/unit/m3/scan_test.py | 155 ++++++++++++++++ tests/unit/m3/train_loop_test.py | 14 +- 8 files changed, 356 insertions(+), 10 deletions(-) create mode 100644 src/maxtext/m3/infra/scan.py create mode 100644 tests/unit/m3/scan_test.py diff --git a/src/maxtext/m3/infra/__init__.py b/src/maxtext/m3/infra/__init__.py index 817af52b1a..33b381deee 100644 --- a/src/maxtext/m3/infra/__init__.py +++ b/src/maxtext/m3/infra/__init__.py @@ -14,9 +14,11 @@ """Post-construction transforms applied from training setup, never from models. +Implemented modules: + * `scan.py` — layer stacking via `jax.lax.scan` + `maybe_apply_scan`. + Planned modules: * `remat.py` — path-transparent remat wrapper + `maybe_apply_remat`. - * `scan.py` — layer stacking via `jax.lax.scan` + `maybe_apply_scan`. * `quantization.py` — qwix int8/fp8 rules + `maybe_quantize`. * `offload.py` — `to_host` for optimizer/parameter offloading. diff --git a/src/maxtext/m3/infra/scan.py b/src/maxtext/m3/infra/scan.py new file mode 100644 index 0000000000..58909b87c3 --- /dev/null +++ b/src/maxtext/m3/infra/scan.py @@ -0,0 +1,169 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Post-construction scanning of repeated decoder layers or blocks. + +Scan units must share an NNX graph definition and matching corresponding +variable shapes; their weight values are independent. A unit may be a single +layer or a model-specific block containing different layer types. + +Models opt in through ``decoder.layer_names`` and define the scan units. +Keep this module free of model-specific branches: architecture, cyclic patterns, +and grouping belong to model definitions, not to the scan transform. +""" + +from collections.abc import Sequence +from typing import Any + +from flax import nnx +import jax +import jax.numpy as jnp +from jax.sharding import NamedSharding +from jax.sharding import PartitionSpec as P + +_SHARDING_KEYS = ("out_sharding", "sharding", "sharding_names") + + +def _is_variable(x: Any) -> bool: + return isinstance(x, nnx.Variable) + + +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)): + entries = list(spec) + entries.extend([None] * max(0, axis - len(entries))) + entries.insert(axis, None) + return P(*entries) if isinstance(spec, P) else tuple(entries) + return spec + + +def _remove_axis_from_spec(spec: Any, axis: int) -> Any: + """Removes the scan-axis entry from a sharding spec when present.""" + if isinstance(spec, (P, tuple)) and len(spec) > axis: + entries = list(spec) + del entries[axis] + return P(*entries) if isinstance(spec, P) else tuple(entries) + return spec + + +def _map_sharding_metadata(state: nnx.State, fn, axis: int) -> nnx.State: + """Applies `fn(spec, axis)` to supported sharding metadata in `state`.""" + + def update(var): + if not _is_variable(var): + return var + meta = var.get_metadata() + effective_axis = axis if var.get_value().ndim >= axis else 0 + changes = {k: fn(meta[k], effective_axis) for k in _SHARDING_KEYS if meta.get(k) is not None} + return var.replace(**changes) if changes else var + + return jax.tree.map(update, state, is_leaf=_is_variable) + + +def _stack_states(states: Sequence[nnx.State], param_axis: int, mesh) -> nnx.State: + """Stacks Params on `param_axis` when rank permits; other variables use axis 0.""" + + def stack(*variables): + first = variables[0] + if not _is_variable(first): + return first + axis = param_axis if isinstance(first, nnx.Param) and first.get_value().ndim >= param_axis else 0 + value = jnp.stack([v.get_value() for v in variables], axis=axis) + meta = first.get_metadata() + changes = {k: _insert_axis_in_spec(meta[k], axis) for k in _SHARDING_KEYS if meta.get(k) is not None} + spec = changes.get("out_sharding", changes.get("sharding")) + if mesh is not None and isinstance(spec, P): + value = jax.device_put(value, NamedSharding(mesh, spec)) + return first.replace(value=value, **changes) + + return jax.tree.map(stack, *states, is_leaf=_is_variable) + + +class ScannedLayers(nnx.Module): + """Scan compatible layers or blocks with independent weights. + + Units must have equal NNX graph definitions and matching corresponding variable + shapes. Calls take an array carry and return an array of the same shape and + dtype; additional arguments are shared across all iterations. Parameters must + remain unchanged during calls; updated non-Param state is written back. + """ + + def __init__(self, layers: Sequence[nnx.Module], *, param_scan_axis: int = 1, mesh=None): + if len(layers) < 2: + raise ValueError(f"ScannedLayers needs at least two layers, got {len(layers)}.") + graphdef = nnx.graphdef(layers[0]) + for i, layer in enumerate(layers[1:], start=1): + if nnx.graphdef(layer) != graphdef: + raise ValueError(f"Layer {i} is not structurally identical to layer 0; cannot scan them together.") + + stacked = nnx.merge( + graphdef, + _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(): + if not name.startswith("_"): + setattr(self, name, value) + self._layer_graphdef = graphdef + self._param_scan_axis = param_scan_axis + self.num_layers = len(layers) + + def __call__(self, y: jax.Array, *args, **kwargs) -> jax.Array: + axis = self._param_scan_axis + _, params, rest = nnx.split(self, nnx.Param, ...) + params = jax.tree.map(lambda x: jnp.moveaxis(x, axis, 0) if x.ndim > axis else x, params) + + def body(carry, xs): + layer_params, layer_rest = xs + layer_params = _map_sharding_metadata(layer_params, _remove_axis_from_spec, axis) + layer_rest = _map_sharding_metadata(layer_rest, _remove_axis_from_spec, 0) + 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, ...) + new_rest = _map_sharding_metadata(new_rest, _insert_axis_in_spec, 0) + return out, new_rest + + y, new_rest = jax.lax.scan(body, y, (params, rest)) + if jax.tree.leaves(new_rest): + nnx.update(self, new_rest) + return y + + +def maybe_apply_scan(model: nnx.Module, config, mesh=None) -> nnx.Module: + """Scan an opted-in decoder stack in place when enabled; no remat is applied. + + Decoders opt in by executing compatible layers or blocks through ordered + ``layer_names``. Missing or None means unsupported; fewer than two is a no-op. + The transform replaces these units with one module named ``layers``. + """ + if not config.scan_layers: + return model + decoder = model.decoder + layer_names = getattr(decoder, "layer_names", None) + if layer_names is None: + raise ValueError( + 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: + return model + layers = [getattr(decoder, name) for name in layer_names] + scanned = ScannedLayers(layers, param_scan_axis=config.param_scan_axis, mesh=mesh) + for name in layer_names: + delattr(decoder, name) + decoder.layers = scanned + decoder.layer_names = ("layers",) + return model diff --git a/src/maxtext/m3/models/qwen3/modeling_qwen3.py b/src/maxtext/m3/models/qwen3/modeling_qwen3.py index d20d73ff62..9578fb64db 100644 --- a/src/maxtext/m3/models/qwen3/modeling_qwen3.py +++ b/src/maxtext/m3/models/qwen3/modeling_qwen3.py @@ -343,6 +343,8 @@ def __init__( rngs=rngs, ) setattr(self, f"layers_{lyr}", layer) + # The scan transform replaces these entries with one scanned module. + self.layer_names = tuple(f"layers_{lyr}" for lyr in range(self.num_layers)) self.decoder_norm = nnx.RMSNorm( num_features=config.emb_dim, @@ -399,8 +401,8 @@ def __call__( """Iterates through all decoder layers and projects to logits via tied embeddings.""" y = shared_embedding(decoder_input_tokens) - for lyr in range(self.num_layers): - layer = getattr(self, f"layers_{lyr}") + for name in self.layer_names: + layer = getattr(self, name) y = layer( y, decoder_segment_ids=decoder_segment_ids, diff --git a/src/maxtext/m3/models/qwen3_vl/modeling_qwen3_vl.py b/src/maxtext/m3/models/qwen3_vl/modeling_qwen3_vl.py index 9ef6b9dc72..91730f1ac4 100644 --- a/src/maxtext/m3/models/qwen3_vl/modeling_qwen3_vl.py +++ b/src/maxtext/m3/models/qwen3_vl/modeling_qwen3_vl.py @@ -897,6 +897,8 @@ def __init__( rngs: nnx.Rngs, ): super().__init__(config, mesh, rngs=rngs, layer_cls=Qwen3VLDecoderLayer) + # Scan does not yet support inter-layer DeepStack injections. + self.layer_names = None def __call__( self, diff --git a/src/maxtext/m3/train/loop.py b/src/maxtext/m3/train/loop.py index d49043ab8c..54443e34ea 100644 --- a/src/maxtext/m3/train/loop.py +++ b/src/maxtext/m3/train/loop.py @@ -23,6 +23,7 @@ model constructors take today. The fields read here are `steps`, `per_device_batch_size`, `max_target_length`, `learning_rate`, the optimizer fields consumed by `maxtext.optimizers.optimizers.get_optimizer`, +`scan_layers` and `param_scan_axis` consumed by the scan transform, `gradient_clipping_threshold`, `init_weights_seed`, and the profiler fields consumed by `maxtext.common.profiler.Profiler` (`profiler`, `skip_first_n_steps_for_profiler`, `profiler_steps`, ...). The trace is @@ -56,6 +57,7 @@ from maxtext.configs import pyconfig from maxtext.m3 import models as m3_models from maxtext.m3.core.sharding import create_mesh +from maxtext.m3.infra import scan as m3_scan from maxtext.optimizers import optimizers from maxtext.utils import max_logging @@ -148,6 +150,8 @@ def train_loop(config) -> float: with jax.set_mesh(mesh): model = m3_models.create_model(config, mesh, rngs=nnx.Rngs(config.init_weights_seed)) + # TODO: Apply remat_policy to the scan body. + model = m3_scan.maybe_apply_scan(model, config, mesh=mesh) optimizer = nnx.Optimizer(model, make_optimizer(config), wrt=nnx.Param) num_params = sum(x.size for x in jax.tree.leaves(nnx.state(model, nnx.Param))) @@ -160,6 +164,7 @@ def train_loop(config) -> float: prof = profiler.Profiler(config, offset_step=0) data_key = jax.random.key(config.init_weights_seed) loss = float("nan") + step_fn = nnx.cached_partial(train_step, model, optimizer) # Steps are logged one iteration late. Reading a loss blocks until its step # finishes, so doing it right after dispatch would leave the device idle # while the host prepares the next step. Reading step N-1's loss after @@ -169,13 +174,13 @@ def train_loop(config) -> float: last_done = time.perf_counter() for step in range(config.steps): batch = make_synthetic_batch(config, mesh, jax.random.fold_in(data_key, step)) - prof.maybe_activate_profiler(step, nnx.state(model)) + prof.maybe_activate_profiler(step, pending) with jax.profiler.StepTraceAnnotation("train", step_num=step): - step_loss = train_step(model, optimizer, batch) + step_loss = step_fn(batch) if pending is not None: loss, last_done = _log_step(*pending, last_done, tokens_per_step) pending = (step, step_loss) - prof.maybe_deactivate_profiler(step, nnx.state(model)) + prof.maybe_deactivate_profiler(step, pending) if pending is not None: loss, last_done = _log_step(*pending, last_done, tokens_per_step) diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index 62be9ce62f..e3652ec11b 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -552,8 +552,11 @@ def create_model( """Instantiates and returns the model object, sharded across the mesh.""" if config.use_m3_model: from maxtext.m3 import models as m3_models # pylint: disable=import-outside-toplevel + from maxtext.m3.infra import scan as m3_scan # pylint: disable=import-outside-toplevel - 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) + return model # Model definition quant = quantizations.configure_quantization(config, quant_mode_str=quant_mode_str) diff --git a/tests/unit/m3/scan_test.py b/tests/unit/m3/scan_test.py new file mode 100644 index 0000000000..234d6a6594 --- /dev/null +++ b/tests/unit/m3/scan_test.py @@ -0,0 +1,155 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""m3 scan transform checks, including Qwen3 output and gradient parity.""" + +from absl.testing import absltest +from flax import nnx +from jax.sharding import PartitionSpec as P +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.configs import pyconfig +from maxtext.m3.infra import scan as m3_scan +from maxtext.m3.models.qwen3.modeling_qwen3 import create_qwen3_model +from maxtext.m3.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLDecoder +from tests.utils.test_helpers import get_test_config_path + + +def _config(scan_layers: bool, num_layers: int = 4, param_scan_axis: int = 1): + """Build a small Qwen3 test configuration with optional layer scanning.""" + args = { + "model_name": "qwen3-0.6b", + "use_m3_model": "true", + "scan_layers": str(scan_layers).lower(), + "param_scan_axis": param_scan_axis, + "override_model_config": "true", + "base_num_decoder_layers": num_layers, + "base_emb_dim": 64, + "base_mlp_dim": 128, + "base_num_query_heads": 4, + "base_num_kv_heads": 2, + "head_dim": 16, + "vocab_size": 256, + "per_device_batch_size": 1, + "max_target_length": 32, + "dtype": "float32", + "weight_dtype": "float32", + "matmul_precision": "highest", + "attention": "dot_product", + "enable_checkpointing": "false", + "run_name": "m3_scan_test", + } + return pyconfig.initialize(["", get_test_config_path()] + [f"{k}={v}" for k, v in args.items()]) + + +class ScanTransformTest(absltest.TestCase): + + def test_qwen3_scanned_matches_unscanned(self): + for axis in (0, 1): + with self.subTest(param_scan_axis=axis): + cfg = _config(scan_layers=True, param_scan_axis=axis) + model = create_qwen3_model(cfg, rngs=nnx.Rngs(0)) + tokens = jax.random.randint(jax.random.key(1), (1, 32), 0, cfg.vocab_size) + positions = jnp.arange(32)[None] + ref = model(tokens, positions) + + m3_scan.maybe_apply_scan(model, cfg) + self.assertEqual(model.decoder.layer_names, ("layers",)) + + out = nnx.jit(lambda m, t, p: m(t, p))(model, tokens, positions) + np.testing.assert_allclose(np.asarray(out), np.asarray(ref), atol=1e-4, rtol=1e-4) + + def test_qwen3_scan_is_noop_when_disabled(self): + cfg = _config(scan_layers=False) + model = create_qwen3_model(cfg, rngs=nnx.Rngs(0)) + m3_scan.maybe_apply_scan(model, cfg) + self.assertEqual(model.decoder.layer_names, tuple(f"layers_{i}" for i in range(4))) + + def test_qwen3_scanned_gradients_match(self): + for axis in (0, 1): + with self.subTest(param_scan_axis=axis): + cfg = _config(scan_layers=True, num_layers=3, param_scan_axis=axis) + tokens = jax.random.randint(jax.random.key(1), (1, 32), 0, cfg.vocab_size) + positions = jnp.arange(32)[None] + + def loss(m, t, p): + return jnp.mean(m(t, p) ** 2) + + ref_model = create_qwen3_model(cfg, rngs=nnx.Rngs(0)) + ref_grads = nnx.grad(loss)(ref_model, tokens, positions) + model = m3_scan.maybe_apply_scan(create_qwen3_model(cfg, rngs=nnx.Rngs(0)), cfg) + grads = nnx.jit(nnx.grad(loss))(model, tokens, positions) + + self._assert_matches_unscanned(grads, ref_grads, cfg) + + def _assert_matches_unscanned(self, scanned, reference, cfg): + reference = nnx.to_pure_dict(reference) + decoder = reference["decoder"] + layers = [decoder.pop(f"layers_{i}") for i in range(cfg.num_decoder_layers)] + decoder["layers"] = jax.tree.map(lambda *xs: jnp.stack(xs, axis=cfg.param_scan_axis), *layers) + actual = nnx.to_pure_dict(scanned) + self.assertEqual(jax.tree.structure(actual), jax.tree.structure(reference)) + for (path, value), expected in zip(jax.tree.leaves_with_path(actual), jax.tree.leaves(reference)): + np.testing.assert_allclose( + np.asarray(value), + np.asarray(expected), + atol=1e-5, + rtol=1e-4, + err_msg=str(path), + ) + + def test_scalar_and_vector_params_with_short_sharding_specs(self): + class Layer(nnx.Module): + """Affine layer with scalar and vector parameters.""" + + def __init__(self, scale): + self.scale = nnx.Param(jnp.array(scale), out_sharding=P()) + self.bias = nnx.Param(jnp.arange(4, dtype=jnp.float32), out_sharding=P()) + + def __call__(self, x): + # Metadata must describe the sliced values, including the scalar. + if len(self.scale.get_metadata()["out_sharding"]) != 0: + raise ValueError("Scalar sharding must have rank zero.") + return x * self.scale[...] + self.bias[...] + + for axis in (0, 1): + with self.subTest(param_scan_axis=axis): + layers = [Layer(float(i + 1)) for i in range(3)] + x = jnp.ones(4) + expected = x + for layer in layers: + expected = layer(expected) + scanned = m3_scan.ScannedLayers(layers, param_scan_axis=axis) + self.assertEqual(scanned.scale.shape, (3,)) + self.assertEqual(scanned.bias.shape, (3, 4) if axis == 0 else (4, 3)) + self.assertEqual(scanned.bias.get_metadata()["out_sharding"], P(None) if axis == 0 else P(None, None)) + actual = nnx.jit(lambda m, v: m(v))(scanned, x) + np.testing.assert_allclose(np.asarray(actual), np.asarray(expected)) + + def test_scan_rejects_unsupported_model_before_mutation(self): + cfg = _config(scan_layers=True) + model = create_qwen3_model(cfg, rngs=nnx.Rngs(0)) + model.decoder = Qwen3VLDecoder(cfg, rngs=nnx.Rngs(0)) + self.assertIsNone(model.decoder.layer_names) + first_layer = model.decoder.layers_0 + with self.assertRaisesRegex(ValueError, "Qwen3VLDecoder does not expose scan units through `layer_names`"): + m3_scan.maybe_apply_scan(model, cfg) + self.assertIs(model.decoder.layers_0, first_layer) + self.assertIsNone(model.decoder.layer_names) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/m3/train_loop_test.py b/tests/unit/m3/train_loop_test.py index 223d59fb65..1d2071d55c 100644 --- a/tests/unit/m3/train_loop_test.py +++ b/tests/unit/m3/train_loop_test.py @@ -77,9 +77,17 @@ def test_synthetic_batch_targets_are_shifted_inputs(self): self.assertEqual(inputs.shape, (config.global_batch_size_to_train_on, config.max_target_length)) np.testing.assert_array_equal(inputs[:, 1:], targets[:, :-1]) - def test_train_loop_returns_finite_loss(self): - loss = loop.train_loop(_tiny_config(self.output_dir)) - self.assertTrue(math.isfinite(loss), f"loss={loss}") + def test_train_loop_losses_match_with_and_without_scan(self): + losses = [] + for scan_layers in (False, True): + with self.subTest(scan_layers=scan_layers): + config = _tiny_config( + self.output_dir, scan_layers=str(scan_layers).lower(), dtype="float32", weight_dtype="float32" + ) + loss = loop.train_loop(config) + self.assertTrue(math.isfinite(loss), f"loss={loss}") + losses.append(loss) + np.testing.assert_allclose(losses[0], losses[1], rtol=1e-4, atol=1e-5) def test_train_loop_writes_xplane_profile(self): config = _tiny_config(self.output_dir, profiler="xplane", skip_first_n_steps_for_profiler=1, profiler_steps=1)