Dev/zirui/flux mlperf#876
Draft
zirui wants to merge 31 commits into
Draft
Conversation
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>
Avoid tying the shared diffusion utils package description to a specific video backend. Co-authored-by: Cursor <cursoragent@cursor.com>
- Build FLUX RoPE position ids in float32 instead of the model compute dtype so grid indices are not silently corrupted (bf16 only represents integers up to 256 exactly), matching the reference FLUX implementation. - Align frozen raw-path encoders (VAE/T5/CLIP) with the DiT dtype before use so non-bf16 runs no longer hit a dtype-mismatch error. - Validate that prompt-dropout empty encodings match the per-sample encoding shape and fail fast with a clear message instead of a cryptic in-place assignment error mid-training. - Add focused regression tests for the above. Co-authored-by: Cursor <cursoragent@cursor.com>
Resolve PR 832 merge conflicts after the Wan diffusion backend landed on main. Keep the FLUX branch versions of overlapping diffusion files because they already include equivalent Wan fixes plus FLUX-specific additions. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Document the FLUX precomputed and raw image-text launch paths with the required data and checkpoint environment variables, and remove Wan-only wording from shared diffusion docs. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This draft PR expands the Primus diffusion backend to support FLUX.1-dev text-to-image training and moves the training/config surface closer to the MLPerf TorchTitan FLUX reference recipe (precomputed CC12M encodings + optional raw bring-up path), including updated prepare-hook validation, example configs, and focused unit tests.
Changes:
- Add a FLUX model implementation + training pipeline (flow-matching loss, RoPE position IDs, optional frozen T5/CLIP/AE for raw mode) and register it in the diffusion backend.
- Add FLUX datasets/processors for precomputed encodings and raw image-text input; update the diffusion prepare hook and requirements to support Hugging Face
datasets/hub resolution. - Add MLPerf-oriented example configs and unit tests covering FLUX backend behaviors and scheduler validation.
Reviewed changes
Copilot reviewed 40 out of 53 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit_tests/core/utils/init.py | Add test package marker. |
| tests/unit_tests/core/trainer/init.py | Add test package marker. |
| tests/unit_tests/core/runtime/init.py | Add test package marker. |
| tests/unit_tests/core/patches/init.py | Add test package marker. |
| tests/unit_tests/core/launcher/init.py | Add test package marker. |
| tests/unit_tests/core/backend/init.py | Add test package marker. |
| tests/unit_tests/core/init.py | Add test package marker. |
| tests/unit_tests/cli/init.py | Add test package marker. |
| tests/unit_tests/backends/torchtitan/init.py | Add test package marker. |
| tests/unit_tests/backends/megatron/init.py | Add test package marker. |
| tests/unit_tests/backends/diffusion/init.py | Add test package marker. |
| tests/unit_tests/backends/init.py | Add test package marker. |
| tests/unit_tests/backends/diffusion/test_flux_backend.py | New unit tests for FLUX arg building, attention dispatch, datasets/processors, and tiny loss smoke checks. |
| tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py | New test for FlowMatchScheduler exponential-shift parameter validation. |
| runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt | Add HF dependencies needed by FLUX (datasets/hub/sentencepiece). |
| runner/helpers/hooks/train/pretrain/diffusion/prepare.py | Extend config validation for FLUX raw/precomputed modes + optional DiT init checkpoint handling. |
| primus/configs/models/diffusion/flux.1_dev_t2i.yaml | New base FLUX model preset with optional pretrained path and encoder references. |
| primus/backends/diffusion/utils/init.py | Update diffusion backend utility module docstring. |
| primus/backends/diffusion/trainers/fsdp2.py | Skip final save when save_strategy disables checkpointing. |
| primus/backends/diffusion/trainers/base.py | Improve local batch-size inference for list batches; add optional debug-step logging hooks. |
| primus/backends/diffusion/trainers/init.py | Update trainer module docstring. |
| primus/backends/diffusion/schedulers/flow_match.py | Add explicit error when exponential shift is enabled without mu/dynamic length. |
| primus/backends/diffusion/registry.py | Register FLUX model/dataset builders in the diffusion backend registry. |
| primus/backends/diffusion/README.md | Document FLUX data/checkpoint modes and new example launch patterns. |
| primus/backends/diffusion/models/registrations/flux.py | New FLUX model builder (DiT init loading + optional frozen encoders/AE). |
| primus/backends/diffusion/models/flux/utils.py | New FLUX latent utilities (mean/logvar sampling, packing, position encoding). |
| primus/backends/diffusion/models/flux/train_pipeline.py | New FLUX flow-matching training pipeline + loss computation for raw/precomputed modes. |
| primus/backends/diffusion/models/flux/model.py | New FLUX DiT backbone implementation (blocks, RoPE, checkpointing). |
| primus/backends/diffusion/models/flux/math.py | New FLUX attention + RoPE helper functions integrated with backend attention selection. |
| primus/backends/diffusion/models/flux/layers.py | New FLUX transformer layer implementations and embeddings. |
| primus/backends/diffusion/models/flux/configuration_flux.py | New config shim type for FLUX training settings. |
| primus/backends/diffusion/models/flux/conditioner.py | New HF-based text encoder wrapper for T5/CLIP. |
| primus/backends/diffusion/models/flux/autoencoder.py | New FLUX autoencoder implementation + checkpoint loader with HF download support. |
| primus/backends/diffusion/models/flux/adapter.py | New Primus GenAIModel adapter wrapping FLUX training forward path. |
| primus/backends/diffusion/models/flux/init.py | Export FLUX public symbols for diffusion backend. |
| primus/backends/diffusion/models/init.py | Update diffusion backend model exports docstring. |
| primus/backends/diffusion/diffusion_pretrain_trainer.py | Generalize naming from Wan-only trainer to diffusion trainer; add FLUX dependency check. |
| primus/backends/diffusion/diffusion_adapter.py | Switch to DiffusionArgBuilder and log the selected diffusion model name. |
| primus/backends/diffusion/data/registrations/wan.py | Refactor Wan dataset registration to use shared dataset/processor modules. |
| primus/backends/diffusion/data/registrations/flux.py | New FLUX dataset registration (raw vs precomputed). |
| primus/backends/diffusion/data/registrations/init.py | Update registrations module docstring. |
| primus/backends/diffusion/data/processor.py | New WanVideoDataProcessor implementation for prompt/token/video preprocessing. |
| primus/backends/diffusion/data/processing_wanvideo.py | New WanVideo processor implementation (image processor + tokenizer). |
| primus/backends/diffusion/data/flux_precomputed.py | New FLUX datasets/processors for precomputed encodings and raw image-text. |
| primus/backends/diffusion/data/dataset.py | New WanVideoDataset implementation using raw collation and backend video loading. |
| primus/backends/diffusion/data/config.py | New Pydantic dataset config schema (video + packing fields). |
| primus/backends/diffusion/data/collator.py | Add raw collator and a vision collator utility. |
| primus/backends/diffusion/data/init.py | Export diffusion data modules (Wan + FLUX). |
| primus/backends/diffusion/argument_builder.py | Generalize to DiffusionArgBuilder with FLUX defaults and additional section mappings. |
| examples/diffusion/README.md | Expand examples documentation to include FLUX usage and data expectations. |
| examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml | New example config for FLUX raw image-text pretrain. |
| examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml | New example config for FLUX precomputed-encodings pretrain. |
| .gitignore | Ensure diffusion data module paths aren’t accidentally ignored; ignore pycache/pyc under that subtree. |
Comment on lines
+49
to
+53
| outputs = self.hf_module( | ||
| input_ids=batch_encoding["input_ids"].to(self.device), | ||
| attention_mask=None, | ||
| output_hidden_states=False, | ||
| ) |
Comment on lines
267
to
271
| ("num_train_epochs",): ("num_train_epochs",), | ||
| ("dataloader_num_workers",): ("dataloader_num_workers",), | ||
| ("resume_from_checkpoint",): ("resume_from_checkpoint",), | ||
| ("dataloader_num_workers",): ("dataloader_num_workers",), | ||
| } |
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
Draft PR to align Primus FLUX training with the MLPerf TorchTitan reference:
https://github.com/pytorch/torchtitan/tree/9603aa83f0e0c69b5d66f9fab61b1f28e9013a2f
This branch is based on
dev/zirui/fluxand depends on #832. It keeps the existing Primus/Megatron FLUX implementation path, while narrowing the training recipe toward MLPerf Flux-schnell convergencebehavior.
Key alignment areas:
validation loss
0.586.Test Plan
Notes
This PR should remain draft until #832 is merged. After #832 lands, this branch can be rebased onto the merged base and promoted to a formal PR.
Environment And Config
flux.1-schnell/ TorchTitanflux-schnellconstant_with_warmup, configuredwarmup_steps=1600, effective warmup=100 after clampDoubleStreamBlock,SingleStreamBlockblock wrapping, explicit module wrapping for input/final layerstorch.compileenabledflash_attn_aiter/data/cc12m_preprocessed, empty encodings/data/empty_encodings/zirui/code/Primus-flux/local_runs/flux_bench_20260618/cc12m_test_repeat/metadata.local.jsonlgoogle/t5-v1_1-xxl, CLIPopenai/clip-vit-large-patch14, VAE/models/ae.safetensorsValidation Status
Performance Metrics
Metrics use steps 2-100 to exclude first-step compile/setup overhead. TorchTitan memory is console-reported reserved memory; Primus reports allocated/reserved/peak.
Primus vs TorchTitan Performance
Training Loss Summary
Loss Alignment Against TorchTitan