Feat/diffusion/wan21 megatron training by jianhan-amd · Pull Request #1132 · AMD-AGI/Primus

jianhan-amd · 2026-09-10T13:54:21Z

WAN 2.1 / 2.2 video DiT training on Megatron-Core

What this adds

End-to-end training support for the WAN video diffusion transformer on
Megatron-Core, covering the model, the data path, the trainer, runnable example
configs, and a checkpoint conversion CLI. Along the way it fixes three defects
that WAN surfaced but that are not WAN-specific, and adds unit coverage for the
invariants the new code depends on.

Validated by training WAN 2.1 T2V 1.3B on a pre-encoded PusaV1 dataset across
three precision arms — TransformerEngine, Primus-Turbo local bf16, and MXFP4 —
which differ only in their precision block so they can be diffed against each
other.

Layout

The 11 commits are grouped so each can be reviewed on its own; the history is
bisectable.

Commits	Area
`fbab4c4`	The `WanTransformerBlock` / DiT backbone on Megatron-Core
`e10fb93`, `55bb06f`	Data: Energon shards, UMT5 / VAE encoders, synthetic providers
`97f2576`	Trainer, forward step, flow-matching scheduler and loss
`f1e1d15`	Example configs, `wan_train.sh`, HF→Primus conversion CLI
`55f98c4`, `13663e9`	Two fixes outside WAN (MXFP4 local spec, TE hipBLASLt workspace)
`c02e66b`, `b850687`, `a194ccc`	Three performance fixes in the WAN path
`df2d5d7`	Unit tests

Points worth a reviewer's attention

WanTransformerBlock is registered as a virtual subclass of
TransformerLayer (c02e66b). Megatron-FSDP shards at the granularity of its
"FSDP unit modules", but nothing passes that list for WAN — training.py builds
the wrapper with no fsdp_unit_modules and DistributedDataParallelConfig has
no field for one, so the adapter falls back to its default of
[TransformerLayer]. WAN's block deliberately is not a Megatron
TransformerLayer, so that default matched nothing and the unit list came out
empty, which is worse than coarse sharding: the hook-registration loop skips
modules inside a registered unit, and with nothing registered that skip never
fires, so every norm, linear and attention takes an unshard hook plus an
fp8-transpose post-hook a bf16 run has no use for. ABCMeta.register affects
isinstance only and leaves the checkpoint key layout alone. The alternatives
were worse — Megatron-LM is an upstream submodule, so plumbing the argument
through training.py means forking it, and real inheritance would change
checkpoint keys.

One YAML key drives two backends' attention kernel selection (b850687,
a194ccc). WAN's sequence length is not divisible by 64, which puts both
Primus-Turbo and TE's ROCm CK backend in the same corner: the kernel families
serving non-padded shapes accumulate dQ through atomics, so with the fp32

The backbone and everything needed to describe it: Conv3d patch embed, 30 transformer blocks with self- and cross-attention, AdaLN modulation from the timestep embedding, and the unpatchify head. Attention and the MLP are built from BackendSpecProvider submodules, so the same model runs on TransformerEngine, Turbo local bf16, Turbo FP8 or MXFP4 without a code change. backend_resolution.py holds the provider ladder. It mirrors the one inlined in flux/layer_spec.get_flux_layer_spec rather than sharing it, which keeps this change off the Flux path entirely; the docstring records that a new precision branch has to be added in both places. The checkpoint converter maps HuggingFace WAN weights onto this layout, including the per-block QKV fusions, and is what the accompanying CLI drives. Nothing here is reachable yet: the data pipeline, the trainer and the configs that select this model arrive in the following commits. Co-authored-by: Cursor <cursoragent@cursor.com>

WAN conditions on UMT5 text embeddings and trains on latents from the WAN video VAE, so both join the encoder registry alongside the Flux encoders, with the config dataclasses the registry needs to build them. The video task encoders cook Energon samples in both directions: raw clips that still need the VAE and text encoder, and shards that were pre-encoded offline, which is what the benchmark runs use. MockWanDataset and its pre-generated variant cover the no-dataset path. Selecting them needed a 'family' key in SyntheticDatasetProvider: the trainer collapses model_type to 'flux' before the provider ever sees it, so with model_type alone 'wan' was unreachable. Co-authored-by: Cursor <cursoragent@cursor.com>

