File size: 6,457 Bytes
0a1df87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | # TT-DiT: Tenstorrent Diffusion Transformers
This directory contains the implementation of the Tenstorrent Diffusion Transformer (TT-DiT) architecture, designed for optimized parallelism of DiTs for image and video generation.
## Supported Models
For detailed information about each model including performance metrics, usage instructions, and specific requirements, see:
- **[Stable Diffusion 3.5 Large](models/StableDiffusion35.md)** - Text-to-image generation
- **[Flux 1](models/Flux1.md)** - Text-to-image generation (schnell & dev variants)
- **[Motif](models/Motif.md)** - Text-to-image generation model
- **[Qwen-Image](models/QwenImage.md)** - Text-to-image generation model
- **[Ideogram 4.0](models/Ideogram4.md)** - Text-to-image generation model (best-in-class typography; Blackhole)
- **[Mochi-1](models/Mochi_1.md)** - Video generation model
- **[Wan2.2-T2V-A14B](models/Wan2_2.md)** - Text-to-video generation model
- See **[experimental/](experimental/)** for in-progress model variants (AniSora, Wan2.2-Distill, LoRA adapters)
- **[LTX-2.3](models/LTX2.md)** - Joint audio-video generation model (Pro and Fast variants)
## Directory Structure
```
tt_dit/
├── layers/ # Core neural network layers
├── models/ # Model architectures and documentation
│ ├── transformers/ # Transformer implementations (SD35, Mochi, Wan, LTX, Flux1, Motif, QwenImage)
│ ├── vae/ # VAE/Autoencoder implementations
│ ├── StableDiffusion35.md # SD3.5 model documentation
│ ├── Flux1.md # Flux 1 model documentation
│ ├── Motif.md # Motif model documentation
│ ├── QwenImage.md # Qwen-Image model documentation
│ ├── Mochi_1.md # Mochi-1 model documentation
│ ├── Wan2_2.md # Wan2.2 model documentation
│ └── LTX2.md # LTX-2.3 model documentation
├── encoders/ # Text encoder implementations
│ ├── clip/ # CLIP encoder
│ ├── gemma/ # Gemma encoder
│ └── t5/ # T5 encoder
├── parallel/ # Parallelization utilities
│ ├── config.py # Parallel configuration
│ └── manager.py # Parallel execution management
├── pipelines/ # End-to-end model pipelines
│ ├── stable_diffusion_35_large/
│ ├── mochi/
│ ├── wan/
│ ├── ltx/
│ ├── flux1/
│ ├── motif/
│ └── qwenimage/
├── tests/ # Test suite
│ ├── models/ # Model-level tests (sd35, mochi, wan2_2, ltx, flux1, motif, qwenimage)
│ ├── encoders/ # Encoder tests
│ ├── blocks/ # Block-level tests
│ └── unit/ # Unit tests for layers
└── utils/ # Utility functions
├── check.py # Validation utilities
├── padding.py # Padding operations
├── substate.py # State management
└── tensor.py # Tensor operations
```
## Core Components
### Layers
- **Embeddings**: Position and token embedding implementations
- **Feedforward**: MLP and feed-forward network components
- **Linear**: Optimized linear transformation layers
- **Normalization**: Layer normalization implementations
### Models
- **Transformers**: DiT transformer implementations for various generative models
- Support for multiple architectures including SD3.5, Mochi, Wan2.2, LTX-2, Flux1, Motif, and Qwen-Image
- Model-specific attention mechanisms and transformer architectures
- **Autoencoders**: VAE implementations for different models
### Parallel Processing
- **Config**: Parallel configuration management with `DiTParallelConfig`
- **Manager**: Parallel execution and device management
### Pipelines
End-to-end pipeline implementations for multiple generative models:
- **Stable Diffusion 3.5 Large**: Text-to-image generation (1024x1024px)
- **Flux 1**: Text-to-image generation (schnell & dev variants, 1024x1024px)
- **Motif**: Text-to-image generation (6B model, 1024x1024px)
- **Qwen-Image**: Text-to-image generation (1024x1024px)
- **Mochi-1**: Video generation model (824x480px, 168 frames)
- **Wan2.2-T2V-A14B**: Text-to-video generation
- **LTX-2.3**: Joint audio-video generation (Pro one-stage and Fast distilled two-stage)
Each pipeline includes:
- Automatic parallel configuration for different device meshes
- Optimized execution on Wormhole systems
- Comprehensive timing collection and profiling
## Testing
The test suite is organized into three main categories:
1. **Unit Tests** (`tests/unit/`):
- Layer-level testing
- Individual component validation
- Coverage for embeddings, feedforward, linear, and normalization layers
2. **Model Tests** (`tests/models/`):
- Smoke tests for model components
- Performance testing
- Pipeline validation
Running tests:
```bash
# Run tests from tt-metal/
# Run unit tests
python -m pytest models/tt_dit/tests/unit/
# Run all model tests
python -m pytest models/tt_dit/tests/models/
# Run specific model pipeline tests
python -m pytest models/tt_dit/tests/models/sd35/test_pipeline_sd35.py -v
python -m pytest models/tt_dit/tests/models/flux1/test_pipeline_flux1.py -v
python -m pytest models/tt_dit/tests/models/motif/test_pipeline_motif.py -v
python -m pytest models/tt_dit/tests/models/qwenimage/test_pipeline_qwenimage.py -v
python -m pytest models/tt_dit/tests/models/mochi/test_pipeline_mochi.py -v
python -m pytest models/tt_dit/tests/models/wan2_2/test_pipeline_wan.py -v
python -m pytest models/tt_dit/tests/models/ltx/test_pipeline_ltx_distilled.py -v
python -m pytest models/tt_dit/tests/models/ltx/test_pipeline_ltx_two_stages.py -v
```
## Key Features
- **Modular Architecture**: Clean separation between layers, models, and parallel processing
- **Optimized Parallelism**:
- Automatic configuration of tensor and sequence parallelism
- Efficient device mesh management
- Dedicated parallel managers for different model components
- **Performance Monitoring**:
- Built-in timing collection
- Performance profiling support
- Tracing capabilities for optimization
## Contributing
When adding new components:
1. Place implementation in appropriate subdirectory
2. Add corresponding unit tests
3. Update model tests if necessary
4. Ensure parallel processing compatibility
|