Add diffusion backend & Wan training support

[zirui]

Summary

This PR integrates diffusion training support into Primus training framework, including WAN2.1 and WAN2.2 models.

The goal is to enable stable single-node and multi-node training for WAN diffusion models with minimal framework changes and clean extensibility.

---

Motivation

WAN diffusion models are now a primary workload for video generation training. This integration enables:

• Unified training interface for diffusion models
• Scalable single-node → multi-node training
• Reuse of existing Primus training infra (optimizer, checkpointing, logging)

---

Scope

Included

• WAN2.1 / WAN2.2 model wrapper integration
• Diffusion training loop adaptation
• Basic dataset pipeline support
• Single-node SFT training validation
• Checkpoint save/load compatibility

Not included (future work)

• Inference pipeline
• Advanced scheduling strategies
• Flash attention / kernel-level optimizations

---

Current Status

Completed

• WAN model integration (WAN2.1-1.3B / WAN2.2-5B)
• Training loop adapted for diffusion objective
• Single-node SFT training verified

In Progress

• Multi-node training validation (2-node cluster test)
• Config refactor (clean separation of model/trainer/data)
• Dataset pipeline cleanup
• Documentation and example scripts

---

Testing

Single-node

• Wan2.1-1.3B SFT training: ✅
• Wan2.2-5B SFT training: ✅

Multi-node

• 2-node test: in progress

Known Issues / Risks

• Dataset preprocessing still partially ad-hoc
• Config system needs refactor for diffusion-specific parameters
• Multi-node stability not fully validated yet

Next Steps

1. Complete 2-node validation
2. Add training performance benchmarks
3. Refactor config structure for diffusion training
4. Add reproducible example scripts

---

Notes

This is a WIP draft PR. Frequent commits will be pushed as development continues.

benchmarks

Primus Wan2.2 TI2V 5B Benchmark - 2026-07-06 10:28:13 UTC

• Summary uses median step time after skipping first 5 logged steps.
• batch_mode=local_accum means effective per-GPU batch size via gradient accumulation with micro batch 1.
• Stability/efficiency fixes enabled: video_backend=decord, dataloader_num_workers=0, PRIMUS_CACHE_RAW_VIDEO_FRAMES=1, and PRIMUS_CACHE_PROCESSED_SAMPLES=1.

engine	model	data	resolution	frames	batch_size	batch_mode	gpus	steps	gpu_mem_GB	step_time_s	step/s	TPS(samples/s/gpu)	status
primus-fsdp2-flash_attn_aiter	wan2.2-ti2v-5b	tiny-video-sample	480P	121	1	local_accum	8	95	27.080	1.760000	0.568182	0.569300	ok
primus-fsdp2-flash_attn_aiter	wan2.2-ti2v-5b	tiny-video-sample	480P	121	8	local_accum	8	95	28.330	14.330000	0.069784	0.558300	ok
primus-fsdp2-flash_attn_aiter	wan2.2-ti2v-5b	tiny-video-sample	480P	121	16	local_accum	8	95	28.330	28.430000	0.035174	0.562800	ok

[Copilot]

Pull request overview

This PR introduces an in-tree Primus diffusion backend to enable training WAN2.1/WAN2.2 diffusion video models using the existing Primus lifecycle (configs, launch, logging, checkpointing), with a PyTorch FSDP2 + optional Ulysses sequence-parallel trainer stack.

Changes:

• Added a new primus.backends.diffusion backend (adapter, registry, trainer, distributed/checkpoint utilities, attention backends).
• Added WAN model components + training pipeline (DiT, VAE, UMT5 encoder, flow-matching scheduler).
• Added runner hooks, example configs, and backend documentation to support pretrain/posttrain flows.

Reviewed changes

