Fix/ideogram4 ragged cu seqlens compile - #960
Draft
amital-amd wants to merge 125 commits into
Draft
Conversation
… 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>
Add an env-gated (PRIMUS_ATTN_NONDETERMINISTIC=1) hook that rebinds the diffusers FLASH/AITER attention backends to aiter flash_attn_func with deterministic=False, trading a reproducible backward for a faster atomic-accumulation backward that is numerically equivalent up to floating-point atomic-add ordering. Off by default (no-op); no diffusers/AutoModel source is modified, and it falls back to the original backend for context-parallel / return_lse / additive mask / dropout / sliding-window calls. Wire installation into NemoAutomodelPretrainTrainer via a small _install_optional_hooks() step that runs before the recipe builds the transformer, alongside the existing env-gated FP8/MXFP4 GEMM and FP8-attention hooks (each a no-op unless its own flag is set). Co-authored-by: Cursor <cursoragent@cursor.com>
FluxTransformer2DModel is not in the AutoModel parallelization-strategy registry, so it falls back to the default strategy whose activation-checkpointing wraps per-layer self_attn/mlp submodules that a FLUX block does not have. As a result fsdp.activation_checkpointing is a silent no-op for FLUX (accepted but nothing is checkpointed), capping the achievable batch size. Add an env-gated (PRIMUS_FLUX_REAL_AC=1, default off) hook that registers a FLUX parallelization strategy via the submodule's register_parallel_strategy entry point (no fork). The strategy wraps every dual-stream and single-stream FLUX transformer block in a non-reentrant checkpoint_wrapper before FSDP2 sharding (only when the recipe requests activation checkpointing), then shards like the in-tree Wan/Hunyuan diffusion strategies. Wired into the same _install_optional_hooks path as the other env-gated hooks. Co-authored-by: Cursor <cursoragent@cursor.com>
…ram-4 Bumps third_party/Automodel 44f2acde (v0.3.0rc4-870) -> ba4ec358 (v0.3.0rc4-998) to pick up the newer flow-matching recipe and MoE / expert-parallel infrastructure needed for Ideogram-4 integration. bf16-focused branch; low-precision (mxfp4) hooks come later. Co-authored-by: Cursor <cursoragent@cursor.com>
…ttention, and activation checkpointing Ideogram-4 is a 9.3B single-stream diffusion transformer. This adds the modeling core as no-fork Primus hooks on the nemo_automodel backend: - ideogram4_adapter.py: flow-matching (rectified-flow) training adapter - ideogram4_varlen_attn.py: var-len flash-attention path (dense fast-path capable) - ideogram_ac.py: real activation checkpointing (full + TorchTitan-style selective) - configs/models/nemo_automodel/ideogram4.yaml: weightless model definition Co-authored-by: Cursor <cursoragent@cursor.com>
… pipeline - ideogram_processor.py: real image/text preprocessing into a training cache - ideogram_cache_data.py: dataloader over the preprocessed cache - ideogram_synthetic_data.py: synthetic dataloader for perf / overfit smokes Co-authored-by: Cursor <cursoragent@cursor.com>
…rofiler hooks - ideogram_zero1.py: DDP + ZeRO-1 distributed optimizer (no-fork) as an A/B point against FSDP2 sharding - ideogram_profile.py: torch.profiler train-loop wrapper for op/comm-share analysis Co-authored-by: Cursor <cursoragent@cursor.com>
…etrain recipes Wire the Ideogram-4 hooks into NemoAutomodelPretrainTrainer's optional-import registry and add the trainer-module configs + example pretrain recipes (FSDP2, DDP+ZeRO-1, and PCAM-cache variants). Co-authored-by: Cursor <cursoragent@cursor.com>
PRIMUS_IDEOGRAM_ATTN_ASSUME_DENSE=1 skips the block-diagonal mask analysis and runs dense flash over the whole row. That is exact only when no row has padding; on a ragged batch it lets pad tokens attend and leaks attention across segments, corrupting training with no error at all. The adapter already materializes text_lengths on the host, so detecting the dangerous configuration is free. Refuse it there rather than let the run produce plausible-looking but wrong losses. Co-authored-by: Cursor <cursoragent@cursor.com>
The attention processor currently recovers the var-len packing by inspecting the block-diagonal mask's values, which needs three device-to-host reads. Dynamo cannot trace those, so each one graph-breaks, and under FSDP2 a mid-forward break splits the region the per-layer collectives are registered around and desyncs them across ranks. The adapter already holds the per-sample text lengths as Python ints, so it can build the packing itself for the cost of a loop and one host-to-device copy. build_cu_seqlens is module-level rather than a method so it stays unit-testable without nemo_automodel present. Reserve one always-pad column so every row keeps a non-empty pad segment. Without it a caption filling the full text width collapses to a single segment, making the segment count -- and therefore cu_seqlens' shape -- vary per batch, which trades a graph break for a recompile per length pattern. The alternative, capping captions at max_text-1, would drop a real token to satisfy the compiler. max_seqlen is the static bound S rather than the true per-batch maximum, since Dynamo guards Python ints by value. The processor consumes these in the next commit; for now the two dict entries are carried and unused. Off switch: PRIMUS_IDEOGRAM_PRECOMPUTE_CU_SEQLENS=0. Co-authored-by: Cursor <cursoragent@cursor.com>
…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>
mvstrauss
force-pushed
the
fix/ideogram4-ragged-cu-seqlens-compile
branch
from
August 11, 2026 08:37
2e8cff1 to
555570f
Compare
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>
mvstrauss
force-pushed
the
fix/ideogram4-ragged-cu-seqlens-compile
branch
from
August 11, 2026 10:19
c44f7c6 to
8c6f204
Compare
…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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.