From 48c478b8b8882dbc48f3ab298d4fd5b9ef702f43 Mon Sep 17 00:00:00 2001 From: mvstrauss Date: Wed, 1 Jul 2026 08:54:09 -0500 Subject: [PATCH 1/2] feat(backend): add NeMo AutoModel as a first-class Primus backend Wraps NVIDIA-NeMo AutoModel's diffusion TrainDiffusionRecipe (Wan 2.2 T2V A14B) behind the Primus BackendAdapter/Trainer contract, following the same pattern as the torchtitan/megatron backends: a pinned third_party/ submodule installed editable on first run. - primus/backends/nemo_automodel/: adapter, pretrain trainer, argument builder, and registry wiring (mirrors torchtitan/megatron). - third_party/Automodel submodule pinned to 44f2acde; primus/_thirdparty.lock regenerated from the submodule pins. - examples/nemo_automodel/prepare.py: 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) + optional aiter version shim so diffusers' aiter attention backend is selectable; idempotent. - configs: Wan 2.2-A14B model preset + optimized nemo_automodel pre_trainer module config, plus MI355X example EXPs (pretrain bs32 AC-on, bs8 AC-off, finetune, minimal default). - tests/unit_tests/backends/nemo_automodel: registration + config-strip guard + unsupported-stage tests. - docs: NeMo AutoModel row in examples/README.md. Co-authored-by: Cursor --- .gitmodules | 4 + examples/README.md | 9 +- .../diffusion/wan2_2_t2v_a14b-finetune.yaml | 71 ++++ .../wan2_2_t2v_a14b-pretrain-bs32-acon.yaml | 52 +++ .../wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml | 52 +++ .../diffusion/wan2_2_t2v_a14b-pretrain.yaml | 43 +++ examples/nemo_automodel/prepare.py | 324 ++++++++++++++++++ primus/_thirdparty.lock | 6 + primus/backends/nemo_automodel/__init__.py | 17 + .../nemo_automodel/argument_builder.py | 100 ++++++ .../nemo_automodel/nemo_automodel_adapter.py | 120 +++++++ .../nemo_automodel_pretrain_trainer.py | 92 +++++ .../nemo_automodel/wan2_2_t2v_a14b.yaml | 23 ++ .../modules/nemo_automodel/pre_trainer.yaml | 98 ++++++ .../test_nemo_automodel_registration.py | 233 +++++++++++++ third_party/Automodel | 1 + 16 files changed, 1241 insertions(+), 4 deletions(-) create mode 100644 examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml create mode 100644 examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml create mode 100644 examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml create mode 100644 examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml create mode 100644 examples/nemo_automodel/prepare.py create mode 100644 primus/backends/nemo_automodel/__init__.py create mode 100644 primus/backends/nemo_automodel/argument_builder.py create mode 100644 primus/backends/nemo_automodel/nemo_automodel_adapter.py create mode 100644 primus/backends/nemo_automodel/nemo_automodel_pretrain_trainer.py create mode 100644 primus/configs/models/nemo_automodel/wan2_2_t2v_a14b.yaml create mode 100644 primus/configs/modules/nemo_automodel/pre_trainer.yaml create mode 100644 tests/unit_tests/backends/nemo_automodel/test_nemo_automodel_registration.py create mode 160000 third_party/Automodel diff --git a/.gitmodules b/.gitmodules index 1d234e782..134e4be84 100644 --- a/.gitmodules +++ b/.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 diff --git a/examples/README.md b/examples/README.md index ce0e36460..97049c744 100644 --- a/examples/README.md +++ b/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 diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml new file mode 100644 index 000000000..0fe2884bd --- /dev/null +++ b/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 diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml new file mode 100644 index 000000000..147275927 --- /dev/null +++ b/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 diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml new file mode 100644 index 000000000..e5b3eeb84 --- /dev/null +++ b/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 diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml new file mode 100644 index 000000000..dd19f295b --- /dev/null +++ b/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 diff --git a/examples/nemo_automodel/prepare.py b/examples/nemo_automodel/prepare.py new file mode 100644 index 000000000..d723ddf77 --- /dev/null +++ b/examples/nemo_automodel/prepare.py @@ -0,0 +1,324 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Backend prepare entry for the NeMo AutoModel framework. + +Mirrors ``examples/torchtitan/prepare.py``: it resolves the backend checkout +(the ``third_party/Automodel`` submodule by default) and installs it editable so +the ``nemo_automodel`` package is importable for the training phase. Two +AutoModel-specific differences from the torchtitan flow: + + 1. **ROCm-safe install.** AutoModel's dependency graph would otherwise let pip + pull CUDA torch wheels on top of the image's ROCm build. We pin the native + ROCm packages (torch/torchvision/torchaudio/triton/aiter/flash-attn) to the + versions already installed and pass them as a pip *constraint*, so the + editable install never replaces the ROCm stack. The ``diffusion`` extras + are installed for the Wan diffusion recipe. + 2. **Idempotent / skip-if-present.** If ``nemo_automodel`` is already importable + (e.g. a base image that ships it pre-installed) we skip the install unless + ``AUTOMODEL_REINSTALL=1`` is set. + +Unlike torchtitan there is no tokenizer asset to pre-download: the Wan diffusion +recipe pulls model weights from HuggingFace / a local ``/models`` mount at +runtime. The heavy lifting (config -> AutoModel ConfigNode -> TrainDiffusionRecipe) +is done later by ``NemoAutomodelPretrainTrainer`` inside the Primus core runtime. +""" + +import argparse +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +from examples.scripts.utils import ( + get_env_case_insensitive, + log_error_and_exit, + log_info, + write_patch_args, +) +from primus.core.launcher.parser import PrimusParser + +# Native ROCm packages to pin so the editable install never swaps them for CUDA +# wheels. +_ROCM_PINS = [ + "torch", + "torchvision", + "torchaudio", + "triton", + "pytorch-triton-rocm", + "amd-aiter", + "flash-attn", + "flash_attn", +] + +# diffusers gates its ``aiter`` attention backend behind a minimum amd-aiter +# version. Some ROCm base images ship a functional but dev-versioned aiter (e.g. +# ``0.1.1.dev*``) that fails that guard even though the kernel works (its +# ``flash_attn_func`` supports ``return_lse``). The configs default to +# ``model.attention_backend: aiter``, so without a fix diffusers would refuse it. +_REQUIRED_AITER_VERSION = "0.1.5" + + +def parse_args(): + parser = argparse.ArgumentParser(description="Prepare Primus NeMo AutoModel backend") + parser.add_argument("--primus_path", type=str, required=True, help="Root path to the Primus project") + parser.add_argument("--data_path", type=str, required=True, help="Path to data directory") + parser.add_argument("--config", type=str, required=True, help="Path to experiment YAML config") + parser.add_argument( + "--patch_args", + type=str, + default="/tmp/primus_patch_args.txt", + help="Path to write additional args (used during training phase)", + ) + parser.add_argument( + "--backend_path", + type=str, + default=None, + help="Optional AutoModel checkout path; overrides AUTOMODEL_PATH/BACKEND_PATH and the default submodule.", + ) + return parser.parse_args() + + +def resolve_backend_path(cli_path, primus_path: Path) -> Path: + """CLI --backend_path > AUTOMODEL_PATH/BACKEND_PATH env > third_party/Automodel.""" + if cli_path: + path = Path(cli_path).resolve() + log_info(f"Using AutoModel path from CLI: {path}") + return path + env_value = get_env_case_insensitive("AUTOMODEL_PATH") or get_env_case_insensitive("BACKEND_PATH") + if env_value: + path = Path(env_value).resolve() + log_info(f"Using AutoModel path from env: {path}") + return path + path = (primus_path / "third_party" / "Automodel").resolve() + log_info(f"Using default AutoModel submodule path: {path}") + return path + + +def write_rocm_constraints() -> str: + """Pin currently-installed ROCm packages so pip treats them as satisfied.""" + import importlib.metadata as md + + lines = [] + for name in _ROCM_PINS: + try: + lines.append(f"{name}=={md.version(name)}") + except md.PackageNotFoundError: + continue + fd, path = tempfile.mkstemp(prefix="primus_rocm_constraints.", suffix=".txt") + with os.fdopen(fd, "w") as f: + f.write("\n".join(lines) + "\n") + log_info(f"ROCm constraints ({len(lines)} pins) -> {path}") + return path + + +def install_automodel_editable(automodel_path: Path): + """ROCm-safe editable install of AutoModel with the diffusion extras.""" + if not (automodel_path / "pyproject.toml").exists() and not (automodel_path / "setup.py").exists(): + log_error_and_exit( + f"AutoModel checkout not found at {automodel_path} (no pyproject.toml/setup.py).\n" + "Initialize the submodule first:\n" + " git submodule update --init third_party/Automodel\n" + "or run `primus-cli deps sync`, or pass --backend_path / set AUTOMODEL_PATH." + ) + + extras = os.environ.get("AUTOMODEL_EXTRAS", "diffusion,diffusion-media") + spec = f"{automodel_path}[{extras}]" if extras else str(automodel_path) + constraints = write_rocm_constraints() + + env = os.environ.copy() + env["PIP_CONSTRAINT"] = constraints + log_info(f"Installing AutoModel (editable, ROCm-pinned): pip install --no-build-isolation -e {spec!r}") + ret = None + try: + ret = subprocess.run( + [sys.executable, "-m", "pip", "install", "--no-build-isolation", "-e", spec, "-q"], + env=env, + check=False, + ) + except OSError as e: + log_error_and_exit(f"Failed to invoke pip for AutoModel editable install: {e}") + finally: + try: + os.unlink(constraints) + except OSError: + pass + if ret is None or ret.returncode != 0: + rc = ret.returncode if ret is not None else "n/a" + log_error_and_exit(f"AutoModel editable install failed (exit {rc}).") + log_info("AutoModel editable install complete.") + + +def maybe_shim_aiter_version(): + """Make diffusers' ``aiter`` backend selectable on dev-versioned aiter images. + + Rewrites only the ``Version:`` field in the installed ``amd-aiter`` dist + metadata so it satisfies diffusers' minimum-version guard. This touches + neither aiter, diffusers, nor AutoModel source; it reverts on image rebuild + and is idempotent (skipped when the installed version already satisfies the + guard, when aiter is absent, or when its metadata is read-only). Disable with + ``AITER_VERSION_SHIM=0``. + + NOTE (single source of truth): the same metadata-rewrite is performed by the + NeMo-AutoModel dev-env ``entrypoint.sh`` (tiger-training-internal #208). The + two must stay in sync -- keep ``_REQUIRED_AITER_VERSION`` and the rewrite + logic identical, or fold both onto a shared helper if the dev-env starts + importing from this repo. Kept duplicated for now because the dev-env image + does not depend on Primus at container-build time. + """ + if os.environ.get("AITER_VERSION_SHIM", "1") != "1": + log_info("AITER version shim disabled (AITER_VERSION_SHIM=0).") + return + + import importlib.metadata as md + import re + + try: + dist = md.distribution("amd-aiter") + cur = dist.version + except md.PackageNotFoundError: + log_info("amd-aiter not installed; AITER version shim skipped.") + return + + try: + from packaging.version import Version + + if Version(cur) >= Version(_REQUIRED_AITER_VERSION): + log_info(f"AITER shim: amd-aiter {cur} already satisfies >={_REQUIRED_AITER_VERSION}; no change.") + return + except Exception: + pass # unparseable version -> attempt the bump below + + # Locate the metadata file via the public API (RECORD-relative paths resolved + # through Distribution.locate_file). Fall back to the .dist-info dir only if + # the distribution ships no RECORD (dist.files is None). + meta_path = None + try: + for f in dist.files or []: + if f.name in ("METADATA", "PKG-INFO"): + located = dist.locate_file(f) + if os.path.exists(located): + meta_path = str(located) + break + except Exception: + meta_path = None + if not meta_path: + try: + base = dist._path # PathDistribution .dist-info dir (fallback only) + for cand in ("METADATA", "PKG-INFO"): + p = os.path.join(str(base), cand) + if os.path.exists(p): + meta_path = p + break + except Exception: + meta_path = None + + if not meta_path or not os.access(meta_path, os.W_OK): + log_info( + f"AITER shim: amd-aiter metadata not writable ({meta_path}); leaving {cur} as-is. " + "If the diffusers aiter guard blocks it, set model.attention_backend: flash." + ) + return + + with open(meta_path, encoding="utf-8") as f: + txt = f.read() + new, n = re.subn(r"(?im)^Version: .*$", f"Version: {_REQUIRED_AITER_VERSION}", txt, count=1) + if n == 0 or new == txt: + log_info( + f"AITER shim: no Version field found in amd-aiter metadata ({meta_path}); leaving {cur} as-is. " + "If the diffusers aiter guard blocks it, set model.attention_backend: flash." + ) + return + with open(meta_path, "w", encoding="utf-8") as f: + f.write(new) + log_info( + f"AITER shim: bumped amd-aiter metadata Version {cur} -> {_REQUIRED_AITER_VERSION} " + "to satisfy the diffusers aiter guard." + ) + + +def ensure_automodel_installed(cli_path, primus_path: Path): + """Install AutoModel from the submodule unless it is already importable.""" + reinstall = os.environ.get("AUTOMODEL_REINSTALL", "0") == "1" + # Only take the skip-if-present fast path when the user did not explicitly + # point us at a checkout (CLI --backend_path / AUTOMODEL_PATH / BACKEND_PATH). + # If they did, honor it and (re)install so an explicit override beats a + # nemo_automodel that the base image happens to ship (matches --backend_path help). + explicit_source = ( + bool(cli_path) + or bool(get_env_case_insensitive("AUTOMODEL_PATH")) + or bool(get_env_case_insensitive("BACKEND_PATH")) + ) + if not reinstall and not explicit_source: + try: + import nemo_automodel # noqa: F401 + + log_info( + f"nemo_automodel already importable ({getattr(nemo_automodel, '__file__', '?')}); " + "skipping install. Set AUTOMODEL_REINSTALL=1 to force a reinstall from third_party." + ) + return + except Exception: + log_info("nemo_automodel not importable; installing from the AutoModel checkout.") + + automodel_path = resolve_backend_path(cli_path, primus_path) + install_automodel_editable(automodel_path) + + +def main(): + args = parse_args() + + primus_path = Path(args.primus_path).resolve() + exp_path = Path(args.config).resolve() + patch_args_file = Path(args.patch_args).resolve() + + log_info(f"PRIMUS_PATH: {primus_path}") + log_info(f"DATA_PATH: {Path(args.data_path).resolve()}") + log_info(f"EXP: {exp_path}") + log_info(f"BACKEND_PATH: {args.backend_path}") + log_info(f"PATCH-ARGS: {patch_args_file}") + + if not exp_path.is_file(): + log_error_and_exit(f"EXP file not found: {exp_path}") + + # 1) Make the nemo_automodel package importable (submodule -> editable install). + ensure_automodel_installed(args.backend_path, primus_path) + + # 1b) Make diffusers' aiter attention backend selectable on images that ship a + # dev-versioned amd-aiter (configs default to attention_backend: aiter). + maybe_shim_aiter_version() + + # 2) Validate the experiment parses and routes to the nemo_automodel backend. + primus_cfg = PrimusParser().parse(args) + try: + pre_trainer_cfg = primus_cfg.get_module_config("pre_trainer") + except Exception: + log_error_and_exit("Missing required module config: pre_trainer") + + framework = getattr(pre_trainer_cfg, "framework", None) + if framework != "nemo_automodel": + log_error_and_exit( + f"pre_trainer.framework must be 'nemo_automodel' (got {framework!r}). " + "Check the experiment config." + ) + + if not getattr(pre_trainer_cfg, "model", None): + log_error_and_exit("Missing required field: pre_trainer.model (model preset)") + + log_info( + f"NeMo AutoModel backend ready (framework={framework}, model={pre_trainer_cfg.model}). " + "Weights resolve via HF cache / /models at runtime." + ) + + # 3) Keep multi-GPU stdout to rank 0 only, consistent with other backends. + write_patch_args(patch_args_file, "torchrun_args", {"local-ranks-filter": "0"}) + + +if __name__ == "__main__": + log_info("========== Prepare NeMo AutoModel backend ==========") + main() diff --git a/primus/_thirdparty.lock b/primus/_thirdparty.lock index 465d3a116..c88e63ea7 100644 --- a/primus/_thirdparty.lock +++ b/primus/_thirdparty.lock @@ -35,6 +35,12 @@ "path": "third_party/HummingbirdXT", "url": "https://github.com/AMD-AGI/HummingbirdXT.git", "commit": "ed7b7bd0dd03ed084f645c5dda1bc7b333f326f0" + }, + { + "name": "Automodel", + "path": "third_party/Automodel", + "url": "https://github.com/NVIDIA-NeMo/Automodel.git", + "commit": "44f2acdee1628b8a5a5acb35ada6ed8c539a22e4" } ] } diff --git a/primus/backends/nemo_automodel/__init__.py b/primus/backends/nemo_automodel/__init__.py new file mode 100644 index 000000000..99759c740 --- /dev/null +++ b/primus/backends/nemo_automodel/__init__.py @@ -0,0 +1,17 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from primus.backends.nemo_automodel.nemo_automodel_adapter import NemoAutomodelAdapter +from primus.backends.nemo_automodel.nemo_automodel_pretrain_trainer import ( + NemoAutomodelPretrainTrainer, +) +from primus.core.backend.backend_registry import BackendRegistry + +# Register adapter +BackendRegistry.register_adapter("nemo_automodel", NemoAutomodelAdapter) + +# Register trainer +BackendRegistry.register_trainer_class(NemoAutomodelPretrainTrainer, "nemo_automodel") diff --git a/primus/backends/nemo_automodel/argument_builder.py b/primus/backends/nemo_automodel/argument_builder.py new file mode 100644 index 000000000..ffa7eb59c --- /dev/null +++ b/primus/backends/nemo_automodel/argument_builder.py @@ -0,0 +1,100 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +NeMo AutoModel configuration builder. + +This is the NeMo AutoModel counterpart of ``MaxTextConfigBuilder`` / +``TorchTitanJobConfigBuilder``. AutoModel (like MaxText) is an external trainer +that owns its own configuration system: a recipe is driven by a single YAML that +``nemo_automodel.components.config._arg_parser.parse_args_and_load_config`` loads +into a ``ConfigNode``. So there is no JobConfig-style dataclass of defaults to +merge against — the Primus ``module_config.params`` (already the AutoModel config +schema: ``model:``/``fsdp:``/``optim:``/``step_scheduler:``/...) *is* the config. + +This builder therefore just normalises ``params`` (a ``SimpleNamespace``) and the +trainer materialises it into an AutoModel ``ConfigNode`` via AutoModel's own +loader (see ``nemo_automodel_pretrain_trainer.py``), keeping us version-agnostic. +""" + +from __future__ import annotations + +import logging +import os +import tempfile +from types import SimpleNamespace +from typing import Any, Dict + +logger = logging.getLogger(__name__) + +# Top-level keys injected by Primus (module_base.yaml / runtime) that are not +# part of the AutoModel config schema. Stripped before handing the dict to +# AutoModel. NOTE: this is TOP-LEVEL only — nested AutoModel keys such as +# ``model.stage`` (the high/low-noise expert window) must be preserved. +# +# This is a denylist: a new top-level key added to module_base.yaml would +# otherwise leak into the AutoModel ConfigNode. ``test_module_base_keys_are_all_stripped`` +# (tests/unit_tests/backends/nemo_automodel) guards against that drift — if it +# fails, add the new key here (or confirm it belongs in the AutoModel recipe). +PRIMUS_ONLY_TOP_KEYS = frozenset( + { + "trainable", + "sink_level", + "file_sink_level", + "stderr_sink_level", + "framework", + "config", + "name", + "stage", # Primus trainer-stage selector; AutoModel uses model.stage + } +) + + +class NemoAutomodelConfigBuilder: + """Absorb Primus ``module_config.params`` and return it for downstream use.""" + + def __init__(self) -> None: + self.config: SimpleNamespace = SimpleNamespace() + + def update(self, params: SimpleNamespace) -> "NemoAutomodelConfigBuilder": + """Take params (already merged with model preset + CLI overrides).""" + self.config = params + return self + + def finalize(self) -> SimpleNamespace: + return self.config + + +def namespace_to_dict(obj: Any) -> Any: + """Recursively convert ``SimpleNamespace`` (and nested containers) to dict.""" + if isinstance(obj, SimpleNamespace): + return {k: namespace_to_dict(v) for k, v in vars(obj).items()} + if isinstance(obj, dict): + return {k: namespace_to_dict(v) for k, v in obj.items()} + if isinstance(obj, list): + return [namespace_to_dict(v) for v in obj] + if isinstance(obj, tuple): + return tuple(namespace_to_dict(v) for v in obj) + return obj + + +def strip_primus_keys(params_dict: Dict[str, Any]) -> Dict[str, Any]: + """Drop Primus-only top-level keys so AutoModel sees a clean recipe config.""" + return {k: v for k, v in params_dict.items() if k not in PRIMUS_ONLY_TOP_KEYS} + + +def export_params_to_yaml(params_dict: Dict[str, Any]) -> str: + """ + Write the (cleaned) config dict to a temporary YAML so AutoModel's own + ``parse_args_and_load_config`` can load it — exactly mirroring the MaxText + backend. The caller is responsible for deleting the file. + """ + import yaml + + fd, yaml_path = tempfile.mkstemp(suffix=".yaml", prefix="primus_automodel_") + with os.fdopen(fd, "w") as f: + yaml.dump(params_dict, f, default_flow_style=False, allow_unicode=True, sort_keys=False) + return yaml_path diff --git a/primus/backends/nemo_automodel/nemo_automodel_adapter.py b/primus/backends/nemo_automodel/nemo_automodel_adapter.py new file mode 100644 index 000000000..6f5a37246 --- /dev/null +++ b/primus/backends/nemo_automodel/nemo_automodel_adapter.py @@ -0,0 +1,120 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +NeMo AutoModel BackendAdapter. + +The AutoModel counterpart of ``TorchTitanAdapter`` / ``MaxTextAdapter``: + + - Resolve the AutoModel backend path (the ``third_party/Automodel`` submodule, + or an already-installed ``nemo_automodel`` if the base image ships one). + - Convert Primus ``module_config.params`` -> AutoModel config namespace. + - Provide the AutoModel trainer class to the Primus runtime. + - Expose a backend version string for diagnostics. +""" + +from __future__ import annotations + +import os +from typing import Any + +from primus.backends.nemo_automodel.argument_builder import NemoAutomodelConfigBuilder +from primus.core.backend.backend_adapter import BackendAdapter +from primus.modules.module_utils import log_rank_0, warning_rank_0 + + +class NemoAutomodelAdapter(BackendAdapter): + """Complete BackendAdapter implementation for NeMo AutoModel.""" + + def __init__(self, framework: str = "nemo_automodel"): + super().__init__(framework) + # Submodule lives at ``third_party/Automodel`` (the upstream repo name), + # matching the megatron/torchtitan convention where third_party_dir_name + # == the checkout dir. The python package it provides is ``nemo_automodel``. + self.third_party_dir_name = "Automodel" + + def setup_backend_path(self, backend_path=None) -> str: + """ + Resolve the AutoModel python path. + + The canonical source is the ``third_party/Automodel`` submodule (init via + ``git submodule update --init third_party/Automodel`` or ``primus-cli deps + sync``), which ``prepare.py`` installs editable on first run -- identical + to torchtitan/megatron. + + As a fast path, if no explicit ``--backend_path`` / ``BACKEND_PATH`` is + given and ``nemo_automodel`` is already importable (e.g. a base image that + ships it pre-installed), we use its install location and skip the + third_party requirement. Otherwise we + fall back to the default resolution (CLI/env/third_party/deps-sync). + """ + if not backend_path and not os.getenv("BACKEND_PATH"): + try: + import nemo_automodel # noqa: F401 + + # Return the *import root* (the directory containing the package), + # not the package dir itself: Primus puts this path on PYTHONPATH + # for the training subprocess, so `import nemo_automodel` must + # resolve from it. For an installed package that is site-packages, + # matching the third_party/ convention where the backend root + # is the repo dir that *contains* the package (e.g. third_party/ + # Megatron-LM for `megatron`). Returning the package dir would put + # the wrong dir on PYTHONPATH and risk shadowing its submodules. + import_root = os.path.dirname(os.path.dirname(os.path.abspath(nemo_automodel.__file__))) + log_rank_0( + f"[Primus:NemoAutomodelAdapter] using installed nemo_automodel import root {import_root}" + ) + return import_root + except Exception as exc: # not installed -> fall back to third_party + warning_rank_0( + f"[Primus:NemoAutomodelAdapter] nemo_automodel not importable ({exc}); " + "falling back to third_party/backend_path resolution." + ) + return super().setup_backend_path(backend_path) + + def convert_config(self, params: Any): + """Convert Primus params -> AutoModel config namespace (pass-through).""" + builder = NemoAutomodelConfigBuilder() + builder.update(params) + cfg = builder.finalize() + log_rank_0("[Primus:NemoAutomodelAdapter] Converted Primus module params -> AutoModel config") + return cfg + + def load_trainer_class(self, stage: str = "pretrain"): + """Resolve the trainer via ``BackendRegistry`` (registered in __init__).""" + from primus.core.backend.backend_registry import BackendRegistry + + # AutoModel diffusion currently exposes a single (pretrain) trainer; treat + # the common ``train`` alias as ``pretrain``. Reject any other stage up + # front so the error names the real cause (unsupported stage) rather than + # implying a broken registration. + if stage not in ("pretrain", "train"): + raise RuntimeError( + f"[Primus:NemoAutomodelAdapter] '{self.framework}' backend only exposes a " + f"'pretrain' stage (got stage={stage!r}). Finetune runs also use the pretrain " + "trainer via model.mode: finetune in the experiment config." + ) + resolved_stage = "pretrain" + try: + return BackendRegistry.get_trainer_class(self.framework, stage=resolved_stage) + except (ValueError, AssertionError) as exc: + raise RuntimeError( + f"[Primus:NemoAutomodelAdapter] '{self.framework}' trainer not registered for " + f"stage {resolved_stage!r}. Ensure primus.backends.nemo_automodel.__init__ " + "registers the trainer via BackendRegistry.register_trainer_class." + ) from exc + + def detect_backend_version(self) -> str: + try: + import importlib.metadata as importlib_metadata + except Exception: # pragma: no cover + return "unknown" + for dist in ("nemo_automodel", "nemo-automodel"): + try: + return importlib_metadata.version(dist) + except Exception: + continue + return "unknown" diff --git a/primus/backends/nemo_automodel/nemo_automodel_pretrain_trainer.py b/primus/backends/nemo_automodel/nemo_automodel_pretrain_trainer.py new file mode 100644 index 000000000..576a3fdb1 --- /dev/null +++ b/primus/backends/nemo_automodel/nemo_automodel_pretrain_trainer.py @@ -0,0 +1,92 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +NemoAutomodelPretrainTrainer: Primus wrapper for NeMo AutoModel diffusion +pre-training. Wan 2.2 (T2V-A14B) is the model that is wired up and tested +here; AutoModel's diffusion recipe also targets other models (e.g. FLUX, +Qwen-Image) but those are not exercised by this backend yet. + +Thin-wrapper pattern (same as ``MaxTextPretrainTrainer`` / +``TorchTitanPretrainTrainer``): AutoModel owns FSDP2, the dataloader, the +optimizer and checkpointing internally, so this trainer only + + backend_args (SimpleNamespace) + -> cleaned dict + -> temp YAML + -> AutoModel ``parse_args_and_load_config`` (-> ConfigNode) + -> ``TrainDiffusionRecipe`` + +and then delegates ``setup()`` / ``run_train_validation_loop()`` to the recipe. + +By routing through AutoModel's own loader we inherit its config semantics +(``_target_``/``_fn`` resolution, the ``wandb.enable`` toggle, ...) and stay +agnostic to AutoModel internals. +""" + +from __future__ import annotations + +import os +from typing import Any, Optional + +from primus.core.trainer.base_trainer import BaseTrainer +from primus.modules.module_utils import error_rank_0, log_rank_0 + + +class NemoAutomodelPretrainTrainer(BaseTrainer): + """Trainer class for NeMo AutoModel diffusion pre-training.""" + + def __init__(self, backend_args: Any): + super().__init__(backend_args=backend_args) + self._recipe: Optional[Any] = None + log_rank_0("Initialized NemoAutomodelPretrainTrainer") + + # ------------------------------------------------------------------ # + # Lifecycle + # ------------------------------------------------------------------ # + def setup(self): + """Optional pre-init phase (kept for API symmetry).""" + log_rank_0("NemoAutomodelPretrainTrainer.setup()") + + def init(self): + """Build the AutoModel recipe from Primus params and set it up.""" + log_rank_0("NemoAutomodelPretrainTrainer.init() - building AutoModel recipe") + + from nemo_automodel.components.config._arg_parser import ( + parse_args_and_load_config, + ) + from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe + + from primus.backends.nemo_automodel.argument_builder import ( + export_params_to_yaml, + namespace_to_dict, + strip_primus_keys, + ) + + params_dict = strip_primus_keys(namespace_to_dict(self.backend_args)) + + # Delegate config materialization to AutoModel's own loader (argv=[] so + # Primus's process argv is never re-parsed). + yaml_path = export_params_to_yaml(params_dict) + try: + cfg = parse_args_and_load_config(yaml_path, argv=[]) + finally: + try: + os.unlink(yaml_path) + except OSError as e: + error_rank_0(f"NemoAutomodelPretrainTrainer: failed to delete temp YAML {yaml_path}: {e}") + + self._recipe = TrainDiffusionRecipe(cfg) + self._recipe.setup() + log_rank_0("AutoModel recipe initialized successfully") + + def train(self): + """Execute the AutoModel train/validation loop.""" + if self._recipe is None: + raise RuntimeError("NemoAutomodelPretrainTrainer.init() must be called before train().") + log_rank_0("Executing AutoModel diffusion pretrain...") + self._recipe.run_train_validation_loop() + log_rank_0("AutoModel diffusion pretrain completed.") diff --git a/primus/configs/models/nemo_automodel/wan2_2_t2v_a14b.yaml b/primus/configs/models/nemo_automodel/wan2_2_t2v_a14b.yaml new file mode 100644 index 000000000..0bbbbec40 --- /dev/null +++ b/primus/configs/models/nemo_automodel/wan2_2_t2v_a14b.yaml @@ -0,0 +1,23 @@ +# NeMo AutoModel model preset: Wan 2.2 T2V A14B (mixture-of-experts DiT). +# +# Merged (as the `model:` block) into the nemo_automodel pre_trainer module +# config by the Primus launcher. `mode: pretrain` builds ONE randomly-initialised +# A14B expert chosen by `pipeline_spec.subfolder`: +# high-noise expert -> transformer (default; stage: high_noise) +# low-noise expert -> transformer_2 (override stage: low_noise too) +model: + pretrained_model_name_or_path: /models/Wan2.2-T2V-A14B-Diffusers + mode: pretrain + stage: high_noise + # REQUIRED in pretrain mode (the from_config pipeline carries no config to + # infer it from). Wan 2.2-A14B ships boundary_ratio = 0.875. + boundary_ratio: 0.875 + # AITER attention backend (gfx950). Falls back to "flash" if needed. + attention_backend: aiter + # fp32 master + bf16 compute (canonical pretraining precision, GB300-comparable). + torch_dtype: float32 + compute_dtype: bfloat16 + pipeline_spec: + transformer_cls: WanTransformer3DModel + subfolder: transformer + load_full_pipeline: false diff --git a/primus/configs/modules/nemo_automodel/pre_trainer.yaml b/primus/configs/modules/nemo_automodel/pre_trainer.yaml new file mode 100644 index 000000000..f8f667e14 --- /dev/null +++ b/primus/configs/modules/nemo_automodel/pre_trainer.yaml @@ -0,0 +1,98 @@ +# NeMo AutoModel pre_trainer module config (Primus backend: nemo_automodel). +# +# This is the AutoModel recipe schema (model:/fsdp:/optim:/step_scheduler:/...); +# the Primus launcher deep-merges the `model:` preset and the EXP `overrides:` +# on top of it, then the nemo_automodel backend materialises the whole tree into +# an AutoModel ConfigNode and runs TrainDiffusionRecipe. The defaults below are +# the optimized pretrain knobs: aiter + torch.compile + bf16 reduce + fused AdamW. +# +# NOTE: do not declare a `model:` block here — it comes from the model preset +# (merge is allow_override=False and would conflict). +extends: + - ../module_base.yaml + +seed: 42 + +wandb: + enable: false + +dist_env: + backend: nccl + timeout_minutes: 30 + +step_scheduler: + global_batch_size: 256 + local_batch_size: 32 + ckpt_every_steps: 1000 + num_epochs: 100 + log_every: 2 + save_checkpoint_every_epoch: false + +data: + dataloader: + _target_: nemo_automodel.components.datasets.diffusion.build_video_multiresolution_dataloader + cache_dir: /dataset + model_type: wan + base_resolution: [512, 512] + dynamic_batch_size: false + shuffle: true + drop_last: false + num_workers: 2 + +optim: + learning_rate: 5e-5 + optimizer: + _target_: torch.optim.AdamW + weight_decay: 0.1 + betas: [0.9, 0.95] + # fused AdamW (mutually exclusive with foreach). + foreach: false + fused: true + +performance: + check_loss: false + grad_clip_foreach: true + +lr_scheduler: + lr_decay_style: cosine + lr_warmup_steps: 0 + min_lr: 1e-6 + +flow_matching: + adapter_type: simple + adapter_kwargs: {} + timestep_sampling: uniform + logit_mean: 0.0 + logit_std: 1.0 + flow_shift: 3.0 + mix_uniform_ratio: 0.1 + num_train_timesteps: 1000 + i2v_prob: 0.3 + use_loss_weighting: true + log_interval: 100 + summary_log_interval: 10 + +fsdp: + tp_size: 1 + cp_size: 1 + pp_size: 1 + dp_replicate_size: 1 + # Single-node 8x MI355X (override dp_size to the GPU/world count for other topos). + dp_size: 8 + defer_fsdp_grad_sync: true + enable_fsdp2_prefetch: true + # Optimized knobs: + fsdp2_backward_prefetch_depth: 3 + fsdp2_forward_prefetch_depth: 2 + reduce_dtype: bfloat16 + enable_compile: true + # Keep AC on for bs>8; flip false via overrides for the bs6-8 AC-off point. + activation_checkpointing: true + +checkpoint: + enabled: false + checkpoint_dir: /checkpoints/wan22_primus_pretrain + model_save_format: safetensors + save_consolidated: final + diffusers_compatible: true + restore_from: null diff --git a/tests/unit_tests/backends/nemo_automodel/test_nemo_automodel_registration.py b/tests/unit_tests/backends/nemo_automodel/test_nemo_automodel_registration.py new file mode 100644 index 000000000..5cac22a13 --- /dev/null +++ b/tests/unit_tests/backends/nemo_automodel/test_nemo_automodel_registration.py @@ -0,0 +1,233 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for NeMo AutoModel backend registration in __init__.py. + +This test ensures that the registration logic in +primus/backends/nemo_automodel/__init__.py correctly registers the backend with +BackendRegistry. Without proper registration, the backend is unavailable to the +runtime. Mirrors tests/unit_tests/backends/megatron/test_megatron_registration.py +(NeMo AutoModel exposes a single pretrain trainer; no SFT stage). + +Importing primus.backends.nemo_automodel does not require the `nemo_automodel` +package to be installed: the adapter/trainer modules import it lazily (only inside +the trainer's init()), so these tests run without the backend present. +""" + +import pytest + +from primus.backends.nemo_automodel.nemo_automodel_adapter import NemoAutomodelAdapter +from primus.backends.nemo_automodel.nemo_automodel_pretrain_trainer import ( + NemoAutomodelPretrainTrainer, +) +from primus.core.backend.backend_registry import BackendRegistry + +_SUPPORTS_TRAINER_CLASS_REGISTRY = all( + hasattr(BackendRegistry, attr) + for attr in ("_trainer_classes", "register_trainer_class", "get_trainer_class", "has_trainer_class") +) + +if not _SUPPORTS_TRAINER_CLASS_REGISTRY: + pytest.skip( + "BackendRegistry trainer-class registration API is not available; " + "skip nemo_automodel registration tests that depend on it.", + allow_module_level=True, + ) + + +class TestNemoAutomodelBackendRegistration: + """Test that the nemo_automodel backend is properly registered via __init__.py.""" + + @pytest.fixture(autouse=True) + def ensure_backend_loaded(self): + """Ensure the nemo_automodel backend module is loaded before each test.""" + # Import the __init__ module to trigger registration. + import primus.backends.nemo_automodel # noqa: F401 + + def test_adapter_is_registered(self): + """Verify that NemoAutomodelAdapter is registered for 'nemo_automodel'.""" + assert BackendRegistry.has_adapter("nemo_automodel"), ( + "NemoAutomodelAdapter not registered. " + "Check BackendRegistry.register_adapter() call in __init__.py" + ) + + adapter_cls = BackendRegistry._adapters.get("nemo_automodel") + assert adapter_cls is NemoAutomodelAdapter, ( + f"Expected NemoAutomodelAdapter class, got {adapter_cls}. " + "Check BackendRegistry.register_adapter('nemo_automodel', NemoAutomodelAdapter) in __init__.py" + ) + + def test_trainer_class_is_registered(self): + """Verify that NemoAutomodelPretrainTrainer is registered for 'nemo_automodel'.""" + assert BackendRegistry.has_trainer_class("nemo_automodel"), ( + "NemoAutomodelPretrainTrainer not registered. " + "Check BackendRegistry.register_trainer_class() call in __init__.py" + ) + + trainer_cls = BackendRegistry.get_trainer_class("nemo_automodel") + assert trainer_cls is NemoAutomodelPretrainTrainer, ( + f"Expected NemoAutomodelPretrainTrainer, got {trainer_cls}. " + "Check BackendRegistry.register_trainer_class(NemoAutomodelPretrainTrainer, 'nemo_automodel')" + ) + + # Explicit pretrain stage should also resolve. + assert BackendRegistry.has_trainer_class("nemo_automodel", stage="pretrain") + + def test_adapter_can_be_instantiated_via_registry(self): + """Verify that get_adapter returns a working NemoAutomodelAdapter instance.""" + from unittest.mock import patch + + # Force the lazy-load path and re-register without relying on module reload + # side effects (see the megatron registration test for rationale). + original_adapters = BackendRegistry._adapters.copy() + original_trainers = BackendRegistry._trainer_classes.copy() + try: + BackendRegistry._adapters.pop("nemo_automodel", None) + BackendRegistry._trainer_classes.pop(("nemo_automodel", "pretrain"), None) + + def _fake_load_backend(_backend: str) -> None: + BackendRegistry.register_adapter("nemo_automodel", NemoAutomodelAdapter) + BackendRegistry.register_trainer_class(NemoAutomodelPretrainTrainer, "nemo_automodel") + + with patch.object(BackendRegistry, "_load_backend", side_effect=_fake_load_backend): + adapter = BackendRegistry.get_adapter("nemo_automodel") + finally: + BackendRegistry._adapters = original_adapters + BackendRegistry._trainer_classes = original_trainers + + assert isinstance( + adapter, NemoAutomodelAdapter + ), f"Expected NemoAutomodelAdapter instance, got {type(adapter).__name__}" + assert ( + adapter.framework == "nemo_automodel" + ), f"Expected framework='nemo_automodel', got '{adapter.framework}'" + + def test_trainer_class_can_be_retrieved_via_adapter(self): + """Verify trainer class retrieval through the adapter, including the 'train' alias.""" + from unittest.mock import patch + + original_adapters = BackendRegistry._adapters.copy() + original_trainers = BackendRegistry._trainer_classes.copy() + try: + BackendRegistry._adapters.pop("nemo_automodel", None) + BackendRegistry._trainer_classes.pop(("nemo_automodel", "pretrain"), None) + + def _fake_load_backend(_backend: str) -> None: + BackendRegistry.register_adapter("nemo_automodel", NemoAutomodelAdapter) + BackendRegistry.register_trainer_class(NemoAutomodelPretrainTrainer, "nemo_automodel") + + with patch.object(BackendRegistry, "_load_backend", side_effect=_fake_load_backend): + adapter = BackendRegistry.get_adapter("nemo_automodel") + finally: + BackendRegistry._adapters = original_adapters + BackendRegistry._trainer_classes = original_trainers + + # The adapter maps the common "train" alias to "pretrain". + with patch("primus.backends.nemo_automodel.nemo_automodel_adapter.log_rank_0"): + assert adapter.load_trainer_class() is NemoAutomodelPretrainTrainer + assert adapter.load_trainer_class(stage="train") is NemoAutomodelPretrainTrainer + + def test_nemo_automodel_in_available_backends_list(self): + """Verify nemo_automodel appears in the list of available backends.""" + available = BackendRegistry.list_available_backends() + assert "nemo_automodel" in available, ( + f"'nemo_automodel' not in available backends: {available}. " + "Registration may have failed in __init__.py" + ) + + +class TestNemoAutomodelRegistrationOrder: + """Test that registration happens at import time and is idempotent.""" + + def test_registration_is_idempotent(self): + """Verify that re-importing __init__ doesn't cause errors.""" + import importlib + + import primus.backends.nemo_automodel + + importlib.reload(primus.backends.nemo_automodel) + + assert BackendRegistry.has_adapter("nemo_automodel") + assert BackendRegistry.has_trainer_class("nemo_automodel") + + +class TestNemoAutomodelRegistrationFailures: + """Test error handling when registration is missing.""" + + def test_missing_trainer_registration_would_fail(self): + """Without trainer registration, get_trainer_class should fail.""" + original = BackendRegistry._trainer_classes.pop(("nemo_automodel", "pretrain"), None) + try: + with pytest.raises(ValueError, match="No trainer class registered"): + BackendRegistry.get_trainer_class("nemo_automodel") + finally: + if original: + BackendRegistry._trainer_classes[("nemo_automodel", "pretrain")] = original + + +class TestStripPrimusKeysGuard: + """Guard the strip_primus_keys denylist against module_base.yaml drift. + + strip_primus_keys() removes Primus-only top-level keys before handing the + config to AutoModel. Because it is a denylist, any key later added to + module_base.yaml would silently leak into the AutoModel ConfigNode. This test + fails loudly if module_base.yaml grows a top-level key not in the strip set. + """ + + def test_module_base_keys_are_all_stripped(self): + from pathlib import Path + + import yaml + + import primus + from primus.backends.nemo_automodel.argument_builder import PRIMUS_ONLY_TOP_KEYS + + module_base = Path(primus.__file__).resolve().parent / "configs" / "modules" / "module_base.yaml" + assert module_base.is_file(), f"module_base.yaml not found at {module_base}" + with open(module_base, encoding="utf-8") as f: + base_cfg = yaml.safe_load(f) or {} + + leaked = set(base_cfg) - set(PRIMUS_ONLY_TOP_KEYS) + assert not leaked, ( + f"module_base.yaml top-level keys {sorted(leaked)} are not in PRIMUS_ONLY_TOP_KEYS, " + "so they would leak into the AutoModel ConfigNode. Add them to PRIMUS_ONLY_TOP_KEYS " + "in primus/backends/nemo_automodel/argument_builder.py (or confirm they belong in " + "the AutoModel recipe)." + ) + + def test_strip_primus_keys_preserves_nested_model_keys(self): + from primus.backends.nemo_automodel.argument_builder import strip_primus_keys + + cleaned = strip_primus_keys( + { + "stage": "pretrain", + "framework": "nemo_automodel", + "trainable": False, + "model": {"stage": "low_noise", "pipeline_spec": {"subfolder": "transformer_2"}}, + "optim": {"learning_rate": 1e-4}, + } + ) + assert "stage" not in cleaned and "framework" not in cleaned and "trainable" not in cleaned + # Nested AutoModel keys (incl. model.stage) must survive. + assert cleaned["model"]["stage"] == "low_noise" + assert cleaned["optim"]["learning_rate"] == 1e-4 + + +class TestUnsupportedStage: + """load_trainer_class should name the real cause for an unsupported stage.""" + + def test_unsupported_stage_raises_clear_error(self): + from unittest.mock import patch + + adapter = NemoAutomodelAdapter() + with patch("primus.backends.nemo_automodel.nemo_automodel_adapter.log_rank_0"): + with pytest.raises(RuntimeError, match=r"only exposes a 'pretrain' stage"): + adapter.load_trainer_class(stage="sft") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/third_party/Automodel b/third_party/Automodel new file mode 160000 index 000000000..44f2acdee --- /dev/null +++ b/third_party/Automodel @@ -0,0 +1 @@ +Subproject commit 44f2acdee1628b8a5a5acb35ada6ed8c539a22e4 From f334162202c783929cdbab7d16524d6658ab2dec Mon Sep 17 00:00:00 2001 From: mvstrauss Date: Wed, 1 Jul 2026 09:38:28 -0500 Subject: [PATCH 2/2] fix(backend): derive AutoModel dp_size from world size and log cleanup failures - pre_trainer.yaml: stop hard-coding fsdp.dp_size=8. AutoModel derives dp = world_size / (tp*cp*pp) when dp_size is unset, so single- and multi-node runs work without manual edits and stay correct under any parallel topology. - prepare.py: replace the empty `except OSError: pass` in the constraints cleanup with a non-fatal log_info so the intentional best-effort cleanup is auditable instead of silently suppressed. Co-authored-by: Cursor --- examples/nemo_automodel/prepare.py | 6 ++++-- primus/configs/modules/nemo_automodel/pre_trainer.yaml | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/examples/nemo_automodel/prepare.py b/examples/nemo_automodel/prepare.py index d723ddf77..72f9ef7dd 100644 --- a/examples/nemo_automodel/prepare.py +++ b/examples/nemo_automodel/prepare.py @@ -146,8 +146,10 @@ def install_automodel_editable(automodel_path: Path): finally: try: os.unlink(constraints) - except OSError: - pass + except OSError as e: + # Best-effort cleanup of the temp constraints file; a failure here is + # non-fatal and must not mask the install result. + log_info(f"Non-fatal: could not remove temp ROCm constraints file {constraints!r}: {e}") if ret is None or ret.returncode != 0: rc = ret.returncode if ret is not None else "n/a" log_error_and_exit(f"AutoModel editable install failed (exit {rc}).") diff --git a/primus/configs/modules/nemo_automodel/pre_trainer.yaml b/primus/configs/modules/nemo_automodel/pre_trainer.yaml index f8f667e14..a67aa324a 100644 --- a/primus/configs/modules/nemo_automodel/pre_trainer.yaml +++ b/primus/configs/modules/nemo_automodel/pre_trainer.yaml @@ -77,8 +77,10 @@ fsdp: cp_size: 1 pp_size: 1 dp_replicate_size: 1 - # Single-node 8x MI355X (override dp_size to the GPU/world count for other topos). - dp_size: 8 + # dp_size is intentionally left unset: AutoModel derives it from the runtime + # world size and the parallel dims (dp = world_size / (tp*cp*pp)), so single- + # and multi-node runs work without edits. Pin `dp_size: ` only to force a + # specific value. defer_fsdp_grad_sync: true enable_fsdp2_prefetch: true # Optimized knobs: