From f3791f83dbde79da7ed1b61675eff64b8f489044 Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 17 Jun 2026 06:54:14 +0000 Subject: [PATCH 01/30] first commit for diffusion training Co-authored-by: Cursor --- examples/diffusion/README.md | 66 ++ .../MI355X/wan2.1_t2v_1.3b-posttrain.yaml | 53 ++ .../MI355X/wan2.1_t2v_1.3b-pretrain.yaml | 53 ++ .../MI355X/wan2.2_ti2v_5b-posttrain.yaml | 53 ++ .../MI355X/wan2.2_ti2v_5b-pretrain.yaml | 53 ++ primus/backends/diffusion/README.md | 245 +++++ primus/backends/diffusion/__init__.py | 12 + primus/backends/diffusion/argument_builder.py | 252 ++++++ .../backends/diffusion/attention/__init__.py | 29 + .../diffusion/attention/_flash_common.py | 141 +++ primus/backends/diffusion/attention/aiter.py | 71 ++ .../backends/diffusion/attention/attention.py | 410 +++++++++ primus/backends/diffusion/attention/flex.py | 280 ++++++ .../backends/diffusion/diffusion_adapter.py | 65 ++ .../diffusion/diffusion_pretrain_trainer.py | 125 +++ .../diffusion/distributed/__init__.py | 32 + .../diffusion/distributed/checkpoint.py | 101 +++ primus/backends/diffusion/distributed/mesh.py | 126 +++ .../backends/diffusion/distributed/ulysses.py | 200 ++++ primus/backends/diffusion/models/__init__.py | 13 + primus/backends/diffusion/models/interface.py | 48 + .../models/registrations/__init__.py | 7 + .../diffusion/models/registrations/wan.py | 234 +++++ .../backends/diffusion/models/wan/__init__.py | 20 + .../backends/diffusion/models/wan/adapter.py | 154 ++++ .../diffusion/models/wan/attention_backend.py | 228 +++++ .../diffusion/models/wan/components.py | 26 + .../models/wan/configuration_wanvideo.py | 84 ++ primus/backends/diffusion/models/wan/t5.py | 299 ++++++ .../diffusion/models/wan/train_pipeline.py | 282 ++++++ .../backends/diffusion/models/wan/vae2_1.py | 638 +++++++++++++ .../backends/diffusion/models/wan/vae2_2.py | 855 ++++++++++++++++++ .../backends/diffusion/models/wan/wan_dit.py | 555 ++++++++++++ .../diffusion/optim/adamw_fp32_state.py | 111 +++ primus/backends/diffusion/registry.py | 66 ++ .../diffusion/schedulers/flow_match.py | 130 +++ .../backends/diffusion/trainers/__init__.py | 11 + primus/backends/diffusion/trainers/base.py | 539 +++++++++++ primus/backends/diffusion/trainers/fsdp2.py | 395 ++++++++ primus/backends/diffusion/utils/__init__.py | 10 + primus/backends/diffusion/utils/data_utils.py | 61 ++ primus/backends/diffusion/utils/log.py | 42 + .../backends/diffusion/utils/train_utils.py | 200 ++++ .../diffusion/utils/vision_process.py | 569 ++++++++++++ .../models/diffusion/wan2.1_t2v_1.3b.yaml | 13 + .../models/diffusion/wan2.1_t2v_1.3b_sft.yaml | 8 + .../models/diffusion/wan2.2_ti2v_5b.yaml | 13 + .../models/diffusion/wan2.2_ti2v_5b_sft.yaml | 8 + .../modules/diffusion/post_trainer.yaml | 5 + .../modules/diffusion/pre_trainer.yaml | 4 + .../train/posttrain/diffusion/prepare.py | 24 + .../diffusion/00_install_requirements.sh | 40 + .../hooks/train/pretrain/diffusion/prepare.py | 132 +++ .../diffusion/requirements-diffusion.txt | 11 + 54 files changed, 8202 insertions(+) create mode 100644 examples/diffusion/README.md create mode 100644 examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml create mode 100644 examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml create mode 100644 examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml create mode 100644 examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml create mode 100644 primus/backends/diffusion/README.md create mode 100644 primus/backends/diffusion/__init__.py create mode 100644 primus/backends/diffusion/argument_builder.py create mode 100644 primus/backends/diffusion/attention/__init__.py create mode 100644 primus/backends/diffusion/attention/_flash_common.py create mode 100644 primus/backends/diffusion/attention/aiter.py create mode 100644 primus/backends/diffusion/attention/attention.py create mode 100644 primus/backends/diffusion/attention/flex.py create mode 100644 primus/backends/diffusion/diffusion_adapter.py create mode 100644 primus/backends/diffusion/diffusion_pretrain_trainer.py create mode 100644 primus/backends/diffusion/distributed/__init__.py create mode 100644 primus/backends/diffusion/distributed/checkpoint.py create mode 100644 primus/backends/diffusion/distributed/mesh.py create mode 100644 primus/backends/diffusion/distributed/ulysses.py create mode 100644 primus/backends/diffusion/models/__init__.py create mode 100644 primus/backends/diffusion/models/interface.py create mode 100644 primus/backends/diffusion/models/registrations/__init__.py create mode 100644 primus/backends/diffusion/models/registrations/wan.py create mode 100644 primus/backends/diffusion/models/wan/__init__.py create mode 100644 primus/backends/diffusion/models/wan/adapter.py create mode 100644 primus/backends/diffusion/models/wan/attention_backend.py create mode 100644 primus/backends/diffusion/models/wan/components.py create mode 100644 primus/backends/diffusion/models/wan/configuration_wanvideo.py create mode 100644 primus/backends/diffusion/models/wan/t5.py create mode 100644 primus/backends/diffusion/models/wan/train_pipeline.py create mode 100644 primus/backends/diffusion/models/wan/vae2_1.py create mode 100644 primus/backends/diffusion/models/wan/vae2_2.py create mode 100644 primus/backends/diffusion/models/wan/wan_dit.py create mode 100644 primus/backends/diffusion/optim/adamw_fp32_state.py create mode 100644 primus/backends/diffusion/registry.py create mode 100644 primus/backends/diffusion/schedulers/flow_match.py create mode 100644 primus/backends/diffusion/trainers/__init__.py create mode 100644 primus/backends/diffusion/trainers/base.py create mode 100644 primus/backends/diffusion/trainers/fsdp2.py create mode 100644 primus/backends/diffusion/utils/__init__.py create mode 100644 primus/backends/diffusion/utils/data_utils.py create mode 100644 primus/backends/diffusion/utils/log.py create mode 100644 primus/backends/diffusion/utils/train_utils.py create mode 100644 primus/backends/diffusion/utils/vision_process.py create mode 100644 primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml create mode 100644 primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml create mode 100644 primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml create mode 100644 primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml create mode 100644 primus/configs/modules/diffusion/post_trainer.yaml create mode 100644 primus/configs/modules/diffusion/pre_trainer.yaml create mode 100644 runner/helpers/hooks/train/posttrain/diffusion/prepare.py create mode 100755 runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh create mode 100644 runner/helpers/hooks/train/pretrain/diffusion/prepare.py create mode 100644 runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md new file mode 100644 index 000000000..ba4fb200a --- /dev/null +++ b/examples/diffusion/README.md @@ -0,0 +1,66 @@ +# Wan Examples + +Wan examples exercise the independent PyTorch Diffusion backend under +`primus/backends/diffusion`. For backend details, data/checkpoint layout, and minimal +configs, see `primus/backends/diffusion/README.md`. + +## Data + +The default smoke-test dataset is `zirui3/tiny-video-samples` on Hugging Face: + +```bash +huggingface-cli download zirui3/tiny-video-samples \ + --repo-type dataset \ + --local-dir /data/tiny-video-samples +``` + +Expected layout: + +```text +/data/tiny-video-samples/ + meta.jsonl + data/*.mp4 +``` + +## Run + +Launch an MI355X example config with 8 GPUs: + +```bash +torchrun --standalone --nproc_per_node=8 \ + -m primus.cli.main train pretrain \ + --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml +``` + +Override common paths and runtime options with environment variables: + +```bash +DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ +DATA_FOLDER=/data/tiny-video-samples/data \ +ATTENTION_BACKEND=sdpa \ +SP_SIZE=1 \ +MAX_STEPS=3 \ +torchrun --standalone --nproc_per_node=8 \ + -m primus.cli.main train pretrain \ + --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml +``` + +Use `SP_SIZE=4` or `SP_SIZE=8` to enable Ulysses sequence parallelism +when the model head count supports it. + +Post-train examples use the same public override shape under `modules.post_trainer`: + +```bash +INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B \ +DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ +DATA_FOLDER=/data/tiny-video-samples/data \ +MAX_STEPS=3 \ +torchrun --standalone --nproc_per_node=8 \ + -m primus.cli.main train posttrain \ + --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml +``` + +The MI355X configs use Primus-style override sections such as `training`, +`data`, `parallelism`, `optimizer`, `runtime`, and `metrics`. The diffusion +adapter normalizes those sections into the Wan model/dataset/trainer +arguments at runtime. diff --git a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml new file mode 100644 index 000000000..dcb8c2482 --- /dev/null +++ b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml @@ -0,0 +1,53 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:wan2.1_t2v_1.3b-posttrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + post_trainer: + framework: diffusion + config: post_trainer.yaml + + # Model preset to fine-tune from INIT_CHECKPOINT. + model: wan2.1_t2v_1.3b_sft.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: 1 + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/wan2.1_t2v_1.3b-posttrain} + save_steps: 0 + run_name: wan2.1_t2v_1.3b-posttrain + + data: + dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} + data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} + frame_num: 81 + video_backend: imageio + text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.1-T2V-1.3B/google/umt5-xxl} + height: 480 + width: 832 + + parallelism: + sp_size: ${SP_SIZE:1} + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: 5.0e-6 + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:sdpa} + report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml new file mode 100644 index 000000000..1411c054a --- /dev/null +++ b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml @@ -0,0 +1,53 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:wan2.1_t2v_1.3b-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + + # Model preset to run. + model: wan2.1_t2v_1.3b.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: 1 + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/wan2.1_t2v_1.3b-pretrain} + save_steps: 0 + run_name: wan2.1_t2v_1.3b-pretrain + + data: + dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} + data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} + frame_num: 81 + video_backend: imageio + text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.1-T2V-1.3B/google/umt5-xxl} + height: 480 + width: 832 + + parallelism: + sp_size: ${SP_SIZE:1} + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: 1.0e-5 + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:sdpa} + report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml new file mode 100644 index 000000000..a12fba375 --- /dev/null +++ b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml @@ -0,0 +1,53 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:wan2.2_ti2v_5b-posttrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + post_trainer: + framework: diffusion + config: post_trainer.yaml + + # Model preset to fine-tune from INIT_CHECKPOINT. + model: wan2.2_ti2v_5b_sft.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: 1 + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/wan2.2_ti2v_5b-posttrain} + save_steps: 0 + run_name: wan2.2_ti2v_5b-posttrain + + data: + dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} + data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} + frame_num: 81 + video_backend: imageio + text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.2-TI2V-5B/google/umt5-xxl} + height: 480 + width: 832 + + parallelism: + sp_size: ${SP_SIZE:1} + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: 5.0e-6 + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:sdpa} + report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml new file mode 100644 index 000000000..6757111be --- /dev/null +++ b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml @@ -0,0 +1,53 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:wan2.2_ti2v_5b-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + + # Model preset to run. + model: wan2.2_ti2v_5b.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: 1 + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/wan2.2_ti2v_5b-pretrain} + save_steps: 0 + run_name: wan2.2_ti2v_5b-pretrain + + data: + dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} + data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} + frame_num: 81 + video_backend: imageio + text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.2-TI2V-5B/google/umt5-xxl} + height: 480 + width: 832 + + parallelism: + sp_size: ${SP_SIZE:1} + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: 1.0e-5 + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:sdpa} + report_to: none diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md new file mode 100644 index 000000000..896e3bb38 --- /dev/null +++ b/primus/backends/diffusion/README.md @@ -0,0 +1,245 @@ +# Primus Diffusion Backend + +`diffusion` integrates PyTorch diffusion-model training as an independent Primus +backend. Primus provides the config/launch entrypoint, while model, dataset, +attention, and FSDP2 training logic are owned by the in-tree Wan implementation +under `primus/backends/diffusion`. + +All runtime code resolves through the Primus namespace, for example +`primus.backends.diffusion.models` and `primus.backends.diffusion.trainers`, +so Wan training is a first-class part of the Primus `diffusion` backend. + +Supported scope: + +- Model implementation: `wan` for Wan2.1 and Wan2.2. +- Trainer: FSDP2 only. +- Sequence parallelism: Ulysses SP via `trainer.args.sp_size`. + +Wan-specific dependencies are kept out of top-level Primus requirements. See +`runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt`. + +## Data + +Wan training reads a jsonl metadata file and a video folder: + +```jsonl +{"prompt": "text prompt", "video": "example.mp4"} +``` + +The default small dataset for smoke tests is the Hugging Face dataset +`zirui3/tiny-video-samples`: + +```bash +huggingface-cli download zirui3/tiny-video-samples \ + --repo-type dataset \ + --local-dir /data/tiny-video-samples +``` + +This produces: + +```text +/data/tiny-video-samples/ + meta.jsonl + data/*.mp4 +``` + +Use these public config fields to point training at another dataset: + +```yaml +data: + dataset_path: /path/to/meta.jsonl + data_folder: /path/to/videos + video_backend: imageio +``` + +## Checkpoints + +Each Wan model is a single Hugging Face repo that already bundles everything Wan +training needs: the DiT weights, the UMT5-XXL text encoder +(`models_t5_umt5-xxl-enc-bf16.pth`), the VAE (`Wan2.1_VAE.pth` / +`Wan2.2_VAE.pth`), and the tokenizer under `google/umt5-xxl`. Download the +model(s) you plan to train: + +| Model | Preset | Hugging Face repo | +| --- | --- | --- | +| Wan2.1-T2V-1.3B | `wan2.1_t2v_1.3b.yaml` | [Wan-AI/Wan2.1-T2V-1.3B](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B) | +| Wan2.2-TI2V-5B | `wan2.2_ti2v_5b.yaml` | [Wan-AI/Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) | +| Wan2.1-T2V-14B | (no shipped preset) | [Wan-AI/Wan2.1-T2V-14B](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B) | + +```bash +# Wan2.1-T2V-1.3B +huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B \ + --local-dir /models/Wan2.1-T2V-1.3B + +# Wan2.2-TI2V-5B +huggingface-cli download Wan-AI/Wan2.2-TI2V-5B \ + --local-dir /models/Wan2.2-TI2V-5B +``` + +A downloaded Wan repo looks like this (the T5 encoder, VAE, and tokenizer are +shipped inside the same repo, so no separate download is required): + +```text +/models/Wan2.1-T2V-1.3B/ + config.json + diffusion_pytorch_model*.safetensors # DiT weights + models_t5_umt5-xxl-enc-bf16.pth # T5 (UMT5-XXL) text encoder + Wan2.1_VAE.pth # VAE (Wan2.2_VAE.pth in the 5B repo) + google/umt5-xxl/ # tokenizer +``` + +The model presets reference these paths by default. Override any asset with +environment variables: + +```bash +export PRETRAINED_PATH=/models/Wan2.1-T2V-1.3B # pretrain DiT init +export INIT_CHECKPOINT=/models/Wan2.1-T2V-1.3B # post-train/SFT DiT init +export TEXT_TOKENIZER=/models/Wan2.1-T2V-1.3B/google/umt5-xxl +export TEXT_ENCODER=/models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth +export VAE_CHECKPOINT=/models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth +``` + +Pretrain presets use `PRETRAINED_PATH`; post-train SFT presets use +`INIT_CHECKPOINT` for DiT initialization. For Wan2.2-5B use the matching +`/models/Wan2.2-TI2V-5B` directory and `Wan2.2_VAE.pth`. + +## Pre-flight validation + +The diffusion prepare hook validates the prepared assets; it does **not** +download them. Download the dataset and checkpoints above first, then run the +hook to confirm the configured dataset, tokenizer, and DiT/T5/VAE paths exist +before launching distributed training: + +```bash +# pretrain config (validates modules.pre_trainer) +python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ + --config examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml + +# post-train config (validates modules.post_trainer) +python3 runner/helpers/hooks/train/posttrain/diffusion/prepare.py \ + --config examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml +``` + +On success it prints `env.PREPARED=1`. Set `SKIP_PREPARE=1` to bypass the check +for debugging. + +## Launch + +Training runs through the Primus CLI (`primus.cli.main train `) +under `torchrun`. + +### Single node (e.g. 8 GPUs) + +```bash +torchrun --standalone --nproc_per_node=8 \ + -m primus.cli.main train pretrain --config /path/to/wan_config.yaml +``` + +Post-train examples use `modules.post_trainer` and the posttrain suite: + +```bash +torchrun --standalone --nproc_per_node=8 \ + -m primus.cli.main train posttrain --config /path/to/wan_posttrain_config.yaml +``` + +### Multi-node + +Run the same command on every node, sharing one rendezvous endpoint +(`MASTER_ADDR`/`MASTER_PORT`) and giving each node a distinct `--node_rank`: + +```bash +# on every node (NNODES total); MASTER_ADDR must be a routable IP of node rank 0 +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port=29577 \ + --nproc_per_node=8 \ + -m primus.cli.main train pretrain --config /path/to/wan_config.yaml +``` + +The config and CLI are identical to single node; only the `torchrun` rendezvous +flags change. Total world size is `NNODES * nproc_per_node`. + +Useful runtime knobs: + +- `trainer.args.attention_backend`: `sdpa` is the most portable ROCm option. +- `trainer.args.sp_size`: Ulysses sequence parallel size. It must divide the + model attention head count; for example Wan2.1-1.3B supports `sp_size=4` but + not `sp_size=8`. +- `trainer.args.dp_replicate`: data parallel replication size. +- `FIXED_TIMESTEP` and `FIXED_SEED`: optional debug variables for reproducible + loss-alignment checks. + +On ROCm clusters, point compiler/cache directories to a large filesystem: + +```bash +export TMPDIR=/path/to/large/tmp +export TRITON_CACHE_DIR=/path/to/large/cache/triton +export TORCHINDUCTOR_CACHE_DIR=/path/to/large/cache/inductor +export AMD_COMGR_CACHE_DIR=/path/to/large/cache/comgr +``` + +## Primus-Style Configs + +New examples should use Primus-style override sections and let the diffusion +adapter normalize them into Wan args. See: + +```text +examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml +examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml +examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml +examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml +``` + +Minimal Wan2.1-1.3B shape: + +```yaml +work_group: local +user_name: local +exp_name: wan2.1_t2v_1.3b-pretrain +workspace: ./output + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + model: wan2.1_t2v_1.3b.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + training: + local_batch_size: 1 + steps: 100 + gradient_accumulation_steps: 1 + output_dir: ./output/wan2.1_t2v_1.3b-pretrain + save_steps: 0 + data: + dataset_path: /data/tiny-video-samples/meta.jsonl + data_folder: /data/tiny-video-samples/data + text_tokenizer: /models/Wan2.1-T2V-1.3B/google/umt5-xxl + height: 480 + width: 832 + parallelism: + sp_size: 1 + dp_replicate: 1 + runtime: + attention_backend: sdpa + report_to: none +``` + +Minimal Wan2.2-5B uses the same shape with `model: wan2.2_ti2v_5b.yaml` and +the Wan2.2 tokenizer path: + +```yaml +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + model: wan2.2_ti2v_5b.yaml + overrides: + data: + text_tokenizer: /models/Wan2.2-TI2V-5B/google/umt5-xxl +``` diff --git a/primus/backends/diffusion/__init__.py b/primus/backends/diffusion/__init__.py new file mode 100644 index 000000000..c47954e23 --- /dev/null +++ b/primus/backends/diffusion/__init__.py @@ -0,0 +1,12 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Primus-owned diffusion backend components for t2i/t2v training.""" + +from primus.backends.diffusion.diffusion_adapter import DiffusionAdapter +from primus.core.backend.backend_registry import BackendRegistry + +BackendRegistry.register_adapter("diffusion", DiffusionAdapter) diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py new file mode 100644 index 000000000..2e1c6e9d2 --- /dev/null +++ b/primus/backends/diffusion/argument_builder.py @@ -0,0 +1,252 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import copy +from types import SimpleNamespace +from typing import Any + +from primus.core.utils.yaml_utils import nested_namespace_to_dict + + +class WanArgBuilder: + """Build the compact config object consumed by the Wan trainer.""" + + DEFAULT_DATASET: dict[str, Any] = { + "name": "wan", + "config": { + "dataset_type": "vision", + "dataset_format": "jsonl", + "dataset_path": "/path/to/meta.jsonl", + "data_folder": "/path/to/videos", + "video_sampling_strategy": "frame_num", + "frame_num": 81, + "shuffle": True, + "video_backend": "imageio", + "processor_config": { + "processor_name": "wanvideo", + "processor_type": "wanvideo", + "max_text_length": 512, + "text_tokenizer": "/path/to/umt5-xxl", + "extra_kwargs": { + "do_resize": True, + "size": { + "height": 480, + "width": 832, + }, + "do_normalize": True, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], + }, + }, + }, + } + DEFAULT_TRAINER: dict[str, Any] = { + "name": "fsdp2", + "args": { + "output_dir": "./output/wan", + "per_device_train_batch_size": 1, + "per_device_eval_batch_size": 1, + "gradient_accumulation_steps": 1, + "gradient_checkpointing": True, + "attention_backend": "sdpa", + "learning_rate": 1.0e-5, + "lr_scheduler_type": "constant", + "warmup_steps": 0, + "weight_decay": 0.01, + "num_train_epochs": 1, + "max_steps": 100, + "logging_steps": 1, + "save_steps": 0, + "dataloader_num_workers": 4, + "report_to": "none", + "run_name": "wan-fsdp2", + "bf16": True, + "seed": 10007, + "optim": "adamw_torch", + "adam_beta1": 0.9, + "adam_beta2": 0.999, + "adam_epsilon": 1.0e-8, + "max_grad_norm": 1.0, + "fsdp2_wrap_target": "dit", + "fsdp_transformer_layer_cls_to_wrap": "DiTBlock", + "fsdp2_reshard_after_forward": True, + "save_strategy": "dit_only", + "sp_size": 1, + "dp_replicate": 1, + "flow_match_scheduler": { + "shift": 5, + "sigma_min": 0.0, + "extra_one_step": True, + "num_train_timesteps": 1000, + }, + }, + } + + def __init__(self) -> None: + self._params: dict[str, Any] = {} + + def update(self, params: Any) -> None: + if isinstance(params, SimpleNamespace): + self._params = nested_namespace_to_dict(params) + elif isinstance(params, dict): + self._params = copy.deepcopy(params) + else: + raise TypeError(f"WanArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}") + + def finalize(self) -> SimpleNamespace: + params = copy.deepcopy(self._params) + if "model" not in params: + raise ValueError("Wan backend config requires a model preset.") + for legacy_section in ("dataset", "trainer"): + if legacy_section in params: + raise ValueError( + f"Wan backend no longer accepts public `{legacy_section}` overrides. " + "Use Primus-style `data`, `training`, `parallelism`, `optimizer`, " + "`runtime`, and `metrics` sections instead." + ) + + params = self._normalize_primus_style_sections(params) + + return SimpleNamespace( + model=params["model"], + dataset=params["dataset"], + trainer=params["trainer"], + stage=params.get("stage", "pretrain"), + primus=params.get("primus", {}), + ) + + @staticmethod + def _set_nested(target: dict[str, Any], path: tuple[str, ...], value: Any) -> None: + cursor = target + for key in path[:-1]: + next_value = cursor.get(key) + if not isinstance(next_value, dict): + next_value = {} + cursor[key] = next_value + cursor = next_value + cursor[path[-1]] = value + + @staticmethod + def _get_any(source: dict[str, Any], *keys: str) -> Any: + for key in keys: + if key in source: + return source[key] + return None + + def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, Any]: + """Translate concise Primus-style Wan sections into trainer arguments. + + Public configs use high-level sections such as `training`, `data`, + `parallelism`, and `runtime`; internal defaults supply the compact + `dataset` and `trainer` objects consumed by the Wan runtime. + """ + + normalized = { + "model": params["model"], + "dataset": copy.deepcopy(self.DEFAULT_DATASET), + "trainer": copy.deepcopy(self.DEFAULT_TRAINER), + "stage": params.get("stage", "pretrain"), + "primus": params.get("primus", {}), + } + + dataset_cfg = normalized["dataset"]["config"] + trainer_args = normalized["trainer"]["args"] + + training = params.get("training") or {} + data = params.get("data") or {} + parallelism = params.get("parallelism") or {} + optimizer = params.get("optimizer") or {} + runtime = params.get("runtime") or {} + metrics = params.get("metrics") or {} + + training_map = { + ("steps",): ("max_steps",), + ("local_batch_size",): ("per_device_train_batch_size",), + ("global_batch_size",): ("global_batch_size",), + ("gradient_accumulation_steps",): ("gradient_accumulation_steps",), + ("output_dir",): ("output_dir",), + ("save_steps",): ("save_steps",), + ("run_name",): ("run_name",), + ("num_train_epochs",): ("num_train_epochs",), + ("resume_from_checkpoint",): ("resume_from_checkpoint",), + } + for source_path, target_path in training_map.items(): + value = self._get_any(training, *source_path) + if value is not None: + self._set_nested(trainer_args, target_path, value) + + data_map = { + ("dataset_path",): ("dataset_path",), + ("data_folder",): ("data_folder",), + ("frame_num",): ("frame_num",), + ("video_backend",): ("video_backend",), + ("text_tokenizer",): ("processor_config", "text_tokenizer"), + ("processor_name",): ("processor_config", "processor_name"), + ("processor_type",): ("processor_config", "processor_type"), + } + for source_path, target_path in data_map.items(): + value = self._get_any(data, *source_path) + if value is not None: + self._set_nested(dataset_cfg, target_path, value) + + height = data.get("height") + width = data.get("width") + if height is not None: + self._set_nested(dataset_cfg, ("processor_config", "extra_kwargs", "size", "height"), height) + if width is not None: + self._set_nested(dataset_cfg, ("processor_config", "extra_kwargs", "size", "width"), width) + + parallelism_map = { + ("sp_size",): ("sp_size",), + ("dp_replicate",): ("dp_replicate",), + } + for source_path, target_path in parallelism_map.items(): + value = self._get_any(parallelism, *source_path) + if value is not None: + self._set_nested(trainer_args, target_path, value) + + optimizer_map = { + ("lr",): ("learning_rate",), + ("learning_rate",): ("learning_rate",), + ("weight_decay",): ("weight_decay",), + ("adam_beta1",): ("adam_beta1",), + ("adam_beta2",): ("adam_beta2",), + ("adam_epsilon",): ("adam_epsilon",), + ("max_grad_norm",): ("max_grad_norm",), + } + for source_path, target_path in optimizer_map.items(): + value = self._get_any(optimizer, *source_path) + if value is not None: + self._set_nested(trainer_args, target_path, value) + + runtime_map = { + ("attention_backend",): ("attention_backend",), + ("report_to",): ("report_to",), + ("seed",): ("seed",), + } + for source_path, target_path in runtime_map.items(): + value = self._get_any(runtime, *source_path) + if value is not None: + self._set_nested(trainer_args, target_path, value) + + log_freq = metrics.get("log_freq") + if log_freq is not None: + self._set_nested(trainer_args, ("logging_steps",), log_freq) + + enable_wandb = metrics.get("enable_wandb") + if enable_wandb is False and runtime.get("report_to") is None: + self._set_nested(trainer_args, ("report_to",), "none") + elif enable_wandb is True and runtime.get("report_to") is None: + self._set_nested(trainer_args, ("report_to",), "wandb") + + checkpoint = params.get("checkpoint") or {} + resume_from_checkpoint = checkpoint.get("resume_from_checkpoint") + if resume_from_checkpoint is not None: + self._set_nested(trainer_args, ("resume_from_checkpoint",), resume_from_checkpoint) + + return normalized diff --git a/primus/backends/diffusion/attention/__init__.py b/primus/backends/diffusion/attention/__init__.py new file mode 100644 index 000000000..a1fea3747 --- /dev/null +++ b/primus/backends/diffusion/attention/__init__.py @@ -0,0 +1,29 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from .attention import ( + AITER_FLASH_ATTN_AVAILABLE, + FLASH_ATTN_2_AVAILABLE, + FLASH_ATTN_3_AVAILABLE, + attention, + attention_fused, + flash_attention, + get_attention_backend, + set_attention_backend, +) +from .flex import FLEX_ATTENTION_AVAILABLE + +__all__ = [ + "AITER_FLASH_ATTN_AVAILABLE", + "FLASH_ATTN_2_AVAILABLE", + "FLASH_ATTN_3_AVAILABLE", + "FLEX_ATTENTION_AVAILABLE", + "attention", + "attention_fused", + "flash_attention", + "get_attention_backend", + "set_attention_backend", +] diff --git a/primus/backends/diffusion/attention/_flash_common.py b/primus/backends/diffusion/attention/_flash_common.py new file mode 100644 index 000000000..7535a48a1 --- /dev/null +++ b/primus/backends/diffusion/attention/_flash_common.py @@ -0,0 +1,141 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from collections.abc import Callable + +import torch + + +def _unwrap_output(x: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor: + return x[0] if isinstance(x, tuple) else x + + +def run_flash_attention_backend( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + *, + fixed_attention: Callable[..., torch.Tensor | tuple[torch.Tensor, ...]], + varlen_attention: Callable[..., torch.Tensor | tuple[torch.Tensor, ...]], + window_size_adapter: Callable[[tuple[int, ...]], tuple[int, ...]] | None = None, + fixed_extra_kwargs: dict | None = None, + varlen_extra_kwargs: dict | None = None, +): + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == "cuda" and q.size(-1) <= 256 + + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + call_window_size = window_size_adapter(window_size) if window_size_adapter is not None else window_size + fixed_extra_kwargs = fixed_extra_kwargs or {} + varlen_extra_kwargs = varlen_extra_kwargs or {} + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + if q_lens is None and k_lens is None: + qh = half(q) + kh = half(k) + vh = half(v) + qh = qh.to(vh.dtype) + kh = kh.to(vh.dtype) + if q_scale is not None: + qh = qh * q_scale + + x = fixed_attention( + qh, + kh, + vh, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=call_window_size, + deterministic=deterministic, + **fixed_extra_kwargs, + ) + return _unwrap_output(x).type(out_dtype) + + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens, strict=False)])) + + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens, strict=False)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens, strict=False)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + if q_scale is not None: + q = q * q_scale + + cu_seqlens_q = ( + torch.cat([q_lens.new_zeros([1]), q_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True) + ) + cu_seqlens_k = ( + torch.cat([k_lens.new_zeros([1]), k_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True) + ) + max_sq = int(q_lens.max()) + max_sk = int(k_lens.max()) + + import os + + import torch.distributed as dist + + if os.environ.get("OMNIFLOW_DEBUG_ATTN") == "1" and (not dist.is_initialized() or dist.get_rank() == 0): + print(f"[DEBUG ATTN] q.shape={q.shape}, k.shape={k.shape}, v.shape={v.shape}") + print(f"[DEBUG ATTN] q_lens={q_lens}, k_lens={k_lens}") + print(f"[DEBUG ATTN] cu_seqlens_q={cu_seqlens_q}, cu_seqlens_k={cu_seqlens_k}") + print(f"[DEBUG ATTN] max_seqlen_q={max_sq} (lq={lq}), max_seqlen_k={max_sk} (lk={lk})") + print(f"[DEBUG ATTN] q.dtype={q.dtype}, causal={causal}, window_size={call_window_size}") + + x = varlen_attention( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_sq, + max_seqlen_k=max_sk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=call_window_size, + deterministic=deterministic, + **varlen_extra_kwargs, + ) + # Pad output back to (b, lq) shape to match the input padded layout + out = _unwrap_output(x) + if max_sq == lq: + return out.unflatten(0, (b, lq)).type(out_dtype) + # Variable-length: need to scatter back into padded tensor + result = q.new_zeros(b, lq, *out.shape[1:]) + offset = 0 + for i in range(b): + sl = int(q_lens[i]) + result[i, :sl] = out[offset : offset + sl] + offset += sl + return result.type(out_dtype) diff --git a/primus/backends/diffusion/attention/aiter.py b/primus/backends/diffusion/attention/aiter.py new file mode 100644 index 000000000..dfe3e655a --- /dev/null +++ b/primus/backends/diffusion/attention/aiter.py @@ -0,0 +1,71 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import torch + +from ._flash_common import run_flash_attention_backend + +AITER_FLASH_ATTN_AVAILABLE = False +aiter = None + +try: + import aiter as _aiter # type: ignore + + if hasattr(_aiter, "flash_attn_func") and hasattr(_aiter, "flash_attn_varlen_func"): + aiter = _aiter + AITER_FLASH_ATTN_AVAILABLE = True +except Exception: + AITER_FLASH_ATTN_AVAILABLE = False + + +def _normalize_window_size(window_size: tuple[int, ...]) -> tuple[int, int, int]: + if len(window_size) == 2: + left, right = window_size + return (left, right, 0) + if len(window_size) == 3: + left, right, sink = window_size + return (left, right, sink) + raise ValueError(f"window_size must have 2 or 3 items, got {window_size}") + + +def aiter_flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, +): + assert AITER_FLASH_ATTN_AVAILABLE and aiter is not None + + need_lse = torch.is_grad_enabled() + return run_flash_attention_backend( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + fixed_attention=aiter.flash_attn_func, + varlen_attention=aiter.flash_attn_varlen_func, + window_size_adapter=_normalize_window_size, + fixed_extra_kwargs={"return_lse": need_lse}, + varlen_extra_kwargs={"return_lse": need_lse}, + ) diff --git a/primus/backends/diffusion/attention/attention.py b/primus/backends/diffusion/attention/attention.py new file mode 100644 index 000000000..4e67b59b9 --- /dev/null +++ b/primus/backends/diffusion/attention/attention.py @@ -0,0 +1,410 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Unified attention entry points. + +Public API: +- `attention(q, k, v, ...)` : Wan2.2-style, q/k/v shape [B, L, N, D] +- `attention_fused(q, k, v, ...)` : Legacy Wan-style, q/k/v shape [B, S, N*D] +""" + +from __future__ import annotations + +import warnings + +import torch + +from ._flash_common import run_flash_attention_backend +from .aiter import AITER_FLASH_ATTN_AVAILABLE, aiter_flash_attention + +FLASH_ATTN_3_AVAILABLE = False +FLASH_ATTN_2_AVAILABLE = False +flash_attn_interface = None +flash_attn = None + +try: + import flash_attn_interface as _flash_attn_interface # type: ignore + + flash_attn_interface = _flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except Exception: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn as _flash_attn # type: ignore + + flash_attn = _flash_attn + FLASH_ATTN_2_AVAILABLE = True +except Exception: + FLASH_ATTN_2_AVAILABLE = False + + +__all__ = [ + "AITER_FLASH_ATTN_AVAILABLE", + "FLASH_ATTN_2_AVAILABLE", + "FLASH_ATTN_3_AVAILABLE", + "attention", + "attention_fused", + "flash_attention", + "get_attention_backend", + "set_attention_backend", +] + + +# --------------------------------------------------------------------------- +# Global backend config (set once at startup, e.g. from YAML) +# --------------------------------------------------------------------------- +_ATTENTION_BACKEND: str = "auto" + + +_VALID_BACKENDS = ("auto", "sdpa", "flex_attention", "flash_attn2", "flash_attn3", "flash_attn_aiter") + + +def set_attention_backend(backend: str) -> None: + """ + Set global attention backend. + + Accepted values (case-insensitive): + - "auto" | "sdpa" | "flex_attention" | "flash_attn2" | "flash_attn3" | "flash_attn_aiter" + """ + global _ATTENTION_BACKEND + backend = (backend or "").strip().lower() + if backend not in _VALID_BACKENDS: + hint = "" + if backend.startswith("flash_atten"): + hint = " (did you mean 'flash_attn2' / 'flash_attn3'?)" + raise ValueError(f"attention_backend must be one of {_VALID_BACKENDS}, got '{backend}'{hint}") + _ATTENTION_BACKEND = backend + + +def get_attention_backend() -> str: + return _ATTENTION_BACKEND + + +def _resolve_flash_version(device_type: str) -> int | None: + """ + Returns: + - None: do not use flash attention (use SDPA) + - 2 or 3: use flash attention, prefer that version (3 may fall back to 2 in `flash_attention`) + """ + backend = _ATTENTION_BACKEND + if backend in ("sdpa", "flex_attention", "flash_attn_aiter"): + return None + + if device_type != "cuda": + # auto: CPU should use SDPA + if backend == "auto": + return None + # explicit flash backend: error (prevents silent slow CPU path) + raise RuntimeError(f"attention_backend='{backend}' requires CUDA tensors.") + + # CUDA path + if backend == "auto": + if FLASH_ATTN_3_AVAILABLE: + return 3 + if FLASH_ATTN_2_AVAILABLE: + return 2 + return None + + if backend == "flash_attn2": + if not FLASH_ATTN_2_AVAILABLE: + raise RuntimeError("attention_backend='flash_attn2' but flash-attn 2 is not available.") + return 2 + + # "flash_attn3" + if not FLASH_ATTN_3_AVAILABLE: + raise RuntimeError("attention_backend='flash_attn3' but flash-attn 3 is not available.") + return 3 + + +# --------------------------------------------------------------------------- +# Flash Attention implementation (varlen + fixed-length fast path) +# --------------------------------------------------------------------------- +def _flash_attn_3_fixed(q, k, v, **kwargs): + return flash_attn_interface.flash_attn_func(q, k, v) + + +def _flash_attn_3_varlen( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + causal, + window_size, + deterministic, +): + return flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_q=None, + seqused_k=None, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + ) + + +def _flash_attn_2_fixed(q, k, v, dropout_p, softmax_scale, causal, window_size, deterministic): + try: + return flash_attn.flash_attn_func( + q, + k, + v, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + ) + except TypeError: + return flash_attn.flash_attn_func(q, k, v) + + +def _flash_attn_2_varlen( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + causal, + window_size, + deterministic, +): + return flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + ) + + +def flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + version=None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn( + "Flash attention 3 is not available, use flash attention 2 instead.", + stacklevel=2, + ) + + if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p and window_size are not supported in FA3 now. + return run_flash_attention_backend( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + fixed_attention=_flash_attn_3_fixed, + varlen_attention=_flash_attn_3_varlen, + ) + assert FLASH_ATTN_2_AVAILABLE + return run_flash_attention_backend( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + fixed_attention=_flash_attn_2_fixed, + varlen_attention=_flash_attn_2_varlen, + ) + + +def attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + fa_version=None, +): + """ + Unified attention for Wan2.2-style tensors. + + q: [B, Lq, N, D] + k: [B, Lk, N, D] + v: [B, Lk, N, D] + """ + if _ATTENTION_BACKEND == "flash_attn_aiter": + if q.device.type != "cuda": + raise RuntimeError("attention_backend='flash_attn_aiter' requires CUDA tensors.") + if not AITER_FLASH_ATTN_AVAILABLE: + raise RuntimeError("attention_backend='flash_attn_aiter' but aiter is not available.") + return aiter_flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + ) + + resolved = _resolve_flash_version(q.device.type) + if resolved is not None: + version = resolved + # Only "auto" allows call-site override (useful for debugging). + if _ATTENTION_BACKEND == "auto" and fa_version is not None: + version = fa_version + return flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=version, + ) + + # FlexAttention path (PyTorch, requires torch.compile for perf) + if _ATTENTION_BACKEND == "flex_attention": + from .flex import flex_attention_fn + + return flex_attention_fn( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + ) + + # SDPA fallback (also used for CPU). + # Currently assumes *uniform-length* training sequences, so we do NOT support padding masks on the SDPA path yet + if q_lens is not None or k_lens is not None: + warnings.warn("Padding mask is disabled when using scaled_dot_product_attention.", stacklevel=2) + out_dtype = q.dtype + q_ = q.transpose(1, 2).to(dtype) + k_ = k.transpose(1, 2).to(dtype) + v_ = v.transpose(1, 2).to(dtype) + if q_scale is not None: + q_ = q_ * q_scale + out = torch.nn.functional.scaled_dot_product_attention( + q_, k_, v_, attn_mask=None, is_causal=causal, dropout_p=dropout_p + ) + return out.transpose(1, 2).contiguous().to(out_dtype) + + +def attention_fused( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + num_heads: int, + dropout_p: float = 0.0, + softmax_scale=None, + causal: bool = False, +): + """ + Legacy Wan fused-head attention. + + q/k/v: [B, S, num_heads * head_dim] + returns: [B, S, num_heads * head_dim] + """ + b, s, c = q.shape + if c % num_heads != 0: + raise ValueError(f"hidden dim {c} not divisible by num_heads {num_heads}") + d = c // num_heads + + q4d = q.view(b, s, num_heads, d) + k4d = k.view(b, -1, num_heads, d) + v4d = v.view(b, -1, num_heads, d) + + # For fused calls, q_lens/k_lens are not used (fixed length). + out = attention( + q=q4d, + k=k4d, + v=v4d, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + dtype=q.dtype if q.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16, + ) + return out.flatten(2) diff --git a/primus/backends/diffusion/attention/flex.py b/primus/backends/diffusion/attention/flex.py new file mode 100644 index 000000000..c5f0ad091 --- /dev/null +++ b/primus/backends/diffusion/attention/flex.py @@ -0,0 +1,280 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +FlexAttention backend — pure PyTorch, relies on torch.compile for performance. + +This module provides a drop-in replacement for the SDPA/FlashAttention paths +using ``torch.nn.attention.flex_attention``. FlexAttention compiles a +user-defined ``score_mod`` function into a fused kernel via ``torch.compile``, +giving FlashAttention-level performance without third-party dependencies. + +Key capabilities over SDPA: + - Variable-length padding masks via ``create_block_mask`` + - Sliding-window attention via ``score_mod`` + - Composable attention score modifications +""" + +from __future__ import annotations + +import torch + +# Import flex_attention from PyTorch (available since PyTorch 2.5+). +# We keep the import at module level so that ``FLEX_ATTENTION_AVAILABLE`` +# reflects the actual environment at init time. +FLEX_ATTENTION_AVAILABLE = False +_flex_attention_compiled = None +_create_block_mask = None + +try: + from torch.nn.attention.flex_attention import create_block_mask, flex_attention + + # torch.compile is essential for FlexAttention performance — without it, + # score_mod/block_mask run as eager Python callbacks (separate kernel per op). + # We attempt compile lazily on first call; if Triton/Inductor fails (e.g. + # missing system libs on ROCm), we fall back to eager with a warning. + _flex_attention_eager = flex_attention + _flex_attention_compiled = None # lazy-init on first call + _create_block_mask = create_block_mask + FLEX_ATTENTION_AVAILABLE = True +except Exception: + FLEX_ATTENTION_AVAILABLE = False + + +def _get_flex_attention(): + """ + Return compiled flex_attention, falling back to eager if compile fails. + + torch.compile is lazy: wrapping succeeds immediately, but actual Triton + codegen happens on first invocation and can fail (e.g. missing system libs + on ROCm). We handle this by catching runtime errors on first call and + permanently switching to eager mode. + """ + global _flex_attention_compiled + if _flex_attention_compiled is not None: + return _flex_attention_compiled + + import logging + + logger = logging.getLogger(__name__) + try: + compiled = torch.compile(_flex_attention_eager) + logger.info("flex_attention: torch.compile wrapper created (lazy — actual compile on first call)") + _flex_attention_compiled = compiled + except Exception as e: + logger.warning("flex_attention: torch.compile wrapping failed (%s), using eager mode", e) + _flex_attention_compiled = _flex_attention_eager + return _flex_attention_compiled + + +def _flex_attention_with_fallback(q, k, v, **kwargs): + """Call flex_attention with runtime fallback if compiled version fails.""" + global _flex_attention_compiled + fn = _get_flex_attention() + try: + return fn(q, k, v, **kwargs) + except Exception as e: + if fn is _flex_attention_eager: + raise # already in eager mode, real error + import logging + + logger = logging.getLogger(__name__) + logger.warning( + "flex_attention: compiled call failed (%s: %s), falling back to eager mode", + type(e).__name__, + e, + ) + _flex_attention_compiled = _flex_attention_eager + return _flex_attention_eager(q, k, v, **kwargs) + + +def _build_score_mod( + causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + softmax_scale: float | None = None, +): + """ + Compose a ``score_mod`` function for ``flex_attention``. + + The returned function has signature ``(score, b, h, q_idx, kv_idx) -> score`` + and is compiled by ``torch.compile`` into fused operations. + """ + mods: list = [] + + if softmax_scale is not None: + # flex_attention already applies 1/sqrt(d) by default, so we only + # inject a custom scale if the caller wants a non-default value. + def scale_mod(score, b, h, q_idx, kv_idx): + return score * softmax_scale + + mods.append(scale_mod) + + if causal: + + def causal_mod(score, b, h, q_idx, kv_idx): + return torch.where(q_idx >= kv_idx, score, float("-inf")) + + mods.append(causal_mod) + + if window_size != (-1, -1): + left, right = window_size + + def window_mod(score, b, h, q_idx, kv_idx): + in_window = True + if left >= 0: + in_window = in_window & (q_idx - kv_idx <= left) + if right >= 0: + in_window = in_window & (kv_idx - q_idx <= right) + return torch.where(in_window, score, float("-inf")) + + mods.append(window_mod) + + if not mods: + return None + + # Compose all mods into one function. + def composed_score_mod(score, b, h, q_idx, kv_idx): + for mod in mods: + score = mod(score, b, h, q_idx, kv_idx) + return score + + return composed_score_mod + + +def _build_block_mask( + B: int, + N: int, + Lq: int, + Lk: int, + q_lens: torch.Tensor | None, + k_lens: torch.Tensor | None, + device: torch.device, + causal: bool = False, + window_size: tuple[int, int] = (-1, -1), +): + """ + Build a ``BlockMask`` for ``flex_attention`` from variable-length seqs. + + If both ``q_lens`` and ``k_lens`` are None (uniform-length), returns None + so that ``flex_attention`` uses a full dense mask. + """ + if q_lens is None and k_lens is None and not causal and window_size == (-1, -1): + return None + + # Ensure length tensors are on the same device as the mask indices + # (create_block_mask uses vmap which creates index tensors on `device`) + if q_lens is not None: + q_lens = q_lens.to(device) + if k_lens is not None: + k_lens = k_lens.to(device) + + # Build the mask function. Closures capture the lengths. + def mask_fn(b, h, q_idx, kv_idx): + mask = True + if q_lens is not None: + mask = mask & (q_idx < q_lens[b]) + if k_lens is not None: + mask = mask & (kv_idx < k_lens[b]) + if causal: + mask = mask & (q_idx >= kv_idx) + if window_size != (-1, -1): + left, right = window_size + if left >= 0: + mask = mask & (q_idx - kv_idx <= left) + if right >= 0: + mask = mask & (kv_idx - q_idx <= right) + return mask + + return _create_block_mask(mask_fn, B=B, H=N, Q_LEN=Lq, KV_LEN=Lk, device=device) + + +def flex_attention_fn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: torch.Tensor | None = None, + k_lens: torch.Tensor | None = None, + dropout_p: float = 0.0, + softmax_scale: float | None = None, + q_scale: float | None = None, + causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + deterministic: bool = False, + dtype: torch.dtype = torch.bfloat16, + fa_version=None, +) -> torch.Tensor: + """ + Attention using ``torch.nn.attention.flex_attention``. + + Drop-in replacement for ``attention()`` — same signature, same tensor + layout (q/k/v: ``[B, L, N, D]``). + + Args: + q: Query tensor ``[B, Lq, N, D]`` + k: Key tensor ``[B, Lk, N, D]`` + v: Value tensor ``[B, Lk, N, D]`` + q_lens: Optional valid lengths per batch for queries ``[B]`` + k_lens: Optional valid lengths per batch for keys ``[B]`` + dropout_p: Dropout probability (passed through) + softmax_scale: Custom softmax scale (None = default 1/sqrt(D)) + q_scale: Pre-multiply q by this scalar + causal: Whether to apply causal mask + window_size: (left, right) sliding window; (-1, -1) = disabled + deterministic: Ignored (for API compat with flash_attention) + dtype: Target dtype for computation + fa_version: Ignored (for API compat) + """ + if not FLEX_ATTENTION_AVAILABLE: + raise RuntimeError( + "flex_attention backend requested but torch.nn.attention.flex_attention " + "is not available. Requires PyTorch >= 2.5." + ) + + out_dtype = q.dtype + B, Lq, N, D = q.shape + Lk = k.shape[1] + + # Cast to compute dtype + half_dtypes = (torch.float16, torch.bfloat16) + q_ = q.to(dtype) if q.dtype not in half_dtypes else q + k_ = k.to(dtype) if k.dtype not in half_dtypes else k + v_ = v.to(dtype) if v.dtype not in half_dtypes else v + q_ = q_.to(v_.dtype) + k_ = k_.to(v_.dtype) + + if q_scale is not None: + q_ = q_ * q_scale + + # flex_attention expects [B, N, L, D] (heads-first) + q_ = q_.transpose(1, 2) # [B, N, Lq, D] + k_ = k_.transpose(1, 2) # [B, N, Lk, D] + v_ = v_.transpose(1, 2) # [B, N, Lk, D] + + # Build block mask (handles padding, causal, window in the mask itself) + block_mask = _build_block_mask( + B=B, + N=N, + Lq=Lq, + Lk=Lk, + q_lens=q_lens, + k_lens=k_lens, + device=q.device, + causal=causal, + window_size=window_size, + ) + + # Build score_mod only for softmax_scale override + # (causal and window are handled in block_mask for better sparsity) + score_mod = None + if softmax_scale is not None: + score_mod = _build_score_mod(softmax_scale=softmax_scale) + + # Call flex_attention (with runtime compile-error fallback) + out = _flex_attention_with_fallback(q_, k_, v_, score_mod=score_mod, block_mask=block_mask) + + # Back to [B, L, N, D] + out = out.transpose(1, 2).contiguous() + return out.to(out_dtype) diff --git a/primus/backends/diffusion/diffusion_adapter.py b/primus/backends/diffusion/diffusion_adapter.py new file mode 100644 index 000000000..262b68b64 --- /dev/null +++ b/primus/backends/diffusion/diffusion_adapter.py @@ -0,0 +1,65 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from primus.backends.diffusion.argument_builder import WanArgBuilder +from primus.core.backend.backend_adapter import BackendAdapter +from primus.modules.module_utils import log_rank_0 + + +class DiffusionAdapter(BackendAdapter): + """Primus adapter for PyTorch diffusion training.""" + + def __init__(self, framework: str = "diffusion"): + super().__init__(framework) + + def setup_backend_path(self, backend_path=None) -> str: + """Validate the Primus-owned in-tree diffusion package.""" + if backend_path: + raise ValueError( + "The diffusion backend is built into Primus and does not support " + f"external backend_path overrides. Got: {backend_path}" + ) + + resolved = Path(__file__).resolve().parent + if not resolved.exists(): + raise FileNotFoundError(f"[Primus:Diffusion] backend package does not exist: {resolved}") + + resolved_str = str(resolved) + try: + log_rank_0(f"[Primus:Diffusion] using in-tree backend package -> {resolved_str}") + except Exception: + pass + + return resolved_str + + def convert_config(self, params: Any): + builder = WanArgBuilder() + builder.update(params) + wan_args = builder.finalize() + # convert_config is also called by the standalone prepare hook, where the + # Primus logger may not be initialized yet; guard the informational log. + try: + log_rank_0("[Primus:DiffusionAdapter] Converted Primus module params -> Wan args") + except Exception: + pass + return wan_args + + def load_trainer_class(self, stage: str = "pretrain"): + if stage in ("pretrain", "posttrain", "sft"): + from primus.backends.diffusion.diffusion_pretrain_trainer import ( + DiffusionPretrainTrainer, + ) + + return DiffusionPretrainTrainer + raise ValueError(f"Invalid stage for Diffusion backend: {stage}") + + def detect_backend_version(self) -> str: + return "in-tree" diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py new file mode 100644 index 000000000..1f1aa24c5 --- /dev/null +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -0,0 +1,125 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import importlib.util +from typing import Any + +from primus.core.trainer.base_trainer import BaseTrainer +from primus.core.utils.yaml_utils import nested_namespace_to_dict +from primus.modules.module_utils import log_rank_0 + + +class DiffusionPretrainTrainer(BaseTrainer): + """Primus lifecycle wrapper for Wan diffusion training.""" + + def __init__(self, backend_args: Any): + super().__init__(backend_args=backend_args) + self.wan_trainer = None + + @staticmethod + def _as_dict(value: Any) -> dict: + if isinstance(value, dict): + return value + return nested_namespace_to_dict(value) + + def setup(self): + trainer_cfg = self._as_dict(self.backend_args.trainer) + dataset_cfg = self._as_dict(self.backend_args.dataset) + trainer_args = trainer_cfg.get("args", {}) + attention_backend = trainer_args.get("attention_backend") + + missing = [ + package + for package in ("torch", "loguru", "safetensors", "transformers", "PIL") + if importlib.util.find_spec(package) is None + ] + video_backend = (dataset_cfg.get("config", {}) or {}).get("video_backend") + if video_backend == "imageio" and importlib.util.find_spec("imageio") is None: + missing.append("imageio") + if video_backend == "decord" and importlib.util.find_spec("decord") is None: + missing.append("decord") + if video_backend == "qwen_vl_utils" and importlib.util.find_spec("torchvision") is None: + missing.append("torchvision") + if missing: + raise RuntimeError( + "Diffusion backend missing required Python packages: " + f"{', '.join(missing)}. Install the Wan diffusion training extras first." + ) + + if attention_backend: + from primus.backends.diffusion.attention import set_attention_backend + + set_attention_backend(attention_backend) + log_rank_0(f"[Primus:Diffusion] attention_backend={attention_backend}") + + if attention_backend == "flash_attn_aiter": + from primus.backends.diffusion.attention.aiter import ( + AITER_FLASH_ATTN_AVAILABLE, + ) + + if not AITER_FLASH_ATTN_AVAILABLE: + raise RuntimeError( + "attention_backend=flash_attn_aiter was requested, but AITER flash attention " + "is unavailable in this environment." + ) + + def init(self): + from primus.backends.diffusion.registry import ( + get_dataset_builder, + get_model_builder, + get_trainer_builder, + ) + + model_cfg = self._as_dict(self.backend_args.model) + dataset_cfg = self._as_dict(self.backend_args.dataset) + trainer_cfg = self._as_dict(self.backend_args.trainer) + + model_name = model_cfg["name"] + dataset_name = dataset_cfg["name"] + trainer_name = trainer_cfg["name"] + + model_config = model_cfg["config"] + dataset_config = dataset_cfg["config"] + trainer_args = trainer_cfg["args"] + + log_rank_0( + f"[Primus:Diffusion] Building model={model_name}, dataset={dataset_name}, trainer={trainer_name}" + ) + model = get_model_builder(model_name)(model_config) + dataset, processor = get_dataset_builder(dataset_name)(dataset_config) + self.wan_trainer = get_trainer_builder(trainer_name)( + model=model, + dataset=dataset, + processor=processor, + trainer_args=trainer_args, + ) + + def train(self): + if self.wan_trainer is None: + raise RuntimeError("DiffusionPretrainTrainer.init() must be called before train().") + + self.wan_trainer.train() + self.wan_trainer.save_model() + + def cleanup(self, on_error: bool = False): + try: + import wandb + + if getattr(wandb, "run", None) is not None: + wandb.finish(exit_code=1 if on_error else 0) + except Exception: + pass + + try: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() + except Exception: + pass diff --git a/primus/backends/diffusion/distributed/__init__.py b/primus/backends/diffusion/distributed/__init__.py new file mode 100644 index 000000000..7fea79b62 --- /dev/null +++ b/primus/backends/diffusion/distributed/__init__.py @@ -0,0 +1,32 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Primus diffusion distributed training utilities. + +Modules: + mesh - Device mesh creation, distributed setup + checkpoint - DTCP sharded checkpoint save/load + ulysses - Ulysses Sequence Parallel primitives +""" + +from .checkpoint import load_checkpoint_dtcp, save_checkpoint_dtcp +from .mesh import create_device_mesh, setup_distributed +from .ulysses import distributed_attention, sp_gather, sp_split, sp_unpad + +__all__ = [ + # mesh + "setup_distributed", + "create_device_mesh", + # checkpoint + "save_checkpoint_dtcp", + "load_checkpoint_dtcp", + # ulysses + "distributed_attention", + "sp_split", + "sp_gather", + "sp_unpad", +] diff --git a/primus/backends/diffusion/distributed/checkpoint.py b/primus/backends/diffusion/distributed/checkpoint.py new file mode 100644 index 000000000..4a4828c6b --- /dev/null +++ b/primus/backends/diffusion/distributed/checkpoint.py @@ -0,0 +1,101 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Distributed Tensor Checkpointing (DTCP) save / load utilities. + +Handles FSDP2 sharded checkpoints without gathering to rank 0. +""" + +from __future__ import annotations + +import os +from typing import Any, Dict, Optional + +import torch +import torch.distributed as dist +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter, load, save +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + set_model_state_dict, + set_optimizer_state_dict, +) + +from primus.backends.diffusion.utils.log import logger + +from .mesh import _ensure_process_group + + +def save_checkpoint_dtcp( + model: torch.nn.Module, + optimizer: Optional[torch.optim.Optimizer], + path: str, + epoch: int, + step: int, + additional_data: Dict[str, Any] = None, + *, + model_state_options: Optional[StateDictOptions] = None, + optim_state_options: Optional[StateDictOptions] = None, +): + """Save checkpoint using DTCP (sharded, no gather to rank 0).""" + os.makedirs(path, exist_ok=True) + _ensure_process_group(backend="nccl" if torch.cuda.is_available() else "gloo") + + model_state_options = model_state_options or StateDictOptions(full_state_dict=False) + optim_state_options = optim_state_options or StateDictOptions(full_state_dict=False) + + model_state = get_model_state_dict(model, options=model_state_options) + optim_state = None + if optimizer is not None: + optim_state = get_optimizer_state_dict(model, optimizer, options=optim_state_options) + + state_dict = { + "model": model_state, + **({"optimizer": optim_state} if optim_state is not None else {}), + "meta": {"epoch": epoch, "step": step, **(additional_data or {})}, + } + + save(state_dict, FileSystemWriter(path)) + if (not dist.is_initialized()) or dist.get_rank() == 0: + logger.info(f"Saved DTCP checkpoint to {path}") + + +def load_checkpoint_dtcp( + model: torch.nn.Module, + optimizer: Optional[torch.optim.Optimizer], + path: str, + *, + model_state_options: Optional[StateDictOptions] = None, + optim_state_options: Optional[StateDictOptions] = None, +) -> Dict[str, Any]: + """Load checkpoint using DTCP. Updates model/optimizer in-place. Returns metadata.""" + if not os.path.exists(path): + raise FileNotFoundError(f"Checkpoint not found at {path}") + + _ensure_process_group(backend="nccl" if torch.cuda.is_available() else "gloo") + + model_state_options = model_state_options or StateDictOptions(full_state_dict=False) + optim_state_options = optim_state_options or StateDictOptions(full_state_dict=False) + + model_state = get_model_state_dict(model, options=model_state_options) + optim_state = None + if optimizer is not None: + optim_state = get_optimizer_state_dict(model, optimizer, options=optim_state_options) + + state_dict = { + "model": model_state, + **({"optimizer": optim_state} if optim_state is not None else {}), + } + + load(state_dict, FileSystemReader(path)) + + set_model_state_dict(model, model_state, options=model_state_options) + if optimizer is not None and optim_state is not None: + set_optimizer_state_dict(model, optimizer, optim_state, options=optim_state_options) + + return state_dict.get("meta", {}) diff --git a/primus/backends/diffusion/distributed/mesh.py b/primus/backends/diffusion/distributed/mesh.py new file mode 100644 index 000000000..b9ee3b7b4 --- /dev/null +++ b/primus/backends/diffusion/distributed/mesh.py @@ -0,0 +1,126 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Device mesh creation and distributed setup. + +Supports: + - Pure FSDP2 (1D mesh: dp_shard) + - FSDP2 + HSDP (2D mesh: dp_replicate × dp_shard) + - FSDP2 + Ulysses SP (mesh: dp_shard × ulysses, flattened into dp_shard_sp) + - FSDP2 + HSDP + Ulysses SP (mesh: dp_replicate × dp_shard × ulysses) + +SP ranks are included in the FSDP sharding mesh so that parameters are +sharded across both DP and SP ranks — this reduces per-rank memory. +""" + +from __future__ import annotations + +import os +from datetime import timedelta +from typing import Optional + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh + +from primus.backends.diffusion.utils.log import logger + + +def _ensure_process_group(*, backend: str) -> None: + """Best-effort init_process_group (lazy, avoids single-GPU hangs).""" + if dist.is_initialized(): + return + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "12345") + dist.init_process_group(backend, timeout=timedelta(minutes=60)) + + +def setup_distributed() -> tuple[int, int, int]: + """ + Initialize distributed process group. + + Returns: (rank, world_size, local_rank) + """ + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + + if world_size > 1 and not dist.is_initialized(): + _ensure_process_group(backend="nccl") + + return rank, world_size, local_rank + + +def create_device_mesh( + world_size: int, + sp_size: int = 1, + dp_replicate: int = 1, +) -> Optional[DeviceMesh]: + """ + Create a DeviceMesh for FSDP2 + optional Ulysses Sequence Parallel. + + Mesh layout (innermost → outermost): + [dp_replicate?] × [dp_shard] × [ulysses?] + + Flattened sub-meshes created automatically: + - "dp_shard_sp": dp_shard × ulysses (used for FSDP2 fully_shard) + - "dp": dp_replicate × dp_shard (used for DistributedSampler) + + Args: + world_size: total number of ranks + sp_size: Ulysses sequence parallel size (must divide world_size) + dp_replicate: HSDP replicate dimension (1 = no HSDP) + + Returns: + DeviceMesh, or None if world_size <= 1 + """ + if world_size <= 1: + return None + + _ensure_process_group(backend="nccl") + + dp_shard = world_size // (sp_size * dp_replicate) + if dp_shard * sp_size * dp_replicate != world_size: + raise ValueError( + f"world_size={world_size} is not divisible by " f"sp_size={sp_size} * dp_replicate={dp_replicate}" + ) + + # Build mesh dimensions + dims = [] + names = [] + if dp_replicate > 1: + dims.append(dp_replicate) + names.append("dp_replicate") + dims.append(dp_shard) + names.append("dp_shard") + if sp_size > 1: + dims.append(sp_size) + names.append("ulysses") + + mesh = init_device_mesh("cuda", tuple(dims), mesh_dim_names=tuple(names)) + + # Flatten composite sub-meshes for convenient access + if sp_size > 1: + mesh["dp_shard", "ulysses"]._flatten("dp_shard_sp") + + if dp_replicate > 1: + mesh["dp_replicate", "dp_shard"]._flatten("dp") + + rank = dist.get_rank() if dist.is_initialized() else 0 + if rank == 0: + logger.info( + f"DeviceMesh created: dims={dict(zip(names, dims))}, " + f"sp_size={sp_size}, dp_replicate={dp_replicate}" + ) + + return mesh diff --git a/primus/backends/diffusion/distributed/ulysses.py b/primus/backends/diffusion/distributed/ulysses.py new file mode 100644 index 000000000..04bf84cfc --- /dev/null +++ b/primus/backends/diffusion/distributed/ulysses.py @@ -0,0 +1,200 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Ulysses Sequence Parallel primitives. + +Follows the Wan2.2 official design: a self-contained ``distributed_attention`` +function that wraps all-to-all + attention + all-to-all, plus ``sp_split`` / +``sp_gather`` / ``sp_unpad`` for slicing model inputs and gathering outputs. + +The model code only needs to: + 1. ``sp_split`` inputs before the blocks + 2. call ``distributed_attention`` instead of ``attention`` in self-attn + 3. ``sp_gather`` + ``sp_unpad`` the output after the head + +Gradient scaling (÷ sp_size in sp_slice backward, × sp_size in sp_gather +backward) ensures correctness with FSDP2 gradient averaging across the +combined DP+SP mesh. +""" + +from __future__ import annotations + +from typing import Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +# --------------------------------------------------------------------------- +# Core autograd primitives +# --------------------------------------------------------------------------- + + +class _SeqAllToAll(torch.autograd.Function): + """All-to-all with autograd. Backward is the inverse (swap dims).""" + + @staticmethod + def forward(ctx, group, x, scatter_dim, gather_dim): + ctx.group = group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + sp_size = dist.get_world_size(group) + input_list = [t.contiguous() for t in x.tensor_split(sp_size, scatter_dim)] + output_list = [torch.empty_like(input_list[0]) for _ in range(sp_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + @staticmethod + def backward(ctx, grad_output): + return ( + None, + _SeqAllToAll.apply(ctx.group, grad_output, ctx.gather_dim, ctx.scatter_dim), + None, + None, + ) + + +class _SliceWithGather(torch.autograd.Function): + """Forward: slice. Backward: all-gather ÷ sp_size (FSDP2 compat).""" + + @staticmethod + def forward(ctx, x, dim, group): + ctx.dim = dim + ctx.group = group + sp_size = dist.get_world_size(group) + sp_rank = dist.get_rank(group) + ctx.sp_size = sp_size + chunk_size = x.shape[dim] // sp_size + return x.narrow(dim, sp_rank * chunk_size, chunk_size).contiguous() + + @staticmethod + def backward(ctx, grad_output): + gathered = [torch.empty_like(grad_output) for _ in range(ctx.sp_size)] + dist.all_gather(gathered, grad_output.contiguous(), group=ctx.group) + return torch.cat(gathered, dim=ctx.dim) / ctx.sp_size, None, None + + +class _GatherWithSlice(torch.autograd.Function): + """Forward: all-gather. Backward: slice × sp_size (FSDP2 compat).""" + + @staticmethod + def forward(ctx, x, dim, group): + ctx.dim = dim + ctx.group = group + sp_size = dist.get_world_size(group) + sp_rank = dist.get_rank(group) + ctx.sp_size = sp_size + ctx.sp_rank = sp_rank + ctx.chunk_size = x.shape[dim] + gathered = [torch.empty_like(x) for _ in range(sp_size)] + dist.all_gather(gathered, x.contiguous(), group=group) + return torch.cat(gathered, dim=dim) + + @staticmethod + def backward(ctx, grad_output): + chunk = grad_output.narrow(ctx.dim, ctx.sp_rank * ctx.chunk_size, ctx.chunk_size).contiguous() + return chunk * ctx.sp_size, None, None + + +# --------------------------------------------------------------------------- +# Public API — high-level +# --------------------------------------------------------------------------- + + +def distributed_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + group: dist.ProcessGroup, + attention_fn: Callable, + **attention_kwargs, +) -> torch.Tensor: + """ + Ulysses distributed attention (DeepSpeed Ulysses, arXiv:2309.14509). + + Wraps ``attention_fn`` with all-to-all communication so that each rank + computes attention on the **full sequence** with a **subset of heads**. + + Input / output shapes: ``[B, S/P, H, D]`` (sharded-seq, full-heads). + + Args: + q, k, v: query / key / value ``[B, S/P, H, D]`` + group: Ulysses SP process group + attention_fn: e.g. ``primus.backends.diffusion.attention.attention`` + **attention_kwargs: forwarded to ``attention_fn(q=, k=, v=, ...)`` + """ + # scatter heads, gather seq: [B, S/P, H, D] → [B, S, H/P, D] + q = _SeqAllToAll.apply(group, q, 2, 1) + k = _SeqAllToAll.apply(group, k, 2, 1) + v = _SeqAllToAll.apply(group, v, 2, 1) + # standard attention on full sequence with partial heads + out = attention_fn(q=q, k=k, v=v, **attention_kwargs) + # scatter seq, gather heads: [B, S, H/P, D] → [B, S/P, H, D] + return _SeqAllToAll.apply(group, out, 1, 2) + + +def sp_split( + tensors: List[torch.Tensor], + dim: int, + group: dist.ProcessGroup, +) -> Tuple[List[torch.Tensor], int]: + """ + Pad + slice multiple tensors for sequence parallelism. + + Returns ``(sliced_tensors, original_size)`` where *original_size* is + the size along *dim* before padding (needed by :func:`sp_unpad`). + """ + sp_size = dist.get_world_size(group) + original_size: Optional[int] = None + results: List[torch.Tensor] = [] + for t in tensors: + t, orig = _sp_pad(t, dim, sp_size) + if original_size is None: + original_size = orig + results.append(_SliceWithGather.apply(t, dim, group)) + assert original_size is not None + return results, original_size + + +# --------------------------------------------------------------------------- +# Public API — low-level (used directly for gathering the output) +# --------------------------------------------------------------------------- + + +def sp_slice(x: torch.Tensor, dim: int, group: dist.ProcessGroup) -> torch.Tensor: + """Slice ``x`` along *dim* for this SP rank (with autograd). Low-level.""" + return _SliceWithGather.apply(x, dim, group) + + +def sp_gather(x: torch.Tensor, dim: int, group: dist.ProcessGroup) -> torch.Tensor: + """All-gather ``x`` along *dim* from all SP ranks (with autograd).""" + return _GatherWithSlice.apply(x, dim, group) + + +def sp_unpad(x: torch.Tensor, dim: int, original_size: int) -> torch.Tensor: + """Remove padding added by :func:`sp_split`.""" + if x.shape[dim] == original_size: + return x + return x.narrow(dim, 0, original_size) + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _sp_pad(x: torch.Tensor, dim: int, sp_size: int) -> Tuple[torch.Tensor, int]: + """Pad *x* along *dim* so its size is divisible by *sp_size*.""" + original_size = x.shape[dim] + remainder = original_size % sp_size + if remainder == 0: + return x, original_size + pad_amount = sp_size - remainder + pad_config = [0] * (2 * x.dim()) + pos = 2 * (x.dim() - 1 - dim) + pad_config[pos + 1] = pad_amount + return F.pad(x, pad_config), original_size diff --git a/primus/backends/diffusion/models/__init__.py b/primus/backends/diffusion/models/__init__.py new file mode 100644 index 000000000..2798dc3fb --- /dev/null +++ b/primus/backends/diffusion/models/__init__.py @@ -0,0 +1,13 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Wan model exports for the Primus Wan backend.""" + +from .wan import WanForTraining + +__all__ = [ + "WanForTraining", +] diff --git a/primus/backends/diffusion/models/interface.py b/primus/backends/diffusion/models/interface.py new file mode 100644 index 000000000..1355c3658 --- /dev/null +++ b/primus/backends/diffusion/models/interface.py @@ -0,0 +1,48 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from abc import ABC, abstractmethod +from typing import Any, Dict + +import torch +import torch.nn as nn + + +class GenAIModel(nn.Module, ABC): + """ + Unified interface for Generative Models (DiT, AR, Hybrid). + Wraps the Backbone (DiT/Transformer), VAE, and TextEncoder. + """ + + @abstractmethod + def forward_train(self, batch: Dict[str, Any], scheduler: Any = None) -> Dict[str, torch.Tensor]: + """ + The single entry point for training. + + Responsibilities: + 1. Process raw batch (Tokenization, VAE Encoding if needed). + 2. Apply Training Recipe (e.g., Add Noise for DiT, Shift Tokens for AR). + 3. Forward pass through Backbone. + 4. Calculate Loss. + + Args: + batch: Raw batch from data loader. + scheduler: Optional scheduler (e.g., FlowMatchScheduler) passed from Trainer. + + Returns: + { + "loss": torch.Tensor, # Main optimization objective + "log_metrics": Dict, # Metrics for WandB (MSE, Accuracy, etc.) + } + """ + + @abstractmethod + def forward_inference(self, batch: Dict[str, Any], **kwargs): + """ + Entry point for validation/inference. + For DiT: Runs the diffusion sampling loop. + For AR: Runs autoregressive generation. + """ diff --git a/primus/backends/diffusion/models/registrations/__init__.py b/primus/backends/diffusion/models/registrations/__init__.py new file mode 100644 index 000000000..98d014572 --- /dev/null +++ b/primus/backends/diffusion/models/registrations/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Model registrations.""" diff --git a/primus/backends/diffusion/models/registrations/wan.py b/primus/backends/diffusion/models/registrations/wan.py new file mode 100644 index 000000000..aa22cc48c --- /dev/null +++ b/primus/backends/diffusion/models/registrations/wan.py @@ -0,0 +1,234 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Register the Wan model builder.""" + +from __future__ import annotations + +import glob +import os +from typing import Any + +import torch +from safetensors.torch import load_file as safe_load_file + +from primus.backends.diffusion.models.wan.adapter import WanForTraining +from primus.backends.diffusion.models.wan.components import WanComponents +from primus.backends.diffusion.models.wan.configuration_wanvideo import WanVideoConfig +from primus.backends.diffusion.models.wan.t5 import umt5_xxl_encoder_from_checkpoint +from primus.backends.diffusion.models.wan.train_pipeline import ( + WanFlowMatchTrainPipeline, +) +from primus.backends.diffusion.models.wan.vae2_1 import Wan2_1_VAE +from primus.backends.diffusion.models.wan.vae2_2 import Wan2_2_VAE +from primus.backends.diffusion.models.wan.wan_dit import WanModel as WanDiT +from primus.backends.diffusion.utils.log import logger +from primus.backends.diffusion.utils.train_utils import count_parameters + + +def _strip_module_prefix(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + for k, v in state_dict.items(): + if k.startswith("module."): + out[k[len("module.") :]] = v + else: + out[k] = v + return out + + +def _load_state_dict(path: str) -> dict[str, torch.Tensor]: + if path.endswith(".safetensors"): + return dict(safe_load_file(path)) + obj = torch.load(path, map_location="cpu") + if isinstance(obj, dict) and "model" in obj and isinstance(obj["model"], dict): + obj = obj["model"] + if not isinstance(obj, dict): + raise ValueError(f"Unsupported checkpoint format at {path}") + return obj + + +def _load_dit_weights_into_module(dit: torch.nn.Module, pretrained_path: str): + """ + Support two common cases: + 1) Wan training export: `dit_model.safetensors` with keys like `blocks.0...` + 2) Official-ish export: `*model*.safetensors|bin` possibly with keys like `dit.blocks.0...` or `blocks.0...` + """ + # Direct file + if os.path.isfile(pretrained_path): + state = _strip_module_prefix(_load_state_dict(pretrained_path)) + # Accept either `dit.*` or plain keys. + if any(k.startswith("dit.") for k in state): + state = {k[len("dit.") :]: v for k, v in state.items() if k.startswith("dit.")} + result = dit.load_state_dict(state, strict=False) + logger.info( + f"Loaded DiT from file. missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}" + ) + return + + # Directory: prefer explicit Wan trainer file name, else fall back to pattern search + candidates: list[str] = [] + for fname in ("dit_model.safetensors", "diffusion_pytorch_model.safetensors", "model.safetensors"): + p = os.path.join(pretrained_path, fname) + if os.path.exists(p): + candidates.append(p) + if not candidates: + candidates = sorted(glob.glob(os.path.join(pretrained_path, "*model*.safetensors"))) + if not candidates: + candidates = sorted(glob.glob(os.path.join(pretrained_path, "*model*.bin"))) + if not candidates: + raise FileNotFoundError(f"No DiT weights found under {pretrained_path}") + + merged: dict[str, torch.Tensor] = {} + for ckpt in candidates: + part = _strip_module_prefix(_load_state_dict(ckpt)) + merged.update(part) + + if any(k.startswith("dit.") for k in merged): + merged = {k[len("dit.") :]: v for k, v in merged.items() if k.startswith("dit.")} + + result = dit.load_state_dict(merged, strict=False) + logger.info( + f"Loaded DiT from dir. files={len(candidates)} missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}" + ) + + +def _load_vae(vae: torch.nn.Module, ckpt_path: str | None): + if not ckpt_path: + return + logger.info(f"Loading VAE from {ckpt_path}") + state = _strip_module_prefix(_load_state_dict(ckpt_path)) + # VAE wrapper may keep real module under `.model` + target = vae.model if hasattr(vae, "model") else vae + result = target.load_state_dict(state, strict=False) + logger.info(f"VAE loaded. missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}") + + +def build_wan_model(model_config: dict): + """ + YAML compatibility: + model_config: + name: wan + load_from_pretrained_path: /path/to/Wan2.2-TI2V-5B (optional) + config: {...} (optional overrides) + encoder: + t5_encoder: ... + autoencoder: ... + """ + encoder_cfg = model_config.get("encoder", {}) if isinstance(model_config.get("encoder"), dict) else {} + cfg_dict: dict[str, Any] = dict(model_config.get("config", {}) or {}) + if encoder_cfg: + cfg_dict.setdefault("encoder", encoder_cfg) + + # 1. Try to load config.json from pretrained path if available + pretrained_path = model_config.get("load_from_pretrained_path") + if pretrained_path: + import json + + cfg_path = os.path.join(pretrained_path, "config.json") + if os.path.exists(cfg_path): + logger.info(f"Loading config from {cfg_path}") + with open(cfg_path) as f: + loaded_cfg = json.load(f) + + # Map WanModel checkpoint config keys to WanVideoConfig keys. + mapping = { + "dim": "dit_hidden_size", + "num_layers": "dit_num_layers", + "num_heads": "dit_num_heads", + "ffn_dim": "dit_intermediate_size", + "in_dim": "dit_in_channels", + "out_dim": "dit_out_channels", + "freq_dim": "dit_freq_dim", + "text_len": "text_len", + } + + for k, v in loaded_cfg.items(): + target_key = mapping.get(k) + # Only infer values the user did not set explicitly in YAML. + if target_key is not None and target_key not in cfg_dict: + cfg_dict[target_key] = v + + # Build a WanVideoConfig-compatible object to maximize reuse of existing fields. + model_cfg = WanVideoConfig(**cfg_dict) + + # Build DiT (close to official Wan2.2 modules/model.py signature). + # model_config.config.model_type: one of {"t2v","i2v","ti2v","s2v"} + dit_task_type = getattr(model_cfg, "model_type", None) + if dit_task_type not in ("t2v", "i2v", "ti2v", "s2v"): + raise ValueError( + "wan requires explicit `model_config.config.model_type` in YAML " + "(one of: t2v / i2v / ti2v / s2v). " + f"Got: {dit_task_type!r}" + ) + dit = WanDiT( + model_type=dit_task_type, + patch_size=tuple(model_cfg.dit_patch_size), + text_len=int(getattr(model_cfg, "text_len", 512)), + in_dim=int(model_cfg.dit_in_channels), + dim=int(model_cfg.dit_hidden_size), + ffn_dim=int(model_cfg.dit_intermediate_size), + freq_dim=int(model_cfg.dit_freq_dim), + text_dim=int(model_cfg.dit_text_dim), + out_dim=int(model_cfg.dit_out_channels), + num_heads=int(model_cfg.dit_num_heads), + num_layers=int(model_cfg.dit_num_layers), + window_size=tuple(getattr(model_cfg, "dit_window_size", (-1, -1))), + qk_norm=bool(getattr(model_cfg, "dit_qk_norm", True)), + cross_attn_norm=bool(getattr(model_cfg, "dit_cross_attn_norm", True)), + eps=float(model_cfg.dit_eps), + ) + + # Build VAE + vae_type = getattr(model_cfg, "vae_type", "wan_video_vae_38") + vae_ckpt = encoder_cfg.get("vae_checkpoint") or encoder_cfg.get("autoencoder") + if vae_type in ("wan2.2", "wan_video_vae_38"): + if not vae_ckpt: + raise ValueError("wan requires `model_config.encoder.autoencoder` (Wan2.2 VAE checkpoint path)") + vae = Wan2_2_VAE(z_dim=48, vae_pth=vae_ckpt) + vae.upsampling_factor = 16 + else: + if not vae_ckpt: + raise ValueError("wan requires `model_config.encoder.autoencoder` (Wan2.1 VAE checkpoint path)") + vae = Wan2_1_VAE(z_dim=16, vae_pth=vae_ckpt) + vae.upsampling_factor = 8 + + # Build text encoder (UMT5-XXL, encoder-only). Tokenization is handled by the dataset processor. + t5_ckpt = encoder_cfg.get("t5_encoder") + if not t5_ckpt: + raise ValueError("wan requires `model_config.encoder.t5_encoder` (UMT5 encoder checkpoint path)") + text_encoder = umt5_xxl_encoder_from_checkpoint(t5_ckpt, dtype=torch.bfloat16, device="cpu") + + # Optionally load pretrained DiT + pretrained_path = model_config.get("load_from_pretrained_path") + if pretrained_path: + logger.info(f"Loading DiT weights from {pretrained_path}") + _load_dit_weights_into_module(dit, pretrained_path) + + # Load encoders (VAE weights loaded into `.model`) + _load_vae(vae, vae_ckpt) + + components = WanComponents(dit=dit, vae=vae, text_encoder=text_encoder, image_encoder=None) + pipeline = WanFlowMatchTrainPipeline() + + model = WanForTraining( + components=components, + train_pipeline=pipeline, + model_config=model_cfg, + raw_config={ + "model_config": model_config, + "resolved_model_cfg": getattr(model_cfg, "__dict__", cfg_dict), + }, + trainable_modules=getattr(model_cfg, "trainable_modules", None), + ) + + total_params, trainable_params = count_parameters(model) + logger.info(f"wan parameters: total={total_params/1e9:.3f}B trainable={trainable_params/1e9:.3f}B") + if hasattr(model, "freeze_except"): + model.freeze_except() + total_params, trainable_params = count_parameters(model) + logger.info(f"wan after freeze: total={total_params/1e9:.3f}B trainable={trainable_params/1e9:.3f}B") + + return model diff --git a/primus/backends/diffusion/models/wan/__init__.py b/primus/backends/diffusion/models/wan/__init__.py new file mode 100644 index 000000000..86c01a615 --- /dev/null +++ b/primus/backends/diffusion/models/wan/__init__.py @@ -0,0 +1,20 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +wan: Components + TaskPipeline design for Wan PyTorch/FSDP training. + +Goal: +- Keep modeling close to official Wan2.2 (pure torch modules). +- Keep trainer generic: trainer calls model(batch, scheduler) -> {"loss": ...}. +- Decouple training/inference workflow (pipeline) from modeling (modules). +""" + +from .adapter import WanForTraining + +__all__ = [ + "WanForTraining", +] diff --git a/primus/backends/diffusion/models/wan/adapter.py b/primus/backends/diffusion/models/wan/adapter.py new file mode 100644 index 000000000..995177d6e --- /dev/null +++ b/primus/backends/diffusion/models/wan/adapter.py @@ -0,0 +1,154 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch +import torch.nn as nn + +from primus.backends.diffusion.models.interface import GenAIModel + +from .components import WanComponents +from .train_pipeline import WanFlowMatchTrainPipeline + + +@dataclass +class WanConfigShim: + """ + A tiny config shim to keep trainer saving behavior happy. + We intentionally do not implement a full HF/Diffusers config system. + """ + + raw: dict + + def save_pretrained(self, save_directory: str): + # Best-effort: keep minimal JSON for debugging/reproducibility. + import json + import os + + os.makedirs(save_directory, exist_ok=True) + path = os.path.join(save_directory, "wan_config.json") + with open(path, "w") as f: + json.dump(self.raw, f, indent=2, sort_keys=True) + + def to_dict(self): + return self.raw + + +class WanForTraining(GenAIModel, nn.Module): + """ + A thin adapter that exposes the call pattern expected by Wan trainers. + + Trainers call: + outputs = model(batch, scheduler) -> {"loss": ...} + + Internally we delegate workflow to the pipeline. + """ + + def __init__( + self, + *, + components: WanComponents, + train_pipeline: WanFlowMatchTrainPipeline, + model_config: Any, + raw_config: Optional[dict] = None, + trainable_modules: Optional[str] = None, + ): + super().__init__() + self.components = components + self.train_pipeline = train_pipeline + self.model_config = model_config + self.trainable_modules = trainable_modules + + # Expose common attribute names expected by trainers/FSDP ignore regex. + self.dit = components.dit + self.vae = components.vae + self.text_encoder = components.text_encoder + self.image_encoder = components.image_encoder + + self.config = WanConfigShim(raw=raw_config or {}) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def to(self, *args, **kwargs): + """ + Override to propagate .to() to non-nn.Module components. + + Wan VAE wrappers (Wan2_1_VAE, Wan2_2_VAE) are plain Python classes + with a custom .to() method, not nn.Module subclasses. PyTorch's + nn.Module.to() only recurses into registered submodules, so the VAE + would be silently skipped — causing device mismatches on multi-GPU. + """ + result = super().to(*args, **kwargs) + for component in (self.vae, self.image_encoder): + if component is not None and not isinstance(component, nn.Module) and hasattr(component, "to"): + component.to(*args, **kwargs) + return result + + def freeze_except(self): + """ + Keep the Wan training behavior: freeze non-trainable modules. + Default: train only DiT. + """ + mode = ( + self.trainable_modules or getattr(self.model_config, "trainable_modules", None) or "dit" + ).lower() + + def freeze(m: nn.Module): + for p in m.parameters(): + p.requires_grad_(False) + + def unfreeze(m: nn.Module): + for p in m.parameters(): + p.requires_grad_(True) + + # Freeze everything first + freeze(self) + + # Unfreeze requested parts + if mode in ("dit", "diffusion", "backbone"): + unfreeze(self.dit) + elif mode in ("all",): + unfreeze(self) + else: + # Conservative default: DiT only + unfreeze(self.dit) + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + """ + Activated by HF Trainer when `gradient_checkpointing=True`. + """ + if self.dit and hasattr(self.dit, "gradient_checkpointing"): + self.dit.gradient_checkpointing = True + + def forward(self, *args, **kwargs): + # Match existing trainer call convention: model(batch, scheduler) + if len(args) >= 1 and isinstance(args[0], dict): + scheduler = args[1] if len(args) > 1 else kwargs.get("scheduler", None) + return self.forward_train(args[0], scheduler=scheduler) + raise TypeError("WanForTraining.forward expects (batch_dict, scheduler)") + + def forward_train(self, batch: Dict[str, Any], scheduler: Any = None) -> Dict[str, torch.Tensor]: + if scheduler is None: + raise ValueError("scheduler must be provided by trainer") + return self.train_pipeline.compute_loss( + components=self.components, + batch=batch, + scheduler=scheduler, + model_config=self.model_config, + ) + + def forward_inference(self, batch: Dict[str, Any], **kwargs): + raise NotImplementedError("wan inference pipeline not wired yet") diff --git a/primus/backends/diffusion/models/wan/attention_backend.py b/primus/backends/diffusion/models/wan/attention_backend.py new file mode 100644 index 000000000..306f8afe4 --- /dev/null +++ b/primus/backends/diffusion/models/wan/attention_backend.py @@ -0,0 +1,228 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +# +# This file is vendored from `Wan2.2/wan/modules/attention.py` with minimal changes: +# - keep PyTorch implementation +# - allow CPU fallback when flash-attn is unavailable or tensors are on CPU + +import warnings + +import torch + +try: + import flash_attn_interface + + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +__all__ = [ + "flash_attention", + "attention", +] + + +def flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + version=None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == "cuda" and q.size(-1) <= 256 + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # Fast path (fixed-length, no padding mask): + # Prefer flash_attn_func over flash_attn_varlen_func to reduce overhead. + if q_lens is None and k_lens is None: + qh = half(q) + kh = half(k) + vh = half(v) + qh = qh.to(vh.dtype) + kh = kh.to(vh.dtype) + if q_scale is not None: + qh = qh * q_scale + + # apply attention + if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + x = flash_attn_interface.flash_attn_func(qh, kh, vh) + if isinstance(x, tuple): + x = x[0] + else: + assert FLASH_ATTN_2_AVAILABLE + # flash_attn_func signature differs across versions; be conservative + try: + x = flash_attn.flash_attn_func( + qh, + kh, + vh, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + ) + except TypeError: + x = flash_attn.flash_attn_func(qh, kh, vh) + return x.type(out_dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + + if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn("Flash attention 3 is not available, use flash attention 2 instead.") + + # apply attention + if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p, window_size are not supported in FA3 now. + x = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + seqused_q=None, + seqused_k=None, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + )[0].unflatten(0, (b, lq)) + else: + assert FLASH_ATTN_2_AVAILABLE + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + ).unflatten(0, (b, lq)) + + # output + return x.type(out_dtype) + + +def attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + fa_version=None, +): + # Minimal change vs upstream: allow CPU fallback. + if (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE) and q.device.type == "cuda": + return flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=fa_version, + ) + + if q_lens is not None or k_lens is not None: + warnings.warn( + "Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance." + ) + attn_mask = None + + out_dtype = q.dtype + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p + ) + + out = out.transpose(1, 2).contiguous().to(out_dtype) + return out diff --git a/primus/backends/diffusion/models/wan/components.py b/primus/backends/diffusion/models/wan/components.py new file mode 100644 index 000000000..3e5a854d2 --- /dev/null +++ b/primus/backends/diffusion/models/wan/components.py @@ -0,0 +1,26 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch.nn as nn + + +@dataclass +class WanComponents: + """ + A thin container for model components. + + Keep this intentionally simple: it's just a bundle of modules that a pipeline can use. + """ + + dit: nn.Module + vae: nn.Module + text_encoder: nn.Module + image_encoder: Optional[nn.Module] = None diff --git a/primus/backends/diffusion/models/wan/configuration_wanvideo.py b/primus/backends/diffusion/models/wan/configuration_wanvideo.py new file mode 100644 index 000000000..dcfb1f1d6 --- /dev/null +++ b/primus/backends/diffusion/models/wan/configuration_wanvideo.py @@ -0,0 +1,84 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +# coding=utf-8 +# Copyright 2024 WanVideo team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict + + +class WanVideoConfig: + model_type = "wanvideo" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + dit_hidden_size: int = 3072, + dit_num_layers: int = 30, + dit_num_heads: int = 24, + dit_intermediate_size: int = 14336, + dit_patch_size: tuple = (1, 2, 2), + dit_in_channels: int = 48, + dit_out_channels: int = 48, + dit_freq_dim: int = 256, + dit_text_dim: int = 4096, + dit_eps: float = 1e-6, + dit_has_image_input: bool = False, + dit_has_image_pos_emb: bool = False, + dit_has_ref_conv: bool = False, + trainable_modules=None, + seperated_timestep: bool = True, + require_clip_embedding: bool = False, + require_vae_embedding: bool = False, + fuse_vae_embedding_in_latents: bool = True, + tie_word_embeddings: bool = False, + vae_type: str = "wan_video_vae_38", + **kwargs, + ): + # DiT configuration + self.vae_type = vae_type + self.dit_hidden_size = dit_hidden_size + self.dit_num_layers = dit_num_layers + self.dit_num_heads = dit_num_heads + self.dit_intermediate_size = dit_intermediate_size + self.dit_patch_size = dit_patch_size + self.dit_in_channels = dit_in_channels + self.dit_out_channels = dit_out_channels + self.dit_freq_dim = dit_freq_dim + self.dit_text_dim = dit_text_dim + self.dit_eps = dit_eps + self.dit_has_image_input = dit_has_image_input + + self.dit_has_image_pos_emb = dit_has_image_pos_emb + self.dit_has_ref_conv = dit_has_ref_conv + + self.seperated_timestep = seperated_timestep + self.require_clip_embedding = require_clip_embedding + self.require_vae_embedding = require_vae_embedding + self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents + + self.trainable_modules = trainable_modules + + for k, v in kwargs.items(): + setattr(self, k, v) + + def to_dict(self) -> Dict[str, Any]: + """Convert configuration to dictionary""" + output = {k: v for k, v in self.__dict__.items() if not k.startswith("_")} + output["model_type"] = self.model_type + return output diff --git a/primus/backends/diffusion/models/wan/t5.py b/primus/backends/diffusion/models/wan/t5.py new file mode 100644 index 000000000..d9cdfab71 --- /dev/null +++ b/primus/backends/diffusion/models/wan/t5.py @@ -0,0 +1,299 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +# +# Notes for Wan training: +# - We keep this file focused on the PyTorch encoder only. +# - We do NOT vendor the tokenizer wrapper (ftfy/regex deps). Training uses +# `input_ids/attention_mask` produced by the existing processor instead. + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = [ + "T5Encoder", + "umt5_xxl_encoder", + "umt5_xxl_encoder_from_checkpoint", +] + + +def fp16_clamp(x): + if x.dtype == torch.float16 and torch.isinf(x).any(): + clamp = torch.finfo(x.dtype).max - 1000 + x = torch.clamp(x, min=-clamp, max=clamp) + return x + + +class GELU(nn.Module): + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + + +class T5LayerNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) + if self.weight.dtype in (torch.float16, torch.bfloat16): + x = x.type_as(self.weight) + return self.weight * x + + +class T5Attention(nn.Module): + def __init__(self, dim, dim_attn, num_heads, dropout=0.1): + assert dim_attn % num_heads == 0 + super().__init__() + self.dim = dim + self.dim_attn = dim_attn + self.num_heads = num_heads + self.head_dim = dim_attn // num_heads + + self.q = nn.Linear(dim, dim_attn, bias=False) + self.k = nn.Linear(dim, dim_attn, bias=False) + self.v = nn.Linear(dim, dim_attn, bias=False) + self.o = nn.Linear(dim_attn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, context=None, mask=None, pos_bias=None): + """ + x: [B, L1, C] + context: [B, L2, C] or None + mask: [B, L2] or [B, L1, L2] or None + """ + context = x if context is None else context + b, n, c = x.size(0), self.num_heads, self.head_dim + + q = self.q(x).view(b, -1, n, c) + k = self.k(context).view(b, -1, n, c) + v = self.v(context).view(b, -1, n, c) + + attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) + if pos_bias is not None: + attn_bias += pos_bias + if mask is not None: + assert mask.ndim in (2, 3) + mask = mask.view(b, 1, 1, -1) if mask.ndim == 2 else mask.unsqueeze(1) + attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min) + + # T5 does not use scaling + attn = torch.einsum("binc,bjnc->bnij", q, k) + attn_bias + attn = F.softmax(attn.float(), dim=-1).type_as(attn) + x = torch.einsum("bnij,bjnc->binc", attn, v) + + x = x.reshape(b, -1, n * c) + x = self.o(x) + x = self.dropout(x) + return x + + +class T5FeedForward(nn.Module): + def __init__(self, dim, dim_ffn, dropout=0.1): + super().__init__() + self.dim = dim + self.dim_ffn = dim_ffn + self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU()) + self.fc1 = nn.Linear(dim, dim_ffn, bias=False) + self.fc2 = nn.Linear(dim_ffn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + x = self.fc1(x) * self.gate(x) + x = self.dropout(x) + x = self.fc2(x) + x = self.dropout(x) + return x + + +class T5RelativeEmbedding(nn.Module): + def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): + super().__init__() + self.num_buckets = num_buckets + self.num_heads = num_heads + self.bidirectional = bidirectional + self.max_dist = max_dist + self.embedding = nn.Embedding(num_buckets, num_heads) + + def forward(self, lq, lk): + device = self.embedding.weight.device + rel_pos = torch.arange(lk, device=device).unsqueeze(0) - torch.arange(lq, device=device).unsqueeze(1) + rel_pos = self._relative_position_bucket(rel_pos) + rel_pos_embeds = self.embedding(rel_pos) + rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(0) # [1, N, Lq, Lk] + return rel_pos_embeds.contiguous() + + def _relative_position_bucket(self, rel_pos): + if self.bidirectional: + num_buckets = self.num_buckets // 2 + rel_buckets = (rel_pos > 0).long() * num_buckets + rel_pos = torch.abs(rel_pos) + else: + num_buckets = self.num_buckets + rel_buckets = 0 + rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) + + max_exact = num_buckets // 2 + rel_pos_large = ( + max_exact + + ( + torch.log(rel_pos.float() / max_exact) + / math.log(self.max_dist / max_exact) + * (num_buckets - max_exact) + ).long() + ) + rel_pos_large = torch.min(rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) + rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) + return rel_buckets + + +class T5SelfAttention(nn.Module): + def __init__(self, dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos=True, dropout=0.1): + super().__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + self.norm1 = T5LayerNorm(dim) + self.attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm2 = T5LayerNorm(dim) + self.ffn = T5FeedForward(dim, dim_ffn, dropout) + self.pos_embedding = ( + None if shared_pos else T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True) + ) + + def forward(self, x, mask=None, pos_bias=None): + e = pos_bias if self.shared_pos else self.pos_embedding(x.size(1), x.size(1)) + x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) + x = fp16_clamp(x + self.ffn(self.norm2(x))) + return x + + +def init_weights(m): + if isinstance(m, T5LayerNorm): + nn.init.ones_(m.weight) + elif isinstance(m, T5FeedForward): + nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5) + nn.init.normal_(m.fc1.weight, std=m.dim**-0.5) + nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5) + elif isinstance(m, T5Attention): + nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn) ** -0.5) + nn.init.normal_(m.k.weight, std=m.dim**-0.5) + nn.init.normal_(m.v.weight, std=m.dim**-0.5) + nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn) ** -0.5) + elif isinstance(m, T5RelativeEmbedding): + nn.init.normal_(m.embedding.weight, std=(2 * m.num_buckets * m.num_heads) ** -0.5) + + +class T5Encoder(nn.Module): + def __init__( + self, vocab, dim, dim_attn, dim_ffn, num_heads, num_layers, num_buckets, shared_pos=True, dropout=0.1 + ): + super().__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_layers = num_layers + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + self.token_embedding = vocab if isinstance(vocab, nn.Embedding) else nn.Embedding(vocab, dim) + self.pos_embedding = ( + T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True) if shared_pos else None + ) + self.dropout = nn.Dropout(dropout) + self.blocks = nn.ModuleList( + [ + T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos, dropout) + for _ in range(num_layers) + ] + ) + self.norm = T5LayerNorm(dim) + self.apply(init_weights) + + def forward(self, ids, mask=None): + x = self.token_embedding(ids) + x = self.dropout(x) + e = self.pos_embedding(x.size(1), x.size(1)) if self.shared_pos else None + for block in self.blocks: + x = block(x, mask, pos_bias=e) + x = self.norm(x) + x = self.dropout(x) + return x + + +def umt5_xxl_encoder(**kwargs) -> T5Encoder: + """ + Official UMT5-XXL encoder-only config. + """ + cfg = dict( + vocab=256384, + dim=4096, + dim_attn=4096, + dim_ffn=10240, + num_heads=64, + num_layers=24, + num_buckets=32, + shared_pos=False, + dropout=0.1, + ) + cfg.update(**kwargs) + model = T5Encoder(**cfg) + model.eval().requires_grad_(False) + return model + + +def umt5_xxl_encoder_from_checkpoint( + checkpoint_path: str, + *, + dtype: torch.dtype = torch.bfloat16, + device: str | torch.device = "cpu", +) -> T5Encoder: + """ + Build UMT5-XXL encoder-only model and load weights without allocating + full parameter tensors upfront. + + This mirrors the official pattern of meta-init + `assign=True` loading. + """ + if not checkpoint_path: + raise ValueError("checkpoint_path is required") + + # Build on meta to avoid allocating a ~1B parameter embedding upfront. + with torch.device("meta"): + model = umt5_xxl_encoder() + + state = torch.load(checkpoint_path, map_location="cpu") + if isinstance(state, dict) and "model" in state and isinstance(state["model"], dict): + state = state["model"] + if not isinstance(state, dict): + raise ValueError(f"Unsupported checkpoint format at {checkpoint_path}") + + # Strip potential DDP prefix. + if any(k.startswith("module.") for k in state.keys()): + state = {k[len("module.") :]: v for k, v in state.items()} + + try: + model.load_state_dict(state, strict=False, assign=True) + except TypeError as exc: + raise RuntimeError( + "Your PyTorch build does not support `assign=True` for loading into meta-initialized modules. " + "Please upgrade torch (>=2.0) or adjust loading strategy." + ) from exc + + model = model.to(device=device, dtype=dtype).eval().requires_grad_(False) + return model diff --git a/primus/backends/diffusion/models/wan/train_pipeline.py b/primus/backends/diffusion/models/wan/train_pipeline.py new file mode 100644 index 000000000..5d6a2dd62 --- /dev/null +++ b/primus/backends/diffusion/models/wan/train_pipeline.py @@ -0,0 +1,282 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F + +from .components import WanComponents + + +@dataclass +class WanFlowMatchTrainPipelineConfig: + """ + Minimal training-pipeline config. + + We intentionally keep this small and derive most behavior from the model config + to maximize YAML reuse. + """ + + # For Wan VAEs, temporal downsample is typically (4, 1): 4n+1 frames. + time_division_factor: int = 4 + time_division_remainder: int = 1 + + +class WanFlowMatchTrainPipeline: + """ + Flow-Matching training pipeline for Wan-style DiT models. + + Contract: + compute_loss(components, batch, scheduler) -> {"loss": Tensor, ...} + """ + + def __init__(self, cfg: Optional[WanFlowMatchTrainPipelineConfig] = None): + self.cfg = cfg or WanFlowMatchTrainPipelineConfig() + # Internal counter for optional per-step profiling logs. + self._profile_step: int = 0 + + @staticmethod + def _encode_prompt(text_encoder: torch.nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor): + # Match existing wan_new behavior: zero-out padded embeddings explicitly. + seq_lens = attention_mask.gt(0).sum(dim=1).long() + prompt_emb = text_encoder(input_ids, attention_mask) + for i, v in enumerate(seq_lens): + prompt_emb[i, v:] = 0 + return prompt_emb + + @staticmethod + def _get_seed_from_env_or_batch(batch: Dict[str, Any]) -> Optional[int]: + # Keep parity with existing scripts/wan_new behavior. + import os + + if os.environ.get("FIXED_SEED"): + try: + return int(os.environ["FIXED_SEED"]) + except ValueError as exc: + raise ValueError(f"Invalid FIXED_SEED value: {os.environ['FIXED_SEED']}") from exc + seed = batch.get("seed", None) + if seed is None: + return None + try: + return int(seed) + except Exception: + return None + + @staticmethod + def _randn_like_on_cpu_then_to( + x: torch.Tensor, *, seed: Optional[int], dtype: torch.dtype, device: torch.device + ) -> torch.Tensor: + generator = None + if seed is not None: + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + noise = torch.randn(x.shape, generator=generator, device="cpu", dtype=torch.float32) + return noise.to(device=device, dtype=dtype) + + @staticmethod + def _vae_encode(vae: torch.nn.Module, videos_bcthw: torch.Tensor) -> torch.Tensor: + # Wan VAEs in this repo use list-of-tensors interface: List[[C,T,H,W]] -> List[[C,F,H',W']] + videos_list = [videos_bcthw[i] for i in range(videos_bcthw.shape[0])] + latents_list = vae.encode(videos_list) + return torch.stack(latents_list) + + @staticmethod + def _get_dit_patch_size(model_config: Any) -> tuple[int, int, int]: + patch_size = tuple(getattr(model_config, "dit_patch_size", (1, 2, 2))) + if len(patch_size) != 3: + raise ValueError(f"Expected 3D dit_patch_size, got {patch_size!r}") + return patch_size + + @staticmethod + def _pad_latents_for_dit( + latents: torch.Tensor, *, patch_size: tuple[int, int, int] + ) -> tuple[torch.Tensor, tuple[int, int]]: + _, _, _, height, width = latents.shape + _, patch_h, patch_w = patch_size + pad_h = (-height) % patch_h + pad_w = (-width) % patch_w + if pad_h == 0 and pad_w == 0: + return latents, (height, width) + + # Pad only the latent-space bottom/right edges so DiT patchify works on any + # VAE output shape. We crop predictions back before computing loss. + latents = F.pad(latents, (0, pad_w, 0, pad_h)) + return latents, (height, width) + + @staticmethod + def _crop_latents(latents: torch.Tensor, *, spatial_size: tuple[int, int]) -> torch.Tensor: + height, width = spatial_size + return latents[..., :height, :width] + + @staticmethod + def _select_timestep(scheduler: Any, device: torch.device) -> torch.Tensor: + """ + Match `wan_new` timestep selection: + - If FIXED_TIMESTEP is set, use that discrete index into [0, num_train_timesteps). + - Else uniform randint over [0, num_train_timesteps). + """ + import os + + if os.environ.get("FIXED_TIMESTEP"): + try: + fixed_step = int(os.environ["FIXED_TIMESTEP"]) + except ValueError as exc: + raise ValueError(f"Invalid FIXED_TIMESTEP value: {os.environ['FIXED_TIMESTEP']}") from exc + max_step = int(scheduler.num_train_timesteps) - 1 + fixed_step = max(0, min(fixed_step, max_step)) + timestep_id = torch.tensor([fixed_step], device=device) + else: + timestep_id = torch.randint(0, int(scheduler.num_train_timesteps), (1,), device=device) + + # scheduler.timesteps live on CPU in this repo; `wan_new` indexes with cpu tensor. + timestep = scheduler.timesteps[timestep_id.cpu()].float() + return timestep.to(device=device) + + @staticmethod + def _maybe_expand_separated_timestep( + *, + timestep: torch.Tensor, + x_list: list[torch.Tensor], + patch_size: tuple[int, int, int], + enabled: bool, + ) -> torch.Tensor: + """ + Match `wan_new.forward_dit` separated-timestep behavior (only used when enabled). + For each sample, build per-token timestep [L] and set first-frame patches to 0. + Returns: + - t: [B] if not enabled + - t: [B, L] if enabled + """ + if not enabled: + if timestep.ndim == 0: + return timestep.unsqueeze(0).repeat(len(x_list)) + if timestep.ndim == 1 and timestep.numel() == 1 and len(x_list) > 1: + return timestep.repeat(len(x_list)) + if timestep.ndim == 1 and timestep.shape[0] != len(x_list): + return timestep.repeat(len(x_list)) + return timestep + + d_f, d_h, d_w = patch_size + if timestep.ndim == 0: + timestep_b = timestep.unsqueeze(0).repeat(len(x_list)) + elif timestep.ndim == 1 and timestep.shape[0] != len(x_list): + timestep_b = timestep.repeat(len(x_list)) + else: + timestep_b = timestep + + t_expand_list: list[torch.Tensor] = [] + for i, x in enumerate(x_list): + # x: [C, F, H, W] in latent space + f, h, w = x.shape[1], x.shape[2], x.shape[3] + seq_len = f * (h // d_h) * (w // d_w) + t_seq = torch.full((seq_len,), timestep_b[i], device=timestep_b.device, dtype=timestep_b.dtype) + spatial_patches = (h // d_h) * (w // d_w) + t_seq[:spatial_patches] = 0 + t_expand_list.append(t_seq) + return torch.stack(t_expand_list) + + def compute_loss( + self, + *, + components: WanComponents, + batch: Dict[str, Any], + scheduler: Any, + model_config: Any, + ) -> Dict[str, torch.Tensor]: + """ + batch requirements (from current dataset/processor): + - video: Tensor [B, C, T, H, W] + - input_ids: Tensor [B, L] + - attention_mask: Tensor [B, L] + """ + video = batch.get("video") + if video is None: + raise ValueError("Batch must contain 'video'") + if not isinstance(video, torch.Tensor) or video.ndim != 5: + raise ValueError( + f"Expected batch['video'] as 5D tensor [B,C,T,H,W], got {type(video)} shape={getattr(video, 'shape', None)}" + ) + + input_ids = batch.get("input_ids") + attention_mask = batch.get("attention_mask") + if input_ids is None or attention_mask is None: + raise ValueError("Batch must contain 'input_ids' and 'attention_mask'") + + device = next(components.dit.parameters()).device + dtype = next(components.dit.parameters()).dtype + + video = video.to(device=device, dtype=dtype, non_blocking=True) + input_ids = input_ids.to(device=device, non_blocking=True) + attention_mask = attention_mask.to(device=device, non_blocking=True) + + # 1) Encode to latents + # Keep VAE/text encoder in eval() by default (consistent with wan_new practice). + components.text_encoder.eval() + try: + components.vae.eval() + except Exception: + pass + + with torch.no_grad(): + input_latents = self._vae_encode(components.vae, video).to(device=device, dtype=dtype) + patch_size = self._get_dit_patch_size(model_config) + input_latents, original_latent_spatial_size = self._pad_latents_for_dit( + input_latents, patch_size=patch_size + ) + + # 2) Sample timestep + noise (match `wan_new`: noise on CPU, timestep from scheduler.timesteps) + timestep = self._select_timestep(scheduler, device=device) # [1] (float) + seed = self._get_seed_from_env_or_batch(batch) + noise = self._randn_like_on_cpu_then_to(input_latents, seed=seed, dtype=dtype, device=device) + + # 3) Diffusion target + noise injection (match `wan_new`) + # Note: in this repo's FlowMatchScheduler, `training_target(sample, noise, t)` uses `noise - sample`. + # `wan_new` passes `input_latents` (not noisy latents). + target = scheduler.training_target(input_latents, noise, timestep) + noisy_latents = scheduler.add_noise(input_latents, noise, timestep=timestep) + + # 4) Text embeddings + with torch.no_grad(): + context = self._encode_prompt(components.text_encoder, input_ids, attention_mask) + + # 5) DiT forward (official interface: List[Tensor] per-sample) + x_list = [noisy_latents[i] for i in range(noisy_latents.shape[0])] + context_list = [context[i] for i in range(context.shape[0])] + + # seq_len computed like wan_new: based on latent shape and patch size + d_f, d_h, d_w = patch_size + max_seq_len = 0 + for x in x_list: + seq_len = x.shape[1] * (x.shape[2] // d_h) * (x.shape[3] // d_w) + max_seq_len = max(max_seq_len, seq_len) + + separated = bool(getattr(model_config, "seperated_timestep", False)) + fuse_flag = bool(getattr(model_config, "fuse_vae_embedding_in_latents", False)) + t = self._maybe_expand_separated_timestep( + timestep=timestep, + x_list=x_list, + patch_size=(d_f, d_h, d_w), + enabled=bool(separated and fuse_flag), + ) + + sp_group = batch.get("sp_group", None) + noise_pred_list = components.dit( + x=x_list, t=t, context=context_list, seq_len=max_seq_len, y=None, sp_group=sp_group + ) + noise_pred = torch.stack(noise_pred_list) + noise_pred = self._crop_latents(noise_pred, spatial_size=original_latent_spatial_size) + target = self._crop_latents(target, spatial_size=original_latent_spatial_size) + + # 6) Loss + loss = F.mse_loss(noise_pred.float(), target.float(), reduction="mean") + # `wan_new` uses scalar `timestep` for weighting (not expanded per-token) + weight = scheduler.training_weight(timestep) + loss = loss * weight + return {"loss": loss} diff --git a/primus/backends/diffusion/models/wan/vae2_1.py b/primus/backends/diffusion/models/wan/vae2_1.py new file mode 100644 index 000000000..fa0b2b78a --- /dev/null +++ b/primus/backends/diffusion/models/wan/vae2_1.py @@ -0,0 +1,638 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Wan2.1 VAE (vendored from official Wan2.2 repo). + +Source: `Wan2.2/wan/modules/vae2_1.py` + +Minimal changes: +- add small `to()/eval()/train()` helpers on the wrapper to fit Wan trainer usage + (HF trainer sometimes calls `model.vae.to(dtype=...)`). +""" + +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import logging +from typing import List + +import torch +import torch.cuda.amp as amp +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +__all__ = [ + "Wan2_1_VAE", + "WanVAE", +] + +CACHE_T = 2 + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = ( + self.padding[2], + self.padding[2], + self.padding[1], + self.padding[1], + 2 * self.padding[0], + 0, + ) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + return super().forward(x) + + +class RMS_norm(nn.Module): + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + def __init__(self, dim, mode): + assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d") + super().__init__() + self.dim = dim + self.mode = mode + + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim // 2, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim // 2, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1, c2) + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[: c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2 :, :, -1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + +class ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), + nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1), + ) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm(x) + q, k, v = ( + self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1) + ) + x = F.scaled_dot_product_attention(q, k, v) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + x = self.proj(x) + x = rearrange(x, "(b t) c h w-> b c t h w", t=t) + return x + identity + + +class Encoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + if i != len(dim_mult) - 1: + mode = "downsample3d" if temperal_downsample[i] else "downsample2d" + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Decoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + self.middle = nn.Sequential( + ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout), + ) + + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + if i == 1 or i == 2 or i == 3: + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + if i != len(dim_mult) - 1: + mode = "upsample3d" if temperal_upsample[i] else "upsample2d" + upsamples.append(Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, 3, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class WanVAE_(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + self.encoder = Encoder3d( + dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout + ) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d( + dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout + ) + + def forward(self, x): + mu, log_var = self.encode(x) + z = self.reparameterize(mu, log_var) + x_recon = self.decode(z) + return x_recon, mu, log_var + + def encode(self, x, scale): + self.clear_cache() + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder( + x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx + ) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1) + else: + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + def decode(self, z, scale): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1) + else: + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx + ) + else: + out_ = self.decoder( + x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx + ) + out = torch.cat([out, out_], 2) + self.clear_cache() + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +def _video_vae(pretrained_path=None, z_dim=None, device="cpu", **kwargs): + cfg = dict( + dim=96, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0, + ) + cfg.update(**kwargs) + + with torch.device("meta"): + model = WanVAE_(**cfg) + + logging.info(f"loading {pretrained_path}") + model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True) + return model + + +class Wan2_1_VAE: + def __init__(self, z_dim=16, vae_pth="cache/vae_step_411000.pth", dtype=torch.float, device="cuda"): + self.dtype = dtype + self.device = device + + mean = [ + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, + ] + std = [ + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, + ] + self.mean = torch.tensor(mean, dtype=dtype, device=device) + self.std = torch.tensor(std, dtype=dtype, device=device) + self.scale = [self.mean, 1.0 / self.std] + + if vae_pth is None: + raise ValueError("Wan2_1_VAE requires `vae_pth` (checkpoint path).") + self.model = _video_vae(pretrained_path=vae_pth, z_dim=z_dim).eval().requires_grad_(False).to(device) + + # --- Wan trainer compatibility helpers (non-upstream) --- + def to(self, *args, **kwargs): + self.model.to(*args, **kwargs) + if "dtype" in kwargs and kwargs["dtype"] is not None: + self.dtype = kwargs["dtype"] + device = kwargs.get("device", None) + if device is None and len(args) > 0: + device = args[0] + if device is not None: + self.device = device + self.mean = self.mean.to(device) + self.std = self.std.to(device) + self.scale = [self.mean, 1.0 / self.std] + return self + + def eval(self): + self.model.eval() + return self + + def train(self, mode: bool = True): + self.model.train(mode) + return self + + def encode(self, videos: List[torch.Tensor]): + """ + videos: A list of videos each with shape [C, T, H, W]. + """ + with amp.autocast(dtype=self.dtype): + # Avoid forcing fp32 outputs: it increases memory and slows training. + return [self.model.encode(u.unsqueeze(0), self.scale).squeeze(0) for u in videos] + + def decode(self, zs: List[torch.Tensor]): + with amp.autocast(dtype=self.dtype): + return [self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0) for u in zs] + + +# Backward-compat alias (older Wan naming) +WanVAE = Wan2_1_VAE diff --git a/primus/backends/diffusion/models/wan/vae2_2.py b/primus/backends/diffusion/models/wan/vae2_2.py new file mode 100644 index 000000000..47eab810a --- /dev/null +++ b/primus/backends/diffusion/models/wan/vae2_2.py @@ -0,0 +1,855 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Wan2.2 VAE (vendored from official Wan2.2 repo). + +Source: `Wan2.2/wan/modules/vae2_2.py` + +Minimal changes: +- add small `to()/eval()/train()` helpers on the wrapper to fit Wan trainer usage + (HF trainer sometimes calls `model.vae.to(dtype=...)`). +""" + +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import logging +from typing import List + +import torch +import torch.cuda.amp as amp +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +__all__ = [ + "Wan2_2_VAE", +] + +CACHE_T = 2 + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = ( + self.padding[2], + self.padding[2], + self.padding[1], + self.padding[1], + 2 * self.padding[0], + 0, + ) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + return super().forward(x) + + +class RMS_norm(nn.Module): + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + def __init__(self, dim, mode): + assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d") + super().__init__() + self.dim = dim + self.mode = mode + + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight.detach().clone() + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1, c2) + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight = nn.Parameter(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data.detach().clone() + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[: c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2 :, :, -1, 0, 0] = init_matrix + conv.weight = nn.Parameter(conv_weight) + nn.init.zeros_(conv.bias.data) + + +class ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), + nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1), + ) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm(x) + q, k, v = ( + self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1) + ) + x = F.scaled_dot_product_attention(q, k, v) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + x = self.proj(x) + x = rearrange(x, "(b t) c h w-> b c t h w", t=t) + return x + identity + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + return x + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size) + return x + + +class AvgDown3D(nn.Module): + def __init__(self, in_channels, out_channels, factor_t, factor_s=1): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + b, c, t, h, w = x.shape + x = x.view( + b, + c, + t // self.factor_t, + self.factor_t, + h // self.factor_s, + self.factor_s, + w // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(b, c * self.factor, t // self.factor_t, h // self.factor_s, w // self.factor_s) + x = x.view( + b, self.out_channels, self.group_size, t // self.factor_t, h // self.factor_s, w // self.factor_s + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + def __init__(self, in_channels: int, out_channels: int, factor_t, factor_s=1): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False): + super().__init__() + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample(out_dim, mode=mode)) + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + return x + self.avg_shortcut(x_copy) + + +class Up_ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False): + super().__init__() + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2 if up_flag else 1, + ) + else: + self.avg_shortcut = None + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + if up_flag: + mode = "upsample3d" if temperal_upsample else "upsample2d" + upsamples.append(Resample(out_dim, mode=mode)) + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + x_main = x.clone() + for module in self.upsamples: + x_main = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + x_shortcut = self.avg_shortcut(x, first_chunk) + return x_main + x_shortcut + return x_main + + +class Encoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + dims = [dim * u for u in [1] + dim_mult] + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = temperal_downsample[i] if i < len(temperal_downsample) else False + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(dim_mult) - 1, + ) + ) + self.downsamples = nn.Sequential(*downsamples) + + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Decoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + ): + super().__init__() + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + self.middle = nn.Sequential( + ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout), + ) + + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False + upsamples.append( + Up_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temperal_upsample=t_up_flag, + up_flag=i != len(dim_mult) - 1, + ) + ) + self.upsamples = nn.Sequential(*upsamples) + + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class WanVAE_(nn.Module): + # NOTE: Keep this identical to upstream for numerical parity. + def __init__( + self, + dim=160, + dec_dim=256, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + self.encoder = Encoder3d( + dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout + ) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d( + dec_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout + ) + + def forward(self, x, scale=[0, 1]): + mu = self.encode(x, scale) + x_recon = self.decode(mu, scale) + return x_recon, mu + + def encode(self, x, scale): + self.clear_cache() + x = patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder( + x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx + ) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1) + else: + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + def decode(self, z, scale): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1) + else: + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True, + ) + else: + out_ = self.decoder( + x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx + ) + out = torch.cat([out, out_], 2) + out = unpatchify(out, patch_size=2) + self.clear_cache() + return out + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs): + cfg = dict( + dim=dim, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, True], + dropout=0.0, + ) + cfg.update(**kwargs) + with torch.device("meta"): + model = WanVAE_(**cfg) + logging.info(f"loading {pretrained_path}") + model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True) + return model + + +class Wan2_2_VAE: + def __init__( + self, + z_dim=48, + c_dim=160, + vae_pth=None, + dim_mult=[1, 2, 4, 4], + temperal_downsample=[False, True, True], + dtype=torch.float, + device="cuda", + ): + self.dtype = dtype + self.device = device + mean = torch.tensor( + [ + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.1570, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.1230, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.0520, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, + ], + dtype=dtype, + device=device, + ) + std = torch.tensor( + [ + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.4990, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.0600, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, + ], + dtype=dtype, + device=device, + ) + self.scale = [mean, 1.0 / std] + + if vae_pth is None: + raise ValueError("Wan2_2_VAE requires `vae_pth` (checkpoint path).") + self.model = ( + _video_vae( + pretrained_path=vae_pth, + z_dim=z_dim, + dim=c_dim, + dim_mult=dim_mult, + temperal_downsample=temperal_downsample, + ) + .eval() + .requires_grad_(False) + .to(device) + ) + + # --- Wan trainer compatibility helpers (non-upstream) --- + def to(self, *args, **kwargs): + self.model.to(*args, **kwargs) + if "dtype" in kwargs and kwargs["dtype"] is not None: + self.dtype = kwargs["dtype"] + device = kwargs.get("device", None) + if device is None and len(args) > 0: + device = args[0] + if device is not None: + self.device = device + dev = next(self.model.parameters()).device + self.scale = [t.to(device=dev) for t in self.scale] + return self + + def eval(self): + self.model.eval() + return self + + def train(self, mode: bool = True): + self.model.train(mode) + return self + + def encode(self, videos: List[torch.Tensor]): + try: + if not isinstance(videos, list): + raise TypeError("videos should be a list") + with amp.autocast(dtype=self.dtype): + # Avoid forcing fp32 outputs: it increases memory and slows training. + # Keep outputs in the model/autocast dtype (typically bf16). + return [self.model.encode(u.unsqueeze(0), self.scale).squeeze(0) for u in videos] + except TypeError as e: + logging.info(e) + return None + + def decode(self, zs: List[torch.Tensor]): + try: + if not isinstance(zs, list): + raise TypeError("zs should be a list") + with amp.autocast(dtype=self.dtype): + return [self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0) for u in zs] + except TypeError as e: + logging.info(e) + return None diff --git a/primus/backends/diffusion/models/wan/wan_dit.py b/primus/backends/diffusion/models/wan/wan_dit.py new file mode 100644 index 000000000..98b878b78 --- /dev/null +++ b/primus/backends/diffusion/models/wan/wan_dit.py @@ -0,0 +1,555 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +# +# Vendored from `Wan2.2/wan/modules/model.py` with minimal changes: +# - remove diffusers dependencies (ModelMixin/ConfigMixin/register_to_config) +# - keep pure PyTorch modeling as close as possible +# - provide a `DiTBlock` alias class so existing Wan configs can keep +# `fsdp_transformer_layer_cls_to_wrap: "DiTBlock"` + +from __future__ import annotations + +import math +from typing import Optional + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint + +from primus.backends.diffusion.attention import attention +from primus.backends.diffusion.distributed.ulysses import ( + distributed_attention, + sp_gather, + sp_split, + sp_unpad, +) + +__all__ = ["WanModel", "DiTBlock"] + + +def sinusoidal_embedding_1d(dim, position): + # preprocess + assert dim % 2 == 0 + half = dim // 2 + position = position.type(torch.float64) + + # calculation + sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +@torch.amp.autocast("cuda", enabled=False) +def rope_params(max_seq_len, dim, theta=10000): + assert dim % 2 == 0 + freqs = torch.outer( + torch.arange(max_seq_len), + 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)), + ) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + +@torch.amp.autocast("cuda", enabled=False) +def rope_apply(x, grid_sizes, freqs, sp_group=None): + """ + Apply 3-D RoPE. When *sp_group* is given the input is assumed to hold + only this rank's local token chunk (S/P) and the correct positional + frequencies are selected automatically (following Wan2.2 official). + """ + n, c = x.size(2), x.size(3) // 2 + s = x.size(1) # local token count (= full seq when no SP) + + if sp_group is not None: + sp_size = dist.get_world_size(sp_group) + sp_rank = dist.get_rank(sp_group) + else: + sp_size, sp_rank = 1, 0 + + # split freqs + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + seq_len = f * h * w + + # tokens to process: local chunk when SP, valid tokens when single-GPU + t = s if sp_size > 1 else seq_len + + # precompute multipliers + x_i = torch.view_as_complex(x[i, :t].to(torch.float64).reshape(t, n, -1, 2)) + freqs_i = torch.cat( + [ + freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + + # SP: pad freqs to padded-full-length, select this rank's range + if sp_size > 1: + full_len = s * sp_size + if seq_len < full_len: + freqs_i = torch.cat( + [ + freqs_i, + torch.ones( + full_len - seq_len, + 1, + freqs_i.size(-1), + dtype=freqs_i.dtype, + device=freqs_i.device, + ), + ], + dim=0, + ) + freqs_i = freqs_i[sp_rank * s : (sp_rank + 1) * s] + + # apply rotary embedding + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, t:]]) + + # append to collection + output.append(x_i) + return torch.stack(output).to(x.dtype) + + +class WanRMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class WanLayerNorm(nn.LayerNorm): + def __init__(self, dim, eps=1e-6, elementwise_affine=False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + # Keep numerical stability by normalizing in fp32, but do NOT rely on + # module parameters being fp32 (FSDP/bf16 can cast weights/bias). + weight = self.weight.float() if self.weight is not None else None + bias = self.bias.float() if self.bias is not None else None + y = F.layer_norm(x.float(), self.normalized_shape, weight, bias, self.eps) + return y.type_as(x) + + +class WanSelfAttention(nn.Module): + def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, seq_lens, grid_sizes, freqs, sp_group: Optional[dist.ProcessGroup] = None): + r""" + Args: + x(Tensor): Shape [B, L, C] (L = S/P when SP enabled) + seq_lens(Tensor): Shape [B] — valid lengths in the *full* sequence + grid_sizes(Tensor): Shape [B, 3], (F, H, W) — full spatial grid + freqs(Tensor): Rope freqs + sp_group: Ulysses SP process group (None = no SP) + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + x_ = x.to(dtype=self.q.weight.dtype) + q = self.norm_q(self.q(x_)).view(b, s, n, d) + k = self.norm_k(self.k(x_)).view(b, s, n, d) + v = self.v(x_).view(b, s, n, d) + + # RoPE on local tokens (SP-aware: uses rank-offset freqs when sp_group is set) + q = rope_apply(q, grid_sizes, freqs, sp_group=sp_group) + k = rope_apply(k, grid_sizes, freqs, sp_group=sp_group) + + attn_kwargs = { + "q_lens": seq_lens, + "k_lens": seq_lens, + "window_size": self.window_size, + "dtype": self.q.weight.dtype, + } + if sp_group is not None: + x = distributed_attention(q, k, v, group=sp_group, attention_fn=attention, **attn_kwargs) + else: + x = attention(q=q, k=k, v=v, **attn_kwargs) + + x = x.to(dtype=self.o.weight.dtype) + x = x.flatten(2) + x = self.o(x) + return x + + +class WanCrossAttention(WanSelfAttention): + def forward(self, x, context, context_lens): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + x = x.to(dtype=self.q.weight.dtype) + context = context.to(dtype=self.k.weight.dtype) + + q = self.norm_q(self.q(x)).view(b, -1, n, d) + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + + x = attention(q, k, v, k_lens=context_lens, dtype=self.q.weight.dtype) + + x = x.to(dtype=self.o.weight.dtype) + x = x.flatten(2) + x = self.o(x) + return x + + +class WanAttentionBlock(nn.Module): + def __init__( + self, + dim, + ffn_dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + ): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps) + self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim) + ) + + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + e, + seq_lens, + grid_sizes, + freqs, + context, + context_lens, + sp_group: Optional[dist.ProcessGroup] = None, + ): + # Memory-critical: + # `e` can be very large ([B, seq_len, 6, dim]). Keeping it in fp32 and + # doing broadcast add in fp32 easily OOMs for Wan2.1 (more tokens). + # Align to x.dtype (bf16/fp16) like `wan_new` implementation. + e = e.to(dtype=x.dtype) + modulation = self.modulation.to(dtype=x.dtype) + e = (modulation.unsqueeze(0) + e).chunk(6, dim=2) + + # self-attention (Ulysses all-to-all happens inside self_attn) + y = self.self_attn( + self.norm1(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2), + seq_lens, + grid_sizes, + freqs, + sp_group=sp_group, + ) + x = x + y * e[2].squeeze(2) + + # cross-attention + FFN (no SP communication needed: + # each rank's visual queries attend to full text keys independently) + def cross_attn_ffn(x_, context_, context_lens_, e_): + x_ = x_ + self.cross_attn(self.norm3(x_), context_, context_lens_) + ffn_in = self.norm2(x_) * (1 + e_[4].squeeze(2)) + e_[3].squeeze(2) + y_ = self.ffn(ffn_in) + x_ = x_ + y_ * e_[5].squeeze(2) + return x_ + + x = cross_attn_ffn(x, context, context_lens, e) + return x + + +# Wan config compatibility: FSDP auto-wrap often uses "DiTBlock". +class DiTBlock(WanAttentionBlock): + pass + + +class Head(nn.Module): + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + # Align modulation embedding dtype to x (saves memory vs fp32 broadcast). + e = e.to(dtype=x.dtype) + modulation = self.modulation.to(dtype=x.dtype) + e = (modulation.unsqueeze(0) + e.unsqueeze(2)).chunk(2, dim=2) + x = self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2)) + return x + + +class WanModel(nn.Module): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + def __init__( + self, + model_type="t2v", + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6, + ): + super().__init__() + + assert model_type in ["t2v", "i2v", "ti2v", "s2v"] + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + # Used by the Wan trainer: it sets `model.dit.gradient_checkpointing = True` + # when `trainer_args.gradient_checkpointing: true`. + self.gradient_checkpointing = False + + # embeddings + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim) + ) + self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks (use DiTBlock for YAML compatibility) + self.blocks = nn.ModuleList( + [ + DiTBlock(dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps) + for _ in range(num_layers) + ] + ) + + # head + self.head = Head(dim, out_dim, patch_size, eps) + + # rope buffers (avoid register_buffer to keep dtype stable under .to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + self.freqs = torch.cat( + [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + ], + dim=1, + ) + + self.init_weights() + + def forward(self, x, t, context, seq_len, y=None, sp_group: Optional[dist.ProcessGroup] = None): + if self.model_type == "i2v": + assert y is not None + + device = self.patch_embedding.weight.device + if self.freqs.device != device: + self.freqs = self.freqs.to(device) + + if y is not None: + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + assert seq_lens.max() <= seq_len + x = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x]) + + # time embeddings + if t.dim() == 1: + t = t.expand(t.size(0), seq_len) + with torch.amp.autocast("cuda", dtype=torch.float32): + bt = t.size(0) + t_flat = t.flatten() + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, t_flat).unflatten(0, (bt, seq_len)).float() + ) + e0 = self.time_projection(e).unflatten(2, (6, self.dim)) + assert e.dtype == torch.float32 and e0.dtype == torch.float32 + + # Cast large modulation tensors back to model dtype early to reduce peak memory. + e = e.to(dtype=x.dtype) + e0 = e0.to(dtype=x.dtype) + + # --- Ulysses SP: split sequence-dim tensors across SP ranks --- + original_seq_len = None + if sp_group is not None: + (x, e, e0), original_seq_len = sp_split([x, e, e0], dim=1, group=sp_group) + + # context (text embeddings — NOT sliced, full text on every SP rank) + context_lens = None + context_in = torch.stack( + [torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context] + ) + context_in = context_in.to(dtype=self.text_embedding[0].weight.dtype) + context = self.text_embedding(context_in) + + freqs = self.freqs + seq_lens_ = seq_lens + grid_sizes_ = grid_sizes + e0_ = e0 + context_ = context + context_lens_ = context_lens + + for block in self.blocks: + if self.training and getattr(self, "gradient_checkpointing", False): + + def create_custom_forward(module, _sp_group): + def custom_forward(x_in, e_in, ctx_in): + return module( + x_in, + e=e_in, + seq_lens=seq_lens_, + grid_sizes=grid_sizes_, + freqs=freqs, + context=ctx_in, + context_lens=context_lens_, + sp_group=_sp_group, + ) + + return custom_forward + + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block, sp_group), + x, + e0_, + context_, + use_reentrant=False, + ) + else: + x = block( + x, + e=e0_, + seq_lens=seq_lens_, + grid_sizes=grid_sizes_, + freqs=freqs, + context=context_, + context_lens=context_lens_, + sp_group=sp_group, + ) + + # head operates on local tokens (like Wan2.2 official) + x = self.head(x, e) + + # --- Ulysses SP: gather output back to full sequence --- + if sp_group is not None: + x = sp_unpad(sp_gather(x, dim=1, group=sp_group), dim=1, original_size=original_seq_len) + + x = self.unpatchify(x, grid_sizes) + # IMPORTANT: + # Returning fp32 here dramatically increases activation/grad memory and slows + # training (DiffSynth returns bf16 here). Keep dtype consistent with model. + return x + + def unpatchify(self, x, grid_sizes): + c = self.out_dim + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[: math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum("fhwpqrc->cfphqwr", u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + def init_weights(self): + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/primus/backends/diffusion/optim/adamw_fp32_state.py b/primus/backends/diffusion/optim/adamw_fp32_state.py new file mode 100644 index 000000000..1403043f2 --- /dev/null +++ b/primus/backends/diffusion/optim/adamw_fp32_state.py @@ -0,0 +1,111 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +AdamW that keeps optimizer state in FP32 even when parameters are BF16/FP16. + +Why: +- DeepSpeed bf16 commonly keeps FP32 master weights / FP32 optimizer states. +- Vanilla torch.optim.AdamW will create exp_avg/exp_avg_sq with the same dtype as the parameter, + so bf16 params -> bf16 states, which changes early-step behavior and can destabilize. + +This optimizer: +- Stores exp_avg/exp_avg_sq in FP32 +- Applies the AdamW update in FP32 on a FP32 view of the param +- Writes the updated value back to the original parameter dtype + +Scope: +- Intended for this repo's bf16 diffusion training (single GPU or FSDP world_size=1). +- Supports common AdamW args: lr, betas, eps, weight_decay. +""" + +from __future__ import annotations + +from typing import Iterable, Optional + +import torch + + +class AdamWFP32State(torch.optim.Optimizer): + def __init__( + self, + params: Iterable[torch.nn.Parameter], + lr: float = 1e-3, + betas: tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0.0, + ): + if lr < 0.0: + raise ValueError(f"Invalid lr: {lr}") + if eps < 0.0: + raise ValueError(f"Invalid eps: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta1: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta2: {betas[1]}") + if weight_decay < 0.0: + raise ValueError(f"Invalid weight_decay: {weight_decay}") + + defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) + super().__init__(params, defaults) + + @torch.no_grad() + def step(self, closure: Optional[callable] = None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + lr: float = group["lr"] + beta1, beta2 = group["betas"] + eps: float = group["eps"] + weight_decay: float = group["weight_decay"] + + for p in group["params"]: + if p.grad is None: + continue + if p.grad.is_sparse: + raise RuntimeError("AdamWFP32State does not support sparse gradients") + + grad = p.grad + state = self.state[p] + + # Initialize master weights in state if not present + if len(state) == 0: + state["step"] = 0 + state["exp_avg"] = torch.zeros_like(grad, dtype=torch.float32) + state["exp_avg_sq"] = torch.zeros_like(grad, dtype=torch.float32) + state["master_param"] = p.detach().clone().float() + + # Always use the persistent master weight for updates + p_fp32 = state["master_param"] + + exp_avg: torch.Tensor = state["exp_avg"] + exp_avg_sq: torch.Tensor = state["exp_avg_sq"] + state["step"] += 1 + step: int = state["step"] + + # Decoupled weight decay (AdamW) + if weight_decay != 0.0: + p_fp32.add_(p_fp32, alpha=-lr * weight_decay) + + # Adam moments + exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) + + bias_correction1 = 1.0 - beta1**step + bias_correction2 = 1.0 - beta2**step + + denom = (exp_avg_sq.sqrt() / (bias_correction2**0.5)).add_(eps) + step_size = lr / bias_correction1 + + p_fp32.addcdiv_(exp_avg, denom, value=-step_size) + + # Write back to original dtype + p.copy_(p_fp32.to(dtype=p.dtype)) + + return loss diff --git a/primus/backends/diffusion/registry.py b/primus/backends/diffusion/registry.py new file mode 100644 index 000000000..b0a728827 --- /dev/null +++ b/primus/backends/diffusion/registry.py @@ -0,0 +1,66 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Explicit factory registry for diffusion model, dataset, and trainer builders.""" + +from __future__ import annotations + +from typing import Any, Callable, Dict, Tuple + + +def _build_wan_model(model_config: dict): + from primus.backends.diffusion.models.registrations.wan import build_wan_model + + return build_wan_model(model_config) + + +def _build_wan_dataset(dataset_config: dict): + from primus.backends.diffusion.data.registrations.wan import build_wan_dataset + + return build_wan_dataset(dataset_config) + + +def _build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict): + from primus.backends.diffusion.trainers.fsdp2 import build_fsdp2_trainer + + return build_fsdp2_trainer( + model=model, + dataset=dataset, + processor=processor, + trainer_args=trainer_args, + ) + + +MODEL_BUILDERS: Dict[str, Callable[[dict], Any]] = { + "wan": _build_wan_model, +} +DATASET_BUILDERS: Dict[str, Callable[[dict], Tuple[Any, Any]]] = { + "wan": _build_wan_dataset, +} +TRAINER_BUILDERS: Dict[str, Callable[..., Any]] = { + "fsdp2": _build_fsdp2_trainer, +} + + +def get_model_builder(name: str) -> Callable[[dict], Any]: + try: + return MODEL_BUILDERS[name] + except KeyError: + raise KeyError(f"Unknown model name: {name}") + + +def get_dataset_builder(name: str) -> Callable[[dict], Tuple[Any, Any]]: + try: + return DATASET_BUILDERS[name] + except KeyError: + raise KeyError(f"Unknown dataset name: {name}") + + +def get_trainer_builder(name: str) -> Callable[..., Any]: + try: + return TRAINER_BUILDERS[name] + except KeyError: + raise KeyError(f"Unknown trainer name: {name}") diff --git a/primus/backends/diffusion/schedulers/flow_match.py b/primus/backends/diffusion/schedulers/flow_match.py new file mode 100644 index 000000000..ec5519574 --- /dev/null +++ b/primus/backends/diffusion/schedulers/flow_match.py @@ -0,0 +1,130 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import math + +import torch + + +class FlowMatchScheduler: + def __init__( + self, + num_inference_steps=100, + num_train_timesteps=1000, + shift=3.0, + sigma_max=1.0, + sigma_min=0.003 / 1.002, + inverse_timesteps=False, + extra_one_step=False, + reverse_sigmas=False, + exponential_shift=False, + exponential_shift_mu=None, + shift_terminal=None, + ): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.exponential_shift = exponential_shift + self.exponential_shift_mu = exponential_shift_mu + self.shift_terminal = shift_terminal + self.set_timesteps(num_inference_steps) + + def set_timesteps( + self, + num_inference_steps=100, + denoising_strength=1.0, + training=False, + shift=None, + dynamic_shift_len=None, + ): + if shift is not None: + self.shift = shift + sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + if self.exponential_shift: + mu = ( + self.calculate_shift(dynamic_shift_len) + if dynamic_shift_len is not None + else self.exponential_shift_mu + ) + self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1)) + else: + self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) + if self.shift_terminal is not None: + one_minus_z = 1 - self.sigmas + scale_factor = one_minus_z[-1] / (1 - self.shift_terminal) + self.sigmas = 1 - (one_minus_z / scale_factor) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + self.training = True + else: + self.training = False + + def step(self, model_output, timestep, sample, to_final=False, **kwargs): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + if to_final or timestep_id + 1 >= len(self.timesteps): + sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 + else: + sigma_ = self.sigmas[timestep_id + 1] + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample + + def return_to_timestep(self, timestep, sample, sample_stablized): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + model_output = (sample - sample_stablized) / sigma + return model_output + + def add_noise(self, original_samples, noise, timestep): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + sample = (1 - sigma) * original_samples + sigma * noise + return sample + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + def training_weight(self, timestep): + timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs()) + weights = self.linear_timesteps_weights[timestep_id] + return weights + + def calculate_shift( + self, + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 8192, + base_shift: float = 0.5, + max_shift: float = 0.9, + ): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu diff --git a/primus/backends/diffusion/trainers/__init__.py b/primus/backends/diffusion/trainers/__init__.py new file mode 100644 index 000000000..d4df6355e --- /dev/null +++ b/primus/backends/diffusion/trainers/__init__.py @@ -0,0 +1,11 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Trainer registrations for the Primus Wan backend.""" + +from .fsdp2 import build_fsdp2_trainer + +__all__ = ["build_fsdp2_trainer"] diff --git a/primus/backends/diffusion/trainers/base.py b/primus/backends/diffusion/trainers/base.py new file mode 100644 index 000000000..d7ffd25ce --- /dev/null +++ b/primus/backends/diffusion/trainers/base.py @@ -0,0 +1,539 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Base trainer with shared logic for Wan PyTorch trainers. + +This module holds functionality shared across FSDP-style Wan trainers: +config parsing, optimizer creation, LR scheduling, training loop, logging, +and W&B integration. Concrete trainers (e.g. FSDP2Trainer) subclass it. +""" + +from __future__ import annotations + +import math +import os +import time +from contextlib import contextmanager + +import torch + +from primus.backends.diffusion.optim.adamw_fp32_state import AdamWFP32State +from primus.backends.diffusion.schedulers.flow_match import FlowMatchScheduler +from primus.backends.diffusion.utils.log import logger +from primus.backends.diffusion.utils.train_utils import ( + get_memory, + resolve_dtype, + set_seed, +) + +try: + import wandb +except ImportError: + wandb = None + + +def create_lr_scheduler(optimizer, scheduler_type, warmup_steps, total_steps): + """ + Create LR scheduler with warmup support. + + Supports: constant, constant_with_warmup, linear, cosine, polynomial. + Shared between FSDP and FSDP2 trainers for consistency. + """ + if total_steps <= 0: + return torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: 1.0) + + def linear_warmup(step): + if warmup_steps == 0: + return 1.0 + return min(1.0, float(step) / float(max(1, warmup_steps))) + + def linear_decay(step): + if step <= warmup_steps: + return linear_warmup(step) + progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps)) + return max(0.0, 1.0 - progress) + + def cosine_decay(step): + if step <= warmup_steps: + return linear_warmup(step) + progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps)) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + def constant_with_warmup(step): + return linear_warmup(step) + + def polynomial_decay(step, power=1.0): + if step <= warmup_steps: + return linear_warmup(step) + progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps)) + return max(0.0, (1.0 - progress) ** power) + + scheduler_type = (scheduler_type or "constant").lower() + lambdas = { + "constant": lambda step: 1.0, + "constant_with_warmup": constant_with_warmup, + "linear": linear_decay, + "cosine": cosine_decay, + "cosine_with_restarts": cosine_decay, + "polynomial": polynomial_decay, + } + lr_lambda = lambdas.get(scheduler_type) + if lr_lambda is None: + logger.warning(f"Unknown lr_scheduler_type={scheduler_type}, falling back to constant.") + + def lr_lambda(step): + return 1.0 + + return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) + + +class BaseWanTrainer: + """ + Shared base class for Wan PyTorch FSDP-style trainers. + + Subclasses must implement: + - _apply_parallelism(): set up distributed wrapping / sharding + - save_model(): save final model + + Subclasses may override: + - _grad_sync_context(is_update_step): gradient sync during accumulation + - _clip_grad_norm(): gradient clipping strategy + - _save_checkpoint(): periodic checkpoint saving + """ + + def __init__( + self, + model: torch.nn.Module, + args: dict, + train_dataset, + data_collator, + processing_class, + rank: int, + world_size: int, + local_rank: int, + ): + self.model = model + self.args = args + self.rank = rank + self.world_size = world_size + self.local_rank = local_rank + self.device = torch.device(f"cuda:{local_rank}") + + # --- Config extraction --- + self.output_dir = self.args.get("output_dir", "./output") + self.logging_steps = int(self.args.get("logging_steps", 1)) + self.save_steps = int(self.args.get("save_steps", 0)) + self.max_steps = int(self.args.get("max_steps", -1) if self.args.get("max_steps") is not None else -1) + self.grad_accum_steps = int(self.args.get("gradient_accumulation_steps", 1)) + self.max_grad_norm = float(self.args.get("max_grad_norm", 1.0)) + self.num_train_epochs = int(self.args.get("num_train_epochs", 1)) + + if self.rank == 0: + os.makedirs(self.output_dir, exist_ok=True) + + # --- Seeding --- + seed = self.args.get("seed") + if seed is not None: + set_seed(int(seed)) + if os.environ.get("FIXED_SEED"): + set_seed(int(os.environ["FIXED_SEED"])) + + # --- Gradient Checkpointing --- + if self.args.get("gradient_checkpointing", False): + if hasattr(self.model, "gradient_checkpointing_enable"): + self.model.gradient_checkpointing_enable() + elif hasattr(self.model, "dit") and hasattr(self.model.dit, "gradient_checkpointing"): + self.model.dit.gradient_checkpointing = True + if self.rank == 0: + logger.info("Gradient checkpointing enabled") + + # --- W&B --- + self._setup_wandb() + + # --- Parallelism (subclass hook) --- + # Subclass sets self.sp_group (Ulysses SP group) if SP is enabled. + self.sp_group = None + self._apply_parallelism() + + # --- DataLoader --- + self.train_dataset = train_dataset + self.processing_class = processing_class + self.data_collator = data_collator + + # When SP is enabled, all ranks in the same SP group process the same sample. + # DistributedSampler should use DP-only rank/size so SP peers get identical data. + self.sp_size = 1 + dp_world_size = world_size + dp_rank = rank + if self.sp_group is not None: + import torch.distributed as dist + + self.sp_size = dist.get_world_size(self.sp_group) + dp_world_size = world_size // self.sp_size + dp_rank = rank // self.sp_size + + self.data_parallel_world_size = dp_world_size + self.per_device_train_batch_size = int(self.args.get("per_device_train_batch_size", 1)) + + self.sampler = torch.utils.data.distributed.DistributedSampler( + train_dataset, + num_replicas=dp_world_size, + rank=dp_rank, + shuffle=self.args.get("shuffle", True), + ) + + num_workers = int(self.args.get("dataloader_num_workers", 4) or 0) + self.dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_size=self.per_device_train_batch_size, + sampler=self.sampler, + num_workers=num_workers, + collate_fn=data_collator, + pin_memory=True, + persistent_workers=num_workers > 0, + prefetch_factor=2 if num_workers > 0 else None, + ) + + # --- Optimizer --- + self.optimizer = self._create_optimizer() + + # --- LR Scheduler --- + steps_per_epoch = math.ceil(len(self.dataloader) / max(1, self.grad_accum_steps)) + self.total_steps = self.max_steps if self.max_steps > 0 else self.num_train_epochs * steps_per_epoch + self.lr_scheduler = create_lr_scheduler( + self.optimizer, + self.args.get("lr_scheduler_type", "constant"), + int(self.args.get("warmup_steps", 0)), + self.total_steps, + ) + + # --- Diffusion Scheduler (configurable from YAML) --- + scheduler_cfg = self.args.get("flow_match_scheduler", {}) or {} + self.scheduler = FlowMatchScheduler( + shift=float(scheduler_cfg.get("shift", 5)), + sigma_min=float(scheduler_cfg.get("sigma_min", 0.0)), + extra_one_step=bool(scheduler_cfg.get("extra_one_step", True)), + ) + self.scheduler.set_timesteps( + int(scheduler_cfg.get("num_train_timesteps", 1000)), + training=True, + ) + + self.global_step = 0 + + # ------------------------------------------------------------------ # + # Subclass hooks # + # ------------------------------------------------------------------ # + + def _apply_parallelism(self): + """Set up distributed parallelism. Called during __init__.""" + raise NotImplementedError + + @contextmanager + def _grad_sync_context(self, is_update_step: bool): + """Context manager for gradient sync control during accumulation.""" + yield + + def _clip_grad_norm(self) -> float: + """Clip gradient norm. Returns the total norm value.""" + if self.max_grad_norm > 0: + norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm) + return norm.item() if isinstance(norm, torch.Tensor) else float(norm) + return 0.0 + + def _save_checkpoint(self): + """Save checkpoint at save_steps intervals. Override for custom strategies.""" + + # ------------------------------------------------------------------ # + # Common methods # + # ------------------------------------------------------------------ # + + def _setup_wandb(self): + self.use_wandb = False + if self.rank != 0: + return + if str(self.args.get("report_to", "")).lower() != "wandb": + return + if self.args.get("use_wandb") is False: + return + if wandb is None: + logger.warning("W&B requested but wandb is not installed.") + return + + project = self.args.get("wandb_project") or os.environ.get("WANDB_PROJECT", "primus-diffusion") + run_name = self.args.get("wandb_name") or self.args.get("run_name") + wandb_dir = self.args.get("wandb_dir") or os.environ.get("WANDB_DIR") + wandb.init(project=project, name=run_name, dir=wandb_dir, config=self.args) + self.use_wandb = True + + def _resolve_dtype(self) -> torch.dtype: + return resolve_dtype(self.args) + + def _create_optimizer(self): + lr = float(self.args.get("learning_rate", 1e-4)) + wd = float(self.args.get("weight_decay", 0.01)) + betas = ( + float(self.args.get("adam_beta1", 0.9)), + float(self.args.get("adam_beta2", 0.999)), + ) + eps = float(self.args.get("adam_epsilon", 1e-8)) + + params = [p for p in self.model.parameters() if p.requires_grad] + optimizer_kwargs = {"params": params, "lr": lr, "betas": betas, "eps": eps, "weight_decay": wd} + + optimizer = None + try: + optimizer = torch.optim.AdamW(**optimizer_kwargs, fused=True) + if self.rank == 0: + logger.info("Optimizer: torch.optim.AdamW(fused=True)") + except TypeError: + try: + optimizer = torch.optim.AdamW(**optimizer_kwargs, foreach=True) + if self.rank == 0: + logger.info("Optimizer: torch.optim.AdamW(foreach=True)") + except TypeError: + optimizer = torch.optim.AdamW(**optimizer_kwargs) + if self.rank == 0: + logger.info("Optimizer: torch.optim.AdamW(default)") + + if (self.args.get("bf16", False) or self.args.get("fp16", False)) and os.getenv( + "FP32_MASTER_WEIGHTS", "0" + ) == "1": + if self.rank == 0: + logger.info( + "FP32_MASTER_WEIGHTS=1: using AdamWFP32State (fp32 master weights + fp32 moments)." + ) + optimizer = AdamWFP32State( + optimizer.param_groups, + lr=lr, + betas=betas, + eps=eps, + weight_decay=wd, + ) + + return optimizer + + def compute_loss(self, batch): + """Prepare batch and compute training loss.""" + prepare_batch = getattr(self.processing_class, "prepare_batch", None) + if callable(prepare_batch): + batch = prepare_batch( + batch=batch, + device=self.device, + dtype=self._resolve_dtype(), + ) + # Ensure all tensors are on the correct device + for k, v in batch.items(): + if isinstance(v, torch.Tensor): + batch[k] = v.to(self.device, non_blocking=True) + + # Pass SP group so model can shard sequences across SP ranks + if self.sp_group is not None: + batch["sp_group"] = self.sp_group + + # Use explicit training entry point if available (GenAIModel interface) + forward_train = getattr(self.model, "forward_train", None) + if callable(forward_train): + outputs = forward_train(batch, scheduler=self.scheduler) + else: + outputs = self.model(batch, self.scheduler) + return outputs["loss"] + + def _infer_batch_size_from_tensors(self, value) -> int | None: + if isinstance(value, torch.Tensor): + return int(value.shape[0]) if value.ndim > 0 else 1 + if isinstance(value, dict): + for item in value.values(): + batch_size = self._infer_batch_size_from_tensors(item) + if batch_size is not None: + return batch_size + if isinstance(value, (list, tuple)): + for item in value: + batch_size = self._infer_batch_size_from_tensors(item) + if batch_size is not None: + return batch_size + return None + + def _infer_batch_size_from_sequences(self, value) -> int | None: + if isinstance(value, dict): + for item in value.values(): + batch_size = self._infer_batch_size_from_sequences(item) + if batch_size is not None: + return batch_size + return None + if isinstance(value, (list, tuple)): + if not value: + return 0 + first = value[0] + if isinstance(first, (dict, list, tuple, torch.Tensor)): + for item in value: + batch_size = self._infer_batch_size_from_sequences(item) + if batch_size is not None: + return batch_size + return None + return len(value) + return None + + def _infer_local_batch_size(self, batch) -> int: + tensor_batch_size = self._infer_batch_size_from_tensors(batch) + if tensor_batch_size is not None: + return tensor_batch_size + + sequence_batch_size = self._infer_batch_size_from_sequences(batch) + if sequence_batch_size is not None: + return sequence_batch_size + + return self.per_device_train_batch_size + + def _compute_samples_per_gpu_per_second( + self, + local_samples: int, + interval_seconds: float | None, + ) -> float | None: + if interval_seconds is None or interval_seconds <= 0 or local_samples <= 0 or self.world_size <= 0: + return None + + global_samples = float(local_samples) * float(self.data_parallel_world_size) + return global_samples / float(self.world_size) / float(interval_seconds) + + def _log_step( + self, + loss_value: float, + grad_norm: float = 0.0, + step_time: float | None = None, + elapsed: float | None = None, + eta_seconds: float | None = None, + throughput_samples_per_gpu_s: float | None = None, + ): + """Log training metrics. Format matches test regex expectations.""" + if self.rank != 0: + return + alloc, res, max_mem = get_memory() + lr = self.optimizer.param_groups[0]["lr"] + + # NOTE: The "step=... loss=... mem=.../...GB" line format is relied on by + # downstream log parsers; keep it stable when editing. + msg = ( + f"step={self.global_step} loss={loss_value:.4f} " + f"mem={alloc:.2f}/{res:.2f}GB peak_mem={max_mem:.2f}GB " + f"gnorm={grad_norm:.4f}" + ) + if step_time is not None: + msg += f" step_time={step_time:.2f}s" + if throughput_samples_per_gpu_s is not None: + msg += f" throughput={throughput_samples_per_gpu_s:.4f}samples/gpu/s" + if elapsed is not None: + msg += f" elapsed={elapsed / 60:.2f}m" + if eta_seconds is not None: + msg += f" eta={eta_seconds / 60:.2f}m" + logger.info(msg) + + if self.use_wandb: + payload = { + "train/loss": loss_value, + "train/step": self.global_step, + "train/grad_norm": grad_norm, + "train/lr": lr, + "mem/allocated_gb": alloc, + "mem/reserved_gb": res, + "mem/max_alloc_gb": max_mem, + } + if step_time is not None: + payload["time/step_s"] = step_time + if throughput_samples_per_gpu_s is not None: + payload["perf/samples_per_gpu_s"] = throughput_samples_per_gpu_s + if elapsed is not None: + payload["time/elapsed_s"] = elapsed + if eta_seconds is not None: + payload["time/eta_s"] = eta_seconds + wandb.log(payload, step=self.global_step) + + def train(self): + if self.rank == 0: + logger.info("Starting training...") + + # Ensure frozen state (idempotent) + core = getattr(self.model, "module", self.model) + if hasattr(core, "freeze_except"): + core.freeze_except() + + self.model.train() + self.optimizer.zero_grad(set_to_none=True) + torch.cuda.reset_peak_memory_stats() + + start_time = time.time() + last_log_time = start_time + local_samples_in_update = 0 + local_samples_since_log = 0 + update_steps_since_log = 0 + + for epoch in range(self.num_train_epochs): + self.sampler.set_epoch(epoch) + + for batch_idx, batch in enumerate(self.dataloader): + is_update_step = ((batch_idx + 1) % max(1, self.grad_accum_steps)) == 0 + local_samples_in_update += self._infer_local_batch_size(batch) + + with self._grad_sync_context(is_update_step): + loss = self.compute_loss(batch) + loss = loss / max(1, self.grad_accum_steps) + loss.backward() + + if is_update_step: + grad_norm = self._clip_grad_norm() + + self.optimizer.step() + self.lr_scheduler.step() + self.optimizer.zero_grad(set_to_none=True) + self.global_step += 1 + update_steps_since_log += 1 + local_samples_since_log += local_samples_in_update + local_samples_in_update = 0 + + # Logging + if self.global_step % self.logging_steps == 0: + loss_val = loss.detach().float().item() * max(1, self.grad_accum_steps) + now = time.time() + log_interval = now - last_log_time + step_time = log_interval / max(1, update_steps_since_log) + last_log_time = now + elapsed = now - start_time + steps_left = max(0, self.total_steps - self.global_step) + eta_seconds = step_time * steps_left + throughput_samples_per_gpu_s = self._compute_samples_per_gpu_per_second( + local_samples=local_samples_since_log, + interval_seconds=log_interval, + ) + self._log_step( + loss_val, + grad_norm=grad_norm, + step_time=step_time, + elapsed=elapsed, + eta_seconds=eta_seconds, + throughput_samples_per_gpu_s=throughput_samples_per_gpu_s, + ) + local_samples_since_log = 0 + update_steps_since_log = 0 + + # Periodic save + if self.save_steps > 0 and self.global_step % self.save_steps == 0: + self._save_checkpoint() + + # Early termination + if self.max_steps > 0 and self.global_step >= self.max_steps: + return + + if self.max_steps > 0 and self.global_step >= self.max_steps: + break + + if self.rank == 0: + elapsed = time.time() - start_time + logger.info(f"Training finished in {elapsed / 60:.2f} min") + + def save_model(self): + """Save final model. Override in subclass.""" + raise NotImplementedError diff --git a/primus/backends/diffusion/trainers/fsdp2.py b/primus/backends/diffusion/trainers/fsdp2.py new file mode 100644 index 000000000..1fa79b177 --- /dev/null +++ b/primus/backends/diffusion/trainers/fsdp2.py @@ -0,0 +1,395 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Wan PyTorch FSDP2 trainer +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager +from importlib import import_module + +import torch +from safetensors.torch import load_file as safe_load_file +from safetensors.torch import save_file as safe_save_file +from torch.distributed._composable.fsdp import MixedPrecisionPolicy, fully_shard +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, +) + +from primus.backends.diffusion.distributed import ( + create_device_mesh, + load_checkpoint_dtcp, + save_checkpoint_dtcp, + setup_distributed, +) +from primus.backends.diffusion.utils.log import logger + +from .base import BaseWanTrainer + + +class FSDP2Trainer(BaseWanTrainer): + def __init__( + self, + model: torch.nn.Module, + args: dict, + train_dataset, + data_collator, + processing_class, + rank: int, + world_size: int, + local_rank: int, + ): + super().__init__( + model=model, + args=args, + train_dataset=train_dataset, + data_collator=data_collator, + processing_class=processing_class, + rank=rank, + world_size=world_size, + local_rank=local_rank, + ) + + # FSDP2-specific: checkpoint strategy + # - "dit_only": save `dit_model.safetensors` (default for T2V training) + # - "dtcp_full": save model + optimizer via DTCP + # - "dtcp_model_only": save only model via DTCP + # - "dtcp_trainable": DTCP save only trainable params + optimizer + self.save_strategy = str(self.args.get("save_strategy", "dit_only")).lower() + + # Checkpoint loading (after optimizer & scheduler are fully initialized) + resume_from = self.args.get("resume_from_checkpoint") + if resume_from: + self._load_checkpoint(resume_from) + + # ------------------------------------------------------------------ # + # Parallelism # + # ------------------------------------------------------------------ # + + def _apply_parallelism(self): + """Set up FSDP2 composable sharding, optionally with Ulysses SP.""" + sp_size = int(self.args.get("sp_size", 1)) + dp_replicate = int(self.args.get("dp_replicate", 1)) + + self.mesh = create_device_mesh(self.world_size, sp_size=sp_size, dp_replicate=dp_replicate) + self.sp_group = self.mesh.get_group("ulysses") if (self.mesh is not None and sp_size > 1) else None + self.model.to(self.device) + + # Freeze non-trainable params BEFORE FSDP and optimizer creation + if hasattr(self.model, "freeze_except"): + self.model.freeze_except() + if self.rank == 0: + logger.info("FSDP2: Applied freeze_except (frozen non-trainable params)") + + self._apply_fsdp2() + + def _apply_fsdp2(self): + """Apply torch.distributed._composable.fsdp.fully_shard to the model.""" + mp_dtype = self._resolve_dtype() + + # ---- Mixed-precision policy (aligned with DeepSpeed bf16) ---- + # 1. Pre-cast params to bf16 BEFORE FSDP wrapping (bf16 storage) + # 2. reduce_dtype = bf16 (gradient reduce matches DeepSpeed bf16 all-reduce) + # 3. AdamWFP32State maintains fp32 master weights and writes back to bf16 + if mp_dtype != torch.float32: + mp_policy = MixedPrecisionPolicy( + param_dtype=mp_dtype, + reduce_dtype=mp_dtype, # bf16 reduce (matches DeepSpeed bf16) + ) + else: + mp_policy = None + + # When SP is enabled, FSDP2 shards across dp_shard_sp (DP + SP combined). + # This reduces per-rank parameter memory. + fsdp_mesh = self.mesh + if self.mesh is not None and self.sp_group is not None: + try: + fsdp_mesh = self.mesh["dp_shard_sp"] + except KeyError: + pass # fallback to full mesh + + wrap_target = str(self.args.get("fsdp2_wrap_target", "") or "").strip() + wrap_root = self._get_module_by_path(self.model, wrap_target) if wrap_target else self.model + + # Pre-cast to bf16 before FSDP wrapping so parameters are stored in bf16, + # matching DiffSynth/DeepSpeed bf16 behavior. + if mp_dtype != torch.float32: + wrap_root.to(dtype=mp_dtype) + if self.rank == 0: + logger.info(f"FSDP2: pre-cast '{wrap_target or ''}' to {mp_dtype} (DiffSynth-aligned)") + + if self.world_size == 1: + if self.rank == 0: + logger.info("FSDP2: world_size=1; skipping composable FSDP wrapping.") + return + + reshard_after_forward = bool(self.args.get("fsdp2_reshard_after_forward", True)) + + # Wrap transformer blocks first for optimal memory management + layer_cls_spec = self.args.get("fsdp_transformer_layer_cls_to_wrap") + if layer_cls_spec: + if isinstance(layer_cls_spec, str): + layer_cls_items = [x.strip() for x in layer_cls_spec.split(",") if x.strip()] + else: + layer_cls_items = [str(x).strip() for x in layer_cls_spec if str(x).strip()] + + cls_objs = set() + cls_names = set() + for item in layer_cls_items: + if "." in item: + mod_path, cls_name = item.rsplit(".", 1) + cls_objs.add(getattr(import_module(mod_path), cls_name)) + else: + cls_names.add(item) + + wrapped_count = 0 + seen = set() + for _, module in wrap_root.named_modules(): + if id(module) in seen: + continue + seen.add(id(module)) + if module is wrap_root: + continue + if (cls_objs and isinstance(module, tuple(cls_objs))) or ( + cls_names and module.__class__.__name__ in cls_names + ): + fully_shard( + module, + mesh=fsdp_mesh, + reshard_after_forward=reshard_after_forward, + mp_policy=mp_policy, + ) + wrapped_count += 1 + + if self.rank == 0: + logger.info(f"FSDP2: wrapped {wrapped_count} submodules under '{wrap_target or ''}'") + + fully_shard( + wrap_root, + mesh=fsdp_mesh, + reshard_after_forward=reshard_after_forward, + mp_policy=mp_policy, + ) + if self.rank == 0: + logger.info(f"FSDP2: applied fully_shard to '{wrap_target or ''}' with mp={mp_dtype}") + + @staticmethod + def _get_module_by_path(root: torch.nn.Module, path: str) -> torch.nn.Module: + """Resolve a dot-separated attribute path on a module.""" + cur = root + if not path: + return cur + for part in path.split("."): + if not hasattr(cur, part): + raise ValueError(f"fsdp2_wrap_target='{path}' is invalid: missing attribute '{part}'") + cur = getattr(cur, part) + if not isinstance(cur, torch.nn.Module): + raise ValueError(f"fsdp2_wrap_target='{path}' did not resolve to a torch.nn.Module") + return cur + + # ------------------------------------------------------------------ # + # Gradient sync # + # ------------------------------------------------------------------ # + + def _set_requires_gradient_sync(self, enabled: bool) -> None: + """ + Best-effort gradient sync control for composable FSDP2. + `fully_shard()` turns modules into FSDPM which exposes + `set_requires_gradient_sync`. For grad accumulation, we disable + sync on non-update micro-steps. + """ + seen = set() + for _, m in self.model.named_modules(): + if id(m) in seen: + continue + seen.add(id(m)) + setter = getattr(m, "set_requires_gradient_sync", None) + if callable(setter): + setter(bool(enabled)) + + @contextmanager + def _grad_sync_context(self, is_update_step: bool): + if self.world_size > 1 and self.grad_accum_steps > 1: + self._set_requires_gradient_sync(is_update_step) + yield + if is_update_step and self.world_size > 1 and self.grad_accum_steps > 1: + self._set_requires_gradient_sync(True) + + # ------------------------------------------------------------------ # + # Checkpointing # + # ------------------------------------------------------------------ # + + def _dtcp_save_args(self): + """Compute DTCP save parameters based on save_strategy.""" + if self.save_strategy == "dtcp_model_only": + return None, None + if self.save_strategy == "dtcp_trainable": + opts = StateDictOptions(full_state_dict=False, ignore_frozen_params=True) + return self.optimizer, opts + # dtcp_full (default non-dit strategy) + return self.optimizer, None + + def _save_dtcp(self, path): + """Save checkpoint via Distributed Tensor Checkpointing.""" + optimizer, opts = self._dtcp_save_args() + kwargs = {} + if opts is not None: + kwargs["model_state_options"] = opts + kwargs["optim_state_options"] = opts + save_checkpoint_dtcp( + self.model, + optimizer, + path, + epoch=0, + step=self.global_step, + additional_data={ + "lr_scheduler": (self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None) + }, + **kwargs, + ) + + def _save_checkpoint(self): + path = os.path.join(self.output_dir, f"checkpoint-{self.global_step}") + if self.save_strategy == "dit_only": + self._save_dit(os.path.join(path, "dit_model.safetensors")) + else: + self._save_dtcp(path) + + def _load_checkpoint(self, path): + if self.rank == 0: + logger.info(f"Loading checkpoint from {path}") + + if self.save_strategy == "dit_only": + core_model = self.model + if not hasattr(core_model, "dit"): + raise ValueError("save_strategy=dit_only requires model.dit to exist for checkpoint loading") + candidate = path + if os.path.isdir(path): + candidate = os.path.join(path, "dit_model.safetensors") + if not os.path.exists(candidate): + raise FileNotFoundError(f"dit_only checkpoint not found at: {candidate}") + state = safe_load_file(candidate) + missing, unexpected = core_model.dit.load_state_dict(state, strict=False) + if self.rank == 0: + logger.info( + f"Loaded DiT weights from {candidate}. " + f"Missing keys: {len(missing)}, " + f"Unexpected keys: {len(unexpected)}" + ) + return + + if self.save_strategy == "dtcp_trainable": + opts = StateDictOptions(full_state_dict=False, ignore_frozen_params=True) + meta = load_checkpoint_dtcp( + self.model, + self.optimizer, + path, + model_state_options=opts, + optim_state_options=opts, + ) + elif self.save_strategy == "dtcp_model_only": + meta = load_checkpoint_dtcp(self.model, None, path) + else: + meta = load_checkpoint_dtcp(self.model, self.optimizer, path) + + if "step" in meta: + self.global_step = meta["step"] + if self.rank == 0: + logger.info(f"Resumed from step {self.global_step}") + if self.lr_scheduler is not None and isinstance(meta, dict) and meta.get("lr_scheduler") is not None: + try: + self.lr_scheduler.load_state_dict(meta["lr_scheduler"]) + if self.rank == 0: + logger.info("Resumed lr_scheduler state from checkpoint meta") + except Exception as exc: + if self.rank == 0: + logger.warning(f"Failed to restore lr_scheduler state: {exc}") + + def _save_dit(self, save_path: str) -> None: + """ + Save `dit` weights as a single safetensors file. + + Notes: + - On composable FSDP (world_size>1), ``get_model_state_dict`` with + ``full_state_dict=True`` returns the full dict **only on rank 0**; + other ranks receive an empty dict. This is expected – do NOT + early-return on the empty-dict check, or non-rank-0 processes will + race ahead and desynchronize from rank 0 (which still needs to + write to disk). + - A ``dist.barrier()`` at the end keeps all ranks in lockstep so + that back-to-back saves (e.g. periodic checkpoint + final save) + never overlap. + """ + import torch.distributed as dist + + core_model = self.model + if not hasattr(core_model, "dit"): + logger.warning("save_model: model has no `dit` attribute; skipping save.") + return + + if self.world_size > 1: + # FSDP-wrapped: use get_model_state_dict to unshard DTensors + full_state = get_model_state_dict( + core_model, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) + dit_state_dict = {k[len("dit.") :]: v for k, v in full_state.items() if k.startswith("dit.")} + else: + # Single GPU (no FSDP wrapping): direct state_dict + dit_state_dict = core_model.dit.state_dict() + if not dit_state_dict: + logger.warning("save_model: DiT state dict is empty; skipping save.") + return + + # Only rank 0 has the full dict in multi-GPU; write from rank 0 only. + if self.rank == 0: + if not dit_state_dict: + logger.warning("save_model: DiT state dict is empty on rank 0; skipping save.") + else: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + safe_save_file(dit_state_dict, save_path) + logger.info(f"Saved DiT weights to {save_path}") + + if hasattr(core_model, "config"): + try: + core_model.config.save_pretrained(os.path.dirname(save_path)) + except Exception as exc: + logger.warning(f"save_model: failed to save config: {exc}") + + # Barrier: ensure all ranks wait for rank 0 to finish writing before + # any rank proceeds to the next operation (e.g. another save or exit). + if self.world_size > 1: + dist.barrier() + + def save_model(self): + """Save final model using the configured strategy.""" + if self.save_strategy == "dit_only": + save_path = os.path.join(self.output_dir, "dit_model.safetensors") + self._save_dit(save_path) + return + + path = os.path.join(self.output_dir, "checkpoint-final") + if self.rank == 0: + logger.info(f"Saving final checkpoint to {path} (strategy={self.save_strategy})") + self._save_dtcp(path) + + +def build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict): + rank, world_size, local_rank = setup_distributed() + return FSDP2Trainer( + model=model, + args=trainer_args, + train_dataset=dataset, + data_collator=dataset.get_collator(), + processing_class=processor, + rank=rank, + world_size=world_size, + local_rank=local_rank, + ) diff --git a/primus/backends/diffusion/utils/__init__.py b/primus/backends/diffusion/utils/__init__.py new file mode 100644 index 000000000..3d7aa5fe7 --- /dev/null +++ b/primus/backends/diffusion/utils/__init__.py @@ -0,0 +1,10 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Utility package for the Primus Wan backend. + +Heavy vision helpers are imported lazily by the qwen_vl_utils video path. +""" diff --git a/primus/backends/diffusion/utils/data_utils.py b/primus/backends/diffusion/utils/data_utils.py new file mode 100644 index 000000000..de149c48b --- /dev/null +++ b/primus/backends/diffusion/utils/data_utils.py @@ -0,0 +1,61 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + + +import math + +FRAME_FACTOR = 2 +FPS = 2.0 +FPS_MIN_FRAMES = 4 +FPS_MAX_FRAMES = 768 + + +def smart_nframes( + total_frames: int, + video_fps: int | float, + fps: int, +) -> int: + """calculate the number of frames for video used for model inputs. + + Args: + ele (dict): a dict contains the configuration of video. + support either `fps` or `nframes`: + - nframes: the number of frames to extract for model inputs. + - fps: the fps to extract frames for model inputs. + - min_frames: the minimum number of frames of the video, only used when fps is provided. + - max_frames: the maximum number of frames of the video, only used when fps is provided. + total_frames (int): the original total number of frames of the video. + video_fps (int | float): the original fps of the video. + + Raises: + ValueError: nframes should in interval [FRAME_FACTOR, total_frames]. + + Returns: + int: the number of frames for video used for model inputs. + """ + min_frames = ceil_by_factor(FPS_MIN_FRAMES, FRAME_FACTOR) + max_frames = floor_by_factor(min(FPS_MAX_FRAMES, total_frames), FRAME_FACTOR) + nframes = total_frames / video_fps * fps + nframes = min(min(max(nframes, min_frames), max_frames), total_frames) + nframes = floor_by_factor(nframes, FRAME_FACTOR) + if not (FRAME_FACTOR <= nframes and nframes <= total_frames): + raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.") + return nframes + + +def round_by_factor(number: int, factor: int) -> int: + """Returns the closest integer to 'number' that is divisible by 'factor'.""" + return round(number / factor) * factor + + +def ceil_by_factor(number: int, factor: int) -> int: + """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'.""" + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int, factor: int) -> int: + """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'.""" + return math.floor(number / factor) * factor diff --git a/primus/backends/diffusion/utils/log.py b/primus/backends/diffusion/utils/log.py new file mode 100644 index 000000000..2193ddebc --- /dev/null +++ b/primus/backends/diffusion/utils/log.py @@ -0,0 +1,42 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Logger proxy for the diffusion backend. + +Primus configures loguru sinks whose format references bound ``extra`` fields +(``rank``, ``world_size``, ``user``, ``team``, ``module_name``, ``node_ip``). +Those fields are only present on the Primus-bound logger created in +``primus.core.utils.logger.setup_logger``. Emitting from the raw, unbound +global loguru logger (``from loguru import logger``) produces records without +those extras, which makes every sink raise "Logging error in Handler". + +This module exposes a ``logger`` proxy that forwards attribute access to the +live Primus-bound logger, so diffusion code keeps the familiar +``logger.info(...)`` style while inheriting the bound extras. Before the Primus +logger is initialized (e.g. standalone tooling), it falls back to the raw +loguru logger, which at that point still uses loguru's default sink. +""" + +from __future__ import annotations + +from primus.core.utils import logger as _primus_logger_module + + +class _LoggerProxy: + @staticmethod + def _resolve(): + bound = getattr(_primus_logger_module, "_logger", None) + if bound is not None: + return bound + from loguru import logger as _raw_logger + + return _raw_logger + + def __getattr__(self, name): + return getattr(self._resolve(), name) + + +logger = _LoggerProxy() diff --git a/primus/backends/diffusion/utils/train_utils.py b/primus/backends/diffusion/utils/train_utils.py new file mode 100644 index 000000000..71deede8f --- /dev/null +++ b/primus/backends/diffusion/utils/train_utils.py @@ -0,0 +1,200 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + + +""" +utils for training +""" + +import hashlib +import os +import random +from contextlib import contextmanager + +import numpy as np +import torch +from safetensors import safe_open + + +def resolve_dtype(config_or_args) -> torch.dtype: + """ + Resolve mixed-precision dtype from either: + - a dict-like config (e.g. trainer_args) + - an object with attributes (e.g. HF TrainingArguments) + + Priority: + bf16 -> fp16 -> fp32 + """ + # Dict-like + if isinstance(config_or_args, dict): + if config_or_args.get("bf16", False): + return torch.bfloat16 + if config_or_args.get("fp16", False): + return torch.float16 + return torch.float32 + + # Attribute-like (HF TrainingArguments etc.) + bf16 = bool(getattr(config_or_args, "bf16", False)) + fp16 = bool(getattr(config_or_args, "fp16", False)) + if bf16: + return torch.bfloat16 + if fp16: + return torch.float16 + return torch.float32 + + +def set_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def count_parameters(model): + total = sum(p.numel() for p in model.parameters()) + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + return total, trainable + + +def get_memory(unit=1e9): + torch.cuda.synchronize() + allocated = torch.cuda.memory_allocated() + reserved = torch.cuda.memory_reserved() + max_alloc = torch.cuda.max_memory_allocated() + return allocated / unit, reserved / unit, max_alloc / unit + + +def print_cuda_memory(prefix="", unit=1e9): + allocated = torch.cuda.memory_allocated() / unit + reserved = torch.cuda.memory_reserved() / unit + max_alloc = torch.cuda.max_memory_allocated() / unit + print( + f"{prefix} " + f"allocated={allocated:.2f}GB, " + f"reserved={reserved:.2f}GB, " + f"max_alloc={max_alloc:.2f}GB" + ) + + +@contextmanager +def init_weights_on_device(device=torch.device("meta"), include_buffers: bool = False): + + old_register_parameter = torch.nn.Module.register_parameter + if include_buffers: + old_register_buffer = torch.nn.Module.register_buffer + + def register_empty_parameter(module, name, param): + old_register_parameter(module, name, param) + if param is not None: + param_cls = type(module._parameters[name]) + kwargs = module._parameters[name].__dict__ + kwargs["requires_grad"] = param.requires_grad + module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) + + def register_empty_buffer(module, name, buffer, persistent=True): + old_register_buffer(module, name, buffer, persistent=persistent) + if buffer is not None: + module._buffers[name] = module._buffers[name].to(device) + + def patch_tensor_constructor(fn): + def wrapper(*args, **kwargs): + kwargs["device"] = device + return fn(*args, **kwargs) + + return wrapper + + if include_buffers: + tensor_constructors_to_patch = { + torch_function_name: getattr(torch, torch_function_name) + for torch_function_name in ["empty", "zeros", "ones", "full"] + } + else: + tensor_constructors_to_patch = {} + + try: + torch.nn.Module.register_parameter = register_empty_parameter + if include_buffers: + torch.nn.Module.register_buffer = register_empty_buffer + for torch_function_name in tensor_constructors_to_patch.keys(): + setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name))) + yield + finally: + torch.nn.Module.register_parameter = old_register_parameter + if include_buffers: + torch.nn.Module.register_buffer = old_register_buffer + for torch_function_name, old_torch_function in tensor_constructors_to_patch.items(): + setattr(torch, torch_function_name, old_torch_function) + + +def load_state_dict_from_folder(file_path, torch_dtype=None): + state_dict = {} + for file_name in os.listdir(file_path): + if "." in file_name and file_name.split(".")[-1] in ["safetensors", "bin", "ckpt", "pth", "pt"]: + state_dict.update(load_state_dict(os.path.join(file_path, file_name), torch_dtype=torch_dtype)) + return state_dict + + +def load_state_dict(file_path, torch_dtype=None, device="cpu"): + if file_path.endswith(".safetensors"): + return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device) + else: + return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device) + + +def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"): + state_dict = {} + with safe_open(file_path, framework="pt", device=str(device)) as f: + for k in f.keys(): + state_dict[k] = f.get_tensor(k) + if torch_dtype is not None: + state_dict[k] = state_dict[k].to(torch_dtype) + return state_dict + + +def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"): + state_dict = torch.load(file_path, map_location=device, weights_only=True) + if torch_dtype is not None: + for i in state_dict: + if isinstance(state_dict[i], torch.Tensor): + state_dict[i] = state_dict[i].to(torch_dtype) + return state_dict + + +def convert_state_dict_keys_to_single_str(state_dict, with_shape=True): + keys = [] + for key, value in state_dict.items(): + if isinstance(key, str): + if isinstance(value, torch.Tensor): + if with_shape: + shape = "_".join(map(str, list(value.shape))) + keys.append(key + ":" + shape) + keys.append(key) + elif isinstance(value, dict): + keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape)) + keys.sort() + keys_str = ",".join(keys) + return keys_str + + +def split_state_dict_with_prefix(state_dict): + keys = sorted([key for key in state_dict if isinstance(key, str)]) + prefix_dict = {} + for key in keys: + prefix = key if "." not in key else key.split(".")[0] + if prefix not in prefix_dict: + prefix_dict[prefix] = [] + prefix_dict[prefix].append(key) + state_dicts = [] + for prefix, keys in prefix_dict.items(): + sub_state_dict = {key: state_dict[key] for key in keys} + state_dicts.append(sub_state_dict) + return state_dicts + + +def hash_state_dict_keys(state_dict, with_shape=True): + keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape) + keys_str = keys_str.encode(encoding="UTF-8") + return hashlib.md5(keys_str).hexdigest() diff --git a/primus/backends/diffusion/utils/vision_process.py b/primus/backends/diffusion/utils/vision_process.py new file mode 100644 index 000000000..ea2aebd36 --- /dev/null +++ b/primus/backends/diffusion/utils/vision_process.py @@ -0,0 +1,569 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +# Adapted from qwen-vl-utils +""" + +import base64 +import copy +import logging +import math +import os +import sys +import time +import warnings +from concurrent.futures import ThreadPoolExecutor +from functools import lru_cache +from io import BytesIO +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import requests +import torch +import torchvision +from packaging import version +from PIL import Image +from torchvision import io, transforms +from torchvision.transforms import InterpolationMode + +MAX_RATIO = 200 +SPATIAL_MERGE_SIZE = 2 +IMAGE_MIN_TOKEN_NUM = 4 +IMAGE_MAX_TOKEN_NUM = 16384 +VIDEO_MIN_TOKEN_NUM = 128 +VIDEO_MAX_TOKEN_NUM = 768 + +FPS = 2.0 +FRAME_FACTOR = 2 +FPS_MIN_FRAMES = 4 +FPS_MAX_FRAMES = 768 +MAX_NUM_WORKERS_FETCH_VIDEO = 8 + +MODEL_SEQ_LEN = int(float(os.environ.get("MODEL_SEQ_LEN", 128000))) +logger = logging.getLogger(__name__) + + +def round_by_factor(number: int, factor: int) -> int: + """Returns the closest integer to 'number' that is divisible by 'factor'.""" + return round(number / factor) * factor + + +def ceil_by_factor(number: int, factor: int) -> int: + """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'.""" + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int, factor: int) -> int: + """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'.""" + return math.floor(number / factor) * factor + + +def smart_resize( + height: int, width: int, factor: int, min_pixels: Optional[int] = None, max_pixels: Optional[int] = None +) -> Tuple[int, int]: + """ + Rescales the image so that the following conditions are met: + + 1. Both dimensions (height and width) are divisible by 'factor'. + 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. + 3. The aspect ratio of the image is maintained as closely as possible. + """ + max_pixels = max_pixels if max_pixels is not None else (IMAGE_MAX_TOKEN_NUM * factor**2) + min_pixels = min_pixels if min_pixels is not None else (IMAGE_MIN_TOKEN_NUM * factor**2) + assert max_pixels >= min_pixels, "The max_pixels of image must be greater than or equal to min_pixels." + if max(height, width) / min(height, width) > MAX_RATIO: + raise ValueError( + f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" + ) + h_bar = max(factor, round_by_factor(height, factor)) + w_bar = max(factor, round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = floor_by_factor(height / beta, factor) + w_bar = floor_by_factor(width / beta, factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = ceil_by_factor(height * beta, factor) + w_bar = ceil_by_factor(width * beta, factor) + return h_bar, w_bar + + +def to_rgb(pil_image: Image.Image) -> Image.Image: + if pil_image.mode == "RGBA": + white_background = Image.new("RGB", pil_image.size, (255, 255, 255)) + white_background.paste(pil_image, mask=pil_image.split()[3]) # Use alpha channel as mask + return white_background + else: + return pil_image.convert("RGB") + + +def fetch_image(ele: Dict[str, Union[str, Image.Image]], image_patch_size: int = 14) -> Image.Image: + if "image" in ele: + image = ele["image"] + else: + image = ele["image_url"] + + image_obj = None + patch_factor = int(image_patch_size * SPATIAL_MERGE_SIZE) + if isinstance(image, Image.Image): + image_obj = image + elif image.startswith("http://") or image.startswith("https://"): + with requests.get(image, stream=True) as response: + response.raise_for_status() + with BytesIO(response.content) as bio: + image_obj = copy.deepcopy(Image.open(bio)) + elif image.startswith("file://"): + image_obj = Image.open(image[7:]) + elif image.startswith("data:image"): + if "base64," in image: + _, base64_data = image.split("base64,", 1) + data = base64.b64decode(base64_data) + with BytesIO(data) as bio: + image_obj = copy.deepcopy(Image.open(bio)) + else: + image_obj = Image.open(image) + if image_obj is None: + raise ValueError( + f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}" + ) + image = to_rgb(image_obj) + + ## resize + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=patch_factor, + ) + else: + width, height = image.size + min_pixels = ele.get("min_pixels", IMAGE_MIN_TOKEN_NUM * patch_factor**2) + max_pixels = ele.get("max_pixels", IMAGE_MAX_TOKEN_NUM * patch_factor**2) + resized_height, resized_width = smart_resize( + height, + width, + factor=patch_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + image = image.resize((resized_width, resized_height)) + return image + + +def smart_nframes( + ele: Dict[str, Any], + total_frames: int, + video_fps: Union[int, float], +) -> int: + """calculate the number of frames for video used for model inputs. + + Args: + ele (dict): a dict contains the configuration of video. + support either `fps` or `nframes`: + - nframes: the number of frames to extract for model inputs. + - fps: the fps to extract frames for model inputs. + - min_frames: the minimum number of frames of the video, only used when fps is provided. + - max_frames: the maximum number of frames of the video, only used when fps is provided. + total_frames (int): the original total number of frames of the video. + video_fps (int | float): the original fps of the video. + + Raises: + ValueError: nframes should in interval [FRAME_FACTOR, total_frames]. + + Returns: + int: the number of frames for video used for model inputs. + """ + assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`" + if "nframes" in ele: + nframes = round_by_factor(ele["nframes"], FRAME_FACTOR) + else: + fps = ele.get("fps", FPS) + min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR) + max_frames = floor_by_factor(ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR) + nframes = total_frames / video_fps * fps + if nframes > total_frames: + logger.warning(f"smart_nframes: nframes[{nframes}] > total_frames[{total_frames}]") + nframes = min(min(max(nframes, min_frames), max_frames), total_frames) + nframes = floor_by_factor(nframes, FRAME_FACTOR) + if not (FRAME_FACTOR <= nframes and nframes <= total_frames): + raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.") + return nframes + + +def _read_video_torchvision( + ele: Dict[str, Any], +) -> Tuple[torch.Tensor, float]: + """read video using torchvision.io.read_video + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + video_path = ele["video"] + if version.parse(torchvision.__version__) < version.parse("0.19.0"): + if "http://" in video_path or "https://" in video_path: + warnings.warn( + "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0." + ) + if "file://" in video_path: + video_path = video_path[7:] + st = time.time() + video, audio, info = io.read_video( + video_path, + start_pts=ele.get("video_start", 0.0), + end_pts=ele.get("video_end", None), + pts_unit="sec", + output_format="TCHW", + ) + total_frames, video_fps = video.size(0), info["video_fps"] + logger.info(f"torchvision: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(0, total_frames - 1, nframes).round().long() + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + video = video[idx] + + video_metadata = dict( + fps=video_fps, + frames_indices=idx, + total_num_frames=total_frames, + video_backend="torchvision", + ) + return video, video_metadata, sample_fps + + +def is_decord_available() -> bool: + import importlib.util + + return importlib.util.find_spec("decord") is not None + + +def calculate_video_frame_range( + ele: Dict[str, Any], + total_frames: int, + video_fps: float, +) -> Tuple[int, int, int]: + """ + Calculate the start and end frame indices based on the given time range. + + Args: + ele (dict): A dictionary containing optional 'video_start' and 'video_end' keys (in seconds). + total_frames (int): Total number of frames in the video. + video_fps (float): Frames per second of the video. + + Returns: + tuple: A tuple containing (start_frame, end_frame, frame_count). + + Raises: + ValueError: If input parameters are invalid or the time range is inconsistent. + """ + # Validate essential parameters + if video_fps <= 0: + raise ValueError("video_fps must be a positive number") + if total_frames <= 0: + raise ValueError("total_frames must be a positive integer") + + # Get start and end time in seconds + video_start = ele.get("video_start", None) + video_end = ele.get("video_end", None) + if video_start is None and video_end is None: + return 0, total_frames - 1, total_frames + + max_duration = total_frames / video_fps + # Process start frame + if video_start is not None: + video_start_clamped = max(0.0, min(video_start, max_duration)) + start_frame = math.ceil(video_start_clamped * video_fps) + else: + start_frame = 0 + # Process end frame + if video_end is not None: + video_end_clamped = max(0.0, min(video_end, max_duration)) + end_frame = math.floor(video_end_clamped * video_fps) + end_frame = min(end_frame, total_frames - 1) + else: + end_frame = total_frames - 1 + + # Validate frame order + if start_frame >= end_frame: + raise ValueError( + f"Invalid time range: Start frame {start_frame} (at {video_start_clamped if video_start is not None else 0}s) " + f"exceeds end frame {end_frame} (at {video_end_clamped if video_end is not None else max_duration}s). " + f"Video duration: {max_duration:.2f}s ({total_frames} frames @ {video_fps}fps)" + ) + + logger.info( + f"calculate video frame range: {start_frame=}, {end_frame=}, {total_frames=} from {video_start=}, {video_end=}, {video_fps=:.3f}" + ) + return start_frame, end_frame, end_frame - start_frame + 1 + + +def _read_video_decord( + ele: Dict[str, Any], +) -> Tuple[torch.Tensor, float]: + """read video using decord.VideoReader + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + import decord + + video_path = ele["video"] + time.time() + vr = decord.VideoReader(video_path) + total_frames, video_fps = len(vr), vr.get_avg_fps() + start_frame, end_frame, total_frames = calculate_video_frame_range( + ele, + total_frames, + video_fps, + ) + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist() + video = vr.get_batch(idx).asnumpy() + video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format + # logger.info(f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + + video_metadata = dict( + fps=video_fps, + frames_indices=idx, + total_num_frames=total_frames, + video_backend="decord", + ) + return video, video_metadata, sample_fps + + +def is_torchcodec_available() -> bool: + import importlib.util + + return importlib.util.find_spec("torchcodec") is not None + + +def _read_video_torchcodec( + ele: Dict[str, Any], +) -> Tuple[torch.Tensor, float]: + """read video using torchcodec.decoders.VideoDecoder + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + from torchcodec.decoders import VideoDecoder + + TORCHCODEC_NUM_THREADS = int(os.environ.get("TORCHCODEC_NUM_THREADS", 8)) + logger.info(f"set TORCHCODEC_NUM_THREADS: {TORCHCODEC_NUM_THREADS}") + video_path = ele["video"] + st = time.time() + decoder = VideoDecoder(video_path, num_ffmpeg_threads=TORCHCODEC_NUM_THREADS) + video_fps = decoder.metadata.average_fps + total_frames = decoder.metadata.num_frames + start_frame, end_frame, total_frames = calculate_video_frame_range( + ele, + total_frames, + video_fps, + ) + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist() + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + video = decoder.get_frames_at(indices=idx).data + logger.info(f"torchcodec: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + + video_metadata = dict( + fps=video_fps, + frames_indices=idx, + total_num_frames=total_frames, + video_backend="torchcodec", + ) + return video, video_metadata, sample_fps + + +VIDEO_READER_BACKENDS = { + "decord": _read_video_decord, + "torchvision": _read_video_torchvision, + "torchcodec": _read_video_torchcodec, +} + +FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None) + + +@lru_cache(maxsize=1) +def get_video_reader_backend() -> str: + if FORCE_QWENVL_VIDEO_READER is not None: + video_reader_backend = FORCE_QWENVL_VIDEO_READER + elif is_torchcodec_available(): + video_reader_backend = "torchcodec" + elif is_decord_available(): + video_reader_backend = "decord" + else: + video_reader_backend = "torchvision" + print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr) + return video_reader_backend + + +def fetch_video( + ele: Dict[str, Any], + image_patch_size: int = 14, + return_video_sample_fps: bool = False, + return_video_metadata: bool = False, +) -> Union[torch.Tensor, List[Image.Image]]: + image_factor = image_patch_size * SPATIAL_MERGE_SIZE + VIDEO_FRAME_MIN_PIXELS = VIDEO_MIN_TOKEN_NUM * image_factor * image_factor + VIDEO_FRAME_MAX_PIXELS = VIDEO_MAX_TOKEN_NUM * image_factor * image_factor + if isinstance(ele["video"], str): + video_reader_backend = get_video_reader_backend() + try: + video, video_metadata, sample_fps = VIDEO_READER_BACKENDS[video_reader_backend](ele) + except Exception as e: + # logger.warning(f"video_reader_backend {video_reader_backend} error, use torchvision as default, msg: {e}") + logger.warning( + f"video_reader_backend {video_reader_backend} error, for {ele=} use torchvision as default, msg: {e}" + ) + video, video_metadata, sample_fps = VIDEO_READER_BACKENDS["torchvision"](ele) + else: + # The input is a list of frames + assert isinstance(ele["video"], (list, tuple)) + process_info = ele.copy() + process_info.pop("type", None) + process_info.pop("video", None) + # use ThreadPoolExecutor to parallel process frames + max_workers = min(MAX_NUM_WORKERS_FETCH_VIDEO, len(ele["video"])) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit(fetch_image, {"image": video_element, **process_info}, image_factor) + for video_element in ele["video"] + ] + image_list = [future.result() for future in futures] + + nframes = ceil_by_factor(len(image_list), FRAME_FACTOR) + if len(image_list) < nframes: + image_list.extend([image_list[-1]] * (nframes - len(image_list))) + + sample_fps = ele.get("sample_fps", 2.0) + video = torch.stack([torch.from_numpy(np.array(image).transpose(2, 0, 1)) for image in image_list]) + + # fake video metadata + raw_fps = process_info.pop("raw_fps", sample_fps) + video_metadata = dict( + fps=raw_fps, + frames_indices=[i for i in range(len(video))], + total_num_frames=(nframes / sample_fps) * raw_fps, + ) + + nframes, _, height, width = video.shape + min_pixels = ele.get("min_pixels", VIDEO_FRAME_MIN_PIXELS) + total_pixels = ele.get("total_pixels", MODEL_SEQ_LEN * image_factor * image_factor * 0.9) + max_pixels = max( + min(VIDEO_FRAME_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05) + ) + max_pixels_supposed = ele.get("max_pixels", max_pixels) + if max_pixels_supposed > max_pixels: + logger.warning(f"The given max_pixels[{max_pixels_supposed}] exceeds limit[{max_pixels}].") + max_pixels = min(max_pixels_supposed, max_pixels) + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=image_factor, + ) + else: + resized_height, resized_width = smart_resize( + height, + width, + factor=image_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + video = transforms.functional.resize( + video, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ).float() + + final_video = (video, video_metadata) if return_video_metadata else video + if return_video_sample_fps: + return final_video, sample_fps + return final_video + + +def extract_vision_info( + conversations: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]] +) -> List[Dict[str, Any]]: + vision_infos = [] + if isinstance(conversations[0], dict): + conversations = [conversations] + for conversation in conversations: + for message in conversation: + if isinstance(message["content"], list): + for ele in message["content"]: + if ( + "image" in ele + or "image_url" in ele + or "video" in ele + or ele.get("type", "text") in ("image", "image_url", "video") + ): + vision_infos.append(ele) + return vision_infos + + +def process_vision_info( + conversations: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]], + return_video_kwargs: bool = False, + return_video_metadata: bool = False, + image_patch_size: int = 14, +) -> Tuple[ + Optional[List[Image.Image]], + Optional[List[Union[torch.Tensor, List[Image.Image]]]], + Optional[Dict[str, Any]], +]: + + vision_infos = extract_vision_info(conversations) + ## Read images or videos + image_inputs = [] + video_inputs = [] + video_sample_fps_list = [] + for vision_info in vision_infos: + if "image" in vision_info or "image_url" in vision_info: + image_inputs.append(fetch_image(vision_info, image_patch_size=image_patch_size)) + elif "video" in vision_info: + video_input, video_sample_fps = fetch_video( + vision_info, + return_video_sample_fps=True, + image_patch_size=image_patch_size, + return_video_metadata=return_video_metadata, + ) + video_sample_fps_list.append(video_sample_fps) + video_inputs.append(video_input) + else: + raise ValueError("image, image_url or video should in content.") + if len(image_inputs) == 0: + image_inputs = None + if len(video_inputs) == 0: + video_inputs = None + + video_kwargs = {"do_sample_frames": False} + if not return_video_metadata: # BC for qwen2.5vl + video_kwargs.update({"fps": video_sample_fps_list}) + + if return_video_kwargs: + return image_inputs, video_inputs, video_kwargs + return image_inputs, video_inputs diff --git a/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml b/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml new file mode 100644 index 000000000..77995b460 --- /dev/null +++ b/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml @@ -0,0 +1,13 @@ +model: + name: wan + config: + load_from_pretrained_path: ${PRETRAINED_PATH:/models/Wan2.1-T2V-1.3B} + config: + model_type: t2v + seperated_timestep: false + fuse_vae_embedding_in_latents: false + trainable_modules: dit + vae_type: wan_video_vae + encoder: + t5_encoder: ${TEXT_ENCODER:/models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth} + autoencoder: ${VAE_CHECKPOINT:/models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth} diff --git a/primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml b/primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml new file mode 100644 index 000000000..d9f089b3a --- /dev/null +++ b/primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml @@ -0,0 +1,8 @@ +extends: + - wan2.1_t2v_1.3b.yaml + +# SFT reuses the Wan2.1-T2V-1.3B preset and only swaps the DiT +# initialization source to INIT_CHECKPOINT (post-train convention). +model: + config: + load_from_pretrained_path: ${INIT_CHECKPOINT:/models/Wan2.1-T2V-1.3B} diff --git a/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml b/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml new file mode 100644 index 000000000..a3fa8ee9f --- /dev/null +++ b/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml @@ -0,0 +1,13 @@ +model: + name: wan + config: + load_from_pretrained_path: ${PRETRAINED_PATH:/models/Wan2.2-TI2V-5B} + config: + model_type: ti2v + seperated_timestep: false + fuse_vae_embedding_in_latents: false + trainable_modules: dit + vae_type: wan_video_vae_38 + encoder: + t5_encoder: ${TEXT_ENCODER:/models/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth} + autoencoder: ${VAE_CHECKPOINT:/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth} diff --git a/primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml b/primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml new file mode 100644 index 000000000..b8b1cd8d8 --- /dev/null +++ b/primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml @@ -0,0 +1,8 @@ +extends: + - wan2.2_ti2v_5b.yaml + +# SFT reuses the Wan2.2-TI2V-5B preset and only swaps the DiT +# initialization source to INIT_CHECKPOINT (post-train convention). +model: + config: + load_from_pretrained_path: ${INIT_CHECKPOINT:/models/Wan2.2-TI2V-5B} diff --git a/primus/configs/modules/diffusion/post_trainer.yaml b/primus/configs/modules/diffusion/post_trainer.yaml new file mode 100644 index 000000000..e89d7b62f --- /dev/null +++ b/primus/configs/modules/diffusion/post_trainer.yaml @@ -0,0 +1,5 @@ +extends: + - ../module_base.yaml + +trainable: true +stage: posttrain diff --git a/primus/configs/modules/diffusion/pre_trainer.yaml b/primus/configs/modules/diffusion/pre_trainer.yaml new file mode 100644 index 000000000..3d9d3cbdb --- /dev/null +++ b/primus/configs/modules/diffusion/pre_trainer.yaml @@ -0,0 +1,4 @@ +extends: + - ../module_base.yaml + +stage: pretrain diff --git a/runner/helpers/hooks/train/posttrain/diffusion/prepare.py b/runner/helpers/hooks/train/posttrain/diffusion/prepare.py new file mode 100644 index 000000000..4f011269f --- /dev/null +++ b/runner/helpers/hooks/train/posttrain/diffusion/prepare.py @@ -0,0 +1,24 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import runpy +import sys +from pathlib import Path + + +def main() -> None: + primus_root = Path(__file__).resolve().parents[6] + hook = primus_root / "runner" / "helpers" / "hooks" / "train" / "pretrain" / "diffusion" / "prepare.py" + sys.argv[0] = str(hook) + if "--module_name" not in sys.argv: + sys.argv.extend(["--module_name", "post_trainer"]) + runpy.run_path(str(hook), run_name="__main__") + + +if __name__ == "__main__": + main() diff --git a/runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh b/runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh new file mode 100755 index 000000000..e0c3c97c0 --- /dev/null +++ b/runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh @@ -0,0 +1,40 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +DEFAULT_PRIMUS_ROOT="$(cd "${SCRIPT_DIR}/../../../../../.." && pwd)" +PRIMUS_ROOT="${PRIMUS_PATH:-${DEFAULT_PRIMUS_ROOT}}" + +while [[ $# -gt 0 ]]; do + case "$1" in + --data_path) + DATA_PATH="$2" + shift 2 + ;; + --primus_path) + PRIMUS_ROOT="$2" + shift 2 + ;; + *) + shift + ;; + esac +done + +DATA_PATH="${DATA_PATH:-${PRIMUS_ROOT}/data}" +PIP_CACHE_DIR="${PIP_CACHE_DIR:-${DATA_PATH}/pip_cache}" + +echo "[INFO] Using pip cache: ${PIP_CACHE_DIR}" +mkdir -p "${PIP_CACHE_DIR}" + +REQ_FILE="${SCRIPT_DIR}/requirements-diffusion.txt" +if [[ -f "${REQ_FILE}" ]] && grep -qE '^[[:space:]]*[^#[:space:]]' "${REQ_FILE}"; then + echo "[+] Installing Diffusion dependencies..." + pip install --cache-dir="${PIP_CACHE_DIR}" -r "${REQ_FILE}" + echo "[OK] Diffusion dependencies installed" +fi diff --git a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py new file mode 100644 index 000000000..b23e6a5ae --- /dev/null +++ b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py @@ -0,0 +1,132 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path +from typing import Any + +PRIMUS_ROOT = Path(__file__).resolve().parents[6] +if str(PRIMUS_ROOT) not in sys.path: + sys.path.insert(0, str(PRIMUS_ROOT)) + +from primus.backends.diffusion.diffusion_adapter import DiffusionAdapter +from primus.core.config.primus_config import get_module_config, load_primus_config +from primus.core.utils.yaml_utils import nested_namespace_to_dict + + +def _log(message: str) -> None: + print(f"[INFO] diffusion prepare: {message}", file=sys.stderr) + + +def _fail(message: str) -> None: + print(f"[ERROR] diffusion prepare: {message}", file=sys.stderr) + raise SystemExit(1) + + +def _as_dict(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return value + return nested_namespace_to_dict(value) + + +def _select_module_name(cfg: Any, requested: str | None) -> str: + if requested: + get_module_config(cfg, requested) + return requested + + for module_name in ("pre_trainer", "post_trainer"): + try: + get_module_config(cfg, module_name) + return module_name + except Exception: + continue + + _fail("config must contain either modules.pre_trainer or modules.post_trainer") + raise AssertionError("unreachable") + + +def _is_placeholder(path: str | None) -> bool: + return not path or path.startswith("/path/to/") + + +def _require_path(path: str | None, description: str, *, kind: str = "any") -> None: + if _is_placeholder(path): + _fail(f"{description} is not configured: {path!r}") + + resolved = Path(path).expanduser() + if kind == "file" and not resolved.is_file(): + _fail(f"{description} file not found: {resolved}") + if kind == "dir" and not resolved.is_dir(): + _fail(f"{description} directory not found: {resolved}") + if kind == "any" and not resolved.exists(): + _fail(f"{description} path not found: {resolved}") + + _log(f"{description}: {resolved}") + + +def validate_diffusion_config(config_path: Path, module_name: str | None = None) -> None: + cfg = load_primus_config(config_path) + selected_module = _select_module_name(cfg, module_name) + module_cfg = get_module_config(cfg, selected_module) + + if getattr(module_cfg, "framework", None) != "diffusion": + _log(f"module {selected_module} framework is not diffusion; skipping") + return + + backend_args = DiffusionAdapter().convert_config(module_cfg.params) + model = _as_dict(backend_args.model) + dataset = _as_dict(backend_args.dataset) + trainer = _as_dict(backend_args.trainer) + + dataset_cfg = dataset.get("config", {}) + processor_cfg = dataset_cfg.get("processor_config", {}) + encoder_cfg = model.get("config", {}).get("encoder", {}) or model.get("encoder", {}) + + _require_path(dataset_cfg.get("dataset_path"), "dataset metadata", kind="file") + _require_path(dataset_cfg.get("data_folder"), "dataset media folder", kind="dir") + _require_path(processor_cfg.get("text_tokenizer"), "text tokenizer", kind="dir") + + model_cfg = model.get("config", {}) + _require_path(model_cfg.get("load_from_pretrained_path"), "DiT initialization checkpoint") + _require_path(encoder_cfg.get("t5_encoder"), "text encoder checkpoint", kind="file") + _require_path(encoder_cfg.get("autoencoder"), "VAE checkpoint", kind="file") + + _log( + "validated " + f"module={selected_module} stage={getattr(backend_args, 'stage', None)} " + f"trainer={trainer.get('name')} model={model.get('name')}" + ) + print("env.PREPARED=1") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Prepare Primus diffusion training environment") + parser.add_argument("--config", type=Path, required=True, help="Experiment YAML config") + parser.add_argument("--data_path", type=Path, required=False, help="Reserved for hook API compatibility") + parser.add_argument( + "--primus_path", type=Path, required=False, help="Reserved for hook API compatibility" + ) + parser.add_argument("--patch_args", type=Path, required=False, help="Reserved for hook API compatibility") + parser.add_argument("--backend_path", type=str, default=None, help="Unused; diffusion is in-tree") + parser.add_argument("--module_name", type=str, default=None, help="Override module name to validate") + args, _unknown = parser.parse_known_args() + + if os.environ.get("SKIP_PREPARE") == "1": + _log("SKIP_PREPARE=1; skipping validation") + return + + if args.backend_path: + _fail("diffusion is an in-tree backend and does not support --backend_path") + + validate_diffusion_config(args.config, module_name=args.module_name) + + +if __name__ == "__main__": + main() diff --git a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt new file mode 100644 index 000000000..a71bd1196 --- /dev/null +++ b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt @@ -0,0 +1,11 @@ +einops +loguru +safetensors +numpy +pillow +decord +imageio[ffmpeg] +pydantic +wandb +transformers==4.50.0 +peft==0.10.0 From 313e7d05d2eae82346d8026ca1338ea235bec45e Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 17 Jun 2026 07:01:08 +0000 Subject: [PATCH 02/30] docs(diffusion): unify single/multi-node torchrun launch in READMEs Co-authored-by: Cursor --- examples/diffusion/README.md | 46 +++++++++++++++++++++++++---- primus/backends/diffusion/README.md | 40 ++++++++++--------------- 2 files changed, 56 insertions(+), 30 deletions(-) diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md index ba4fb200a..edc238a9a 100644 --- a/examples/diffusion/README.md +++ b/examples/diffusion/README.md @@ -24,15 +24,42 @@ Expected layout: ## Run -Launch an MI355X example config with 8 GPUs: +Single-node and multi-node use the **same** launch command: one `torchrun` +invocation per node, parameterized by the standard rendezvous variables. The +defaults below give a single-node 8-GPU run. ```bash -torchrun --standalone --nproc_per_node=8 \ +# distributed knobs (defaults = single node, 8 GPUs) +export NNODES=${NNODES:-1} +export NODE_RANK=${NODE_RANK:-0} +export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} +export MASTER_PORT=${MASTER_PORT:-29500} +export GPUS_PER_NODE=${GPUS_PER_NODE:-8} + +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ -m primus.cli.main train pretrain \ --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml ``` -Override common paths and runtime options with environment variables: +For **multi-node**, run that exact command on every node with a shared +`MASTER_ADDR`/`MASTER_PORT` (a routable IP of node rank 0) and a distinct +`NODE_RANK` per node. For example, 2 nodes: + +```bash +# node 0 +NNODES=2 NODE_RANK=0 MASTER_ADDR= MASTER_PORT=29500 +# node 1 +NNODES=2 NODE_RANK=1 MASTER_ADDR= MASTER_PORT=29500 +``` + +World size is `NNODES * GPUS_PER_NODE`. Single node is just `NNODES=1`, +`NODE_RANK=0`, `MASTER_ADDR=127.0.0.1` (the defaults above). + +Override common paths and runtime options with environment variables (these +prepend to the same `torchrun` command): ```bash DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ @@ -40,7 +67,10 @@ DATA_FOLDER=/data/tiny-video-samples/data \ ATTENTION_BACKEND=sdpa \ SP_SIZE=1 \ MAX_STEPS=3 \ -torchrun --standalone --nproc_per_node=8 \ +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ -m primus.cli.main train pretrain \ --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml ``` @@ -48,14 +78,18 @@ torchrun --standalone --nproc_per_node=8 \ Use `SP_SIZE=4` or `SP_SIZE=8` to enable Ulysses sequence parallelism when the model head count supports it. -Post-train examples use the same public override shape under `modules.post_trainer`: +Post-train examples use the same public override shape under `modules.post_trainer` +and the same unified launch command (just `train posttrain`): ```bash INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B \ DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ DATA_FOLDER=/data/tiny-video-samples/data \ MAX_STEPS=3 \ -torchrun --standalone --nproc_per_node=8 \ +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ -m primus.cli.main train posttrain \ --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml ``` diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md index 896e3bb38..87b6328dd 100644 --- a/primus/backends/diffusion/README.md +++ b/primus/backends/diffusion/README.md @@ -126,38 +126,30 @@ for debugging. ## Launch Training runs through the Primus CLI (`primus.cli.main train `) -under `torchrun`. - -### Single node (e.g. 8 GPUs) - -```bash -torchrun --standalone --nproc_per_node=8 \ - -m primus.cli.main train pretrain --config /path/to/wan_config.yaml -``` - -Post-train examples use `modules.post_trainer` and the posttrain suite: +under `torchrun`. Single-node and multi-node share **one** launch command: the +same `torchrun` invocation runs on every node, parameterized by the standard +rendezvous variables. The defaults below give a single-node 8-GPU run. ```bash -torchrun --standalone --nproc_per_node=8 \ - -m primus.cli.main train posttrain --config /path/to/wan_posttrain_config.yaml -``` - -### Multi-node +# distributed knobs (defaults = single node, 8 GPUs) +export NNODES=${NNODES:-1} +export NODE_RANK=${NODE_RANK:-0} +export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} +export MASTER_PORT=${MASTER_PORT:-29500} +export GPUS_PER_NODE=${GPUS_PER_NODE:-8} -Run the same command on every node, sharing one rendezvous endpoint -(`MASTER_ADDR`/`MASTER_PORT`) and giving each node a distinct `--node_rank`: - -```bash -# on every node (NNODES total); MASTER_ADDR must be a routable IP of node rank 0 torchrun \ --nnodes="$NNODES" --node_rank="$NODE_RANK" \ - --master_addr="$MASTER_ADDR" --master_port=29577 \ - --nproc_per_node=8 \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ -m primus.cli.main train pretrain --config /path/to/wan_config.yaml ``` -The config and CLI are identical to single node; only the `torchrun` rendezvous -flags change. Total world size is `NNODES * nproc_per_node`. +- **Single node**: run as-is (the defaults above). +- **Multi-node**: run the same command on each node with a shared + `MASTER_ADDR`/`MASTER_PORT` (a routable IP of node rank 0) and a distinct + `NODE_RANK` per node. World size is `NNODES * GPUS_PER_NODE`. +- **Post-train**: identical command with `train posttrain` and a posttrain config. Useful runtime knobs: From 38830fea1db04c5752ec9f5d99d455d97fa49288 Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 17 Jun 2026 07:34:51 +0000 Subject: [PATCH 03/30] docs(diffusion): simplify unified run instructions --- examples/diffusion/README.md | 34 ++++------------------------------ 1 file changed, 4 insertions(+), 30 deletions(-) diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md index edc238a9a..16bb46622 100644 --- a/examples/diffusion/README.md +++ b/examples/diffusion/README.md @@ -24,44 +24,18 @@ Expected layout: ## Run -Single-node and multi-node use the **same** launch command: one `torchrun` -invocation per node, parameterized by the standard rendezvous variables. The -defaults below give a single-node 8-GPU run. +Single-node and multi-node use the same `torchrun` command. The defaults below +run a single-node 8-GPU smoke test. For multi-node, run the same command on +every node with a shared `MASTER_ADDR`/`MASTER_PORT` and a distinct `NODE_RANK`. +World size is `NNODES * GPUS_PER_NODE`. ```bash -# distributed knobs (defaults = single node, 8 GPUs) export NNODES=${NNODES:-1} export NODE_RANK=${NODE_RANK:-0} export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} export MASTER_PORT=${MASTER_PORT:-29500} export GPUS_PER_NODE=${GPUS_PER_NODE:-8} -torchrun \ - --nnodes="$NNODES" --node_rank="$NODE_RANK" \ - --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ - --nproc_per_node="$GPUS_PER_NODE" \ - -m primus.cli.main train pretrain \ - --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml -``` - -For **multi-node**, run that exact command on every node with a shared -`MASTER_ADDR`/`MASTER_PORT` (a routable IP of node rank 0) and a distinct -`NODE_RANK` per node. For example, 2 nodes: - -```bash -# node 0 -NNODES=2 NODE_RANK=0 MASTER_ADDR= MASTER_PORT=29500 -# node 1 -NNODES=2 NODE_RANK=1 MASTER_ADDR= MASTER_PORT=29500 -``` - -World size is `NNODES * GPUS_PER_NODE`. Single node is just `NNODES=1`, -`NODE_RANK=0`, `MASTER_ADDR=127.0.0.1` (the defaults above). - -Override common paths and runtime options with environment variables (these -prepend to the same `torchrun` command): - -```bash DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ DATA_FOLDER=/data/tiny-video-samples/data \ ATTENTION_BACKEND=sdpa \ From 5f94542d4bc24510596a2b3c45e74e457ed5d830 Mon Sep 17 00:00:00 2001 From: zirui Date: Thu, 18 Jun 2026 07:27:29 +0000 Subject: [PATCH 04/30] fix(diffusion): broadcast Wan timesteps per sample Co-authored-by: Cursor --- primus/backends/diffusion/models/wan/wan_dit.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/primus/backends/diffusion/models/wan/wan_dit.py b/primus/backends/diffusion/models/wan/wan_dit.py index 98b878b78..cca581d3b 100644 --- a/primus/backends/diffusion/models/wan/wan_dit.py +++ b/primus/backends/diffusion/models/wan/wan_dit.py @@ -440,7 +440,10 @@ def forward(self, x, t, context, seq_len, y=None, sp_group: Optional[dist.Proces # time embeddings if t.dim() == 1: - t = t.expand(t.size(0), seq_len) + # [B] -> [B, seq_len]: per-sample timestep broadcast across the sequence. + # Must unsqueeze before expand, otherwise the batch dim (B) is wrongly + # aligned with seq_len and breaks for batch_size > 1. + t = t.unsqueeze(1).expand(t.size(0), seq_len) with torch.amp.autocast("cuda", dtype=torch.float32): bt = t.size(0) t_flat = t.flatten() From b107de26aa0461801421620d31076473b019d6c7 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 30 Jun 2026 09:40:31 +0000 Subject: [PATCH 05/30] Remove unused pip packages and clean up README.md --- examples/diffusion/README.md | 20 ++++++++++--------- .../diffusion/requirements-diffusion.txt | 3 +-- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md index 16bb46622..eaabb373d 100644 --- a/examples/diffusion/README.md +++ b/examples/diffusion/README.md @@ -24,10 +24,7 @@ Expected layout: ## Run -Single-node and multi-node use the same `torchrun` command. The defaults below -run a single-node 8-GPU smoke test. For multi-node, run the same command on -every node with a shared `MASTER_ADDR`/`MASTER_PORT` and a distinct `NODE_RANK`. -World size is `NNODES * GPUS_PER_NODE`. +Set the shared `torchrun` environment first: ```bash export NNODES=${NNODES:-1} @@ -35,10 +32,15 @@ export NODE_RANK=${NODE_RANK:-0} export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} export MASTER_PORT=${MASTER_PORT:-29500} export GPUS_PER_NODE=${GPUS_PER_NODE:-8} +``` + +### Pretrain +```bash DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ DATA_FOLDER=/data/tiny-video-samples/data \ ATTENTION_BACKEND=sdpa \ +CONFIG=examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml \ SP_SIZE=1 \ MAX_STEPS=3 \ torchrun \ @@ -46,29 +48,29 @@ torchrun \ --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ --nproc_per_node="$GPUS_PER_NODE" \ -m primus.cli.main train pretrain \ - --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml + --config "$CONFIG" ``` Use `SP_SIZE=4` or `SP_SIZE=8` to enable Ulysses sequence parallelism when the model head count supports it. -Post-train examples use the same public override shape under `modules.post_trainer` -and the same unified launch command (just `train posttrain`): +### Posttrain ```bash INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B \ DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ DATA_FOLDER=/data/tiny-video-samples/data \ +CONFIG=examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml \ MAX_STEPS=3 \ torchrun \ --nnodes="$NNODES" --node_rank="$NODE_RANK" \ --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ --nproc_per_node="$GPUS_PER_NODE" \ -m primus.cli.main train posttrain \ - --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml + --config "$CONFIG" ``` The MI355X configs use Primus-style override sections such as `training`, `data`, `parallelism`, `optimizer`, `runtime`, and `metrics`. The diffusion adapter normalizes those sections into the Wan model/dataset/trainer -arguments at runtime. +arguments at runtime. \ No newline at end of file diff --git a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt index a71bd1196..d392d1e1b 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt +++ b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt @@ -5,7 +5,6 @@ numpy pillow decord imageio[ffmpeg] -pydantic wandb +pydantic transformers==4.50.0 -peft==0.10.0 From b3cf03478ecc851a22405376cab5f422241ab14d Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 1 Jul 2026 03:57:38 +0000 Subject: [PATCH 06/30] fix launch command --- examples/diffusion/README.md | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md index eaabb373d..126c54056 100644 --- a/examples/diffusion/README.md +++ b/examples/diffusion/README.md @@ -40,15 +40,14 @@ export GPUS_PER_NODE=${GPUS_PER_NODE:-8} DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ DATA_FOLDER=/data/tiny-video-samples/data \ ATTENTION_BACKEND=sdpa \ -CONFIG=examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml \ SP_SIZE=1 \ -MAX_STEPS=3 \ +MAX_STEPS=10 \ torchrun \ --nnodes="$NNODES" --node_rank="$NODE_RANK" \ --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ --nproc_per_node="$GPUS_PER_NODE" \ -m primus.cli.main train pretrain \ - --config "$CONFIG" + --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml ``` Use `SP_SIZE=4` or `SP_SIZE=8` to enable Ulysses sequence parallelism @@ -60,14 +59,13 @@ when the model head count supports it. INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B \ DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ DATA_FOLDER=/data/tiny-video-samples/data \ -CONFIG=examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml \ -MAX_STEPS=3 \ +MAX_STEPS=10 \ torchrun \ --nnodes="$NNODES" --node_rank="$NODE_RANK" \ --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ --nproc_per_node="$GPUS_PER_NODE" \ -m primus.cli.main train posttrain \ - --config "$CONFIG" + --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml ``` The MI355X configs use Primus-style override sections such as `training`, From 38a04c1fb9de262781731f81b74c301a8174a11c Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 1 Jul 2026 11:22:48 +0000 Subject: [PATCH 07/30] set aiter as default attention backend --- .../diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml | 2 +- examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml | 2 +- examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml | 2 +- examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml | 2 +- primus/backends/diffusion/argument_builder.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml index dcb8c2482..c7f93c421 100644 --- a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml +++ b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml @@ -49,5 +49,5 @@ modules: weight_decay: 0.01 runtime: - attention_backend: ${ATTENTION_BACKEND:sdpa} + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml index 1411c054a..5f17780f1 100644 --- a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml +++ b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml @@ -49,5 +49,5 @@ modules: weight_decay: 0.01 runtime: - attention_backend: ${ATTENTION_BACKEND:sdpa} + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml index a12fba375..203aa13ff 100644 --- a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml +++ b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml @@ -49,5 +49,5 @@ modules: weight_decay: 0.01 runtime: - attention_backend: ${ATTENTION_BACKEND:sdpa} + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml index 6757111be..7eed61bdd 100644 --- a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml +++ b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml @@ -49,5 +49,5 @@ modules: weight_decay: 0.01 runtime: - attention_backend: ${ATTENTION_BACKEND:sdpa} + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} report_to: none diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py index 2e1c6e9d2..d907f65a5 100644 --- a/primus/backends/diffusion/argument_builder.py +++ b/primus/backends/diffusion/argument_builder.py @@ -53,7 +53,7 @@ class WanArgBuilder: "per_device_eval_batch_size": 1, "gradient_accumulation_steps": 1, "gradient_checkpointing": True, - "attention_backend": "sdpa", + "attention_backend": "flash_attn_aiter", "learning_rate": 1.0e-5, "lr_scheduler_type": "constant", "warmup_steps": 0, From 696748a17e19a36a99675b25206f4bfa4b9d11f0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 6 Jul 2026 11:13:08 +0000 Subject: [PATCH 08/30] fix: add missing newline at end of examples/diffusion/README.md --- examples/diffusion/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md index 126c54056..997c9ad5f 100644 --- a/examples/diffusion/README.md +++ b/examples/diffusion/README.md @@ -71,4 +71,4 @@ torchrun \ The MI355X configs use Primus-style override sections such as `training`, `data`, `parallelism`, `optimizer`, `runtime`, and `metrics`. The diffusion adapter normalizes those sections into the Wan model/dataset/trainer -arguments at runtime. \ No newline at end of file +arguments at runtime. From ed3327f94ed80b2365257132b9fbad3bdf942f1a Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 03:21:08 +0000 Subject: [PATCH 09/30] set aiter as default attention backend & add set seperated_timestep=True for wan2.2-5B --- examples/diffusion/README.md | 2 +- primus/backends/diffusion/README.md | 4 ++-- primus/backends/diffusion/argument_builder.py | 2 ++ .../diffusion/diffusion_pretrain_trainer.py | 1 - .../diffusion/models/wan/train_pipeline.py | 2 -- primus/backends/diffusion/models/wan/vae2_2.py | 16 +++++++++++----- .../configs/models/diffusion/wan2.2_ti2v_5b.yaml | 4 ++-- 7 files changed, 18 insertions(+), 13 deletions(-) diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md index 997c9ad5f..2df6093f7 100644 --- a/examples/diffusion/README.md +++ b/examples/diffusion/README.md @@ -39,7 +39,7 @@ export GPUS_PER_NODE=${GPUS_PER_NODE:-8} ```bash DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ DATA_FOLDER=/data/tiny-video-samples/data \ -ATTENTION_BACKEND=sdpa \ +ATTENTION_BACKEND=flash_attn_aiter \ SP_SIZE=1 \ MAX_STEPS=10 \ torchrun \ diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md index 87b6328dd..9b7db9704 100644 --- a/primus/backends/diffusion/README.md +++ b/primus/backends/diffusion/README.md @@ -153,7 +153,7 @@ torchrun \ Useful runtime knobs: -- `trainer.args.attention_backend`: `sdpa` is the most portable ROCm option. +- `trainer.args.attention_backend`: defaults to `flash_attn_aiter` for Wan training; use `sdpa` as the portable fallback or baseline. - `trainer.args.sp_size`: Ulysses sequence parallel size. It must divide the model attention head count; for example Wan2.1-1.3B supports `sp_size=4` but not `sp_size=8`. @@ -218,7 +218,7 @@ modules: sp_size: 1 dp_replicate: 1 runtime: - attention_backend: sdpa + attention_backend: flash_attn_aiter report_to: none ``` diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py index d907f65a5..9814ed140 100644 --- a/primus/backends/diffusion/argument_builder.py +++ b/primus/backends/diffusion/argument_builder.py @@ -173,6 +173,7 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, ("save_steps",): ("save_steps",), ("run_name",): ("run_name",), ("num_train_epochs",): ("num_train_epochs",), + ("dataloader_num_workers",): ("dataloader_num_workers",), ("resume_from_checkpoint",): ("resume_from_checkpoint",), } for source_path, target_path in training_map.items(): @@ -228,6 +229,7 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, ("attention_backend",): ("attention_backend",), ("report_to",): ("report_to",), ("seed",): ("seed",), + ("fsdp2_reshard_after_forward",): ("fsdp2_reshard_after_forward",), } for source_path, target_path in runtime_map.items(): value = self._get_any(runtime, *source_path) diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index 1f1aa24c5..45ffb3a0e 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -119,7 +119,6 @@ def cleanup(self, on_error: bool = False): import torch.distributed as dist if dist.is_available() and dist.is_initialized(): - dist.barrier() dist.destroy_process_group() except Exception: pass diff --git a/primus/backends/diffusion/models/wan/train_pipeline.py b/primus/backends/diffusion/models/wan/train_pipeline.py index 5d6a2dd62..38bfc87ab 100644 --- a/primus/backends/diffusion/models/wan/train_pipeline.py +++ b/primus/backends/diffusion/models/wan/train_pipeline.py @@ -39,8 +39,6 @@ class WanFlowMatchTrainPipeline: def __init__(self, cfg: Optional[WanFlowMatchTrainPipelineConfig] = None): self.cfg = cfg or WanFlowMatchTrainPipelineConfig() - # Internal counter for optional per-step profiling logs. - self._profile_step: int = 0 @staticmethod def _encode_prompt(text_encoder: torch.nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor): diff --git a/primus/backends/diffusion/models/wan/vae2_2.py b/primus/backends/diffusion/models/wan/vae2_2.py index 47eab810a..05e4db1d0 100644 --- a/primus/backends/diffusion/models/wan/vae2_2.py +++ b/primus/backends/diffusion/models/wan/vae2_2.py @@ -680,7 +680,7 @@ def __init__( vae_pth=None, dim_mult=[1, 2, 4, 4], temperal_downsample=[False, True, True], - dtype=torch.float, + dtype=torch.bfloat16, device="cuda", ): self.dtype = dtype @@ -812,16 +812,22 @@ def __init__( # --- Wan trainer compatibility helpers (non-upstream) --- def to(self, *args, **kwargs): - self.model.to(*args, **kwargs) - if "dtype" in kwargs and kwargs["dtype"] is not None: - self.dtype = kwargs["dtype"] + dtype = kwargs.get("dtype", None) + if dtype is None: + dtype = self.dtype + else: + self.dtype = dtype device = kwargs.get("device", None) if device is None and len(args) > 0: device = args[0] + if device is not None: + self.model.to(device=device, dtype=dtype) + else: + self.model.to(*args, **kwargs) if device is not None: self.device = device dev = next(self.model.parameters()).device - self.scale = [t.to(device=dev) for t in self.scale] + self.scale = [t.to(device=dev, dtype=dtype) for t in self.scale] return self def eval(self): diff --git a/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml b/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml index a3fa8ee9f..a9dd68567 100644 --- a/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml +++ b/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml @@ -4,8 +4,8 @@ model: load_from_pretrained_path: ${PRETRAINED_PATH:/models/Wan2.2-TI2V-5B} config: model_type: ti2v - seperated_timestep: false - fuse_vae_embedding_in_latents: false + seperated_timestep: true + fuse_vae_embedding_in_latents: true trainable_modules: dit vae_type: wan_video_vae_38 encoder: From e80b7033286f597560821e2b24953b13bf5d0449 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 03:46:10 +0000 Subject: [PATCH 10/30] fix(diffusion): address Wan code quality findings Co-authored-by: Cursor --- .../backends/diffusion/models/wan/vae2_1.py | 16 +++++++++---- .../backends/diffusion/models/wan/vae2_2.py | 24 ++++++++++++++----- .../backends/diffusion/models/wan/wan_dit.py | 4 +++- 3 files changed, 33 insertions(+), 11 deletions(-) diff --git a/primus/backends/diffusion/models/wan/vae2_1.py b/primus/backends/diffusion/models/wan/vae2_1.py index fa0b2b78a..ac6b44760 100644 --- a/primus/backends/diffusion/models/wan/vae2_1.py +++ b/primus/backends/diffusion/models/wan/vae2_1.py @@ -108,7 +108,9 @@ def __init__(self, dim, mode): else: self.resample = nn.Identity() - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] b, c, t, h, w = x.size() if self.mode == "upsample3d": if feat_cache is not None: @@ -189,7 +191,9 @@ def __init__(self, in_dim, out_dim, dropout=0.0): ) self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] h = self.shortcut(x) for layer in self.residual: if isinstance(layer, CausalConv3d) and feat_cache is not None: @@ -284,7 +288,9 @@ def __init__( CausalConv3d(out_dim, z_dim, 3, padding=1), ) - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() @@ -376,7 +382,9 @@ def __init__( CausalConv3d(out_dim, 3, 3, padding=1), ) - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() diff --git a/primus/backends/diffusion/models/wan/vae2_2.py b/primus/backends/diffusion/models/wan/vae2_2.py index 05e4db1d0..05998db45 100644 --- a/primus/backends/diffusion/models/wan/vae2_2.py +++ b/primus/backends/diffusion/models/wan/vae2_2.py @@ -107,7 +107,9 @@ def __init__(self, dim, mode): else: self.resample = nn.Identity() - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] b, c, t, h, w = x.size() if self.mode == "upsample3d": if feat_cache is not None: @@ -188,7 +190,9 @@ def __init__(self, in_dim, out_dim, dropout=0.0): ) self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] h = self.shortcut(x) for layer in self.residual: if isinstance(layer, CausalConv3d) and feat_cache is not None: @@ -345,7 +349,9 @@ def __init__(self, in_dim, out_dim, dropout, mult, temperal_downsample=False, do downsamples.append(Resample(out_dim, mode=mode)) self.downsamples = nn.Sequential(*downsamples) - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] x_copy = x.clone() for module in self.downsamples: x = module(x, feat_cache, feat_idx) @@ -373,7 +379,9 @@ def __init__(self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_f upsamples.append(Resample(out_dim, mode=mode)) self.upsamples = nn.Sequential(*upsamples) - def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False): + if feat_idx is None: + feat_idx = [0] x_main = x.clone() for module in self.upsamples: x_main = module(x_main, feat_cache, feat_idx) @@ -424,7 +432,9 @@ def __init__( CausalConv3d(out_dim, z_dim, 3, padding=1), ) - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() @@ -507,7 +517,9 @@ def __init__( CausalConv3d(out_dim, 12, 3, padding=1), ) - def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False): + if feat_idx is None: + feat_idx = [0] if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() diff --git a/primus/backends/diffusion/models/wan/wan_dit.py b/primus/backends/diffusion/models/wan/wan_dit.py index cca581d3b..2e545e9a4 100644 --- a/primus/backends/diffusion/models/wan/wan_dit.py +++ b/primus/backends/diffusion/models/wan/wan_dit.py @@ -186,7 +186,9 @@ def forward(self, x, seq_lens, grid_sizes, freqs, sp_group: Optional[dist.Proces freqs(Tensor): Rope freqs sp_group: Ulysses SP process group (None = no SP) """ - b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + b, s = x.shape[:2] + n = self.num_heads + d = self.head_dim x_ = x.to(dtype=self.q.weight.dtype) q = self.norm_q(self.q(x_)).view(b, s, n, d) From 8ec64195a85b6e452dc76ac6918967acb2b19a35 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 04:04:44 +0000 Subject: [PATCH 11/30] fix(diffusion): resolve remaining code quality findings Co-authored-by: Cursor --- .../backends/diffusion/diffusion_adapter.py | 2 ++ .../diffusion/diffusion_pretrain_trainer.py | 8 +++---- .../diffusion/models/wan/train_pipeline.py | 4 +++- .../backends/diffusion/models/wan/wan_dit.py | 24 +++++++++++++++++-- 4 files changed, 31 insertions(+), 7 deletions(-) diff --git a/primus/backends/diffusion/diffusion_adapter.py b/primus/backends/diffusion/diffusion_adapter.py index 262b68b64..202c3914a 100644 --- a/primus/backends/diffusion/diffusion_adapter.py +++ b/primus/backends/diffusion/diffusion_adapter.py @@ -36,6 +36,7 @@ def setup_backend_path(self, backend_path=None) -> str: try: log_rank_0(f"[Primus:Diffusion] using in-tree backend package -> {resolved_str}") except Exception: + # Best-effort startup logging should not block backend setup. pass return resolved_str @@ -49,6 +50,7 @@ def convert_config(self, params: Any): try: log_rank_0("[Primus:DiffusionAdapter] Converted Primus module params -> Wan args") except Exception: + # Standalone prepare hooks may run before the Primus logger is bound. pass return wan_args diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index 45ffb3a0e..586a2d9ba 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -112,13 +112,13 @@ def cleanup(self, on_error: bool = False): if getattr(wandb, "run", None) is not None: wandb.finish(exit_code=1 if on_error else 0) - except Exception: - pass + except Exception as exc: + log_rank_0(f"[Primus:Diffusion] wandb cleanup failed: {exc}") try: import torch.distributed as dist if dist.is_available() and dist.is_initialized(): dist.destroy_process_group() - except Exception: - pass + except Exception as exc: + log_rank_0(f"[Primus:Diffusion] distributed cleanup failed: {exc}") diff --git a/primus/backends/diffusion/models/wan/train_pipeline.py b/primus/backends/diffusion/models/wan/train_pipeline.py index 38bfc87ab..e48d40987 100644 --- a/primus/backends/diffusion/models/wan/train_pipeline.py +++ b/primus/backends/diffusion/models/wan/train_pipeline.py @@ -219,7 +219,9 @@ def compute_loss( components.text_encoder.eval() try: components.vae.eval() - except Exception: + except AttributeError: + # Some VAE compatibility wrappers may not expose eval(); real eval + # failures should still surface. pass with torch.no_grad(): diff --git a/primus/backends/diffusion/models/wan/wan_dit.py b/primus/backends/diffusion/models/wan/wan_dit.py index 2e545e9a4..0d975b250 100644 --- a/primus/backends/diffusion/models/wan/wan_dit.py +++ b/primus/backends/diffusion/models/wan/wan_dit.py @@ -217,13 +217,29 @@ def forward(self, x, seq_lens, grid_sizes, freqs, sp_group: Optional[dist.Proces class WanCrossAttention(WanSelfAttention): - def forward(self, x, context, context_lens): + def forward( + self, + x, + seq_lens=None, + grid_sizes=None, + freqs=None, + sp_group: Optional[dist.ProcessGroup] = None, + context=None, + context_lens=None, + ): r""" Args: x(Tensor): Shape [B, L1, C] + seq_lens/grid_sizes/freqs/sp_group: accepted for WanSelfAttention + signature compatibility; cross-attention uses text context only. context(Tensor): Shape [B, L2, C] context_lens(Tensor): Shape [B] """ + if context is None: + # Backward-compatible positional form: forward(x, context, context_lens). + context = seq_lens + context_lens = grid_sizes if context_lens is None else context_lens + b, n, d = x.size(0), self.num_heads, self.head_dim x = x.to(dtype=self.q.weight.dtype) @@ -304,7 +320,11 @@ def forward( # cross-attention + FFN (no SP communication needed: # each rank's visual queries attend to full text keys independently) def cross_attn_ffn(x_, context_, context_lens_, e_): - x_ = x_ + self.cross_attn(self.norm3(x_), context_, context_lens_) + x_ = x_ + self.cross_attn( + self.norm3(x_), + context=context_, + context_lens=context_lens_, + ) ffn_in = self.norm2(x_) * (1 + e_[4].squeeze(2)) + e_[3].squeeze(2) y_ = self.ffn(ffn_in) x_ = x_ + y_ * e_[5].squeeze(2) From 176545a859321e004478d1d951acf134b9eb65a1 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 05:11:02 +0000 Subject: [PATCH 12/30] fix(diffusion): address Copilot PR review findings Co-authored-by: Cursor --- primus/backends/diffusion/data/__init__.py | 7 + .../diffusion/data/registrations/__init__.py | 7 + .../diffusion/data/registrations/wan.py | 168 ++++++++++++++++++ .../diffusion/distributed/checkpoint.py | 18 +- .../models/wan/configuration_wanvideo.py | 12 +- primus/backends/diffusion/utils/data_utils.py | 9 +- .../backends/diffusion/utils/train_utils.py | 3 +- .../diffusion/utils/vision_process.py | 30 ++-- 8 files changed, 232 insertions(+), 22 deletions(-) create mode 100644 primus/backends/diffusion/data/__init__.py create mode 100644 primus/backends/diffusion/data/registrations/__init__.py create mode 100644 primus/backends/diffusion/data/registrations/wan.py diff --git a/primus/backends/diffusion/data/__init__.py b/primus/backends/diffusion/data/__init__.py new file mode 100644 index 000000000..18ae84901 --- /dev/null +++ b/primus/backends/diffusion/data/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Data builders for the diffusion backend.""" diff --git a/primus/backends/diffusion/data/registrations/__init__.py b/primus/backends/diffusion/data/registrations/__init__.py new file mode 100644 index 000000000..b5a3705a2 --- /dev/null +++ b/primus/backends/diffusion/data/registrations/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Registered diffusion dataset builders.""" diff --git a/primus/backends/diffusion/data/registrations/wan.py b/primus/backends/diffusion/data/registrations/wan.py new file mode 100644 index 000000000..f487ec8ca --- /dev/null +++ b/primus/backends/diffusion/data/registrations/wan.py @@ -0,0 +1,168 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image +from torch.utils.data import Dataset +from torchvision.transforms import functional as F +from transformers import AutoTokenizer + +from primus.backends.diffusion.utils.vision_process import fetch_video + + +class WanVideoProcessor: + """Tokenize prompts and normalize video tensors for Wan training.""" + + def __init__(self, config: dict[str, Any]): + self.config = config + self.max_text_length = int(config.get("max_text_length", 512)) + tokenizer_path = config.get("text_tokenizer") + if not tokenizer_path: + raise ValueError("Wan dataset processor requires `text_tokenizer`.") + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) + + extra_kwargs = config.get("extra_kwargs", {}) or {} + size = extra_kwargs.get("size", {}) or {} + self.height = int(size.get("height", 480)) + self.width = int(size.get("width", 832)) + self.image_mean = torch.tensor(extra_kwargs.get("image_mean", [0.5, 0.5, 0.5])).view(3, 1, 1, 1) + self.image_std = torch.tensor(extra_kwargs.get("image_std", [0.5, 0.5, 0.5])).view(3, 1, 1, 1) + + def tokenize(self, prompt: str) -> dict[str, torch.Tensor]: + encoded = self.tokenizer( + prompt, + max_length=self.max_text_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ) + return { + "input_ids": encoded["input_ids"].squeeze(0).long(), + "attention_mask": encoded["attention_mask"].squeeze(0).long(), + } + + def normalize_video(self, video_tchw: torch.Tensor) -> torch.Tensor: + if video_tchw.ndim != 4: + raise ValueError(f"Expected video tensor [T,C,H,W], got shape={tuple(video_tchw.shape)}") + video_tchw = video_tchw[:, :3].float() + if video_tchw.max() > 2: + video_tchw = video_tchw / 255.0 + video_tchw = F.resize(video_tchw, [self.height, self.width], antialias=True) + video_cthw = video_tchw.permute(1, 0, 2, 3).contiguous() + return (video_cthw - self.image_mean) / self.image_std + + def prepare_batch(self, *, batch: dict[str, Any], device: torch.device, dtype: torch.dtype) -> dict[str, Any]: + return batch + + +class WanVideoDataset(Dataset): + def __init__(self, config: dict[str, Any], processor: WanVideoProcessor): + self.config = config + self.processor = processor + self.dataset_path = Path(config.get("dataset_path", "")) + self.data_folder = Path(config.get("data_folder", "")) + self.frame_num = int(config.get("frame_num", 81)) + self.video_backend = str(config.get("video_backend", "imageio")).lower() + + if not self.dataset_path.exists(): + raise FileNotFoundError(f"Wan dataset metadata not found: {self.dataset_path}") + with self.dataset_path.open(encoding="utf-8") as f: + self.samples = [json.loads(line) for line in f if line.strip()] + if not self.samples: + raise ValueError(f"Wan dataset metadata is empty: {self.dataset_path}") + + def __len__(self) -> int: + return len(self.samples) + + def _resolve_video_path(self, value: str) -> str: + if value.startswith(("http://", "https://", "file://")): + return value + path = Path(value) + if not path.is_absolute() and self.data_folder: + path = self.data_folder / path + return str(path) + + def _sample_video(self, video_tchw: torch.Tensor) -> torch.Tensor: + total_frames = int(video_tchw.shape[0]) + if total_frames <= 0: + raise ValueError("Video contains no frames.") + idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long() + return video_tchw[idx] + + def _read_video_imageio(self, path: str) -> torch.Tensor: + import imageio.v3 as iio + + frames = [] + for frame in iio.imiter(path): + image = Image.fromarray(frame).convert("RGB") + frames.append(torch.from_numpy(np.asarray(image).copy())) + if not frames: + raise ValueError(f"Video contains no frames: {path}") + return torch.stack(frames).permute(0, 3, 1, 2) + + def _read_video_decord(self, path: str) -> torch.Tensor: + import decord + + vr = decord.VideoReader(path) + total_frames = len(vr) + idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long().tolist() + return torch.from_numpy(vr.get_batch(idx).asnumpy()).permute(0, 3, 1, 2) + + def _read_video(self, path: str) -> torch.Tensor: + if self.video_backend == "imageio": + video = self._sample_video(self._read_video_imageio(path)) + elif self.video_backend == "decord": + video = self._read_video_decord(path) + else: + video = fetch_video( + { + "video": path, + "nframes": self.frame_num, + "resized_height": self.processor.height, + "resized_width": self.processor.width, + } + ) + return self.processor.normalize_video(video) + + def __getitem__(self, index: int) -> dict[str, Any]: + sample = self.samples[index] + prompt = sample.get("prompt") or sample.get("text") or sample.get("caption") or "" + video_key = sample.get("video") or sample.get("video_path") + if not video_key: + raise ValueError(f"Wan dataset sample missing `video`: index={index}") + + item = self.processor.tokenize(str(prompt)) + item["video"] = self._read_video(self._resolve_video_path(str(video_key))) + if "seed" in sample: + item["seed"] = int(sample["seed"]) + return item + + @staticmethod + def get_collator(): + def collate(samples: list[dict[str, Any]]) -> dict[str, Any]: + batch = { + "video": torch.stack([sample["video"] for sample in samples]), + "input_ids": torch.stack([sample["input_ids"] for sample in samples]), + "attention_mask": torch.stack([sample["attention_mask"] for sample in samples]), + } + if any("seed" in sample for sample in samples): + batch["seed"] = torch.tensor([sample.get("seed", 0) for sample in samples], dtype=torch.long) + return batch + + return collate + + +def build_wan_dataset(config: dict[str, Any]): + processor = WanVideoProcessor(config.get("processor_config", {}) or {}) + dataset = WanVideoDataset(config, processor) + return dataset, processor diff --git a/primus/backends/diffusion/distributed/checkpoint.py b/primus/backends/diffusion/distributed/checkpoint.py index 4a4828c6b..9c4047f2f 100644 --- a/primus/backends/diffusion/distributed/checkpoint.py +++ b/primus/backends/diffusion/distributed/checkpoint.py @@ -12,6 +12,7 @@ from __future__ import annotations +import json import os from typing import Any, Dict, Optional @@ -31,6 +32,9 @@ from .mesh import _ensure_process_group +_META_FILENAME = "meta.json" + + def save_checkpoint_dtcp( model: torch.nn.Module, optimizer: Optional[torch.optim.Optimizer], @@ -54,14 +58,17 @@ def save_checkpoint_dtcp( if optimizer is not None: optim_state = get_optimizer_state_dict(model, optimizer, options=optim_state_options) + meta = {"epoch": epoch, "step": step, **(additional_data or {})} state_dict = { "model": model_state, **({"optimizer": optim_state} if optim_state is not None else {}), - "meta": {"epoch": epoch, "step": step, **(additional_data or {})}, + "meta": meta, } save(state_dict, FileSystemWriter(path)) if (not dist.is_initialized()) or dist.get_rank() == 0: + with open(os.path.join(path, _META_FILENAME), "w", encoding="utf-8") as f: + json.dump(meta, f, indent=2, sort_keys=True) logger.info(f"Saved DTCP checkpoint to {path}") @@ -90,6 +97,7 @@ def load_checkpoint_dtcp( state_dict = { "model": model_state, **({"optimizer": optim_state} if optim_state is not None else {}), + "meta": {}, } load(state_dict, FileSystemReader(path)) @@ -98,4 +106,10 @@ def load_checkpoint_dtcp( if optimizer is not None and optim_state is not None: set_optimizer_state_dict(model, optimizer, optim_state, options=optim_state_options) - return state_dict.get("meta", {}) + meta = state_dict.get("meta", {}) + if not meta: + meta_path = os.path.join(path, _META_FILENAME) + if os.path.exists(meta_path): + with open(meta_path, encoding="utf-8") as f: + meta = json.load(f) + return meta diff --git a/primus/backends/diffusion/models/wan/configuration_wanvideo.py b/primus/backends/diffusion/models/wan/configuration_wanvideo.py index dcfb1f1d6..f7d3e4506 100644 --- a/primus/backends/diffusion/models/wan/configuration_wanvideo.py +++ b/primus/backends/diffusion/models/wan/configuration_wanvideo.py @@ -42,7 +42,7 @@ def __init__( dit_has_image_pos_emb: bool = False, dit_has_ref_conv: bool = False, trainable_modules=None, - seperated_timestep: bool = True, + seperated_timestep: bool | None = None, require_clip_embedding: bool = False, require_vae_embedding: bool = False, fuse_vae_embedding_in_latents: bool = True, @@ -50,6 +50,13 @@ def __init__( vae_type: str = "wan_video_vae_38", **kwargs, ): + separated_timestep = kwargs.pop("separated_timestep", None) + if separated_timestep is not None and seperated_timestep is not None: + if bool(separated_timestep) != bool(seperated_timestep): + raise ValueError("`separated_timestep` and `seperated_timestep` must match when both are set.") + if seperated_timestep is None: + seperated_timestep = True if separated_timestep is None else bool(separated_timestep) + # DiT configuration self.vae_type = vae_type self.dit_hidden_size = dit_hidden_size @@ -67,7 +74,8 @@ def __init__( self.dit_has_image_pos_emb = dit_has_image_pos_emb self.dit_has_ref_conv = dit_has_ref_conv - self.seperated_timestep = seperated_timestep + self.seperated_timestep = bool(seperated_timestep) + self.separated_timestep = self.seperated_timestep self.require_clip_embedding = require_clip_embedding self.require_vae_embedding = require_vae_embedding self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents diff --git a/primus/backends/diffusion/utils/data_utils.py b/primus/backends/diffusion/utils/data_utils.py index de149c48b..4e78671da 100644 --- a/primus/backends/diffusion/utils/data_utils.py +++ b/primus/backends/diffusion/utils/data_utils.py @@ -16,19 +16,14 @@ def smart_nframes( total_frames: int, video_fps: int | float, - fps: int, + fps: int | float, ) -> int: """calculate the number of frames for video used for model inputs. Args: - ele (dict): a dict contains the configuration of video. - support either `fps` or `nframes`: - - nframes: the number of frames to extract for model inputs. - - fps: the fps to extract frames for model inputs. - - min_frames: the minimum number of frames of the video, only used when fps is provided. - - max_frames: the maximum number of frames of the video, only used when fps is provided. total_frames (int): the original total number of frames of the video. video_fps (int | float): the original fps of the video. + fps (int | float): the target sampling fps for model inputs. Raises: ValueError: nframes should in interval [FRAME_FACTOR, total_frames]. diff --git a/primus/backends/diffusion/utils/train_utils.py b/primus/backends/diffusion/utils/train_utils.py index 71deede8f..315ee9e67 100644 --- a/primus/backends/diffusion/utils/train_utils.py +++ b/primus/backends/diffusion/utils/train_utils.py @@ -171,7 +171,8 @@ def convert_state_dict_keys_to_single_str(state_dict, with_shape=True): if with_shape: shape = "_".join(map(str, list(value.shape))) keys.append(key + ":" + shape) - keys.append(key) + else: + keys.append(key) elif isinstance(value, dict): keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape)) keys.sort() diff --git a/primus/backends/diffusion/utils/vision_process.py b/primus/backends/diffusion/utils/vision_process.py index ea2aebd36..84f785569 100644 --- a/primus/backends/diffusion/utils/vision_process.py +++ b/primus/backends/diffusion/utils/vision_process.py @@ -42,10 +42,20 @@ FPS_MIN_FRAMES = 4 FPS_MAX_FRAMES = 768 MAX_NUM_WORKERS_FETCH_VIDEO = 8 +REMOTE_FETCH_TIMEOUT = (5, 30) MODEL_SEQ_LEN = int(float(os.environ.get("MODEL_SEQ_LEN", 128000))) logger = logging.getLogger(__name__) +VideoMetadata = Dict[str, Any] +VideoReaderOutput = Tuple[torch.Tensor, VideoMetadata, float] +FetchVideoOutput = Union[ + torch.Tensor, + Tuple[torch.Tensor, VideoMetadata], + Tuple[torch.Tensor, float], + Tuple[Tuple[torch.Tensor, VideoMetadata], float], +] + def round_by_factor(number: int, factor: int) -> int: """Returns the closest integer to 'number' that is divisible by 'factor'.""" @@ -112,7 +122,7 @@ def fetch_image(ele: Dict[str, Union[str, Image.Image]], image_patch_size: int = if isinstance(image, Image.Image): image_obj = image elif image.startswith("http://") or image.startswith("https://"): - with requests.get(image, stream=True) as response: + with requests.get(image, stream=True, timeout=REMOTE_FETCH_TIMEOUT) as response: response.raise_for_status() with BytesIO(response.content) as bio: image_obj = copy.deepcopy(Image.open(bio)) @@ -196,7 +206,7 @@ def smart_nframes( def _read_video_torchvision( ele: Dict[str, Any], -) -> Tuple[torch.Tensor, float]: +) -> VideoReaderOutput: """read video using torchvision.io.read_video Args: @@ -206,7 +216,7 @@ def _read_video_torchvision( - video_start: the start time of video. - video_end: the end time of video. Returns: - torch.Tensor: the video tensor with shape (T, C, H, W). + Tuple of video tensor (T, C, H, W), metadata, and sampled FPS. """ video_path = ele["video"] if version.parse(torchvision.__version__) < version.parse("0.19.0"): @@ -308,7 +318,7 @@ def calculate_video_frame_range( def _read_video_decord( ele: Dict[str, Any], -) -> Tuple[torch.Tensor, float]: +) -> VideoReaderOutput: """read video using decord.VideoReader Args: @@ -318,12 +328,12 @@ def _read_video_decord( - video_start: the start time of video. - video_end: the end time of video. Returns: - torch.Tensor: the video tensor with shape (T, C, H, W). + Tuple of video tensor (T, C, H, W), metadata, and sampled FPS. """ import decord video_path = ele["video"] - time.time() + st = time.time() vr = decord.VideoReader(video_path) total_frames, video_fps = len(vr), vr.get_avg_fps() start_frame, end_frame, total_frames = calculate_video_frame_range( @@ -335,7 +345,7 @@ def _read_video_decord( idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist() video = vr.get_batch(idx).asnumpy() video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format - # logger.info(f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + logger.info(f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") sample_fps = nframes / max(total_frames, 1e-6) * video_fps video_metadata = dict( @@ -355,7 +365,7 @@ def is_torchcodec_available() -> bool: def _read_video_torchcodec( ele: Dict[str, Any], -) -> Tuple[torch.Tensor, float]: +) -> VideoReaderOutput: """read video using torchcodec.decoders.VideoDecoder Args: @@ -365,7 +375,7 @@ def _read_video_torchcodec( - video_start: the start time of video. - video_end: the end time of video. Returns: - torch.Tensor: the video tensor with shape (T, C, H, W). + Tuple of video tensor (T, C, H, W), metadata, and sampled FPS. """ from torchcodec.decoders import VideoDecoder @@ -424,7 +434,7 @@ def fetch_video( image_patch_size: int = 14, return_video_sample_fps: bool = False, return_video_metadata: bool = False, -) -> Union[torch.Tensor, List[Image.Image]]: +) -> FetchVideoOutput: image_factor = image_patch_size * SPATIAL_MERGE_SIZE VIDEO_FRAME_MIN_PIXELS = VIDEO_MIN_TOKEN_NUM * image_factor * image_factor VIDEO_FRAME_MAX_PIXELS = VIDEO_MAX_TOKEN_NUM * image_factor * image_factor From 03d67a5096622d0a40add870225cdac874563db4 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 06:20:55 +0000 Subject: [PATCH 13/30] style(diffusion): apply pre-commit formatting Co-authored-by: Cursor --- primus/backends/diffusion/data/registrations/wan.py | 4 +++- primus/backends/diffusion/distributed/checkpoint.py | 1 - .../backends/diffusion/models/wan/configuration_wanvideo.py | 4 +++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/primus/backends/diffusion/data/registrations/wan.py b/primus/backends/diffusion/data/registrations/wan.py index f487ec8ca..1f7f72e2e 100644 --- a/primus/backends/diffusion/data/registrations/wan.py +++ b/primus/backends/diffusion/data/registrations/wan.py @@ -61,7 +61,9 @@ def normalize_video(self, video_tchw: torch.Tensor) -> torch.Tensor: video_cthw = video_tchw.permute(1, 0, 2, 3).contiguous() return (video_cthw - self.image_mean) / self.image_std - def prepare_batch(self, *, batch: dict[str, Any], device: torch.device, dtype: torch.dtype) -> dict[str, Any]: + def prepare_batch( + self, *, batch: dict[str, Any], device: torch.device, dtype: torch.dtype + ) -> dict[str, Any]: return batch diff --git a/primus/backends/diffusion/distributed/checkpoint.py b/primus/backends/diffusion/distributed/checkpoint.py index 9c4047f2f..5ca2f32ae 100644 --- a/primus/backends/diffusion/distributed/checkpoint.py +++ b/primus/backends/diffusion/distributed/checkpoint.py @@ -31,7 +31,6 @@ from .mesh import _ensure_process_group - _META_FILENAME = "meta.json" diff --git a/primus/backends/diffusion/models/wan/configuration_wanvideo.py b/primus/backends/diffusion/models/wan/configuration_wanvideo.py index f7d3e4506..689c7449d 100644 --- a/primus/backends/diffusion/models/wan/configuration_wanvideo.py +++ b/primus/backends/diffusion/models/wan/configuration_wanvideo.py @@ -53,7 +53,9 @@ def __init__( separated_timestep = kwargs.pop("separated_timestep", None) if separated_timestep is not None and seperated_timestep is not None: if bool(separated_timestep) != bool(seperated_timestep): - raise ValueError("`separated_timestep` and `seperated_timestep` must match when both are set.") + raise ValueError( + "`separated_timestep` and `seperated_timestep` must match when both are set." + ) if seperated_timestep is None: seperated_timestep = True if separated_timestep is None else bool(separated_timestep) From c7d18edb1b41aba451dc0747fe0ac6359461eda1 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 06:36:35 +0000 Subject: [PATCH 14/30] fix(diffusion): address remaining PR review comments Co-authored-by: Cursor --- primus/backends/diffusion/data/registrations/wan.py | 8 ++++++-- primus/backends/diffusion/diffusion_pretrain_trainer.py | 4 +--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/primus/backends/diffusion/data/registrations/wan.py b/primus/backends/diffusion/data/registrations/wan.py index 1f7f72e2e..46393aae1 100644 --- a/primus/backends/diffusion/data/registrations/wan.py +++ b/primus/backends/diffusion/data/registrations/wan.py @@ -29,7 +29,11 @@ def __init__(self, config: dict[str, Any]): tokenizer_path = config.get("text_tokenizer") if not tokenizer_path: raise ValueError("Wan dataset processor requires `text_tokenizer`.") - self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) + trust_remote_code = bool(config.get("trust_remote_code", False)) + self.tokenizer = AutoTokenizer.from_pretrained( + tokenizer_path, + trust_remote_code=trust_remote_code, + ) extra_kwargs = config.get("extra_kwargs", {}) or {} size = extra_kwargs.get("size", {}) or {} @@ -141,7 +145,7 @@ def __getitem__(self, index: int) -> dict[str, Any]: prompt = sample.get("prompt") or sample.get("text") or sample.get("caption") or "" video_key = sample.get("video") or sample.get("video_path") if not video_key: - raise ValueError(f"Wan dataset sample missing `video`: index={index}") + raise KeyError(f"Wan dataset sample missing `video`: index={index}") item = self.processor.tokenize(str(prompt)) item["video"] = self._read_video(self._resolve_video_path(str(video_key))) diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index 586a2d9ba..0bec5f520 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -35,7 +35,7 @@ def setup(self): missing = [ package - for package in ("torch", "loguru", "safetensors", "transformers", "PIL") + for package in ("torch", "loguru", "safetensors", "transformers", "PIL", "torchvision") if importlib.util.find_spec(package) is None ] video_backend = (dataset_cfg.get("config", {}) or {}).get("video_backend") @@ -43,8 +43,6 @@ def setup(self): missing.append("imageio") if video_backend == "decord" and importlib.util.find_spec("decord") is None: missing.append("decord") - if video_backend == "qwen_vl_utils" and importlib.util.find_spec("torchvision") is None: - missing.append("torchvision") if missing: raise RuntimeError( "Diffusion backend missing required Python packages: " From 5a974fd550db8f216e3f04bf79e8a48fa0a13562 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 07:28:13 +0000 Subject: [PATCH 15/30] fix(diffusion): align Wan timestep config naming Co-authored-by: Cursor --- .../diffusion/models/wan/configuration_wanvideo.py | 12 +----------- .../backends/diffusion/models/wan/train_pipeline.py | 4 ++-- 2 files changed, 3 insertions(+), 13 deletions(-) diff --git a/primus/backends/diffusion/models/wan/configuration_wanvideo.py b/primus/backends/diffusion/models/wan/configuration_wanvideo.py index 689c7449d..2eff5413d 100644 --- a/primus/backends/diffusion/models/wan/configuration_wanvideo.py +++ b/primus/backends/diffusion/models/wan/configuration_wanvideo.py @@ -42,7 +42,7 @@ def __init__( dit_has_image_pos_emb: bool = False, dit_has_ref_conv: bool = False, trainable_modules=None, - seperated_timestep: bool | None = None, + seperated_timestep: bool = True, require_clip_embedding: bool = False, require_vae_embedding: bool = False, fuse_vae_embedding_in_latents: bool = True, @@ -50,15 +50,6 @@ def __init__( vae_type: str = "wan_video_vae_38", **kwargs, ): - separated_timestep = kwargs.pop("separated_timestep", None) - if separated_timestep is not None and seperated_timestep is not None: - if bool(separated_timestep) != bool(seperated_timestep): - raise ValueError( - "`separated_timestep` and `seperated_timestep` must match when both are set." - ) - if seperated_timestep is None: - seperated_timestep = True if separated_timestep is None else bool(separated_timestep) - # DiT configuration self.vae_type = vae_type self.dit_hidden_size = dit_hidden_size @@ -77,7 +68,6 @@ def __init__( self.dit_has_ref_conv = dit_has_ref_conv self.seperated_timestep = bool(seperated_timestep) - self.separated_timestep = self.seperated_timestep self.require_clip_embedding = require_clip_embedding self.require_vae_embedding = require_vae_embedding self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents diff --git a/primus/backends/diffusion/models/wan/train_pipeline.py b/primus/backends/diffusion/models/wan/train_pipeline.py index e48d40987..55bb49ebe 100644 --- a/primus/backends/diffusion/models/wan/train_pipeline.py +++ b/primus/backends/diffusion/models/wan/train_pipeline.py @@ -138,7 +138,7 @@ def _select_timestep(scheduler: Any, device: torch.device) -> torch.Tensor: return timestep.to(device=device) @staticmethod - def _maybe_expand_separated_timestep( + def _maybe_expand_seperated_timestep( *, timestep: torch.Tensor, x_list: list[torch.Tensor], @@ -259,7 +259,7 @@ def compute_loss( separated = bool(getattr(model_config, "seperated_timestep", False)) fuse_flag = bool(getattr(model_config, "fuse_vae_embedding_in_latents", False)) - t = self._maybe_expand_separated_timestep( + t = self._maybe_expand_seperated_timestep( timestep=timestep, x_list=x_list, patch_size=(d_f, d_h, d_w), From 814cfbc64a0bdb1a3f68ed28bd4bc33b759c984a Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 07:54:02 +0000 Subject: [PATCH 16/30] fix(diffusion): resolve Wan training review issues Co-authored-by: Cursor --- .../models/wan/configuration_wanvideo.py | 4 +-- .../diffusion/models/wan/train_pipeline.py | 30 ++++++++++++------- .../models/diffusion/wan2.1_t2v_1.3b.yaml | 2 +- .../models/diffusion/wan2.2_ti2v_5b.yaml | 2 +- 4 files changed, 24 insertions(+), 14 deletions(-) diff --git a/primus/backends/diffusion/models/wan/configuration_wanvideo.py b/primus/backends/diffusion/models/wan/configuration_wanvideo.py index 2eff5413d..0df88d0ec 100644 --- a/primus/backends/diffusion/models/wan/configuration_wanvideo.py +++ b/primus/backends/diffusion/models/wan/configuration_wanvideo.py @@ -42,7 +42,7 @@ def __init__( dit_has_image_pos_emb: bool = False, dit_has_ref_conv: bool = False, trainable_modules=None, - seperated_timestep: bool = True, + separated_timestep: bool = True, require_clip_embedding: bool = False, require_vae_embedding: bool = False, fuse_vae_embedding_in_latents: bool = True, @@ -67,7 +67,7 @@ def __init__( self.dit_has_image_pos_emb = dit_has_image_pos_emb self.dit_has_ref_conv = dit_has_ref_conv - self.seperated_timestep = bool(seperated_timestep) + self.separated_timestep = bool(separated_timestep) self.require_clip_embedding = require_clip_embedding self.require_vae_embedding = require_vae_embedding self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents diff --git a/primus/backends/diffusion/models/wan/train_pipeline.py b/primus/backends/diffusion/models/wan/train_pipeline.py index 55bb49ebe..e0501466f 100644 --- a/primus/backends/diffusion/models/wan/train_pipeline.py +++ b/primus/backends/diffusion/models/wan/train_pipeline.py @@ -62,10 +62,20 @@ def _get_seed_from_env_or_batch(batch: Dict[str, Any]) -> Optional[int]: seed = batch.get("seed", None) if seed is None: return None - try: - return int(seed) - except Exception: - return None + if isinstance(seed, torch.Tensor): + flat_seed = seed.detach().reshape(-1).cpu() + if flat_seed.numel() == 0: + return None + if flat_seed.numel() > 1 and not bool((flat_seed == flat_seed[0]).all()): + raise ValueError("Per-sample dataset seeds are not supported; got different seeds in one batch.") + seed = flat_seed[0].item() + elif isinstance(seed, (list, tuple)): + if not seed: + return None + if len(seed) > 1 and any(value != seed[0] for value in seed): + raise ValueError("Per-sample dataset seeds are not supported; got different seeds in one batch.") + seed = seed[0] + return int(seed) @staticmethod def _randn_like_on_cpu_then_to( @@ -138,7 +148,7 @@ def _select_timestep(scheduler: Any, device: torch.device) -> torch.Tensor: return timestep.to(device=device) @staticmethod - def _maybe_expand_seperated_timestep( + def _maybe_expand_separated_timestep( *, timestep: torch.Tensor, x_list: list[torch.Tensor], @@ -173,7 +183,7 @@ def _maybe_expand_seperated_timestep( for i, x in enumerate(x_list): # x: [C, F, H, W] in latent space f, h, w = x.shape[1], x.shape[2], x.shape[3] - seq_len = f * (h // d_h) * (w // d_w) + seq_len = (f // d_f) * (h // d_h) * (w // d_w) t_seq = torch.full((seq_len,), timestep_b[i], device=timestep_b.device, dtype=timestep_b.dtype) spatial_patches = (h // d_h) * (w // d_w) t_seq[:spatial_patches] = 0 @@ -250,16 +260,16 @@ def compute_loss( x_list = [noisy_latents[i] for i in range(noisy_latents.shape[0])] context_list = [context[i] for i in range(context.shape[0])] - # seq_len computed like wan_new: based on latent shape and patch size + # seq_len matches DiT Conv3d patchification over [F, H, W]. d_f, d_h, d_w = patch_size max_seq_len = 0 for x in x_list: - seq_len = x.shape[1] * (x.shape[2] // d_h) * (x.shape[3] // d_w) + seq_len = (x.shape[1] // d_f) * (x.shape[2] // d_h) * (x.shape[3] // d_w) max_seq_len = max(max_seq_len, seq_len) - separated = bool(getattr(model_config, "seperated_timestep", False)) + separated = bool(getattr(model_config, "separated_timestep", False)) fuse_flag = bool(getattr(model_config, "fuse_vae_embedding_in_latents", False)) - t = self._maybe_expand_seperated_timestep( + t = self._maybe_expand_separated_timestep( timestep=timestep, x_list=x_list, patch_size=(d_f, d_h, d_w), diff --git a/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml b/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml index 77995b460..023776ecd 100644 --- a/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml +++ b/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml @@ -4,7 +4,7 @@ model: load_from_pretrained_path: ${PRETRAINED_PATH:/models/Wan2.1-T2V-1.3B} config: model_type: t2v - seperated_timestep: false + separated_timestep: false fuse_vae_embedding_in_latents: false trainable_modules: dit vae_type: wan_video_vae diff --git a/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml b/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml index a9dd68567..0f23001e2 100644 --- a/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml +++ b/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml @@ -4,7 +4,7 @@ model: load_from_pretrained_path: ${PRETRAINED_PATH:/models/Wan2.2-TI2V-5B} config: model_type: ti2v - seperated_timestep: true + separated_timestep: true fuse_vae_embedding_in_latents: true trainable_modules: dit vae_type: wan_video_vae_38 From fefe3977be414288d7fc36dbd86399c2c89b1654 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 07:59:30 +0000 Subject: [PATCH 17/30] Fix black formatting in train_pipeline.py to pass code-lint CI --- .../diffusion/models/wan/train_pipeline.py | 70 +++++++++++++++---- 1 file changed, 55 insertions(+), 15 deletions(-) diff --git a/primus/backends/diffusion/models/wan/train_pipeline.py b/primus/backends/diffusion/models/wan/train_pipeline.py index e0501466f..b932d3279 100644 --- a/primus/backends/diffusion/models/wan/train_pipeline.py +++ b/primus/backends/diffusion/models/wan/train_pipeline.py @@ -41,7 +41,11 @@ def __init__(self, cfg: Optional[WanFlowMatchTrainPipelineConfig] = None): self.cfg = cfg or WanFlowMatchTrainPipelineConfig() @staticmethod - def _encode_prompt(text_encoder: torch.nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor): + def _encode_prompt( + text_encoder: torch.nn.Module, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + ): # Match existing wan_new behavior: zero-out padded embeddings explicitly. seq_lens = attention_mask.gt(0).sum(dim=1).long() prompt_emb = text_encoder(input_ids, attention_mask) @@ -58,7 +62,9 @@ def _get_seed_from_env_or_batch(batch: Dict[str, Any]) -> Optional[int]: try: return int(os.environ["FIXED_SEED"]) except ValueError as exc: - raise ValueError(f"Invalid FIXED_SEED value: {os.environ['FIXED_SEED']}") from exc + raise ValueError( + f"Invalid FIXED_SEED value: {os.environ['FIXED_SEED']}" + ) from exc seed = batch.get("seed", None) if seed is None: return None @@ -67,25 +73,35 @@ def _get_seed_from_env_or_batch(batch: Dict[str, Any]) -> Optional[int]: if flat_seed.numel() == 0: return None if flat_seed.numel() > 1 and not bool((flat_seed == flat_seed[0]).all()): - raise ValueError("Per-sample dataset seeds are not supported; got different seeds in one batch.") + raise ValueError( + "Per-sample dataset seeds are not supported; got different seeds in one batch." + ) seed = flat_seed[0].item() elif isinstance(seed, (list, tuple)): if not seed: return None if len(seed) > 1 and any(value != seed[0] for value in seed): - raise ValueError("Per-sample dataset seeds are not supported; got different seeds in one batch.") + raise ValueError( + "Per-sample dataset seeds are not supported; got different seeds in one batch." + ) seed = seed[0] return int(seed) @staticmethod def _randn_like_on_cpu_then_to( - x: torch.Tensor, *, seed: Optional[int], dtype: torch.dtype, device: torch.device + x: torch.Tensor, + *, + seed: Optional[int], + dtype: torch.dtype, + device: torch.device, ) -> torch.Tensor: generator = None if seed is not None: generator = torch.Generator(device="cpu") generator.manual_seed(seed) - noise = torch.randn(x.shape, generator=generator, device="cpu", dtype=torch.float32) + noise = torch.randn( + x.shape, generator=generator, device="cpu", dtype=torch.float32 + ) return noise.to(device=device, dtype=dtype) @staticmethod @@ -119,7 +135,9 @@ def _pad_latents_for_dit( return latents, (height, width) @staticmethod - def _crop_latents(latents: torch.Tensor, *, spatial_size: tuple[int, int]) -> torch.Tensor: + def _crop_latents( + latents: torch.Tensor, *, spatial_size: tuple[int, int] + ) -> torch.Tensor: height, width = spatial_size return latents[..., :height, :width] @@ -136,12 +154,16 @@ def _select_timestep(scheduler: Any, device: torch.device) -> torch.Tensor: try: fixed_step = int(os.environ["FIXED_TIMESTEP"]) except ValueError as exc: - raise ValueError(f"Invalid FIXED_TIMESTEP value: {os.environ['FIXED_TIMESTEP']}") from exc + raise ValueError( + f"Invalid FIXED_TIMESTEP value: {os.environ['FIXED_TIMESTEP']}" + ) from exc max_step = int(scheduler.num_train_timesteps) - 1 fixed_step = max(0, min(fixed_step, max_step)) timestep_id = torch.tensor([fixed_step], device=device) else: - timestep_id = torch.randint(0, int(scheduler.num_train_timesteps), (1,), device=device) + timestep_id = torch.randint( + 0, int(scheduler.num_train_timesteps), (1,), device=device + ) # scheduler.timesteps live on CPU in this repo; `wan_new` indexes with cpu tensor. timestep = scheduler.timesteps[timestep_id.cpu()].float() @@ -184,7 +206,12 @@ def _maybe_expand_separated_timestep( # x: [C, F, H, W] in latent space f, h, w = x.shape[1], x.shape[2], x.shape[3] seq_len = (f // d_f) * (h // d_h) * (w // d_w) - t_seq = torch.full((seq_len,), timestep_b[i], device=timestep_b.device, dtype=timestep_b.dtype) + t_seq = torch.full( + (seq_len,), + timestep_b[i], + device=timestep_b.device, + dtype=timestep_b.dtype, + ) spatial_patches = (h // d_h) * (w // d_w) t_seq[:spatial_patches] = 0 t_expand_list.append(t_seq) @@ -235,7 +262,9 @@ def compute_loss( pass with torch.no_grad(): - input_latents = self._vae_encode(components.vae, video).to(device=device, dtype=dtype) + input_latents = self._vae_encode(components.vae, video).to( + device=device, dtype=dtype + ) patch_size = self._get_dit_patch_size(model_config) input_latents, original_latent_spatial_size = self._pad_latents_for_dit( input_latents, patch_size=patch_size @@ -244,7 +273,9 @@ def compute_loss( # 2) Sample timestep + noise (match `wan_new`: noise on CPU, timestep from scheduler.timesteps) timestep = self._select_timestep(scheduler, device=device) # [1] (float) seed = self._get_seed_from_env_or_batch(batch) - noise = self._randn_like_on_cpu_then_to(input_latents, seed=seed, dtype=dtype, device=device) + noise = self._randn_like_on_cpu_then_to( + input_latents, seed=seed, dtype=dtype, device=device + ) # 3) Diffusion target + noise injection (match `wan_new`) # Note: in this repo's FlowMatchScheduler, `training_target(sample, noise, t)` uses `noise - sample`. @@ -254,7 +285,9 @@ def compute_loss( # 4) Text embeddings with torch.no_grad(): - context = self._encode_prompt(components.text_encoder, input_ids, attention_mask) + context = self._encode_prompt( + components.text_encoder, input_ids, attention_mask + ) # 5) DiT forward (official interface: List[Tensor] per-sample) x_list = [noisy_latents[i] for i in range(noisy_latents.shape[0])] @@ -278,10 +311,17 @@ def compute_loss( sp_group = batch.get("sp_group", None) noise_pred_list = components.dit( - x=x_list, t=t, context=context_list, seq_len=max_seq_len, y=None, sp_group=sp_group + x=x_list, + t=t, + context=context_list, + seq_len=max_seq_len, + y=None, + sp_group=sp_group, ) noise_pred = torch.stack(noise_pred_list) - noise_pred = self._crop_latents(noise_pred, spatial_size=original_latent_spatial_size) + noise_pred = self._crop_latents( + noise_pred, spatial_size=original_latent_spatial_size + ) target = self._crop_latents(target, spatial_size=original_latent_spatial_size) # 6) Loss From 12a87d0715fc07328b90a8ede608cd2b982c7a0a Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 08:21:15 +0000 Subject: [PATCH 18/30] fix(diffusion): remove Wan attention debug leftovers Co-authored-by: Cursor --- .../diffusion/attention/_flash_common.py | 11 ------ .../diffusion/models/wan/train_pipeline.py | 36 +++++-------------- 2 files changed, 9 insertions(+), 38 deletions(-) diff --git a/primus/backends/diffusion/attention/_flash_common.py b/primus/backends/diffusion/attention/_flash_common.py index 7535a48a1..c1bf0f7a4 100644 --- a/primus/backends/diffusion/attention/_flash_common.py +++ b/primus/backends/diffusion/attention/_flash_common.py @@ -101,17 +101,6 @@ def half(x): max_sq = int(q_lens.max()) max_sk = int(k_lens.max()) - import os - - import torch.distributed as dist - - if os.environ.get("OMNIFLOW_DEBUG_ATTN") == "1" and (not dist.is_initialized() or dist.get_rank() == 0): - print(f"[DEBUG ATTN] q.shape={q.shape}, k.shape={k.shape}, v.shape={v.shape}") - print(f"[DEBUG ATTN] q_lens={q_lens}, k_lens={k_lens}") - print(f"[DEBUG ATTN] cu_seqlens_q={cu_seqlens_q}, cu_seqlens_k={cu_seqlens_k}") - print(f"[DEBUG ATTN] max_seqlen_q={max_sq} (lq={lq}), max_seqlen_k={max_sk} (lk={lk})") - print(f"[DEBUG ATTN] q.dtype={q.dtype}, causal={causal}, window_size={call_window_size}") - x = varlen_attention( q=q, k=k, diff --git a/primus/backends/diffusion/models/wan/train_pipeline.py b/primus/backends/diffusion/models/wan/train_pipeline.py index b932d3279..3f8a4d1b8 100644 --- a/primus/backends/diffusion/models/wan/train_pipeline.py +++ b/primus/backends/diffusion/models/wan/train_pipeline.py @@ -62,9 +62,7 @@ def _get_seed_from_env_or_batch(batch: Dict[str, Any]) -> Optional[int]: try: return int(os.environ["FIXED_SEED"]) except ValueError as exc: - raise ValueError( - f"Invalid FIXED_SEED value: {os.environ['FIXED_SEED']}" - ) from exc + raise ValueError(f"Invalid FIXED_SEED value: {os.environ['FIXED_SEED']}") from exc seed = batch.get("seed", None) if seed is None: return None @@ -99,9 +97,7 @@ def _randn_like_on_cpu_then_to( if seed is not None: generator = torch.Generator(device="cpu") generator.manual_seed(seed) - noise = torch.randn( - x.shape, generator=generator, device="cpu", dtype=torch.float32 - ) + noise = torch.randn(x.shape, generator=generator, device="cpu", dtype=torch.float32) return noise.to(device=device, dtype=dtype) @staticmethod @@ -135,9 +131,7 @@ def _pad_latents_for_dit( return latents, (height, width) @staticmethod - def _crop_latents( - latents: torch.Tensor, *, spatial_size: tuple[int, int] - ) -> torch.Tensor: + def _crop_latents(latents: torch.Tensor, *, spatial_size: tuple[int, int]) -> torch.Tensor: height, width = spatial_size return latents[..., :height, :width] @@ -154,16 +148,12 @@ def _select_timestep(scheduler: Any, device: torch.device) -> torch.Tensor: try: fixed_step = int(os.environ["FIXED_TIMESTEP"]) except ValueError as exc: - raise ValueError( - f"Invalid FIXED_TIMESTEP value: {os.environ['FIXED_TIMESTEP']}" - ) from exc + raise ValueError(f"Invalid FIXED_TIMESTEP value: {os.environ['FIXED_TIMESTEP']}") from exc max_step = int(scheduler.num_train_timesteps) - 1 fixed_step = max(0, min(fixed_step, max_step)) timestep_id = torch.tensor([fixed_step], device=device) else: - timestep_id = torch.randint( - 0, int(scheduler.num_train_timesteps), (1,), device=device - ) + timestep_id = torch.randint(0, int(scheduler.num_train_timesteps), (1,), device=device) # scheduler.timesteps live on CPU in this repo; `wan_new` indexes with cpu tensor. timestep = scheduler.timesteps[timestep_id.cpu()].float() @@ -262,9 +252,7 @@ def compute_loss( pass with torch.no_grad(): - input_latents = self._vae_encode(components.vae, video).to( - device=device, dtype=dtype - ) + input_latents = self._vae_encode(components.vae, video).to(device=device, dtype=dtype) patch_size = self._get_dit_patch_size(model_config) input_latents, original_latent_spatial_size = self._pad_latents_for_dit( input_latents, patch_size=patch_size @@ -273,9 +261,7 @@ def compute_loss( # 2) Sample timestep + noise (match `wan_new`: noise on CPU, timestep from scheduler.timesteps) timestep = self._select_timestep(scheduler, device=device) # [1] (float) seed = self._get_seed_from_env_or_batch(batch) - noise = self._randn_like_on_cpu_then_to( - input_latents, seed=seed, dtype=dtype, device=device - ) + noise = self._randn_like_on_cpu_then_to(input_latents, seed=seed, dtype=dtype, device=device) # 3) Diffusion target + noise injection (match `wan_new`) # Note: in this repo's FlowMatchScheduler, `training_target(sample, noise, t)` uses `noise - sample`. @@ -285,9 +271,7 @@ def compute_loss( # 4) Text embeddings with torch.no_grad(): - context = self._encode_prompt( - components.text_encoder, input_ids, attention_mask - ) + context = self._encode_prompt(components.text_encoder, input_ids, attention_mask) # 5) DiT forward (official interface: List[Tensor] per-sample) x_list = [noisy_latents[i] for i in range(noisy_latents.shape[0])] @@ -319,9 +303,7 @@ def compute_loss( sp_group=sp_group, ) noise_pred = torch.stack(noise_pred_list) - noise_pred = self._crop_latents( - noise_pred, spatial_size=original_latent_spatial_size - ) + noise_pred = self._crop_latents(noise_pred, spatial_size=original_latent_spatial_size) target = self._crop_latents(target, spatial_size=original_latent_spatial_size) # 6) Loss From 1c71546bff6e66f024cc100108470425dab3224b Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 08:44:24 +0000 Subject: [PATCH 19/30] fix(diffusion): resolve Wan review follow-ups Remove redundant VAE checkpoint loading, harden AdamW optimizer fallback, and add focused tests for Wan argument normalization and optimizer fallback behavior. Co-authored-by: Cursor --- primus/backends/diffusion/argument_builder.py | 12 +- .../diffusion/models/registrations/wan.py | 24 +-- primus/backends/diffusion/trainers/base.py | 12 +- .../diffusion/test_wan_argument_builder.py | 194 ++++++++++++++++++ .../diffusion/test_wan_trainer_optimizer.py | 61 ++++++ 5 files changed, 282 insertions(+), 21 deletions(-) create mode 100644 tests/unit_tests/backends/diffusion/test_wan_argument_builder.py create mode 100644 tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py index 9814ed140..b6b444ceb 100644 --- a/primus/backends/diffusion/argument_builder.py +++ b/primus/backends/diffusion/argument_builder.py @@ -198,9 +198,17 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, height = data.get("height") width = data.get("width") if height is not None: - self._set_nested(dataset_cfg, ("processor_config", "extra_kwargs", "size", "height"), height) + self._set_nested( + dataset_cfg, + ("processor_config", "extra_kwargs", "size", "height"), + height, + ) if width is not None: - self._set_nested(dataset_cfg, ("processor_config", "extra_kwargs", "size", "width"), width) + self._set_nested( + dataset_cfg, + ("processor_config", "extra_kwargs", "size", "width"), + width, + ) parallelism_map = { ("sp_size",): ("sp_size",), diff --git a/primus/backends/diffusion/models/registrations/wan.py b/primus/backends/diffusion/models/registrations/wan.py index aa22cc48c..491284ff7 100644 --- a/primus/backends/diffusion/models/registrations/wan.py +++ b/primus/backends/diffusion/models/registrations/wan.py @@ -29,7 +29,9 @@ from primus.backends.diffusion.utils.train_utils import count_parameters -def _strip_module_prefix(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: +def _strip_module_prefix( + state_dict: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: out: dict[str, torch.Tensor] = {} for k, v in state_dict.items(): if k.startswith("module."): @@ -70,7 +72,11 @@ def _load_dit_weights_into_module(dit: torch.nn.Module, pretrained_path: str): # Directory: prefer explicit Wan trainer file name, else fall back to pattern search candidates: list[str] = [] - for fname in ("dit_model.safetensors", "diffusion_pytorch_model.safetensors", "model.safetensors"): + for fname in ( + "dit_model.safetensors", + "diffusion_pytorch_model.safetensors", + "model.safetensors", + ): p = os.path.join(pretrained_path, fname) if os.path.exists(p): candidates.append(p) @@ -95,17 +101,6 @@ def _load_dit_weights_into_module(dit: torch.nn.Module, pretrained_path: str): ) -def _load_vae(vae: torch.nn.Module, ckpt_path: str | None): - if not ckpt_path: - return - logger.info(f"Loading VAE from {ckpt_path}") - state = _strip_module_prefix(_load_state_dict(ckpt_path)) - # VAE wrapper may keep real module under `.model` - target = vae.model if hasattr(vae, "model") else vae - result = target.load_state_dict(state, strict=False) - logger.info(f"VAE loaded. missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}") - - def build_wan_model(model_config: dict): """ YAML compatibility: @@ -207,9 +202,6 @@ def build_wan_model(model_config: dict): logger.info(f"Loading DiT weights from {pretrained_path}") _load_dit_weights_into_module(dit, pretrained_path) - # Load encoders (VAE weights loaded into `.model`) - _load_vae(vae, vae_ckpt) - components = WanComponents(dit=dit, vae=vae, text_encoder=text_encoder, image_encoder=None) pipeline = WanFlowMatchTrainPipeline() diff --git a/primus/backends/diffusion/trainers/base.py b/primus/backends/diffusion/trainers/base.py index d7ffd25ce..3b036732d 100644 --- a/primus/backends/diffusion/trainers/base.py +++ b/primus/backends/diffusion/trainers/base.py @@ -283,19 +283,25 @@ def _create_optimizer(self): eps = float(self.args.get("adam_epsilon", 1e-8)) params = [p for p in self.model.parameters() if p.requires_grad] - optimizer_kwargs = {"params": params, "lr": lr, "betas": betas, "eps": eps, "weight_decay": wd} + optimizer_kwargs = { + "params": params, + "lr": lr, + "betas": betas, + "eps": eps, + "weight_decay": wd, + } optimizer = None try: optimizer = torch.optim.AdamW(**optimizer_kwargs, fused=True) if self.rank == 0: logger.info("Optimizer: torch.optim.AdamW(fused=True)") - except TypeError: + except (TypeError, RuntimeError): try: optimizer = torch.optim.AdamW(**optimizer_kwargs, foreach=True) if self.rank == 0: logger.info("Optimizer: torch.optim.AdamW(foreach=True)") - except TypeError: + except (TypeError, RuntimeError): optimizer = torch.optim.AdamW(**optimizer_kwargs) if self.rank == 0: logger.info("Optimizer: torch.optim.AdamW(default)") diff --git a/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py b/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py new file mode 100644 index 000000000..9d07d4f07 --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py @@ -0,0 +1,194 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from types import SimpleNamespace + +import pytest + +from primus.backends.diffusion.argument_builder import WanArgBuilder + + +def _minimal_params(): + return { + "model": { + "name": "wan", + "config": { + "model_type": "t2v", + }, + }, + } + + +def test_rejects_legacy_public_dataset_override(): + params = _minimal_params() + params["dataset"] = {"name": "wan"} + + builder = WanArgBuilder() + builder.update(params) + + with pytest.raises(ValueError, match="no longer accepts public `dataset` overrides"): + builder.finalize() + + +def test_rejects_legacy_public_trainer_override(): + params = _minimal_params() + params["trainer"] = {"name": "fsdp2"} + + builder = WanArgBuilder() + builder.update(params) + + with pytest.raises(ValueError, match="no longer accepts public `trainer` overrides"): + builder.finalize() + + +def test_maps_primus_style_sections_to_wan_runtime_config(): + params = { + **_minimal_params(), + "stage": "posttrain", + "primus": {"experiment": "wan-smoke"}, + "training": { + "steps": 7, + "local_batch_size": 2, + "global_batch_size": 16, + "gradient_accumulation_steps": 4, + "output_dir": "/tmp/wan-out", + "save_steps": 3, + "run_name": "wan-test", + "num_train_epochs": 9, + "dataloader_num_workers": 0, + "resume_from_checkpoint": "/tmp/resume", + }, + "data": { + "dataset_path": "/data/meta.jsonl", + "data_folder": "/data/videos", + "frame_num": 17, + "video_backend": "decord", + "text_tokenizer": "/models/umt5", + "height": 256, + "width": 384, + }, + "parallelism": { + "sp_size": 4, + "dp_replicate": 2, + }, + "optimizer": { + "lr": 2.0e-5, + "weight_decay": 0.02, + "adam_beta1": 0.8, + "adam_beta2": 0.95, + "adam_epsilon": 1.0e-7, + "max_grad_norm": 0.5, + }, + "runtime": { + "attention_backend": "sdpa", + "report_to": "none", + "seed": 1234, + "fsdp2_reshard_after_forward": False, + }, + "metrics": { + "log_freq": 5, + "enable_wandb": False, + }, + } + + builder = WanArgBuilder() + builder.update(params) + result = builder.finalize() + + dataset_cfg = result.dataset["config"] + processor_cfg = dataset_cfg["processor_config"] + trainer_args = result.trainer["args"] + + assert result.stage == "posttrain" + assert result.primus == {"experiment": "wan-smoke"} + assert result.model == params["model"] + + assert dataset_cfg["dataset_path"] == "/data/meta.jsonl" + assert dataset_cfg["data_folder"] == "/data/videos" + assert dataset_cfg["frame_num"] == 17 + assert dataset_cfg["video_backend"] == "decord" + assert processor_cfg["text_tokenizer"] == "/models/umt5" + assert processor_cfg["extra_kwargs"]["size"] == {"height": 256, "width": 384} + + assert trainer_args["max_steps"] == 7 + assert trainer_args["per_device_train_batch_size"] == 2 + assert trainer_args["global_batch_size"] == 16 + assert trainer_args["gradient_accumulation_steps"] == 4 + assert trainer_args["output_dir"] == "/tmp/wan-out" + assert trainer_args["save_steps"] == 3 + assert trainer_args["run_name"] == "wan-test" + assert trainer_args["num_train_epochs"] == 9 + assert trainer_args["dataloader_num_workers"] == 0 + assert trainer_args["resume_from_checkpoint"] == "/tmp/resume" + + assert trainer_args["sp_size"] == 4 + assert trainer_args["dp_replicate"] == 2 + assert trainer_args["learning_rate"] == 2.0e-5 + assert trainer_args["weight_decay"] == 0.02 + assert trainer_args["adam_beta1"] == 0.8 + assert trainer_args["adam_beta2"] == 0.95 + assert trainer_args["adam_epsilon"] == 1.0e-7 + assert trainer_args["max_grad_norm"] == 0.5 + assert trainer_args["attention_backend"] == "sdpa" + assert trainer_args["report_to"] == "none" + assert trainer_args["seed"] == 1234 + assert trainer_args["fsdp2_reshard_after_forward"] is False + assert trainer_args["logging_steps"] == 5 + + +def test_defaults_propagate_when_optional_sections_are_omitted(): + builder = WanArgBuilder() + builder.update(_minimal_params()) + result = builder.finalize() + + dataset_cfg = result.dataset["config"] + trainer_args = result.trainer["args"] + + assert result.stage == "pretrain" + assert dataset_cfg["dataset_path"] == "/path/to/meta.jsonl" + assert dataset_cfg["data_folder"] == "/path/to/videos" + assert dataset_cfg["processor_config"]["text_tokenizer"] == "/path/to/umt5-xxl" + assert result.trainer["name"] == "fsdp2" + assert trainer_args["max_steps"] == 100 + assert trainer_args["attention_backend"] == "flash_attn_aiter" + assert trainer_args["sp_size"] == 1 + assert trainer_args["dp_replicate"] == 1 + assert trainer_args["report_to"] == "none" + + +def test_metrics_enable_wandb_sets_report_to_when_runtime_omits_it(): + builder = WanArgBuilder() + builder.update( + { + **_minimal_params(), + "metrics": { + "enable_wandb": True, + }, + } + ) + + result = builder.finalize() + + assert result.trainer["args"]["report_to"] == "wandb" + + +def test_update_accepts_simplenamespace_input(): + builder = WanArgBuilder() + builder.update( + SimpleNamespace( + model=SimpleNamespace( + name="wan", + config=SimpleNamespace(model_type="t2v"), + ), + data=SimpleNamespace(dataset_path="/data/meta.jsonl"), + ) + ) + + result = builder.finalize() + + assert result.model["name"] == "wan" + assert result.model["config"]["model_type"] == "t2v" + assert result.dataset["config"]["dataset_path"] == "/data/meta.jsonl" diff --git a/tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py b/tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py new file mode 100644 index 000000000..6636942e9 --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py @@ -0,0 +1,61 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import torch + +from primus.backends.diffusion.trainers.base import BaseWanTrainer + + +def _make_trainer(): + trainer = BaseWanTrainer.__new__(BaseWanTrainer) + trainer.model = torch.nn.Linear(2, 2) + trainer.rank = 0 + trainer.args = { + "learning_rate": 1.0e-4, + "weight_decay": 0.01, + "adam_beta1": 0.9, + "adam_beta2": 0.999, + "adam_epsilon": 1.0e-8, + } + return trainer + + +def test_optimizer_falls_back_to_foreach_when_fused_raises_runtime_error(monkeypatch): + calls = [] + + def fake_adamw(**kwargs): + calls.append(kwargs) + if kwargs.get("fused"): + raise RuntimeError("fused AdamW is unsupported") + return kwargs + + monkeypatch.setattr(torch.optim, "AdamW", fake_adamw) + + optimizer = _make_trainer()._create_optimizer() + + assert calls[0]["fused"] is True + assert calls[1]["foreach"] is True + assert optimizer is calls[1] + + +def test_optimizer_falls_back_to_default_when_foreach_raises_runtime_error(monkeypatch): + calls = [] + + def fake_adamw(**kwargs): + calls.append(kwargs) + if kwargs.get("fused") or kwargs.get("foreach"): + raise RuntimeError("optimized AdamW path is unsupported") + return kwargs + + monkeypatch.setattr(torch.optim, "AdamW", fake_adamw) + + optimizer = _make_trainer()._create_optimizer() + + assert calls[0]["fused"] is True + assert calls[1]["foreach"] is True + assert "fused" not in calls[2] + assert "foreach" not in calls[2] + assert optimizer is calls[2] From 78053faf2ccaaffe7dcee56aba15e8c19d0c7f9b Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 09:46:08 +0000 Subject: [PATCH 20/30] fix(diffusion): address latest Wan review comments Add missing diffusion dependency checks, normalize file:// video paths across readers, report accumulated loss correctly, and validate Ulysses all-to-all divisibility before collectives. Co-authored-by: Cursor --- primus/backends/diffusion/data/registrations/wan.py | 4 +++- .../backends/diffusion/diffusion_pretrain_trainer.py | 11 ++++++++++- primus/backends/diffusion/distributed/ulysses.py | 11 +++++++++++ primus/backends/diffusion/trainers/base.py | 12 +++++++++--- primus/backends/diffusion/utils/vision_process.py | 12 ++++++++---- .../pretrain/diffusion/requirements-diffusion.txt | 2 ++ 6 files changed, 43 insertions(+), 9 deletions(-) diff --git a/primus/backends/diffusion/data/registrations/wan.py b/primus/backends/diffusion/data/registrations/wan.py index 46393aae1..76492cb2b 100644 --- a/primus/backends/diffusion/data/registrations/wan.py +++ b/primus/backends/diffusion/data/registrations/wan.py @@ -17,7 +17,7 @@ from torchvision.transforms import functional as F from transformers import AutoTokenizer -from primus.backends.diffusion.utils.vision_process import fetch_video +from primus.backends.diffusion.utils.vision_process import fetch_video, strip_file_uri class WanVideoProcessor: @@ -108,6 +108,7 @@ def _sample_video(self, video_tchw: torch.Tensor) -> torch.Tensor: def _read_video_imageio(self, path: str) -> torch.Tensor: import imageio.v3 as iio + path = strip_file_uri(path) frames = [] for frame in iio.imiter(path): image = Image.fromarray(frame).convert("RGB") @@ -119,6 +120,7 @@ def _read_video_imageio(self, path: str) -> torch.Tensor: def _read_video_decord(self, path: str) -> torch.Tensor: import decord + path = strip_file_uri(path) vr = decord.VideoReader(path) total_frames = len(vr) idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long().tolist() diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index 0bec5f520..ddbb70d49 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -35,7 +35,16 @@ def setup(self): missing = [ package - for package in ("torch", "loguru", "safetensors", "transformers", "PIL", "torchvision") + for package in ( + "torch", + "loguru", + "safetensors", + "transformers", + "PIL", + "torchvision", + "requests", + "packaging", + ) if importlib.util.find_spec(package) is None ] video_backend = (dataset_cfg.get("config", {}) or {}).get("video_backend") diff --git a/primus/backends/diffusion/distributed/ulysses.py b/primus/backends/diffusion/distributed/ulysses.py index 04bf84cfc..b61ae80e2 100644 --- a/primus/backends/diffusion/distributed/ulysses.py +++ b/primus/backends/diffusion/distributed/ulysses.py @@ -34,6 +34,15 @@ # --------------------------------------------------------------------------- +def _require_divisible_dim(x: torch.Tensor, dim: int, divisor: int, op_name: str) -> None: + size = x.shape[dim] + if size % divisor != 0: + raise ValueError( + f"{op_name} requires tensor dimension {dim} (size={size}) to be divisible by " + f"sequence parallel size {divisor}. Got shape={tuple(x.shape)}." + ) + + class _SeqAllToAll(torch.autograd.Function): """All-to-all with autograd. Backward is the inverse (swap dims).""" @@ -43,6 +52,7 @@ def forward(ctx, group, x, scatter_dim, gather_dim): ctx.scatter_dim = scatter_dim ctx.gather_dim = gather_dim sp_size = dist.get_world_size(group) + _require_divisible_dim(x, scatter_dim, sp_size, "_SeqAllToAll") input_list = [t.contiguous() for t in x.tensor_split(sp_size, scatter_dim)] output_list = [torch.empty_like(input_list[0]) for _ in range(sp_size)] dist.all_to_all(output_list, input_list, group=group) @@ -68,6 +78,7 @@ def forward(ctx, x, dim, group): sp_size = dist.get_world_size(group) sp_rank = dist.get_rank(group) ctx.sp_size = sp_size + _require_divisible_dim(x, dim, sp_size, "_SliceWithGather") chunk_size = x.shape[dim] // sp_size return x.narrow(dim, sp_rank * chunk_size, chunk_size).contiguous() diff --git a/primus/backends/diffusion/trainers/base.py b/primus/backends/diffusion/trainers/base.py index 3b036732d..c5d5227fd 100644 --- a/primus/backends/diffusion/trainers/base.py +++ b/primus/backends/diffusion/trainers/base.py @@ -476,6 +476,8 @@ def train(self): local_samples_in_update = 0 local_samples_since_log = 0 update_steps_since_log = 0 + update_loss_sum = 0.0 + update_loss_count = 0 for epoch in range(self.num_train_epochs): self.sampler.set_epoch(epoch) @@ -485,11 +487,16 @@ def train(self): local_samples_in_update += self._infer_local_batch_size(batch) with self._grad_sync_context(is_update_step): - loss = self.compute_loss(batch) - loss = loss / max(1, self.grad_accum_steps) + raw_loss = self.compute_loss(batch) + update_loss_sum += raw_loss.detach().float().item() + update_loss_count += 1 + loss = raw_loss / max(1, self.grad_accum_steps) loss.backward() if is_update_step: + loss_val = update_loss_sum / max(1, update_loss_count) + update_loss_sum = 0.0 + update_loss_count = 0 grad_norm = self._clip_grad_norm() self.optimizer.step() @@ -502,7 +509,6 @@ def train(self): # Logging if self.global_step % self.logging_steps == 0: - loss_val = loss.detach().float().item() * max(1, self.grad_accum_steps) now = time.time() log_interval = now - last_log_time step_time = log_interval / max(1, update_steps_since_log) diff --git a/primus/backends/diffusion/utils/vision_process.py b/primus/backends/diffusion/utils/vision_process.py index 84f785569..2058179be 100644 --- a/primus/backends/diffusion/utils/vision_process.py +++ b/primus/backends/diffusion/utils/vision_process.py @@ -57,6 +57,11 @@ ] +def strip_file_uri(path: str) -> str: + """Return a filesystem path for local file URIs.""" + return path[7:] if path.startswith("file://") else path + + def round_by_factor(number: int, factor: int) -> int: """Returns the closest integer to 'number' that is divisible by 'factor'.""" return round(number / factor) * factor @@ -224,8 +229,7 @@ def _read_video_torchvision( warnings.warn( "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0." ) - if "file://" in video_path: - video_path = video_path[7:] + video_path = strip_file_uri(video_path) st = time.time() video, audio, info = io.read_video( video_path, @@ -332,7 +336,7 @@ def _read_video_decord( """ import decord - video_path = ele["video"] + video_path = strip_file_uri(ele["video"]) st = time.time() vr = decord.VideoReader(video_path) total_frames, video_fps = len(vr), vr.get_avg_fps() @@ -381,7 +385,7 @@ def _read_video_torchcodec( TORCHCODEC_NUM_THREADS = int(os.environ.get("TORCHCODEC_NUM_THREADS", 8)) logger.info(f"set TORCHCODEC_NUM_THREADS: {TORCHCODEC_NUM_THREADS}") - video_path = ele["video"] + video_path = strip_file_uri(ele["video"]) st = time.time() decoder = VideoDecoder(video_path, num_ffmpeg_threads=TORCHCODEC_NUM_THREADS) video_fps = decoder.metadata.average_fps diff --git a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt index d392d1e1b..38fabbeef 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt +++ b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt @@ -3,6 +3,8 @@ loguru safetensors numpy pillow +packaging +requests decord imageio[ffmpeg] wandb From 69263af81add0d436e5eccf437dfaed9be558d18 Mon Sep 17 00:00:00 2001 From: zirui Date: Thu, 18 Jun 2026 13:59:02 +0000 Subject: [PATCH 21/30] feat(diffusion): add FLUX training support Add FLUX.1-dev pretraining support for precomputed and raw image-text data so the diffusion backend can run FLUX workloads without TorchTitan framework dependencies. Co-authored-by: Cursor --- .gitignore | 4 + .../MI355X/flux.1_dev_t2i-pretrain.yaml | 52 +++ .../MI355X/flux.1_dev_t2i-raw-pretrain.yaml | 63 +++ primus/backends/diffusion/README.md | 133 +++++- primus/backends/diffusion/argument_builder.py | 138 +++++- primus/backends/diffusion/data/__init__.py | 22 +- primus/backends/diffusion/data/collator.py | 88 ++++ primus/backends/diffusion/data/config.py | 89 ++++ primus/backends/diffusion/data/dataset.py | 301 +++++++++++++ .../diffusion/data/flux_precomputed.py | 312 +++++++++++++ .../diffusion/data/processing_wanvideo.py | 416 ++++++++++++++++++ primus/backends/diffusion/data/processor.py | 180 ++++++++ .../diffusion/data/registrations/__init__.py | 2 +- .../diffusion/data/registrations/flux.py | 32 ++ .../diffusion/data/registrations/wan.py | 190 +------- .../backends/diffusion/diffusion_adapter.py | 11 +- .../diffusion/diffusion_pretrain_trainer.py | 2 + .../diffusion/models/flux/__init__.py | 11 + .../backends/diffusion/models/flux/adapter.py | 102 +++++ .../diffusion/models/flux/autoencoder.py | 301 +++++++++++++ .../diffusion/models/flux/conditioner.py | 48 ++ .../models/flux/configuration_flux.py | 8 + .../backends/diffusion/models/flux/layers.py | 203 +++++++++ primus/backends/diffusion/models/flux/math.py | 31 ++ .../backends/diffusion/models/flux/model.py | 151 +++++++ .../diffusion/models/flux/train_pipeline.py | 149 +++++++ .../backends/diffusion/models/flux/utils.py | 44 ++ .../diffusion/models/registrations/flux.py | 199 +++++++++ primus/backends/diffusion/registry.py | 14 + primus/backends/diffusion/trainers/base.py | 14 + primus/backends/diffusion/trainers/fsdp2.py | 4 + .../models/diffusion/flux.1_dev_t2i.yaml | 18 + .../hooks/train/pretrain/diffusion/prepare.py | 69 ++- .../diffusion/requirements-diffusion.txt | 3 + .../backends/diffusion/test_flux_backend.py | 213 +++++++++ 35 files changed, 3419 insertions(+), 198 deletions(-) create mode 100644 examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml create mode 100644 examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml create mode 100644 primus/backends/diffusion/data/collator.py create mode 100644 primus/backends/diffusion/data/config.py create mode 100644 primus/backends/diffusion/data/dataset.py create mode 100644 primus/backends/diffusion/data/flux_precomputed.py create mode 100644 primus/backends/diffusion/data/processing_wanvideo.py create mode 100644 primus/backends/diffusion/data/processor.py create mode 100644 primus/backends/diffusion/data/registrations/flux.py create mode 100644 primus/backends/diffusion/models/flux/__init__.py create mode 100644 primus/backends/diffusion/models/flux/adapter.py create mode 100644 primus/backends/diffusion/models/flux/autoencoder.py create mode 100644 primus/backends/diffusion/models/flux/conditioner.py create mode 100644 primus/backends/diffusion/models/flux/configuration_flux.py create mode 100644 primus/backends/diffusion/models/flux/layers.py create mode 100644 primus/backends/diffusion/models/flux/math.py create mode 100644 primus/backends/diffusion/models/flux/model.py create mode 100644 primus/backends/diffusion/models/flux/train_pipeline.py create mode 100644 primus/backends/diffusion/models/flux/utils.py create mode 100644 primus/backends/diffusion/models/registrations/flux.py create mode 100644 primus/configs/models/diffusion/flux.1_dev_t2i.yaml create mode 100644 tests/unit_tests/backends/diffusion/test_flux_backend.py diff --git a/.gitignore b/.gitignore index 3353a8bdb..126a777d9 100644 --- a/.gitignore +++ b/.gitignore @@ -12,4 +12,8 @@ local/ output experiment /data/* +!primus/backends/diffusion/data/ +!primus/backends/diffusion/data/** +primus/backends/diffusion/data/**/__pycache__/ +primus/backends/diffusion/data/**/*.pyc pp_simulation_result diff --git a/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml new file mode 100644 index 000000000..561c05acc --- /dev/null +++ b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml @@ -0,0 +1,52 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:flux.1_dev_t2i-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + + model: flux.1_dev_t2i.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: ${LOCAL_BATCH_SIZE:1} + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/flux.1_dev_t2i-pretrain} + save_steps: 0 + save_strategy: ${SAVE_STRATEGY:dit_only} + run_name: flux.1_dev_t2i-pretrain + + data: + dataset_path: ${DATASET_PATH:/mnt/vast/zirui/data/cc12m_preprocessed} + empty_encodings_path: ${EMPTY_ENCODINGS_PATH:/mnt/vast/zirui/data/empty_encodings} + prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1} + img_size: ${IMG_SIZE:256} + + parallelism: + sp_size: 1 + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: ${LR:2.0e-4} + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:sdpa} + gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false} + fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true} + report_to: none diff --git a/examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml new file mode 100644 index 000000000..11b59ac4a --- /dev/null +++ b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml @@ -0,0 +1,63 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:flux.1_dev_t2i-raw-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + + model: flux.1_dev_t2i.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + model: + config: + encoder: + t5_encoder: ${T5_ENCODER:google/t5-v1_1-xxl} + clip_encoder: ${CLIP_ENCODER:openai/clip-vit-large-patch14} + autoencoder: ${VAE_CHECKPOINT:black-forest-labs/FLUX.1-dev/ae.safetensors} + max_t5_length: ${MAX_T5_LENGTH:256} + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: ${LOCAL_BATCH_SIZE:1} + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/flux.1_dev_t2i-raw-pretrain} + save_steps: 0 + save_strategy: ${SAVE_STRATEGY:dit_only} + run_name: flux.1_dev_t2i-raw-pretrain + + data: + dataset_type: raw + dataset: ${DATASET:cc12m-test} + dataset_format: ${DATASET_FORMAT:webdataset} + dataset_path: ${DATASET_PATH:} + prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1} + img_size: ${IMG_SIZE:256} + skip_low_resolution: false + + parallelism: + sp_size: 1 + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: ${LR:2.0e-4} + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:sdpa} + gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false} + fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true} + report_to: none diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md index 9b7db9704..a8a786c89 100644 --- a/primus/backends/diffusion/README.md +++ b/primus/backends/diffusion/README.md @@ -11,9 +11,10 @@ so Wan training is a first-class part of the Primus `diffusion` backend. Supported scope: -- Model implementation: `wan` for Wan2.1 and Wan2.2. +- Model implementation: `wan` for Wan2.1/Wan2.2 and `flux` for FLUX.1-dev T2I. - Trainer: FSDP2 only. -- Sequence parallelism: Ulysses SP via `trainer.args.sp_size`. +- Sequence parallelism: Ulysses SP via `trainer.args.sp_size` for Wan. FLUX + currently requires `sp_size=1`. Wan-specific dependencies are kept out of top-level Primus requirements. See `runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt`. @@ -52,6 +53,126 @@ data: video_backend: imageio ``` +### FLUX Checkpoints And Data + +FLUX supports two training-data modes: + +- `dataset_type: precomputed`: samples already contain T5 encodings, CLIP + encodings, and VAE latent statistics. This is the fastest path and matches the + MLPerf/TorchTitan preprocessed recipe. +- `dataset_type: raw`: samples contain original image-text pairs. Primus loads + frozen T5, CLIP, and FLUX autoencoder modules and computes the encodings online + during training. This is useful for bring-up and custom datasets, but it is + slower and uses more memory. + +#### FLUX Pretrain Checkpoints + +For FLUX pretraining, the DiT checkpoint is optional: + +```yaml +model: + config: + load_from_pretrained_path: "" # empty means random DiT init +``` + +Raw image-text training needs frozen encoders/autoencoder. These can be local +paths or Hugging Face identifiers in the YAML: + +```yaml +model: + config: + encoder: + t5_encoder: google/t5-v1_1-xxl + clip_encoder: openai/clip-vit-large-patch14 + autoencoder: black-forest-labs/FLUX.1-dev/ae.safetensors +``` + +`t5_encoder` and `clip_encoder` are loaded with Hugging Face `from_pretrained`. +`autoencoder` accepts either a local safetensors path or `repo_id/filename`; +for example `black-forest-labs/FLUX.1-dev/ae.safetensors`. + +The precomputed MLPerf data on this machine is: + +```text +/mnt/vast/zirui/data/ + cc12m_preprocessed/ # train, precomputed text + VAE mean/logvar + coco_preprocessed/ # validation-style precomputed data + empty_encodings/ # t5_empty.npy, clip_empty.npy +``` + +The shipped FLUX precomputed example uses these paths by default: + +```text +examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml +``` + +#### FLUX Raw Data + +Raw FLUX data is selected from YAML. The supported built-in dataset names mirror +TorchTitan's FLUX recipe: + +```yaml +data: + dataset_type: raw + dataset: cc12m-test # small smoke-test webdataset +``` + +For full CC12M webdataset training, use: + +```yaml +data: + dataset_type: raw + dataset: cc12m-wds # maps to Hugging Face dataset pixparse/cc12m-wds +``` + +Do not use full `cc12m-wds` as the default smoke test; it is large. The raw FLUX +example defaults to `cc12m-test`, which maps to TorchTitan's small test tar: + +```text +/mnt/shared/zirui/code/torchtitan-main/tests/assets/cc12m_test/ + cc12m-train-0000.tar +``` + +The raw FLUX example is: + +```text +examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml +``` + +Raw FLUX data can also be a local JSONL file: + +```jsonl +{"image": "000001.jpg", "prompt": "a caption"} +{"image": "000002.jpg", "txt": "another caption"} +``` + +Use: + +```yaml +data: + dataset_type: raw + dataset_format: jsonl + dataset_path: /path/to/meta.jsonl + data_folder: /path/to/images + img_size: 256 +``` + +You can also point `dataset_path` at a local webdataset directory or a Hugging +Face dataset repo, depending on `dataset_format`: + +```yaml +data: + dataset_type: raw + dataset_format: hf_repo + dataset_path: pixparse/cc12m-wds +``` + +In short, the intended design is: checkpoint and dataset references live in YAML; +each value can be a local path when assets are staged manually, or a Hugging Face +identifier when `transformers`, `huggingface_hub`, or `datasets` should resolve +it. This keeps runs reproducible while still allowing local cache/download +behavior. + ## Checkpoints Each Wan model is a single Hugging Face repo that already bundles everything Wan @@ -115,6 +236,14 @@ before launching distributed training: python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ --config examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml +# FLUX precomputed config; DiT checkpoint is optional +python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ + --config examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml + +# FLUX raw image-text config; requires T5/CLIP access and FLUX AE checkpoint +python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ + --config examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml + # post-train config (validates modules.post_trainer) python3 runner/helpers/hooks/train/posttrain/diffusion/prepare.py \ --config examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py index b6b444ceb..95b2fcf25 100644 --- a/primus/backends/diffusion/argument_builder.py +++ b/primus/backends/diffusion/argument_builder.py @@ -13,8 +13,8 @@ from primus.core.utils.yaml_utils import nested_namespace_to_dict -class WanArgBuilder: - """Build the compact config object consumed by the Wan trainer.""" +class DiffusionArgBuilder: + """Build the compact config object consumed by diffusion trainers.""" DEFAULT_DATASET: dict[str, Any] = { "name": "wan", @@ -45,7 +45,7 @@ class WanArgBuilder: }, }, } - DEFAULT_TRAINER: dict[str, Any] = { + DEFAULT_WAN_TRAINER: dict[str, Any] = { "name": "fsdp2", "args": { "output_dir": "./output/wan", @@ -86,6 +86,65 @@ class WanArgBuilder: }, }, } + DEFAULT_FLUX_DATASET: dict[str, Any] = { + "name": "flux", + "config": { + "dataset_type": "precomputed", + "dataset_format": "hf_dataset", + "dataset_path": "/path/to/flux_precomputed_dataset", + "dataset": None, + "shuffle": True, + "processor_config": { + "processor_name": "flux_precomputed", + "processor_type": "flux_precomputed", + "prompt_dropout_prob": 0.0, + "empty_encodings_path": None, + "img_size": 256, + "skip_low_resolution": True, + }, + }, + } + DEFAULT_FLUX_TRAINER: dict[str, Any] = { + "name": "fsdp2", + "args": { + "output_dir": "./output/flux", + "per_device_train_batch_size": 1, + "per_device_eval_batch_size": 1, + "gradient_accumulation_steps": 1, + "gradient_checkpointing": False, + "attention_backend": "sdpa", + "learning_rate": 2.0e-4, + "lr_scheduler_type": "constant", + "warmup_steps": 0, + "weight_decay": 0.01, + "num_train_epochs": 1, + "max_steps": 100, + "logging_steps": 1, + "save_steps": 0, + "dataloader_num_workers": 4, + "report_to": "none", + "run_name": "flux-fsdp2", + "bf16": True, + "seed": 10007, + "optim": "adamw_torch", + "adam_beta1": 0.9, + "adam_beta2": 0.999, + "adam_epsilon": 1.0e-8, + "max_grad_norm": 1.0, + "fsdp2_wrap_target": "dit", + "fsdp_transformer_layer_cls_to_wrap": "DoubleStreamBlock,SingleStreamBlock", + "fsdp2_reshard_after_forward": True, + "save_strategy": "dit_only", + "sp_size": 1, + "dp_replicate": 1, + "flow_match_scheduler": { + "shift": 1, + "sigma_min": 0.0, + "extra_one_step": False, + "num_train_timesteps": 1000, + }, + }, + } def __init__(self) -> None: self._params: dict[str, Any] = {} @@ -96,16 +155,16 @@ def update(self, params: Any) -> None: elif isinstance(params, dict): self._params = copy.deepcopy(params) else: - raise TypeError(f"WanArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}") + raise TypeError(f"DiffusionArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}") def finalize(self) -> SimpleNamespace: params = copy.deepcopy(self._params) if "model" not in params: - raise ValueError("Wan backend config requires a model preset.") + raise ValueError("Diffusion backend config requires a model preset.") for legacy_section in ("dataset", "trainer"): if legacy_section in params: raise ValueError( - f"Wan backend no longer accepts public `{legacy_section}` overrides. " + f"Diffusion backend no longer accepts public `{legacy_section}` overrides. " "Use Primus-style `data`, `training`, `parallelism`, `optimizer`, " "`runtime`, and `metrics` sections instead." ) @@ -138,18 +197,48 @@ def _get_any(source: dict[str, Any], *keys: str) -> Any: return source[key] return None + @staticmethod + def _coerce_bool(value: Any) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + return bool(value) + + def _model_family(self, params: dict[str, Any]) -> str: + model = params.get("model") or {} + if not isinstance(model, dict): + raise TypeError(f"Diffusion model preset must be a dict, got {type(model).__name__}") + name = str(model.get("name", "")).strip().lower() + if not name: + raise ValueError("Diffusion model preset requires `model.name`.") + return name + + def _defaults_for_model(self, model_name: str) -> tuple[dict[str, Any], dict[str, Any]]: + if model_name == "wan": + return self.DEFAULT_DATASET, self.DEFAULT_WAN_TRAINER + if model_name == "flux": + return self.DEFAULT_FLUX_DATASET, self.DEFAULT_FLUX_TRAINER + raise ValueError(f"Unsupported diffusion model name: {model_name!r}") + def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, Any]: - """Translate concise Primus-style Wan sections into trainer arguments. + """Translate concise Primus-style sections into trainer arguments. Public configs use high-level sections such as `training`, `data`, - `parallelism`, and `runtime`; internal defaults supply the compact - `dataset` and `trainer` objects consumed by the Wan runtime. + `parallelism`, and `runtime`; internal defaults supply compact + `dataset` and `trainer` objects consumed by the diffusion runtime. """ + model_name = self._model_family(params) + default_dataset, default_trainer = self._defaults_for_model(model_name) normalized = { "model": params["model"], - "dataset": copy.deepcopy(self.DEFAULT_DATASET), - "trainer": copy.deepcopy(self.DEFAULT_TRAINER), + "dataset": copy.deepcopy(default_dataset), + "trainer": copy.deepcopy(default_trainer), "stage": params.get("stage", "pretrain"), "primus": params.get("primus", {}), } @@ -171,10 +260,12 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, ("gradient_accumulation_steps",): ("gradient_accumulation_steps",), ("output_dir",): ("output_dir",), ("save_steps",): ("save_steps",), + ("save_strategy",): ("save_strategy",), ("run_name",): ("run_name",), ("num_train_epochs",): ("num_train_epochs",), ("dataloader_num_workers",): ("dataloader_num_workers",), ("resume_from_checkpoint",): ("resume_from_checkpoint",), + ("dataloader_num_workers",): ("dataloader_num_workers",), } for source_path, target_path in training_map.items(): value = self._get_any(training, *source_path) @@ -189,6 +280,14 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, ("text_tokenizer",): ("processor_config", "text_tokenizer"), ("processor_name",): ("processor_config", "processor_name"), ("processor_type",): ("processor_config", "processor_type"), + ("dataset_format",): ("dataset_format",), + ("dataset_type",): ("dataset_type",), + ("dataset",): ("dataset",), + ("shuffle",): ("shuffle",), + ("empty_encodings_path",): ("processor_config", "empty_encodings_path"), + ("prompt_dropout_prob",): ("processor_config", "prompt_dropout_prob"), + ("img_size",): ("processor_config", "img_size"), + ("skip_low_resolution",): ("processor_config", "skip_low_resolution"), } for source_path, target_path in data_map.items(): value = self._get_any(data, *source_path) @@ -237,11 +336,21 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, ("attention_backend",): ("attention_backend",), ("report_to",): ("report_to",), ("seed",): ("seed",), + ("bf16",): ("bf16",), + ("fp16",): ("fp16",), + ("gradient_checkpointing",): ("gradient_checkpointing",), ("fsdp2_reshard_after_forward",): ("fsdp2_reshard_after_forward",), } for source_path, target_path in runtime_map.items(): value = self._get_any(runtime, *source_path) if value is not None: + if target_path in { + ("bf16",), + ("fp16",), + ("gradient_checkpointing",), + ("fsdp2_reshard_after_forward",), + }: + value = self._coerce_bool(value) self._set_nested(trainer_args, target_path, value) log_freq = metrics.get("log_freq") @@ -259,4 +368,11 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, if resume_from_checkpoint is not None: self._set_nested(trainer_args, ("resume_from_checkpoint",), resume_from_checkpoint) + if model_name == "flux" and int(trainer_args.get("sp_size", 1)) != 1: + raise ValueError("FLUX diffusion training currently requires `parallelism.sp_size: 1`.") + return normalized + + +# Backwards-compatible import path for existing Wan-only callers. +WanArgBuilder = DiffusionArgBuilder diff --git a/primus/backends/diffusion/data/__init__.py b/primus/backends/diffusion/data/__init__.py index 18ae84901..b72452684 100644 --- a/primus/backends/diffusion/data/__init__.py +++ b/primus/backends/diffusion/data/__init__.py @@ -4,4 +4,24 @@ # See LICENSE for license information. ############################################################################### -"""Data builders for the diffusion backend.""" +"""Data modules for the Primus diffusion backend.""" + +from .config import DatasetConfig +from .dataset import WanVideoDataset +from .flux_precomputed import ( + FluxPrecomputedDataset, + FluxPrecomputedProcessor, + FluxRawImageTextDataset, + FluxRawImageTextProcessor, +) +from .processor import WanVideoDataProcessor + +__all__ = [ + "DatasetConfig", + "FluxPrecomputedDataset", + "FluxPrecomputedProcessor", + "FluxRawImageTextDataset", + "FluxRawImageTextProcessor", + "WanVideoDataProcessor", + "WanVideoDataset", +] diff --git a/primus/backends/diffusion/data/collator.py b/primus/backends/diffusion/data/collator.py new file mode 100644 index 000000000..9dc07d814 --- /dev/null +++ b/primus/backends/diffusion/data/collator.py @@ -0,0 +1,88 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import collections +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass +class VisionCollator: + processor: Any + + def pad_sequence(self, input_ids, batch_first, padding_value): + if self.processor.tokenizer.padding_side == "left": + input_ids = [torch.flip(_input_ids, [0]) for _input_ids in input_ids] + input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=batch_first, padding_value=padding_value) + if self.processor.tokenizer.padding_side == "left": + input_ids = torch.flip(input_ids, [1]) + return input_ids + + def __call__(self, instances: Sequence[dict]) -> dict[str, torch.Tensor]: + if isinstance(instances[0], list): + instances = [inst for instance in instances for inst in instance] + inputs = collections.defaultdict(list) + for instance in instances: + for key, values in instance.items(): + inputs[key].append(values) + + batched_inputs = {} + if "input_ids" in inputs: + input_ids = inputs.pop("input_ids") + input_ids = self.pad_sequence( + input_ids, + batch_first=True, + padding_value=self.processor.tokenizer.pad_token_id, + ) + batched_inputs["input_ids"] = input_ids + if "labels" in inputs: + labels = inputs.pop("labels") + labels = self.pad_sequence( + labels, + batch_first=True, + padding_value=-100, + ) + batched_inputs["labels"] = labels + + if "attention_mask" in inputs: + inputs.pop("attention_mask") + + attention_mask = input_ids.ne(self.processor.tokenizer.pad_token_id).long() + batched_inputs["attention_mask"] = attention_mask + + # for the other keys + for key, values in inputs.items(): + # Handle scalar/boolean values ( use_audio_in_video) + if isinstance(values[0], bool) or ( + isinstance(values[0], (int, float)) and not isinstance(values[0], torch.Tensor) + ): + batched_inputs[key] = values[0] + else: + batched_inputs[key] = torch.stack(values, dim=0) + return batched_inputs + + @property + def image_token_id(self): + return self.processor.tokenizer.convert_tokens_to_ids(self.processor.image_token) + + +@dataclass +class RawBatchCollator: + """ + A minimal collator that returns raw samples as a list of dicts. + + This is useful when model-specific padding/encoding happens in a separate + batch preparation step (e.g. processor.prepare_batch / trainer.prepare_batch), + keeping the DataLoader and trainer model-agnostic. + """ + + def __call__(self, instances: Sequence[dict]) -> list[dict]: + if isinstance(instances[0], list): + instances = [inst for instance in instances for inst in instance] + return list(instances) diff --git a/primus/backends/diffusion/data/config.py b/primus/backends/diffusion/data/config.py new file mode 100644 index 000000000..229145758 --- /dev/null +++ b/primus/backends/diffusion/data/config.py @@ -0,0 +1,89 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from typing import Any, Literal + +from pydantic import BaseModel, field_validator + + +class Args(BaseModel): + extra_kwargs: dict[str, Any] = {} + + def to_dict(self): + return self.model_dump() + + def to_json(self): + return self.model_dump_json() + + +class ProcessorConfig(Args): + processor_name: str + processor_type: str + + +class DatasetConfig(Args): + dataset_type: str + data_folder: str + dataset_format: Literal["json", "jsonl", "csv", "yaml", "hf_dataset", "arrow"] + processor_config: dict | ProcessorConfig + + # Dataset configuration + dataset_path: str | None = None # Optional - used for external files + datasets: list[dict] | None = None # Optional - used for inline YAML definitions + shuffle: bool = True + data_seed: int | None = 42 + eval_dataset_path: str | None = None + + # Object storage configuration + object_storage: Literal["azure", "gcs", "none"] | None = "none" + bucket_name: str | None = None + + # Packing configuration + packing: bool | None = False + packing_strategy: str | None = None + packing_length: int | None = 32000 + filter_overlong: bool | None = True + filter_overlong_workers: int | None = 8 + max_length: int | None = None + + # Video configuration + video_sampling_strategy: Literal["fps", "frame_num"] | None = "fps" + video_max_pixels: int | None = 768 * 28 * 28 + video_max_frames: int | None = 768 + video_min_pixels: int | None = 3136 + frame_num: int | None = 64 + fps: int | None = 1 + video_backend: Literal["decord", "qwen_vl_utils", "qwen_omni_utils", "imageio"] | None = "qwen_vl_utils" + + @field_validator( + "video_max_pixels", + "video_max_frames", + "frame_num", + "fps", + "packing_length", + "max_length", + "filter_overlong_workers", + ) + @classmethod + def validate_positive_values(cls, v, info): + """Validate that numeric video and packing parameters are positive.""" + if v is not None and v <= 0: + field_name = info.field_name + raise ValueError(f"{field_name} must be positive, got {v}") + return v + + @field_validator("video_backend") + @classmethod + def validate_video_backend_migration(cls, v): + """Provide migration warning for deprecated torchvision backend.""" + if v == "torchvision": + raise ValueError( + "The 'torchvision' video backend has been removed. " + "Please use 'decord', 'qwen_vl_utils', or 'qwen_omni_utils' instead. " + "Migration guide: If you were using torchvision, 'decord' provides " + "similar functionality with better performance." + ) + return v \ No newline at end of file diff --git a/primus/backends/diffusion/data/dataset.py b/primus/backends/diffusion/data/dataset.py new file mode 100644 index 000000000..a528b3db1 --- /dev/null +++ b/primus/backends/diffusion/data/dataset.py @@ -0,0 +1,301 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Dataset class for loading video data from JSONL/CSV files. +""" + +import json +import os +from io import BytesIO +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +from primus.backends.diffusion.utils.data_utils import smart_nframes +from primus.backends.diffusion.utils.vision_process import strip_file_uri + +from .collator import RawBatchCollator + + +def _get_decord_vr( + video_path: str, + *, + num_threads: int, +): + """Construct a decord.VideoReader.""" + from decord import VideoReader, cpu + + return VideoReader(video_path, ctx=cpu(0), num_threads=num_threads) + + +class BaseDataset(Dataset): + """Base dataset with minimal shared interface.""" + + def __init__(self, config, **kwargs) -> None: + super().__init__() + self.config = config + self.samples = [] + + def __len__(self): + return len(self.samples) + + +class WanVideoDataset(BaseDataset): + """Dataset for WanVideo training from JSONL/CSV.""" + + def __init__(self, processor, config={}): + """ + Initialize WanVideo dataset. + + Args: + processor: WanVideoDataProcessor instance + video_backend: Backend for video loading ('qwen_vl_utils' or 'decord') + """ + super().__init__(config) + self.config = config + self.data_path = Path(self.config.dataset_path) + self.processor = processor + + # Load metadata + self.samples = self._load_metadata() + self._sync_processor_video_limits() + + def _sync_processor_video_limits(self): + image_processor = None + if hasattr(self.processor, "processor") and hasattr(self.processor.processor, "image_processor"): + image_processor = self.processor.processor.image_processor + if image_processor is None: + return + if ( + getattr(image_processor, "max_pixels", None) is None + and getattr(self.config, "video_max_pixels", None) is not None + ): + image_processor.max_pixels = self.config.video_max_pixels + + def _load_metadata(self) -> list[dict]: + """Load metadata from JSONL or CSV file.""" + samples = [] + # TODO: add dummy data for debugging + if self.data_path.name == "dummy.jsonl": + return samples + + if self.data_path.suffix == ".jsonl": + with open(self.data_path) as f: + for line in f: + samples.append(json.loads(line.strip())) + elif self.data_path.suffix == ".json": + with open(self.data_path) as f: + samples = json.load(f) + elif self.data_path.suffix == ".csv": + import pandas as pd + + df = pd.read_csv(self.data_path) + samples = df.to_dict("records") + else: + raise ValueError(f"Unsupported file format: {self.data_path=}") + + return samples + + def _load_video_frames(self, video_path: str, data_folder=None, fps: int = 1) -> tuple[np.ndarray, float]: + """Load video frames using the specified backend.""" + if self.config.data_folder is not None: + video_path = os.path.join(self.config.data_folder, video_path) + + if self.config.video_backend == "decord": + return self.load_video_decord(video_path, fps) + elif self.config.video_backend == "qwen_vl_utils": + return self.load_video_qwen_vl_utils(video_path, fps) + elif self.config.video_backend == "imageio": + return self.load_video_imageio(video_path, fps) + else: + raise ValueError(f"Unsupported video backend: {self.config.video_backend}") + + def load_video_imageio(self, video_path, fps): + import imageio + + video_path = strip_file_uri(video_path) + reader = imageio.get_reader(video_path) + + # Sampling Strategy + total_frames = reader.count_frames() + total_frames = int(total_frames) + + if self.config.video_sampling_strategy == "frame_num": + nframes = self.config.frame_num + # Enforce VAE divisibility: (n - 1) % 4 == 0 + actual_nframes = min(nframes, total_frames) + + valid_nframes = (actual_nframes - 1) // 4 * 4 + 1 if actual_nframes > 1 else 1 + + # DiffSynth Sequential Reading + frames = [] + for i, frame in enumerate(reader): + if i >= valid_nframes: + break + frames.append(frame) + + # Stack to numpy (T, H, W, C) + frames = np.array(frames) + sample_fps = fps # Simplification + + reader.close() + else: + reader.close() + raise NotImplementedError("Only frame_num strategy implemented for imageio backend") + + return frames, sample_fps + + def load_video_decord( + self, + video_path: str | list[str] | BytesIO, + fps: int, + ) -> tuple[np.ndarray, float]: + """ + Load video using Decord backend. + + Args: + video_path: Path to video file or BytesIO object + fps: Target frames per second + + Returns: + Tuple of (video frames, sample fps) + """ + # Keep dataset logic simple: use a fixed thread count here. + from decord import VideoReader, cpu + + num_threads = 4 + align = os.getenv("ALIGN_WITH_DIFFSYNTH") == "1" + + if isinstance(video_path, BytesIO): + vr = VideoReader(video_path, ctx=cpu(0), num_threads=num_threads) + elif isinstance(video_path, list): + vr = _get_decord_vr(strip_file_uri(video_path[0]), num_threads=num_threads) + elif isinstance(video_path, str): + vr = _get_decord_vr(strip_file_uri(video_path), num_threads=num_threads) + else: + raise ValueError(f"Unsupported video path type: {type(video_path)}") + + total_frames, video_fps = len(vr), vr.get_avg_fps() + if self.config.video_sampling_strategy == "fps": + nframes = smart_nframes(total_frames, video_fps=video_fps, fps=fps) + # Maintain uniform sampling for FPS strategy + uniform_sampled_frames = np.linspace(0, total_frames - 1, nframes, dtype=int) + elif self.config.video_sampling_strategy == "frame_num": + nframes = self.config.frame_num + # Enforce VAE divisibility: (n - 1) % 4 == 0 + actual_nframes = min(nframes, total_frames) + + valid_nframes = (actual_nframes - 1) // 4 * 4 + 1 if actual_nframes > 1 else 1 + + if align: + # sequential sampling align with diffsynth + uniform_sampled_frames = np.arange(valid_nframes, dtype=int) + else: + # uniform sampling + uniform_sampled_frames = np.linspace(0, total_frames - 1, valid_nframes, dtype=int) + else: + raise ValueError(f"Invalid video sampling strategy: {self.config.video_sampling_strategy}") + + frame_idx = uniform_sampled_frames.tolist() + spare_frames = vr.get_batch(frame_idx).asnumpy() + # spare_frames = torch.tensor(spare_frames).permute(0, 3, 1, 2) # Convert to TCHW format + + # Calculate sample_fps + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + + # Return HWC numpy array to match processor expectations + return spare_frames, sample_fps # (frames, height, width, channels) + + def load_video_qwen_vl_utils( + self, + video_path: str, + fps: int, + ) -> tuple[np.ndarray, float]: + """ + Load video using Qwen VL utils. + + Args: + video_path: Path to video file + fps: Target frames per second + + Returns: + Tuple of (video frames, sample fps) + """ + from primus.backends.diffusion.utils.vision_process import fetch_video + + video_dict = { + "type": "video", + "video": f"file://{video_path}", + "min_frames": 1, + "max_pixels": self.config.video_max_pixels, + "max_frames": self.config.video_max_frames, + "min_pixels": self.config.video_min_pixels, + } + + if self.config.video_sampling_strategy == "frame_num": + is_even = self.config.frame_num % 2 == 0 + n_frames = self.config.frame_num if is_even else self.config.frame_num + 1 + video_dict["nframes"] = n_frames + + frames, sample_fps = fetch_video(video_dict, return_video_sample_fps=True) + frames = frames.numpy() + + # if is_even: + # return frames, sample_fps + # else: + # return frames[:-1], sample_fps + + # Enforce VAE divisibility constraint + actual_n = len(frames) + if actual_n > 1: + valid_n = ((actual_n - 1) // 4) * 4 + 1 + frames = frames[:valid_n] + # else: keep 1 frame (or handle error) + + return frames, sample_fps + elif self.config.video_sampling_strategy == "fps": + video_dict["fps"] = fps + frames, sample_fps = fetch_video(video_dict, return_video_sample_fps=True) + frames = frames.numpy() + + # Also enforce for fps strategy + actual_n = len(frames) + if actual_n > 1: + valid_n = ((actual_n - 1) // 4) * 4 + 1 + frames = frames[:valid_n] + + return frames, sample_fps + else: + raise ValueError(f"Invalid video sampling strategy: {self.config.video_sampling_strategy}") + + def __len__(self) -> int: + return len(self.samples) + + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: + """Get a single sample.""" + sample = self.samples[idx] + + # Load video frames + video_path = sample["video"] + video_frames, fps = self._load_video_frames(video_path) + + # Get prompt + prompt = sample.get("prompt", "") + # Return raw sample. + return { + "video_frames": video_frames, # np.ndarray, typically (T, H, W, C) + "prompt": prompt, + "fps": fps, + "num_frames": int(getattr(self.config, "frame_num", 0) or 0), + "video_path": str(video_path), + } + + def get_collator(self): + # Prefer raw collation; model-specific processing should happen in processor.prepare_batch. + return RawBatchCollator() diff --git a/primus/backends/diffusion/data/flux_precomputed.py b/primus/backends/diffusion/data/flux_precomputed.py new file mode 100644 index 000000000..e896f22e9 --- /dev/null +++ b/primus/backends/diffusion/data/flux_precomputed.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +import io +import json +import math +import os +from pathlib import Path +from dataclasses import dataclass +from typing import Any, Sequence + +import numpy as np +import torch +from PIL import Image, ImageFile +from torch.utils.data import Dataset + +from primus.backends.diffusion.data.collator import RawBatchCollator + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +def _tensor_from_serialized_numpy(value: bytes) -> torch.Tensor: + array = np.load(io.BytesIO(value)) + tensor = torch.from_numpy(array) + if tensor.dtype == torch.uint16: + tensor = tensor.view(torch.bfloat16) + return tensor + + +def _to_tensor(value: Any) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (bytes, bytearray)): + return _tensor_from_serialized_numpy(bytes(value)) + if isinstance(value, np.ndarray): + tensor = torch.from_numpy(value) + if tensor.dtype == torch.uint16: + tensor = tensor.view(torch.bfloat16) + return tensor + return torch.as_tensor(value) + + +class FluxPrecomputedDataset(Dataset): + """Map-style dataset for precomputed FLUX text and VAE encodings.""" + + required_fields = ("t5_encodings", "clip_encodings", "mean", "logvar") + + def __init__(self, dataset_path: str): + if not dataset_path: + raise ValueError("FLUX precomputed dataset requires `dataset_path`.") + if not os.path.isdir(dataset_path): + raise FileNotFoundError(f"FLUX precomputed dataset directory not found: {dataset_path}") + from datasets import load_from_disk + + self.dataset_path = dataset_path + self.dataset = load_from_disk(dataset_path) + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: + sample = self.dataset[idx] + missing = [field for field in self.required_fields if field not in sample] + if missing: + raise KeyError(f"FLUX precomputed sample missing fields: {missing}") + return {field: _to_tensor(sample[field]) for field in self.required_fields} + + def get_collator(self): + return RawBatchCollator() + + +class FluxRawImageTextDataset(Dataset): + """Map-style raw image-text dataset for online FLUX encoding.""" + + def __init__( + self, + *, + dataset_path: str | None, + dataset_format: str = "webdataset", + dataset_name: str | None = None, + data_folder: str | None = None, + ): + dataset_name = dataset_name or None + dataset_path, dataset_format = self._resolve_dataset(dataset_name, dataset_path, dataset_format) + self.dataset_path = dataset_path + self.dataset_format = dataset_format + self.dataset_name = dataset_name + self.data_folder = data_folder + self._records: list[dict[str, Any]] | None = None + + if dataset_format == "jsonl": + path = Path(dataset_path) + if not path.is_file(): + raise FileNotFoundError(f"FLUX raw jsonl metadata not found: {dataset_path}") + self._records = [json.loads(line) for line in path.read_text().splitlines() if line.strip()] + self.dataset = None + elif dataset_format == "hf_dataset": + if not os.path.isdir(dataset_path): + raise FileNotFoundError(f"FLUX raw HF dataset directory not found: {dataset_path}") + from datasets import load_from_disk + + self.dataset = load_from_disk(dataset_path) + elif dataset_format == "hf_repo": + from datasets import load_dataset + + self.dataset = load_dataset(dataset_path, split="train") + elif dataset_format == "webdataset": + if not os.path.isdir(dataset_path): + raise FileNotFoundError(f"FLUX raw webdataset directory not found: {dataset_path}") + from datasets import load_dataset + + self.dataset = load_dataset( + "webdataset", + split="train", + data_dir=dataset_path, + data_files={"train": "*.tar"}, + ) + else: + raise ValueError("FLUX raw dataset_format must be one of: jsonl, hf_dataset, hf_repo, webdataset") + + @staticmethod + def _resolve_dataset( + dataset_name: str | None, + dataset_path: str | None, + dataset_format: str, + ) -> tuple[str, str]: + if dataset_name == "cc12m-test": + return ( + dataset_path or "/mnt/shared/zirui/code/torchtitan-main/tests/assets/cc12m_test", + "webdataset", + ) + if dataset_name == "cc12m-wds": + return dataset_path or "pixparse/cc12m-wds", "hf_repo" + if not dataset_path: + raise ValueError("FLUX raw dataset requires either `dataset` or `dataset_path`.") + return dataset_path, dataset_format + + def __len__(self) -> int: + if self._records is not None: + return len(self._records) + return len(self.dataset) + + @staticmethod + def _prompt_from_sample(sample: dict[str, Any]) -> str: + for key in ("txt", "caption", "prompt", "text"): + if key in sample: + value = sample[key] + if isinstance(value, list): + value = value[0] + return str(value) + raise KeyError("FLUX raw sample requires one of: txt, caption, prompt, text") + + def _image_from_jsonl_sample(self, sample: dict[str, Any]) -> Image.Image: + image_key = "image" if "image" in sample else "jpg" if "jpg" in sample else "png" + image_path = Path(str(sample[image_key])) + if self.data_folder and not image_path.is_absolute(): + image_path = Path(self.data_folder) / image_path + return Image.open(image_path).convert("RGB") + + @staticmethod + def _image_from_dataset_sample(sample: dict[str, Any]) -> Image.Image: + if "image" in sample: + image = sample["image"] + elif "jpg" in sample: + image = sample["jpg"] + elif "png" in sample: + image = sample["png"] + else: + raise KeyError("FLUX raw dataset sample requires image, jpg, or png") + if isinstance(image, Image.Image): + return image.convert("RGB") + if isinstance(image, (bytes, bytearray)): + return Image.open(io.BytesIO(image)).convert("RGB") + return image.convert("RGB") + + def __getitem__(self, idx: int) -> dict[str, Any]: + if self._records is not None: + sample = self._records[idx] + image = self._image_from_jsonl_sample(sample) + else: + sample = self.dataset[idx] + image = self._image_from_dataset_sample(sample) + return {"image": image, "prompt": self._prompt_from_sample(sample)} + + def get_collator(self): + return RawBatchCollator() + + +@dataclass +class FluxPrecomputedProcessor: + config: dict[str, Any] + + def __post_init__(self): + self.prompt_dropout_prob = float(self.config.get("prompt_dropout_prob", 0.0) or 0.0) + self.empty_encodings_path = self.config.get("empty_encodings_path") + self.img_size = int(self.config.get("img_size", 256) or 256) + self._empty_t5: torch.Tensor | None = None + self._empty_clip: torch.Tensor | None = None + + def build(self): + if self.prompt_dropout_prob > 0.0: + self._load_empty_encodings() + + @staticmethod + def _normalize_empty_encoding(tensor: torch.Tensor) -> torch.Tensor: + if tensor.ndim >= 1 and tensor.shape[0] == 1: + return tensor[0] + return tensor + + def _load_empty_encodings(self): + if self._empty_t5 is not None and self._empty_clip is not None: + return + if not self.empty_encodings_path: + raise ValueError("FLUX prompt dropout requires `data.empty_encodings_path`.") + t5_path = os.path.join(self.empty_encodings_path, "t5_empty.npy") + clip_path = os.path.join(self.empty_encodings_path, "clip_empty.npy") + if not os.path.isfile(t5_path) or not os.path.isfile(clip_path): + raise FileNotFoundError( + "FLUX empty encodings must contain t5_empty.npy and clip_empty.npy " + f"under {self.empty_encodings_path}" + ) + self._empty_t5 = self._normalize_empty_encoding(torch.from_numpy(np.load(t5_path))) + self._empty_clip = self._normalize_empty_encoding(torch.from_numpy(np.load(clip_path))) + + @staticmethod + def _collate_raw(batch: Any) -> dict[str, torch.Tensor]: + if isinstance(batch, dict): + return {key: _to_tensor(value) for key, value in batch.items()} + if not isinstance(batch, Sequence) or not batch: + raise ValueError(f"FLUX prepare_batch expected non-empty sequence, got {type(batch).__name__}") + + keys = ("t5_encodings", "clip_encodings", "mean", "logvar") + collated: dict[str, torch.Tensor] = {} + for key in keys: + collated[key] = torch.stack([_to_tensor(sample[key]) for sample in batch], dim=0) + return collated + + def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) -> dict[str, torch.Tensor]: + tensors = self._collate_raw(batch) + for key, value in tensors.items(): + tensors[key] = value.to(device=device, dtype=dtype, non_blocking=True) + + if self.prompt_dropout_prob > 0.0: + self._load_empty_encodings() + assert self._empty_t5 is not None and self._empty_clip is not None + bsz = tensors["t5_encodings"].shape[0] + drop_mask = torch.rand((bsz,), device=device) < self.prompt_dropout_prob + if drop_mask.any(): + tensors["t5_encodings"][drop_mask] = self._empty_t5.to(device=device, dtype=dtype) + tensors["clip_encodings"][drop_mask] = self._empty_clip.to(device=device, dtype=dtype) + + return tensors + + +@dataclass +class FluxRawImageTextProcessor: + config: dict[str, Any] + + def __post_init__(self): + self.prompt_dropout_prob = float(self.config.get("prompt_dropout_prob", 0.0) or 0.0) + self.img_size = int(self.config.get("img_size", 256) or 256) + self.skip_low_resolution = bool(self.config.get("skip_low_resolution", True)) + + def build(self): + return + + def _process_image(self, image: Image.Image) -> torch.Tensor | None: + width, height = image.size + if self.skip_low_resolution and (width < self.img_size or height < self.img_size): + return None + + if width == self.img_size and height == self.img_size: + resized = image + elif width >= height: + new_width, new_height = math.ceil(self.img_size / height * width), self.img_size + image = image.resize((new_width, new_height), Image.Resampling.BICUBIC) + left = int(torch.randint(0, new_width - self.img_size + 1, (1,)).item()) + resized = image.crop((left, 0, left + self.img_size, self.img_size)) + else: + new_width, new_height = self.img_size, math.ceil(self.img_size / width * height) + image = image.resize((new_width, new_height), Image.Resampling.BICUBIC) + lower = int(torch.randint(0, new_height - self.img_size + 1, (1,)).item()) + resized = image.crop((0, lower, self.img_size, lower + self.img_size)) + + if resized.mode != "RGB": + resized = resized.convert("RGB") + np_img = np.array(resized).transpose((2, 0, 1)) + return torch.tensor(np_img).float() / 255.0 * 2.0 - 1.0 + + def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) -> dict[str, Any]: + if isinstance(batch, dict): + batch = [batch] + if not isinstance(batch, Sequence) or not batch: + raise ValueError(f"FLUX raw prepare_batch expected non-empty sequence, got {type(batch).__name__}") + + images: list[torch.Tensor] = [] + prompts: list[str] = [] + for sample in batch: + image = self._process_image(sample["image"]) + if image is None: + continue + prompt = str(sample.get("prompt", "")) + if self.prompt_dropout_prob > 0.0 and torch.rand(1).item() < self.prompt_dropout_prob: + prompt = "" + images.append(image) + prompts.append(prompt) + + if not images: + raise ValueError("FLUX raw batch contained no usable images after preprocessing") + return { + "image": torch.stack(images, dim=0).to(device=device, dtype=dtype, non_blocking=True), + "prompts": prompts, + } diff --git a/primus/backends/diffusion/data/processing_wanvideo.py b/primus/backends/diffusion/data/processing_wanvideo.py new file mode 100644 index 000000000..64ffddb5f --- /dev/null +++ b/primus/backends/diffusion/data/processing_wanvideo.py @@ -0,0 +1,416 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import warnings +from typing import Any, Dict, List, Optional, Union, Tuple + +import os +import numpy as np +import torch +from PIL import Image +from transformers import AutoTokenizer +from transformers.image_processing_utils import BaseImageProcessor +from transformers.image_utils import ImageInput +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput +from transformers.utils import TensorType, logging + +logger = logging.get_logger(__name__) + + +class WanVideoImageProcessor(BaseImageProcessor): + """ + Image/Video processor for WanVideo models. + + Args: + do_resize: Whether to resize the image/video frames. + size: Target size for resizing. + do_center_crop: Whether to center crop. + crop_size: Size for center cropping. + do_normalize: Whether to normalize pixel values. + image_mean: Mean values for normalization. + image_std: Standard deviation values for normalization. + do_convert_rgb: Whether to convert to RGB. + """ + + model_input_names = ["pixel_values"] + + def __init__( + self, + do_resize: bool = True, + size: Optional[Dict[str, int]] = None, + do_center_crop: bool = True, + crop_size: Optional[Dict[str, int]] = None, + do_normalize: bool = True, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + do_convert_rgb: bool = True, + max_pixels: Optional[int] = None, + height_division_factor: int = 16, + width_division_factor: int = 16, + **kwargs, + ): + super().__init__(**kwargs) + + self.do_resize = do_resize + self.size = size + self.do_center_crop = do_center_crop + self.crop_size = crop_size + self.do_normalize = do_normalize + self.image_mean = image_mean or [0.5, 0.5, 0.5] + self.image_std = image_std or [0.5, 0.5, 0.5] + self.do_convert_rgb = do_convert_rgb + self.max_pixels = max_pixels + self.height_division_factor = height_division_factor + self.width_division_factor = width_division_factor + + def resize( + self, + image: np.ndarray, + size: Dict[str, int], + **kwargs, + ) -> np.ndarray: + """Resize image or video frame.""" + from PIL import Image as PILImage + + image = PILImage.fromarray(image.astype(np.uint8)) + # TODO: Here, we align with DiffSynth's resize logic for debugging(may remove in the future) + if os.getenv("ALIGN_WITH_DIFFSYNTH") == "1": + # Match DiffSynth logic: scale based on max dimension ratio + width, height = image.size + target_width = size["width"] + target_height = size["height"] + + scale = max(target_width / width, target_height / height) + new_width = round(width * scale) + new_height = round(height * scale) + + from torchvision.transforms import functional as F + from torchvision.transforms import InterpolationMode + + # DiffSynth uses torchvision.transforms.resize with BILINEAR + # We must use F.resize to match exactly (antialias behavior etc) + image = F.resize(image, (new_height, new_width), interpolation=InterpolationMode.BILINEAR) + else: + image = image.resize((size["width"], size["height"]), PILImage.LANCZOS) + return np.array(image) + + def _extract_hw(self, image: Union[np.ndarray, Image.Image]) -> Tuple[int, int]: + if isinstance(image, Image.Image): + width, height = image.size + return height, width + if isinstance(image, np.ndarray): + if image.ndim == 3: + return image.shape[0], image.shape[1] + if image.ndim == 2: + return image.shape[0], image.shape[1] + raise ValueError("Unsupported image type for size extraction.") + + def _dynamic_target_size(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, int]: + height, width = self._extract_hw(image) + if self.max_pixels is not None and width * height > self.max_pixels: + scale = (width * height / self.max_pixels) ** 0.5 + height = int(height / scale) + width = int(width / scale) + height = height // self.height_division_factor * self.height_division_factor + width = width // self.width_division_factor * self.width_division_factor + height = max(self.height_division_factor, height) + width = max(self.width_division_factor, width) + return {"height": height, "width": width} + + def _resolve_sizes(self, image: Union[np.ndarray, Image.Image]) -> Tuple[Dict[str, int], Dict[str, int]]: + size = self.size + crop_size = self.crop_size + if size is None and crop_size is None: + size = self._dynamic_target_size(image) + crop_size = dict(size) + elif size is None: + size = dict(crop_size) + elif crop_size is None: + crop_size = dict(size) + size = { + "height": self._ceil_to_factor(size["height"], self.height_division_factor), + "width": self._ceil_to_factor(size["width"], self.width_division_factor), + } + crop_size = { + "height": self._ceil_to_factor(crop_size["height"], self.height_division_factor), + "width": self._ceil_to_factor(crop_size["width"], self.width_division_factor), + } + return size, crop_size + + @staticmethod + def _ceil_to_factor(value: int, factor: int) -> int: + if value % factor == 0: + return value + return (value + factor - 1) // factor * factor + + def center_crop( + self, + image: np.ndarray, + crop_size: Dict[str, int], + **kwargs, + ) -> np.ndarray: + """Center crop image or video frame.""" + h, w = image.shape[:2] + crop_h, crop_w = crop_size["height"], crop_size["width"] + + top = (h - crop_h) // 2 + left = (w - crop_w) // 2 + if image.ndim == 3: + return image[top : top + crop_h, left : left + crop_w, :] + else: + return image[top : top + crop_h, left : left + crop_w] + + def normalize( + self, + image: np.ndarray, + mean: List[float], + std: List[float], + **kwargs, + ) -> np.ndarray: + """Normalize image or video frame.""" + image = image.astype(np.float32) / 255.0 + mean = np.array(mean).reshape(1, 1, -1) + std = np.array(std).reshape(1, 1, -1) + return (image - mean) / std + + def preprocess( + self, + images: ImageInput, + do_resize: Optional[bool] = None, + size: Optional[Dict[str, int]] = None, + do_center_crop: Optional[bool] = None, + crop_size: Optional[Dict[str, int]] = None, + do_normalize: Optional[bool] = None, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + do_convert_rgb: Optional[bool] = None, + num_frames: Optional[int] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + **kwargs, + ) -> Dict[str, Any]: + """ + Preprocess images or video frames. + + Args: + images: Input images or video frames. + return_tensors: Type of tensors to return ("pt" for PyTorch). + + Returns: + Dictionary with preprocessed pixel values. + """ + do_resize = do_resize if do_resize is not None else self.do_resize + size = size if size is not None else self.size + do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop + crop_size = crop_size if crop_size is not None else self.crop_size + do_normalize = do_normalize if do_normalize is not None else self.do_normalize + image_mean = image_mean if image_mean is not None else self.image_mean + image_std = image_std if image_std is not None else self.image_std + do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb + + # Handle single image or list of images (video frames) + if not isinstance(images, list): + images = [images] + + processed_images = [] + for image in images: + if isinstance(image, Image.Image): + image = np.array(image) + + # Handle 4D tensor case (batch of video frames) + if isinstance(image, np.ndarray) and image.ndim == 4: + # Process each frame in the batch + batch_processed_frames = [] + for i in range(image.shape[0]): + frame = image[i] # Get single frame + if frame.ndim == 3 and (frame.shape[0] == 3 or frame.shape[0] == 1): + frame = np.transpose(frame, (1, 2, 0)) # (C, H, W) -> (H, W, C) + + # Convert to RGB if needed + if do_convert_rgb and frame.shape[-1] != 3: + if len(frame.shape) == 2: # Grayscale + frame = np.stack([frame] * 3, axis=-1) + elif frame.shape[-1] == 4: # RGBA + frame = frame[..., :3] + + size_, crop_size_ = self._resolve_sizes(frame) + if do_resize: + frame = self.resize(frame, size_) + if do_center_crop: + frame = self.center_crop(frame, crop_size_) + + # Normalize + if do_normalize: + frame = self.normalize(frame, image_mean, image_std) + + batch_processed_frames.append(frame) + + # Stack frames back into 4D tensor + processed_image = np.stack(batch_processed_frames, axis=0) + else: + # Handle single image case (3D or 2D) + # Convert to RGB if needed + if do_convert_rgb and image.shape[-1] != 3: + if len(image.shape) == 2: # Grayscale + image = np.stack([image] * 3, axis=-1) + elif image.shape[-1] == 4: # RGBA + image = image[..., :3] + + size_, crop_size_ = self._resolve_sizes(image) + if do_resize: + image = self.resize(image, size_) + if do_center_crop: + image = self.center_crop(image, crop_size_) + + # Normalize + if do_normalize: + image = self.normalize(image, image_mean, image_std) + + processed_image = image + + processed_images.append(processed_image) + + # Stack frames for video + processed_images = np.stack(processed_images, axis=0) # B, T, H, W, C + + # Temporal Handling (Interpolate or Truncate) + if num_frames is not None: + current_frames = processed_images.shape[1] + if current_frames > num_frames: + logger.info(f"Truncating video frames from {current_frames} to {num_frames}") + processed_images = processed_images[:, :num_frames, ...] + elif current_frames < num_frames: + logger.info(f"Interpolating video frames from {current_frames} to {num_frames}") + # Interpolate requires (B, C, T, H, W) or (B, C, H, W) - we have (B, T, H, W, C) + # Permute to (B, C, T, H, W) for interpolate + vid_tensor = torch.from_numpy(processed_images).permute(0, 4, 1, 2, 3) + + # Interpolate + vid_tensor = torch.nn.functional.interpolate( + vid_tensor, + size=(num_frames, vid_tensor.shape[3], vid_tensor.shape[4]), + mode='trilinear', + align_corners=False + ) + + # Permute back to (B, T, H, W, C) and convert to numpy + processed_images = vid_tensor.permute(0, 2, 3, 4, 1).numpy() + + # Convert to tensor if requested + if return_tensors == "pt": + processed_images = torch.from_numpy(processed_images) + # Rearrange to (B, C, T, H, W) for video (since input was B, T, H, W, C) + if processed_images.ndim == 5: + processed_images = processed_images.permute(0, 4, 1, 2, 3) + elif processed_images.ndim == 4: + # (T, H, W, C) -> (C, T, H, W) if it was list of frames + processed_images = processed_images.permute(3, 0, 1, 2) + # Add batch dimension + # processed_images = processed_images.unsqueeze(0) + + return {"pixel_values": processed_images} + + +class WanVideoProcessor: + """ + Processor for WanVideo models, combining image/video processing and text tokenization. + + Args: + image_processor: Image/video processor instance. + tokenizer: Text tokenizer instance. + """ + + attributes = ["tokenizer", "image_processor"] + valid_kwargs = [ + "chat_template", + ] + tokenizer_class = "AutoTokenizer" + image_processor_class = "WanVideoImageProcessor" + + def __init__(self, image_processor=None, tokenizer=None, max_text_length: Optional[int] = None, **kwargs): + if image_processor is None: + image_processor = WanVideoImageProcessor(**kwargs) + if tokenizer is None: + # Default to T5 tokenizer for text encoding + try: + print("[Debug-processor], loading default tokenizer: google/umt5-xxl") + tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl") + except: + logger.warning("Could not load default tokenizer, using None") + tokenizer = None + + self.image_processor = image_processor + self.tokenizer = tokenizer + self.max_text_length = max_text_length + + def __call__( + self, + text: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None, + images: Optional[ImageInput] = None, + videos: Optional[ImageInput] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + **kwargs, + ) -> Dict[str, Any]: + """ + Process text and image/video inputs. + + Args: + text: Text input(s) to tokenize. + images: Image input(s) to process. + videos: Video frames to process. + return_tensors: Type of tensors to return. + + Returns: + Dictionary with processed inputs. + """ + if text is None and images is None and videos is None: + raise ValueError("You must provide either text, images, or videos.") + + data = {} + + # Process text + if text is not None and self.tokenizer is not None: + max_length = kwargs.pop("max_length", None) + if max_length is None: + max_length = self.max_text_length or getattr(self.tokenizer, "model_max_length", 256) + text_inputs = self.tokenizer( + text, + return_tensors=return_tensors, + padding=True, + truncation=True, + max_length=max_length, + **kwargs, + ) + data.update(text_inputs) + + # Process images or video + if images is not None or videos is not None: + image_inputs = self.image_processor( + images or videos, + return_tensors=return_tensors, + **kwargs, + ) + data.update(image_inputs) + + return data + + def batch_decode(self, *args, **kwargs): + """Delegate to tokenizer's batch_decode method.""" + if self.tokenizer is None: + raise ValueError("No tokenizer available for decoding.") + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + """Delegate to tokenizer's decode method.""" + if self.tokenizer is None: + raise ValueError("No tokenizer available for decoding.") + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + """Get model input names from components.""" + tokenizer_input_names = self.tokenizer.model_input_names if self.tokenizer else [] + image_processor_input_names = self.image_processor.model_input_names + return list(set(tokenizer_input_names + image_processor_input_names)) diff --git a/primus/backends/diffusion/data/processor.py b/primus/backends/diffusion/data/processor.py new file mode 100644 index 000000000..9b940d118 --- /dev/null +++ b/primus/backends/diffusion/data/processor.py @@ -0,0 +1,180 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Processes video and text data for training. +""" + +from collections.abc import Sequence +from typing import Any + +import torch + + +class WanVideoDataProcessor: + """Standalone data processor for WanVideo training.""" + + def __init__(self, config, model_id=None): + self.config = config + self.model_id = model_id + self.processor = None + self.tokenizer = None + + def apply_prompt_template(self, hf_messages: str) -> str: + """Apply prompt template for WanVideo.""" + # WanVideo uses direct prompts without special formatting. + # Keep backward compatibility: + # - old path passed `hf_messages` list[dict] + # - new path passes plain prompt string + if isinstance(hf_messages, str): + return hf_messages + try: + return hf_messages[0]["content"][1]["text"] + except Exception: + # Best-effort fallback + return str(hf_messages) + + def save_pretrained(self, save_directory: str): + pass + + def build(self): + """Initialize the processor and tokenizer.""" + if self.processor is not None and self.tokenizer is not None: + return + + from transformers import AutoTokenizer + + from primus.backends.diffusion.data.processing_wanvideo import WanVideoProcessor as WanVideoModelProcessor + + wanvideo_kwargs = self.config.get("extra_kwargs", {}) + max_text_length = self.config.get("max_text_length") + if max_text_length is not None: + wanvideo_kwargs.setdefault("max_text_length", max_text_length) + + # Load tokenizer if specified + if self.config.get("text_tokenizer", None) is not None: + self.tokenizer = AutoTokenizer.from_pretrained(self.config.get("text_tokenizer")) + else: + self.tokenizer = None + self.processor = WanVideoModelProcessor(**wanvideo_kwargs, tokenizer=self.tokenizer) + + if self.tokenizer is None: + self.tokenizer = self.processor.tokenizer + + def _normalize_raw_batch(self, batch: Any) -> tuple[list[str], list[Any], int | None]: + """ + Normalize dataloader output into: + - prompts: list[str] + - frames_list: list[Any] (each item is typically np.ndarray (T,H,W,C)) + - num_frames: Optional[int] (only if consistent across samples) + """ + if isinstance(batch, dict): + prompts = batch.get("prompt") + frames_list = batch.get("video_frames") + num_frames = batch.get("num_frames", None) + + if prompts is None or frames_list is None: + raise ValueError("prepare_batch(dict) requires keys: 'prompt' and 'video_frames'") + if isinstance(prompts, str): + prompts = [prompts] + if not isinstance(prompts, (list, tuple)): + raise TypeError(f"'prompt' must be str or list[str], got {type(prompts)}") + if not isinstance(frames_list, (list, tuple)): + raise TypeError(f"'video_frames' must be list, got {type(frames_list)}") + + if isinstance(num_frames, (list, tuple)): + nfs = [int(x) for x in num_frames if x] + num_frames = nfs[0] if nfs and all(int(x) == int(nfs[0]) for x in nfs) else None + elif num_frames: + num_frames = int(num_frames) + else: + num_frames = None + + return list(prompts), list(frames_list), num_frames + + # RawBatchCollator returns list[dict] + if not isinstance(batch, (list, tuple)) or not batch: + raise ValueError(f"prepare_batch expected non-empty list/tuple, got {type(batch)}") + + prompts = [str(ex.get("prompt", "")) for ex in batch] + frames_list = [ex.get("video_frames") for ex in batch] + + # Prefer per-example num_frames if present and consistent; else None. + nfs = [ex.get("num_frames") for ex in batch if ex.get("num_frames")] + num_frames = nfs[0] if nfs and all(int(x) == int(nfs[0]) for x in nfs) else None + num_frames = int(num_frames) if num_frames else None + + return prompts, frames_list, num_frames + + def _tokenize_prompts(self, prompts: Sequence[str]) -> dict[str, torch.Tensor]: + formatted_prompts = [self.apply_prompt_template(p) for p in prompts] + return self.tokenizer( + formatted_prompts, + return_tensors="pt", + padding=self.config.get("padding_strategy", "max_length"), + truncation=True, + max_length=self.config.get("max_text_length", 512), + ) + + def _preprocess_videos(self, frames_list: Sequence[Any], num_frames: int | None) -> torch.Tensor: + if any(v is None for v in frames_list): + raise ValueError("prepare_batch got None in 'video_frames'") + video_inputs = self.processor.image_processor.preprocess( + list(frames_list), + num_frames=num_frames, + return_tensors="pt", + ) + if "pixel_values" not in video_inputs: + raise KeyError("image_processor.preprocess must return dict with key 'pixel_values'") + pixel_values = video_inputs["pixel_values"] + if not isinstance(pixel_values, torch.Tensor): + raise TypeError(f"pixel_values must be torch.Tensor, got {type(pixel_values)}") + if pixel_values.ndim != 5: + raise ValueError(f"Expected pixel_values shape [B,C,T,H,W], got {tuple(pixel_values.shape)}") + return pixel_values + + def _assemble_model_batch( + self, *, pixel_values: torch.Tensor, text_inputs: dict[str, torch.Tensor] + ) -> dict[str, Any]: + if "input_ids" not in text_inputs or "attention_mask" not in text_inputs: + raise KeyError("tokenizer output must include 'input_ids' and 'attention_mask'") + + out: dict[str, Any] = { + "video": pixel_values, # [B,C,T,H,W] + "input_ids": text_inputs["input_ids"], + "attention_mask": text_inputs["attention_mask"], + } + + # Legacy convenience fields (some models expect these keys). + out.setdefault("num_frames", int(pixel_values.shape[2])) + out.setdefault("height", int(pixel_values.shape[3])) + out.setdefault("width", int(pixel_values.shape[4])) + return out + + def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) -> dict[str, Any]: + """ + Convert raw dataloader batch into model inputs. + + Expected raw batch from `RawBatchCollator`: + - batch: list[dict] with keys: + - video_frames: np.ndarray (T,H,W,C) or equivalent + - prompt: str + - num_frames: int (optional) + + Output matches current training pipelines: + - video: Tensor [B, C, T, H, W] + - input_ids: Tensor [B, L] + - attention_mask: Tensor [B, L] + - (optional) height/width/num_frames for legacy wan/wan_new paths + """ + if self.processor is None or self.tokenizer is None: + # Lazy init to keep behavior robust across call sites. + self.build() + + prompts, frames_list, num_frames = self._normalize_raw_batch(batch) + text_inputs = self._tokenize_prompts(prompts) + pixel_values = self._preprocess_videos(frames_list, num_frames=num_frames) + return self._assemble_model_batch(pixel_values=pixel_values, text_inputs=text_inputs) \ No newline at end of file diff --git a/primus/backends/diffusion/data/registrations/__init__.py b/primus/backends/diffusion/data/registrations/__init__.py index b5a3705a2..4573e0a65 100644 --- a/primus/backends/diffusion/data/registrations/__init__.py +++ b/primus/backends/diffusion/data/registrations/__init__.py @@ -4,4 +4,4 @@ # See LICENSE for license information. ############################################################################### -"""Registered diffusion dataset builders.""" +"""Dataset registrations.""" diff --git a/primus/backends/diffusion/data/registrations/flux.py b/primus/backends/diffusion/data/registrations/flux.py new file mode 100644 index 000000000..a03e25a1a --- /dev/null +++ b/primus/backends/diffusion/data/registrations/flux.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from primus.backends.diffusion.data.flux_precomputed import ( + FluxPrecomputedDataset, + FluxPrecomputedProcessor, + FluxRawImageTextDataset, + FluxRawImageTextProcessor, +) +from primus.backends.diffusion.utils.log import logger + + +def build_flux_dataset(dataset_config: dict): + processor_config = dataset_config.get("processor_config", {}) or {} + dataset_type = str(dataset_config.get("dataset_type", "precomputed")).lower() + if dataset_type == "raw": + processor = FluxRawImageTextProcessor(processor_config) + processor.build() + dataset = FluxRawImageTextDataset( + dataset_path=dataset_config.get("dataset_path"), + dataset_format=dataset_config.get("dataset_format", "webdataset"), + dataset_name=dataset_config.get("dataset"), + data_folder=dataset_config.get("data_folder"), + ) + logger.info(f"Built FLUX raw image-text dataset with {len(dataset)} samples") + elif dataset_type == "precomputed": + processor = FluxPrecomputedProcessor(processor_config) + processor.build() + dataset = FluxPrecomputedDataset(dataset_config["dataset_path"]) + logger.info(f"Built FLUX precomputed dataset with {len(dataset)} samples") + else: + raise ValueError("FLUX dataset_type must be either 'precomputed' or 'raw'") + return dataset, processor diff --git a/primus/backends/diffusion/data/registrations/wan.py b/primus/backends/diffusion/data/registrations/wan.py index 76492cb2b..f832f352a 100644 --- a/primus/backends/diffusion/data/registrations/wan.py +++ b/primus/backends/diffusion/data/registrations/wan.py @@ -4,173 +4,25 @@ # See LICENSE for license information. ############################################################################### -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any - -import numpy as np -import torch -from PIL import Image -from torch.utils.data import Dataset -from torchvision.transforms import functional as F -from transformers import AutoTokenizer - -from primus.backends.diffusion.utils.vision_process import fetch_video, strip_file_uri - - -class WanVideoProcessor: - """Tokenize prompts and normalize video tensors for Wan training.""" - - def __init__(self, config: dict[str, Any]): - self.config = config - self.max_text_length = int(config.get("max_text_length", 512)) - tokenizer_path = config.get("text_tokenizer") - if not tokenizer_path: - raise ValueError("Wan dataset processor requires `text_tokenizer`.") - trust_remote_code = bool(config.get("trust_remote_code", False)) - self.tokenizer = AutoTokenizer.from_pretrained( - tokenizer_path, - trust_remote_code=trust_remote_code, - ) - - extra_kwargs = config.get("extra_kwargs", {}) or {} - size = extra_kwargs.get("size", {}) or {} - self.height = int(size.get("height", 480)) - self.width = int(size.get("width", 832)) - self.image_mean = torch.tensor(extra_kwargs.get("image_mean", [0.5, 0.5, 0.5])).view(3, 1, 1, 1) - self.image_std = torch.tensor(extra_kwargs.get("image_std", [0.5, 0.5, 0.5])).view(3, 1, 1, 1) - - def tokenize(self, prompt: str) -> dict[str, torch.Tensor]: - encoded = self.tokenizer( - prompt, - max_length=self.max_text_length, - padding="max_length", - truncation=True, - return_tensors="pt", - ) - return { - "input_ids": encoded["input_ids"].squeeze(0).long(), - "attention_mask": encoded["attention_mask"].squeeze(0).long(), - } - - def normalize_video(self, video_tchw: torch.Tensor) -> torch.Tensor: - if video_tchw.ndim != 4: - raise ValueError(f"Expected video tensor [T,C,H,W], got shape={tuple(video_tchw.shape)}") - video_tchw = video_tchw[:, :3].float() - if video_tchw.max() > 2: - video_tchw = video_tchw / 255.0 - video_tchw = F.resize(video_tchw, [self.height, self.width], antialias=True) - video_cthw = video_tchw.permute(1, 0, 2, 3).contiguous() - return (video_cthw - self.image_mean) / self.image_std - - def prepare_batch( - self, *, batch: dict[str, Any], device: torch.device, dtype: torch.dtype - ) -> dict[str, Any]: - return batch - - -class WanVideoDataset(Dataset): - def __init__(self, config: dict[str, Any], processor: WanVideoProcessor): - self.config = config - self.processor = processor - self.dataset_path = Path(config.get("dataset_path", "")) - self.data_folder = Path(config.get("data_folder", "")) - self.frame_num = int(config.get("frame_num", 81)) - self.video_backend = str(config.get("video_backend", "imageio")).lower() - - if not self.dataset_path.exists(): - raise FileNotFoundError(f"Wan dataset metadata not found: {self.dataset_path}") - with self.dataset_path.open(encoding="utf-8") as f: - self.samples = [json.loads(line) for line in f if line.strip()] - if not self.samples: - raise ValueError(f"Wan dataset metadata is empty: {self.dataset_path}") - - def __len__(self) -> int: - return len(self.samples) - - def _resolve_video_path(self, value: str) -> str: - if value.startswith(("http://", "https://", "file://")): - return value - path = Path(value) - if not path.is_absolute() and self.data_folder: - path = self.data_folder / path - return str(path) - - def _sample_video(self, video_tchw: torch.Tensor) -> torch.Tensor: - total_frames = int(video_tchw.shape[0]) - if total_frames <= 0: - raise ValueError("Video contains no frames.") - idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long() - return video_tchw[idx] - - def _read_video_imageio(self, path: str) -> torch.Tensor: - import imageio.v3 as iio - - path = strip_file_uri(path) - frames = [] - for frame in iio.imiter(path): - image = Image.fromarray(frame).convert("RGB") - frames.append(torch.from_numpy(np.asarray(image).copy())) - if not frames: - raise ValueError(f"Video contains no frames: {path}") - return torch.stack(frames).permute(0, 3, 1, 2) - - def _read_video_decord(self, path: str) -> torch.Tensor: - import decord - - path = strip_file_uri(path) - vr = decord.VideoReader(path) - total_frames = len(vr) - idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long().tolist() - return torch.from_numpy(vr.get_batch(idx).asnumpy()).permute(0, 3, 1, 2) - - def _read_video(self, path: str) -> torch.Tensor: - if self.video_backend == "imageio": - video = self._sample_video(self._read_video_imageio(path)) - elif self.video_backend == "decord": - video = self._read_video_decord(path) - else: - video = fetch_video( - { - "video": path, - "nframes": self.frame_num, - "resized_height": self.processor.height, - "resized_width": self.processor.width, - } - ) - return self.processor.normalize_video(video) - - def __getitem__(self, index: int) -> dict[str, Any]: - sample = self.samples[index] - prompt = sample.get("prompt") or sample.get("text") or sample.get("caption") or "" - video_key = sample.get("video") or sample.get("video_path") - if not video_key: - raise KeyError(f"Wan dataset sample missing `video`: index={index}") - - item = self.processor.tokenize(str(prompt)) - item["video"] = self._read_video(self._resolve_video_path(str(video_key))) - if "seed" in sample: - item["seed"] = int(sample["seed"]) - return item - - @staticmethod - def get_collator(): - def collate(samples: list[dict[str, Any]]) -> dict[str, Any]: - batch = { - "video": torch.stack([sample["video"] for sample in samples]), - "input_ids": torch.stack([sample["input_ids"] for sample in samples]), - "attention_mask": torch.stack([sample["attention_mask"] for sample in samples]), - } - if any("seed" in sample for sample in samples): - batch["seed"] = torch.tensor([sample.get("seed", 0) for sample in samples], dtype=torch.long) - return batch - - return collate - - -def build_wan_dataset(config: dict[str, Any]): - processor = WanVideoProcessor(config.get("processor_config", {}) or {}) - dataset = WanVideoDataset(config, processor) +"""Register Wan dataset builder.""" + +from primus.backends.diffusion.data import ( + DatasetConfig, + WanVideoDataProcessor, + WanVideoDataset, +) +from primus.backends.diffusion.utils.log import logger + + +def build_wan_dataset(dataset_config: dict): + processor_config = dataset_config["processor_config"] + processor = WanVideoDataProcessor(processor_config) + processor.build() + logger.info("Built data processor") + + dataset = WanVideoDataset( + processor=processor, + config=DatasetConfig(**dataset_config), + ) + logger.info(f"Built dataset with {len(dataset)} samples") return dataset, processor diff --git a/primus/backends/diffusion/diffusion_adapter.py b/primus/backends/diffusion/diffusion_adapter.py index 202c3914a..7e358db02 100644 --- a/primus/backends/diffusion/diffusion_adapter.py +++ b/primus/backends/diffusion/diffusion_adapter.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Any -from primus.backends.diffusion.argument_builder import WanArgBuilder +from primus.backends.diffusion.argument_builder import DiffusionArgBuilder from primus.core.backend.backend_adapter import BackendAdapter from primus.modules.module_utils import log_rank_0 @@ -42,17 +42,18 @@ def setup_backend_path(self, backend_path=None) -> str: return resolved_str def convert_config(self, params: Any): - builder = WanArgBuilder() + builder = DiffusionArgBuilder() builder.update(params) - wan_args = builder.finalize() + diffusion_args = builder.finalize() # convert_config is also called by the standalone prepare hook, where the # Primus logger may not be initialized yet; guard the informational log. try: - log_rank_0("[Primus:DiffusionAdapter] Converted Primus module params -> Wan args") + model_name = getattr(getattr(diffusion_args, "model", {}), "get", lambda *_: None)("name") + log_rank_0(f"[Primus:DiffusionAdapter] Converted Primus module params -> {model_name} args") except Exception: # Standalone prepare hooks may run before the Primus logger is bound. pass - return wan_args + return diffusion_args def load_trainer_class(self, stage: str = "pretrain"): if stage in ("pretrain", "posttrain", "sft"): diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index ddbb70d49..366042e24 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -48,6 +48,8 @@ def setup(self): if importlib.util.find_spec(package) is None ] video_backend = (dataset_cfg.get("config", {}) or {}).get("video_backend") + if dataset_cfg.get("name") == "flux" and importlib.util.find_spec("datasets") is None: + missing.append("datasets") if video_backend == "imageio" and importlib.util.find_spec("imageio") is None: missing.append("imageio") if video_backend == "decord" and importlib.util.find_spec("decord") is None: diff --git a/primus/backends/diffusion/models/flux/__init__.py b/primus/backends/diffusion/models/flux/__init__.py new file mode 100644 index 000000000..d09c9fd7d --- /dev/null +++ b/primus/backends/diffusion/models/flux/__init__.py @@ -0,0 +1,11 @@ +from primus.backends.diffusion.models.flux.adapter import FluxForTraining +from primus.backends.diffusion.models.flux.model import Flux, FluxParams, flux_1_dev_params +from primus.backends.diffusion.models.flux.train_pipeline import FluxFlowMatchTrainPipeline + +__all__ = [ + "Flux", + "FluxForTraining", + "FluxFlowMatchTrainPipeline", + "FluxParams", + "flux_1_dev_params", +] diff --git a/primus/backends/diffusion/models/flux/adapter.py b/primus/backends/diffusion/models/flux/adapter.py new file mode 100644 index 000000000..764a6ae71 --- /dev/null +++ b/primus/backends/diffusion/models/flux/adapter.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn + +from primus.backends.diffusion.models.interface import GenAIModel +from primus.backends.diffusion.models.flux.train_pipeline import FluxFlowMatchTrainPipeline + + +@dataclass +class FluxConfigShim: + raw: dict + + def save_pretrained(self, save_directory: str): + import json + import os + + os.makedirs(save_directory, exist_ok=True) + with open(os.path.join(save_directory, "flux_config.json"), "w") as f: + json.dump(self.raw, f, indent=2, sort_keys=True) + + def to_dict(self): + return self.raw + + +class FluxForTraining(GenAIModel, nn.Module): + """Thin Primus adapter around the FLUX DiT backbone.""" + + def __init__( + self, + *, + dit: nn.Module, + train_pipeline: FluxFlowMatchTrainPipeline, + model_config: Any, + autoencoder: nn.Module | None = None, + t5_encoder: nn.Module | None = None, + clip_encoder: nn.Module | None = None, + raw_config: dict | None = None, + trainable_modules: str | None = None, + ): + super().__init__() + self.dit = dit + self.autoencoder = autoencoder + self.t5_encoder = t5_encoder + self.clip_encoder = clip_encoder + self.train_pipeline = train_pipeline + self.model_config = model_config + self.trainable_modules = trainable_modules + self.config = FluxConfigShim(raw=raw_config or {}) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def freeze_except(self): + mode = (self.trainable_modules or getattr(self.model_config, "trainable_modules", None) or "dit").lower() + + def freeze(module: nn.Module): + for param in module.parameters(): + param.requires_grad_(False) + + def unfreeze(module: nn.Module): + for param in module.parameters(): + param.requires_grad_(True) + + freeze(self) + if mode in ("dit", "diffusion", "backbone"): + unfreeze(self.dit) + elif mode == "all": + unfreeze(self) + else: + unfreeze(self.dit) + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + if hasattr(self.dit, "gradient_checkpointing"): + self.dit.gradient_checkpointing = True + + def forward(self, *args, **kwargs): + if len(args) >= 1 and isinstance(args[0], dict): + return self.forward_train(args[0], scheduler=kwargs.get("scheduler")) + raise TypeError("FluxForTraining.forward expects (batch_dict, scheduler)") + + def forward_train(self, batch: dict[str, Any], scheduler: Any = None) -> dict[str, torch.Tensor]: + del scheduler + return self.train_pipeline.compute_loss( + dit=self.dit, + autoencoder=self.autoencoder, + t5_encoder=self.t5_encoder, + clip_encoder=self.clip_encoder, + batch=batch, + model_config=self.model_config, + ) + + def forward_inference(self, batch: dict[str, Any], **kwargs): + raise NotImplementedError("FLUX inference pipeline is not wired for Primus training yet") diff --git a/primus/backends/diffusion/models/flux/autoencoder.py b/primus/backends/diffusion/models/flux/autoencoder.py new file mode 100644 index 000000000..dddb454ec --- /dev/null +++ b/primus/backends/diffusion/models/flux/autoencoder.py @@ -0,0 +1,301 @@ +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import torch +from einops import rearrange +from safetensors.torch import load_file as load_sft +from torch import Tensor, nn + + +@dataclass +class AutoEncoderParams: + resolution: int = 256 + in_channels: int = 3 + ch: int = 128 + out_ch: int = 3 + ch_mult: list[int] | None = None + num_res_blocks: int = 2 + z_channels: int = 16 + scale_factor: float = 0.3611 + shift_factor: float = 0.1159 + + def __post_init__(self): + if self.ch_mult is None: + self.ch_mult = [1, 2, 4, 4] + + +def swish(x: Tensor) -> Tensor: + return x * torch.sigmoid(x) + + +class AttnBlock(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + bsz, channels, height, width = q.shape + q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous() + k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous() + v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous() + h_ = nn.functional.scaled_dot_product_attention(q, k, v) + return rearrange(h_, "b 1 (h w) c -> b c h w", h=height, w=width, c=channels, b=bsz) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x: Tensor) -> Tensor: + h = self.conv1(swish(self.norm1(x))) + h = self.conv2(swish(self.norm2(h))) + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class Downsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) + + def forward(self, x: Tensor) -> Tensor: + return self.conv(nn.functional.pad(x, (0, 1, 0, 1), mode="constant", value=0)) + + +class Upsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor) -> Tensor: + x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + return self.conv(x) + + +class Encoder(nn.Module): + def __init__( + self, + resolution: int, + in_channels: int, + ch: int, + ch_mult: list[int], + num_res_blocks: int, + z_channels: int, + ): + super().__init__() + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1) + + in_ch_mult = (1,) + tuple(ch_mult) + self.down = nn.ModuleList() + block_in = self.ch + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + down.downsample = Downsample(block_in) + self.down.append(down) + + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor) -> Tensor: + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions - 1: + hs.append(self.down[i_level].downsample(hs[-1])) + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + return self.conv_out(swish(self.norm_out(h))) + + +class Decoder(nn.Module): + def __init__( + self, + ch: int, + out_ch: int, + ch_mult: list[int], + num_res_blocks: int, + in_channels: int, + resolution: int, + z_channels: int, + ): + super().__init__() + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + block_in = ch * ch_mult[self.num_resolutions - 1] + self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in) + self.up.insert(0, up) + + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1) + + def forward(self, z: Tensor) -> Tensor: + upscale_dtype = next(self.up.parameters()).dtype + h = self.conv_in(z) + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h).to(upscale_dtype) + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + return self.conv_out(swish(self.norm_out(h))) + + +class DiagonalGaussian(nn.Module): + def __init__(self, sample: bool = True, chunk_dim: int = 1): + super().__init__() + self.sample = sample + self.chunk_dim = chunk_dim + + def forward(self, z: Tensor) -> Tensor: + mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim) + if not self.sample: + return mean + std = torch.exp(0.5 * logvar) + return mean + std * torch.randn_like(mean) + + +class AutoEncoder(nn.Module): + def __init__(self, params: AutoEncoderParams, sample_z: bool = True): + super().__init__() + self.params = params + self.encoder = Encoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + ch_mult=params.ch_mult or [1, 2, 4, 4], + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.decoder = Decoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + out_ch=params.out_ch, + ch_mult=params.ch_mult or [1, 2, 4, 4], + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.reg = DiagonalGaussian(sample=sample_z) + self.scale_factor = params.scale_factor + self.shift_factor = params.shift_factor + + def encode(self, x: Tensor) -> Tensor: + z = self.reg(self.encoder(x)) + return self.scale_factor * (z - self.shift_factor) + + def decode(self, z: Tensor) -> Tensor: + z = z / self.scale_factor + self.shift_factor + return self.decoder(z) + + def forward(self, x: Tensor) -> Tensor: + return self.decode(self.encode(x)) + + +def load_autoencoder( + ckpt_path: str | None, + params: AutoEncoderParams | None = None, + *, + dtype: torch.dtype = torch.bfloat16, + sample_z: bool = True, +) -> AutoEncoder: + ae = AutoEncoder(params or AutoEncoderParams(), sample_z=sample_z) + if ckpt_path: + ckpt_path = _resolve_checkpoint_path(ckpt_path, default_filename="ae.safetensors") + state = load_sft(ckpt_path, device="cpu") + missing, unexpected = ae.load_state_dict(state, strict=False) + if missing: + raise ValueError(f"FLUX autoencoder checkpoint missing {len(missing)} keys; first={missing[:3]}") + if unexpected: + raise ValueError(f"FLUX autoencoder checkpoint has {len(unexpected)} unexpected keys; first={unexpected[:3]}") + return ae.to(dtype=dtype) + + +def _resolve_checkpoint_path(path_or_repo_file: str, *, default_filename: str) -> str: + path = Path(path_or_repo_file).expanduser() + if path.exists(): + return str(path) + + parts = path_or_repo_file.split("/") + if len(parts) >= 3: + repo_id = "/".join(parts[:2]) + filename = "/".join(parts[2:]) + elif len(parts) == 2: + repo_id = path_or_repo_file + filename = default_filename + else: + raise FileNotFoundError(f"FLUX checkpoint not found locally and is not a HF repo path: {path_or_repo_file}") + + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id=repo_id, filename=filename) diff --git a/primus/backends/diffusion/models/flux/conditioner.py b/primus/backends/diffusion/models/flux/conditioner.py new file mode 100644 index 000000000..eac5092aa --- /dev/null +++ b/primus/backends/diffusion/models/flux/conditioner.py @@ -0,0 +1,48 @@ +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +import torch +from torch import Tensor, nn +from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer + + +class HFEmbedder(nn.Module): + def __init__(self, version: str, max_length: int, **hf_kwargs): + super().__init__() + self.is_clip = "clip" in version.lower() or version.startswith("openai") + self.max_length = max_length + self.output_key = "pooler_output" if self.is_clip else "last_hidden_state" + + if self.is_clip: + self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length) + self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs) + else: + self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length) + self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs) + + self.hf_module = self.hf_module.eval().requires_grad_(False) + for param in self.hf_module.parameters(): + if not param.is_contiguous(): + param.data = param.data.contiguous() + + @property + def device(self) -> torch.device: + return next(self.hf_module.parameters()).device + + def forward(self, text: list[str]) -> Tensor: + batch_encoding = self.tokenizer( + text, + truncation=True, + max_length=self.max_length, + return_length=False, + return_overflowing_tokens=False, + padding="max_length", + return_tensors="pt", + ) + outputs = self.hf_module( + input_ids=batch_encoding["input_ids"].to(self.device), + attention_mask=None, + output_hidden_states=False, + ) + return outputs[self.output_key] diff --git a/primus/backends/diffusion/models/flux/configuration_flux.py b/primus/backends/diffusion/models/flux/configuration_flux.py new file mode 100644 index 000000000..6c0ff908f --- /dev/null +++ b/primus/backends/diffusion/models/flux/configuration_flux.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +from types import SimpleNamespace + + +class FluxTrainingConfig(SimpleNamespace): + def to_dict(self) -> dict: + return dict(self.__dict__) diff --git a/primus/backends/diffusion/models/flux/layers.py b/primus/backends/diffusion/models/flux/layers.py new file mode 100644 index 000000000..373502198 --- /dev/null +++ b/primus/backends/diffusion/models/flux/layers.py @@ -0,0 +1,203 @@ +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch +from einops import rearrange +from torch import Tensor, nn + +from primus.backends.diffusion.models.flux.math import attention, rope + + +class EmbedND(nn.Module): + def __init__(self, dim: int, theta: int, axes_dim: list[int]): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: Tensor) -> Tensor: + n_axes = ids.shape[-1] + emb = torch.cat([rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], dim=-3) + return emb.unsqueeze(1) + + +def timestep_embedding(t: Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0) -> Tensor: + t = time_factor * t + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + if torch.is_floating_point(t): + embedding = embedding.to(t) + return embedding + + +class MLPEmbedder(nn.Module): + def __init__(self, in_dim: int, hidden_dim: int): + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) + + def forward(self, x: Tensor) -> Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.scale = nn.Parameter(torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + x_dtype = x.dtype + x = x.float() + rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6) + return (x * rrms).to(dtype=x_dtype) * self.scale + + +class QKNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.query_norm = RMSNorm(dim) + self.key_norm = RMSNorm(dim) + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]: + q = self.query_norm(q) + k = self.key_norm(k) + return q.to(v), k.to(v) + + +class SelfAttention(nn.Module): + def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.norm = QKNorm(head_dim) + self.proj = nn.Linear(dim, dim) + + def forward(self, x: Tensor, pe: Tensor) -> Tensor: + qkv = self.qkv(x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + x = attention(q, k, v, pe=pe) + return self.proj(x) + + +@dataclass +class ModulationOut: + shift: Tensor + scale: Tensor + gate: Tensor + + +class Modulation(nn.Module): + def __init__(self, dim: int, double: bool): + super().__init__() + self.is_double = double + self.multiplier = 6 if double else 3 + self.lin = nn.Linear(dim, self.multiplier * dim, bias=True) + + def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]: + out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1) + return ModulationOut(*out[:3]), ModulationOut(*out[3:]) if self.is_double else None + + +class DoubleStreamBlock(nn.Module): + def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False): + super().__init__() + mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.num_heads = num_heads + self.hidden_size = hidden_size + self.img_mod = Modulation(hidden_size, double=True) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_hidden_dim, hidden_size, bias=True), + ) + self.txt_mod = Modulation(hidden_size, double=True) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_hidden_dim, hidden_size, bias=True), + ) + + def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]: + img_mod1, img_mod2 = self.img_mod(vec) + txt_mod1, txt_mod2 = self.txt_mod(vec) + + img_modulated = (1 + img_mod1.scale) * self.img_norm1(img) + img_mod1.shift + img_qkv = self.img_attn.qkv(img_modulated) + img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + img_q, img_k = self.img_attn.norm(img_q, img_k, img_v) + + txt_modulated = (1 + txt_mod1.scale) * self.txt_norm1(txt) + txt_mod1.shift + txt_qkv = self.txt_attn.qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v) + + q = torch.cat((txt_q, img_q), dim=2) + k = torch.cat((txt_k, img_k), dim=2) + v = torch.cat((txt_v, img_v), dim=2) + attn = attention(q, k, v, pe=pe) + txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :] + + img = img + img_mod1.gate * self.img_attn.proj(img_attn) + img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift) + txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn) + txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift) + return img, txt + + +class SingleStreamBlock(nn.Module): + def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0, qk_scale: float | None = None): + super().__init__() + self.hidden_dim = hidden_size + self.num_heads = num_heads + head_dim = hidden_size // num_heads + self.scale = qk_scale or head_dim**-0.5 + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim) + self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size) + self.norm = QKNorm(head_dim) + self.hidden_size = hidden_size + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.mlp_act = nn.GELU(approximate="tanh") + self.modulation = Modulation(hidden_size, double=False) + + def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor: + mod, _ = self.modulation(vec) + x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift + qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + attn = attention(q, k, v, pe=pe) + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + mod.gate * output + + +class LastLayer(nn.Module): + def __init__(self, hidden_size: int, patch_size: int, out_channels: int): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) + + def forward(self, x: Tensor, vec: Tensor) -> Tensor: + shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1) + x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :] + return self.linear(x) diff --git a/primus/backends/diffusion/models/flux/math.py b/primus/backends/diffusion/models/flux/math.py new file mode 100644 index 000000000..c197c4370 --- /dev/null +++ b/primus/backends/diffusion/models/flux/math.py @@ -0,0 +1,31 @@ +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +import torch +from einops import rearrange +from torch import Tensor + + +def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor: + q, k = apply_rope(q, k, pe) + x = torch.nn.functional.scaled_dot_product_attention(q, k, v) + return rearrange(x, "B H L D -> B L (H D)") + + +def rope(pos: Tensor, dim: int, theta: int) -> Tensor: + if dim % 2 != 0: + raise ValueError(f"RoPE dimension must be even, got {dim}") + scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim + omega = 1.0 / (theta**scale) + out = torch.einsum("...n,d->...nd", pos, omega) + out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1) + return rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2).float() + + +def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]: + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) diff --git a/primus/backends/diffusion/models/flux/model.py b/primus/backends/diffusion/models/flux/model.py new file mode 100644 index 000000000..eee47972b --- /dev/null +++ b/primus/backends/diffusion/models/flux/model.py @@ -0,0 +1,151 @@ +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +from dataclasses import asdict, dataclass + +import torch +from torch import Tensor, nn + +from primus.backends.diffusion.models.flux.layers import ( + DoubleStreamBlock, + EmbedND, + LastLayer, + MLPEmbedder, + SingleStreamBlock, + timestep_embedding, +) + + +@dataclass +class FluxParams: + in_channels: int + out_channels: int + vec_in_dim: int + context_in_dim: int + hidden_size: int + mlp_ratio: float + num_heads: int + depth: int + depth_single_blocks: int + axes_dim: list[int] + theta: int + qkv_bias: bool + guidance_embed: bool + + def to_dict(self) -> dict: + return asdict(self) + + +def flux_1_dev_params(**overrides) -> FluxParams: + values = { + "in_channels": 64, + "out_channels": 64, + "vec_in_dim": 768, + "context_in_dim": 4096, + "hidden_size": 3072, + "mlp_ratio": 4.0, + "num_heads": 24, + "depth": 19, + "depth_single_blocks": 38, + "axes_dim": [16, 56, 56], + "theta": 10000, + "qkv_bias": True, + "guidance_embed": True, + } + values.update(overrides) + return FluxParams(**values) + + +class Flux(nn.Module): + """Transformer model for flow matching on packed latent sequences.""" + + def __init__(self, params: FluxParams): + super().__init__() + self.params = params + self.in_channels = params.in_channels + self.out_channels = params.out_channels + if params.hidden_size % params.num_heads != 0: + raise ValueError(f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}") + pe_dim = params.hidden_size // params.num_heads + if sum(params.axes_dim) != pe_dim: + raise ValueError(f"Got axes_dim={params.axes_dim}, expected sum={pe_dim}") + + self.hidden_size = params.hidden_size + self.num_heads = params.num_heads + self.gradient_checkpointing = False + self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim) + self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True) + self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) + self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size) + self.guidance_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity() + self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size) + self.double_blocks = nn.ModuleList( + [ + DoubleStreamBlock( + self.hidden_size, + self.num_heads, + mlp_ratio=params.mlp_ratio, + qkv_bias=params.qkv_bias, + ) + for _ in range(params.depth) + ] + ) + self.single_blocks = nn.ModuleList( + [ + SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio) + for _ in range(params.depth_single_blocks) + ] + ) + self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels) + + def _checkpoint_double(self, block: nn.Module, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor): + import torch.utils.checkpoint as checkpoint_utils + + return checkpoint_utils.checkpoint(block, img, txt, vec, pe, use_reentrant=False) + + def _checkpoint_single(self, block: nn.Module, img: Tensor, vec: Tensor, pe: Tensor): + import torch.utils.checkpoint as checkpoint_utils + + return checkpoint_utils.checkpoint(block, img, vec, pe, use_reentrant=False) + + def forward( + self, + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + timesteps: Tensor, + y: Tensor, + guidance: Tensor | None = None, + ) -> Tensor: + if img.ndim != 3 or txt.ndim != 3: + raise ValueError("Input img and txt tensors must have 3 dimensions.") + + img = self.img_in(img) + vec = self.time_in(timestep_embedding(timesteps, 256)) + if self.params.guidance_embed: + if guidance is None: + raise ValueError("FLUX guidance-distilled models require a guidance tensor.") + vec = vec + self.guidance_in(timestep_embedding(guidance, 256)) + vec = vec + self.vector_in(y) + txt = self.txt_in(txt) + + ids = torch.cat((txt_ids, img_ids), dim=1) + pe = self.pe_embedder(ids) + + use_checkpoint = self.training and self.gradient_checkpointing + for block in self.double_blocks: + if use_checkpoint: + img, txt = self._checkpoint_double(block, img, txt, vec, pe) + else: + img, txt = block(img=img, txt=txt, vec=vec, pe=pe) + + img = torch.cat((txt, img), 1) + for block in self.single_blocks: + if use_checkpoint: + img = self._checkpoint_single(block, img, vec, pe) + else: + img = block(img, vec=vec, pe=pe) + img = img[:, txt.shape[1] :, ...] + return self.final_layer(img, vec) diff --git a/primus/backends/diffusion/models/flux/train_pipeline.py b/primus/backends/diffusion/models/flux/train_pipeline.py new file mode 100644 index 000000000..cef4bd87d --- /dev/null +++ b/primus/backends/diffusion/models/flux/train_pipeline.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F + +from primus.backends.diffusion.models.flux.utils import ( + create_position_encoding_for_latents, + generate_latent_from_mean_logvar, + pack_latents, +) + + +@dataclass +class FluxFlowMatchTrainPipelineConfig: + autoencoder_scale_factor: float = 0.3611 + autoencoder_shift_factor: float = 0.1159 + guidance: float = 1.0 + + +class FluxFlowMatchTrainPipeline: + """Flow-matching loss for FLUX using precomputed or online encodings.""" + + def __init__(self, cfg: FluxFlowMatchTrainPipelineConfig | None = None): + self.cfg = cfg or FluxFlowMatchTrainPipelineConfig() + + @staticmethod + def _require_module(module: torch.nn.Module | None, name: str) -> torch.nn.Module: + if module is None: + raise ValueError(f"FLUX raw image-text training requires `{name}` to be configured.") + return module + + def _prepare_precomputed( + self, + *, + batch: dict[str, Any], + device: torch.device, + dtype: torch.dtype, + model_config: Any, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + t5_encodings = batch["t5_encodings"].to(device=device, dtype=dtype, non_blocking=True) + clip_encodings = batch["clip_encodings"].to(device=device, dtype=dtype, non_blocking=True) + mean = batch["mean"].to(device=device, dtype=dtype, non_blocking=True) + logvar = batch["logvar"].to(device=device, dtype=dtype, non_blocking=True) + + latents = generate_latent_from_mean_logvar(mean, logvar) + scale = float(getattr(model_config, "autoencoder_scale_factor", self.cfg.autoencoder_scale_factor)) + shift = float(getattr(model_config, "autoencoder_shift_factor", self.cfg.autoencoder_shift_factor)) + labels = (latents - shift) * scale + return labels, t5_encodings, clip_encodings + + def _prepare_raw( + self, + *, + batch: dict[str, Any], + autoencoder: torch.nn.Module | None, + t5_encoder: torch.nn.Module | None, + clip_encoder: torch.nn.Module | None, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ae = self._require_module(autoencoder, "model.config.encoder.autoencoder") + t5 = self._require_module(t5_encoder, "model.config.encoder.t5_encoder") + clip = self._require_module(clip_encoder, "model.config.encoder.clip_encoder") + image = batch["image"].to(device=device, dtype=dtype, non_blocking=True) + prompts = batch.get("prompts") + if not isinstance(prompts, list): + raise ValueError("FLUX raw batch requires `prompts` as list[str].") + + ae.eval() + t5.eval() + clip.eval() + with torch.no_grad(): + labels = ae.encode(image).to(device=device, dtype=dtype) + t5_encodings = t5(prompts).to(device=device, dtype=dtype) + clip_encodings = clip(prompts).to(device=device, dtype=dtype) + return labels, t5_encodings, clip_encodings + + def compute_loss( + self, + *, + dit: torch.nn.Module, + batch: dict[str, Any], + model_config: Any, + autoencoder: torch.nn.Module | None = None, + t5_encoder: torch.nn.Module | None = None, + clip_encoder: torch.nn.Module | None = None, + ) -> dict[str, torch.Tensor]: + if batch.get("sp_group") is not None: + raise ValueError("FLUX diffusion training currently requires `parallelism.sp_size: 1`.") + + device = next(dit.parameters()).device + dtype = next(dit.parameters()).dtype + if "image" in batch: + labels, t5_encodings, clip_encodings = self._prepare_raw( + batch=batch, + autoencoder=autoencoder, + t5_encoder=t5_encoder, + clip_encoder=clip_encoder, + device=device, + dtype=dtype, + ) + else: + required = ("t5_encodings", "clip_encodings", "mean", "logvar") + missing = [key for key in required if key not in batch] + if missing: + raise ValueError(f"FLUX precomputed batch is missing required keys: {missing}") + labels, t5_encodings, clip_encodings = self._prepare_precomputed( + batch=batch, + device=device, + dtype=dtype, + model_config=model_config, + ) + + bsz = labels.shape[0] + noise = torch.randn_like(labels) + timesteps = torch.rand((bsz,), device=device, dtype=dtype) + sigmas = timesteps.view(-1, 1, 1, 1) + noisy_latents = (1 - sigmas) * labels + sigmas * noise + target = noise - labels + + _, _, latent_height, latent_width = noisy_latents.shape + img_ids = create_position_encoding_for_latents( + bsz, + latent_height, + latent_width, + position_dim=3, + device=device, + dtype=dtype, + ) + txt_ids = torch.zeros(bsz, t5_encodings.shape[1], 3, device=device, dtype=dtype) + noisy_latents = pack_latents(noisy_latents) + target = pack_latents(target) + + guidance_value = float(getattr(model_config, "guidance", self.cfg.guidance)) + guidance = torch.full((bsz,), guidance_value, device=device, dtype=dtype) + pred = dit( + img=noisy_latents, + img_ids=img_ids, + txt=t5_encodings, + txt_ids=txt_ids, + y=clip_encodings, + timesteps=timesteps, + guidance=guidance, + ) + loss = F.mse_loss(pred.float(), target.float(), reduction="sum") / target.numel() + return {"loss": loss, "log_metrics": {"mse": loss.detach()}} diff --git a/primus/backends/diffusion/models/flux/utils.py b/primus/backends/diffusion/models/flux/utils.py new file mode 100644 index 000000000..7e4ddf167 --- /dev/null +++ b/primus/backends/diffusion/models/flux/utils.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import torch +from torch import Tensor + +PATCH_HEIGHT = 2 +PATCH_WIDTH = 2 +LATENT_CHANNELS = 16 +IMAGE_LATENT_SIZE_RATIO = 8 + + +def generate_latent_from_mean_logvar(mean: Tensor, logvar: Tensor) -> Tensor: + return mean + torch.exp(0.5 * logvar) * torch.randn_like(mean) + + +def create_position_encoding_for_latents( + bsz: int, + latent_height: int, + latent_width: int, + position_dim: int = 3, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, +) -> Tensor: + height = latent_height // PATCH_HEIGHT + width = latent_width // PATCH_WIDTH + position_encoding = torch.zeros(height, width, position_dim, device=device, dtype=dtype) + position_encoding[:, :, 1] = torch.arange(height, device=device, dtype=dtype).unsqueeze(1) + position_encoding[:, :, 2] = torch.arange(width, device=device, dtype=dtype).unsqueeze(0) + return position_encoding.view(1, height * width, position_dim).repeat(bsz, 1, 1) + + +def pack_latents(x: Tensor) -> Tensor: + bsz, channels, latent_height, latent_width = x.shape + if latent_height % PATCH_HEIGHT != 0 or latent_width % PATCH_WIDTH != 0: + raise ValueError( + "FLUX latents must have height and width divisible by 2, " + f"got shape={tuple(x.shape)}" + ) + height = latent_height // PATCH_HEIGHT + width = latent_width // PATCH_WIDTH + x = x.unfold(2, PATCH_HEIGHT, PATCH_HEIGHT).unfold(3, PATCH_WIDTH, PATCH_WIDTH) + x = x.permute(0, 2, 3, 1, 4, 5) + return x.reshape(bsz, height * width, channels * PATCH_HEIGHT * PATCH_WIDTH) diff --git a/primus/backends/diffusion/models/registrations/flux.py b/primus/backends/diffusion/models/registrations/flux.py new file mode 100644 index 000000000..145f6f599 --- /dev/null +++ b/primus/backends/diffusion/models/registrations/flux.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import glob +import os +from typing import Any + +import torch +from safetensors.torch import load_file as safe_load_file + +from primus.backends.diffusion.models.flux.adapter import FluxForTraining +from primus.backends.diffusion.models.flux.autoencoder import AutoEncoderParams, load_autoencoder +from primus.backends.diffusion.models.flux.configuration_flux import FluxTrainingConfig +from primus.backends.diffusion.models.flux.conditioner import HFEmbedder +from primus.backends.diffusion.models.flux.model import Flux, flux_1_dev_params +from primus.backends.diffusion.models.flux.train_pipeline import ( + FluxFlowMatchTrainPipeline, + FluxFlowMatchTrainPipelineConfig, +) +from primus.backends.diffusion.utils.log import logger +from primus.backends.diffusion.utils.train_utils import count_parameters + + +def _strip_known_prefixes(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + prefixes = ("module.", "dit.", "model.") + out: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + stripped = key + changed = True + while changed: + changed = False + for prefix in prefixes: + if stripped.startswith(prefix): + stripped = stripped[len(prefix) :] + changed = True + out[stripped] = value + return out + + +def _load_state_dict(path: str) -> dict[str, torch.Tensor]: + if path.endswith(".safetensors"): + return dict(safe_load_file(path)) + obj = torch.load(path, map_location="cpu") + if isinstance(obj, dict) and "model" in obj and isinstance(obj["model"], dict): + obj = obj["model"] + if not isinstance(obj, dict): + raise ValueError(f"Unsupported checkpoint format at {path}") + return obj + + +def _candidate_weight_files(path: str) -> list[str]: + if os.path.isfile(path): + return [path] + if not os.path.exists(path): + resolved = _resolve_hf_checkpoint(path, default_filename="flux1-dev.safetensors") + if resolved: + return [resolved] + candidates: list[str] = [] + for fname in ("flux1-dev.safetensors", "dit_model.safetensors", "model.safetensors"): + candidate = os.path.join(path, fname) + if os.path.exists(candidate): + candidates.append(candidate) + if not candidates: + candidates = sorted(glob.glob(os.path.join(path, "*.safetensors"))) + if not candidates: + candidates = sorted(glob.glob(os.path.join(path, "*.bin"))) + return candidates + + +def _resolve_hf_checkpoint(path_or_repo_file: str, *, default_filename: str) -> str | None: + parts = path_or_repo_file.split("/") + if len(parts) < 2: + return None + if len(parts) >= 3: + repo_id = "/".join(parts[:2]) + filename = "/".join(parts[2:]) + else: + repo_id = path_or_repo_file + filename = default_filename + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id=repo_id, filename=filename) + + +def _load_flux_weights(dit: torch.nn.Module, pretrained_path: str) -> None: + candidates = _candidate_weight_files(pretrained_path) + if not candidates: + raise FileNotFoundError(f"No FLUX DiT weights found under {pretrained_path}") + + merged: dict[str, torch.Tensor] = {} + for ckpt in candidates: + merged.update(_strip_known_prefixes(_load_state_dict(ckpt))) + + result = dit.load_state_dict(merged, strict=False) + logger.info( + "Loaded FLUX DiT weights. " + f"files={len(candidates)} missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}" + ) + + +def _build_flux_dit(params) -> Flux: + local_rank = os.environ.get("LOCAL_RANK") + use_cuda = local_rank is not None and torch.cuda.is_available() + device = torch.device(f"cuda:{local_rank}") if use_cuda else torch.device("cpu") + old_dtype = torch.get_default_dtype() + try: + if use_cuda: + torch.set_default_dtype(torch.bfloat16) + with torch.device(device): + dit = Flux(params) + finally: + torch.set_default_dtype(old_dtype) + return dit + + +def build_flux_model(model_config: dict[str, Any]): + """ + YAML shape: + model_config: + load_from_pretrained_path: /path/to/flux1-dev.safetensors + config: + model_variant: flux-dev + trainable_modules: dit + guidance: 1.0 + params: {... optional FluxParams overrides ...} + """ + cfg_dict: dict[str, Any] = dict(model_config.get("config", {}) or {}) + variant = cfg_dict.get("model_variant", "flux-dev") + if variant != "flux-dev": + raise ValueError(f"Only FLUX.1-dev is currently supported, got model_variant={variant!r}") + + params_overrides = dict(cfg_dict.get("params", {}) or {}) + params = flux_1_dev_params(**params_overrides) + dit = _build_flux_dit(params) + + pretrained_path = model_config.get("load_from_pretrained_path") or model_config.get("pretrained_path") + if pretrained_path: + logger.info(f"Loading FLUX DiT weights from {pretrained_path}") + _load_flux_weights(dit, pretrained_path) + + encoder_cfg = dict(model_config.get("encoder", {}) or cfg_dict.get("encoder", {}) or {}) + dtype = torch.bfloat16 + t5_encoder = None + clip_encoder = None + autoencoder = None + if encoder_cfg.get("t5_encoder"): + t5_encoder = HFEmbedder( + str(encoder_cfg["t5_encoder"]), + max_length=int(encoder_cfg.get("max_t5_length", 256)), + torch_dtype=dtype, + ) + if encoder_cfg.get("clip_encoder"): + clip_encoder = HFEmbedder( + str(encoder_cfg["clip_encoder"]), + max_length=int(encoder_cfg.get("max_clip_length", 77)), + torch_dtype=dtype, + ) + if encoder_cfg.get("autoencoder"): + ae_params = AutoEncoderParams( + resolution=int(encoder_cfg.get("resolution", 256)), + scale_factor=float(cfg_dict.get("autoencoder_scale_factor", 0.3611)), + shift_factor=float(cfg_dict.get("autoencoder_shift_factor", 0.1159)), + ) + autoencoder = load_autoencoder( + str(encoder_cfg["autoencoder"]), + ae_params, + dtype=dtype, + sample_z=bool(encoder_cfg.get("sample_z", True)), + ) + + training_cfg = FluxTrainingConfig( + model_variant=variant, + trainable_modules=cfg_dict.get("trainable_modules", "dit"), + guidance=float(cfg_dict.get("guidance", 1.0)), + autoencoder_scale_factor=float(cfg_dict.get("autoencoder_scale_factor", 0.3611)), + autoencoder_shift_factor=float(cfg_dict.get("autoencoder_shift_factor", 0.1159)), + ) + pipeline = FluxFlowMatchTrainPipeline( + FluxFlowMatchTrainPipelineConfig( + autoencoder_scale_factor=training_cfg.autoencoder_scale_factor, + autoencoder_shift_factor=training_cfg.autoencoder_shift_factor, + guidance=training_cfg.guidance, + ) + ) + model = FluxForTraining( + dit=dit, + train_pipeline=pipeline, + model_config=training_cfg, + autoencoder=autoencoder, + t5_encoder=t5_encoder, + clip_encoder=clip_encoder, + raw_config={ + "model_config": model_config, + "flux_params": params.to_dict(), + }, + trainable_modules=training_cfg.trainable_modules, + ) + total_params, trainable_params = count_parameters(model) + logger.info(f"Built FLUX model: total={total_params:,} trainable={trainable_params:,}") + return model diff --git a/primus/backends/diffusion/registry.py b/primus/backends/diffusion/registry.py index b0a728827..72bce2776 100644 --- a/primus/backends/diffusion/registry.py +++ b/primus/backends/diffusion/registry.py @@ -23,6 +23,18 @@ def _build_wan_dataset(dataset_config: dict): return build_wan_dataset(dataset_config) +def _build_flux_model(model_config: dict): + from primus.backends.diffusion.models.registrations.flux import build_flux_model + + return build_flux_model(model_config) + + +def _build_flux_dataset(dataset_config: dict): + from primus.backends.diffusion.data.registrations.flux import build_flux_dataset + + return build_flux_dataset(dataset_config) + + def _build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict): from primus.backends.diffusion.trainers.fsdp2 import build_fsdp2_trainer @@ -35,9 +47,11 @@ def _build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict): MODEL_BUILDERS: Dict[str, Callable[[dict], Any]] = { + "flux": _build_flux_model, "wan": _build_wan_model, } DATASET_BUILDERS: Dict[str, Callable[[dict], Tuple[Any, Any]]] = { + "flux": _build_flux_dataset, "wan": _build_wan_dataset, } TRAINER_BUILDERS: Dict[str, Callable[..., Any]] = { diff --git a/primus/backends/diffusion/trainers/base.py b/primus/backends/diffusion/trainers/base.py index c5d5227fd..89e3b764e 100644 --- a/primus/backends/diffusion/trainers/base.py +++ b/primus/backends/diffusion/trainers/base.py @@ -385,6 +385,9 @@ def _infer_batch_size_from_sequences(self, value) -> int | None: return None def _infer_local_batch_size(self, batch) -> int: + if isinstance(batch, (list, tuple)): + return len(batch) + tensor_batch_size = self._infer_batch_size_from_tensors(batch) if tensor_batch_size is not None: return tensor_batch_size @@ -470,6 +473,11 @@ def train(self): self.model.train() self.optimizer.zero_grad(set_to_none=True) torch.cuda.reset_peak_memory_stats() + debug_steps = os.getenv("DIFFUSION_DEBUG_STEPS", "0") == "1" + + def debug_phase(message: str) -> None: + if debug_steps: + logger.info(f"debug_step rank={self.rank} step={self.global_step} {message}") start_time = time.time() last_log_time = start_time @@ -485,21 +493,27 @@ def train(self): for batch_idx, batch in enumerate(self.dataloader): is_update_step = ((batch_idx + 1) % max(1, self.grad_accum_steps)) == 0 local_samples_in_update += self._infer_local_batch_size(batch) + debug_phase(f"batch_idx={batch_idx} before_compute_loss update={is_update_step}") with self._grad_sync_context(is_update_step): raw_loss = self.compute_loss(batch) + debug_phase("after_compute_loss") update_loss_sum += raw_loss.detach().float().item() update_loss_count += 1 loss = raw_loss / max(1, self.grad_accum_steps) loss.backward() + debug_phase("after_backward") if is_update_step: loss_val = update_loss_sum / max(1, update_loss_count) update_loss_sum = 0.0 update_loss_count = 0 + debug_phase("before_clip_grad_norm") grad_norm = self._clip_grad_norm() + debug_phase("before_optimizer_step") self.optimizer.step() + debug_phase("after_optimizer_step") self.lr_scheduler.step() self.optimizer.zero_grad(set_to_none=True) self.global_step += 1 diff --git a/primus/backends/diffusion/trainers/fsdp2.py b/primus/backends/diffusion/trainers/fsdp2.py index 1fa79b177..e3bd72a70 100644 --- a/primus/backends/diffusion/trainers/fsdp2.py +++ b/primus/backends/diffusion/trainers/fsdp2.py @@ -370,6 +370,10 @@ def _save_dit(self, save_path: str) -> None: def save_model(self): """Save final model using the configured strategy.""" + if self.save_strategy in ("none", "skip", "disabled"): + if self.rank == 0: + logger.info(f"Skipping final checkpoint save (strategy={self.save_strategy})") + return if self.save_strategy == "dit_only": save_path = os.path.join(self.output_dir, "dit_model.safetensors") self._save_dit(save_path) diff --git a/primus/configs/models/diffusion/flux.1_dev_t2i.yaml b/primus/configs/models/diffusion/flux.1_dev_t2i.yaml new file mode 100644 index 000000000..64d5287bd --- /dev/null +++ b/primus/configs/models/diffusion/flux.1_dev_t2i.yaml @@ -0,0 +1,18 @@ +model: + name: flux + config: + # Optional: set PRETRAINED_PATH to initialize from existing FLUX DiT weights. + load_from_pretrained_path: ${PRETRAINED_PATH:} + config: + model_variant: flux-dev + trainable_modules: dit + guidance: ${FLUX_GUIDANCE:1.0} + autoencoder_scale_factor: 0.3611 + autoencoder_shift_factor: 0.1159 + # Optional for precomputed data, required for raw image-text data. + encoder: + t5_encoder: ${T5_ENCODER:} + clip_encoder: ${CLIP_ENCODER:} + autoencoder: ${VAE_CHECKPOINT:} + max_t5_length: ${MAX_T5_LENGTH:256} + max_clip_length: 77 diff --git a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py index b23e6a5ae..a59415f59 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py +++ b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py @@ -71,6 +71,23 @@ def _require_path(path: str | None, description: str, *, kind: str = "any") -> N _log(f"{description}: {resolved}") +def _require_optional_path(path: str | None, description: str, *, kind: str = "any") -> None: + if _is_placeholder(path): + _log(f"{description}: not configured; skipping optional initialization") + return + _require_path(path, description, kind=kind) + + +def _validate_local_or_hf_id(value: str | None, description: str) -> None: + if _is_placeholder(value): + _fail(f"{description} is not configured: {value!r}") + path = Path(str(value)).expanduser() + if path.exists(): + _log(f"{description}: {path}") + else: + _log(f"{description}: {value} (assuming Hugging Face model id)") + + def validate_diffusion_config(config_path: Path, module_name: str | None = None) -> None: cfg = load_primus_config(config_path) selected_module = _select_module_name(cfg, module_name) @@ -87,16 +104,50 @@ def validate_diffusion_config(config_path: Path, module_name: str | None = None) dataset_cfg = dataset.get("config", {}) processor_cfg = dataset_cfg.get("processor_config", {}) - encoder_cfg = model.get("config", {}).get("encoder", {}) or model.get("encoder", {}) - - _require_path(dataset_cfg.get("dataset_path"), "dataset metadata", kind="file") - _require_path(dataset_cfg.get("data_folder"), "dataset media folder", kind="dir") - _require_path(processor_cfg.get("text_tokenizer"), "text tokenizer", kind="dir") - model_cfg = model.get("config", {}) - _require_path(model_cfg.get("load_from_pretrained_path"), "DiT initialization checkpoint") - _require_path(encoder_cfg.get("t5_encoder"), "text encoder checkpoint", kind="file") - _require_path(encoder_cfg.get("autoencoder"), "VAE checkpoint", kind="file") + model_name = model.get("name") + + if model_name == "wan": + encoder_cfg = model_cfg.get("encoder", {}) or model.get("encoder", {}) + _require_path(dataset_cfg.get("dataset_path"), "dataset metadata", kind="file") + _require_path(dataset_cfg.get("data_folder"), "dataset media folder", kind="dir") + _require_path(processor_cfg.get("text_tokenizer"), "text tokenizer", kind="dir") + _require_path(model_cfg.get("load_from_pretrained_path"), "DiT initialization checkpoint") + _require_path(encoder_cfg.get("t5_encoder"), "text encoder checkpoint", kind="file") + _require_path(encoder_cfg.get("autoencoder"), "VAE checkpoint", kind="file") + elif model_name == "flux": + dataset_type = str(dataset_cfg.get("dataset_type", "precomputed")).lower() + if dataset_type == "raw": + dataset_name = dataset_cfg.get("dataset") + dataset_format = str(dataset_cfg.get("dataset_format", "webdataset")).lower() + if dataset_name == "cc12m-test" and _is_placeholder(dataset_cfg.get("dataset_path")): + default_path = "/mnt/shared/zirui/code/torchtitan-main/tests/assets/cc12m_test" + _require_path(default_path, "FLUX raw image-text dataset cc12m-test", kind="dir") + elif dataset_name == "cc12m-wds" and _is_placeholder(dataset_cfg.get("dataset_path")): + _log("FLUX raw image-text dataset: pixparse/cc12m-wds (Hugging Face dataset)") + else: + kind = "file" if dataset_format == "jsonl" else "dir" + _require_path(dataset_cfg.get("dataset_path"), "FLUX raw image-text dataset", kind=kind) + encoder_cfg = model_cfg.get("encoder", {}) or model.get("encoder", {}) + _validate_local_or_hf_id(encoder_cfg.get("t5_encoder"), "FLUX T5 encoder") + _validate_local_or_hf_id(encoder_cfg.get("clip_encoder"), "FLUX CLIP encoder") + _validate_local_or_hf_id(encoder_cfg.get("autoencoder"), "FLUX autoencoder checkpoint") + elif dataset_type == "precomputed": + _require_path(dataset_cfg.get("dataset_path"), "FLUX precomputed dataset", kind="dir") + prompt_dropout_prob = float(processor_cfg.get("prompt_dropout_prob", 0.0) or 0.0) + if prompt_dropout_prob > 0.0: + empty_dir = processor_cfg.get("empty_encodings_path") + _require_path(empty_dir, "FLUX empty encodings directory", kind="dir") + _require_path(str(Path(empty_dir) / "t5_empty.npy"), "FLUX empty T5 encoding", kind="file") + _require_path(str(Path(empty_dir) / "clip_empty.npy"), "FLUX empty CLIP encoding", kind="file") + else: + _fail(f"unsupported FLUX dataset_type: {dataset_type!r}") + _require_optional_path( + model_cfg.get("load_from_pretrained_path"), + "FLUX DiT initialization checkpoint", + ) + else: + _fail(f"unsupported diffusion model for prepare hook: {model_name!r}") _log( "validated " diff --git a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt index 38fabbeef..1e4547a8d 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt +++ b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt @@ -1,6 +1,9 @@ einops +datasets +huggingface_hub loguru safetensors +sentencepiece numpy pillow packaging diff --git a/tests/unit_tests/backends/diffusion/test_flux_backend.py b/tests/unit_tests/backends/diffusion/test_flux_backend.py new file mode 100644 index 000000000..f680b5eda --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_flux_backend.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from primus.backends.diffusion.argument_builder import DiffusionArgBuilder +from primus.backends.diffusion.data.flux_precomputed import ( + FluxPrecomputedProcessor, + FluxRawImageTextDataset, + FluxRawImageTextProcessor, +) +from primus.backends.diffusion.models.registrations.flux import build_flux_model + + +def _finalize(params: dict): + builder = DiffusionArgBuilder() + builder.update(params) + return builder.finalize() + + +def test_flux_argument_builder_selects_flux_defaults(): + args = _finalize( + { + "model": {"name": "flux", "config": {"config": {"model_variant": "flux-dev"}}}, + "training": {"steps": 7, "local_batch_size": 3}, + "data": { + "dataset_path": "/tmp/precomputed", + "dataset_type": "precomputed", + "empty_encodings_path": "/tmp/empty", + "prompt_dropout_prob": 0.25, + }, + } + ) + + assert args.model["name"] == "flux" + assert args.dataset["name"] == "flux" + assert args.dataset["config"]["dataset_path"] == "/tmp/precomputed" + assert args.dataset["config"]["processor_config"]["empty_encodings_path"] == "/tmp/empty" + assert args.dataset["config"]["processor_config"]["prompt_dropout_prob"] == 0.25 + assert args.trainer["args"]["max_steps"] == 7 + assert args.trainer["args"]["per_device_train_batch_size"] == 3 + assert args.trainer["args"]["fsdp_transformer_layer_cls_to_wrap"] == "DoubleStreamBlock,SingleStreamBlock" + + +def test_flux_argument_builder_maps_raw_dataset_type(): + args = _finalize( + { + "model": {"name": "flux", "config": {"config": {"model_variant": "flux-dev"}}}, + "data": { + "dataset_type": "raw", + "dataset_format": "webdataset", + "dataset_path": "/tmp/cc12m_test", + "prompt_dropout_prob": 0.1, + "img_size": 128, + }, + } + ) + + assert args.dataset["config"]["dataset_type"] == "raw" + assert args.dataset["config"]["dataset_format"] == "webdataset" + assert args.dataset["config"]["processor_config"]["img_size"] == 128 + + +def test_flux_raw_dataset_name_defaults(): + path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-test", None, "webdataset") + assert path.endswith("torchtitan-main/tests/assets/cc12m_test") + assert fmt == "webdataset" + + path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-wds", None, "webdataset") + assert path == "pixparse/cc12m-wds" + assert fmt == "hf_repo" + + +def test_flux_argument_builder_rejects_sequence_parallelism(): + with pytest.raises(ValueError, match="sp_size"): + _finalize( + { + "model": {"name": "flux", "config": {"config": {"model_variant": "flux-dev"}}}, + "parallelism": {"sp_size": 2}, + } + ) + + +def test_flux_precomputed_processor_stacks_and_drops_empty_encodings(tmp_path): + empty_dir = tmp_path / "empty" + empty_dir.mkdir() + np.save(empty_dir / "t5_empty.npy", np.zeros((1, 3, 8), dtype=np.float32)) + np.save(empty_dir / "clip_empty.npy", np.zeros((1, 4), dtype=np.float32)) + + processor = FluxPrecomputedProcessor( + { + "prompt_dropout_prob": 1.0, + "empty_encodings_path": str(empty_dir), + } + ) + batch = [ + { + "t5_encodings": torch.ones(3, 8), + "clip_encodings": torch.ones(4), + "mean": torch.zeros(1, 2, 2), + "logvar": torch.zeros(1, 2, 2), + }, + { + "t5_encodings": torch.ones(3, 8), + "clip_encodings": torch.ones(4), + "mean": torch.zeros(1, 2, 2), + "logvar": torch.zeros(1, 2, 2), + }, + ] + + out = processor.prepare_batch(batch=batch, device=torch.device("cpu"), dtype=torch.float32) + + assert out["t5_encodings"].shape == (2, 3, 8) + assert out["clip_encodings"].shape == (2, 4) + assert torch.count_nonzero(out["t5_encodings"]) == 0 + assert torch.count_nonzero(out["clip_encodings"]) == 0 + + +def test_flux_raw_processor_prepares_images_and_prompts(): + from PIL import Image + + processor = FluxRawImageTextProcessor({"img_size": 8, "prompt_dropout_prob": 1.0, "skip_low_resolution": False}) + image = Image.fromarray(np.full((6, 10, 3), 127, dtype=np.uint8)) + + out = processor.prepare_batch( + batch=[{"image": image, "prompt": "a test image"}], + device=torch.device("cpu"), + dtype=torch.float32, + ) + + assert out["image"].shape == (1, 3, 8, 8) + assert out["prompts"] == [""] + + +def test_tiny_flux_model_computes_precomputed_loss(): + model = build_flux_model( + { + "config": { + "model_variant": "flux-dev", + "guidance": 1.0, + "params": { + "in_channels": 4, + "out_channels": 4, + "vec_in_dim": 4, + "context_in_dim": 8, + "hidden_size": 12, + "num_heads": 2, + "depth": 1, + "depth_single_blocks": 1, + "axes_dim": [2, 2, 2], + }, + } + } + ) + batch = { + "t5_encodings": torch.randn(2, 3, 8), + "clip_encodings": torch.randn(2, 4), + "mean": torch.randn(2, 1, 2, 2), + "logvar": torch.zeros(2, 1, 2, 2), + } + + outputs = model.forward_train(batch) + + assert outputs["loss"].ndim == 0 + assert torch.isfinite(outputs["loss"]) + + +def test_tiny_flux_model_computes_raw_loss_with_dummy_encoders(): + class DummyAutoencoder(torch.nn.Module): + def encode(self, image): + return image[:, :1, :2, :2] + + class DummyT5(torch.nn.Module): + def forward(self, prompts): + return torch.zeros(len(prompts), 3, 8) + + class DummyClip(torch.nn.Module): + def forward(self, prompts): + return torch.zeros(len(prompts), 4) + + model = build_flux_model( + { + "config": { + "model_variant": "flux-dev", + "guidance": 1.0, + "params": { + "in_channels": 4, + "out_channels": 4, + "vec_in_dim": 4, + "context_in_dim": 8, + "hidden_size": 12, + "num_heads": 2, + "depth": 1, + "depth_single_blocks": 1, + "axes_dim": [2, 2, 2], + }, + } + } + ) + model.autoencoder = DummyAutoencoder() + model.t5_encoder = DummyT5() + model.clip_encoder = DummyClip() + + outputs = model.forward_train( + { + "image": torch.randn(2, 3, 4, 4), + "prompts": ["cat", "dog"], + } + ) + + assert outputs["loss"].ndim == 0 + assert torch.isfinite(outputs["loss"]) From d0dbbcbc57dff24256e195199e2e3e29fba50e96 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 23 Jun 2026 10:13:23 +0000 Subject: [PATCH 22/30] fix(diffusion): make FLUX examples portable Remove local asset assumptions from FLUX configs/docs, preserve AMD and BFL attribution, and tighten checkpoint/dataset validation. Co-authored-by: Cursor --- .../MI355X/flux.1_dev_t2i-pretrain.yaml | 4 +-- primus/backends/diffusion/README.md | 19 +++++----- .../diffusion/data/flux_precomputed.py | 13 ++++--- .../diffusion/data/registrations/flux.py | 6 ++++ .../diffusion/models/flux/__init__.py | 6 ++++ .../backends/diffusion/models/flux/adapter.py | 6 ++++ .../diffusion/models/flux/autoencoder.py | 10 ++++++ .../diffusion/models/flux/conditioner.py | 6 ++++ .../models/flux/configuration_flux.py | 6 ++++ .../backends/diffusion/models/flux/layers.py | 6 ++++ primus/backends/diffusion/models/flux/math.py | 6 ++++ .../backends/diffusion/models/flux/model.py | 6 ++++ .../diffusion/models/flux/train_pipeline.py | 6 ++++ .../backends/diffusion/models/flux/utils.py | 6 ++++ .../diffusion/models/registrations/flux.py | 10 ++++++ .../hooks/train/pretrain/diffusion/prepare.py | 35 +++++++++++++++---- .../backends/diffusion/test_flux_backend.py | 7 ++-- 17 files changed, 134 insertions(+), 24 deletions(-) diff --git a/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml index 561c05acc..9235979a9 100644 --- a/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml +++ b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml @@ -32,8 +32,8 @@ modules: run_name: flux.1_dev_t2i-pretrain data: - dataset_path: ${DATASET_PATH:/mnt/vast/zirui/data/cc12m_preprocessed} - empty_encodings_path: ${EMPTY_ENCODINGS_PATH:/mnt/vast/zirui/data/empty_encodings} + dataset_path: ${DATASET_PATH:} + empty_encodings_path: ${EMPTY_ENCODINGS_PATH:} prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1} img_size: ${IMG_SIZE:256} diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md index a8a786c89..fd81a0816 100644 --- a/primus/backends/diffusion/README.md +++ b/primus/backends/diffusion/README.md @@ -91,16 +91,15 @@ model: `autoencoder` accepts either a local safetensors path or `repo_id/filename`; for example `black-forest-labs/FLUX.1-dev/ae.safetensors`. -The precomputed MLPerf data on this machine is: +The shipped precomputed example expects these paths through environment +variables: ```text -/mnt/vast/zirui/data/ - cc12m_preprocessed/ # train, precomputed text + VAE mean/logvar - coco_preprocessed/ # validation-style precomputed data - empty_encodings/ # t5_empty.npy, clip_empty.npy +DATASET_PATH=/path/to/cc12m_preprocessed +EMPTY_ENCODINGS_PATH=/path/to/empty_encodings ``` -The shipped FLUX precomputed example uses these paths by default: +The example config is: ```text examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml @@ -125,12 +124,12 @@ data: dataset: cc12m-wds # maps to Hugging Face dataset pixparse/cc12m-wds ``` -Do not use full `cc12m-wds` as the default smoke test; it is large. The raw FLUX -example defaults to `cc12m-test`, which maps to TorchTitan's small test tar: +Do not use full `cc12m-wds` as the default smoke test; it is large. For +`cc12m-test`, pass a local webdataset directory explicitly: ```text -/mnt/shared/zirui/code/torchtitan-main/tests/assets/cc12m_test/ - cc12m-train-0000.tar +DATASET=cc12m-test +DATASET_PATH=/path/to/cc12m_test ``` The raw FLUX example is: diff --git a/primus/backends/diffusion/data/flux_precomputed.py b/primus/backends/diffusion/data/flux_precomputed.py index e896f22e9..adc80e6fa 100644 --- a/primus/backends/diffusion/data/flux_precomputed.py +++ b/primus/backends/diffusion/data/flux_precomputed.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + from __future__ import annotations import io @@ -124,10 +130,9 @@ def _resolve_dataset( dataset_format: str, ) -> tuple[str, str]: if dataset_name == "cc12m-test": - return ( - dataset_path or "/mnt/shared/zirui/code/torchtitan-main/tests/assets/cc12m_test", - "webdataset", - ) + if not dataset_path: + raise ValueError("FLUX raw dataset `cc12m-test` requires `dataset_path` to be set.") + return dataset_path, "webdataset" if dataset_name == "cc12m-wds": return dataset_path or "pixparse/cc12m-wds", "hf_repo" if not dataset_path: diff --git a/primus/backends/diffusion/data/registrations/flux.py b/primus/backends/diffusion/data/registrations/flux.py index a03e25a1a..234868970 100644 --- a/primus/backends/diffusion/data/registrations/flux.py +++ b/primus/backends/diffusion/data/registrations/flux.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + from __future__ import annotations from primus.backends.diffusion.data.flux_precomputed import ( diff --git a/primus/backends/diffusion/models/flux/__init__.py b/primus/backends/diffusion/models/flux/__init__.py index d09c9fd7d..efd786e05 100644 --- a/primus/backends/diffusion/models/flux/__init__.py +++ b/primus/backends/diffusion/models/flux/__init__.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + from primus.backends.diffusion.models.flux.adapter import FluxForTraining from primus.backends.diffusion.models.flux.model import Flux, FluxParams, flux_1_dev_params from primus.backends.diffusion.models.flux.train_pipeline import FluxFlowMatchTrainPipeline diff --git a/primus/backends/diffusion/models/flux/adapter.py b/primus/backends/diffusion/models/flux/adapter.py index 764a6ae71..50bf5c188 100644 --- a/primus/backends/diffusion/models/flux/adapter.py +++ b/primus/backends/diffusion/models/flux/adapter.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + from __future__ import annotations from dataclasses import dataclass diff --git a/primus/backends/diffusion/models/flux/autoencoder.py b/primus/backends/diffusion/models/flux/autoencoder.py index dddb454ec..703ffc781 100644 --- a/primus/backends/diffusion/models/flux/autoencoder.py +++ b/primus/backends/diffusion/models/flux/autoencoder.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# # Adapted from Black Forest Labs FLUX official implementation. from __future__ import annotations @@ -285,8 +291,12 @@ def _resolve_checkpoint_path(path_or_repo_file: str, *, default_filename: str) - path = Path(path_or_repo_file).expanduser() if path.exists(): return str(path) + if path_or_repo_file.startswith(("/", "./", "../", "~")): + raise FileNotFoundError(f"FLUX checkpoint path not found: {path}") parts = path_or_repo_file.split("/") + if len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")): + raise FileNotFoundError(f"FLUX checkpoint path not found: {path}") if len(parts) >= 3: repo_id = "/".join(parts[:2]) filename = "/".join(parts[2:]) diff --git a/primus/backends/diffusion/models/flux/conditioner.py b/primus/backends/diffusion/models/flux/conditioner.py index eac5092aa..d63779c7c 100644 --- a/primus/backends/diffusion/models/flux/conditioner.py +++ b/primus/backends/diffusion/models/flux/conditioner.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# # Adapted from Black Forest Labs FLUX official implementation. from __future__ import annotations diff --git a/primus/backends/diffusion/models/flux/configuration_flux.py b/primus/backends/diffusion/models/flux/configuration_flux.py index 6c0ff908f..6808fcac7 100644 --- a/primus/backends/diffusion/models/flux/configuration_flux.py +++ b/primus/backends/diffusion/models/flux/configuration_flux.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + from __future__ import annotations from types import SimpleNamespace diff --git a/primus/backends/diffusion/models/flux/layers.py b/primus/backends/diffusion/models/flux/layers.py index 373502198..cf057db90 100644 --- a/primus/backends/diffusion/models/flux/layers.py +++ b/primus/backends/diffusion/models/flux/layers.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# # Adapted from Black Forest Labs FLUX official implementation. from __future__ import annotations diff --git a/primus/backends/diffusion/models/flux/math.py b/primus/backends/diffusion/models/flux/math.py index c197c4370..0b28c51e9 100644 --- a/primus/backends/diffusion/models/flux/math.py +++ b/primus/backends/diffusion/models/flux/math.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# # Adapted from Black Forest Labs FLUX official implementation. from __future__ import annotations diff --git a/primus/backends/diffusion/models/flux/model.py b/primus/backends/diffusion/models/flux/model.py index eee47972b..38460efb8 100644 --- a/primus/backends/diffusion/models/flux/model.py +++ b/primus/backends/diffusion/models/flux/model.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# # Adapted from Black Forest Labs FLUX official implementation. from __future__ import annotations diff --git a/primus/backends/diffusion/models/flux/train_pipeline.py b/primus/backends/diffusion/models/flux/train_pipeline.py index cef4bd87d..d803b1629 100644 --- a/primus/backends/diffusion/models/flux/train_pipeline.py +++ b/primus/backends/diffusion/models/flux/train_pipeline.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + from __future__ import annotations from dataclasses import dataclass diff --git a/primus/backends/diffusion/models/flux/utils.py b/primus/backends/diffusion/models/flux/utils.py index 7e4ddf167..13b07c434 100644 --- a/primus/backends/diffusion/models/flux/utils.py +++ b/primus/backends/diffusion/models/flux/utils.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + from __future__ import annotations import torch diff --git a/primus/backends/diffusion/models/registrations/flux.py b/primus/backends/diffusion/models/registrations/flux.py index 145f6f599..b37dc57e9 100644 --- a/primus/backends/diffusion/models/registrations/flux.py +++ b/primus/backends/diffusion/models/registrations/flux.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + from __future__ import annotations import glob @@ -67,7 +73,11 @@ def _candidate_weight_files(path: str) -> list[str]: def _resolve_hf_checkpoint(path_or_repo_file: str, *, default_filename: str) -> str | None: + if path_or_repo_file.startswith(("/", "./", "../", "~")): + return None parts = path_or_repo_file.split("/") + if len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")): + return None if len(parts) < 2: return None if len(parts) >= 3: diff --git a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py index a59415f59..089f945eb 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py +++ b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py @@ -78,14 +78,33 @@ def _require_optional_path(path: str | None, description: str, *, kind: str = "a _require_path(path, description, kind=kind) +def _looks_like_local_path(value: str) -> bool: + if value.startswith(("/", "./", "../", "~")): + return True + parts = value.split("/") + return len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")) + + +def _looks_like_hf_reference(value: str) -> bool: + if _looks_like_local_path(value): + return False + parts = value.split("/") + return len(parts) >= 2 and all(part for part in parts[:2]) + + def _validate_local_or_hf_id(value: str | None, description: str) -> None: if _is_placeholder(value): _fail(f"{description} is not configured: {value!r}") - path = Path(str(value)).expanduser() + value = str(value) + path = Path(value).expanduser() if path.exists(): _log(f"{description}: {path}") + elif _looks_like_local_path(value): + _fail(f"{description} path not found: {path}") + elif _looks_like_hf_reference(value): + _log(f"{description}: {value} (assuming Hugging Face reference)") else: - _log(f"{description}: {value} (assuming Hugging Face model id)") + _fail(f"{description} is neither an existing path nor a Hugging Face reference: {value!r}") def validate_diffusion_config(config_path: Path, module_name: str | None = None) -> None: @@ -121,10 +140,14 @@ def validate_diffusion_config(config_path: Path, module_name: str | None = None) dataset_name = dataset_cfg.get("dataset") dataset_format = str(dataset_cfg.get("dataset_format", "webdataset")).lower() if dataset_name == "cc12m-test" and _is_placeholder(dataset_cfg.get("dataset_path")): - default_path = "/mnt/shared/zirui/code/torchtitan-main/tests/assets/cc12m_test" - _require_path(default_path, "FLUX raw image-text dataset cc12m-test", kind="dir") - elif dataset_name == "cc12m-wds" and _is_placeholder(dataset_cfg.get("dataset_path")): - _log("FLUX raw image-text dataset: pixparse/cc12m-wds (Hugging Face dataset)") + _fail("FLUX raw dataset `cc12m-test` requires `dataset_path` to point at local webdataset tars") + elif dataset_name == "cc12m-wds": + if _is_placeholder(dataset_cfg.get("dataset_path")): + _log("FLUX raw image-text dataset: pixparse/cc12m-wds (Hugging Face dataset)") + else: + _validate_local_or_hf_id(dataset_cfg.get("dataset_path"), "FLUX raw Hugging Face dataset") + elif dataset_format == "hf_repo": + _validate_local_or_hf_id(dataset_cfg.get("dataset_path"), "FLUX raw Hugging Face dataset") else: kind = "file" if dataset_format == "jsonl" else "dir" _require_path(dataset_cfg.get("dataset_path"), "FLUX raw image-text dataset", kind=kind) diff --git a/tests/unit_tests/backends/diffusion/test_flux_backend.py b/tests/unit_tests/backends/diffusion/test_flux_backend.py index f680b5eda..0710505e3 100644 --- a/tests/unit_tests/backends/diffusion/test_flux_backend.py +++ b/tests/unit_tests/backends/diffusion/test_flux_backend.py @@ -63,8 +63,11 @@ def test_flux_argument_builder_maps_raw_dataset_type(): def test_flux_raw_dataset_name_defaults(): - path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-test", None, "webdataset") - assert path.endswith("torchtitan-main/tests/assets/cc12m_test") + with pytest.raises(ValueError, match="cc12m-test.*dataset_path"): + FluxRawImageTextDataset._resolve_dataset("cc12m-test", None, "webdataset") + + path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-test", "/tmp/cc12m_test", "webdataset") + assert path == "/tmp/cc12m_test" assert fmt == "webdataset" path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-wds", None, "webdataset") From 626c075050f8d749fc2bbb0b2d4c47c3e1107750 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 23 Jun 2026 11:01:48 +0000 Subject: [PATCH 23/30] style(diffusion): apply pre-commit formatting Normalize diffusion backend formatting after running pre-commit across the repository. Co-authored-by: Cursor --- primus/backends/diffusion/argument_builder.py | 4 +- primus/backends/diffusion/data/collator.py | 4 +- primus/backends/diffusion/data/config.py | 2 +- .../diffusion/data/flux_precomputed.py | 10 ++-- .../diffusion/data/processing_wanvideo.py | 53 +++++++++---------- primus/backends/diffusion/data/processor.py | 6 ++- .../diffusion/models/flux/__init__.py | 10 +++- .../backends/diffusion/models/flux/adapter.py | 8 ++- .../diffusion/models/flux/autoencoder.py | 8 ++- .../backends/diffusion/models/flux/layers.py | 4 +- .../backends/diffusion/models/flux/model.py | 8 ++- .../backends/diffusion/models/flux/utils.py | 3 +- .../diffusion/models/registrations/flux.py | 7 ++- .../hooks/train/pretrain/diffusion/prepare.py | 8 ++- .../backends/diffusion/test_flux_backend.py | 4 +- 15 files changed, 88 insertions(+), 51 deletions(-) diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py index 95b2fcf25..6d7739b9c 100644 --- a/primus/backends/diffusion/argument_builder.py +++ b/primus/backends/diffusion/argument_builder.py @@ -155,7 +155,9 @@ def update(self, params: Any) -> None: elif isinstance(params, dict): self._params = copy.deepcopy(params) else: - raise TypeError(f"DiffusionArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}") + raise TypeError( + f"DiffusionArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}" + ) def finalize(self) -> SimpleNamespace: params = copy.deepcopy(self._params) diff --git a/primus/backends/diffusion/data/collator.py b/primus/backends/diffusion/data/collator.py index 9dc07d814..c1fdb9e09 100644 --- a/primus/backends/diffusion/data/collator.py +++ b/primus/backends/diffusion/data/collator.py @@ -19,7 +19,9 @@ class VisionCollator: def pad_sequence(self, input_ids, batch_first, padding_value): if self.processor.tokenizer.padding_side == "left": input_ids = [torch.flip(_input_ids, [0]) for _input_ids in input_ids] - input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=batch_first, padding_value=padding_value) + input_ids = torch.nn.utils.rnn.pad_sequence( + input_ids, batch_first=batch_first, padding_value=padding_value + ) if self.processor.tokenizer.padding_side == "left": input_ids = torch.flip(input_ids, [1]) return input_ids diff --git a/primus/backends/diffusion/data/config.py b/primus/backends/diffusion/data/config.py index 229145758..e425149ba 100644 --- a/primus/backends/diffusion/data/config.py +++ b/primus/backends/diffusion/data/config.py @@ -86,4 +86,4 @@ def validate_video_backend_migration(cls, v): "Migration guide: If you were using torchvision, 'decord' provides " "similar functionality with better performance." ) - return v \ No newline at end of file + return v diff --git a/primus/backends/diffusion/data/flux_precomputed.py b/primus/backends/diffusion/data/flux_precomputed.py index adc80e6fa..7b3e2e510 100644 --- a/primus/backends/diffusion/data/flux_precomputed.py +++ b/primus/backends/diffusion/data/flux_precomputed.py @@ -10,8 +10,8 @@ import json import math import os -from pathlib import Path from dataclasses import dataclass +from pathlib import Path from typing import Any, Sequence import numpy as np @@ -239,7 +239,9 @@ def _collate_raw(batch: Any) -> dict[str, torch.Tensor]: collated[key] = torch.stack([_to_tensor(sample[key]) for sample in batch], dim=0) return collated - def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) -> dict[str, torch.Tensor]: + def prepare_batch( + self, *, batch: Any, device: torch.device, dtype: torch.dtype + ) -> dict[str, torch.Tensor]: tensors = self._collate_raw(batch) for key, value in tensors.items(): tensors[key] = value.to(device=device, dtype=dtype, non_blocking=True) @@ -295,7 +297,9 @@ def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) if isinstance(batch, dict): batch = [batch] if not isinstance(batch, Sequence) or not batch: - raise ValueError(f"FLUX raw prepare_batch expected non-empty sequence, got {type(batch).__name__}") + raise ValueError( + f"FLUX raw prepare_batch expected non-empty sequence, got {type(batch).__name__}" + ) images: list[torch.Tensor] = [] prompts: list[str] = [] diff --git a/primus/backends/diffusion/data/processing_wanvideo.py b/primus/backends/diffusion/data/processing_wanvideo.py index 64ffddb5f..446411d98 100644 --- a/primus/backends/diffusion/data/processing_wanvideo.py +++ b/primus/backends/diffusion/data/processing_wanvideo.py @@ -4,10 +4,9 @@ # See LICENSE for license information. ############################################################################### -import warnings -from typing import Any, Dict, List, Optional, Union, Tuple - import os +from typing import Any, Dict, List, Optional, Tuple, Union + import numpy as np import torch from PIL import Image @@ -82,14 +81,14 @@ def resize( width, height = image.size target_width = size["width"] target_height = size["height"] - + scale = max(target_width / width, target_height / height) new_width = round(width * scale) new_height = round(height * scale) - - from torchvision.transforms import functional as F + from torchvision.transforms import InterpolationMode - + from torchvision.transforms import functional as F + # DiffSynth uses torchvision.transforms.resize with BILINEAR # We must use F.resize to match exactly (antialias behavior etc) image = F.resize(image, (new_height, new_width), interpolation=InterpolationMode.BILINEAR) @@ -277,26 +276,26 @@ def preprocess( # Temporal Handling (Interpolate or Truncate) if num_frames is not None: - current_frames = processed_images.shape[1] - if current_frames > num_frames: - logger.info(f"Truncating video frames from {current_frames} to {num_frames}") - processed_images = processed_images[:, :num_frames, ...] - elif current_frames < num_frames: - logger.info(f"Interpolating video frames from {current_frames} to {num_frames}") - # Interpolate requires (B, C, T, H, W) or (B, C, H, W) - we have (B, T, H, W, C) - # Permute to (B, C, T, H, W) for interpolate - vid_tensor = torch.from_numpy(processed_images).permute(0, 4, 1, 2, 3) - - # Interpolate - vid_tensor = torch.nn.functional.interpolate( - vid_tensor, - size=(num_frames, vid_tensor.shape[3], vid_tensor.shape[4]), - mode='trilinear', - align_corners=False - ) - - # Permute back to (B, T, H, W, C) and convert to numpy - processed_images = vid_tensor.permute(0, 2, 3, 4, 1).numpy() + current_frames = processed_images.shape[1] + if current_frames > num_frames: + logger.info(f"Truncating video frames from {current_frames} to {num_frames}") + processed_images = processed_images[:, :num_frames, ...] + elif current_frames < num_frames: + logger.info(f"Interpolating video frames from {current_frames} to {num_frames}") + # Interpolate requires (B, C, T, H, W) or (B, C, H, W) - we have (B, T, H, W, C) + # Permute to (B, C, T, H, W) for interpolate + vid_tensor = torch.from_numpy(processed_images).permute(0, 4, 1, 2, 3) + + # Interpolate + vid_tensor = torch.nn.functional.interpolate( + vid_tensor, + size=(num_frames, vid_tensor.shape[3], vid_tensor.shape[4]), + mode="trilinear", + align_corners=False, + ) + + # Permute back to (B, T, H, W, C) and convert to numpy + processed_images = vid_tensor.permute(0, 2, 3, 4, 1).numpy() # Convert to tensor if requested if return_tensors == "pt": diff --git a/primus/backends/diffusion/data/processor.py b/primus/backends/diffusion/data/processor.py index 9b940d118..6434d24a6 100644 --- a/primus/backends/diffusion/data/processor.py +++ b/primus/backends/diffusion/data/processor.py @@ -47,7 +47,9 @@ def build(self): from transformers import AutoTokenizer - from primus.backends.diffusion.data.processing_wanvideo import WanVideoProcessor as WanVideoModelProcessor + from primus.backends.diffusion.data.processing_wanvideo import ( + WanVideoProcessor as WanVideoModelProcessor, + ) wanvideo_kwargs = self.config.get("extra_kwargs", {}) max_text_length = self.config.get("max_text_length") @@ -177,4 +179,4 @@ def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) prompts, frames_list, num_frames = self._normalize_raw_batch(batch) text_inputs = self._tokenize_prompts(prompts) pixel_values = self._preprocess_videos(frames_list, num_frames=num_frames) - return self._assemble_model_batch(pixel_values=pixel_values, text_inputs=text_inputs) \ No newline at end of file + return self._assemble_model_batch(pixel_values=pixel_values, text_inputs=text_inputs) diff --git a/primus/backends/diffusion/models/flux/__init__.py b/primus/backends/diffusion/models/flux/__init__.py index efd786e05..7f4567d8a 100644 --- a/primus/backends/diffusion/models/flux/__init__.py +++ b/primus/backends/diffusion/models/flux/__init__.py @@ -5,8 +5,14 @@ ############################################################################### from primus.backends.diffusion.models.flux.adapter import FluxForTraining -from primus.backends.diffusion.models.flux.model import Flux, FluxParams, flux_1_dev_params -from primus.backends.diffusion.models.flux.train_pipeline import FluxFlowMatchTrainPipeline +from primus.backends.diffusion.models.flux.model import ( + Flux, + FluxParams, + flux_1_dev_params, +) +from primus.backends.diffusion.models.flux.train_pipeline import ( + FluxFlowMatchTrainPipeline, +) __all__ = [ "Flux", diff --git a/primus/backends/diffusion/models/flux/adapter.py b/primus/backends/diffusion/models/flux/adapter.py index 50bf5c188..a19d28a24 100644 --- a/primus/backends/diffusion/models/flux/adapter.py +++ b/primus/backends/diffusion/models/flux/adapter.py @@ -12,8 +12,10 @@ import torch import torch.nn as nn +from primus.backends.diffusion.models.flux.train_pipeline import ( + FluxFlowMatchTrainPipeline, +) from primus.backends.diffusion.models.interface import GenAIModel -from primus.backends.diffusion.models.flux.train_pipeline import FluxFlowMatchTrainPipeline @dataclass @@ -66,7 +68,9 @@ def dtype(self): return next(self.parameters()).dtype def freeze_except(self): - mode = (self.trainable_modules or getattr(self.model_config, "trainable_modules", None) or "dit").lower() + mode = ( + self.trainable_modules or getattr(self.model_config, "trainable_modules", None) or "dit" + ).lower() def freeze(module: nn.Module): for param in module.parameters(): diff --git a/primus/backends/diffusion/models/flux/autoencoder.py b/primus/backends/diffusion/models/flux/autoencoder.py index 703ffc781..9ff6f3e56 100644 --- a/primus/backends/diffusion/models/flux/autoencoder.py +++ b/primus/backends/diffusion/models/flux/autoencoder.py @@ -283,7 +283,9 @@ def load_autoencoder( if missing: raise ValueError(f"FLUX autoencoder checkpoint missing {len(missing)} keys; first={missing[:3]}") if unexpected: - raise ValueError(f"FLUX autoencoder checkpoint has {len(unexpected)} unexpected keys; first={unexpected[:3]}") + raise ValueError( + f"FLUX autoencoder checkpoint has {len(unexpected)} unexpected keys; first={unexpected[:3]}" + ) return ae.to(dtype=dtype) @@ -304,7 +306,9 @@ def _resolve_checkpoint_path(path_or_repo_file: str, *, default_filename: str) - repo_id = path_or_repo_file filename = default_filename else: - raise FileNotFoundError(f"FLUX checkpoint not found locally and is not a HF repo path: {path_or_repo_file}") + raise FileNotFoundError( + f"FLUX checkpoint not found locally and is not a HF repo path: {path_or_repo_file}" + ) from huggingface_hub import hf_hub_download diff --git a/primus/backends/diffusion/models/flux/layers.py b/primus/backends/diffusion/models/flux/layers.py index cf057db90..475be057b 100644 --- a/primus/backends/diffusion/models/flux/layers.py +++ b/primus/backends/diffusion/models/flux/layers.py @@ -170,7 +170,9 @@ def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Te class SingleStreamBlock(nn.Module): - def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0, qk_scale: float | None = None): + def __init__( + self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0, qk_scale: float | None = None + ): super().__init__() self.hidden_dim = hidden_size self.num_heads = num_heads diff --git a/primus/backends/diffusion/models/flux/model.py b/primus/backends/diffusion/models/flux/model.py index 38460efb8..e5ebcce5a 100644 --- a/primus/backends/diffusion/models/flux/model.py +++ b/primus/backends/diffusion/models/flux/model.py @@ -72,7 +72,9 @@ def __init__(self, params: FluxParams): self.in_channels = params.in_channels self.out_channels = params.out_channels if params.hidden_size % params.num_heads != 0: - raise ValueError(f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}") + raise ValueError( + f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}" + ) pe_dim = params.hidden_size // params.num_heads if sum(params.axes_dim) != pe_dim: raise ValueError(f"Got axes_dim={params.axes_dim}, expected sum={pe_dim}") @@ -84,7 +86,9 @@ def __init__(self, params: FluxParams): self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True) self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size) - self.guidance_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity() + self.guidance_in = ( + MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity() + ) self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size) self.double_blocks = nn.ModuleList( [ diff --git a/primus/backends/diffusion/models/flux/utils.py b/primus/backends/diffusion/models/flux/utils.py index 13b07c434..00d0a2acf 100644 --- a/primus/backends/diffusion/models/flux/utils.py +++ b/primus/backends/diffusion/models/flux/utils.py @@ -40,8 +40,7 @@ def pack_latents(x: Tensor) -> Tensor: bsz, channels, latent_height, latent_width = x.shape if latent_height % PATCH_HEIGHT != 0 or latent_width % PATCH_WIDTH != 0: raise ValueError( - "FLUX latents must have height and width divisible by 2, " - f"got shape={tuple(x.shape)}" + "FLUX latents must have height and width divisible by 2, " f"got shape={tuple(x.shape)}" ) height = latent_height // PATCH_HEIGHT width = latent_width // PATCH_WIDTH diff --git a/primus/backends/diffusion/models/registrations/flux.py b/primus/backends/diffusion/models/registrations/flux.py index b37dc57e9..89e8f3e23 100644 --- a/primus/backends/diffusion/models/registrations/flux.py +++ b/primus/backends/diffusion/models/registrations/flux.py @@ -14,9 +14,12 @@ from safetensors.torch import load_file as safe_load_file from primus.backends.diffusion.models.flux.adapter import FluxForTraining -from primus.backends.diffusion.models.flux.autoencoder import AutoEncoderParams, load_autoencoder -from primus.backends.diffusion.models.flux.configuration_flux import FluxTrainingConfig +from primus.backends.diffusion.models.flux.autoencoder import ( + AutoEncoderParams, + load_autoencoder, +) from primus.backends.diffusion.models.flux.conditioner import HFEmbedder +from primus.backends.diffusion.models.flux.configuration_flux import FluxTrainingConfig from primus.backends.diffusion.models.flux.model import Flux, flux_1_dev_params from primus.backends.diffusion.models.flux.train_pipeline import ( FluxFlowMatchTrainPipeline, diff --git a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py index 089f945eb..92bf5e678 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py +++ b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py @@ -140,7 +140,9 @@ def validate_diffusion_config(config_path: Path, module_name: str | None = None) dataset_name = dataset_cfg.get("dataset") dataset_format = str(dataset_cfg.get("dataset_format", "webdataset")).lower() if dataset_name == "cc12m-test" and _is_placeholder(dataset_cfg.get("dataset_path")): - _fail("FLUX raw dataset `cc12m-test` requires `dataset_path` to point at local webdataset tars") + _fail( + "FLUX raw dataset `cc12m-test` requires `dataset_path` to point at local webdataset tars" + ) elif dataset_name == "cc12m-wds": if _is_placeholder(dataset_cfg.get("dataset_path")): _log("FLUX raw image-text dataset: pixparse/cc12m-wds (Hugging Face dataset)") @@ -162,7 +164,9 @@ def validate_diffusion_config(config_path: Path, module_name: str | None = None) empty_dir = processor_cfg.get("empty_encodings_path") _require_path(empty_dir, "FLUX empty encodings directory", kind="dir") _require_path(str(Path(empty_dir) / "t5_empty.npy"), "FLUX empty T5 encoding", kind="file") - _require_path(str(Path(empty_dir) / "clip_empty.npy"), "FLUX empty CLIP encoding", kind="file") + _require_path( + str(Path(empty_dir) / "clip_empty.npy"), "FLUX empty CLIP encoding", kind="file" + ) else: _fail(f"unsupported FLUX dataset_type: {dataset_type!r}") _require_optional_path( diff --git a/tests/unit_tests/backends/diffusion/test_flux_backend.py b/tests/unit_tests/backends/diffusion/test_flux_backend.py index 0710505e3..35db038ff 100644 --- a/tests/unit_tests/backends/diffusion/test_flux_backend.py +++ b/tests/unit_tests/backends/diffusion/test_flux_backend.py @@ -123,7 +123,9 @@ def test_flux_precomputed_processor_stacks_and_drops_empty_encodings(tmp_path): def test_flux_raw_processor_prepares_images_and_prompts(): from PIL import Image - processor = FluxRawImageTextProcessor({"img_size": 8, "prompt_dropout_prob": 1.0, "skip_low_resolution": False}) + processor = FluxRawImageTextProcessor( + {"img_size": 8, "prompt_dropout_prob": 1.0, "skip_low_resolution": False} + ) image = Image.fromarray(np.full((6, 10, 3), 127, dtype=np.uint8)) out = processor.prepare_batch( From 2b1016687fceb21a2a54e2e4dbb4137efbf355a5 Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 8 Jul 2026 00:06:18 +0000 Subject: [PATCH 24/30] fix(diffusion): use AITER attention for FLUX Route FLUX attention through the diffusion backend dispatcher and align the default runtime backend with ROCm Wan defaults while preserving SDPA coverage in tests. Co-authored-by: Cursor --- .../MI355X/flux.1_dev_t2i-pretrain.yaml | 2 +- .../MI355X/flux.1_dev_t2i-raw-pretrain.yaml | 2 +- primus/backends/diffusion/README.md | 4 ++- primus/backends/diffusion/argument_builder.py | 2 +- primus/backends/diffusion/models/flux/math.py | 11 +++++-- .../backends/diffusion/test_flux_backend.py | 29 +++++++++++++++++++ 6 files changed, 44 insertions(+), 6 deletions(-) diff --git a/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml index 9235979a9..5310f2dc3 100644 --- a/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml +++ b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml @@ -46,7 +46,7 @@ modules: weight_decay: 0.01 runtime: - attention_backend: ${ATTENTION_BACKEND:sdpa} + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false} fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true} report_to: none diff --git a/examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml index 11b59ac4a..e3689a489 100644 --- a/examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml +++ b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml @@ -57,7 +57,7 @@ modules: weight_decay: 0.01 runtime: - attention_backend: ${ATTENTION_BACKEND:sdpa} + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false} fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true} report_to: none diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md index fd81a0816..b84b552f6 100644 --- a/primus/backends/diffusion/README.md +++ b/primus/backends/diffusion/README.md @@ -281,7 +281,9 @@ torchrun \ Useful runtime knobs: -- `trainer.args.attention_backend`: defaults to `flash_attn_aiter` for Wan training; use `sdpa` as the portable fallback or baseline. +- `trainer.args.attention_backend`: defaults to `flash_attn_aiter` for Wan + training on ROCm. Use `sdpa` as the portable fallback when AITER flash + attention is unavailable. - `trainer.args.sp_size`: Ulysses sequence parallel size. It must divide the model attention head count; for example Wan2.1-1.3B supports `sp_size=4` but not `sp_size=8`. diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py index 6d7739b9c..ca352a003 100644 --- a/primus/backends/diffusion/argument_builder.py +++ b/primus/backends/diffusion/argument_builder.py @@ -112,7 +112,7 @@ class DiffusionArgBuilder: "per_device_eval_batch_size": 1, "gradient_accumulation_steps": 1, "gradient_checkpointing": False, - "attention_backend": "sdpa", + "attention_backend": "flash_attn_aiter", "learning_rate": 2.0e-4, "lr_scheduler_type": "constant", "warmup_steps": 0, diff --git a/primus/backends/diffusion/models/flux/math.py b/primus/backends/diffusion/models/flux/math.py index 0b28c51e9..829d2f395 100644 --- a/primus/backends/diffusion/models/flux/math.py +++ b/primus/backends/diffusion/models/flux/math.py @@ -12,11 +12,18 @@ from einops import rearrange from torch import Tensor +from primus.backends.diffusion.attention import attention as backend_attention + def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor: q, k = apply_rope(q, k, pe) - x = torch.nn.functional.scaled_dot_product_attention(q, k, v) - return rearrange(x, "B H L D -> B L (H D)") + x = backend_attention( + q=rearrange(q, "B H L D -> B L H D"), + k=rearrange(k, "B H L D -> B L H D"), + v=rearrange(v, "B H L D -> B L H D"), + dtype=q.dtype if q.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16, + ) + return rearrange(x, "B L H D -> B L (H D)") def rope(pos: Tensor, dim: int, theta: int) -> Tensor: diff --git a/tests/unit_tests/backends/diffusion/test_flux_backend.py b/tests/unit_tests/backends/diffusion/test_flux_backend.py index 35db038ff..d4f087fea 100644 --- a/tests/unit_tests/backends/diffusion/test_flux_backend.py +++ b/tests/unit_tests/backends/diffusion/test_flux_backend.py @@ -5,11 +5,18 @@ import torch from primus.backends.diffusion.argument_builder import DiffusionArgBuilder +from primus.backends.diffusion.attention import ( + get_attention_backend, + set_attention_backend, +) from primus.backends.diffusion.data.flux_precomputed import ( FluxPrecomputedProcessor, FluxRawImageTextDataset, FluxRawImageTextProcessor, ) +from primus.backends.diffusion.models.flux.math import apply_rope +from primus.backends.diffusion.models.flux.math import attention as flux_attention +from primus.backends.diffusion.models.flux.math import rope from primus.backends.diffusion.models.registrations.flux import build_flux_model @@ -40,6 +47,7 @@ def test_flux_argument_builder_selects_flux_defaults(): assert args.dataset["config"]["processor_config"]["prompt_dropout_prob"] == 0.25 assert args.trainer["args"]["max_steps"] == 7 assert args.trainer["args"]["per_device_train_batch_size"] == 3 + assert args.trainer["args"]["attention_backend"] == "flash_attn_aiter" assert args.trainer["args"]["fsdp_transformer_layer_cls_to_wrap"] == "DoubleStreamBlock,SingleStreamBlock" @@ -85,6 +93,27 @@ def test_flux_argument_builder_rejects_sequence_parallelism(): ) +def test_flux_attention_dispatch_matches_sdpa_layout(): + previous_backend = get_attention_backend() + set_attention_backend("sdpa") + try: + torch.manual_seed(7) + q = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16) + k = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16) + v = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16) + pos = torch.arange(5, dtype=torch.float32).repeat(2, 1) + pe = rope(pos, dim=4, theta=10000).unsqueeze(1) + + actual = flux_attention(q, k, v, pe=pe) + q_rope, k_rope = apply_rope(q, k, pe) + expected = torch.nn.functional.scaled_dot_product_attention(q_rope, k_rope, v) + expected = expected.transpose(1, 2).flatten(2) + + torch.testing.assert_close(actual, expected) + finally: + set_attention_backend(previous_backend) + + def test_flux_precomputed_processor_stacks_and_drops_empty_encodings(tmp_path): empty_dir = tmp_path / "empty" empty_dir.mkdir() From 377d09eafa67ebeedb0b09561caaee22e318f781 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 6 Jul 2026 11:20:52 +0000 Subject: [PATCH 25/30] fix(tests): add __init__.py to test dirs to prevent module name conflicts The CI was failing with a pytest collection error: ERROR collecting tests/unit_tests/tools/test_utils.py import file mismatch: imported module 'test_utils' has this __file__ attribute: tests/unit_tests/core/patches/test_utils.py HINT: remove __pycache__ / .pyc files and/or use a unique basename for your test file modules Root cause: two files named test_utils.py existed in different sub-directories of tests/unit_tests/ (core/patches/ and tools/), but their parent directories lacked __init__.py files. Without __init__.py, pytest imports every test file as a top-level module using just its basename, so both were imported as 'test_utils', causing the conflict. Fix: add empty __init__.py to all test sub-directories that were missing one, making each directory a proper Python package. pytest then imports tests with fully-qualified package paths (e.g. tests.unit_tests.core.patches.test_utils vs tests.unit_tests.tools.test_utils), eliminating the name collision. --- tests/unit_tests/backends/__init__.py | 0 tests/unit_tests/backends/diffusion/__init__.py | 0 tests/unit_tests/backends/megatron/__init__.py | 0 tests/unit_tests/backends/torchtitan/__init__.py | 0 tests/unit_tests/cli/__init__.py | 0 tests/unit_tests/core/__init__.py | 0 tests/unit_tests/core/backend/__init__.py | 0 tests/unit_tests/core/launcher/__init__.py | 0 tests/unit_tests/core/patches/__init__.py | 0 tests/unit_tests/core/runtime/__init__.py | 0 tests/unit_tests/core/trainer/__init__.py | 0 tests/unit_tests/core/utils/__init__.py | 0 12 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit_tests/backends/__init__.py create mode 100644 tests/unit_tests/backends/diffusion/__init__.py create mode 100644 tests/unit_tests/backends/megatron/__init__.py create mode 100644 tests/unit_tests/backends/torchtitan/__init__.py create mode 100644 tests/unit_tests/cli/__init__.py create mode 100644 tests/unit_tests/core/__init__.py create mode 100644 tests/unit_tests/core/backend/__init__.py create mode 100644 tests/unit_tests/core/launcher/__init__.py create mode 100644 tests/unit_tests/core/patches/__init__.py create mode 100644 tests/unit_tests/core/runtime/__init__.py create mode 100644 tests/unit_tests/core/trainer/__init__.py create mode 100644 tests/unit_tests/core/utils/__init__.py diff --git a/tests/unit_tests/backends/__init__.py b/tests/unit_tests/backends/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/backends/diffusion/__init__.py b/tests/unit_tests/backends/diffusion/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/backends/megatron/__init__.py b/tests/unit_tests/backends/megatron/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/backends/torchtitan/__init__.py b/tests/unit_tests/backends/torchtitan/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/cli/__init__.py b/tests/unit_tests/cli/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/__init__.py b/tests/unit_tests/core/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/backend/__init__.py b/tests/unit_tests/core/backend/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/launcher/__init__.py b/tests/unit_tests/core/launcher/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/patches/__init__.py b/tests/unit_tests/core/patches/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/runtime/__init__.py b/tests/unit_tests/core/runtime/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/trainer/__init__.py b/tests/unit_tests/core/trainer/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/utils/__init__.py b/tests/unit_tests/core/utils/__init__.py new file mode 100644 index 000000000..e69de29bb From 0b9d37c4280e735fdb2779edebcfbe05b1ce976d Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 8 Jul 2026 01:36:14 +0000 Subject: [PATCH 26/30] fix(diffusion): address PR review comments Clarify diffusion-generic docs and error messages, tighten Wan preprocessing quality issues, and fail fast for missing exponential shift parameters. Co-authored-by: Cursor --- examples/diffusion/README.md | 9 +++++---- primus/backends/diffusion/data/dataset.py | 5 ----- .../diffusion/data/processing_wanvideo.py | 17 +++++++++++------ .../diffusion/diffusion_pretrain_trainer.py | 14 +++++++------- primus/backends/diffusion/models/__init__.py | 2 +- .../backends/diffusion/schedulers/flow_match.py | 5 +++++ primus/backends/diffusion/trainers/__init__.py | 2 +- primus/backends/diffusion/utils/__init__.py | 2 +- .../diffusion/test_flow_match_scheduler.py | 14 ++++++++++++++ 9 files changed, 45 insertions(+), 25 deletions(-) create mode 100644 tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md index 2df6093f7..fc0c30606 100644 --- a/examples/diffusion/README.md +++ b/examples/diffusion/README.md @@ -1,8 +1,9 @@ -# Wan Examples +# Diffusion Examples -Wan examples exercise the independent PyTorch Diffusion backend under -`primus/backends/diffusion`. For backend details, data/checkpoint layout, and minimal -configs, see `primus/backends/diffusion/README.md`. +These examples exercise the independent PyTorch diffusion backend under +`primus/backends/diffusion`, including Wan video training and FLUX text-to-image +training. For backend details, data/checkpoint layout, and minimal configs, see +`primus/backends/diffusion/README.md`. ## Data diff --git a/primus/backends/diffusion/data/dataset.py b/primus/backends/diffusion/data/dataset.py index a528b3db1..139209923 100644 --- a/primus/backends/diffusion/data/dataset.py +++ b/primus/backends/diffusion/data/dataset.py @@ -246,11 +246,6 @@ def load_video_qwen_vl_utils( frames, sample_fps = fetch_video(video_dict, return_video_sample_fps=True) frames = frames.numpy() - # if is_even: - # return frames, sample_fps - # else: - # return frames[:-1], sample_fps - # Enforce VAE divisibility constraint actual_n = len(frames) if actual_n > 1: diff --git a/primus/backends/diffusion/data/processing_wanvideo.py b/primus/backends/diffusion/data/processing_wanvideo.py index 446411d98..675f3baa4 100644 --- a/primus/backends/diffusion/data/processing_wanvideo.py +++ b/primus/backends/diffusion/data/processing_wanvideo.py @@ -119,9 +119,14 @@ def _dynamic_target_size(self, image: Union[np.ndarray, Image.Image]) -> Dict[st width = max(self.width_division_factor, width) return {"height": height, "width": width} - def _resolve_sizes(self, image: Union[np.ndarray, Image.Image]) -> Tuple[Dict[str, int], Dict[str, int]]: - size = self.size - crop_size = self.crop_size + def _resolve_sizes( + self, + image: Union[np.ndarray, Image.Image], + size: Optional[Dict[str, int]] = None, + crop_size: Optional[Dict[str, int]] = None, + ) -> Tuple[Dict[str, int], Dict[str, int]]: + size = size if size is not None else self.size + crop_size = crop_size if crop_size is not None else self.crop_size if size is None and crop_size is None: size = self._dynamic_target_size(image) crop_size = dict(size) @@ -234,7 +239,7 @@ def preprocess( elif frame.shape[-1] == 4: # RGBA frame = frame[..., :3] - size_, crop_size_ = self._resolve_sizes(frame) + size_, crop_size_ = self._resolve_sizes(frame, size=size, crop_size=crop_size) if do_resize: frame = self.resize(frame, size_) if do_center_crop: @@ -257,7 +262,7 @@ def preprocess( elif image.shape[-1] == 4: # RGBA image = image[..., :3] - size_, crop_size_ = self._resolve_sizes(image) + size_, crop_size_ = self._resolve_sizes(image, size=size, crop_size=crop_size) if do_resize: image = self.resize(image, size_) if do_center_crop: @@ -336,7 +341,7 @@ def __init__(self, image_processor=None, tokenizer=None, max_text_length: Option try: print("[Debug-processor], loading default tokenizer: google/umt5-xxl") tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl") - except: + except Exception: logger.warning("Could not load default tokenizer, using None") tokenizer = None diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index 366042e24..be1d0111d 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -15,11 +15,11 @@ class DiffusionPretrainTrainer(BaseTrainer): - """Primus lifecycle wrapper for Wan diffusion training.""" + """Primus lifecycle wrapper for diffusion backend training.""" def __init__(self, backend_args: Any): super().__init__(backend_args=backend_args) - self.wan_trainer = None + self.diffusion_trainer = None @staticmethod def _as_dict(value: Any) -> dict: @@ -57,7 +57,7 @@ def setup(self): if missing: raise RuntimeError( "Diffusion backend missing required Python packages: " - f"{', '.join(missing)}. Install the Wan diffusion training extras first." + f"{', '.join(missing)}. Install the diffusion training extras first." ) if attention_backend: @@ -101,7 +101,7 @@ def init(self): ) model = get_model_builder(model_name)(model_config) dataset, processor = get_dataset_builder(dataset_name)(dataset_config) - self.wan_trainer = get_trainer_builder(trainer_name)( + self.diffusion_trainer = get_trainer_builder(trainer_name)( model=model, dataset=dataset, processor=processor, @@ -109,11 +109,11 @@ def init(self): ) def train(self): - if self.wan_trainer is None: + if self.diffusion_trainer is None: raise RuntimeError("DiffusionPretrainTrainer.init() must be called before train().") - self.wan_trainer.train() - self.wan_trainer.save_model() + self.diffusion_trainer.train() + self.diffusion_trainer.save_model() def cleanup(self, on_error: bool = False): try: diff --git a/primus/backends/diffusion/models/__init__.py b/primus/backends/diffusion/models/__init__.py index 2798dc3fb..7503f134c 100644 --- a/primus/backends/diffusion/models/__init__.py +++ b/primus/backends/diffusion/models/__init__.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Wan model exports for the Primus Wan backend.""" +"""Model exports for the Primus diffusion backend.""" from .wan import WanForTraining diff --git a/primus/backends/diffusion/schedulers/flow_match.py b/primus/backends/diffusion/schedulers/flow_match.py index ec5519574..cdd3f2f9b 100644 --- a/primus/backends/diffusion/schedulers/flow_match.py +++ b/primus/backends/diffusion/schedulers/flow_match.py @@ -59,6 +59,11 @@ def set_timesteps( if dynamic_shift_len is not None else self.exponential_shift_mu ) + if mu is None: + raise ValueError( + "`exponential_shift=True` requires either `dynamic_shift_len` " + "or `exponential_shift_mu`." + ) self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1)) else: self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) diff --git a/primus/backends/diffusion/trainers/__init__.py b/primus/backends/diffusion/trainers/__init__.py index d4df6355e..f1fbeeb37 100644 --- a/primus/backends/diffusion/trainers/__init__.py +++ b/primus/backends/diffusion/trainers/__init__.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Trainer registrations for the Primus Wan backend.""" +"""Trainer registrations for the Primus diffusion backend.""" from .fsdp2 import build_fsdp2_trainer diff --git a/primus/backends/diffusion/utils/__init__.py b/primus/backends/diffusion/utils/__init__.py index 3d7aa5fe7..5c11ee313 100644 --- a/primus/backends/diffusion/utils/__init__.py +++ b/primus/backends/diffusion/utils/__init__.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Utility package for the Primus Wan backend. +"""Utility package for the Primus diffusion backend. Heavy vision helpers are imported lazily by the qwen_vl_utils video path. """ diff --git a/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py b/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py new file mode 100644 index 000000000..9a97445dd --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py @@ -0,0 +1,14 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import pytest + +from primus.backends.diffusion.schedulers.flow_match import FlowMatchScheduler + + +def test_exponential_shift_requires_mu_or_dynamic_shift_len(): + with pytest.raises(ValueError, match="exponential_shift=True"): + FlowMatchScheduler(exponential_shift=True) From b8566cd1cf13f4ba4fc74abfdcfba9221a65c333 Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 8 Jul 2026 03:27:58 +0000 Subject: [PATCH 27/30] docs(diffusion): generalize utils docstring Avoid tying the shared diffusion utils package description to a specific video backend. Co-authored-by: Cursor --- primus/backends/diffusion/utils/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/primus/backends/diffusion/utils/__init__.py b/primus/backends/diffusion/utils/__init__.py index 5c11ee313..b4ecc5389 100644 --- a/primus/backends/diffusion/utils/__init__.py +++ b/primus/backends/diffusion/utils/__init__.py @@ -6,5 +6,5 @@ """Utility package for the Primus diffusion backend. -Heavy vision helpers are imported lazily by the qwen_vl_utils video path. +Heavy vision helpers are imported lazily by the video dataset path. """ From c8fa9e022b698a145f343ac13c52a5aa296a091c Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 14 Jul 2026 15:02:03 +0000 Subject: [PATCH 28/30] fix(diffusion): harden FLUX training numerics and edge cases - Build FLUX RoPE position ids in float32 instead of the model compute dtype so grid indices are not silently corrupted (bf16 only represents integers up to 256 exactly), matching the reference FLUX implementation. - Align frozen raw-path encoders (VAE/T5/CLIP) with the DiT dtype before use so non-bf16 runs no longer hit a dtype-mismatch error. - Validate that prompt-dropout empty encodings match the per-sample encoding shape and fail fast with a clear message instead of a cryptic in-place assignment error mid-training. - Add focused regression tests for the above. Co-authored-by: Cursor --- .../diffusion/data/flux_precomputed.py | 21 ++++++ .../diffusion/models/flux/train_pipeline.py | 29 +++++++- .../backends/diffusion/test_flux_backend.py | 72 +++++++++++++++++++ 3 files changed, 120 insertions(+), 2 deletions(-) diff --git a/primus/backends/diffusion/data/flux_precomputed.py b/primus/backends/diffusion/data/flux_precomputed.py index 7b3e2e510..7e54ebc2c 100644 --- a/primus/backends/diffusion/data/flux_precomputed.py +++ b/primus/backends/diffusion/data/flux_precomputed.py @@ -211,6 +211,26 @@ def _normalize_empty_encoding(tensor: torch.Tensor) -> torch.Tensor: return tensor[0] return tensor + def _check_empty_encoding_shapes(self, t5_encodings: torch.Tensor, clip_encodings: torch.Tensor) -> None: + # The empty encodings are broadcast-assigned into the per-sample encodings + # for dropped prompts, so their trailing (non-batch) shape must match the + # dataset encodings exactly. Fail fast with a clear message instead of a + # cryptic in-place assignment shape error mid-training. + assert self._empty_t5 is not None and self._empty_clip is not None + if self._empty_t5.shape != t5_encodings.shape[1:]: + raise ValueError( + "FLUX prompt dropout: empty T5 encoding shape " + f"{tuple(self._empty_t5.shape)} does not match the per-sample T5 encoding shape " + f"{tuple(t5_encodings.shape[1:])}. Regenerate t5_empty.npy with a matching " + "sequence length / hidden size." + ) + if self._empty_clip.shape != clip_encodings.shape[1:]: + raise ValueError( + "FLUX prompt dropout: empty CLIP encoding shape " + f"{tuple(self._empty_clip.shape)} does not match the per-sample CLIP encoding shape " + f"{tuple(clip_encodings.shape[1:])}. Regenerate clip_empty.npy with a matching shape." + ) + def _load_empty_encodings(self): if self._empty_t5 is not None and self._empty_clip is not None: return @@ -249,6 +269,7 @@ def prepare_batch( if self.prompt_dropout_prob > 0.0: self._load_empty_encodings() assert self._empty_t5 is not None and self._empty_clip is not None + self._check_empty_encoding_shapes(tensors["t5_encodings"], tensors["clip_encodings"]) bsz = tensors["t5_encodings"].shape[0] drop_mask = torch.rand((bsz,), device=device) < self.prompt_dropout_prob if drop_mask.any(): diff --git a/primus/backends/diffusion/models/flux/train_pipeline.py b/primus/backends/diffusion/models/flux/train_pipeline.py index d803b1629..4bea0e954 100644 --- a/primus/backends/diffusion/models/flux/train_pipeline.py +++ b/primus/backends/diffusion/models/flux/train_pipeline.py @@ -38,6 +38,21 @@ def _require_module(module: torch.nn.Module | None, name: str) -> torch.nn.Modul raise ValueError(f"FLUX raw image-text training requires `{name}` to be configured.") return module + @staticmethod + def _align_module_dtype(module: torch.nn.Module, *, device: torch.device, dtype: torch.dtype) -> None: + """Move a frozen encoder to the target device/dtype only when it differs. + + `nn.Module.to(dtype=...)` casts floating-point params/buffers only, leaving + integer buffers (e.g. token position ids) untouched. We inspect the first + floating-point parameter to avoid re-casting on every step. + """ + current = next((p for p in module.parameters() if p.is_floating_point()), None) + if current is None: + module.to(device=device) + return + if current.dtype != dtype or current.device != device: + module.to(device=device, dtype=dtype) + def _prepare_precomputed( self, *, @@ -75,6 +90,12 @@ def _prepare_raw( if not isinstance(prompts, list): raise ValueError("FLUX raw batch requires `prompts` as list[str].") + # The frozen encoders are built in bf16 regardless of the training compute + # dtype; align them with the DiT dtype so running them on `image`/inputs of + # `dtype` does not raise a dtype-mismatch error (e.g. for fp32 runs). + self._align_module_dtype(ae, device=device, dtype=dtype) + self._align_module_dtype(t5, device=device, dtype=dtype) + self._align_module_dtype(clip, device=device, dtype=dtype) ae.eval() t5.eval() clip.eval() @@ -128,15 +149,19 @@ def compute_loss( target = noise - labels _, _, latent_height, latent_width = noisy_latents.shape + # Position ids are integer grid indices consumed by RoPE; build them in + # float32 (independent of the model compute dtype) so that indices remain + # exactly representable. bf16 only represents integers up to 256 exactly, + # which would silently corrupt positions for larger latent grids. img_ids = create_position_encoding_for_latents( bsz, latent_height, latent_width, position_dim=3, device=device, - dtype=dtype, + dtype=torch.float32, ) - txt_ids = torch.zeros(bsz, t5_encodings.shape[1], 3, device=device, dtype=dtype) + txt_ids = torch.zeros(bsz, t5_encodings.shape[1], 3, device=device, dtype=torch.float32) noisy_latents = pack_latents(noisy_latents) target = pack_latents(target) diff --git a/tests/unit_tests/backends/diffusion/test_flux_backend.py b/tests/unit_tests/backends/diffusion/test_flux_backend.py index d4f087fea..b73e146c1 100644 --- a/tests/unit_tests/backends/diffusion/test_flux_backend.py +++ b/tests/unit_tests/backends/diffusion/test_flux_backend.py @@ -149,6 +149,32 @@ def test_flux_precomputed_processor_stacks_and_drops_empty_encodings(tmp_path): assert torch.count_nonzero(out["clip_encodings"]) == 0 +def test_flux_precomputed_processor_rejects_mismatched_empty_encoding(tmp_path): + empty_dir = tmp_path / "empty" + empty_dir.mkdir() + # Empty T5 encoding has sequence length 5, but the batch samples use length 3. + np.save(empty_dir / "t5_empty.npy", np.zeros((1, 5, 8), dtype=np.float32)) + np.save(empty_dir / "clip_empty.npy", np.zeros((1, 4), dtype=np.float32)) + + processor = FluxPrecomputedProcessor( + { + "prompt_dropout_prob": 1.0, + "empty_encodings_path": str(empty_dir), + } + ) + batch = [ + { + "t5_encodings": torch.ones(3, 8), + "clip_encodings": torch.ones(4), + "mean": torch.zeros(1, 2, 2), + "logvar": torch.zeros(1, 2, 2), + } + ] + + with pytest.raises(ValueError, match="empty T5 encoding shape"): + processor.prepare_batch(batch=batch, device=torch.device("cpu"), dtype=torch.float32) + + def test_flux_raw_processor_prepares_images_and_prompts(): from PIL import Image @@ -200,6 +226,52 @@ def test_tiny_flux_model_computes_precomputed_loss(): assert torch.isfinite(outputs["loss"]) +def test_flux_position_ids_are_float32_regardless_of_model_dtype(): + model = build_flux_model( + { + "config": { + "model_variant": "flux-dev", + "guidance": 1.0, + "params": { + "in_channels": 4, + "out_channels": 4, + "vec_in_dim": 4, + "context_in_dim": 8, + "hidden_size": 12, + "num_heads": 2, + "depth": 1, + "depth_single_blocks": 1, + "axes_dim": [2, 2, 2], + }, + } + } + ) + # Force a low-precision compute dtype; position ids must stay float32 so RoPE + # grid indices are not corrupted (bf16 only represents integers up to 256). + model.dit = model.dit.to(dtype=torch.bfloat16) + + captured = {} + original_forward = model.dit.forward + + def capturing_forward(*args, **kwargs): + captured["img_ids_dtype"] = kwargs["img_ids"].dtype + captured["txt_ids_dtype"] = kwargs["txt_ids"].dtype + return original_forward(*args, **kwargs) + + model.dit.forward = capturing_forward + model.forward_train( + { + "t5_encodings": torch.randn(2, 3, 8, dtype=torch.bfloat16), + "clip_encodings": torch.randn(2, 4, dtype=torch.bfloat16), + "mean": torch.randn(2, 1, 2, 2, dtype=torch.bfloat16), + "logvar": torch.zeros(2, 1, 2, 2, dtype=torch.bfloat16), + } + ) + + assert captured["img_ids_dtype"] == torch.float32 + assert captured["txt_ids_dtype"] == torch.float32 + + def test_tiny_flux_model_computes_raw_loss_with_dummy_encoders(): class DummyAutoencoder(torch.nn.Module): def encode(self, image): From 288442f6dbe0203a6fc364f3a475a32bb4d00806 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 14 Jul 2026 15:23:32 +0000 Subject: [PATCH 29/30] fix(diffusion): update module utils import after main merge Co-authored-by: Cursor --- primus/backends/diffusion/diffusion_adapter.py | 2 +- primus/backends/diffusion/diffusion_pretrain_trainer.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/primus/backends/diffusion/diffusion_adapter.py b/primus/backends/diffusion/diffusion_adapter.py index 7e358db02..878ea3fea 100644 --- a/primus/backends/diffusion/diffusion_adapter.py +++ b/primus/backends/diffusion/diffusion_adapter.py @@ -11,7 +11,7 @@ from primus.backends.diffusion.argument_builder import DiffusionArgBuilder from primus.core.backend.backend_adapter import BackendAdapter -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class DiffusionAdapter(BackendAdapter): diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index be1d0111d..b1b7fd650 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -10,8 +10,8 @@ from typing import Any from primus.core.trainer.base_trainer import BaseTrainer +from primus.core.utils.module_utils import log_rank_0 from primus.core.utils.yaml_utils import nested_namespace_to_dict -from primus.modules.module_utils import log_rank_0 class DiffusionPretrainTrainer(BaseTrainer): From cf894d4764711402b3b26cd9d70a4f8d28bff634 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 14 Jul 2026 15:49:37 +0000 Subject: [PATCH 30/30] docs(diffusion): clarify FLUX pretraining launch Document the FLUX precomputed and raw image-text launch paths with the required data and checkpoint environment variables, and remove Wan-only wording from shared diffusion docs. Co-authored-by: Cursor --- examples/diffusion/README.md | 55 ++++++++++++++++++++++++++--- primus/backends/diffusion/README.md | 49 ++++++++++++++++--------- 2 files changed, 84 insertions(+), 20 deletions(-) diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md index fc0c30606..5134b6068 100644 --- a/examples/diffusion/README.md +++ b/examples/diffusion/README.md @@ -7,7 +7,7 @@ training. For backend details, data/checkpoint layout, and minimal configs, see ## Data -The default smoke-test dataset is `zirui3/tiny-video-samples` on Hugging Face: +Wan examples use the `zirui3/tiny-video-samples` smoke-test dataset: ```bash huggingface-cli download zirui3/tiny-video-samples \ @@ -23,6 +23,15 @@ Expected layout: data/*.mp4 ``` +FLUX examples use either: + +- precomputed encodings: a Hugging Face `datasets` directory with + `t5_encodings`, `clip_encodings`, `mean`, and `logvar` fields. +- raw image-text data: a local webdataset directory, a local JSONL file, a local + Hugging Face dataset, or a Hugging Face dataset repo. + +See `primus/backends/diffusion/README.md` for checkpoint and data layout details. + ## Run Set the shared `torchrun` environment first: @@ -35,7 +44,7 @@ export MASTER_PORT=${MASTER_PORT:-29500} export GPUS_PER_NODE=${GPUS_PER_NODE:-8} ``` -### Pretrain +### Wan Pretrain ```bash DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ @@ -54,7 +63,45 @@ torchrun \ Use `SP_SIZE=4` or `SP_SIZE=8` to enable Ulysses sequence parallelism when the model head count supports it. -### Posttrain +### FLUX Precomputed Pretrain + +```bash +DATASET_PATH=/data/flux_precomputed \ +EMPTY_ENCODINGS_PATH=/data/flux_empty_encodings \ +MAX_STEPS=10 \ +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ + -m primus.cli.main train pretrain \ + --config examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml +``` + +Set `PRETRAINED_PATH=/path/to/flux1-dev.safetensors` only when you want to +initialize the DiT from existing FLUX weights. Empty `PRETRAINED_PATH` means +random initialization. + +### FLUX Raw Image-Text Pretrain + +```bash +DATASET=cc12m-test \ +DATASET_PATH=/data/cc12m_test \ +T5_ENCODER=google/t5-v1_1-xxl \ +CLIP_ENCODER=openai/clip-vit-large-patch14 \ +VAE_CHECKPOINT=black-forest-labs/FLUX.1-dev/ae.safetensors \ +MAX_STEPS=10 \ +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ + -m primus.cli.main train pretrain \ + --config examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml +``` + +For raw training, `T5_ENCODER`, `CLIP_ENCODER`, and `VAE_CHECKPOINT` can be local +paths or Hugging Face identifiers. + +### Wan Posttrain ```bash INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B \ @@ -71,5 +118,5 @@ torchrun \ The MI355X configs use Primus-style override sections such as `training`, `data`, `parallelism`, `optimizer`, `runtime`, and `metrics`. The diffusion -adapter normalizes those sections into the Wan model/dataset/trainer +adapter normalizes those sections into diffusion model/dataset/trainer arguments at runtime. diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md index b84b552f6..951eee590 100644 --- a/primus/backends/diffusion/README.md +++ b/primus/backends/diffusion/README.md @@ -91,10 +91,12 @@ model: `autoencoder` accepts either a local safetensors path or `repo_id/filename`; for example `black-forest-labs/FLUX.1-dev/ae.safetensors`. -The shipped precomputed example expects these paths through environment -variables: +The shipped precomputed example uses random DiT initialization unless +`PRETRAINED_PATH` points at FLUX DiT weights. It expects preprocessed data and, +when prompt dropout is enabled, empty prompt encodings: ```text +PRETRAINED_PATH=/path/to/flux1-dev.safetensors # optional DATASET_PATH=/path/to/cc12m_preprocessed EMPTY_ENCODINGS_PATH=/path/to/empty_encodings ``` @@ -132,12 +134,20 @@ DATASET=cc12m-test DATASET_PATH=/path/to/cc12m_test ``` -The raw FLUX example is: +The raw FLUX example loads frozen encoders online. The defaults use Hugging Face +IDs for T5/CLIP/AE, but each value can also be a local path: ```text -examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml +DATASET=cc12m-test +DATASET_PATH=/path/to/cc12m_test +T5_ENCODER=google/t5-v1_1-xxl +CLIP_ENCODER=openai/clip-vit-large-patch14 +VAE_CHECKPOINT=black-forest-labs/FLUX.1-dev/ae.safetensors ``` +The raw example config is +`examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml`. + Raw FLUX data can also be a local JSONL file: ```jsonl @@ -254,9 +264,9 @@ for debugging. ## Launch Training runs through the Primus CLI (`primus.cli.main train `) -under `torchrun`. Single-node and multi-node share **one** launch command: the -same `torchrun` invocation runs on every node, parameterized by the standard -rendezvous variables. The defaults below give a single-node 8-GPU run. +under `torchrun`. Single-node and multi-node share one launch shape; choose the +config path for the workload you want to run. The defaults below give a +single-node 8-GPU FLUX precomputed run. ```bash # distributed knobs (defaults = single node, 8 GPUs) @@ -266,27 +276,32 @@ export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} export MASTER_PORT=${MASTER_PORT:-29500} export GPUS_PER_NODE=${GPUS_PER_NODE:-8} +# FLUX precomputed. Set DATASET_PATH and EMPTY_ENCODINGS_PATH first. +export CONFIG=examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml + torchrun \ --nnodes="$NNODES" --node_rank="$NODE_RANK" \ --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ --nproc_per_node="$GPUS_PER_NODE" \ - -m primus.cli.main train pretrain --config /path/to/wan_config.yaml + -m primus.cli.main train pretrain --config "$CONFIG" ``` - **Single node**: run as-is (the defaults above). - **Multi-node**: run the same command on each node with a shared `MASTER_ADDR`/`MASTER_PORT` (a routable IP of node rank 0) and a distinct `NODE_RANK` per node. World size is `NNODES * GPUS_PER_NODE`. -- **Post-train**: identical command with `train posttrain` and a posttrain config. +- **FLUX raw image-text**: set `CONFIG` to + `examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml` and set + `DATASET_PATH`, `T5_ENCODER`, `CLIP_ENCODER`, and `VAE_CHECKPOINT` as needed. +- **Wan/post-train**: use the same command with the matching Wan config; for + post-train use `train posttrain` and a posttrain config. Useful runtime knobs: -- `trainer.args.attention_backend`: defaults to `flash_attn_aiter` for Wan - training on ROCm. Use `sdpa` as the portable fallback when AITER flash - attention is unavailable. -- `trainer.args.sp_size`: Ulysses sequence parallel size. It must divide the - model attention head count; for example Wan2.1-1.3B supports `sp_size=4` but - not `sp_size=8`. +- `trainer.args.attention_backend`: defaults to `flash_attn_aiter` on ROCm. Use + `sdpa` as the portable fallback when AITER flash attention is unavailable. +- `trainer.args.sp_size`: Ulysses sequence parallel size for Wan. FLUX currently + requires `sp_size=1`. - `trainer.args.dp_replicate`: data parallel replication size. - `FIXED_TIMESTEP` and `FIXED_SEED`: optional debug variables for reproducible loss-alignment checks. @@ -303,9 +318,11 @@ export AMD_COMGR_CACHE_DIR=/path/to/large/cache/comgr ## Primus-Style Configs New examples should use Primus-style override sections and let the diffusion -adapter normalize them into Wan args. See: +adapter normalize them into diffusion model/dataset/trainer args. See: ```text +examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml +examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml