[perf]:

[SolitaryThinker]

Summary

This PR is the cumulative LTX-2 BF16 training-efficiency stack. It keeps registered/master parameters and Adam state in FP32, uses BF16 only for gathered working parameters, compute, gradients, and reductions, and adds no dependency.

Current tracker head: 52f1114dd95934903e49620338650780d8152b8d

Measured optimization head: 20c36acefc97e8b743f79a5c52883561853a7d85

Final validation commits: 0e60a0e9ccbd0c62d730fcfcf78f7d4a4d0554e7, 3f3f06541c1a4ab01ddb81c4840d42040fd02a38

The optimized recipe reaches 40.810314% MFU on 4x GB200 and 43.623053% MFU on the healthy 8x/two-tray allocation. A replacement 8x allocation was power/clock limited and measured 30.478319% MFU; both 8x results are reported because the absolute difference is allocation-specific, not attributable to a source change. The 50% target is not reached.

What changed

• Add opt-in fused AdamW, BF16 FSDP reductions, no-forward-reshard, native symmetric-memory collectives, accumulation-aware synchronization/parameter retention, and public FSDP2 module grouping.
• Regionally compile the transformer, select the video-only LTX-2 path for video training, batch gradient clipping, defer gradient-norm host materialization, and keep uniform timesteps singleton until sequence-parallel expansion is required.
• Repeat/prefetch the overfit input and persistently pack video self-QKV and text cross-KV projections while preserving strict split-HF load/export, sequence-parallel parity, and FP32 registered/master weights and optimizer state.
• Keep validation inference out of the training compile cache, skip unused audio work for video-only denoising, and unload the validation pipeline before training resumes.
• Enable the measured options in overfit_ltx2_t2v.yaml; defaults preserve existing behavior for other recipes.

Ongoing experiment tracker

This PR and branch are now the cumulative scratch space for MFU optimization attempts. Later we will extract focused, independently mergeable PRs from it.

• scripts/train/ltx2_mfu/README.md is the resume handoff, benchmark contract, current stopping point, and decision index.
• REPORT.md preserves the chronological accepted/rejected experiment history.
• harness/, runners/, probes/, and reports/ preserve the generated source-like artifacts. run_current.sh is the maintained 4x/8x MFU entrypoint; run_observed.sh is the maintained 4x W&B/validation entrypoint.
• Raw logs, traces, checkpoints, videos, telemetry dumps, credentials, generated binaries/CUDA, and copied third-party sources remain outside Git; the report records run/job IDs, versions, and hashes instead.

Final MFU

Common contract: FastVideo/LTX2-Distilled-Diffusers, 81x480x832, dense FA4, regional compile, fused AdamW, repeated/prefetched input, singleton timestep, no validation inside the MFU timer, and the slowest-rank median of 10 warmup + 20 measured steps. Completed rows use no activation checkpointing; the 1x capacity audit used full checkpointing. MFU uses 0.353881 PFLOP/sample and NVIDIA's 2.450 PFLOP/s dense BF16 peak per GB200.

GPUs	Local/global batch	Median step	Aggregate throughput	MFU	Result
1	1 / 1	-	-	-	No valid result: OOM during first-step FP32 Adam-state creation
4	2 / 8	0.707865914446 s	11.301575 samples/s	40.810314%	Final measured stack at 20c36acef
8	3 / 24	0.993335918058 s	24.161011 samples/s	43.623053%	Healthy/fast two-tray allocation
8	3 / 24	1.421743279090 s	16.880685 samples/s	30.478319%	Replacement allocation, power/clock degraded

The 4x B2 result improves over equal-global-batch B1/gradient-accumulation-2 controls at 0.743365705741 s / 38.861435% MFU: -35.499791 ms / -4.775549%, +5.014942% throughput, +1.948878 MFU points, with 0.198028% control drift.

On the degraded 8x pair, group 2 is neutral against its same-pair 1.419745997933 s / 30.521241% control midpoint: +1.997281 ms / +0.140679% and -0.042922 MFU points, smaller than the 0.245257% control drift. The pair nevertheless passed an eight-rank MNNVL/NVLS health gate at 463.195 GB/s derived bus bandwidth; telemetry showed clocks capped near 1200 MHz. The healthy and degraded rows must not be used as a source A/B comparison.

Validation

• Changed files pass pre-commit; the final focused suite passes 31 tests (16 warnings) on GB200.
• Final 4x/B2 validation gate passed 51 optimizer steps with validation at steps 0 and 50, then resumed through step 51 without a second compile or OOM. The run used the exact final file bytes staged on measured head 20c36acef; those bytes are committed at final source head 3f3f06541. W&B run iopu4dwm; log SHA-256 d84ae1a648d37d413a8e6116dbbad855d62cec0f30c3e016a506a4b4fb17c5f1.
• A separate 4x/B1 overfit completed 300 optimizer steps with validations at steps 0/50/100/150/200/250/300 and final loss 0.01216: W&B run.
• The exact 4x MFU log SHA-256 is 645c68e27f3967b8352efdc2ed4690a59a66cad5b644f41350ee7b0acb4df75e.
• The 1x capacity log SHA-256 is aa260691ce2847aa1e5cf5c852fec8c2af67de306b4bf393557bdbc3a4477d54.
• Healthy 8x packed-candidate node log SHA-256 values are 5e1f03976d834465cd60072459f960f40b4d0ac4dfba7d24e7b4856f4e615bcc and 65ea51b8b98c21da82eb0a93d64dbc341b748dc5ecefdf6b4d2f033de56a4324.
• Degraded 8x group-2 node log SHA-256 values are 0277cbcaac04bac957bf1cfda000a365ccde4a15b5ae7b947b61c2a4237baefd and 30c058a42a9f35329ec191afba476aed5f1b3f3d33b4957046a6c384eab7d82e.
• Packing passed focused, distributed, export, and sequence-parallel checks; strict FP32 export/reload split 192 packed tensors into 480 bit-exact HF projections without changing parameter elements.
• Dense FA4 forward/backward parity and the CUDA 13 dependency fix are covered by [bugfix]: align FA4 dependency stacks #1629.

Risks and rollback

• The 1x standard-precision recipe does not fit a GB200 because FP32 Adam-state creation exhausts memory. Optimizer offload or reduced-precision state would change the benchmark contract and is intentionally not included.
• The 8x absolute result depends materially on allocation power/clock behavior; use matched same-allocation controls for attribution.
• Packed projections and communication/compile policies are guarded opt-ins. Disable the corresponding recipe options to return to split projections and legacy FSDP behavior.
• 50% MFU likely requires a materially faster complete projection-training path or a productionized parameter arena. Scratch low-precision kernels and the fixed arena did not meet the correctness, lifecycle, or integration gates, so neither is shipped here.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success~=pre-commit

This rule is failing.

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success~=pre-commit
• check-success=fastcheck-passed
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request adds support for the fused AdamW optimizer parameter across configurations, parsing, and optimizer initialization. Feedback highlights a critical issue in seed_optimizer_state_for_resume where the step state tensor shape must be (1,) instead of () when fused or capturable is enabled to prevent shape mismatch errors during Distributed Checkpoint (DCP) restoration. It is also recommended to update the unit tests to assert these correct tensor shapes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

When fused or capturable is enabled, PyTorch's AdamW (and Adam) initializes the step state tensor as a 1D tensor of shape (1,) (to allow in-place updates during CUDA graph replay or fused kernel execution). For the default (non-fused, non-capturable) path, it initializes step as a 0D scalar tensor of shape ().

Because Distributed Checkpoint (DCP) requires exact shape matching between the pre-seeded optimizer state and the saved checkpoint, initializing step as a 0D tensor when fused or capturable is enabled will cause a shape mismatch RuntimeError when resuming from a checkpoint.

Additionally, PyTorch forces step to be torch.float32 when either fused or capturable is enabled, so step_dtype should also check capturable to avoid using torch.float64 if the default dtype is set to float64.

Suggested change

	step_device = p.device if fused or capturable else torch.device("cpu")
	step_dtype = (torch.float64
	if not fused and torch.get_default_dtype() == torch.float64 else torch.float32)
	opt.state[p] = {
	"step": torch.tensor(0.0),
	"step": torch.zeros((), dtype=step_dtype, device=step_device),
	"exp_avg": torch.zeros_like(p),
	"exp_avg_sq": torch.zeros_like(p),
	}
	step_device = p.device if fused or capturable else torch.device("cpu")
	step_dtype = (torch.float64
	if not (fused or capturable) and torch.get_default_dtype() == torch.float64 else torch.float32)
	step_shape = (1,) if fused or capturable else ()
	opt.state[p] = {
	"step": torch.zeros(step_shape, dtype=step_dtype, device=step_device),
	"exp_avg": torch.zeros_like(p),
	"exp_avg_sq": torch.zeros_like(p),
	}

[gemini-code-assist]

Update the resume seeding tests to assert the correct tensor shapes for both the fused/capturable ((1,)) and default (()) optimizer step states to prevent future regressions.

def test_resume_seed_places_fused_step_with_parameter() -> None:
    parameter = torch.nn.Parameter(torch.empty(1, device="meta"))
    optimizer = torch.optim.AdamW([parameter], fused=True)
    method = SimpleNamespace(get_optimizers=lambda _: [optimizer])

    TrainingMethod.seed_optimizer_state_for_resume(method)

    assert optimizer.state[parameter]["step"].device == parameter.device
    assert optimizer.state[parameter]["step"].dtype == torch.float32
    assert optimizer.state[parameter]["step"].shape == (1,)


def test_resume_seed_keeps_default_step_on_cpu() -> None:
    parameter = torch.nn.Parameter(torch.empty(1, device="meta"))
    optimizer = torch.optim.AdamW([parameter])
    method = SimpleNamespace(get_optimizers=lambda _: [optimizer])

    TrainingMethod.seed_optimizer_state_for_resume(method)

    assert optimizer.state[parameter]["step"].device.type == "cpu"
    assert optimizer.state[parameter]["step"].shape == ()

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease