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

[mvstrauss]

Summary

Adds NVIDIA-NeMo AutoModel as a first-class Primus training backend, wrapping
AutoModel's diffusion TrainDiffusionRecipe (Wan 2.2 T2V A14B) behind the Primus
BackendAdapter/Trainer contract. It follows the exact same integration pattern as
the existing torchtitan/megatron backends — a pinned third_party/ submodule that
is installed editable on first run — so a stock rocm/primus user runs it via the
standard git submodule update --init + run_pretrain.sh flow.

What's included

• Backend (primus/backends/nemo_automodel/): adapter, pretrain_trainer,
  argument_builder, and registry wiring (__init__.py), mirroring torchtitan.
  The trainer converts Primus module_config.params → AutoModel ConfigNode and
  delegates to TrainDiffusionRecipe (AutoModel owns FSDP2 / dataloader / optimizer /
  checkpointing), keeping us version-agnostic.
• Submodule: third_party/Automodel pinned to 44f2acde (.gitmodules,
  branch = main), and primus/_thirdparty.lock regenerated from the submodule pins
  via tools/pip_index/gen_thirdparty_lock.py (the thirdparty-lock pre-commit hook).
• examples/nemo_automodel/prepare.py: resolves the submodule and does a
  ROCm-safe editable install — pins the image's
  torch/torchvision/torchaudio/triton/aiter/flash-attn via a pip constraint so CUDA
  wheels can't be pulled — with the diffusion,diffusion-media extras. Idempotent
  (skips if nemo_automodel is already importable). Also applies an optional, opt-out
  AITER version shim so diffusers' aiter attention backend is selectable on
  images that ship a dev-versioned amd-aiter (disable with AITER_VERSION_SHIM=0).
• Configs:
  • Model preset primus/configs/models/nemo_automodel/wan2_2_t2v_a14b.yaml
  • Module config primus/configs/modules/nemo_automodel/pre_trainer.yaml
    (optimized: AITER attention + torch.compile + fused AdamW + bf16 grad-reduce
    • deeper FSDP2 prefetch, fp32 master / bf16 compute)
  • MI355X example EXPs (examples/nemo_automodel/configs/MI355X/diffusion/):
    -pretrain-bs32-acon (headline AC-on), -pretrain-bs8-acoff (fastest AC-off),
    -finetune (mode: finetune, bf16 master), and a minimal -pretrain stub.
• Tests (tests/unit_tests/backends/nemo_automodel/): registration, a
  config-strip guard (fails if module_base.yaml grows a key that would leak into the
  AutoModel config), and unsupported-stage handling.
• Docs: NeMo AutoModel row in examples/README.md.

Measured results (8× MI355X, 512×512, real PusaV1, fp32 master / bf16 compute)

Reproduced end-to-end through the full Primus runtime (run_pretrain.sh →
prepare_experiment.py → torchrun … train pretrain → PrimusRuntime → backend →
AutoModel recipe):

Operating point	Config	samples/s/GPU	peak mem	notes
pretrain bs32, AC on	wan2_2_t2v_a14b-pretrain-bs32-acon	~3.12	~84.9 GB	headline plateau (+~46% vs baseline)
pretrain bs8, AC off	wan2_2_t2v_a14b-pretrain-bs8-acoff	~3.69	~236.8 GB	fastest point (was OOM un-optimized)

torch.compile runs cleanly with the AITER custom-op backend (no graph break).

Fresh-environment reproduction (this branch)

The headline point was re-verified from scratch: a fresh rocm/primus:v26.3
container + a fresh clone of this branch (only git submodule update --init and
the base image), no pre-existing state. Steady-state (bs32 AC-on, 8× MI355X):

[TRAIN] step=1 loss=15.99 step_time=10.245s samples_per_sec=24.99 samples_per_sec_per_gpu=3.12 mem=84.95GB
[TRAIN] step=2 loss=7.18  step_time=10.209s samples_per_sec=25.07 samples_per_sec_per_gpu=3.13 mem=84.95GB
[TRAIN] step=3 loss=3.94  step_time=10.255s samples_per_sec=24.96 samples_per_sec_per_gpu=3.12 mem=84.95GB

(step 0 is a ~116 s warmup: torch.compile + first-time AITER gfx950 kernel JIT.)
Matches the documented ~3.10–3.12 s/s/GPU @ ~84.9 GB; run exited 0 with a final
checkpoint written.

How to reproduce on a fresh node

Everything below needs only this branch + the rocm/primus base image + public
Hugging Face artifacts. Adjust host paths to your machine.

1. Container

/data/primus is an empty host directory (the branch is cloned into it in step 2);
/data/models, /data/wan22_cache, /data/hf are roomy host paths for weights, the
dataset cache, and the HF cache.

docker run -d --name automodel_test \
  --network host --ipc host --privileged \
  --security-opt seccomp=unconfined --security-opt label=disable \
  --device /dev/kfd --device /dev/dri --shm-size 64g \
  -v /data/primus:/workspace/primus:rw \
  -v /data/models:/models:rw \
  -v /data/wan22_cache:/dataset:rw \
  -v /data/hf:/hf:rw \
  -e HF_HOME=/hf -w /workspace/primus \
  rocm/primus:v26.3 sleep infinity
docker exec -it automodel_test bash

2. Branch + submodule (inside the container)

git clone -b feat/backend/nemo_automodel https://github.com/AMD-AGI/Primus.git /workspace/primus
cd /workspace/primus
git submodule update --init third_party/Automodel   # pinned to 44f2acde

prepare.py does the editable, ROCm-pinned install automatically on the first
run_pretrain.sh; nothing else to install by hand.

3. Model weights (public)

pip install -U "huggingface_hub[cli]"
huggingface-cli download Wan-AI/Wan2.2-T2V-A14B-Diffusers \
  --local-dir /models/Wan2.2-T2V-A14B-Diffusers

(The configs point model.pretrained_model_name_or_path at
/models/Wan2.2-T2V-A14B-Diffusers.)

4. Dataset — build the preprocessed cache

The recipe trains on a preprocessed latent/embedding cache produced by AutoModel's
own preprocessor (tools.diffusion.preprocessing_multiprocess, shipped in the
submodule). Two options:

(a) Real data — RaphaelLiu/PusaV1_training (public, ~79 GB / ~3,860 clips):

# download
huggingface-cli download RaphaelLiu/PusaV1_training --repo-type dataset \
  --local-dir /data/PusaV1_training

# captions: metadata.csv (file_name,text) -> train/meta.json  (self-contained)
python3 - <<'PY'
import csv, json, os
root="/data/PusaV1_training"
rows=[{"file_name":r["file_name"],"caption":r["text"]}
      for r in csv.DictReader(open(os.path.join(root,"metadata.csv")))]
json.dump(rows, open(os.path.join(root,"train","meta.json"),"w"))
print("wrote", len(rows), "captions")
PY

# encode -> AutoModel .meta cache (GPU: Wan VAE + UMT5). Writes into /dataset.
cd /workspace/primus/third_party/Automodel
python -m tools.diffusion.preprocessing_multiprocess video \
  --video_dir /data/PusaV1_training/train \
  --output_dir /dataset/pusa_v1_512sq \
  --processor wan2.2 --model_name /models/Wan2.2-T2V-A14B-Diffusers \
  --resolution_preset 512p --mode video \
  --caption_format meta_json --caption_field caption

(b) Throughput-only smoke: point the preprocessor at any small folder of 512p
.mp4 clips with sidecar .txt captions (--caption_format sidecar); a dozen clips
are enough for a steady-state throughput number.

Either way, set the EXP's data.dataloader.cache_dir to the produced cache
(/dataset/pusa_v1_512sq in the examples).

5. Launch

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

For a quick smoke, override step_scheduler.max_steps to a few steps. On non-8-GPU
topologies set GPUS_PER_NODE and fsdp.dp_size to your world size.

Test plan

• End-to-end from a fresh container + fresh clone (see "Fresh-environment
  reproduction" above): submodule init → ROCm-safe editable install (ROCm torch
  preserved, no CUDA wheels) → AITER shim → config parse → backend routing → 8×
  MI355X training to steady state, exit 0 + final checkpoint written.
• pre-commit clean on all changed files (trailing-whitespace / eof / check-yaml /
  isort / autoflake / black / shellcheck / thirdparty-lock).
• Unit tests: tests/unit_tests/backends/nemo_automodel/ — 10 passed (adapter +
  trainer registration, train→pretrain alias, config-strip guard,
  unsupported-stage error).
• All example EXPs parse via PrimusParser; overrides land (mode/precision/AC/
  batch/optimizer).
• CI green on the PR. The run-unittest-torch job checks out submodules recursively
  and runs the change-selected pytest targets, so it will exercise the new
  tests/unit_tests/backends/nemo_automodel/ suite. Left unchecked because it can
  only run once the PR triggers CI (these tests already pass locally, above).

Notes for reviewers

• AutoModel is not in the stock base image; the submodule + prepare.py editable
  install is the supported path. The adapter keeps a "prefer already-installed" fast
  path for images that do ship it.
• Expert selection (Wan 2.2-A14B's two high/low-noise denoising experts) is by model.stage
  (high_noise → transformer, low_noise → transformer_2). pipeline_spec is only
  consumed in mode: pretrain; it is ignored in mode: finetune (which loads via
  from_pretrained with the matching active_transformer).
• The MI355X example configs carry AMD-specific env/benchmark commentary; happy to trim
  or relocate if you'd prefer the examples stay hardware-agnostic.

[Copilot]

Pull request overview

Adds a new first-class Primus backend that routes Primus’ BackendAdapter/Trainer contract into NVIDIA NeMo AutoModel’s diffusion training recipe (Wan 2.2 T2V A14B), following the existing “third_party submodule + editable install on first run” integration pattern used by other backends.

Changes:

• Introduces the primus/backends/nemo_automodel adapter, trainer, config builder, and registry wiring for backend discovery and trainer resolution.
• Adds model/module presets for Wan 2.2 diffusion and MI355X example experiment configs (pretrain + tuned operating points + finetune).
• Pins the Automodel submodule in .gitmodules and primus/_thirdparty.lock, plus adds a backend-specific examples/nemo_automodel/prepare.py for ROCm-safe editable installation (and optional aiter version shim).

Reviewed changes

Copilot reviewed 15 out of 16 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/unit_tests/backends/nemo_automodel/test_nemo_automodel_registration.py	Unit tests validating backend/adapter/trainer registration and a config-strip guard.
primus/backends/nemo_automodel/init.py	Registers the adapter and trainer class in BackendRegistry.
primus/backends/nemo_automodel/nemo_automodel_adapter.py	BackendAdapter implementation: backend path resolution, config pass-through, trainer resolution, version detection.
primus/backends/nemo_automodel/nemo_automodel_pretrain_trainer.py	Trainer that converts params → temp YAML → AutoModel config loader → TrainDiffusionRecipe.
primus/backends/nemo_automodel/argument_builder.py	Config normalization helpers and Primus-only top-level key stripping before handing off to AutoModel.
primus/configs/modules/nemo_automodel/pre_trainer.yaml	AutoModel diffusion recipe-shaped module config defaults for Primus pre_trainer.
primus/configs/models/nemo_automodel/wan2_2_t2v_a14b.yaml	Wan 2.2 T2V A14B preset merged into the module config under model:.
examples/nemo_automodel/prepare.py	Backend prepare step: editable install with ROCm constraints + optional amd-aiter metadata shim.
examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml	Baseline MI355X pretrain experiment wiring for the new backend.
examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml	MI355X tuned “AC-on bs32” benchmark config overrides.
examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml	MI355X tuned “AC-off bs8” benchmark config overrides.
examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml	MI355X finetune experiment config (uses model.mode: finetune).
examples/README.md	Documents NeMo AutoModel in the supported backends table.
.gitmodules	Adds third_party/Automodel submodule entry.
primus/_thirdparty.lock	Records the Automodel submodule pin for dependency locking.

[mvstrauss]

This PR adds NeMo AutoModel as first-class Primus backend. Initially added explicit configs for BF16 pretrain of Wan2.2, per request from DCGPU team on behalf of customer. We will later extend this with FP8 and MXFP4 support.