Architecture overview

Architecture overview#

This document provides a detailed overview of the diffusion model architecture in Primus, including design decisions, directory structure, and implementation patterns.

Table of contents#

Design philosophy
Directory structure
Architectural decisions
Component hierarchy
Data flow
Comparison with alternative implementations

Design philosophy#

Core principles#

Megatron-Core Native: Built on Megatron-Core patterns and conventions
Extensibility: Easy to add new models (DiT, MovieGen, custom)
Clarity: Clear separation between shared and model-specific code
Reusability: Shared components across multiple models
Performance: Support for precalculated data and multi-GPU training

Architectural advantages#

Aspect	Primus Design Choice
Model Location	`core/models/diffusion/` (Megatron-Core convention)
Layer Organization	Unified `TransformerBlock` with heterogeneous specs
Shared Code	Dedicated `common/` directory
Encoders	Hierarchical `encoders/{type}/{variant}/`
Energon	Shared `data/energon/` for all models
Mock Data	Synthetic providers under `data/synthetic/`
Framework	Pure Megatron (no PyTorch Lightning dependency)

Directory structure#

High-level organization#

primus/backends/megatron/
├── core/models/diffusion/          # Core model implementations
├── training/diffusion/             # Training utilities (schedulers, etc.)
└── data/
    ├── energon/                    # Shared Energon infrastructure
    └── diffusion/                  # Diffusion-specific data

Detailed breakdown#

1. Core models (`core/models/diffusion/`)#

Following Megatron-Core convention (megatron/core/models/gpt/, megatron/core/models/multimodal/):

core/models/diffusion/
├── common/                         # Shared building blocks
│   ├── __init__.py
│   ├── config.py                   # BaseDiffusionConfig
│   ├── embeddings.py               # TimeStepEmbedder, Timesteps, MLPEmbedder
│   └── normalization.py            # RMSNorm, AdaLN, AdaLNContinuous
│
└── flux/                           # Flux-specific components
    ├── __init__.py
    ├── config.py                   # FluxConfig
    ├── model.py                    # Flux (main model class)
    ├── layer_spec.py               # MMDiTLayer, FluxSingleTransformerBlock
    ├── attention.py                # JointSelfAttention, FluxSingleAttention
    ├── layers.py                   # EmbedND (RoPE), embedders
    ├── checkpoint_converter.py     # HF <-> Megatron conversion
    └── utils.py                    # image position ids, helpers

Rationale:

common/: shared building blocks (base config, timestep/positional embeddings, normalization) reusable across diffusion models
flux/: Flux-specific code (model class, MMDiT and single-block layer specs, joint attention, RoPE embedders)
Clear separation keeps model-specific code isolated and shared code reusable for future models

2. Training (`training/diffusion/`)#

training/diffusion/
├── __init__.py
└── schedulers/
    ├── __init__.py
    ├── base.py                     # BaseScheduler
    ├── flow_matching.py            # FlowMatchEulerDiscreteScheduler
    ├── ddpm.py                     # [Future] DDPMScheduler
    ├── edm.py                      # [Future] EDMScheduler
    └── euler.py                    # [Future] EulerDiscreteScheduler

Rationale:

Schedulers define noise schedules and training targets
Separate from model code for clarity
Easy to add new schedulers (DDPM, EDM, etc.)

3. Data pipeline (`data/`)#

data/
├── energon/                        # Shared Energon utilities
│   └── __init__.py
├── dataloader.py                   # MegatronDataloaderWrapper (wraps any iterable)
│
└── diffusion/
    ├── __init__.py
    ├── encoders/                   # Hierarchical encoder registry
    │   ├── __init__.py
    │   ├── registry.py             # EncoderRegistry
    │   ├── image/
    │   │   └── vae/
    │   │       ├── __init__.py
    │   │       ├── autoencoder_kl.py   # AutoencoderKL
    │   │       └── vqvae.py            # VQVAE
    │   └── text/
    │       ├── t5/
    │       │   ├── __init__.py
    │       │   ├── t5_xxl.py           # T5-XXL
    │       │   └── t5_large.py         # T5-Large
    │       └── clip/
    │           ├── __init__.py
    │           ├── clip_l.py           # CLIP-L
    │           └── clip_h.py           # CLIP-H
    │
    ├── preprocessing/
    │   ├── __init__.py
    │   └── image/
    │       ├── __init__.py
    │       ├── transforms.py           # Resizing, normalization
    │       └── augmentation.py         # Data augmentation
    │
    └── task_encoders/
        ├── __init__.py
        └── image.py  # EncodedDiffusionTaskEncoder

