Add Wan2.2 T2V-A14B diffusion training presets

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: ${VIDEO_BACKEND:decord}
+        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

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: ${VIDEO_BACKEND:decord}
+        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

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**

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)

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

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,13 +223,28 @@ 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():
                 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
+            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()

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."""