Fix/ideogram4 ragged cu seqlens compile by amital-amd · Pull Request #960 · AMD-AGI/Primus

amital-amd · 2026-08-10T05:31:27Z

No description provided.

… on draft PRs (#806) Root of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. This PR targets `main` and is the **first to merge**. The content branches are cut from this branch (`feat/flux/ci-env`), not raw `main`, so they carry the bumped pins + draft guard on their own heads — which lets the rest of the stack open as drafts immediately, before this PR merges. ## What this changes Updates the public-CI dependency pins to the versions the Flux diffusion feature needs: - Bumps `PRIMUS_TURBO_COMMIT` to a recent Primus-Turbo `main` build that exposes the mxfp4 `gemm_fp4_impl(..., preshuffled=...)` fast path, and `PRIMUS_TURBO_AITER_COMMIT` to AITER `v0.1.14.post1`. Mirrors the same turbo commit into `benchmark.yaml` (the docker build picks both up through `ci.yaml`'s build-args). - Adds a **draft-skip guard** to the expensive jobs (`build-docker`, `run-unittest-torch`, `run-unittest-jax`) so the rest of the stack can be opened as Draft PRs without firing the full docker-build + GPU pipeline on every open/sync; `code-lint` stays ungated. Two parts: add `ready_for_review` to the `pull_request` trigger `types` (the default `opened/synchronize/reopened` set omits it, so without this the draft→ready flip triggers nothing) **and** add `if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }}` to the three heavy jobs. ## Why it merges first (but need not block opening) On the current `main` pins the mxfp4 path fails (`gemm_fp4_impl` rejects the `preshuffled` arg), so content PRs can only go green if their head branch carries the new pins — which stacking on `feat/flux/ci-env` provides without a merge. Merging this first matters at *merge* time: once it's on `main`, children auto-retargeted to `main` inherit the fix and the guard. No `MEGATRON_PATH` change is needed: the diffusion unit tests run in-process and `tests/conftest.py` already puts the `submodules: recursive`-checked-out `third_party/Megatron-LM` on `sys.path` when `MEGATRON_PATH` is unset. ## Dependencies None — this is the root prerequisite. ## Test plan No unit tests of its own. After the runner rebuilds the image, confirm a smoke run is green: `pytest tests/unit_tests/backends/megatron/diffusion -k mxfp4`. The guard can be sanity-checked by opening a throwaway draft PR (expect only `code-lint`) and marking it ready (expect the full pipeline). ## Files 2 (`.github/workflows/ci.yaml` — pins + draft guard; `.github/workflows/benchmark.yaml` — mirrored turbo pin). --------- Co-authored-by: Flux Split <flux-split@local> Co-authored-by: WangLingxun <linxwang@amd.com> Co-authored-by: Xiaoming-AMD <Xiaoming.Peng@amd.com> Co-authored-by: luiza-amd <Luiza.Sayfullina@amd.com>

## made some minor changes to the logging format of Megatron backend for training steps 1. Currently, we label TFLOP/s/GPU as throughput. Per customer's feedback, throughput is ambiguous and not the best to describe TFLOP/s/GPU. Change it to `Compute per GPU`. 2. For the tokens/s/GPU, we report two numbers in the format of `###/###`. It's not clear what they are. The number before `/` is perf of the step, the number after `/` is arithmetic mean of many steps. Since we measure tokens/s/GPU, which is a rate, it's better to use harmonic mean. So, change this to harmonic mean. **Before change** <img width="1449" height="221" alt="image" src="https://github.com/user-attachments/assets/3e4331d9-6fa6-4477-bd68-ce2eec726227" /> **After change** <img width="1452" height="220" alt="image" src="https://github.com/user-attachments/assets/6d6d6a97-429b-44c8-b3db-f673f3586187" />

…ing (#856) **Important**: The content of this PR was approved and merged #807, but not to main (auto-target to main was not triggered and the lesson was taken into account for further PRs). **The purpose of this PR is to finalize the merge to main and no new changes were introduced.** Base of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/ci-env` and auto-retargets to `main` once that merges. Every other content PR stacks on this one. ## What this changes The shared runtime scaffolding the rest of the feature builds on: core runtime state + train-runtime wiring, the Megatron adapter and base/pretrain trainers, and a patch auto-loader (`patches/__init__.py`) that imports every `*_patches.py` in the package on import, so each later layer only drops in its own patch file with no registry edit. Also carries the shared test root (`tests/conftest.py`, `tests/utils.py`) and a one-line `.gitignore` change (the bare `data` ignore → root-anchored `/data/*`, and nothing else) that stops Git from ignoring the in-repo `data/` source and config directories the later layers add. Kept deliberately small so it can land first. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); no functional dependency on the turbo bump. ## Test plan `pytest tests/unit_tests/core tests/unit_tests/backends/megatron`; lint/pre-commit clean. Validated locally on an AMD GPU container: 58 passed. ## Files 39 (core runtime, Megatron adapter/trainers, base patch loader, shared test root, `.gitignore`). --------- Co-authored-by: Flux Split <flux-split@local> Co-authored-by: WangLingxun <linxwang@amd.com> Co-authored-by: Xiaoming-AMD <Xiaoming.Peng@amd.com> Co-authored-by: Flux Split Trial <flux-split-trial@local>

**Important:** The content of this PR was approved and merged #823, but not to main (auto-target to main was not triggered and the lesson was taken into account for further PRs). The purpose of this PR is to finalize the merge to main and not to introduce any new changes. Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Content-independent — can review/merge in any order. Cut from `feat/flux/ci-env` (so its head carries the draft guard) and targets that branch; it auto-retargets to `main` when the CI-pins PR merges. ## What this changes The diffusion documentation set: the `docs/backends/megatron/diffusion/*` pages, the CLI / top-level / examples READMEs, the example `run_pretrain.sh`, and the requirement-check runner hook. Content reflects the current curated layout (rewritten fp8/mxfp4/structure/data docs). ## Dependencies No content parents (docs-only); cut from `feat/flux/ci-env` only to carry the draft guard, with no functional dependency on the turbo bump. ## Test plan Lint/pre-commit plus a link / relocated-path check; no runtime tests. ## Files 17 (diffusion docs, READMEs, example launcher, requirement-check hook). Co-authored-by: Flux Split Trial <flux-split-trial@local>

Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/core` — review after it. The diff here is only this layer. ## What this changes The FSDP2 optimization layer used by Flux training: fp32 and bf16-master-weight optimizer variants, incremental grad-norm, the FSDP2 fp8 all-gather path, and the related torch-FSDP2 / fp8-cache / optimizer-registration patches. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/core`. It is the parent of the turbo layer, whose float8 extension lazily imports this layer's fp8 all-gather. ## Test plan `pytest tests/unit_tests/optimizer tests/unit_tests/backends/megatron/diffusion/distributed`. Validated locally on an AMD GPU container: 87 passed. ## Files 14 (FSDP2 optimizers, fp8 all-gather, optimizer/FSDP2 patches + tests). Co-authored-by: Flux Split Trial <flux-split-trial@local>

Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/core` — review after it. Can open early; its example runs go green once the runtime PRs land. ## What this changes The curated diffusion config set — the MI300X/MI355X example configs plus the `primus/configs/{data,models,modules}/megatron/diffusion*` definitions and the torch_compile / trainer_base module configs. ## Why it's stacked here It adds files under `primus/configs/data/**`, which public `main`'s broad `data` `.gitignore` rule would otherwise ignore — so it needs the core PR's `.gitignore` fix. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/core`. ## Test plan `yaml.safe_load` parse-check per file (no pytest); optionally one `run_pretrain.sh --dry-run` per GPU family. ## Files ~43 (example + model/data/module diffusion configs). Co-authored-by: Flux Split Trial <flux-split-trial@local>

… into Primus (#847) ## Summary Migrate the MLPerf GPT-OSS-20B pretraining flow and its optimizations from the standalone mlperf source tree into Primus, so it runs through the native `primus-cli ... train pretrain` path (`stage: mlperf_pretrain`) instead of a separate entrypoint/wheel. - **MLPerf trainer & logging** integrated into the BaseTrainer architecture (`primus/backends/megatron/mlperf/`: `mlperf_pretrain_trainer.py`, `mlperf_logger.py`, `warmup.py`), registered as the `mlperf_pretrain` stage. - **Source patches migrated to `register_patch`** (`primus/backends/megatron/patches/`): MoE skip-identity-sort, SDMA param all-gather, TE BSHD-layout, turbo fused-residual-norm. - Honor `MLLOG_TRAIN_LOSS_LOG_FREQ`; add MLPerf log suppression (`mlperf_log_suppression.py`); fix a tensor-keyed `WeakKeyDictionary` in the MoE skip-identity-sort patch. ## Changes - 13 files, +2926 (additive). New `primus/backends/megatron/mlperf/` and `primus/backends/megatron/patches/{moe,parallelism,te,turbo}_patches` modules; `sdma_param_gather.py`, `fused_residual_rmsnorm.py`; `cli/main.py` wiring. ## Test plan - [x] End-to-end on MI355X (1 node × 8 GPUs), image `tasimage/primus:pr-830`, config `gpt_oss_20B-pretrain-fp8.yaml`, EP=1, fp8(e4m3, tensorwise), `use_turbo_grouped_gemm=false` (TE grouped GEMM). - [x] Trains cleanly, ~580 TFLOP/s/GPU, no NaN; **train loss 11.85 → 3.34**, **eval loss 4.57 → 3.35** (approaching the MLPerf target 3.34); eval + `:::MLLOG` events emitted correctly. ## Notes / known limitation (not in this PR) - With `use_turbo_grouped_gemm=true` on gfx950 (MI350/MI355), the Primus-Turbo fp8 tensorwise grouped-GEMM backward hits `K mismatch (5760 vs 2880)` on the non-square expert fc1, because the gfx950 NT-layout backward consumes the extension's pre-quantized `b_t` (col-wise, non-transposed) directly. The raw op is fine standalone; this is a framework/turbo interop issue tracked separately. TE grouped GEMM is the working fp8 path on MI355 for now. --------- Co-authored-by: HuangWei-95 <weihuan@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Wei Huang <weihuan@smc300x-ccs-aus-a16-19.prov.aus.ccs.cpe.ice.amd.com> Co-authored-by: Wei Huang <weihuan@dccs-1334-slurm.prov.aus.ccs.cpe.ice.amd.com>

# Description This PR adds a new global runner hook that sets `AITER_LOG_LEVEL=ERROR` to suppress the verbose AITER logs during training runs. AITER emits a large amount of log output by default, which clutters the run logs and makes it harder to spot the relevant training information. The new hook raises the AITER log level to `ERROR` so that only errors are surfaced. The hook emits an `env.*` line that is exported by `execute_hooks.sh`, following the existing global-hook convention under `runner/helpers/hooks/`. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [x] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Add `runner/helpers/hooks/03_enable_aiter.sh`, a global hook that exports `AITER_LOG_LEVEL=ERROR` to suppress the verbose AITER logs. # Checklist: - [x] The functionality is complete - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes

…ped gemm (#850) # Description This PR fixes duplicated GPU memory usage in `PrimusTurboGroupedLinear` when turbo grouped GEMM is enabled for MoE expert layers. `PrimusTurboGroupedLinear` consolidates per-expert `weight{i}` parameters into a single `self.weights` tensor for grouped GEMM execution. The previous implementation had two issues that left an extra copy of the consolidated weights resident on GPU: 1. `buffer.clone()` was used when registering `self.weights`, allocating a redundant tensor. 2. Per-expert `weight{i}` views were registered immediately in `__init__`. Those views pinned the pre-DDP-remap storage. After the distributed optimizer remapped `self.weights` into the param buffer, both the old pinned storage and the remapped buffer remained on GPU. This change registers `self.weights` directly from the consolidation buffer and defers `weight{i}` view creation until after DDP param-buffer remapping. Views are created lazily on the first forward pass (via a forward pre-hook) or when `state_dict()` is called, preserving checkpoint and legacy `weight{i}` lookup compatibility without retaining duplicate weight storage. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Remove the unnecessary `buffer.clone()` when registering the consolidated `self.weights` parameter in `PrimusTurboGroupedLinear.__init__`. - Defer per-expert `weight{i}` view registration until after DDP distributed-optimizer param-buffer remapping, avoiding pinned pre-remap storage. - Add `_ensure_weight_views()` with lazy registration triggered by a forward pre-hook and overridden `state_dict()`. - Preserve per-expert weight metadata via `_saved_weight_attrs` so checkpoint and state-dict code paths continue to work. # Checklist: - [x] The functionality is complete - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes

# Description This PR removes an unnecessary GPU-to-CPU (d2h) synchronization in `PrimusGroupedMLP.forward()` when applying the final expert output bias. Previously, `tokens_per_expert.tolist()` was called unconditionally at the end of `forward()` before `_apply_bias`, forcing a device sync on every MoE forward pass even when `output_bias` is `None` (the common case when bias is already fused into TE GroupedLinear output). The fix overrides `_apply_bias` to defer the `.tolist()` conversion until bias application is actually needed, and to skip it entirely when `bias_parallel` is `None`. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [x] Code refactoring ## Changes Please list the changes introduced in this PR: - Add a `PrimusGroupedMLP._apply_bias` static override that returns early when `bias_parallel` is `None`, avoiding the d2h sync in the no-bias path. - Move `tokens_per_expert.tolist()` from `forward()` into `_apply_bias`, so the CPU conversion only runs when bias must be applied. - Update the `forward()` call site to pass the GPU `tokens_per_expert` tensor directly to `_apply_bias` instead of a pre-materialized Python list. # Checklist: - [x] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes

## Summary This PR integrates diffusion training support into Primus training framework, including WAN2.1 and WAN2.2 models. The goal is to enable stable single-node and multi-node training for WAN diffusion models with minimal framework changes and clean extensibility. --- ## Motivation WAN diffusion models are now a primary workload for video generation training. This integration enables: - Unified training interface for diffusion models - Scalable single-node → multi-node training - Reuse of existing Primus training infra (optimizer, checkpointing, logging) --- ## Scope ### Included - WAN2.1 / WAN2.2 model wrapper integration - Diffusion training loop adaptation - Basic dataset pipeline support - Single-node SFT training validation - Checkpoint save/load compatibility ### Not included (future work) - Inference pipeline - Advanced scheduling strategies - Flash attention / kernel-level optimizations --- ## Current Status ### Completed - [x] WAN model integration (WAN2.1-1.3B / WAN2.2-5B) - [x] Training loop adapted for diffusion objective - [x] Single-node SFT training verified ### In Progress - [ ] Multi-node training validation (2-node cluster test) - [x] Config refactor (clean separation of model/trainer/data) - [x] Dataset pipeline cleanup - [x] Documentation and example scripts --- ## Testing ### Single-node - Wan2.1-1.3B SFT training: ✅ - Wan2.2-5B SFT training: ✅ ### Multi-node - 2-node test: in progress ## Known Issues / Risks - Dataset preprocessing still partially ad-hoc - Config system needs refactor for diffusion-specific parameters - Multi-node stability not fully validated yet ## Next Steps 1. Complete 2-node validation 2. Add training performance benchmarks 3. Refactor config structure for diffusion training 4. Add reproducible example scripts --- ## Notes This is a WIP draft PR. Frequent commits will be pushed as development continues. ## benchmarks ## Primus Wan2.2 TI2V 5B Benchmark - 2026-07-06 10:28:13 UTC - Summary uses median step time after skipping first 5 logged steps. - `batch_mode=local_accum` means effective per-GPU batch size via gradient accumulation with micro batch 1. - Stability/efficiency fixes enabled: `video_backend=decord`, `dataloader_num_workers=0`, `PRIMUS_CACHE_RAW_VIDEO_FRAMES=1`, and `PRIMUS_CACHE_PROCESSED_SAMPLES=1`. | engine | model | data | resolution | frames | batch_size | batch_mode | gpus | steps | gpu_mem_GB | step_time_s | step/s | TPS(samples/s/gpu) | status | | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | | primus-fsdp2-flash_attn_aiter | wan2.2-ti2v-5b | tiny-video-sample | 480P | 121 | 1 | local_accum | 8 | 95 | 27.080 | 1.760000 | 0.568182 | 0.569300 | ok | | primus-fsdp2-flash_attn_aiter | wan2.2-ti2v-5b | tiny-video-sample | 480P | 121 | 8 | local_accum | 8 | 95 | 28.330 | 14.330000 | 0.069784 | 0.558300 | ok | | primus-fsdp2-flash_attn_aiter | wan2.2-ti2v-5b | tiny-video-sample | 480P | 121 | 16 | local_accum | 8 | 95 | 28.330 | 28.430000 | 0.035174 | 0.562800 | ok | --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>

# Description This PR removes the redundant `use_turbo_fp4_autocast` flag and simplifies FP4 autocast routing in Megatron. Previously, enabling Primus-Turbo FP4 autocast required both `enable_primus_turbo` and `use_turbo_fp4_autocast` to be set. Actually the Primus-Turbo autocast is compatible with TE autocast. That duplicated control was confusing and easy to misconfigure (e.g. Turbo GEMM/attention enabled while FP4 still fell back to Transformer Engine). With this change, MXFP4 training uses the Primus-Turbo FP4 autocast path whenever `enable_primus_turbo` is enabled, which aligns FP4 behavior with other Primus-Turbo features. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Remove `use_turbo_fp4_autocast` from `primus/configs/modules/megatron/primus_turbo.yaml`. - Update `_primus_turbo_enabled()` in `fp4_utils.py` to gate the Turbo FP4 path on `enable_primus_turbo` only, instead of requiring both `enable_primus_turbo` and `use_turbo_fp4_autocast`. # Checklist: - [x] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes

…/backends (#851) ## Summary Remove the legacy `primus/modules/` package entirely. Code still used by the current core-runtime training path is migrated/refactored into its natural home in `primus/core` and `primus/backends/megatron`; the legacy `PRIMUS_TRAIN_RUNTIME=legacy` flow and dead trainer code are deleted. ## What changed - **core**: `module_utils` / `base_module` moved out of `primus/modules` into `primus/core/utils/module_utils.py` and `primus/core/base_module.py`; all importers updated. - **megatron backend**: still-used trainer utilities rehomed — PP-visualization wrappers to `core/pipeline_parallel/pp_visualizer.py`; ROCm arg validation to `patches/args/rocm_arg_validation.py`; batch loader (`DataLoaderStore`) to `data_loader_store.py`; `is_v_schedule_enabled` to `training/utils.py`. - **legacy runtime removed**: dropped the `PRIMUS_TRAIN_RUNTIME=legacy` branch/resolver in the train subcommand and the `load_backend_trainer` / `launch_pretrain_*` entrypoints in `primus/pretrain.py` (kept `setup_backend_path` / `setup_env`, still used by projection/runner/examples). Removed torchtitan/maxtext legacy pretrainers and dead code (`sft_trainer`, torchtitan `parse_utils`). - **projection**: performance/memory layer-benchmark now builds its model via a new training-neutral `PrimusRuntime.setup_model_only()` / `MegatronPretrainTrainer.setup_model_only()` instead of the legacy `MegatronPretrainTrainer`; the entire `primus/backends/megatron/pretrainer/` package is deleted. - **bug fixes surfaced during migration**: restore the FSDP2 availability probe import (autoflake had reduced the guard to `pass`, making `HAVE_FSDP2` unconditionally True), and make `get_batch_func` return a consistent 5-tuple. --------- Co-authored-by: Xiaoming-AMD <Xiaoming.Peng@amd.com>

…lock) (#810) Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/core` — review after it. Parent of the Flux model PR. ## What this changes The model-agnostic diffusion building blocks shared by all diffusion models: the common diffusion module, diffusion config/embeddings/normalization, and the diffusion transformer (DiT) block. Also lands the shared diffusion test scaffolding (`conftest`/`helpers`/`constants`) the later diffusion PRs reuse. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/core`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "embeddings or normalization"`. Validated locally on an AMD GPU container: 7 passed. ## Files 14 (common diffusion module, embeddings/normalization/config, DiT block, shared diffusion test scaffolding). --------- Co-authored-by: Flux Split Trial <flux-split-trial@local> Co-authored-by: luiza-amd <Luiza.Sayfullina@amd.com> Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>

Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/opt` — review after it. Parent of the fp8/mxfp4/compile layers. ## What this changes The Primus-Turbo integration layer: the float8 "local" extension, the turbo local-spec layer wiring, the Triton fp8-cast kernels, the native fp8 layout, and fp8 utilities. ## Why it's stacked here The float8 extension lazily imports the FSDP2 fp8 all-gather added in `feat/flux/opt`, and a turbo test exercises that path — so it bases on `feat/flux/opt`, not `feat/flux/core`. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); its float8/fp8 unit tests are green on the current CI pin (no turbo-bump dependency). Builds on `feat/flux/opt`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "turbo or native_fp8"`. Validated locally on an AMD GPU container: 22 passed. ## Files 8 (turbo float8 + local-spec extensions, Triton fp8-cast kernels, fp8 utils + tests). --------- Co-authored-by: Flux Split Trial <flux-split-trial@local> Co-authored-by: luiza-amd <Luiza.Sayfullina@amd.com> Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>

## Summary Fix `ModuleNotFoundError: No module named 'primus.modules'` that currently breaks test collection / `import primus` on `main`. An earlier refactor removed `primus/modules` and migrated the logging helpers to `primus.core.utils.module_utils`. Independently-merged diffusion and flux changes still imported `log_rank_0` from the old `primus.modules.module_utils` path, which now no longer exists — so importing the diffusion backend (e.g. `tests/unit_tests/backends/diffusion/test_wan_argument_builder.py`) fails at collection. ## Changes Repoint the three remaining offenders to the migrated location (`primus.modules.module_utils` → `primus.core.utils.module_utils`): - `primus/backends/diffusion/diffusion_adapter.py` - `primus/backends/diffusion/diffusion_pretrain_trainer.py` - `primus/backends/megatron/core/utils.py` The repo is now free of `primus.modules` references (grep-clean). ## Verification - `pre-commit run --all-files`: all hooks pass. - Import smoke: `import primus.backends.diffusion.argument_builder`, `primus.cli.main`, `train_runtime`, megatron/torchtitan adapters all OK. - `pytest tests/unit_tests/backends/diffusion/`: 8 passed (including the previously-failing `test_wan_argument_builder.py`). - Core unit subset (adapter/runtime/backend/config/base_trainer): 48 passed, no new failures.

…811) Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/model-common` — review after it. This is the central node several later PRs branch from. ## What this changes The Flux model itself: config, layer-spec, layers, attention, the model module, utils, and the checkpoint converter. The converter ships here because the package's `__init__` eagerly imports it (the package will not import without it). ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/model-common`. The data, training-primitives, compile, checkpoint-tools, and trainer PRs all descend from this one. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k flux`. Validated locally on an AMD GPU container: 29 passed. ## Files 15 (Flux config/layers/attention/model/utils, checkpoint converter + tests). --------- Co-authored-by: Flux Split Trial <flux-split-trial@local> Co-authored-by: Luiza Sayfullina <luiza.sayfullina@amd.com> Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>

Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/turbo` — review after it. ## What this changes The mxfp4 (4-bit) local-spec turbo extension plus the supporting fp4 utils and enums. ## Dependencies Builds on the CI-pins PR (`feat/flux/ci-env`) — this is the path that hard-needs the bumped Primus-Turbo (its head carries that pin; that PR merges first): on the old pin it fails with `gemm_fp4_impl(...)` "expected at most 10 args but received 11" (the concrete motivation for the CI-pins PR). Also builds on `feat/flux/turbo`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "mxfp4 or fp4_utils"`. Validated locally on an AMD GPU container: 7 passed. ## Files 5 (mxfp4 local-spec extension, fp4 utils, enums + tests). Co-authored-by: Flux Split Trial <flux-split-trial@local> Co-authored-by: luiza-amd <Luiza.Sayfullina@amd.com>

Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/turbo` — review after it. One of the parents of the trainers PR. ## What this changes The delayed-fp8-scaling patch set plus the consolidated Transformer-Engine dot-product-attention (DPA) prologue patch. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); its fp8 unit tests are green on the current CI pin (no turbo-bump dependency). Builds on `feat/flux/turbo`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "delayed_fp8 or fused_delayed"`. Validated locally on an AMD GPU container: 24 passed. ## Files 5 (delayed-fp8 patches, TE DPA prologue patch + tests). --------- Co-authored-by: Flux Split Trial <flux-split-trial@local> Co-authored-by: luiza-amd <Luiza.Sayfullina@amd.com> Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>

* Refined the auto-benchmark tool to support Rock images. * Consolidated the source code into a single `metrics.py` file. * Unified the implementation so it works with both the Megatron and TorchTitan backends via Primus

Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/flux` and also merges `feat/flux/turbo` — review after both. ## What this changes The `torch.compile` and DDP-overlap-compile patches for the Flux model, plus the compile / graph-break tests and the two backend-selection / TE-vs-local-spec tests (which top-level-import turbo modules). ## Why it has two parents The source only needs the turbo layer, but the co-located tests run the full compiled model under FSDP2 — importing the Flux model (`feat/flux/flux`) and the fp8 all-gather (via `feat/flux/turbo` → `feat/flux/opt`). ## Dependencies Builds on the CI-pins PR (`feat/flux/ci-env`) — it carries the required bumped Primus-Turbo pin for the CK launch path; that PR merges first. Also builds on `feat/flux/flux` + `feat/flux/turbo`. (No `MEGATRON_PATH` dependency: the compiled-vs-eager test runs in-process via a Dynamo reset before each compiled build, and the conftest puts the recursively-checked-out Megatron submodule on `sys.path`.) ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "compile or backend_selection or te_vs_local"`. Validated locally on an AMD GPU container: 41 passed (incl. the in-process `test_compiled_local_vs_eager`/`test_te_eager_vs_local_compiled`, confirmed passing with no Megatron entry on `PYTHONPATH`). ## Files 7 (compile + DDP-overlap patches, compile/graph-break + backend-selection tests). --------- Co-authored-by: Flux Split Trial <flux-split-trial@local> Co-authored-by: luiza-amd <Luiza.Sayfullina@amd.com>

…oders) (#812) Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/flux` — review after it. ## What this changes The diffusion data layer: dataloader + CUDA prefetch, energon and synthetic dataset providers, the image (VAE) and text (CLIP-L, T5-XXL) encoders, and the image task encoder. Also moves the energon/webdataset pins into `requirements.txt` here, since this layer's source and tests need them at runtime. ## Why it's stacked here One synthetic-dataset module top-level-imports `flux.utils`, so it bases on `feat/flux/flux` (this also transitively re-parents the prep/trainers/mlperf layers onto the Flux branch). ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/flux`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion/data` (needs `megatron-energon`/`webdataset` from the moved `requirements.txt`). Validated locally on an AMD GPU container: 50 passed. ## Files 32 (dataloader/prefetch, dataset providers, VAE/CLIP/T5 encoders, task encoder, `requirements.txt` + tests). --------- Co-authored-by: Flux Split Trial <flux-split-trial@local> Co-authored-by: eshaw2 <elenshaw@amd.com> Co-authored-by: luiza-amd <Luiza.Sayfullina@amd.com> Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>

…loss) (#816) Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/flux` — review after it. One of the parents of the trainers PR. ## What this changes The diffusion training primitives: the forward step, loss computation, noise utilities, timestep sampling, and the flow-matching schedulers. ## Why it's stacked here The forward step imports `flux.utils`, so it bases on `feat/flux/flux` rather than `feat/flux/model-common`. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/flux`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion/training -k "forward_step or loss or scheduler"`. Validated locally on an AMD GPU container: 20 passed. ## Files 13 (forward step, loss, noise/timestep sampling, flow-matching schedulers + tests). Co-authored-by: Flux Split Trial <flux-split-trial@local> Co-authored-by: luiza-amd <Luiza.Sayfullina@amd.com>

working on production documentation --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Peter Park <peter.park@amd.com>

## Summary Adds MLPerf-style pretraining examples for **Llama 3.1 8B** and **GPT-OSS 20B** on AMD MI355X, covering low-precision recipes (MXFP4 for Llama, FP8 for GPT-OSS), plus the supporting Megatron/TE patches needed to run them end-to-end. ## What's included ### Llama 3.1 8B (MI355X, MXFP4/FP4) - `examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml` — FP4 pretrain config - `examples/mlperf/llama3.1_8b/config_MI355X_1x8x1.sh` — single-node 8-GPU launch config (TP1/PP1/EP1) - `examples/mlperf/llama3.1_8b/run_and_time.sh` — MLPerf run+timing wrapper - `examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv` — tuned GEMM configs - `examples/mlperf/llama3.1_8b/README.md` ### GPT-OSS 20B (MI355X, FP8) - `examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml` — FP8 pretrain config - `examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh` — launch config (gbs 32) - `examples/mlperf/gpt_oss_20b/run_and_time.sh` - `examples/mlperf/gpt_oss_20b/tune_gemm_results.txt` - `examples/mlperf/gpt_oss_20b/README.md` ### Framework patches - `primus/backends/megatron/te_patches/fused_bias_swiglu_patches.py` — fused bias-SwiGLU support - `primus/backends/megatron/patches/validation_data_sampling_patches.py` — validation data sampling - `primus/backends/megatron/training/evaluator.py` — evaluator updates - `primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py` — trainer tweaks ## Notes - Paths in the launch/run scripts are relative to `${PRIMUS_PATH}`. - Passes pre-commit (`shellcheck`, YAML lint, `thirdparty-lock`). - Stats: 14 files, +1166 / −5. ## Test plan - [x] `bash examples/mlperf/llama3.1_8b/run_and_time.sh` completes on 1×MI355X (8 GPU), MXFP4 - [x] `bash examples/mlperf/gpt_oss_20b/run_and_time.sh` completes on 1×MI355X (8 GPU), FP8 - [x] Loss/throughput sanity-checked - [train.mlperfpretrain.llama3.1-8b.log](https://github.com/user-attachments/files/29991139/train.mlperfpretrain.llama3.1-8b.log) ; [train.mlperfpretrain.gptoss.log](https://github.com/user-attachments/files/29991140/train.mlperfpretrain.gptoss.log) - [x] Docker image - `unifiedtrainingdockers.azurecr.io/utd/ci:primus_the_rock_ci_2e8745e_20260709` - [x] For GPT OSS 20 FP8 use atomic FP32 `NVTE_CK_IS_V3_ATOMIC_FP32=1` to avoid NaN and `use_turbo_grouped_gemm: false` ## Results (1×MI355X, 8 GPU, GBS=32) | Model | Precision | Target eval loss | Samples to target | Final eval loss | Time to train | Throughput | |---|---|---|---|---|---|---| | Llama 3.1 8B | MXFP4 | ≤ 3.30 | 196,608 | 3.2807 | 102.42 min (6,145.34 s) | 31.99 samples/s | | GPT-OSS 20B | FP8 | ≤ 3.34 | 233,472 | 3.3297 | 153.33 min (9,199.68 s) | 25.38 samples/s | <sub>Metrics from MLPerf MLLOG `run_start`→`run_stop` (both runs converged, `status: success`). "Samples to target" = `samples_count` at the first eval where `eval_accuracy ≤ MLLOG_TARGET_EVAL_LOSS`. Throughput = `overall_throughput` (samples/s). Target loss set via `MLLOG_TARGET_EVAL_LOSS` (3.3 / 3.34).</sub> --------- Co-authored-by: Vidushi Goyal <vidgoyal@dccs-1334-slurm.prov.aus.ccs.cpe.ice.amd.com>

Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/flux` — review after it. ## What this changes Standalone tooling — the HF→Primus Flux checkpoint converter and the empty-encoding generator, plus the converter test. The converter module itself ships in the Flux model PR; this PR is just the CLI tools that use it. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/flux`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion/test_flux_checkpoint_converter.py`. Validated locally on an AMD GPU container: 5 passed. ## Files 3 (checkpoint converter tool, empty-encoding generator + test). Co-authored-by: Flux Split Trial <flux-split-trial@local> Co-authored-by: luiza-amd <Luiza.Sayfullina@amd.com>

…raining example on MI355X (#877) ## Summary Adds an end-to-end MLPerf Training 6.0 **Llama2-70B LoRA** post-training example targeting **AMD MI355X** (8× GPU, 1 node), driven through Megatron-Bridge and `primus-cli`. - Dataset: [GovReport](https://gov-report-data.github.io/) (SCROLLS `gov_report`), packed to **8192** tokens. - Model: **meta-llama/Llama-2-70b-hf** with LoRA (rank 16, alpha 32). - Precision: **MXFP4** + BF16, with **FP8 delayed scaling** after healing at step 340. - Quality target: eval loss **< 0.925**. ## Changes ### Example (`examples/mlperf/llama2_70b/`) - `README.md` — full run instructions (container launch, data/checkpoint paths, config reference). - `config_MI355X_1x8x1.sh` — MLPerf 6.0 MI355X env (MXFP4, AITER, NCCL, MLLOG, 550 iters, lr=0.0006). - `configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml` — post-train overrides. - `a4w4_tuned_gemms.csv` — tuned AITER A4W4 GEMM configs. - `run_and_time.sh` / `run_in_container.sh` — one-shot MLPerf runners via `primus-cli`. ### Megatron-Bridge runtime patches (`primus/backends/megatron_bridge/patches/mlperf_llama2_70b/`) Applied only when the MLPerf Llama2-70B run is selected (see `conditions.py`), replacing prior git patches to `third_party/Megatron-Bridge`: - `megatron_patches.py` — MXFP4 recipe + phase tracking, optional TE SwiGLU (`USE_TE_SWIGLU=1`). - `bridge_patches.py` — data loaders, deterministic eval reset, SFT attention-mask cache, NeMo-style step timing. - `lora.py` — NeMo-stable LoRA (`use_te_fused_lora=False`). - `resettable_data_iterator.py` — deterministic validation iterator. ### MLPerf recipe (`primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/`) - `llama2_custom.py` — `llama2_70b_lora_mxfp4_config` recipe + custom training/eval loop. - `pre_quantize_mxfp4.py` — `PRE_QUANTIZED_MODEL=True` pre-quantization (FP8 stash on CPU, MXFP4 weight swap). - `mxfp4_healing.py` — MXFP4→FP8 healing at `HEALING_ITER=340`. - `nemo_loss.py` — NeMo-equivalent `MaskedTokenLossReduction`. - `_log_suppression.py` — non-MLLOG log suppression for clean submission logs. ### Core / plumbing - `config_utils.py` — recipe resolution now supports direct custom module paths (e.g. `primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.llama2_custom`) with a fallback. - Nested config override logic (`_apply_nested_overrides`) moved from the pretrain trainer to the shared base trainer and wired into the post-train trainer; `runtime_config_update` now invoked before finetune. - `transformer_engine_spec_provider.py` — optional eager-attention fallback. - `cli/main.py` — prefer the git checkout over an installed wheel for in-tree Primus modules; make MLPerf log suppression import optional. - `train_runtime.py` — guard against empty `data_path`. ### Hooks (`runner/helpers/hooks/train/posttrain/megatron_bridge/`) - `00_install_requirements.sh` — container-safe pip cache path; pin `fsspec`. - `01_convert_checkpoints.sh` — container-aware data root / HF cache resolution; HF→Megatron conversion with correct attention env handling. - `02_prepare_mlperf_dataset.sh` (new) — download/convert SCROLLS gov-report and build packed `.npy` + metadata. ### Top-level dataset utilities - `download_dataset.py`, `convert_dataset.py`, `create_metadata.py`, `hash.py` — dataset download, packing, metadata, and integrity hashing. ## MLPerf configuration | Parameter | Value | |-----------|-------| | `train_iters` | 550 | | `global_batch_size` | 8 | | `micro_batch_size` | 1 | | `seq_length` | 8192 | | `lr` | 0.0006 | | `eval_interval` / `eval_iters` | 48 / 24 | | Parallelism | TP=1, PP=1, CP=1 (8 GPUs data parallel) | | Quality target | eval loss < 0.925 | ## Results Measured on **MI355X 1×8×1** (8 GPUs, 1 node) via `examples/mlperf/llama2_70b/run_and_time.sh`. Log: `logs/log_20260715_014540.txt`. | Metric | Value | |--------|-------| | **Time to train** | **517.1 s (8.62 min)** training-loop wall time | | **Final iteration** | 384 / 550 (early exit — quality target reached) | | **Final eval loss** | **0.9242** (target: < 0.925) ✓ | | **Final eval PPL** | 2.52 | | **Consumed train samples** | 3,072 | | **MXFP4→FP8 healing** | Applied at step 340 (`DelayedScaling`, 320 FP8 weights restored) | ### Throughput Steady-state training throughput from `logs/log_20260715_014540.txt` (NeMo-style train-step wall clock, `global_batch_size=8`, `seq_length=8192`): | Phase | Steps | Step time | Model TFLOP/s/GPU | Model TFLOP/s (8× GPU) | Tokens/s/GPU | Tokens/s (8× GPU) | Samples/s | |-------|-------|-----------|-------------------|------------------------|--------------|-------------------|-----------| | Warmup | 10 | 1.30 s | 2,797 | 22,376 | 6,292 | 50,336 | 6.15 | | **MXFP4** | 20–339 | **1.15 s** | **~3,165** | **~25,320** | **~7,120** | **~56,960** | **~6.96** | | **FP8** (post-healing) | 350–380 | **1.43 s** | **~2,540** | **~20,320** | **~5,715** | **~45,720** | **~5.59** | | Eval (iter 384) | — | 13.2 s | — | — | — | — | **5.40** | Cluster tokens/s = `tokens/s/GPU × 8`. Samples/s = `global_batch_size / step_time`. **Eval loss progression** (validation at eval intervals): | Iteration | Eval loss | |-----------|-----------| | 192 | 0.9685 | | 240 | 0.9597 | | 288 | 0.9514 | | 336 | 0.9491 | | **384** | **0.9242** | Training exited early at iteration 384 when eval loss dropped below the MLPerf quality target (0.925). --- ⚡ *Prepared with [PR Pundit](https://github.com/AMD-AGI/pr-pundit) — AMD OSS Agent* --------- Co-authored-by: vidushi8 <vidgoyal@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: shekhar <shekhar.pandey@amd.com> Co-authored-by: Xiaoming-AMD <Xiaoming.Peng@amd.com>

Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/data` — review after it. ## What this changes The offline data-preprocessing layer (auth / download / finalize / validate plus the raw/ingest/encoded pipeline stages), wired into a new `primus data` CLI subcommand. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/data` (uses its providers and inherits the energon/webdataset deps). ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion/data/preprocessing tests/unit_tests/cli/test_data_config.py`. Validated locally on an AMD GPU container: 63 passed. ## Files 20 (preprocessing pipelines, data CLI subcommand + tests). Co-authored-by: Flux Split Trial <flux-split-trial@local>

…4 processor Adds cu_seqlens/max_seqlen to the processor and makes them the first path tried, so a ragged batch runs exact var-len flash with no device-to-host read anywhere in the compiled region. They are declared as named parameters, not **kwargs, on purpose. Diffusers' attention module filters forwarded kwargs against inspect.signature(self.processor.__call__).parameters, so a **kwargs-only processor receives nothing at all and the only trace is a logger.warning that is easy to lose under torchrun. The branch tests whether anything was provided, which is a property of the run rather than of the batch, so it costs a guard and not a graph break. The mask is not read on this path; the model still builds it, which is now dead weight worth removing separately. The legacy mask-derived path stays as the fallback for callers that do not precompute and as the reference the unit test checks against. Co-authored-by: Cursor <cursoragent@cursor.com>

…thetic presets cu_seqlens carries 2*B+1 entries, so B has to be constant for its shape to be constant. With drop_last false the last batch of an epoch is short, which changes that shape (and hidden_states') and recompiles the graph for one batch in every epoch - a shape-stability hole independent of caption raggedness. Co-authored-by: Cursor <cursoragent@cursor.com>

…xt width The collate left-padded llm_features to the longest caption in each batch, so the packed sequence length S moved whenever a batch's longest caption differed. torch.compile keys its graphs on input shapes, so that recompiled on sequence length alone and blocked compile on real captions no matter how cu_seqlens was produced. Default to the longest caption in the whole cache, read from metadata.json without loading samples. It is a dataset-level constant and identical on every rank, which matters because a per-rank width would desync the sharded shapes. Nothing is truncated at the default; an explicit smaller max_text_tokens caps captions and warns once, and -1 restores the old per-batch behaviour. collate_fn is bound with functools.partial rather than a closure so it stays picklable for the dataloader workers. Co-authored-by: Cursor <cursoragent@cursor.com>

… buffer attention_kwargs is a dead channel for this model in diffusers 0.39.0: the LoRA decorator keeps only scale, both block call sites pass four positional arguments, and Ideogram4TransformerBlock.forward has no **kwargs to forward. The precomputed cu_seqlens therefore never arrived, and every layer silently fell back to deriving the packing from the mask - the host-syncing, graph-breaking path the precompute existed to avoid. Route it on the module instead. The processor is already handed its Ideogram4Attention as its first argument, so one int32 non-persistent buffer - the same tensor object on all 34 attention modules - is registered there and the adapter publishes with a single copy_ per step. No diffusers fork and nothing copied from upstream, so a version bump cannot silently diverge from a vendored forward. Dynamo lifts the buffer as a graph input rather than baking its values in, and FSDP2 leaves it as a plain tensor outside the state dict. The processor clones before the kernel call. aiter's var-len op treats cu_seqlens as mutable: it saves the tensor for its backward and then writes it, bumping the version counter once per call, so one buffer shared by 34 layers moves its version 34 times per forward while each layer's backward still expects the version it saved, and the step dies with "IntTensor[5] is at version 35; expected 34". The legacy path never saw this because every layer built its own tensor from the mask. A published packing outlives the step, so a caller that bypasses the adapter could attend on a stale one. The processor rejects any cu_seqlens whose length is not 2B+1, which compares static shape metadata only and costs a guard rather than a host sync, and publish_packing(required=True) raises when nothing can read the packing - on a subset of ranks that would otherwise average two attention paths into one gradient with nothing in the logs to say so. precompute_cu_seqlens_active() now gates on both switches. Without PRIMUS_IDEOGRAM_VARLEN_ATTN the stock SDPA processor has no cu_seqlens parameter, so building the packing and reserving the pad column's token position every step buys nothing. Also corrects the no-diffusers fallback constants, which were wrong in every value that matters (pad 0 vs -1, llm 1 vs 3, offset 4096 vs 65536) and would mislabel every token with no error; they now warn when used. The cu_seqlens test imports the padding id from the adapter rather than hardcoding it, since it runs both with and without diffusers present. Co-authored-by: Cursor <cursoragent@cursor.com>

Re-pin third_party/Automodel to upstream main HEAD (2026-08-10, v0.3.0rc4-1201-gc852b16ff) as part of the new-node Ideogram-4 bring-up. Motivation: - #3157 native diffusion Context Parallelism (pure-Ulysses), a capability Ideogram-4 requires. - #3122 typed RecipeConfig.build() diffusion recipe refactor. - #3411 avoid duplicate FSDP2 prefetch all-gathers. NOTE: this pin moves the diffusion recipe/parallelizer API out from under the no-fork hooks. The Ideogram-4/FLUX hooks and YAML configs still target the old recipe and MUST be re-ported + re-validated before training will run on this pin. The FSDP2 ZeRO-2 bug is confirmed still present upstream at this pin: reshard_after_forward is dropped by the keyword whitelist when the parallel manager is created, so a ZeRO-2 request silently runs as ZeRO-3. Co-authored-by: Cursor <cursoragent@cursor.com>

… trainer The core runtime instantiates every trainer with BaseModule-style context kwargs (module_name, primus_config, module_rank, module_world_size, module_master_addr, module_master_port) alongside backend_args, but this trainer's __init__ only accepted backend_args, so construction died with a TypeError before the trainer logged anything. Nothing surfaced the cause: the CLI's traceback goes to a stderr the logger has already taken over, so the run just exited 1 right after "Loaded trainer class", on one GPU and on eight alike. TorchTitanPretrainTrainer already carries this fix; the AutoModel backend landed via #835 without it. BaseTrainer filters the kwargs cooperatively (it forwards them only when BaseModule is in the MRO), so accepting and forwarding **kwargs is all that is needed and the trainer keeps its stage-only behaviour. This is a bug on main, not something the Ideogram branch introduced - it reproduces on any nemo_automodel pretrain run through primus/cli/main.py. Co-authored-by: Cursor <cursoragent@cursor.com>

…zer schema Automodel c852b16ff (#3122) unified the diffusion recipe onto the LLM/VLM YAML schema and added _reject_removed_diffusion_keys, which fails construction when a config still carries the diffusion-only keys. The rejection on `optim` is unconditional - the presence of the block alone is enough - so every diffusion config in the tree stopped building against the new pin. Apply the upstream mapping: optim.learning_rate -> optimizer.lr optim.optimizer -> optimizer (keeping the explicit _target_) step_scheduler.log_every -> step_scheduler.log_remote_every_steps The nested optim.optimizer block collapses into the top-level optimizer, so the AdamW knobs (weight_decay, betas, foreach/fused) move up one level and keep their values. No behaviour is intended to change: the Ideogram-4 synthetic 8-GPU smoke reproduces its loss trajectory exactly, before and after the migration. One semantic difference remains to resolve: log_remote_every_steps also gates the local [TRAIN] step line, which now prints once per run rather than once per step at the same value of 1. Tracked with the rest of the recipe-API re-port. Co-authored-by: Cursor <cursoragent@cursor.com>

Automodel c852b16ff (#3122) removed the module-level nemo_automodel.recipes.diffusion.train._build_optimizer that this hook patched, so PRIMUS_IDEOGRAM_ZERO1 had become silently dead: the hook registry catches install failures so a default run keeps working, and the flag is off by default, so nothing in a normal run pointed at it. It would have surfaced as "ZeRO-1 measured no different from DDP" in a bench A/B. The optimizer is now built by OptimizerConfig.build(model, device_mesh=, is_peft=), which returns one optimizer per model part rather than a single one, so the wrap maps over the returned list. The learning rate and per-group defaults are read off the built optimizer instead of being passed in, since the new seam hands us the result rather than the ingredients; the ctor-signature filtering that keeps ZeRO from forwarding non-constructor defaults (AdamW's decoupled_weight_decay) is unchanged. Verified by forcing every hook's gate on and calling install() against the new pin: all nine install cleanly, and this was the only one that had broken. Co-authored-by: Cursor <cursoragent@cursor.com>

…mments/defaults The Ideogram-4 integration grew up alongside an internal investigation, and its docstrings and config headers still carried artefacts of that context that do not belong in a public repository: - a reference to the specific engagement the work was done for, - measured performance figures quoted as fact (FSDP comm share, compile speedup, activation-checkpoint recompute tax, attention backward workspace size), - references to internal planning documents and their phase labels, which read as dangling pointers to anyone outside that context, - references to helper scripts that live outside this repository, and - cluster-specific filesystem defaults for dataset and profiler output, which cannot resolve anywhere else. Each is replaced by the underlying technical reason, which is what the reader actually needs: why a dense attention mask forgoes flash, why torch.compile is wired only on the FSDP2 path, why the DDP path trails on a single node, and which sibling config to use for real data. Dataset defaults now follow the /dataset convention already used by the FLUX and Wan presets, and the profiler falls back to ./output. Comments, docstrings and path defaults only; no behaviour change. All affected YAML parses and all modules compile. Co-authored-by: Cursor <cursoragent@cursor.com>

… packages The backend had grown to a flat directory of 17 modules mixing four unrelated concerns, which made it hard to tell what a given file applied to and impossible to review or upstream one model's support on its own. Three groups now have their own home: quantization/ FP8 GEMM, FP8 attention, non-deterministic bf16 attention. Model-agnostic, so it stays out of models/. models/flux/ FLUX parallelization strategy. models/ideogram4/ adapter, var-len attention, packing transport, parallelization strategy, ZeRO-1, profiler, offline processor, and data/ (synthetic + real cache loaders). A model package must not import another's; anything shared belongs in quantization/. That isolation is the point of the split. Two files are renamed rather than just moved: flux_ac.py and ideogram_ac.py become models/*/parallelize.py, because both register a full parallelization strategy (activation-checkpoint wrap *and* FSDP2 sharding) rather than only AC. The trainer's hook registry is now grouped by owner, so adding or removing one model's hooks does not touch another's. The new __init__.py files deliberately re-export nothing: hooks are imported lazily by dotted path precisely so that a missing optional dependency degrades to a skipped hook, and re-exporting would undo that. A README records the layout, the hook contract, and the fact that hooks fail silently when an upstream pin moves their patch target. Pure moves plus import-path updates; no logic changed. Verified: 57 backend unit tests pass, all 9 registered hooks plus the 4 remaining moved modules import at their new paths, and the 8-GPU synthetic smoke reproduces its loss trajectory bit-identically over 60 steps. Co-authored-by: Cursor <cursoragent@cursor.com>

…iffusion caches The AutoModel diffusion recipe can train from a pre-encoded flat cache, so the VAE and text encoder run once offline and training needs neither their weights nor their memory. Building that cache was previously a standalone script living outside the repository, which made it undiscoverable and left the training config pointing at something no in-tree command could produce. Adds it as a sibling of the existing 'primus data diffusion-*' commands rather than a mode of them: those emit Energon WebDataset shards for the Megatron backend, this emits a flat per-sample cache for AutoModel. The two share no on-disk format, and folding them together would put two unrelated pipelines behind one set of flags. primus data automodel-cache --model ideogram4 \ --image-dir <images> --caption-dir <captions> --output-dir <cache> Models register through primus/backends/nemo_automodel/data/registry.py as "<module>:<callable>" strings, imported only when that model is requested. The CLI can therefore list what is available, and print --help, without importing torch or any model's optional dependencies. The builder itself lives with its model, in models/ideogram4/data/build.py, so the per-model isolation of the backend still holds. Selection is deterministic (sorted order, no substitution for skipped samples) so the same source and arguments reproduce the same cache; --shuffle opts out. Captions that exceed the token budget are skipped rather than truncated, because the loader assumes every cached sample fits. Verified end to end on a multi-GPU node: the CLI builds a cache, and that cache then trains through Ideogram4CacheDataloaderConfig with a decreasing loss. That is also the first exercise of the real-data loader against the current Automodel pin; only the synthetic loader had been covered before. The unit tests resolve every registry entry, since lazy import means a stale dotted string would otherwise stay invisible until a GPU-bound cache build. Co-authored-by: Cursor <cursoragent@cursor.com>

… the diffusion path The setting was parsed from YAML and then discarded twice, so it never had any effect and every transformer block reshared after forward regardless of config. Blocks were re-all-gathered during backward, giving ZeRO-3 communication volume (~3P bytes/step) where ZeRO-2 (~2P) was configured. Nothing failed and nothing warned; the only symptom was throughput. Drop 1, upstream: the recipe parses the value into manager_args, but _create_parallel_manager rebuilds FSDP2Config from an explicit keyword whitelist that omits it, so FSDP2Manager.reshard_after_forward was always None. With None, the per-layer heuristic in apply_fsdp2_sharding_recursively reshards every block but the last. Repaired by a new unconditional hook that re-applies the key from manager_args after the manager is built. It repairs rather than invents - the value still comes from YAML - and is behaviour-neutral when the key is absent, so there is nothing to gate. If upstream later adds the field to the whitelist the hook writes the value the manager already holds and degrades to a no-op. Drop 2, ours: both parallelization strategies called apply_fsdp2_sharding_recursively with seven positional arguments, stopping one short of the eighth, which is reshard_after_forward. The value arrived from upstream as a keyword, landed in **kwargs and was never read. Both strategies now accept and forward it, and the trailing arguments are passed by keyword so this class of bug cannot recur. Fixing either drop alone changes nothing, which is why both are here. The hook is model-agnostic, so it lives in a new distributed/ subpackage rather than under models/ - the diffusion whitelist affects every diffusion model on this path, not only ours. Worth reporting upstream for that reason. The strategies log the received value alongside the value the hook applied, and error if the hook is not installed. That guard exists because the trainer swallows hook exceptions: without it, a failed install would produce one log line and then a full run at the wrong communication volume. The hook's state is read through a module reference rather than a from-import, since a from-import binds the initial False and would make the guard fire on every healthy run. Presets now set reshard_after_forward explicitly, with a note that the prefetch knobs are inert while it is false - there is no backward all-gather left to prefetch, so they would otherwise look active while doing nothing. Verified on 8 GPUs against both the synthetic and the pre-encoded-cache configs: the loss trajectory is bit-identical, confirming this changes memory and communication rather than numerics, while peak memory rises by about one parameter copy, which is the independent evidence it took effect. Tests assert the plumbing, not the behaviour, because the bug was a value that was never passed. Reintroducing either drop fails them. Co-authored-by: Cursor <cursoragent@cursor.com>

… bare optimizer The Ideogram-4 profiler wrapper drove prof.step() from self.optimizer.step, but OptimizerConfig.build returns one optimizer per model part, so self.optimizer is a list. Enabling the profiler therefore raised AttributeError the moment the training loop started, and the CLI loses tracebacks to its own logger, so it presented as an immediate exit with no explanation. Hooks the LAST optimizer in the list. The recipe steps each in turn, so this advances the profiler schedule exactly once per optimization step; hooking all of them would advance it once per model part and capture the wrong steps. Worth noting how this was missed: the hook's install() succeeds, because the stale assumption is in the wrapper body rather than at install time. The post-bump audit probed every hook by forcing its gate on and calling install(), which is why this one looked healthy. Installing is not evidence that a hook works. Same root cause as the ZeRO-1 hook re-point, and the last of the known casualties of the optimizer-list change. Co-authored-by: Cursor <cursoragent@cursor.com>

AutoModel's diffusion CP path refuses to enable context parallelism unless the transformer declares a diffusers _cp_plan. Everything else CP needs already exists upstream, so that one missing class attribute was the entire blocker for Ideogram-4. Supply it Primus-side; diffusers stays pristine. Ideogram-4 is single-stream: hidden_states, encoder_hidden_states, position_ids, segment_ids and indicator all run along the same packed sequence axis. So the first three per-token inputs split together at the root, because the root forward combines them before the first block. Two entries are deliberate omissions: - segment_ids stays whole. After the Ulysses all-to-all each rank holds the full sequence with a subset of heads, so attention wants a full-length mask, which is what the in-forward segment_ids comparison yields only while it is unsplit. - position_ids stays whole and MRoPE's outputs are split instead, since RoPE is applied to q/k before the all-to-all and must therefore be local. Guard the combination with the var-len flash processor. That processor calls aiter directly and never reaches dispatch_attention_fn, which is where the all-to-all happens, so under CP each rank would attend only within its own shard and train on wrong gradients without raising anything. The timestep assertion rides on a forward_pre_hook rather than a forward wrapper: diffusers resolves the root plan entries against the forward signature, so a *args/**kwargs wrapper hides encoder_hidden_states and indicator from it and silently leaves them unsplit. Verified on 2 ranks against the non-CP reference on a ragged batch (max abs diff 5.4e-07, evenly spread across shards). A negative control that splits segment_ids is rejected loudly, which also confirms the all-to-all really ran.

Every way this plan can be wrong is silent: the model still runs and the loss still falls, only the numbers are wrong. So these assert the structure rather than any behaviour. Three regressions are pinned, all of which were live while writing it. Wrapping forward, which hides encoder_hidden_states and indicator from the signature diffusers resolves the root plan against (verified the test fails when the wrapper is put back). Splitting segment_ids or position_ids, whose absence is load-bearing. And enabling CP alongside the var-len processor, which bypasses the dispatch that performs the all-to-all. No GPU or distributed init: the plan is a class attribute, and the model is instantiated tiny on CPU only to resolve module names.

…eset builds The hook patched OptimizerConfig.build, but a YAML `_target_: torch.optim.AdamW` is not an OptimizerConfig subclass, so it is routed through OptimizerFromFactoryConfig -- which overrides build and never chains to super(). The patched method was therefore never called, and a run would train with fully replicated optimizer state while reporting a healthy ZeRO-1 setup. Nothing raised and no fallback warning fired, because the code that logs those was itself never reached. Walk the OptimizerConfig hierarchy and wrap every class defining its own build, rather than naming the two current overriders, so a future override cannot re-open the hole. Make the wrap idempotent in case an override does chain to super(), and log the patched classes so coverage is visible in the run log. Verified by unit test, including a negative control that reproduces the bypass. Co-authored-by: Cursor <cursoragent@cursor.com>

The hook exposed record_shapes and with_stack but not with_modules, so a trace could carry op shapes with no way to attribute a kernel to the part of the model that issued it -- which is usually the half a reader needs. with_modules adds the nn.Module frames to the python call tree that with_stack emits. Both stay off by default since they inflate trace size and CPU-side step time; enable them deliberately when a trace is for someone to read. Co-authored-by: Cursor <cursoragent@cursor.com>

…Ideogram-4 loader At a realistic caption width the synthetic dataloader generated a large [max_text_tokens, 53248] tensor per sample. On the CPU side that dominates the micro-batch, so a throughput run ends up measuring the dataloader rather than the model -- the wider the caption budget, the worse it gets. Add share_text_features: generate that buffer once and hand the same tensor to every index. Throughput depends on the shapes and on the per-sample text lengths, not on each sample holding different noise, and the lengths still vary per index. The aliasing contract is the one cache_in_memory already relies on: repeat reads return the identical tensor object. Off by default. The overfit smoke needs distinct samples for its loss-decrease signal to mean anything, so this is a benchmarking lever only, and the tests pin that default along with the properties that make sharing safe (unchanged shape and dtype, unchanged per-sample lengths, latents still distinct). Co-authored-by: Cursor <cursoragent@cursor.com>

…dder configs Adds the FLUX equivalent of the Ideogram-4 synthetic data path so throughput, memory and MFU can be measured without the input pipeline in the loop, letting the FLUX ladder be compared against the published Ideogram-4 numbers. The synthetic loader emits the post-collate batch that the FLUX recipe expects: image_latents [B,16,h,w] fp16, text_embeddings [B,512,4096] bf16 and pooled_prompt_embeds [B,768] bf16, with dtypes matched to what the upstream preprocessor actually writes into the real cache. Unlike the Ideogram-4 loader there is no text_lengths key. FLUX pads every caption to the full 512 T5 positions and its joint attention is dense, so the padding is genuinely computed and the caption-length distribution has no effect on cost. Recording per-sample lengths would imply a saving that does not exist. share_text_features reuses one text buffer across samples. It is a perf-only switch: generating 512x4096 per sample is CPU-bound and would make the loader, not the GPU, the thing being measured. The synthetic module config is kept separate from flux_pre_trainer.yaml rather than overriding its dataloader in place, because the two dataloaders take different keys and a deep merge would leave the real cache's keys behind. Co-authored-by: Cursor <cursoragent@cursor.com>

…e AC for FLUX Two problems on the FLUX real-AC path, both found while building the perf ladder. The flag was tested with a bare truthy check. Config and CLI paths can forward it as a raw string, and a non-empty "false" is truthy in Python, so a run whose whole purpose was to measure activation checkpointing OFF would silently checkpoint all 57 blocks. That corrupts the AC axis in the direction least likely to be noticed, because the run still succeeds and merely looks slower and smaller than it should. False-like strings are now normalized exactly as the Ideogram-4 strategy does. FLUX also only had whole-block AC, while Ideogram-4 has off/selective/full. The missing middle point is the interesting one -- op-level partial AC keeps attention and half the matmuls and recomputes only the cheap ops -- and without it the two studies' AC axes could not be compared. This uses the shared Automodel selective machinery that the Ideogram-4 strategy already calls, so it is the same lever rather than a FLUX-specific reimplementation. The helper replaces blocks by identity, so FLUX's two separate ModuleLists can be passed in one call. Measured at 256, mbs 2, 8x MI355X, confirming the modes are distinct and ordered: full 0.489 s 40.61 GB selective 0.481 s 44.66 GB off 0.438 s 50.12 GB and that a string "false" wraps nothing. Co-authored-by: Cursor <cursoragent@cursor.com>

…UX too The wrapper patches TrainDiffusionRecipe.run_train_validation_loop and finds its per-step boundary through the recipe's optimizer, so nothing in it is specific to Ideogram-4 -- it lives under models/ideogram4 only because that model needed it first. Duplicating 190 lines to profile FLUX would have been the worse trade, so the gate now also accepts a model-agnostic PRIMUS_DIFFUSION_PROFILE. PRIMUS_IDEOGRAM_PROFILE keeps working: the published Ideogram-4 sweep was driven by it and its runbook still names it. The module is left in place for the same reason, rather than moved somewhere more honest about its scope. Co-authored-by: Cursor <cursoragent@cursor.com>

…ram-4 Makes aiter's Triton MHA the default kernel for the Ideogram-4 var-len flash path, worth -20% step time at 8 ranks / mbs=8 / 1024px / compile=true (ZeRO-2 3.125 -> 2.499 s, ZeRO-3 3.274 -> 2.625 s) with identical peak memory and identical loss. The kernels were never the problem: aiter configures its Triton MHA head-dimension-BLIND (_get_config() takes no arguments), so at Ideogram-4's hd=256 the blocks overflow the register file (512 regs/lane, 212 spilled) and the shipped Triton path loses badly to CK-tile. Five block-size values turn that around on one MI355X at the production packing: CK-tile 5.25 fwd + 39.35 bwd = 44.60 ms Triton, aiter config 22.23 + 66.83 = 89.07 ms Triton, tuned here 6.45 + 14.90 = 21.35 ms 2.09x vs CK None of them is reachable through a public aiter API: the autograd wrapper between the two halves drops the config on the way to the backward, where 24 of the 27 ms live. triton_varlen_attn.py therefore drives aiter's forward and one-kernel backward itself and owns only the CONFIG, keeping aiter's shipped dict as the base so an upstream retune of any other field is inherited. Each half is a torch.library custom op rather than a torch.autograd.Function: Dynamo cannot trace aiter's Triton launch path (num_ctas passed straight to the kernel, triton._C.libtriton, hasattr on the kernel object), which produced 57 graph breaks inside a compiled transformer block, and a break there splits the region FSDP2 registers its per-layer collectives around. As opaque ops the path traces fullgraph with zero breaks and zero recompiles, matching CK, so this does not block enable_compile. Two silent failure modes are guarded: a config with BLOCK_M2 < BLOCK_N1 leaves the tail of Q at its zero init and returns a wrong dq with correct dk/dv (four sweep candidates hit it, dq rel-err 0.37-0.62), and aiter's lru_cached config dict must not be mutated in place or every other Triton MHA caller in the process is retuned. Output and all three gradients match CK element-wise at the production shape (1.5e-3, the bf16 floor). The block sizes are measured for (gfx950, hd=256) only. Elsewhere the kernels would fall back to aiter's head-dim-blind config, which is SLOWER than CK, so an unswept GPU resolves back to impl=ck at install time instead of onto a silent regression. Selection is a model-preset key (model.varlen_attn_impl: triton | ck), which needed a way for an install hook to read resolved YAML: hooks run before the recipe exists and take no config argument, so the trainer publishes the merged params and the hook reads model.varlen_attn_impl out of them. PRIMUS_IDEOGRAM_VARLEN_ATTN_IMPL overrides it for A/B runs, and the individual block sizes stay env-overridable for a re-sweep. Refs: AMD-AGI/tiger-training-internal#267

The shipped num_warps=8, num_stages=2 was tuned jointly with BLOCK_N on a single benchmark shape. Re-sweeping across packed sequence lengths from 4k to 20k tokens shows that once BLOCK_N=16 is in force, aiter's original 4/1 is 3.7-11.2% faster, widening with sequence length, and bit-identical in output and gradients. BLOCK_N=16 is confirmed optimal and the backward is already at its config optimum, so this closes the tuning question for this head dim rather than opening it. Co-authored-by: Cursor <cursoragent@cursor.com>

…r configs Wan A14B had no Primus-side model package, so throughput could only be measured on a real encoded cache — which made the ladder wait on a 23 GB video encode and exposed every row to loader stalls. The synthetic loader mirrors the Ideogram-4 one: a Primus dataclass with build() that the YAML _target_ points at, emitting the same batch keys collate_fn_video and FlowMatchingPipeline.step expect. The perf-ladder EXPs pin the study shape (512x512x81 -> latent [16,21,64,64], 21,504 video tokens) and keep the synthetic and real-cache paths separate so the S0/S6 convergence arm can be run against the same model preset. Co-authored-by: Cursor <cursoragent@cursor.com>

…k for Wan Automodel's in-tree WanParallelizationStrategy silently ignores two FSDP2 config values, so both axes were unmeasurable on the Wan 2.2 ladder: - activation_checkpointing is a bare truthiness test, so the non-empty string "selective" takes the full-AC branch. There is no selective path for Wan at all, which is why selective and full produced bit-identical step time and memory. - apply_fsdp2_sharding_recursively is called with seven positional args and stops one short of reshard_after_forward, so zero2/zero3/zero3rs4 all fall through to the same per-layer heuristic. This is the strategy-side half of the bug distributed/fsdp2_reshard.py already repairs on the manager side; fixing either half alone changes nothing. Neither is visible in the logs: the config echo shows the requested value because the value was parsed, just never applied. Subclass the in-tree strategy rather than reimplement it, so its TP plan stays authoritative across submodule bumps, and re-apply the registry entry after _diffusers/auto_diffusion_pipeline._init_parallelizer(), which overwrites the Wan and Hunyuan entries with fresh in-tree instances every time it runs. Registering at hook time alone is silently undone. FLUX and Ideogram-4 are unaffected only because they are absent from that hard-coded list. Env-gated by PRIMUS_WAN_PARALLELIZE_FIX=1, off by default, since it swaps a registry entry other Wan users may rely on. Measured at mbs 1, 512x512x81, 8x MI355X: selective now costs 140.9 GB against 63.1 GB for full and runs 7.4% faster, the trade selective AC is supposed to make. Applying reshard_after_forward=False moves full AC from 40 to 63 GB, matching the ~24 GB of unsharded bf16 params ZeRO-2 keeps resident. Co-authored-by: Cursor <cursoragent@cursor.com>

…lens-compile Brings the FLUX perf-ladder work onto the shared working branch so Ideogram-4, Wan 2.2 and FLUX all live together while the studies are in flight: - FLUX.1-dev synthetic dataloader and perf-ladder configs - selective AC for FLUX, plus false-like AC string normalization - the diffusion profiler wrapper generalized to serve FLUX as well The selective-AC commit matters beyond FLUX: it fixes the same bare-truthiness bug just fixed for Wan in 5aed3c1, so keeping it on a side branch left the same defect live in two of the three models on this branch. This branch is a staging area; it will be split into isolated PRs later.

…fault The FLUX example launch lines now set PRIMUS_ATTN_NONDETERMINISTIC=1. It was previously described as one arm of an A/B, so the default path was the slower one; the uplift grows with attention's share of the step, which at the longer sequence lengths is most of it. Also records two things that were easy to get wrong here: - The previous comment called the non-deterministic backward "numerically equivalent". It is unbiased but not bitwise reproducible, since the dQ/dK/dV reduction order is not fixed. Dropping the variable is the way back to a deterministic backward. - attention_backend is not the same axis and is not a substitute. Setting it to `aiter` on its own inherits aiter's deterministic=True default, which diffusers' aiter backend never overrides, and the resulting split-K dQ accumulator OOMs at ordinary training shapes. The preset therefore stays on `flash`, which the hook rebinds onto the aiter kernel anyway when the variable is set, and which is a working deterministic backward when it is not. Co-authored-by: Cursor <cursoragent@cursor.com>

off/selective/full are the only granularities available, and at long sequence lengths none of them is the right size: 2048/mbs 2 needs to give back ~38 GB to reach a safe occupancy, while selective sheds 115 GB and full sheds 205 GB. Both charge recompute for the whole overshoot. Which blocks get wrapped is a second axis and it is the one sized to the problem. All 34 blocks' activations are live at the end of forward, so checkpointing k of them sheds k/34 of the peak for k/34 of the recompute -- measured flat at 6.3 GB and 0.76% of step time per block, linear from 0 to 12. PRIMUS_IDEOGRAM_AC_EVERY=n wraps blocks 0, n, 2n, ... inside the existing full-AC path; unset or 1 leaves every current configuration byte-identical. At 2048 this replaces selective/mbs 4 (12.930 s) with 6-of-34/mbs 2 at 5.729 s, +12.8% samples/s at 77% of HBM instead of 73%. Co-authored-by: Cursor <cursoragent@cursor.com>

clairesonglee and others added 30 commits July 2, 2026 08:19

ci: revert to fla 0.4.x (#836)

d5609ac

add mixtral-8x22B config files for maxtext backend (#841)

73cd33d

Auto benchmark tool refinement (#872)

68236e7

* Refined the auto-benchmark tool to support Rock images. * Consolidated the source code into a single `metrics.py` file. * Unified the implementation so it works with both the Megatron and TorchTitan backends via Primus

Dev/production doc (#824)

c8e5c09

working on production documentation --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Peter Park <peter.park@amd.com>

Anant Mital and others added 4 commits August 11, 2026 01:34

mvstrauss force-pushed the fix/ideogram4-ragged-cu-seqlens-compile branch from 2e8cff1 to 555570f Compare August 11, 2026 08:37

mvstrauss and others added 6 commits August 11, 2026 05:18

mvstrauss force-pushed the fix/ideogram4-ragged-cu-seqlens-compile branch from c44f7c6 to 8c6f204 Compare August 11, 2026 10:19

mvstrauss and others added 18 commits August 11, 2026 05:47

Navigation Menu

Uh oh!

Fix/ideogram4 ragged cu seqlens compile - #960

Fix/ideogram4 ragged cu seqlens compile#960
amital-amd wants to merge 135 commits into
feat/backend/nemo_automodel-ideogramfrom
fix/ideogram4-ragged-cu-seqlens-compile

amital-amd commented Aug 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

20 participants

Uh oh!

Conversation

amital-amd commented Aug 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

20 participants