Rationale:

Energon shared: VLM and other models can reuse Energon utilities
Hierarchical encoders: Organized by modality and variant type
Registry pattern: Config-driven encoder selection
Preprocessing separated: Clear pipeline stages

Architectural decisions#

Decision 1: Models under `core/models/`#

Choice: primus/backends/megatron/core/models/diffusion/ Not: primus/backends/megatron/models/diffusion/

Reasoning:

Aligns with Megatron-Core structure
Easier upstream tracking
Clear that these are Megatron-Core compatible
Consistent with existing Primus structure (core/models/gpt/)

Decision 2: Separate `common/` directory#

Choice: Shared components in common/ Not: Everything in flux/ or flat structure

Reasoning:

Standard approach puts shared code in model-specific directories
Primus: common/ makes it explicit what’s shared
Future DiT implementation trivial (reuse from common/)
Clear contract: if in common/, must work for all models

Shared Components:

JointSelfAttention: Used by DiT and Flux joint layers
FluxSingleAttention: Used by DiT and Flux single layers
MMDiTLayer: Joint (multimodal) transformer block
FluxSingleTransformerBlock: Single-modality transformer block

Flux-Only Components:

EmbedND: 3D RoPE position embedding (Flux-specific)
Flux model class

Decision 3: Hierarchical encoder structure#

Choice: encoders/image/vae/, encoders/text/t5/, encoders/text/clip/ Not: Flat encoders/conditioner.py

Reasoning:

Traditional approach uses flat structure with all encoders in one file
User requirement: support 5+ variants per modality
Registry pattern enables config-driven selection
Easy to add new encoders without modifying existing files

Registry Pattern:

from primus.backends.megatron.data.diffusion.encoders import get_encoder

# Config-driven selection
vae = get_encoder('autoencoder_kl', config=vae_config)
t5 = get_encoder('t5_xxl', config=t5_config)
clip = get_encoder('clip_l', config=clip_config)

Decision 4: Shared Energon infrastructure#

Choice: data/energon/ for shared utilities Not: Nested under data/diffusion/

Reasoning:

Energon is general-purpose (VLM, diffusion, future models)
Traditional approach nests under model-specific directories
Megatron-LM has Energon at example level (not in core)
Primus approach: shared infra + model-specific TaskEncoders

Pattern:

# Shared: data/dataloader.py
# MegatronDataloaderWrapper wraps any iterable (Energon loader, PyTorch DataLoader, etc.)
wrapper = MegatronDataloaderWrapper(dataloader)

# Model-specific: data/diffusion/task_encoders/image.py
class EncodedDiffusionTaskEncoder:
    """Diffusion-specific encoding logic"""
    pass

Decision 5: Model provider at adapter level#

Choice: construct the model through provider functions in the Megatron trainer/adapter layer under primus/backends/megatron/, above core/models/.

Reasoning:

Model providers are adapter/wrapper functions rather than a standalone module; they live alongside the trainers (for example primus/backends/megatron/megatron_pretrain_trainer.py)
They sit above the core models to add Primus-specific functionality
Example: wrap the model with a custom loss, precision handling, or checkpoint logic

Decision 6: No PyTorch lightning#

Choice: Pure Megatron patterns Not: PyTorch Lightning DataModules

Reasoning:

Primus doesn’t use PyTorch Lightning
Framework-specific implementations reduce flexibility
Better integration with Megatron training loop
Follows Megatron-LM’s MegatronDataloaderWrapper wrapper pattern

Component hierarchy#

1. Model hierarchy#

nn.Module (PyTorch)
└── MegatronModule
    └── DiffusionModule (abstract)
        └── Flux (concrete)
        ├── Joint layers: MMDiTLayer × num_joint_layers
        │   └── JointSelfAttention (shared)
        ├── Single layers: FluxSingleTransformerBlock × num_single_layers
        │   └── FluxSingleAttention (shared)
        └── Embeddings:
            ├── TimeStepEmbedder (shared)
            ├── MLPEmbedder (shared)
            └── EmbedND (Flux-specific, 3D RoPE)

2. Configuration hierarchy#

TransformerConfig (Megatron-Core)
└── BaseDiffusionConfig
    └── FluxConfig
        ├── flux_535m() factory
        └── flux_12b() factory

3. Scheduler hierarchy#

BaseScheduler (abstract)
├── FlowMatchEulerDiscreteScheduler (Flux)
├── DDPMScheduler (future)
├── EDMScheduler (future)
└── EulerDiscreteScheduler (future)

4. Encoder hierarchy#

BaseEncoder (abstract)
├── ImageEncoder
│   └── VAE
│       ├── AutoencoderKL (Flux)
│       └── VQVAE (future)
└── TextEncoder
    ├── T5
    │   ├── T5-XXL (Flux)
    │   └── T5-Large (future)
    └── CLIP
        ├── CLIP-L (Flux)
        └── CLIP-H (future)

Data flow#

Training pipeline#

Raw Data (images + captions)
    ↓
[Optional] Precalculation
    ├─ VAE → latents [B, 64, H/8, W/8]
    ├─ T5-XXL → embeddings [B, 512, 4096]
    └─ CLIP-L → pooled [B, 768]
    ↓
WebDataset/Energon Format (.tar files)
    ↓
MegatronDataloaderWrapper (from data/dataloader.py)
    ↓
EncodedDiffusionTaskEncoder (from data/diffusion/task_encoders/)
    ├─ Load precalculated data, OR
    └─ Encode on-the-fly (slower)
    ↓
Training Batch:
    ├─ latents: [B, 64, H, W]
    ├─ t5_embeddings: [B, S, 4096]
    └─ clip_pooled: [B, 768]
    ↓
FlowMatchEulerDiscreteScheduler
    ├─ Sample timesteps: t ~ U(0, 1)
    ├─ Sample noise: ε ~ N(0, I)
    ├─ Add noise: x_t = (1-t)*ε + t*x_0
    └─ Compute target: v = x_0 - ε
    ↓
Flux Model Forward Pass
    ├─ Embed timesteps
    ├─ Embed pooled text (CLIP)
    ├─ Joint layers (process latents + T5 embeddings)
    ├─ Single layers (process latents only)
    └─ Output: v_pred [B, 64, H, W]
    ↓
Loss Computation: MSE(v_pred, v_target)
    ↓
Backward Pass & Optimizer Step

Inference pipeline#

Text Prompt
    ↓
Text Encoders
    ├─ T5-XXL → embeddings [1, S, 4096]
    └─ CLIP-L → pooled [1, 768]
    ↓
Initialize Noise: x_0 ~ N(0, I)
    ↓
Sampling Loop (t = 1.0 → 0.0)
    ├─ Model forward: v_t = Flux(x_t, t, embeddings)
    ├─ Update: x_{t-dt} = x_t + v_t * dt
    └─ Repeat until t = 0
    ↓
Latents: x_0 [1, 64, H, W]
    ↓
VAE Decoder
    ↓
Generated Image [1, 3, H*8, W*8]

Comparison with alternative implementations#

Primus architectural advantages#

1. TransformerBlock architecture (Primus innovation)#

Primus: Unified TransformerBlock with heterogeneous layer specs
Others: Separate nn.ModuleList containers for double/single blocks
Benefit: Better PP slicing, unified checkpointing, future-proof

2. Megatron-core native#

Primus: Pure Megatron-Core, no framework dependencies
Others: Often integrated with PyTorch Lightning or other frameworks
Benefit: Tighter integration, simpler training loops

3. Checkpoint format#

Primus: Unified transformer.layers.{0-56} structure
Others: Separate double_blocks.{i} and single_blocks.{j}
Benefit: Simpler distributed checkpointing

4. Encoder architecture#

Primus: Registry-based, hierarchical organization
Others: Direct imports from monolithic files
Benefit: Easy extensibility for new encoder variants

File organization comparison#

Component	Standard Location	Primus Location	Improvement
Model	`models/diffusion/flux/`	`core/models/diffusion/flux/`	Megatron-Core convention
Layers	Mixed locations	`common/` for shared, `flux/` for specific	Clear boundaries
Encoders	Single file	`data/diffusion/encoders/{type}/{variant}/`	Hierarchical, extensible
Tests	Mixed with code	`tests/unit_tests/backends/megatron/diffusion/`	Proper separation

