Dev/jax maxtext v26.6 primus cli

[amd-fuyuajin]

Summary

Two independent pieces of work. Both are confined to examples/ and
primus/backends/; no third_party/ submodule is touched.

MI325X maxtext configs (22 files). Copied from the MI300X set, then
per_device_batch_size tuned on tus1-p3-g26 (8x MI325X, 256GB HBM vs 192GB on
MI300X) so the suite uses the memory the larger cards actually have instead of
inheriting MI300X limits. Each value is the largest that completed 3 training
steps against the 230.4GB per-device budget. 12 of the 19 tuned configs are
pinned by a real OOM one step higher; the remaining 7 stopped at the sweep's
attempt cap while still passing and are annotated in-file as not OOM-pinned.
grok1, gemma4_26B and gemma4_31B nanoo_fp8 were excluded from the sweep and keep
their MI300X batch sizes.

MaxDiffusion throughput metrics. The step line carried step time, FLOP rate
and loss but no throughput, so benchmark runs had no sample or token rate to
compare across resolutions and clip lengths.

Before:

completed step: 18, seconds: 11.722, TFLOP/s/device: 348.548, loss: 1.540

After:

completed step: 18, seconds: 11.716, TFLOP/s/device: 348.750, Tokens/s/device: 6760.245, Samples/s/device: 0.0854, Frames/s/device: 7.255, loss: 1.540

• Samples/s/device — videos or images per second per GPU. Defined for every
  model family.
• Frames/s/device — Samples/s/device * num_frames, emitted only for families
  whose config describes frames.
• Tokens/s/device — patchified latent tokens per second per GPU: 79,200 for a
  720x1280x85 WAN video, 1,024 for a 512px FLUX image. Rebuilt from the same
  latent geometry calculate_{wan,flux}_tflops bills its FLOPs for, and omitted
  for the UNet families (SD 1.x/2.x, SDXL) which have no equivalent notion.
  Named to match MaxText's step line so existing log parsers apply.

This is implemented as a before_train patch under
primus/backends/maxdiffusion/patches/, the same mechanism as the existing
logger patch, rather than editing third_party/maxdiffusion, so it survives
submodule updates. The values are added to metrics["scalar"] before any
consumer runs, so they also reach TensorBoard and metrics_file alongside the
cluster-wide perf/samples_per_second, perf/frames_per_second and
perf/tokens_per_second totals.

Two follow-ups in the same area: the tensorboard hint is now logged once per run
instead of on every log_period (the configs pin log_period: 1, so it was
doubling the length of the step log), and the log_period comment in the six
maxdiffusion configs is corrected — it claimed the value gated the
completed step: line, which it never did.

Verification

• Recomputing the FLOP model from the checkpoint geometry with the patch's token
  count reproduces the trainer's own Calculated TFLOPs per pass: 4085.7969 to a
  relative error of 3.7e-9, confirming the token rate and the FLOP rate describe
  the same sequence.
• Reported step time averages 11.719s against 11.750s of wall clock between
  consecutive log lines, so 0.26% falls outside the measured window.
• Metric injection happens before the metrics-file write and the TensorBoard
  queue, and the double-buffered step-to-metric pairing is preserved.
• Degradation paths: an unknown model family reports samples only; an unusable
  per_device_batch_size warns once, disables, and falls back to the
  byte-identical upstream line.

Test plan

• wan2.1_1.3b on MI300X, 20 steps: 348.88 TFLOP/s/device, 6762.7
  Tokens/s/device, 7.258 Frames/s/device

• MI325X 1-node suite, steady-state means over steps 2-18:

  wan2.1_1.3b   11.722s   348.56 TFLOP/s/dev   6756.5 Tok/s/dev   7.251 Frames/s/dev
  wan2.1_14b    65.730s   322.62 TFLOP/s/dev   1204.9 Tok/s/dev   1.293 Frames/s/dev
  flux_dev       1.194s   432.30 TFLOP/s/dev   6862.2 Tok/s/dev   (no frames, image model)
  

• MI355X maxdiffusion configs not re-run; only a comment changed there.

Note for whoever aggregates these numbers

Warmup must be excluded. Steps 0 and 1 are compilation (368s and 1032s in the
MI300X run), and averaging them in skews step time by +628% and the rates by
-10.3%. tools/daily/daily_report.py currently averages every matched step.

[amd-fuyuajin]

@yeandy Please help to review this PR.