diff --git a/src/maxtext/m3/train/loop.py b/src/maxtext/m3/train/loop.py new file mode 100644 index 0000000000..4da66a788e --- /dev/null +++ b/src/maxtext/m3/train/loop.py @@ -0,0 +1,178 @@ +# Copyright 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. + +"""Minimal m3 training loop, intended for collecting profiles. + +Deliberately bare: synthetic data, AdamW at a constant learning rate, plain +next-token cross entropy, no checkpointing, no eval, no metrics writer. It +exists so an m3 model can be stepped end to end and traced; it is not a +replacement for `maxtext.trainers.pre_train.train`. + +Configuration reuses the legacy `pyconfig` surface, because that is what m3 +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`, +`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 +written under `config.tensorboard_dir`. `gradient_accumulation_steps > 1` and +batch sizes that do not split over the mesh are rejected rather than emulated. + +Example (CPU or TPU):: + + python3 -m maxtext.m3.train.loop src/maxtext/configs/base.yml \\ + model_name=qwen3-0.6b use_m3_model=true scan_layers=false \\ + per_device_batch_size=1 max_target_length=256 steps=8 \\ + profiler=xplane skip_first_n_steps_for_profiler=3 profiler_steps=3 \\ + base_output_directory=/tmp/m3 run_name=qwen3_profile \\ + enable_checkpointing=false +""" + +import math +import time +from collections.abc import Sequence + +from absl import app +from flax import nnx +import jax +import jax.numpy as jnp +from jax.sharding import NamedSharding +from jax.sharding import PartitionSpec as P +import optax + +from maxtext.common import profiler +from maxtext.configs import pyconfig +from maxtext.m3 import models as m3_models +from maxtext.m3.core.sharding import create_mesh +from maxtext.optimizers import optimizers +from maxtext.utils import max_logging + +# Mesh axes the batch dimension is split over. Matches the defaults of +# `create_mesh` (fsdp within a slice, data across slices). +_BATCH_AXES = ("data", "fsdp") + + +def make_optimizer(config) -> optax.GradientTransformation: + """`optimizers.get_optimizer` plus global-norm clipping, which it leaves to the caller.""" + tx = optimizers.get_optimizer(config, config.learning_rate) + if config.gradient_clipping_threshold > 0: + return optax.chain(optax.clip_by_global_norm(config.gradient_clipping_threshold), tx) + return tx + + +def make_synthetic_batch(config, mesh, key: jax.Array) -> dict[str, jax.Array]: + """Random token batch, sharded along the batch dimension. + + Targets are the inputs shifted left by one, so the loss is ordinary + next-token prediction rather than a constant. + """ + batch_size = config.global_batch_size_to_train_on + seq_len = config.max_target_length + tokens = jax.random.randint(key, (batch_size, seq_len + 1), 0, config.vocab_size, dtype=jnp.int32) + positions = jnp.broadcast_to(jnp.arange(seq_len, dtype=jnp.int32), (batch_size, seq_len)) + batch = {"inputs": tokens[:, :-1], "targets": tokens[:, 1:], "positions": positions} + return jax.device_put(batch, NamedSharding(mesh, P(_BATCH_AXES))) + + +def loss_fn(model: nnx.Module, batch: dict[str, jax.Array]) -> jax.Array: + """Mean next-token cross entropy, computed in float32.""" + logits = model(batch["inputs"], batch["positions"]) + losses = optax.softmax_cross_entropy_with_integer_labels(logits.astype(jnp.float32), batch["targets"]) + return jnp.mean(losses) + + +@nnx.jit(donate_argnums=(0, 1)) +def train_step(model: nnx.Module, optimizer: nnx.Optimizer, batch: dict[str, jax.Array]) -> jax.Array: + """One forward/backward/update. Mutates `model` and `optimizer` in place.""" + loss, grads = nnx.value_and_grad(loss_fn)(model, batch) + optimizer.update(model, grads) + return loss + + +def _log_step(step: int, loss: jax.Array, last_done: float, tokens_per_step: int) -> tuple[float, float]: + """Blocks until `step` finishes, logs it, and returns its loss and completion time.""" + loss = float(loss) + done = time.perf_counter() + step_time = done - last_done + max_logging.log(f"step {step}: loss {loss:.4f}, {step_time:.3f}s, {tokens_per_step / step_time:,.0f} tokens/s") + return loss, done + + +def train_loop(config) -> float: + """Builds the model and optimizer, runs `config.steps` steps, and returns the last loss.""" + # `create_mesh()` defaults (fsdp within a slice, data across slices) are used as + # is; the `ici_*`/`dcn_*` config fields are not wired through yet. + mesh = create_mesh() + max_logging.log(f"m3 loop: mesh {dict(mesh.shape)} over {len(jax.devices())} device(s)") + + # `global_batch_size_to_train_on` already includes gradient accumulation, and + # the loop does none, so GA > 1 would silently run one GA-times-larger step. + if config.gradient_accumulation_steps > 1: + raise ValueError( + "m3 loop does not accumulate gradients yet; " + f"got gradient_accumulation_steps={config.gradient_accumulation_steps}." + ) + batch_size = config.global_batch_size_to_train_on + batch_shards = math.prod(mesh.shape[axis] for axis in _BATCH_AXES) + if batch_size == 0 or batch_size % batch_shards: + raise ValueError( + f"per_device_batch_size={config.per_device_batch_size} gives a global batch of {batch_size}, which does not " + f"split over the {batch_shards} shards of mesh axes {_BATCH_AXES}." + ) + + with jax.set_mesh(mesh): + model = m3_models.create_model(config, mesh, rngs=nnx.Rngs(config.init_weights_seed)) + 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))) + tokens_per_step = config.global_batch_size_to_train_on * config.max_target_length + max_logging.log( + f"m3 loop: {config.model_name}, {num_params / 1e9:.3f}B params, " + f"batch {config.global_batch_size_to_train_on} x {config.max_target_length} tokens" + ) + + prof = profiler.Profiler(config, offset_step=0) + data_key = jax.random.key(config.init_weights_seed) + loss = float("nan") + # 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 + # dispatching step N keeps one step queued. Step time is the interval + # between consecutive completions, so step 0's includes compilation. + pending = None # (step, loss) dispatched but not yet logged. + 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)) + with jax.profiler.StepTraceAnnotation("train", step_num=step): + step_loss = train_step(model, optimizer, 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)) + if pending is not None: + loss, last_done = _log_step(*pending, last_done, tokens_per_step) + + if config.profiler: + max_logging.log(f"m3 loop: profile written under {config.tensorboard_dir}") + return loss + + +def main(argv: Sequence[str]) -> None: + config = pyconfig.initialize(argv) + train_loop(config) + + +if __name__ == "__main__": + app.run(main) diff --git a/tests/unit/m3/train_loop_test.py b/tests/unit/m3/train_loop_test.py new file mode 100644 index 0000000000..223d59fb65 --- /dev/null +++ b/tests/unit/m3/train_loop_test.py @@ -0,0 +1,101 @@ +# Copyright 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. + +"""Tests for the minimal m3 training loop.""" + +import glob +import math +import os + +from absl.testing import absltest +import jax +import numpy as np + +from maxtext.configs import pyconfig +from maxtext.m3.core.sharding import create_mesh +from maxtext.m3.train import loop +from tests.utils.test_helpers import get_test_config_path + + +def _tiny_config(output_dir: str, **overrides): + """A 2-layer qwen3-0.6b variant small enough to train on CPU in seconds. + + The `base_*` keys are what the model YAML sets; overriding the derived names + (`emb_dim`, `mlp_dim`, ...) is silently ignored. + """ + args = { + "model_name": "qwen3-0.6b", + "use_m3_model": "true", + "scan_layers": "false", + "override_model_config": "true", + "base_num_decoder_layers": "2", + "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": "2", + "max_target_length": "16", + "steps": "3", + "enable_checkpointing": "false", + "base_output_directory": output_dir, + "run_name": "m3_loop_test", + } + args.update({k: str(v) for k, v in overrides.items()}) + return pyconfig.initialize(["", get_test_config_path()] + [f"{k}={v}" for k, v in args.items()]) + + +class TrainLoopTest(absltest.TestCase): + """Runs the loop end to end on a tiny model.""" + + def setUp(self): + super().setUp() + self.output_dir = self.create_tempdir().full_path + + def test_tiny_config_takes_effect(self): + """Guards the `base_*` overrides; a full-size model here would make the suite slow.""" + config = _tiny_config(self.output_dir) + self.assertEqual(config.emb_dim, 64) + self.assertEqual(config.num_decoder_layers, 2) + + def test_synthetic_batch_targets_are_shifted_inputs(self): + config = _tiny_config(self.output_dir) + batch = loop.make_synthetic_batch(config, create_mesh(), jax.random.key(0)) + inputs, targets = np.asarray(batch["inputs"]), np.asarray(batch["targets"]) + 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_writes_xplane_profile(self): + config = _tiny_config(self.output_dir, profiler="xplane", skip_first_n_steps_for_profiler=1, profiler_steps=1) + loop.train_loop(config) + traces = glob.glob(os.path.join(config.tensorboard_dir, "**", "*.xplane.pb"), recursive=True) + self.assertNotEmpty(traces, f"no .xplane.pb under {config.tensorboard_dir}") + + def test_rejects_gradient_accumulation(self): + with self.assertRaisesRegex(ValueError, "gradient_accumulation_steps=2"): + loop.train_loop(_tiny_config(self.output_dir, gradient_accumulation_steps=2)) + + def test_rejects_batch_that_does_not_split_over_mesh(self): + """`per_device_batch_size=0.5` never splits: batch 0 on one device, n // 2 on n devices.""" + with self.assertRaisesRegex(ValueError, "does not split over"): + loop.train_loop(_tiny_config(self.output_dir, per_device_batch_size=0.5)) + + +if __name__ == "__main__": + absltest.main()