WanPretrainTrainer wires the model and the data pipeline into the Megatron training loop and is registered with the adapter, so a config can now ask for it by name. The forward step takes VAE latents and text embeddings, samples a timestep, asks the scheduler for the noised input and the training target, and reduces to a loss. WanFlowMatchScheduler supplies the sigma schedule and those targets. The objective needs a per-timestep weight on the squared error, which the existing compute_flow_matching_loss does not offer, so compute_weighted_flow_matching_loss joins it: the target comes from the caller rather than being derived as noise - latents, and weight=1.0 recovers the unweighted loss exactly. Co-authored-by: Cursor <cursoragent@cursor.com>

EnergonDatasetProvider hardcoded shuffle_over_epochs_multiplier and parallel_shard_iters to Energon's defaults, so no config could reach them. All four shuffle knobs have to be off together to get a deterministic sample order -- max_samples_per_sequence and shuffle_buffer_size null, the multiplier null, and parallel_shard_iters 1, the last required once shard shuffling is off. That order is what makes a step-by-step loss comparison against another implementation meaningful: with shuffling on, two runs draw different samples and the losses cannot be compared no matter how correct both are. The defaults are unchanged for every config that does not set them. Co-authored-by: Cursor <cursoragent@cursor.com>

Three MI355X arms for WAN 2.1 T2V 1.3B on a pre-encoded PusaV1 dataset -- TransformerEngine, Turbo local bf16, and MXFP4 -- differing only in the precision block, so they can be diffed against each other. Paths come from PRIMUS_DIFFUSION_DATA_PATH, PRIMUS_TEAM and PRIMUS_USER, following the Flux examples. wan_train.sh runs them and pins two settings that are wrong for WAN by default: AITER for the FP4 matmuls, because the HipBLASLt/rocRoller path runs a CPU-bound per-shape solution search that stalls on long-sequence shapes, and MIOPEN_FIND_MODE=5, because run_pretrain.sh pins Fast mode, which is inert for LLMs but resolves WAN's Conv3d patch-embed backward to a naive im2col path instead of the fused Composable Kernel wgrad solver -- on its own the dominant cost of a training step. Mode 5 is MIOpen's own default and searches once per new conv shape, so only the first iteration pays. prepare.py returns early for diffusion runs, which read Energon shards and have no tokenizer_model, so the bookcorpus tokenisation flow must not fire. attention_backend is declared in diffusion_model.yaml so a config can override it: "auto" refuses to run when the image pins NVTE_FLASH_ATTN or NVTE_FUSED_ATTN, as these containers do. Co-authored-by: Cursor <cursoragent@cursor.com>

… dtypes Three defects in the MXFP4 local spec, all of which surfaced running WAN but none of which are WAN-specific. Unbacked scale blocks. A shuffled colwise scale buffer covers cdiv(cdiv(d, 32), 8) * 8 E8M0 blocks while the packed FP4 data covers only cdiv(d, 128) * 128. When cdiv(d, 128) is odd, four blocks per logical row have no backing data: the quantizer never writes them and the wgrad GEMM sums whatever the allocator left behind, which shows up as a gradient norm many orders of magnitude too large that varies run to run. Any axis an FP4 GEMM reduces over is now padded to 256, which is a multiple of the 16 that AITER's dispatcher demands and makes cdiv(256t, 128) even, so every block is backed. All three of M, K and N get the treatment, since each is the reduction axis of one of the three GEMMs behind a linear. Padding stays conditional because it is not free: _pad2d copies the whole tensor, not just the appended rows, at every FP4 linear every step. Padding unconditionally gives back the entire MXFP4 win, so axes that already satisfy both rules are left alone -- which for WAN 1.3B means nothing is padded at all. Dtype conflation. The forward emitted at the activation dtype and the backward reused one ctx.out_dtype for both gradients. Diffusion backbones feed FP32 activations, which the quantize kernels reject and the FP4 backends cannot emit. GEMMs now run at the parameter dtype, and each gradient is cast back to the dtype it belongs to: grad_input to the forward input's, grad_weight to the parameter's, so FSDP and the optimizer see what they expect. Older Primus-Turbo. quantize_mxfp4_dual gained padding_align_size and gemm_fp4_impl gained preshuffled= after rocm/primus:v26.3, so the modern signatures raise there. Both are probed from the registered schema rather than inspect.signature, which reports False on every image for gemm_fp4_impl -- it is a CustomOpDef taking (*args, **kwargs) -- and would silently drop preshuffled on modern images, handing unshuffled operands to a kernel expecting shuffled ones. Co-authored-by: Cursor <cursoragent@cursor.com>

TE fixes its GEMM workspace at 64 MiB on gfx950 with no environment override. For a weight-gradient GEMM with a long reduction and a small square output, hipBLASLt answers with a split-K solution whose per-split accumulators grow with the reduction length; once they no longer fit, the call fails with "HIPBLASLT Error: 6" rather than falling back to a solution that does fit. WAN 2.1 T2V-1.3B hits this on its square 1536 -> 1536 wgrad once the token count passes ~70k, which is micro_batch_size 7 at 10,920 tokens per sample. The probe was 10,920 / 21,840 / 43,680 / 65,520 / 70,000 passing and 75,000 / 80,000 / 87,360 / 98,280 / 109,200 failing. The wider linears in the same model (1536 -> 8960 and 8960 -> 1536) never take the split-K path and are fine at every size. The patch rebinds the size and clears the memoised per-device allocation, and is gated on PRIMUS_TE_GEMM_WORKSPACE_MIB so it applies only where a workload asks for it. Unset, TE's own default stands, so no other workload silently gets a different hipBLASLt solution. wan_train.sh sets 128 MiB, which clears every WAN shape. Observed on rocm/primus:v26.5 (TE 2.15.0.dev0). Co-authored-by: Cursor <cursoragent@cursor.com>

Megatron-FSDP gathers and releases parameters at the granularity of its "FSDP unit modules". Nothing passes that list for WAN: megatron/training/training.py builds the wrapper with no fsdp_unit_modules argument and DistributedDataParallelConfig has no field for one, so the adapter falls back to its default of [TransformerLayer] for the optim_grads_params strategy WAN runs under. WanTransformerBlock is deliberately not a Megatron TransformerLayer, so that default matched zero modules and the unit list came out empty. An empty list is worse than coarse sharding, because the hook-registration loop skips modules that live inside a registered unit and with nothing registered that skip never fires. Every norm, linear and attention in every block then takes a pre-forward unshard hook plus an fp8-transpose-cache post-hook that a bf16 run has no use for, instead of one hook per block. Any hook forces nn.Module._call_impl off its no-hook fast path, so the whole model was also gathered as one unit and the Dynamo graph broke at every submodule boundary inside a block rather than only at the block boundary. Registering the block as a virtual subclass makes Megatron's own default match it. ABCMeta reaches TransformerLayer through BaseTransformerLayer(ABC), so register() is available and affects isinstance only. The alternatives were worse: Megatron-LM is an upstream submodule, so the argument cannot be plumbed through training.py without forking it, and real inheritance from TransformerLayer would change the checkpoint key layout. Eager trades a little step time for peak memory, because per-block all-gathers cost more collective time than a single whole-model gather. Compiled improves on both, since the de-fragmented graph lets Inductor elide casts it could not move across a break, and both the AOTAutograd graph count and the number of Triton kernels fall substantially. Loss is unchanged. The other Megatron sites that test for TransformerLayer are either unreachable for WAN (CUDA graphs, Mamba and hybrid blocks, GPT callables) or keyed on layers.N parameter names where WAN's are blocks.N. layer_number is now carried on the block for the one attribute the FSDP-DTensor checkpoint path reads. Co-authored-by: Cursor <cursoragent@cursor.com>