Copilot reviewed 54 out of 54 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt	Extra Python deps for diffusion training hook installation.
runner/helpers/hooks/train/pretrain/diffusion/prepare.py	Pre-flight validation script for diffusion module configs/assets.
runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh	Hook script to install diffusion requirements with pip caching.
runner/helpers/hooks/train/posttrain/diffusion/prepare.py	Posttrain prepare wrapper invoking the pretrain validator with post_trainer.
primus/configs/modules/diffusion/pre_trainer.yaml	Diffusion pretrain module preset (stage wiring).
primus/configs/modules/diffusion/post_trainer.yaml	Diffusion posttrain module preset (stage/trainable wiring).
primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml	WAN2.2 TI2V model preset (paths + encoder assets).
primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml	WAN2.2 SFT preset overriding init checkpoint convention.
primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml	WAN2.1 T2V model preset (paths + encoder assets).
primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml	WAN2.1 SFT preset overriding init checkpoint convention.
primus/backends/diffusion/utils/vision_process.py	Vendored vision/video processing utilities (qwen-vl-utils derived).
primus/backends/diffusion/utils/train_utils.py	Training helper utilities (dtype/seed/state-dict utilities).
primus/backends/diffusion/utils/log.py	Loguru proxy to ensure Primus-bound extra fields are present.
primus/backends/diffusion/utils/data_utils.py	Small video frame-count helpers.
primus/backends/diffusion/utils/init.py	Diffusion utils package init/documentation.
primus/backends/diffusion/trainers/fsdp2.py	WAN FSDP2 trainer implementation (wrapping, grad sync, checkpointing).
primus/backends/diffusion/trainers/base.py	Shared training loop + optimizer/scheduler plumbing for WAN trainers.
primus/backends/diffusion/trainers/init.py	Trainer builder exports.
primus/backends/diffusion/schedulers/flow_match.py	Flow-matching scheduler used by WAN training.
primus/backends/diffusion/registry.py	Explicit diffusion factory registry (model/dataset/trainer builders).
primus/backends/diffusion/README.md	Backend-level documentation (data, checkpoints, launch, configs).
primus/backends/diffusion/optim/adamw_fp32_state.py	AdamW variant maintaining FP32 optimizer state for bf16/fp16 params.
primus/backends/diffusion/models/wan/wan_dit.py	Vendored WAN DiT backbone with SP-aware RoPE/attention integration.
primus/backends/diffusion/models/wan/vae2_1.py	Vendored WAN2.1 VAE wrapper with minor Primus compatibility helpers.
primus/backends/diffusion/models/wan/train_pipeline.py	Training pipeline computing flow-match loss over VAE latents.
primus/backends/diffusion/models/wan/t5.py	UMT5-XXL encoder-only implementation + meta-init checkpoint loader.
primus/backends/diffusion/models/wan/configuration_wanvideo.py	WAN config container used to drive component construction.
primus/backends/diffusion/models/wan/components.py	Container dataclass for WAN components (dit/vae/text encoder).
primus/backends/diffusion/models/wan/attention_backend.py	Vendored attention backend shim with CPU fallback behavior.
primus/backends/diffusion/models/wan/adapter.py	GenAIModel adapter exposing forward_train/freeze/to() semantics.
primus/backends/diffusion/models/wan/init.py	WAN model package exports.
primus/backends/diffusion/models/registrations/wan.py	WAN model builder/loader (config, DiT/VAE/T5 wiring, checkpoint load).
primus/backends/diffusion/models/registrations/init.py	Model registrations package init.
primus/backends/diffusion/models/interface.py	GenAIModel interface definition for train/inference entrypoints.
primus/backends/diffusion/models/init.py	Diffusion model exports.
primus/backends/diffusion/distributed/ulysses.py	Ulysses sequence-parallel collectives + autograd-safe split/gather.
primus/backends/diffusion/distributed/mesh.py	DeviceMesh creation + distributed init utilities for DP/SP/HSDP shapes.
primus/backends/diffusion/distributed/checkpoint.py	DTCP save/load utilities for sharded checkpointing.
primus/backends/diffusion/distributed/init.py	Exports for diffusion distributed utilities.
primus/backends/diffusion/diffusion_pretrain_trainer.py	Primus BaseTrainer wrapper that builds model/dataset/trainer and runs training.
primus/backends/diffusion/diffusion_adapter.py	Primus backend adapter converting module params into WAN args + trainer selection.
primus/backends/diffusion/attention/flex.py	FlexAttention backend integration (torch.compile-based).
primus/backends/diffusion/attention/attention.py	Unified attention API selecting SDPA/Flash/Flex/AITER backends.
primus/backends/diffusion/attention/aiter.py	AITER flash-attn backend integration.
primus/backends/diffusion/attention/_flash_common.py	Shared flash-attn varlen/fixed-length glue.
primus/backends/diffusion/attention/init.py	Attention package exports.
primus/backends/diffusion/argument_builder.py	Primus-style override normalization into compact WAN runtime args.
primus/backends/diffusion/init.py	Registers the diffusion adapter in BackendRegistry.
examples/diffusion/README.md	Example-level usage documentation for WAN diffusion configs.
examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml	MI355X WAN2.2 pretrain example config using Primus-style overrides.
examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml	MI355X WAN2.2 posttrain example config using Primus-style overrides.
examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml	MI355X WAN2.1 pretrain example config using Primus-style overrides.
examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml	MI355X WAN2.1 posttrain example config using Primus-style overrides.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review is ineligible. To be eligible to request a review, you need a paid Copilot license, or your organization must enable Copilot code review.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review is ineligible. To be eligible to request a review, you need a paid Copilot license, or your organization must enable Copilot code review.

[Copilot]

Pull request overview

Copilot reviewed 54 out of 54 changed files in this pull request and generated 8 comments.

[Copilot]

Pull request overview

Copilot reviewed 54 out of 54 changed files in this pull request and generated 7 comments.

[Copilot]

Pull request overview

Copilot reviewed 57 out of 57 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 57 out of 57 changed files in this pull request and generated 7 comments.

[Copilot]

Pull request overview

Copilot reviewed 57 out of 57 changed files in this pull request and generated 4 comments.

[Copilot]

Pull request overview

Copilot reviewed 59 out of 59 changed files in this pull request and generated 7 comments.