Add flux.1 to diffusion backend#832
Draft
zirui wants to merge 27 commits into
Draft
Conversation
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends Primus’s in-tree diffusion backend to support FLUX.1-dev text-to-image pretraining, reusing the existing diffusion training entrypoint, FSDP2 trainer, config normalization, and hook/launch infrastructure (building on the WAN diffusion backend introduced in #779).
Changes:
- Added FLUX model + training pipeline (flow-matching) with support for both precomputed-encoding datasets and raw image-text datasets (online frozen encoders).
- Expanded diffusion backend infrastructure: registries, distributed utilities (DTCP checkpoints + mesh/SP helpers), attention backends, optimizer variant(s), and prepare hooks + requirements.
- Added example configs and focused unit tests validating config conversion, dataset processing, and a tiny FLUX loss computation.
Reviewed changes
Copilot reviewed 79 out of 80 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit_tests/backends/diffusion/test_flux_backend.py | Unit tests for FLUX arg mapping, datasets/processors, and tiny loss. |
| runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt | Python deps for diffusion hook environments. |
| runner/helpers/hooks/train/pretrain/diffusion/prepare.py | Pretrain diffusion config validation/prepare hook (WAN + FLUX). |
| runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh | Hook step to install diffusion requirements with cache. |
| runner/helpers/hooks/train/posttrain/diffusion/prepare.py | Posttrain wrapper that reuses the pretrain diffusion prepare hook. |
| primus/configs/modules/diffusion/pre_trainer.yaml | Adds diffusion module preset (pretrain stage). |
| primus/configs/modules/diffusion/post_trainer.yaml | Adds diffusion module preset (posttrain stage). |
| primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml | WAN2.2 model preset for diffusion backend. |
| primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml | WAN2.2 SFT/posttrain preset override. |
| primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml | WAN2.1 model preset for diffusion backend. |
| primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml | WAN2.1 SFT/posttrain preset override. |
| primus/configs/models/diffusion/flux.1_dev_t2i.yaml | FLUX.1-dev T2I model preset (encoders + scaling params). |
| primus/backends/diffusion/utils/train_utils.py | Training utilities (dtype/seed/state_dict helpers). |
| primus/backends/diffusion/utils/log.py | Loguru proxy to ensure Primus-bound extras are present. |
| primus/backends/diffusion/utils/data_utils.py | Video frame-count utility helpers (WAN path). |
| primus/backends/diffusion/utils/init.py | Diffusion utils package metadata/docstring. |
| primus/backends/diffusion/trainers/init.py | Diffusion trainer exports. |
| primus/backends/diffusion/schedulers/flow_match.py | Flow-matching scheduler used by diffusion training. |
| primus/backends/diffusion/registry.py | Central registry for diffusion model/dataset/trainer builders. |
| primus/backends/diffusion/optim/adamw_fp32_state.py | FP32-state AdamW optimizer for bf16/fp16 params. |
| primus/backends/diffusion/models/wan/train_pipeline.py | WAN flow-matching training pipeline implementation. |
| primus/backends/diffusion/models/wan/t5.py | WAN T5 encoder implementation + checkpoint loader. |
| primus/backends/diffusion/models/wan/configuration_wanvideo.py | WAN config shim for training. |
| primus/backends/diffusion/models/wan/components.py | Container type for WAN model components. |
| primus/backends/diffusion/models/wan/attention_backend.py | WAN attention backend wrapper with CPU fallback. |
| primus/backends/diffusion/models/wan/adapter.py | WAN training adapter implementing Primus model interface. |
| primus/backends/diffusion/models/wan/init.py | WAN model exports. |
| primus/backends/diffusion/models/registrations/wan.py | WAN model builder + checkpoint loading logic. |
| primus/backends/diffusion/models/registrations/flux.py | FLUX model builder + checkpoint/HF download logic. |
| primus/backends/diffusion/models/registrations/init.py | Package marker for model registrations. |
| primus/backends/diffusion/models/interface.py | Abstract diffusion model interface for trainer integration. |
| primus/backends/diffusion/models/flux/utils.py | FLUX latent utilities (sampling, packing, position ids). |
| primus/backends/diffusion/models/flux/train_pipeline.py | FLUX flow-matching loss for precomputed + raw modes. |
| primus/backends/diffusion/models/flux/model.py | FLUX DiT backbone implementation (+ checkpointing toggles). |
| primus/backends/diffusion/models/flux/math.py | FLUX attention math + RoPE helpers. |
| primus/backends/diffusion/models/flux/layers.py | FLUX transformer blocks and embeddings. |
| primus/backends/diffusion/models/flux/configuration_flux.py | Simple config container for FLUX training. |
| primus/backends/diffusion/models/flux/conditioner.py | HF-based CLIP/T5 text embedders for raw-mode training. |
| primus/backends/diffusion/models/flux/autoencoder.py | FLUX autoencoder implementation + checkpoint resolver. |
| primus/backends/diffusion/models/flux/adapter.py | Primus adapter around FLUX backbone + pipeline. |
| primus/backends/diffusion/models/flux/init.py | FLUX model exports. |
| primus/backends/diffusion/models/init.py | Diffusion model package exports. |
| primus/backends/diffusion/distributed/ulysses.py | Ulysses sequence-parallel primitives. |
| primus/backends/diffusion/distributed/mesh.py | Device-mesh creation and distributed setup. |
| primus/backends/diffusion/distributed/checkpoint.py | DTCP save/load helpers for sharded checkpoints. |
| primus/backends/diffusion/distributed/init.py | Distributed utilities exports. |
| primus/backends/diffusion/diffusion_pretrain_trainer.py | Primus trainer wrapper wiring diffusion registry builders. |
| primus/backends/diffusion/diffusion_adapter.py | Backend adapter converting Primus module params to diffusion args. |
| primus/backends/diffusion/data/registrations/wan.py | WAN dataset/processor builder registration. |
| primus/backends/diffusion/data/registrations/flux.py | FLUX dataset/processor builder registration. |
| primus/backends/diffusion/data/registrations/init.py | Package marker for dataset registrations. |
| primus/backends/diffusion/data/processor.py | WAN video processor implementation (tokenize + video preprocess). |
| primus/backends/diffusion/data/flux_precomputed.py | FLUX precomputed + raw datasets and processors. |
| primus/backends/diffusion/data/dataset.py | WAN video dataset implementation (video backends). |
| primus/backends/diffusion/data/config.py | Pydantic dataset config schema (WAN path). |
| primus/backends/diffusion/data/collator.py | Raw and vision collators used by diffusion dataloaders. |
| primus/backends/diffusion/data/init.py | Diffusion data package exports. |
| primus/backends/diffusion/attention/flex.py | FlexAttention backend wrapper with compile fallback. |
| primus/backends/diffusion/attention/attention.py | Unified attention API + backend selection. |
| primus/backends/diffusion/attention/aiter.py | AITER flash-attn backend integration. |
| primus/backends/diffusion/attention/_flash_common.py | Shared flash-attn backend glue (fixed/varlen paths). |
| primus/backends/diffusion/attention/init.py | Attention package exports. |
| primus/backends/diffusion/init.py | Registers diffusion backend adapter. |
| examples/diffusion/README.md | Example usage docs for diffusion training entrypoints. |
| examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml | MI355X WAN2.2 pretrain example config. |
| examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml | MI355X WAN2.2 posttrain example config. |
| examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml | MI355X WAN2.1 pretrain example config. |
| examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml | MI355X WAN2.1 posttrain example config. |
| examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml | MI355X FLUX raw image-text pretrain example config. |
| examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml | MI355X FLUX precomputed pretrain example config. |
| .gitignore | Ensures diffusion data package isn’t inadvertently ignored. |
araina-amd
force-pushed
the
dev/zirui/flux
branch
from
79e3866 to
fabe03c
Compare
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
…rue for wan2.2-5B
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Remove redundant VAE checkpoint loading, harden AdamW optimizer fallback, and add focused tests for Wan argument normalization and optimizer fallback behavior. Co-authored-by: Cursor <cursoragent@cursor.com>
Add missing diffusion dependency checks, normalize file:// video paths across readers, report accumulated loss correctly, and validate Ulysses all-to-all divisibility before collectives. Co-authored-by: Cursor <cursoragent@cursor.com>
Add FLUX.1-dev pretraining support for precomputed and raw image-text data so the diffusion backend can run FLUX workloads without TorchTitan framework dependencies. Co-authored-by: Cursor <cursoragent@cursor.com>
Remove local asset assumptions from FLUX configs/docs, preserve AMD and BFL attribution, and tighten checkpoint/dataset validation. Co-authored-by: Cursor <cursoragent@cursor.com>
Normalize diffusion backend formatting after running pre-commit across the repository. Co-authored-by: Cursor <cursoragent@cursor.com>
Route FLUX attention through the diffusion backend dispatcher and align the default runtime backend with ROCm Wan defaults while preserving SDPA coverage in tests. Co-authored-by: Cursor <cursoragent@cursor.com>
…icts
The CI was failing with a pytest collection error:
ERROR collecting tests/unit_tests/tools/test_utils.py
import file mismatch: imported module 'test_utils' has this __file__ attribute:
tests/unit_tests/core/patches/test_utils.py
HINT: remove __pycache__ / .pyc files and/or use a unique basename for
your test file modules
Root cause: two files named test_utils.py existed in different
sub-directories of tests/unit_tests/ (core/patches/ and tools/),
but their parent directories lacked __init__.py files. Without
__init__.py, pytest imports every test file as a top-level module
using just its basename, so both were imported as 'test_utils',
causing the conflict.
Fix: add empty __init__.py to all test sub-directories that were
missing one, making each directory a proper Python package. pytest
then imports tests with fully-qualified package paths
(e.g. tests.unit_tests.core.patches.test_utils vs
tests.unit_tests.tools.test_utils), eliminating the name collision.
Clarify diffusion-generic docs and error messages, tighten Wan preprocessing quality issues, and fail fast for missing exponential shift parameters. Co-authored-by: Cursor <cursoragent@cursor.com>
araina-amd
force-pushed
the
dev/zirui/flux
branch
from
fabe03c to
0b9d37c
Compare
Avoid tying the shared diffusion utils package description to a specific video backend. 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.
Summary
This PR adds FLUX.1-dev text-to-image training support on top of the Primus
diffusionbackend introduced in #779.The goal is to make FLUX T2I pretraining runnable through the existing Primus training entrypoint, reusing the diffusion FSDP2 trainer, config system, checkpoint flow, logging, and launch infrastructure.
Motivation
FLUX is a primary diffusion workload for text-to-image training. This integration enables:
Scope
Included
Not included / future work
sp_size=1Current Status
Completed
In Progress / Follow-up
Testing
Unit Tests
python -m pytest tests/unit_tests/backends/diffusion/test_flux_backend.pySingle-node Validation
Validated on 8x AMD Instinct MI355X with
torchrun --standalone --nproc_per_node=8:Known Issues / Risks
sp_size=1gradient_checkpointing=false; enabling it needs more investigation with FSDP2Next Steps
Notes
This is a WIP draft PR building on the diffusion backend and Wan training support from #779.