Primus-Turbo selects its flash-attention backward variant from PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32, re-reading it on every call. Turbo's own default is "1", but examples/run_pretrain.sh and runner/helpers/envs/ base_env.sh both export ${PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32:-0}, so every launched run gets "0" as a blanket accuracy default. For WAN that default does not pick a slower-but-safer aiter kernel, it leaves aiter entirely. The sequence length is not divisible by 64, so the non-padded ASM family is out, and the psskddv family that would otherwise serve these shapes requires the atomic accumulator. Dispatch falls through to ck_tile without reporting that it did. Turning it on replaces ck_tile's backward with aiter::fmha_bwd_hd128_bf16_a32_rtna_psskddv over the same number of calls, and the whole attention backward moves to the ASM family. Atomic accumulation is order-dependent, so this is not free in principle; over a 20-step run the loss stayed within bf16 reduction noise rather than showing a behaviour change. The opt-out is the new attn_atomic_fp32 model key rather than the environment variable, because the launchers export that variable unconditionally and so every WAN run reaches the trainer with it already set to "0". Reading it back could not distinguish a deliberate choice from the blanket default that this change works around. Set attn_atomic_fp32: false for the deterministic split-dQ path. Scoped to the WAN trainer, so no other model's kernel selection changes. Co-authored-by: Cursor <cursoragent@cursor.com>

attn_atomic_fp32 only steered Primus-Turbo, so it did nothing on the te_spec path: TransformerEngine's ROCm CK backend keeps its own copy of the same choice in NVTE_CK_IS_V3_ATOMIC_FP32, and the release Dockerfile pins that to 0, which is the value gfx950 wants. On gfx942 the CK v3 backward needs the fp32 atomics, so v3 was unreachable and CK served the backward from ck_tile without saying so -- the same silent fallback the local spec already had, one layer down. Set both variables from the one YAML key, and force NVTE_CK_USES_BWD_V3 on when enabling so a launcher that disabled v3 cannot quietly undo the atomics. tools/installation/env.sh already prescribes this pairing for gfx942 on the bare-metal path; this brings the containerised WAN runs in line with it. Confirmed on a synthetic WAN 2.1 1.3B run with the whole attention backward changing family: the ck_tile dQ/dK/dV, OGradDotO and ConvertQGrad kernels give way to their aiter counterparts over the same call counts, and the forward is untouched. No Inf or NaN in either arm. The accumulation order is non-deterministic; set attn_atomic_fp32: false for the deterministic split-dQ backward on either spec. Co-authored-by: Cursor <cursoragent@cursor.com>

Five invariants the preceding commits rely on that nothing checked. The FP4 reduction-axis padding is the load-bearing one. It exists to stop the wgrad GEMM summing scale blocks with no backing data, whose symptom is a gradient norm that varies run to run, and it is conditional because padding copies the whole operand at every FP4 linear. Both halves of that rule need holding down, so the alignment is tested as integer math: the dims that must pass through untouched, the ones that must be padded up, per-axis independence for M, K and N, and a sweep asserting that whatever comes out clears the dispatcher and backs every block. The cross-validation against FP4GemmMXFunction also gains a shape that pads nothing. Every other numeric test in that file runs at 128 rows, which the rule pads to 256, so the unpadded path -- the common case for Wan's own dims -- had no coverage against the reference at all. The rest: weight=1 recovers the unweighted flow-matching objective exactly, including on the masked reduction path, so Wan and Flux runs stay comparable; WanTransformerBlock answers isinstance(TransformerLayer) without entering its MRO, which is what makes Megatron-FSDP's default unit list match it while leaving the checkpoint key layout alone; attn_atomic_fp32 reaches both backends and overrides the launchers' blanket "0", and opting out does not disable TE's v3 backward as a side effect; and a 'family' key selects a Wan mock dataset while its absence leaves every existing config resolving exactly as before. PrimusUT is a unittest.TestCase, which pytest cannot parametrize, so the provider cases are spelled out rather than swept. Ran on 8x gfx942: 60 passed, 7 skipped. The skips are the gfx950-only MXFP4 numeric tests, which includes the two added here. Co-authored-by: Cursor <cursoragent@cursor.com>

Copilot

🟡 Changes recommended

There are confirmed correctness issues that will break unit tests and can break FP32 runs (tensor boolean assertion and unconditional CUDA autocast dtype handling).

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds end-to-end WAN 2.1/2.2 video DiT training support on Primus Megatron-Core, including the WAN backbone, scheduler/loss plumbing, data/task-encoder path, example configs/launcher, and a HuggingFace→Primus checkpoint conversion tool.

Changes:

Introduces WAN model stack on Megatron-Core (config, layer specs, attention, backbone) plus WAN-specific forward step, scheduler, and weighted loss.
Adds training entrypoints and operational knobs (new WanPretrainTrainer, launcher script, TE hipBLASLt workspace patch, example YAML configs).
Expands diffusion mock/synthetic + Energon support and adds unit tests for WAN invariants (layer spec correctness, attention env plumbing, MXFP4 padding invariants).

File summaries

File	Description
`wan_train.sh`	Adds a WAN-focused training launcher wrapper around `examples/run_pretrain.sh`.
`tools/checkpoint_conversion/convert_wan_hf_to_primus.py`	CLI to convert HF Diffusers WAN checkpoints into Primus/Megatron-compatible format.
`tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py`	Adds MXFP4 reduction-axis alignment and padding-path tests.
`tests/unit_tests/backends/megatron/diffusion/training/test_loss_computation.py`	Adds unit coverage for the new weighted flow-matching loss helper.
`tests/unit_tests/backends/megatron/diffusion/test_wan_layer_spec.py`	Tests WAN layer-spec invariants (q/k RMSNorm, virtual TransformerLayer registration, key tables).
`tests/unit_tests/backends/megatron/diffusion/test_wan_attn_atomic_fp32.py`	Tests WAN trainer’s env plumbing for fp32-atomic attention backward selection.
`tests/unit_tests/backends/megatron/diffusion/data/test_synthetic_datasets.py`	Tests new `family` key for selecting WAN synthetic datasets.
`primus/configs/models/megatron/diffusion/wan2.1_t2v_1.3b.yaml`	Adds WAN 2.1 T2V 1.3B model config (variant overrides).
`primus/configs/models/megatron/diffusion/wan_base.yaml`	Adds shared WAN base model config and WAN-specific knobs (loss weighting, atomic backward, routing).
`primus/configs/models/megatron/diffusion_model.yaml`	Exposes `attention_backend` so diffusion configs can override TE backend choice.
`primus/backends/megatron/wan_pretrain_trainer.py`	Adds WAN-specific trainer (weighted loss, WAN scheduler, attention env overrides, WAN config build).
`primus/backends/megatron/training/diffusion/wan_forward_step.py`	Adds WAN forward step (5D latents + scheduler target/weight return).
`primus/backends/megatron/training/diffusion/schedulers/wan_flow_matching.py`	Adds WAN-template flow-matching scheduler (target/weight/timestep helpers).
`primus/backends/megatron/training/diffusion/schedulers/__init__.py`	Exports `WanFlowMatchScheduler`.
`primus/backends/megatron/training/diffusion/loss_computation.py`	Adds `compute_weighted_flow_matching_loss` helper.
`primus/backends/megatron/patches/te_patches/hipblaslt_workspace_patches.py`	Adds TE hipBLASLt workspace override patch driven by env var.
`primus/backends/megatron/megatron_adapter.py`	Registers `WanPretrainTrainer` in trainer class loader.
`primus/backends/megatron/data/synthetic/mock_datasets.py`	Adds synthetic WAN datasets (on-the-fly + pre-generated).
`primus/backends/megatron/data/synthetic/__init__.py`	Exports WAN synthetic datasets.
`primus/backends/megatron/data/synthetic_dataset_provider.py`	Adds `family` key and WAN defaults for synthetic dataset selection.
`primus/backends/megatron/data/energon_dataset_provider.py`	Threads additional Energon shuffle controls for deterministic ordering.
`primus/backends/megatron/data/diffusion/task_encoders/video.py`	Adds WAN video TaskEncoders (pre-encoded and raw).
`primus/backends/megatron/data/diffusion/task_encoders/__init__.py`	Exports WAN TaskEncoders.
`primus/backends/megatron/data/diffusion/encoders/video/vae/wan/autoencoder_kl_wan.py`	Adds diffusers `AutoencoderKLWan` wrapper for WAN video VAE.
`primus/backends/megatron/data/diffusion/encoders/video/vae/wan/__init__.py`	Packages WAN VAE wrapper.
`primus/backends/megatron/data/diffusion/encoders/video/vae/__init__.py`	Exports video VAE implementations.
`primus/backends/megatron/data/diffusion/encoders/video/__init__.py`	Exports video encoder package entries.
`primus/backends/megatron/data/diffusion/encoders/text/umt5.py`	Adds UMT5 text encoder wrapper for WAN conditioning.
`primus/backends/megatron/data/diffusion/encoders/text/__init__.py`	Exports `UMT5Encoder`.
`primus/backends/megatron/data/diffusion/encoders/config.py`	Adds WAN encoder configs (`WanVAEConfig`, `UMT5Config`, `WanEncoderConfig`).
`primus/backends/megatron/data/diffusion/encoders/__init__.py`	Auto-discovers/exports WAN encoders + config types.
`primus/backends/megatron/core/models/diffusion/wan/utils.py`	Adds WAN 3D latent packing/unpacking helpers.
`primus/backends/megatron/core/models/diffusion/wan/model.py`	Implements WAN backbone + single/dual-expert models and FSDP-unit virtual subclass registration.
`primus/backends/megatron/core/models/diffusion/wan/layers.py`	Adds WAN 3D RoPE + condition embedder + fp32 attention parity path.
`primus/backends/megatron/core/models/diffusion/wan/layer_spec.py`	Resolves WAN block specs from a backend provider (TE/local/FP4/FP8).
`primus/backends/megatron/core/models/diffusion/wan/config.py`	Adds `WanConfig` with routing/stage/window validation and presets.
`primus/backends/megatron/core/models/diffusion/wan/attention.py`	Implements WAN self/cross attention modules (TE packed THD vs local SBHD path).
`primus/backends/megatron/core/models/diffusion/wan/__init__.py`	Exposes WAN model components via lazy imports.
`primus/backends/megatron/core/models/diffusion/common/backend_resolution.py`	Adds shared backend resolution helper used by WAN layer specs.
`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`	Fixes/extends MXFP4 local linear behavior (shape alignment, dtype handling, capability probes).
`examples/megatron/prepare.py`	Skips bookcorpus tokenization flow for diffusion model types (including WAN).
`examples/megatron/configs/MI355X/diffusion/wan2.1_t2v_1.3b_pretrain_pusa_te_spec.yaml`	Adds WAN TE-spec training recipe config.
`examples/megatron/configs/MI355X/diffusion/wan2.1_t2v_1.3b_pretrain_pusa_local_spec.yaml`	Adds WAN local-spec training recipe config.
`examples/megatron/configs/MI355X/diffusion/wan2.1_t2v_1.3b_pretrain_pusa_local_spec_mxfp4.yaml`	Adds WAN local-spec MXFP4 training recipe config.

