# Flux diffusion infrastructure - directory structure **Created**: December 5, 2025 **Status**: ✓ Implementation Complete ## Overview This document describes the directory structure created for Flux diffusion model support in Primus, following Megatron-Core conventions with production-ready enhancements. --- ## Directory tree ``` Primus/ ├── primus/backends/megatron/ │ ├── core/models/ │ │ ├── common/diffusion_module/ # DiffusionModule base class │ │ │ └── diffusion_module.py │ │ └── diffusion/ # Diffusion models (Megatron-Core convention) │ │ ├── common/ # Shared building blocks (config, embeddings, normalization) │ │ │ ├── __init__.py │ │ │ ├── config.py # ✓ BaseDiffusionConfig │ │ │ ├── embeddings.py # ✓ TimeStepEmbedder, MLPEmbedder │ │ │ └── normalization.py # ✓ AdaLN, AdaLNContinuous, RMSNorm │ │ ├── flux/ # Flux-specific components │ │ │ ├── __init__.py │ │ │ ├── config.py # ✓ FluxConfig (with factory methods) │ │ │ ├── model.py # ✓ Flux model │ │ │ ├── layers.py # ✓ EmbedND, embedders │ │ │ ├── layer_spec.py # ✓ get_flux_layer_spec, get_flux_*_spec_for_backend, MMDiTLayer │ │ │ ├── attention.py # ✓ JointSelfAttention, FluxSingleAttention │ │ │ ├── utils.py # ✓ generate_image_position_ids │ │ │ └── checkpoint_converter.py # ✓ HF <-> Megatron conversion │ │ └── __init__.py │ │ │ ├── training/diffusion/ # Training utilities │ │ ├── schedulers/ │ │ │ ├── __init__.py │ │ │ ├── base.py # ✓ BaseScheduler │ │ │ └── flow_matching.py # ✓ FlowMatchEulerDiscreteScheduler │ │ ├── noise_utils.py # ✓ apply_flow_matching_noise, apply_ddpm_noise │ │ ├── loss_computation.py # ✓ compute_flow_matching_loss, etc. │ │ ├── timestep_sampling.py # ✓ LogitNormalSampler, UniformSampler │ │ └── __init__.py │ │ │ └── data/ │ ├── energon/ # Shared Energon infrastructure │ │ └── __init__.py # ✓ Energon wrappers │ │ │ └── diffusion/ # Diffusion-specific data │ ├── encoders/ # Hierarchical encoder registry │ │ ├── image/ │ │ │ ├── vae/ # VAE variants │ │ │ │ └── __init__.py # ✓ AutoencoderKL, VQVAE, etc. │ │ │ └── __init__.py │ │ ├── text/ │ │ │ ├── t5/ # T5 variants │ │ │ │ └── __init__.py # ✓ T5-XXL, T5-Large, etc. │ │ │ ├── clip/ # CLIP variants │ │ │ │ └── __init__.py # ✓ CLIP-L, CLIP-H, etc. │ │ │ └── __init__.py │ │ └── __init__.py # ✓ EncoderRegistry │ │ │ ├── preprocessing/ │ │ ├── image/ │ │ │ └── __init__.py # ✓ Resizing, augmentation │ │ └── __init__.py │ │ │ ├── task_encoders/ # Energon TaskEncoders │ │ ├── __init__.py │ │ └── image.py # ✓ EncodedDiffusionTaskEncoder, RawDiffusionTaskEncoder │ │ │ └── __init__.py │ ├── primus/backends/megatron/ │ └── megatron_pretrain_trainer.py # ✓ Shared Megatron pretrain trainer (drives diffusion pretraining) │ ├── primus/configs/models/megatron/ │ └── diffusion/ # YAML configs │ ├── __init__.py │ ├── flux_535m.yaml # ✓ Flux 535M config │ ├── flux_12b.yaml # ✓ Flux 12B config │ └── encoders.yaml # ✓ Encoder configs │ ├── examples/megatron/ │ ├── diffusion/ │ │ └── README.md # ✓ Training guide (consolidated) │ ├── configs/MI300X/diffusion/ # MI300X training configs │ │ ├── flux_535m_pretrain.yaml │ │ ├── flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml │ │ ├── flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml │ │ └── ... │ ├── configs/MI355X/diffusion/ # MI355X training configs (mirrors MI300X + MXFP4/MLPerf) │ │ ├── flux_12b_ddp_energon_schnell_resample_*.yaml │ │ ├── flux_12b_fsdp2_energon_schnell_resample_*.yaml │ │ └── ... │ └── prepare.py │ ├── examples/run_pretrain.sh # Main training script │ ├── tests/ │ ├── unit_tests/backends/megatron/diffusion/ # Unit test suite │ │ ├── test_flux_model.py │ │ ├── test_flux_config.py │ │ ├── test_flux_layers.py │ │ ├── test_flux_embeddings.py │ │ ├── test_flux_normalization.py │ │ ├── test_flux_utils.py │ │ ├── test_flux_checkpoint_converter.py │ │ ├── test_flux_checkpoint_utils.py │ │ ├── test_flux_layer_spec_backend_selection.py │ │ ├── test_flux_compile_checkpoint_keys.py │ │ ├── training/ │ │ ├── data/ │ │ └── distributed/ │ └── integration_tests/backends/megatron/diffusion/ │ ├── data/ │ └── distributed/ │ └── docs/backends/megatron/ └── diffusion/ # Documentation ├── README.md # ✓ Overview ├── STRUCTURE.md # ✓ This file ├── architecture_overview.md # ✓ Design details ├── data_preprocessing.md # ✓ Data guide (includes Flux-specific section) ├── energon_integration.md # ✓ Energon patterns ├── flux_architecture.md # ✓ Flux deep dive ├── fp8_training.md # ✓ FP8 training guide ├── api_reference.md # ✓ API documentation └── adding_new_models.md # ✓ Extension guide ``` --- ## Completed components ### ✓ Base classes 1. **DiffusionModule** (`core/models/common/diffusion_module/diffusion_module.py`) - Base class for all diffusion models (extends MegatronModule) - Provides Megatron-Core integration - Required methods: `forward()` - Loss computation: Use standalone functions from `loss_computation.py` - Utility methods: `get_num_params()`, `set_requires_grad()` 2. **BaseDiffusionConfig** (`common/config.py`) - Extends `megatron.core.transformer.transformer_config.TransformerConfig` - Common parameters: `in_channels`, `out_channels`, `patch_size` - Validation method for configuration integrity 3. **FluxConfig** (`flux/config.py`) - Flux-specific configuration - Parameters: `num_joint_layers`, `num_single_layers`, `context_dim`, `vec_in_dim` - Factory methods: `flux_535m()`, `flux_12b()` - 3D RoPE configuration: `axes_dim`, `theta` 3. **BaseScheduler** (`schedulers/base.py`) - Abstract base for diffusion schedulers - Required: `add_noise()`, `get_velocity_target()`, `sample_timesteps()` - Optional: `scale_model_input()`, `get_snr()`, `get_alpha()`, `get_sigma()` 4. **FlowMatchEulerDiscreteScheduler** (`schedulers/flow_matching.py`) - Concrete implementation for Flux - Linear interpolation: `x_t = (1-t)*noise + t*data` - Velocity target: `v = data - noise` ### ✓ Directory structure - **25 `__init__.py` files** with comprehensive docstrings - **Multiple implementation files** (models, configs, schedulers, data pipeline) - **Complete test suite** with fixtures and helpers --- ## Architectural decisions ### 1. Models under `core/models/` - Follows Megatron-Core convention (`megatron/core/models/gpt/`, etc.) - Easier upstream tracking when Megatron-Core adds diffusion support ### 2. Shared components in `common/` - Standard approach stores shared code in model-specific directories - Primus: `common/` for shared config, embeddings, and normalization - Flux-specific: model class, MMDiT/single-block layer specs, joint attention, and `EmbedND` ### 3. Hierarchical encoder structure - `encoders/image/vae/`, `encoders/text/t5/`, `encoders/text/clip/` - Registry pattern for config-driven selection - Easy to add new encoder variants (5+ planned per modality) ### 4. Shared Energon infrastructure - `data/energon/` for cross-model utilities (VLM, diffusion, future) - `data/diffusion/task_encoders/` for diffusion-specific TaskEncoders - Traditional approach nests Energon under model-specific directories ### 5. Synthetic (mock) data - Synthetic datasets live in `primus/backends/megatron/data/synthetic/mock_datasets.py`, wired through `primus/backends/megatron/data/synthetic_dataset_provider.py`, so training can run without real data - Unit tests exercise them under `tests/unit_tests/backends/megatron/diffusion/data/` ### 6. No PyTorch lightning - Pure Megatron patterns (no PTL DataModules) - Better integration with Megatron training loop --- ## Import examples ```python # Base classes from primus.backends.megatron.core.models.diffusion.common import ( BaseDiffusionConfig, ) # Flux configuration from primus.backends.megatron.core.models.diffusion.flux import FluxConfig # Create configs config_535m = FluxConfig.flux_535m() config_12b = FluxConfig.flux_12b() # Schedulers from primus.backends.megatron.training.diffusion.schedulers import ( BaseScheduler, FlowMatchEulerDiscreteScheduler, ) # Create scheduler scheduler = FlowMatchEulerDiscreteScheduler() timesteps = scheduler.sample_timesteps(batch_size=8, device='cuda') ``` --- ## Validation status ✓ All Python files syntactically correct ✓ No linter errors detected ✓ All imports properly structured ✓ Comprehensive docstrings ✓ Copyright headers applied (AMD 2025, Apache 2.0) --- ## Files summary All infrastructure files, model implementations, data pipeline components, tests, and documentation are complete and ready for production use. --- **End of Structure Document**