Implementation status#

Core infrastructure ✅#

✅ Directory structure
✅ Base classes (DiffusionModule, BaseDiffusionConfig, BaseScheduler)
✅ DiffusionModule with Megatron-Core integration
✅ FluxConfig with factory methods
✅ FlowMatchEulerDiscreteScheduler implementation
✅ Configuration files (YAML)
✅ Testing framework (290+ tests)
✅ Documentation structure

Flux model implementation ✅#

✅ Flux model architecture
✅ MMDiT layers and attention
✅ Embeddings (RoPE, timestep, vector)
✅ Encoder registry and loaders
✅ TaskEncoder for Energon
✅ Data pipeline

Extension points#

Adding a new model (e.g., DiT)#

Create model directory: core/models/diffusion/dit/
Add config: Extend BaseDiffusionConfig
Implement model: Extend DiffusionModule (which extends MegatronModule)
- Inherit process group management
- Get distributed checkpointing support
- Access attention backend configuration
Reuse shared components: Import from common/
Add tests: tests/unit_tests/backends/megatron/diffusion/test_dit_model.py
Update configs: Add dit_config.yaml

Adding a new encoder variant#

Create encoder file: e.g., data/diffusion/encoders/text/t5/t5_large.py
Implement encoder class: Extend BaseEncoder
Register: Add to ENCODER_REGISTRY
Add config: Update encoders.yaml
Add tests: Test in tests/unit_tests/backends/megatron/diffusion/data/encoders/

Adding a new scheduler#

Create scheduler file: training/diffusion/schedulers/ddpm.py
Implement: Extend BaseScheduler
Export: Add to __init__.py
Add tests: tests/unit_tests/backends/megatron/diffusion/training/test_scheduler.py
Document: Update this file and README

Performance considerations#

Memory optimization#

Precalculated data: 5-10x faster, lower memory
Frozen encoders: Only train diffusion model
Gradient checkpointing: Trade compute for memory
Mixed precision: bf16 on MI300X (compatible with H100/A100)

Multi-GPU scaling#

Tensor Parallelism: Split model across GPUs
Pipeline Parallelism: Split layers across GPUs
Data Parallelism: Replicate model, split data
Sequence Parallelism: For very long sequences

Best practices#

Use precalculated mode for training
Freeze encoders (standard practice)
Use bf16 on modern hardware
Start with TP=1, PP=1, scale as needed
Profile before optimizing

References#

Megatron-Core: nvidia/Megatron-LM transformer patterns
Flux: black-forest-labs/FLUX.1
Flow Matching: Rectified flow and flow matching papers
NeMo: nvidia/NeMo - Alternative diffusion implementation

Last Updated: December 2025

Architecture overview

Contents

Architecture overview#

Table of contents#

Design philosophy#

Core principles#

Architectural advantages#

Directory structure#

High-level organization#

Detailed breakdown#

1. Core models (core/models/diffusion/)#

2. Training (training/diffusion/)#

3. Data pipeline (data/)#

Architectural decisions#

Decision 1: Models under core/models/#

Decision 2: Separate common/ directory#

Decision 3: Hierarchical encoder structure#

Decision 4: Shared Energon infrastructure#

Decision 5: Model provider at adapter level#

Decision 6: No PyTorch lightning#

Component hierarchy#

1. Model hierarchy#

2. Configuration hierarchy#

3. Scheduler hierarchy#

4. Encoder hierarchy#

Data flow#

Training pipeline#

Inference pipeline#

Comparison with alternative implementations#

Primus architectural advantages#

1. TransformerBlock architecture (Primus innovation)#

2. Megatron-core native#

3. Checkpoint format#

4. Encoder architecture#

File organization comparison#

Implementation status#

Core infrastructure ✅#

Flux model implementation ✅#

Extension points#

Adding a new model (e.g., DiT)#

Adding a new encoder variant#

Adding a new scheduler#

Performance considerations#

Memory optimization#

Multi-GPU scaling#

Best practices#

References#

1. Core models (`core/models/diffusion/`)#

2. Training (`training/diffusion/`)#

3. Data pipeline (`data/`)#

Decision 1: Models under `core/models/`#

Decision 2: Separate `common/` directory#