Add minimal m3 training loop for profiling

[bvandermoon]

Description

Adds a minimal training loop for stepping an m3 model end to end and collecting an xplane profile: synthetic next-token data, AdamW at a constant LR, no checkpointing or eval. This is a minimal setup to start collecting profiles.

Each step's loss is read after the next step is dispatched, so the host never stalls the device on a per-step sync. Step time is completion to completion, as in legacy, so profiler start/stop lands in the adjacent step's time.

Known gaps: m3 models carry no sharding annotations yet, so under the default
create_mesh() (fsdp=-1) parameters and optimizer state are replicated and the step is pure data
parallelism. ici_* config fields are ignored. The loop is still host-bound at small batch sizes.

Tests

Collected profile: https://xprof.corp.google.com/overview_page/bvandermoon-4143076710901394919

Qwen3-0.6B on v6e-8, per_device_batch_size=3 max_target_length=2048. per_device_batch_size=4
OOMs at compile by 18 MB (full [B,S,vocab] logits and [B,H,S,S] attention scores per layer).

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=3 max_target_length=2048 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

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces a minimal training loop for m3 models in src/maxtext/m3/train/loop.py along with unit tests in tests/unit/m3/train_loop_test.py. The loop is designed for profiling and uses synthetic data, AdamW, and next-token cross entropy. Feedback on the changes includes JIT-compiling the synthetic batch generation function to avoid eager host dispatch overhead, and removing donate_argnums=(0, 1) from the train_step JIT decorator to prevent duplicate buffer donation errors since nnx.Optimizer already references the model.

[codecov]

Codecov Report

❌ Patch coverage is 91.56627% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/m3/train/loop.py	91.56%	4 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!

[hengtaoguo]

LGTM!

[github-actions]

🤖 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]

📋 Review Summary

Verdict: Spec met; the loop successfully overlaps execution to avoid per-step syncs, but eager batch generation artificially starves the device.

Top issues:

1. [Perf] Eager generation of the synthetic batch causes expensive host-device transfers and syncs, creating artificial host bottlenecks that pollute the profile trace. → JIT the batch generation with out_shardings to generate the data directly distributed on devices.

Inferred spec & validation

• Sources: PR description, code diff.
• Problem: Needs a minimal m3 training loop to step an m3 model end-to-end to collect a clean xplane profile without host stalls.

Spec item	Implementation	Validating test
R1 (stated): Synthetic next-token data	✅	test_synthetic_batch_targets_are_shifted_inputs ✅
R2 (stated): AdamW at constant LR	✅	test_train_loop_returns_finite_loss ✅
R3 (stated): No checkpointing or eval	✅	No tests (omitted by design) ❌
R4 (stated): Host never stalls device on a per-step sync	⚠️ partial (starved by eager data gen)	test_train_loop_returns_finite_loss ⚠️ superficial
R5 (stated): Profiler writes an xplane profile	✅	test_train_loop_writes_xplane_profile ✅

[RissyRan]

It seems there re 4 must fix issues, and please consider to address. Thanks!

[bvandermoon]

It seems there re 4 must fix issues, and please consider to address. Thanks!

Done, thanks @RissyRan