Review details

Files reviewed: 45/45 changed files
Comments generated: 2
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

jianhan-amd · 2026-09-10T14:52:38Z

+        assert padded.shape == (256, 256)
+        assert torch.equal(padded[:128], tensor)
+        assert not padded[128:].any()
+


This is a false positive, so I've left the assertion as-is.

RuntimeError: Boolean value of Tensor is ambiguous only applies to tensors with more than one element. Tensor.any() returns a 0-dim tensor, and bool() on a 0-dim tensor is well defined, so not padded[128:].any() evaluates normally. Both tests pass on gfx942 as written (2 passed, 21 deselected), and were passing in df2d5d7. Noting for the record that the review summary's "confirmed correctness issues that will break unit tests" doesn't hold for this thread.

Three review findings, none of which change behaviour on the paths the benchmark runs take. The WAN forward step enabled CUDA autocast unconditionally at compute_dtype, which is params_dtype when neither bf16 nor fp16 is set. Asking autocast to cast float32 to float32 is a no-op only by convention, so an fp32 run relied on that convention holding; it is now enabled only for bf16 and fp16. The lazily exported model names sit in __all__ but are bound by __getattr__, which static analysis cannot see. They are now declared under TYPE_CHECKING, which satisfies the checker without binding anything at module scope -- a real assignment would shadow __getattr__, since it only runs when normal lookup fails, and every one of these names would resolve to the placeholder instead of the class. decode_wan_caption swallowed OSError and ValueError silently, so a shard whose caption paths were unreadable degraded to filenames with nothing logged. The fallback is unchanged and now says why at debug level. Co-authored-by: Cursor <cursoragent@cursor.com>

