From 60f3d602c06b3029739eb28d23d6e9470d012a6a Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 11:32:23 +0000 Subject: [PATCH 1/4] feat(diffusion): add Wan2.2 T2V A14B presets Support Diffusers two-transformer Wan2.2 T2V-A14B DiT loading, add AutoModel-aligned high/low-noise SFT configs, and cover scheduler/key mapping with focused tests. Co-authored-by: Cursor --- .../wan2.2_t2v_a14b-posttrain-high-noise.yaml | 62 ++++++++++ .../wan2.2_t2v_a14b-posttrain-low-noise.yaml | 62 ++++++++++ primus/backends/diffusion/README.md | 19 +++ primus/backends/diffusion/argument_builder.py | 19 +++ .../diffusion/models/registrations/wan.py | 108 +++++++++++++++++- .../models/diffusion/wan2.2_t2v_a14b.yaml | 32 ++++++ .../wan2.2_t2v_a14b_high_noise_sft.yaml | 10 ++ .../wan2.2_t2v_a14b_low_noise_sft.yaml | 10 ++ .../diffusion/test_wan_argument_builder.py | 10 ++ .../diffusion/test_wan_model_registration.py | 69 +++++++++++ 10 files changed, 399 insertions(+), 2 deletions(-) create mode 100644 examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml create mode 100644 examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml create mode 100644 primus/configs/models/diffusion/wan2.2_t2v_a14b.yaml create mode 100644 primus/configs/models/diffusion/wan2.2_t2v_a14b_high_noise_sft.yaml create mode 100644 primus/configs/models/diffusion/wan2.2_t2v_a14b_low_noise_sft.yaml create mode 100644 tests/unit_tests/backends/diffusion/test_wan_model_registration.py diff --git a/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml b/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml new file mode 100644 index 000000000..4984e3fcc --- /dev/null +++ b/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml @@ -0,0 +1,62 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:wan2.2_t2v_a14b-posttrain-high-noise} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + post_trainer: + framework: diffusion + config: post_trainer.yaml + + # Wan2.2-T2V-A14B high-noise stage. This matches AutoModel's + # model.stage=high_noise behavior: train Diffusers `transformer` on + # sigma in [boundary_ratio, 1.0]. + model: wan2.2_t2v_a14b_high_noise_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_t2v_a14b-posttrain-high-noise} + save_steps: 0 + run_name: wan2.2_t2v_a14b-posttrain-high-noise + + data: + dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} + data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} + frame_num: ${FRAME_NUM:81} + video_backend: imageio + text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.2-T2V-A14B-Diffusers/tokenizer} + height: ${HEIGHT:480} + width: ${WIDTH:832} + + parallelism: + sp_size: ${SP_SIZE:1} + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: 5.0e-6 + weight_decay: 0.01 + + scheduler: + shift: 3.0 + sigma_min: ${SIGMA_MIN:0.875} + sigma_max: ${SIGMA_MAX:1.0} + extra_one_step: true + num_train_timesteps: 1000 + + runtime: + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} + report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml b/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml new file mode 100644 index 000000000..f1db23fed --- /dev/null +++ b/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml @@ -0,0 +1,62 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:wan2.2_t2v_a14b-posttrain-low-noise} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + post_trainer: + framework: diffusion + config: post_trainer.yaml + + # Wan2.2-T2V-A14B low-noise stage. This matches AutoModel's + # model.stage=low_noise behavior: train Diffusers `transformer_2` on + # sigma in [0.0, boundary_ratio]. + model: wan2.2_t2v_a14b_low_noise_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_t2v_a14b-posttrain-low-noise} + save_steps: 0 + run_name: wan2.2_t2v_a14b-posttrain-low-noise + + data: + dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} + data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} + frame_num: ${FRAME_NUM:81} + video_backend: imageio + text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.2-T2V-A14B-Diffusers/tokenizer} + height: ${HEIGHT:480} + width: ${WIDTH:832} + + parallelism: + sp_size: ${SP_SIZE:1} + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: 5.0e-6 + weight_decay: 0.01 + + scheduler: + shift: 3.0 + sigma_min: ${SIGMA_MIN:0.0} + sigma_max: ${SIGMA_MAX:0.875} + extra_one_step: true + num_train_timesteps: 1000 + + runtime: + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} + report_to: none diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md index 9b7db9704..61d90fc03 100644 --- a/primus/backends/diffusion/README.md +++ b/primus/backends/diffusion/README.md @@ -64,6 +64,7 @@ model(s) you plan to train: | --- | --- | --- | | 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.2-T2V-A14B | `wan2.2_t2v_a14b.yaml` | [Wan-AI/Wan2.2-T2V-A14B-Diffusers](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B-Diffusers) | | Wan2.1-T2V-14B | (no shipped preset) | [Wan-AI/Wan2.1-T2V-14B](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B) | ```bash @@ -74,6 +75,10 @@ huggingface-cli download Wan-AI/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 + +# Wan2.2-T2V-A14B Diffusers DiT pipeline +huggingface-cli download Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --local-dir /models/Wan2.2-T2V-A14B-Diffusers ``` A downloaded Wan repo looks like this (the T5 encoder, VAE, and tokenizer are @@ -103,6 +108,20 @@ 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`. +Wan2.2-T2V-A14B is different: it is published as a Diffusers two-transformer +pipeline. Primus selects the active DiT subfolder with +`model.config.config.active_transformer`: + +- `transformer`: high-noise stage, trained on sigma `[boundary_ratio, 1.0]`. +- `transformer_2`: low-noise stage, trained on sigma `[0.0, boundary_ratio]`. + +The shipped examples use `boundary_ratio=0.875` and `shift=3.0`, matching the +NeMo AutoModel Wan2.2-T2V-A14B recipe. The native Primus Wan path still expects +the frozen UMT5/VAE encoder assets in original Wan `.pth` format; the A14B +presets default to the compatible Wan2.1-14B `models_t5_umt5-xxl-enc-bf16.pth` +and `Wan2.1_VAE.pth`, while the tokenizer defaults to the A14B Diffusers +`tokenizer` subfolder. + ## Pre-flight validation The diffusion prepare hook validates the prepared assets; it does **not** diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py index b6b444ceb..e5ec4527b 100644 --- a/primus/backends/diffusion/argument_builder.py +++ b/primus/backends/diffusion/argument_builder.py @@ -163,6 +163,7 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, optimizer = params.get("optimizer") or {} runtime = params.get("runtime") or {} metrics = params.get("metrics") or {} + scheduler = params.get("scheduler") or params.get("flow_matching") or {} training_map = { ("steps",): ("max_steps",), @@ -244,6 +245,24 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, if value is not None: self._set_nested(trainer_args, target_path, value) + scheduler_cfg = copy.deepcopy(runtime.get("flow_match_scheduler") or {}) + scheduler_cfg.update(scheduler) + for key in ( + "shift", + "sigma_min", + "sigma_max", + "extra_one_step", + "num_train_timesteps", + "inverse_timesteps", + "reverse_sigmas", + "exponential_shift", + "exponential_shift_mu", + "shift_terminal", + ): + value = scheduler_cfg.get(key) + if value is not None: + self._set_nested(trainer_args, ("flow_match_scheduler", key), value) + log_freq = metrics.get("log_freq") if log_freq is not None: self._set_nested(trainer_args, ("logging_steps",), log_freq) diff --git a/primus/backends/diffusion/models/registrations/wan.py b/primus/backends/diffusion/models/registrations/wan.py index 491284ff7..45846b9c7 100644 --- a/primus/backends/diffusion/models/registrations/wan.py +++ b/primus/backends/diffusion/models/registrations/wan.py @@ -52,15 +52,95 @@ def _load_state_dict(path: str) -> dict[str, torch.Tensor]: return obj -def _load_dit_weights_into_module(dit: torch.nn.Module, pretrained_path: str): +def _resolve_dit_checkpoint_path(pretrained_path: str, active_transformer: str | None = None) -> str: + if os.path.isfile(pretrained_path): + return pretrained_path + + model_index = os.path.join(pretrained_path, "model_index.json") + if os.path.exists(model_index): + transformer_name = active_transformer or "transformer" + if transformer_name not in ("transformer", "transformer_2"): + raise ValueError( + "Wan2.2 A14B Diffusers checkpoints support active_transformer=" + f"'transformer' or 'transformer_2', got {transformer_name!r}" + ) + checkpoint_path = os.path.join(pretrained_path, transformer_name) + if not os.path.isdir(checkpoint_path): + raise FileNotFoundError( + f"Diffusers checkpoint at {pretrained_path} does not contain {transformer_name!r}" + ) + logger.info(f"Using Diffusers Wan transformer subfolder: {checkpoint_path}") + return checkpoint_path + + return pretrained_path + + +def _convert_diffusers_wan_dit_state_dict( + state_dict: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + """Map Diffusers WanTransformer3DModel keys to Primus WanDiT keys.""" + converted: dict[str, torch.Tensor] = {} + replacements = ( + ("condition_embedder.text_embedder.linear_1.", "text_embedding.0."), + ("condition_embedder.text_embedder.linear_2.", "text_embedding.2."), + ("condition_embedder.time_embedder.linear_1.", "time_embedding.0."), + ("condition_embedder.time_embedder.linear_2.", "time_embedding.2."), + ("condition_embedder.time_proj.", "time_projection.1."), + ("proj_out.", "head.head."), + (".attn1.to_q.", ".self_attn.q."), + (".attn1.to_k.", ".self_attn.k."), + (".attn1.to_v.", ".self_attn.v."), + (".attn1.to_out.0.", ".self_attn.o."), + (".attn1.norm_q.", ".self_attn.norm_q."), + (".attn1.norm_k.", ".self_attn.norm_k."), + (".attn2.to_q.", ".cross_attn.q."), + (".attn2.to_k.", ".cross_attn.k."), + (".attn2.to_v.", ".cross_attn.v."), + (".attn2.to_out.0.", ".cross_attn.o."), + (".attn2.norm_q.", ".cross_attn.norm_q."), + (".attn2.norm_k.", ".cross_attn.norm_k."), + (".ffn.net.0.proj.", ".ffn.0."), + (".ffn.net.2.", ".ffn.2."), + (".norm2.", ".norm3."), + (".scale_shift_table", ".modulation"), + ) + + for key, value in state_dict.items(): + if key == "scale_shift_table": + converted["head.modulation"] = value + continue + new_key = key + for old, new in replacements: + new_key = new_key.replace(old, new) + converted[new_key] = value + return converted + + +def _maybe_convert_diffusers_wan_dit_state_dict( + state_dict: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + if any(key.startswith("condition_embedder.") or ".attn1." in key for key in state_dict): + return _convert_diffusers_wan_dit_state_dict(state_dict) + return state_dict + + +def _load_dit_weights_into_module( + dit: torch.nn.Module, + pretrained_path: str, + *, + active_transformer: str | None = None, +): """ 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...` """ + pretrained_path = _resolve_dit_checkpoint_path(pretrained_path, active_transformer) + # Direct file if os.path.isfile(pretrained_path): state = _strip_module_prefix(_load_state_dict(pretrained_path)) + state = _maybe_convert_diffusers_wan_dit_state_dict(state) # 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.")} @@ -92,6 +172,7 @@ def _load_dit_weights_into_module(dit: torch.nn.Module, pretrained_path: str): part = _strip_module_prefix(_load_state_dict(ckpt)) merged.update(part) + merged = _maybe_convert_diffusers_wan_dit_state_dict(merged) if any(k.startswith("dit.") for k in merged): merged = {k[len("dit.") :]: v for k, v in merged.items() if k.startswith("dit.")} @@ -123,6 +204,10 @@ def build_wan_model(model_config: dict): import json cfg_path = os.path.join(pretrained_path, "config.json") + active_transformer = cfg_dict.get("active_transformer") + if not os.path.exists(cfg_path) and os.path.exists(os.path.join(pretrained_path, "model_index.json")): + transformer_name = active_transformer or "transformer" + cfg_path = os.path.join(pretrained_path, transformer_name, "config.json") if os.path.exists(cfg_path): logger.info(f"Loading config from {cfg_path}") with open(cfg_path) as f: @@ -138,6 +223,13 @@ def build_wan_model(model_config: dict): "out_dim": "dit_out_channels", "freq_dim": "dit_freq_dim", "text_len": "text_len", + "num_attention_heads": "dit_num_heads", + "ffn_dim": "dit_intermediate_size", + "in_channels": "dit_in_channels", + "out_channels": "dit_out_channels", + "patch_size": "dit_patch_size", + "text_dim": "dit_text_dim", + "eps": "dit_eps", } for k, v in loaded_cfg.items(): @@ -145,6 +237,14 @@ def build_wan_model(model_config: dict): # 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 + if ( + "dit_hidden_size" not in cfg_dict + and "num_attention_heads" in loaded_cfg + and "attention_head_dim" in loaded_cfg + ): + cfg_dict["dit_hidden_size"] = int(loaded_cfg["num_attention_heads"]) * int( + loaded_cfg["attention_head_dim"] + ) # Build a WanVideoConfig-compatible object to maximize reuse of existing fields. model_cfg = WanVideoConfig(**cfg_dict) @@ -200,7 +300,11 @@ def build_wan_model(model_config: dict): 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_dit_weights_into_module( + dit, + pretrained_path, + active_transformer=getattr(model_cfg, "active_transformer", None), + ) components = WanComponents(dit=dit, vae=vae, text_encoder=text_encoder, image_encoder=None) pipeline = WanFlowMatchTrainPipeline() diff --git a/primus/configs/models/diffusion/wan2.2_t2v_a14b.yaml b/primus/configs/models/diffusion/wan2.2_t2v_a14b.yaml new file mode 100644 index 000000000..99e5b4208 --- /dev/null +++ b/primus/configs/models/diffusion/wan2.2_t2v_a14b.yaml @@ -0,0 +1,32 @@ +model: + name: wan + config: + # Wan2.2-T2V-A14B is distributed as a Diffusers two-transformer pipeline. + # `active_transformer` selects which stage DiT to load: + # - transformer: high-noise stage + # - transformer_2: low-noise stage + load_from_pretrained_path: ${PRETRAINED_PATH:/models/Wan2.2-T2V-A14B-Diffusers} + config: + model_type: t2v + active_transformer: ${ACTIVE_TRANSFORMER:transformer} + boundary_ratio: ${BOUNDARY_RATIO:0.875} + separated_timestep: false + fuse_vae_embedding_in_latents: false + trainable_modules: dit + vae_type: wan_video_vae + dit_hidden_size: 5120 + dit_num_layers: 40 + dit_num_heads: 40 + dit_intermediate_size: 13824 + dit_patch_size: [1, 2, 2] + dit_in_channels: 16 + dit_out_channels: 16 + dit_freq_dim: 256 + dit_text_dim: 4096 + dit_eps: 1.0e-6 + encoder: + # Wan2.2 A14B Diffusers ships VAE/T5 as Diffusers/Transformers components. + # Primus' native Wan path currently consumes the original Wan .pth encoder assets; + # Wan2.1-14B uses the compatible UMT5 encoder and z_dim=16 VAE. + t5_encoder: ${TEXT_ENCODER:/models/Wan2.1-T2V-14B/models_t5_umt5-xxl-enc-bf16.pth} + autoencoder: ${VAE_CHECKPOINT:/models/Wan2.1-T2V-14B/Wan2.1_VAE.pth} diff --git a/primus/configs/models/diffusion/wan2.2_t2v_a14b_high_noise_sft.yaml b/primus/configs/models/diffusion/wan2.2_t2v_a14b_high_noise_sft.yaml new file mode 100644 index 000000000..ad67c7388 --- /dev/null +++ b/primus/configs/models/diffusion/wan2.2_t2v_a14b_high_noise_sft.yaml @@ -0,0 +1,10 @@ +extends: + - wan2.2_t2v_a14b.yaml + +# High-noise SFT stage, aligned with NeMo AutoModel Wan2.2-T2V-A14B: +# trains the first Diffusers transformer on sigma in [boundary_ratio, 1.0]. +model: + config: + load_from_pretrained_path: ${INIT_CHECKPOINT:/models/Wan2.2-T2V-A14B-Diffusers} + config: + active_transformer: transformer diff --git a/primus/configs/models/diffusion/wan2.2_t2v_a14b_low_noise_sft.yaml b/primus/configs/models/diffusion/wan2.2_t2v_a14b_low_noise_sft.yaml new file mode 100644 index 000000000..67ece71f4 --- /dev/null +++ b/primus/configs/models/diffusion/wan2.2_t2v_a14b_low_noise_sft.yaml @@ -0,0 +1,10 @@ +extends: + - wan2.2_t2v_a14b.yaml + +# Low-noise SFT stage, aligned with NeMo AutoModel Wan2.2-T2V-A14B: +# trains the second Diffusers transformer on sigma in [0.0, boundary_ratio]. +model: + config: + load_from_pretrained_path: ${INIT_CHECKPOINT:/models/Wan2.2-T2V-A14B-Diffusers} + config: + active_transformer: transformer_2 diff --git a/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py b/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py index 9d07d4f07..6672d46dc 100644 --- a/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py +++ b/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py @@ -88,6 +88,12 @@ def test_maps_primus_style_sections_to_wan_runtime_config(): "seed": 1234, "fsdp2_reshard_after_forward": False, }, + "scheduler": { + "shift": 3.0, + "sigma_min": 0.875, + "sigma_max": 1.0, + "extra_one_step": True, + }, "metrics": { "log_freq": 5, "enable_wandb": False, @@ -137,6 +143,10 @@ def test_maps_primus_style_sections_to_wan_runtime_config(): assert trainer_args["seed"] == 1234 assert trainer_args["fsdp2_reshard_after_forward"] is False assert trainer_args["logging_steps"] == 5 + assert trainer_args["flow_match_scheduler"]["shift"] == 3.0 + assert trainer_args["flow_match_scheduler"]["sigma_min"] == 0.875 + assert trainer_args["flow_match_scheduler"]["sigma_max"] == 1.0 + assert trainer_args["flow_match_scheduler"]["extra_one_step"] is True def test_defaults_propagate_when_optional_sections_are_omitted(): diff --git a/tests/unit_tests/backends/diffusion/test_wan_model_registration.py b/tests/unit_tests/backends/diffusion/test_wan_model_registration.py new file mode 100644 index 000000000..3772f4fbb --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_wan_model_registration.py @@ -0,0 +1,69 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from pathlib import Path + +import torch +import pytest + +from primus.backends.diffusion.models.registrations.wan import ( + _convert_diffusers_wan_dit_state_dict, + _resolve_dit_checkpoint_path, +) + + +def test_resolve_diffusers_transformer_subfolder(tmp_path: Path): + (tmp_path / "model_index.json").write_text("{}", encoding="utf-8") + (tmp_path / "transformer").mkdir() + (tmp_path / "transformer_2").mkdir() + + assert _resolve_dit_checkpoint_path(str(tmp_path), "transformer") == str(tmp_path / "transformer") + assert _resolve_dit_checkpoint_path(str(tmp_path), "transformer_2") == str(tmp_path / "transformer_2") + + +def test_resolve_diffusers_transformer_rejects_unknown_stage(tmp_path: Path): + (tmp_path / "model_index.json").write_text("{}", encoding="utf-8") + + with pytest.raises(ValueError, match="active_transformer"): + _resolve_dit_checkpoint_path(str(tmp_path), "middle_noise") + + +def test_convert_diffusers_wan_dit_keys_to_primus_keys(): + state = { + "condition_embedder.text_embedder.linear_1.weight": torch.empty(1), + "condition_embedder.text_embedder.linear_2.bias": torch.empty(1), + "condition_embedder.time_embedder.linear_1.weight": torch.empty(1), + "condition_embedder.time_proj.bias": torch.empty(1), + "patch_embedding.weight": torch.empty(1), + "scale_shift_table": torch.empty(1), + "proj_out.weight": torch.empty(1), + "blocks.0.scale_shift_table": torch.empty(1), + "blocks.0.attn1.to_q.weight": torch.empty(1), + "blocks.0.attn1.to_out.0.bias": torch.empty(1), + "blocks.0.attn2.to_k.weight": torch.empty(1), + "blocks.0.attn2.norm_q.weight": torch.empty(1), + "blocks.0.ffn.net.0.proj.weight": torch.empty(1), + "blocks.0.ffn.net.2.bias": torch.empty(1), + "blocks.0.norm2.weight": torch.empty(1), + } + + converted = _convert_diffusers_wan_dit_state_dict(state) + + assert "text_embedding.0.weight" in converted + assert "text_embedding.2.bias" in converted + assert "time_embedding.0.weight" in converted + assert "time_projection.1.bias" in converted + assert "patch_embedding.weight" in converted + assert "head.modulation" in converted + assert "head.head.weight" in converted + assert "blocks.0.modulation" in converted + assert "blocks.0.self_attn.q.weight" in converted + assert "blocks.0.self_attn.o.bias" in converted + assert "blocks.0.cross_attn.k.weight" in converted + assert "blocks.0.cross_attn.norm_q.weight" in converted + assert "blocks.0.ffn.0.weight" in converted + assert "blocks.0.ffn.2.bias" in converted + assert "blocks.0.norm3.weight" in converted From 25600096afbc451c0d5a7d01fa7f899eab7987a6 Mon Sep 17 00:00:00 2001 From: zirui Date: Tue, 7 Jul 2026 11:40:53 +0000 Subject: [PATCH 2/4] fix isort issue --- .../backends/diffusion/test_wan_model_registration.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/backends/diffusion/test_wan_model_registration.py b/tests/unit_tests/backends/diffusion/test_wan_model_registration.py index 3772f4fbb..700b30355 100644 --- a/tests/unit_tests/backends/diffusion/test_wan_model_registration.py +++ b/tests/unit_tests/backends/diffusion/test_wan_model_registration.py @@ -6,8 +6,8 @@ from pathlib import Path -import torch import pytest +import torch from primus.backends.diffusion.models.registrations.wan import ( _convert_diffusers_wan_dit_state_dict, From 48ca729b072e8535fd0b09f0dad7e36ee5096072 Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 8 Jul 2026 04:09:53 +0000 Subject: [PATCH 3/4] docs(diffusion): update Wan A14B training plan --- wan2.2-t2v-14b-plan.md | 443 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 443 insertions(+) create mode 100644 wan2.2-t2v-14b-plan.md diff --git a/wan2.2-t2v-14b-plan.md b/wan2.2-t2v-14b-plan.md new file mode 100644 index 000000000..b2ffe68c5 --- /dev/null +++ b/wan2.2-t2v-14b-plan.md @@ -0,0 +1,443 @@ +# Wan2.2-T2V-A14B Training Support Plan + +## Goal + +Add Wan2.2-T2V-A14B training support on top of the Primus diffusion backend introduced by PR #779. + +The implementation should stay incremental and should validate that the current diffusion backend can be extended to a larger Wan variant without adding a separate training stack. + +## Base Branch Plan + +- Current base PR: https://github.com/AMD-AGI/Primus/pull/779 (`dev/zirui/wan`). +- Draft A14B PR: https://github.com/AMD-AGI/Primus/pull/861 (`dev/zirui/wan-A2.2-A14B`). +- Keep the A14B PR based on `dev/zirui/wan` while #779 is open so the diff remains small. +- After #779 lands in `main`, rebase or retarget the A14B branch to `main` before final review. + +## Proposed Scope + +### In Scope + +- Add model presets for Wan2.2-T2V-A14B. +- Add high-noise and low-noise SFT stage presets. +- Add MI355X example configs for both stages. +- Align stage behavior with Nvidia/NeMo AutoModel where practical: + - `high_noise`: train Diffusers `transformer` on sigma range `[boundary_ratio, 1.0]`. + - `low_noise`: train Diffusers `transformer_2` on sigma range `[0.0, boundary_ratio]`. + - Default `boundary_ratio=0.875`. + - Default flow shift `3.0`. +- Reuse existing Primus diffusion backend components: + - `WanArgBuilder`. + - `WanForTraining`. + - `WanDiT`. + - `FlowMatchScheduler`. + - FSDP2 trainer. + - Existing Wan dataset and processor. +- Add unit tests around config mapping and A14B-specific checkpoint key handling. + +### Out of Scope for the First Formal PR + +- Full Diffusers pipeline loader support. +- Native Diffusers VAE loading. +- Native Diffusers UMT5/text encoder loading. +- Automatic two-stage orchestration in a single Primus run. +- Generation/inference support for combining high-noise and low-noise trained checkpoints. +- Performance tuning or benchmark parity claims before E2E smoke tests pass. + +## Checkpoint Format Decision + +Wan2.2-T2V-A14B is published as a Diffusers pipeline with: + +- `model_index.json` +- `transformer/` +- `transformer_2/` +- `vae/` +- `text_encoder/` +- `tokenizer/` + +The current Primus diffusion backend uses native Wan-style assets for frozen encoders: + +- DiT `.safetensors` or `.bin` files. +- UMT5 encoder `.pth` file. +- Wan VAE `.pth` file. +- Tokenizer directory. + +For the initial A14B PR, keep the checkpoint-format scope intentionally narrow: + +- Support loading the A14B DiT from the Diffusers `transformer` or `transformer_2` subfolder. +- Do not require full Diffusers pipeline loading. +- Continue using native `.pth` UMT5/VAE assets for the frozen text encoder and VAE. +- Default the A14B presets to compatible Wan2.1-14B encoder assets: + - `/models/Wan2.1-T2V-14B/models_t5_umt5-xxl-enc-bf16.pth` + - `/models/Wan2.1-T2V-14B/Wan2.1_VAE.pth` +- Use the A14B Diffusers tokenizer subfolder: + - `/models/Wan2.2-T2V-A14B-Diffusers/tokenizer` + +Leave the broader checkpoint format decision as a follow-up TODO: + +- Option A: keep Primus training on native Wan-style `.pth` encoder assets and only add a lightweight Diffusers DiT key converter. +- Option B: add full Diffusers component loading for DiT, VAE, text encoder, and tokenizer. +- Option C: provide an offline conversion tool from Diffusers layout to Primus native Wan layout. + +The first formal PR should explicitly document whichever temporary behavior is used. + +## Current Draft Implementation Summary + +Draft PR #861 currently adds: + +- `primus/configs/models/diffusion/wan2.2_t2v_a14b.yaml` +- `primus/configs/models/diffusion/wan2.2_t2v_a14b_high_noise_sft.yaml` +- `primus/configs/models/diffusion/wan2.2_t2v_a14b_low_noise_sft.yaml` +- `examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml` +- `examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml` +- Diffusers Wan DiT key remapping in `primus/backends/diffusion/models/registrations/wan.py`. +- Scheduler override plumbing in `WanArgBuilder`. +- README documentation for the A14B two-stage setup. +- Focused unit tests for scheduler mapping, Diffusers subfolder selection, and key remapping. + +The key-level smoke check passed locally: + +- A14B Diffusers `transformer` state dict keys: 1095. +- Meta-initialized Primus A14B `WanDiT` state dict keys: 1095. +- Missing keys after remap: 0. +- Unexpected keys after remap: 0. + +Prepare-hook validation passed locally for both high-noise and low-noise example configs using: + +- `/mnt/shared/zirui/models/Wan2.2-T2V-A14B-Diffusers` +- `/mnt/shared/zirui/models/Wan2.1-T2V-14B/models_t5_umt5-xxl-enc-bf16.pth` +- `/mnt/shared/zirui/models/Wan2.1-T2V-14B/Wan2.1_VAE.pth` +- `/mnt/shared/zirui/data/tiny-video-samples` + +Execution update on 2026-07-07: + +- Fixed trainer-side `sigma_max` plumbing in the A14B worktree. +- Added a focused unit test for trainer scheduler `sigma_min`/`sigma_max` propagation. +- Ran focused Wan scheduler/builder tests: 9 passed. +- Ran full diffusion backend unit tests: 12 passed. +- Added `/zirui/code/Primus-wan-A2.2-A14B/local_runs/run.sh` as a local high/low/both E2E smoke launcher. +- Added rank-0 trainer logging for the resolved `FlowMatchScheduler` sigma range. +- Re-ran `python -m pytest tests/unit_tests/backends/diffusion`: 12 passed. +- Ran `PREPARE_ONLY=1 bash local_runs/run.sh both`; both high-noise and low-noise prepare hooks passed with `/models` and `/data` inputs. + +Execution update on 2026-07-08: + +- Ran a real single-node MI355X 8-GPU high-noise A14B training smoke on `chi2761`. +- Command used the A14B worktree, `GPUS_PER_NODE=8`, `SP_SIZE=1`, `MAX_STEPS=1`, `FRAME_NUM=1`, `HEIGHT=128`, `WIDTH=128`, `ATTENTION_BACKEND=flash_attn_aiter`, `VIDEO_BACKEND=decord`, and `BLAS_BACKEND=hipblas`. +- Data path on `chi2761`: `/zirui/data/tiny-video-samples/meta.jsonl` and `/zirui/data/tiny-video-samples/data`. +- Full A14B DiT load passed on all ranks with `files=12 missing=0 unexpected=0`. +- FSDP2 initialized with `DeviceMesh created: dims={'dp_shard': 8}`, wrapped 40 DiT blocks, and moved the sharded model to GPU after sharding. +- Training completed one forward/backward/optimizer step: `step=1 loss=0.5935`, `gnorm=19.2500`, `step_time=146.22s`. +- Final DiT save succeeded: `output/wan2.2_t2v_a14b-posttrain-high-noise/dit_model.safetensors` (about 27G). +- Log file: `/zirui/code/Primus-wan-A2.2-A14B/e2e_logs/high-noise-8gpu-tiny-sp1-aiter-hipblas-chi2761.log`. +- Runtime completed cleanly with `Training completed.` and `Cleanup completed.`; GPUs were released afterward. + +## Findings and Issues + +### 1. The Current Approach Does Extend the Diffusion Backend + +The draft implementation is based on the existing diffusion backend rather than a new backend: + +- It reuses `WanArgBuilder` for Primus-style config normalization. +- It reuses the existing Wan model registration path. +- It reuses `WanDiT` rather than importing Diffusers' runtime model. +- It reuses `WanFlowMatchTrainPipeline` and `FlowMatchScheduler`. +- It reuses FSDP2 trainer and the current dataset path. + +This is a good test of backend extensibility. + +### 2. The Code Is Mostly Lightweight, But Key Remapping Needs Care + +The main implementation is compact, but the Diffusers-to-Primus key remapping is string-replacement based. + +This is acceptable for a Draft PR if it remains tightly scoped to `WanTransformer3DModel` and is covered by tests, but formal PR should: + +- Name the helper clearly as Wan-specific, not generic Diffusers conversion. +- Keep full-key smoke tests or unit tests that protect the mapping. +- Fail loudly if missing/unexpected keys are non-zero for A14B loading. + +### 3. `sigma_max` Trainer Plumbing Was Fixed + +The A14B configs set stage-specific `sigma_min` and `sigma_max`. The trainer previously only passed these values to `FlowMatchScheduler` partially. + +Original issue: + +- `WanArgBuilder` maps `scheduler.sigma_max` into `trainer.args.flow_match_scheduler.sigma_max`. +- `BaseWanTrainer` currently passes `shift`, `sigma_min`, and `extra_one_step` into `FlowMatchScheduler`. +- `BaseWanTrainer` does not pass `sigma_max`. + +Fix landed in the A14B worktree: + +- `BaseWanTrainer` now passes `sigma_max=float(scheduler_cfg.get("sigma_max", 1.0))` into `FlowMatchScheduler`. +- Scheduler construction is covered by a focused trainer unit test. +- `python -m pytest tests/unit_tests/backends/diffusion` passes locally. + +Remaining validation: + +- Low-noise E2E smoke training still needs to verify that the stage loads `transformer_2` and samples from the intended sigma window. + +### 4. PR #861 CI State Needs Cleanup + +PR #861 had a CI `code-lint (3.12)` failure due to isort changing import order in `test_wan_model_registration.py`. + +Local worktree currently has an unpushed commit: + +- `25600096 fix isort issue` + +Before using #861 as a formal PR: + +- Push that commit, or squash it into the feature commit. +- Re-run CI and verify code-lint passes. + +### 5. High-Noise E2E Training Smoke Passed + +Completed local validation: + +- Unit tests pass locally. +- `python -m pytest tests/unit_tests/backends/diffusion` passed locally after the scheduler fix. +- `pre-commit run --all-files` passed locally. +- A14B high-noise prepare hook passed. +- A14B low-noise prepare hook passed. +- Full A14B key-remap smoke check passed. +- Local E2E runner exists and prepare-only validation passes for both stages. +- High-noise real 8-GPU A14B smoke passed on `chi2761` for one training step. + +Still not yet completed: + +- Low-noise real 8-GPU smoke test. +- Multi-step high-noise or low-noise smoke test. +- Larger video shape validation beyond the tiny `1x128x128` smoke. +- Any loss comparison with AutoModel. + +Until the low-noise stage also passes, claim only high-noise E2E smoke coverage, not full two-stage E2E support. + +## Implementation Plan + +### Phase 1: Stabilize the Minimal A14B Draft + +1. Keep PR #861 as a Draft PR based on `dev/zirui/wan`. +2. Done: fix scheduler plumbing: + - Pass `sigma_max` through `BaseWanTrainer` into `FlowMatchScheduler`. + - Add a focused unit test. +3. Done: run full diffusion backend unit tests locally. +4. Push or squash the isort-only commit and the scheduler fix. +5. Keep checkpoint support scoped to DiT Diffusers subfolders only. +6. Update PR description to say explicitly: + - Diffusers DiT loading is supported. + - Full Diffusers VAE/T5 loading is not yet supported. + - A14B presets currently use native Wan2.1-14B `.pth` encoder assets. +7. Add the 2026-07-08 `chi2761` high-noise 8-GPU smoke result to the PR notes. + +### Phase 2: E2E Smoke Validation + +AutoModel reference: + +- `/zirui/code/Automodel/examples/diffusion/finetune/wan2_2_t2v_flow.yaml` +- `/zirui/code/Automodel/run_wan2.2-t2v-a14b-benchmark.sh` + +AutoModel's A14B recipe assumes a single 8-GPU node for the benchmark path: + +- `GPUS_PER_NODE=8` +- `fsdp.dp_size=8` +- `activation_checkpointing=true` +- `local_batch_size=1` +- `global_batch_size=8` +- `flow_shift=3.0` +- `boundary_ratio=0.875` +- One stage per run: + - `high_noise`: active `transformer`, sigma range `[0.875, 1.0]` + - `low_noise`: active `transformer_2`, sigma range `[0.0, 0.875]` + +Primus should treat E2E as a single-node MI355X 8-GPU smoke first. A single MI355X GPU is not a good initial E2E target because `world_size=1` skips FSDP2 sharding in the current trainer; the 14B DiT is expected to need 8-way sharding plus activation checkpointing. + +#### E2E Preconditions + +Before launching training: + +1. Run on the A14B branch/worktree with the `sigma_max` trainer fix applied. +2. Confirm local inputs exist: + - `INIT_CHECKPOINT=/models/Wan2.2-T2V-A14B-Diffusers` + - `TEXT_ENCODER=/models/Wan2.1-T2V-14B/models_t5_umt5-xxl-enc-bf16.pth` + - `VAE_CHECKPOINT=/models/Wan2.1-T2V-14B/Wan2.1_VAE.pth` + - `TEXT_TOKENIZER=/models/Wan2.2-T2V-A14B-Diffusers/tokenizer` + - `DATASET_PATH=/data/tiny-video-samples/meta.jsonl` + - `DATA_FOLDER=/data/tiny-video-samples/data` + - On `chi2761`, the tiny sample data path used for the passing high-noise smoke was `/zirui/data/tiny-video-samples`. +3. Use `GPUS_PER_NODE=8`. +4. Start with `MAX_STEPS=1` or `2`. +5. Start with `FRAME_NUM=81`, `HEIGHT=480`, `WIDTH=832`; test `FRAME_NUM=121` only after the shorter smoke passes. +6. Use `ATTENTION_BACKEND=flash_attn_aiter` first. If kernel/runtime issues block launch, retry with `ATTENTION_BACKEND=sdpa` only as a debugging fallback. +7. Use `SP_SIZE=8` for the first 14B smoke if memory is the concern. A14B has 40 attention heads, so `SP_SIZE=4` and `SP_SIZE=8` both divide the head count. `SP_SIZE=1` is useful as a simpler code path but is likely less memory-friendly. +8. Use `BLAS_BACKEND=hipblas` on ROCm if hipBLASLt reports `HIPBLAS_STATUS_INTERNAL_ERROR` during backward. + +#### Local E2E Runner + +`/zirui/code/Primus-wan-A2.2-A14B/local_runs/run.sh` now wraps the two-stage +smoke path and defaults to the mounted `/models` and `/data` prefixes: + +```bash +cd /zirui/code/Primus-wan-A2.2-A14B +bash local_runs/run.sh both +``` + +Useful overrides: + +- `STAGE=high` or `STAGE=low` to run a single stage. +- `MAX_STEPS=2` or `5` after the 1-step smoke passes. +- `RUN_PREPARE=1` to run the posttrain diffusion prepare hook before `torchrun`. +- `ATTENTION_BACKEND=sdpa` as a debugging fallback if `flash_attn_aiter` blocks launch. +- `BLAS_BACKEND=hipblas` to force PyTorch BLAS away from hipBLASLt on ROCm. + +#### High-Noise Smoke Command + +Passing `chi2761` tiny smoke command from 2026-07-08: + +```bash +cd /zirui/code/Primus-wan-A2.2-A14B +mkdir -p e2e_logs +HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ +NCCL_IB_DISABLE=1 RCCL_IB_DISABLE=1 \ +BLAS_BACKEND=hipblas \ +STAGE=high GPUS_PER_NODE=8 SP_SIZE=1 \ +VIDEO_BACKEND=decord \ +DATASET_PATH=/zirui/data/tiny-video-samples/meta.jsonl \ +DATA_FOLDER=/zirui/data/tiny-video-samples/data \ +MAX_STEPS=1 FRAME_NUM=1 HEIGHT=128 WIDTH=128 \ +ATTENTION_BACKEND=flash_attn_aiter MASTER_PORT=29580 \ +bash local_runs/run.sh high 2>&1 | tee e2e_logs/high-noise-8gpu-tiny-sp1-aiter-hipblas-chi2761.log +``` + +Expected pass markers from that run: + +- `Init completed.` +- `Training started.` +- `step=1 loss=0.5935` +- `Saved DiT weights to ./output/wan2.2_t2v_a14b-posttrain-high-noise/dit_model.safetensors` +- `Training completed.` +- `Cleanup completed.` + +Larger-shape high-noise smoke target: + +```bash +NNODES=1 NODE_RANK=0 MASTER_ADDR=127.0.0.1 MASTER_PORT=29500 GPUS_PER_NODE=8 \ +INIT_CHECKPOINT=/models/Wan2.2-T2V-A14B-Diffusers \ +TEXT_ENCODER=/models/Wan2.1-T2V-14B/models_t5_umt5-xxl-enc-bf16.pth \ +VAE_CHECKPOINT=/models/Wan2.1-T2V-14B/Wan2.1_VAE.pth \ +TEXT_TOKENIZER=/models/Wan2.2-T2V-A14B-Diffusers/tokenizer \ +DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ +DATA_FOLDER=/data/tiny-video-samples/data \ +MAX_STEPS=1 FRAME_NUM=81 HEIGHT=480 WIDTH=832 SP_SIZE=8 ATTENTION_BACKEND=flash_attn_aiter \ +torchrun --nnodes=1 --node_rank=0 --master_addr=127.0.0.1 --master_port=29500 --nproc_per_node=8 \ + -m primus.cli.main train posttrain \ + --config examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml +``` + +Expected high-noise behavior: + +- Loads Diffusers `transformer`. +- Does not load/train `transformer_2`. +- Uses scheduler sigma range `[0.875, 1.0]`. +- Runs at least one forward/backward/optimizer step. +- Logs finite loss. + +#### Low-Noise Smoke Command + +Use the same environment and command shape, but switch the config: + +```bash +NNODES=1 NODE_RANK=0 MASTER_ADDR=127.0.0.1 MASTER_PORT=29501 GPUS_PER_NODE=8 \ +INIT_CHECKPOINT=/models/Wan2.2-T2V-A14B-Diffusers \ +TEXT_ENCODER=/models/Wan2.1-T2V-14B/models_t5_umt5-xxl-enc-bf16.pth \ +VAE_CHECKPOINT=/models/Wan2.1-T2V-14B/Wan2.1_VAE.pth \ +TEXT_TOKENIZER=/models/Wan2.2-T2V-A14B-Diffusers/tokenizer \ +DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ +DATA_FOLDER=/data/tiny-video-samples/data \ +MAX_STEPS=1 FRAME_NUM=81 HEIGHT=480 WIDTH=832 SP_SIZE=8 ATTENTION_BACKEND=flash_attn_aiter \ +torchrun --nnodes=1 --node_rank=0 --master_addr=127.0.0.1 --master_port=29501 --nproc_per_node=8 \ + -m primus.cli.main train posttrain \ + --config examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml +``` + +Expected low-noise behavior: + +- Loads Diffusers `transformer_2`. +- Does not load/train `transformer`. +- Uses scheduler sigma range `[0.0, 0.875]`. +- Runs at least one forward/backward/optimizer step. +- Logs finite loss. + +#### E2E Acceptance Criteria + +Minimum smoke pass: + +- High-noise job launches on one 8-GPU MI355X node. Done on `chi2761` with the tiny shape. +- Low-noise job launches on one 8-GPU MI355X node. +- Both jobs load full A14B DiT weights with zero missing/unexpected keys after conversion. Done for high-noise. +- Both jobs complete at least one train step. Done for high-noise. +- Loss is finite for both jobs. Done for high-noise. +- Logs show the intended active transformer and sigma range. Done for high-noise. +- No OOM, distributed hang, or FSDP2 wrapping error. Done for high-noise. + +Stronger validation after smoke: + +- Increase `MAX_STEPS` to `2` or `5`. +- Retry with `FRAME_NUM=81`, then `FRAME_NUM=121`, to match AutoModel benchmark shape more closely. +- Compare Primus loss scale/trend against AutoModel on the same tiny-video samples, without claiming benchmark parity. +- Enable a small checkpoint save and verify the `dit_only` output can be reloaded. + +#### Failure Triage Order + +1. If launch fails before model load, check environment and dependency/kernel availability. +2. If model load fails, inspect Diffusers key remapping and missing/unexpected keys. +3. If OOM occurs, keep `SP_SIZE=8`, `FRAME_NUM=81`, `local_batch_size=1`, `MAX_STEPS=1`, and confirm gradient checkpointing is enabled. +4. If attention kernels fail, retry with `ATTENTION_BACKEND=sdpa` to separate kernel issues from model/trainer issues. +5. If ROCm GEMM fails with `HIPBLAS_STATUS_INTERNAL_ERROR`, retry with `BLAS_BACKEND=hipblas` to avoid hipBLASLt. +6. If distributed training hangs, retry with a fresh `MASTER_PORT`, verify `GPUS_PER_NODE=8`, and check whether all ranks reached FSDP2 wrapping. +7. If loss is NaN/Inf, use `FIXED_SEED` and `FIXED_TIMESTEP` for a reproducible debug run. + +### Phase 3: Decide Checkpoint Format Direction + +After minimal E2E passes, decide one of: + +1. Keep hybrid support: + - Diffusers DiT subfolder loading. + - Native `.pth` frozen encoder assets. +2. Add full Diffusers loading: + - Load Diffusers VAE. + - Load Diffusers UMT5/text encoder. + - Keep tokenizer from Diffusers path. +3. Add conversion tooling: + - Convert A14B Diffusers checkpoint into native Primus/Wan layout. + - Keep training runtime simple. + +This decision should be made before marking the PR ready for final review. + +## Formal PR Checklist + +Before converting the A14B Draft PR to a formal PR: + +- [ ] Rebase or retarget from `dev/zirui/wan` to `main` after PR #779 lands. +- [x] Fix `sigma_max` scheduler plumbing. +- [ ] Push/squash the isort fix. +- [ ] Ensure CI code-lint is green. +- [ ] Run `pre-commit run --all-files`. +- [x] Run `python -m pytest tests/unit_tests/backends/diffusion`. +- [x] Add or keep a unit test for scheduler `sigma_min/sigma_max` propagation. +- [x] Keep key-remap unit tests. +- [ ] Add a model-build smoke test if practical without loading full weights in CI. +- [x] Add a local A14B E2E smoke runner. +- [x] Run local high-noise prepare hook. +- [x] Run local low-noise prepare hook. +- [x] Run high-noise E2E smoke training. +- [ ] Run low-noise E2E smoke training. +- [ ] Document checkpoint format limitations clearly. +- [x] Document exact model/data environment variables for both stages. +- [ ] Avoid claiming benchmark/performance parity before benchmark runs exist. + +## Open TODOs + +- Decide whether A14B should support full Diffusers pipeline loading or remain hybrid. +- Decide whether to add a checkpoint conversion tool. +- Decide whether Primus should orchestrate high-noise and low-noise stages in one top-level config or keep them as separate explicit runs. +- Validate that using Wan2.1-14B VAE/T5 assets with Wan2.2-T2V-A14B DiT is acceptable for training quality, not just code execution. +- Compare loss behavior against AutoModel after E2E smoke succeeds. +- Add generation/inference support for combining two stage checkpoints if needed. From 59046e21fe7a094dfc5a47f687b374c9a98d1815 Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 8 Jul 2026 04:40:42 +0000 Subject: [PATCH 4/4] fix(diffusion): unblock Wan A14B smoke training Co-authored-by: Cursor --- .../MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml | 2 +- .../MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml | 2 +- primus/backends/diffusion/diffusion_pretrain_trainer.py | 8 ++++++++ primus/backends/diffusion/trainers/fsdp2.py | 6 +++++- 4 files changed, 15 insertions(+), 3 deletions(-) diff --git a/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml b/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml index 4984e3fcc..bb38e18dd 100644 --- a/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml +++ b/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-high-noise.yaml @@ -37,7 +37,7 @@ modules: dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} frame_num: ${FRAME_NUM:81} - video_backend: imageio + video_backend: ${VIDEO_BACKEND:decord} text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.2-T2V-A14B-Diffusers/tokenizer} height: ${HEIGHT:480} width: ${WIDTH:832} diff --git a/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml b/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml index f1db23fed..ab01703a1 100644 --- a/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml +++ b/examples/diffusion/configs/MI355X/wan2.2_t2v_a14b-posttrain-low-noise.yaml @@ -37,7 +37,7 @@ modules: dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} frame_num: ${FRAME_NUM:81} - video_backend: imageio + video_backend: ${VIDEO_BACKEND:decord} text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.2-T2V-A14B-Diffusers/tokenizer} height: ${HEIGHT:480} width: ${WIDTH:832} diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index ddbb70d49..fbe2519b5 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -7,6 +7,7 @@ from __future__ import annotations import importlib.util +import os from typing import Any from primus.core.trainer.base_trainer import BaseTrainer @@ -32,6 +33,13 @@ def setup(self): dataset_cfg = self._as_dict(self.backend_args.dataset) trainer_args = trainer_cfg.get("args", {}) attention_backend = trainer_args.get("attention_backend") + blas_backend = os.environ.get("BLAS_BACKEND") + + if blas_backend: + import torch + + torch.backends.cuda.preferred_blas_library(blas_backend) + log_rank_0(f"[Primus:Diffusion] BLAS_BACKEND={blas_backend}") missing = [ package diff --git a/primus/backends/diffusion/trainers/fsdp2.py b/primus/backends/diffusion/trainers/fsdp2.py index 1fa79b177..84261ab41 100644 --- a/primus/backends/diffusion/trainers/fsdp2.py +++ b/primus/backends/diffusion/trainers/fsdp2.py @@ -80,7 +80,6 @@ def _apply_parallelism(self): 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"): @@ -89,6 +88,11 @@ def _apply_parallelism(self): logger.info("FSDP2: Applied freeze_except (frozen non-trainable params)") self._apply_fsdp2() + if self.rank == 0: + logger.info(f"FSDP2: moving sharded model to {self.device}") + self.model.to(self.device) + if self.rank == 0: + logger.info(f"FSDP2: sharded model moved to {self.device}") def _apply_fsdp2(self): """Apply torch.distributed._composable.fsdp.fully_shard to the model."""