feat(backend): add NeMo AutoModel (Wan 2.2 diffusion) as a first-class Primus backend

.gitmodules

@@ -19,3 +19,7 @@
 [submodule "third_party/HummingbirdXT"]
 	path = third_party/HummingbirdXT
 	url = https://github.com/AMD-AGI/HummingbirdXT.git
+[submodule "third_party/Automodel"]
+	path = third_party/Automodel
+	url = https://github.com/NVIDIA-NeMo/Automodel.git
+	branch = main

examples/README.md

@@ -36,10 +36,11 @@ It supports both **single-node** and **multi-node** training, and includes optio
 
 Primus supports multiple backends.
 
-| Backend    | Description                                                  |
-| ---------- | ------------------------------------------------------------ |
-| Megatron   | Open-source framework for large-scale transformer training   |
-| TorchTitan | PyTorch-compatible framework developed for training at scale |
+| Backend        | Description                                                  |
+| -------------- | ------------------------------------------------------------ |
+| Megatron       | Open-source framework for large-scale transformer training   |
+| TorchTitan     | PyTorch-compatible framework developed for training at scale |
+| NeMo AutoModel | NVIDIA-NeMo AutoModel (diffusion: Wan 2.2 T2V); `third_party/Automodel` submodule, installed editable on first run |
 
 
 ## 🖥️ Single Node Training

examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml

@@ -0,0 +1,71 @@
+# Primus EXP — Wan 2.2 T2V A14B FINETUNE via the NeMo AutoModel backend on MI355X.
+# =============================================================================
+# Full finetune of the pretrained Wan 2.2-A14B weights (loaded from
+# model.pretrained_model_name_or_path via from_pretrained). Differs from the
+# pretrain EXPs by selecting AutoModel's finetune path:
+#   model.mode: finetune        -> build_model_and_optimizer(finetune_mode=True);
+#                                  pipeline_spec is ignored in this path.
+#   model.torch_dtype: bfloat16 -> bf16 master weights (canonical finetune
+#                                  precision; pretrain uses fp32 master).
+#   gentle LR (5e-6) + standard AdamW betas/decay (foreach, not fused).
+#
+# Runs through the SAME diffusion recipe and the SAME launcher as pretrain
+# (AutoModel selects pretrain vs finetune from model.mode), so it launches via
+# `examples/run_pretrain.sh` like the other diffusion EXPs.
+#
+# NOTE: the finetune recipe is AutoModel's original Wan 2.2 path, provided here
+# as a complete, first-class example. The throughput numbers documented in the
+# pretrain EXPs were measured on the pretrain operating points.
+#
+# One-time setup (like any Primus backend): initialize the submodule; the
+# launcher's prepare step installs it editable (ROCm-pinned) on first run.
+#   git submodule update --init third_party/Automodel
+#   # (auto-skipped if the image already ships nemo_automodel)
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml \
+#   bash examples/run_pretrain.sh
+#
+# Run once per expert. The expert is chosen solely by model.stage (finetune loads
+# via from_pretrained with the matching active_transformer; pipeline_spec/subfolder
+# is ignored in this path -- it only applies to pretraining's from_config):
+#   high_noise (default) -> transformer;  model.stage=low_noise -> transformer_2.
+# Remove step_scheduler.max_steps to finetune for full epochs.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2_2_t2v_a14b-finetune}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: pre_trainer.yaml
+    model: wan2_2_t2v_a14b.yaml
+
+    overrides:
+      model:
+        mode: finetune
+        torch_dtype: bfloat16    # bf16 master weights for finetuning
+      step_scheduler:
+        max_steps: 25            # cap for a quick test; remove for a real finetune
+        local_batch_size: 8
+        global_batch_size: 64    # 8 x dp8
+      optim:
+        learning_rate: 5e-6      # gentle finetune LR
+        optimizer:
+          weight_decay: 0.01
+          betas: [0.9, 0.999]
+          # finetune uses foreach AdamW (fused is the pretrain throughput knob)
+          foreach: true
+          fused: false
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /dataset/pusa_v1_512sq
+      checkpoint:
+        enabled: true
+        checkpoint_dir: /checkpoints/wan22_finetune_high
+        restore_from: null       # null -> finetune the base pretrained weights

examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml

@@ -0,0 +1,52 @@
+# Primus EXP — Wan 2.2 T2V A14B PRETRAIN, optimized AC-on bs32.
+# =============================================================================
+# This is the "best AC-on" benchmark point reproduced end-to-end through the
+# full Primus runtime (NeMo AutoModel backend) on 8x MI355X. It layers the
+# optimized knobs (already baked into the pre_trainer module config:
+# torch.compile + fused AdamW + bf16 grad-reduce + deeper FSDP2 prefetch, AITER
+# attention, fp32 master / bf16 compute) on top of the Wan 2.2-A14B preset.
+#
+# Measured (8x MI355X, 512x512, real PusaV1, fp32 master / bf16 compute):
+#   ~3.10-3.12 samples/s/GPU (~24.9 total) at mem ~84.9 GB
+#   == +~46% vs the un-optimized baseline (2.14 s/s/GPU). torch.compile runs
+#      cleanly with the AITER custom-op backend (no graph break).
+#
+# One-time setup (like any Primus backend): initialize the submodule; the
+# launcher's prepare step installs it editable (ROCm-pinned) on first run.
+#   git submodule update --init third_party/Automodel
+#   # (auto-skipped if the image already ships nemo_automodel)
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml \
+#   bash examples/run_pretrain.sh
+#
+# low-noise expert: override model.stage=low_noise and
+#   model.pipeline_spec.subfolder=transformer_2.
+# Remove step_scheduler.max_steps to train full epochs instead of a benchmark cap.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2_2_t2v_a14b-pretrain-bs32-acon}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: pre_trainer.yaml
+    model: wan2_2_t2v_a14b.yaml
+
+    overrides:
+      # Headline AC-on plateau point. activation_checkpointing stays TRUE
+      # (the pre_trainer module default) so bs32 fits with headroom.
+      step_scheduler:
+        max_steps: 25            # benchmark cap; remove for a real pretrain run
+        local_batch_size: 32
+        global_batch_size: 256   # 32 x dp8
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /dataset/pusa_v1_512sq
+      checkpoint:
+        enabled: false           # set true + checkpoint_dir for a real run

examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml

@@ -0,0 +1,52 @@
+# Primus EXP — Wan 2.2 T2V A14B PRETRAIN, optimized AC-off bs8.
+# =============================================================================
+# Fastest measured operating point. With torch.compile cutting activation
+# memory, activation checkpointing can be turned OFF at bs8 (which OOM'd in the
+# un-optimized baseline), giving the best throughput. This mirrors the Wan 2.1
+# GB300 small-batch / AC-off reference setup.
+#
+# Measured (8x MI355X, 512x512, real PusaV1, fp32 master / bf16 compute):
+#   ~3.68-3.69 samples/s/GPU (~29.4 total) at mem ~236.8 GB (AC OFF, bs8)
+#   == new best point; un-optimized AC-off OOM'd at bs8 (ceiling was bs6 @ ~240 GB).
+# (bs6 AC-off, the more conservative GB300-comparable point, measures
+#  ~3.54 s/s/GPU @ ~183.5 GB — set local_batch_size=6 / global_batch_size=48.)
+#
+# One-time setup (like any Primus backend): initialize the submodule; the
+# launcher's prepare step installs it editable (ROCm-pinned) on first run.
+#   git submodule update --init third_party/Automodel
+#   # (auto-skipped if the image already ships nemo_automodel)
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml \
+#   bash examples/run_pretrain.sh
+#
+# low-noise expert: override model.stage=low_noise and
+#   model.pipeline_spec.subfolder=transformer_2.
+# Remove step_scheduler.max_steps to train full epochs instead of a benchmark cap.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2_2_t2v_a14b-pretrain-bs8-acoff}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: pre_trainer.yaml
+    model: wan2_2_t2v_a14b.yaml
+
+    overrides:
+      # AC-OFF best point: compile buys enough memory to drop activation
+      # checkpointing at bs8 (~237 GB of the 288 GB HBM3e).
+      step_scheduler:
+        max_steps: 25            # benchmark cap; remove for a real pretrain run
+        local_batch_size: 8
+        global_batch_size: 64    # 8 x dp8
+      fsdp:
+        activation_checkpointing: false
+      data:
+        dataloader:
+          cache_dir: /dataset/pusa_v1_512sq
+      checkpoint:
+        enabled: false           # set true + checkpoint_dir for a real run

examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml

@@ -0,0 +1,43 @@
+# Primus EXP config — Wan 2.2 T2V A14B pretrain via the NeMo AutoModel backend
+# on MI355X. Selects framework: nemo_automodel, consistent
+# with how the megatron diffusion FLUX EXPs select framework: megatron.
+#
+# One-time setup (like any other Primus backend): initialize the submodule. The
+# launcher's prepare step then installs it editable (ROCm-pinned) on first run.
+#   git submodule update --init third_party/Automodel
+#   # (skipped automatically if the image already ships nemo_automodel)
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml \
+#   bash examples/run_pretrain.sh
+#   # AC-off operating points (Wan 2.1 GB300-comparable) via CLI overrides, e.g.:
+#   #   ... bash examples/run_pretrain.sh \
+#   #       fsdp.activation_checkpointing=false \
+#   #       step_scheduler.local_batch_size=8 step_scheduler.global_batch_size=64
+#
+# low-noise expert: override model.stage=low_noise and
+#   model.pipeline_spec.subfolder=transformer_2.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2_2_t2v_a14b-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: pre_trainer.yaml
+    model: wan2_2_t2v_a14b.yaml
+
+    overrides:
+      # Throughput-benchmark sizing (optimized headline point: bs32 AC-on).
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 32
+        global_batch_size: 256
+      data:
+        dataloader:
+          cache_dir: /dataset/pusa_v1_512sq
+      checkpoint:
+        enabled: false