Copilot

🔵 Needs a closer look

The review found correctness issues in WAN validation loss reporting/cloning, a potentially incompatible attention-core call signature, and unreachable subfolder validation logic in new encoders that should be fixed before approval.

Review details

Suppressed comments (5)

Previously missed (5) — in code that hasn't changed since the last review.

primus/backends/megatron/core/models/diffusion/wan/attention.py:130

WanAttentionBase._core calls self.core_attention(...) with positional arguments on the non-packed path. Other diffusion attention implementations in this repo call core_attention with keyword arguments (e.g., attn_mask_type=...), which is safer because some core attention implementations use keyword-only parameters. Using positional args here risks a runtime failure or silently mis-bound parameters when swapping attention backends.
primus/backends/megatron/training/diffusion/loss_computation.py:153
compute_weighted_flow_matching_loss claims that unit weights recover the unweighted objective exactly, but the masked path uses loss_mask.sum().clamp(min=1.0) whereas compute_flow_matching_loss divides by loss_mask.sum() without clamping. If a caller ever passes an all-zero mask, the weighted loss will return 0 while the unweighted loss will produce NaN/Inf, violating the stated equivalence and potentially hiding data/packing bugs.
primus/backends/megatron/wan_pretrain_trainer.py:231
In the validation path, the returned metric pair should follow the same contract as DiffusionPretrainTrainer: the reported value should be a SUM (with a sample_count weight) and it should be cloned to avoid Megatron’s in-place rescaling corrupting the logged value. Returning the mean loss as both the function loss and the metric value will skew aggregated validation loss, and missing .clone() can make the logged loss depend on microbatching.
primus/backends/megatron/data/diffusion/encoders/text/umt5.py:130
This subfolder validation condition will never fire because EncoderConfig always defines subfolder, so not hasattr(config, "subfolder") is always false. If you want to require explicit subfolder configuration only when config=None was passed, track that explicitly so the error is reachable and actionable.
primus/backends/megatron/data/diffusion/encoders/video/vae/wan/autoencoder_kl_wan.py:129
The subfolder validation guard is effectively dead: EncoderConfig always has a subfolder attribute, so not hasattr(config, "subfolder") is never true. That means the intended early, user-friendly error will never trigger, and missing subfolder will fail later in a less actionable way. If the goal is to require callers to be explicit when they pass config=None, track that explicitly.

Files reviewed: 45/45 changed files
Comments generated: 0 new
Review effort level: Lite

jianhan-amd and others added 11 commits September 8, 2026 11:01

jianhan-amd requested review from jasainio and a lite review from Copilot September 10, 2026 13:54

Copilot started reviewing on behalf of jianhan-amd September 10, 2026 13:54 View session

github-code-quality Bot found potential problems Sep 10, 2026

View reviewed changes

Copilot AI reviewed Sep 10, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings September 10, 2026 14:28

Copilot started reviewing on behalf of jianhan-amd September 10, 2026 14:29 View session

Copilot AI reviewed Sep 10, 2026

View reviewed changes

jianhan-amd requested review from amital-amd, eshaw2, gphuang, luiza-amd, mvstrauss and olehtika September 14, 2026 08:17

Navigation Menu

Feat/diffusion/wan21 megatron training - #1132

Feat/diffusion/wan21 megatron training#1132
jianhan-amd wants to merge 12 commits into
mainfrom
feat/diffusion/wan21-megatron-training

jianhan-amd commented Sep 10, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

jianhan-amd Sep 10, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

jianhan-amd commented Sep 10, 2026

WAN 2.1 / 2.2 video DiT training on Megatron-Core

What this adds

Layout

Points worth a reviewer's attention

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

jianhan-amd Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Suppressed comments (5)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants