Skip to content

Commit 8256f35

Browse files
Perseus14Perseus14
committed
Add LTX2 model to Flax/MaxDiffusion
1 parent 7a4b7a3 commit 8256f35

28 files changed

Lines changed: 4220 additions & 777 deletions

‎docs/sharding_strategy_design.md‎

Lines changed: 255 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,255 @@
1+
# Design Doc: Centralized and Configuration-Driven Sharding Strategy
2+
3+
## Objective
4+
To centralize the sharding logic in MaxDiffusion, enabling hardware-specific optimizations (e.g., for TPU v6e vs v7x) without hardcoding checks in model layers or polluting constructors with sharding parameters.
5+
6+
## Background
7+
Currently, sharding specifications are often hardcoded within model layers or determined by ad-hoc hardware checks (e.g., checking `jax.devices()[0].device_kind`). This makes the code:
8+
- **Hard to maintain and extend**: Adding support for new hardware requires modifying multiple files.
9+
- **Difficult to test**: It's hard to test different sharding strategies on the same hardware for debugging or benchmarking.
10+
- **Cluttered**: Model definition code is mixed with hardware-specific execution policies.
11+
12+
In the `prisha/ltx2_opt` branch, we see initial attempts to address this by abstracting TPU type detection, but the sharding specs themselves are still hardcoded based on the detected hardware in [attention_ltx2.py](https://github.com/AI-Hypercomputer/maxdiffusion/blob/main/src/maxdiffusion/models/ltx2/attention_ltx2.py).
13+
14+
### Proposed Design
15+
16+
We propose a design that combines **Discrete Logical Rulesets** and **Explicit Parameter Passing at the Top Level** to achieve a clean separation of concerns while adhering to JAX and Flax NNX best practices.
17+
18+
### 1. Configuration
19+
We will add a `sharding` section to the YAML configuration files, allowing independent overrides for different model components (e.g., Transformer, VAE).
20+
21+
Example in `ltx2_video.yml`:
22+
```yaml
23+
sharding:
24+
transformer: 'ironwood'
25+
vae: 'default'
26+
text_encoder: 'default'
27+
```
28+
29+
#### Auto-Detection & Backward Compatibility
30+
To improve usability and ensure backward compatibility:
31+
- **Auto-Detection**: Specifying the sharding strategy is **optional**. If omitted (or if a legacy config file lacks the `sharding` block), `pyconfig.py` will auto-detect the TPU hardware generation at startup and set the strategy to the optimal default for that chip (e.g., `'ironwood'` for v7x).
32+
- **Logging**: The resolved strategy will be explicitly logged to maintain transparency.
33+
- **Overrides**: Users can always override this auto-detection by explicitly setting the strategy in the YAML file or via CLI.
34+
35+
36+
### 2. Discrete Logical Rulesets (Model-Specific File)
37+
To keep the code simple and avoid file clutter, we organize the sharding specs into a single file per model, located in the model's directory. This keeps the sharding logic close to the model code for better readability by model developers.
38+
39+
For LTX2, this file will be `src/maxdiffusion/models/ltx2/logical_sharding_ltx2.py`.
40+
41+
This file will contain the discrete specs, the registry, and the factory function:
42+
43+
```python
44+
from dataclasses import dataclass
45+
from typing import Any, Optional
46+
47+
48+
# --- Discrete Specs ---
49+
@dataclass
50+
class LTX2DiTShardingSpecs:
51+
"""Sharding specs for the LTX2 Diffusion Transformer."""
52+
53+
qkv_kernel: tuple
54+
out_kernel: tuple
55+
out_bias: tuple
56+
norm_scale: tuple = ("norm",)
57+
embed_bias: tuple = ("embed",)
58+
59+
60+
@dataclass
61+
class TextEncoderShardingSpecs:
62+
"""Specs for the Text Encoder execution."""
63+
64+
use_batched_text_encoder: bool = False
65+
text_encoder_kernel: Optional[tuple] = None
66+
67+
68+
@dataclass
69+
class VAEShardingSpecs:
70+
"""Sharding specs for the VAE."""
71+
72+
vae_conv_kernel: Optional[tuple] = None
73+
74+
75+
# --- Unified Registry for LTX2 ---
76+
STRATEGIES = {
77+
"ironwood": {
78+
"ltx2_dit": LTX2DiTShardingSpecs(
79+
qkv_kernel=(None, "heads"),
80+
out_kernel=("heads", None),
81+
out_bias=(None,),
82+
),
83+
"text_encoder": TextEncoderShardingSpecs(
84+
use_batched_text_encoder=True,
85+
text_encoder_kernel=(None, "embed"),
86+
),
87+
"vae": VAEShardingSpecs(vae_conv_kernel=("batch", None, None, None)),
88+
},
89+
"trillium": {
90+
"ltx2_dit": LTX2DiTShardingSpecs(
91+
qkv_kernel=("embed", "heads"),
92+
out_kernel=("heads", "embed"),
93+
out_bias=("embed",),
94+
),
95+
"text_encoder": TextEncoderShardingSpecs(
96+
use_batched_text_encoder=False,
97+
text_encoder_kernel=(None, "embed"),
98+
),
99+
"vae": VAEShardingSpecs(vae_conv_kernel=(None, None, None, None)),
100+
},
101+
}
102+
103+
104+
def get_sharding_specs(strategy_name: str, component_name: str) -> Any:
105+
"""Unified factory to get specs for any component."""
106+
hardware_profile = STRATEGIES.get(strategy_name, STRATEGIES["trillium"])
107+
specs = hardware_profile.get(component_name)
108+
if specs is None:
109+
raise ValueError(f"Component {component_name} not found in strategy {strategy_name}")
110+
return specs
111+
```
112+
113+
114+
### 3. Application (Unpacking at the Top Level)
115+
116+
To avoid coupling low-level layers to model-specific strategy objects, the top-level model (e.g., `LTX2VideoTransformer3DModel`) will accept the specs object, but will **unpack** it and pass only the specific tuples or `PartitionSpec`s down to the leaf nodes (like `LTX2Attention`).
117+
118+
#### In the Pipeline
119+
The pipeline file (e.g., `ltx2_pipeline.py`) reads the strategy name from the config, retrieves the specific specs object for each component, and passes it to the respective top-level model.
120+
121+
```python
122+
# 1. Read component-specific strategy names from config
123+
sharding_config = getattr(self.config, "sharding", {})
124+
transformer_strategy = sharding_config.get("transformer", "default")
125+
te_strategy = sharding_config.get("text_encoder", "default")
126+
127+
# 2. Get the specific specs for components
128+
dit_specs = get_sharding_specs(transformer_strategy, "ltx2_dit")
129+
te_specs = get_sharding_specs(te_strategy, "text_encoder")
130+
131+
# 3. Use for pipeline execution choices
132+
if te_specs.use_batched_text_encoder:
133+
# ...
134+
135+
# 4. Pass to the top-level model
136+
self.transformer = LTX2VideoTransformer3DModel(
137+
# ...
138+
sharding_specs=dit_specs,
139+
)
140+
```
141+
142+
#### In Model Layers
143+
The top-level model receives the specs object and unpacks it for its children.
144+
145+
Example in `LTX2VideoTransformer3DModel`:
146+
```python
147+
class LTX2VideoTransformer3DModel(nnx.Module):
148+
149+
def __init__(self, ..., sharding_specs: LTX2DiTShardingSpecs):
150+
# Unpack and pass specific tuples to blocks
151+
self.block = LTX2VideoTransformerBlock(
152+
...,
153+
qkv_sharding_spec=sharding_specs.qkv_kernel,
154+
out_sharding_spec=sharding_specs.out_kernel,
155+
out_bias_sharding_spec=sharding_specs.out_bias,
156+
)
157+
```
158+
159+
Example in `LTX2Attention` (Leaf Node):
160+
```python
161+
class LTX2Attention(nnx.Module):
162+
163+
def __init__(
164+
self,
165+
...,
166+
qkv_sharding_spec: tuple,
167+
out_sharding_spec: tuple,
168+
out_bias_sharding_spec: tuple,
169+
):
170+
# Use the specific tuples directly, completely agnostic to the parent strategy
171+
self.qkv_sharding_spec = qkv_sharding_spec
172+
# ...
173+
```
174+
175+
### 4. Logical-to-Physical Mesh Mapping
176+
Logical axis names like `"heads"` and `"embed"` must be bound to a physical JAX Mesh.
177+
178+
In MaxDiffusion, this mapping is handled at the top level via `logical_axis_rules` (typically defined in the YAML config file). These rules map logical axis names to physical mesh axes (e.g., `"data"`, `"model"`, `"fsdp"`).
179+
180+
Different TPU topologies (v6e vs v7x) have different optimal physical mesh dimensions. We handle this by selecting the appropriate config file via the CLI, or by overriding the `logical_axis_rules` directly from the CLI.
181+
182+
Example of overriding `logical_axis_rules` directly via CLI:
183+
184+
```bash
185+
python src/maxdiffusion/generate_ltx2.py src/maxdiffusion/configs/ltx2_video.yml logical_axis_rules="[('heads', 'model'), ('embed', 'data')]"
186+
```
187+
188+
### 5. Startup Validation
189+
To ensure that the configuration and code are in sync, we propose adding a validation step at startup (e.g., in the pipeline or `pyconfig.py`).
190+
191+
**Problem**: If a logical sharding spec uses an axis name (e.g., `"heads"`) that is not defined in the active `logical_axis_rules`, JAX might fail late or silently fall back to suboptimal sharding.
192+
193+
**Solution**:
194+
1. Collect all logical axis names used in the active sharding strategies.
195+
2. Cross-reference them with the keys in `logical_axis_rules`.
196+
3. If any logical axis name is missing from `logical_axis_rules`, raise a `ValueError` to fail fast.
197+
4. Allow users to bypass this check with a `--skip_sharding_validation` flag if they explicitly want to proceed with potential defaults.
198+
199+
## Performance Considerations
200+
- This is purely a code-structuring change and does not introduce any runtime overhead.
201+
- The specs returned by the factory are static strings or tuples of strings, which are perfectly traced by JAX and compiled by XLA.
202+
203+
## Alternatives Considered
204+
205+
### 1. Hardcoded Hardware Checks
206+
Checking `device_kind` directly in the model components.
207+
- **Why rejected**: Scattered checks make the code hard to maintain, extend, and test.
208+
209+
### 2. Excessive Configuration/Plumbing (Pure YAML or Individual Constructor Arguments)
210+
Putting all specs in YAML or passing every spec individually through all constructors.
211+
- **Why rejected**: Leads to either bloated configuration files or polluted constructors in intermediate layers. We struck a balance by using dataclasses at the top level and unpacking them for leaf nodes.
212+
213+
### 3. Monolithic or Global State Objects
214+
Using a class-based strategy per hardware or a global singleton manager.
215+
- **Why rejected**: Leads to class explosion or violates JAX functional purity principles by introducing global state.
216+
217+
## Prototype Plan: LTX2
218+
We will use the LTX2 model as a prototype to validate this design.
219+
220+
1. **Create** the `src/maxdiffusion/models/ltx2/logical_sharding_ltx2.py` file with the specs and factory.
221+
2. **Update** [ltx2_pipeline.py](https://github.com/AI-Hypercomputer/maxdiffusion/blob/main/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py) to read the config and get the strategy.
222+
3. **Update** `transformer_ltx2.py` and `attention_ltx2.py` to accept and use the strategy object.
223+
4. **Verify** by:
224+
* Adding unit tests for the factory and strategy objects.
225+
* Running existing LTX2 integration tests with both `ironwood` and `trillium` strategies to ensure no regressions.
226+
227+
## Shared Components (e.g., `attention_flax.py`)
228+
For components shared across different models (like `NNXSimpleFeedForward` in [attention_flax.py](https://github.com/AI-Hypercomputer/maxdiffusion/blob/main/src/maxdiffusion/models/attention_flax.py)), we will pass the specific sharding specs as arguments to their constructors, and the LTX2-specific caller will fetch those values from the respective specs object.
229+
230+
## Future Expansion
231+
If the prototype succeeds on LTX2, we plan to expand this pattern to other models like **WAN** and **Flux** by adding corresponding strategies and factories.
232+
233+
## Actual Implementation (LTX2 Prototype)
234+
235+
During the implementation of the LTX2 prototype, we made some adjustments to the original design to avoid code bloat and maintain clean interfaces:
236+
237+
### 1. Passing Dataclass Objects Directly
238+
The original design suggested unpacking the `LTX2DiTShardingSpecs` at the top level and passing individual tuples to intermediate layers. As we included more parameters (MLP, embeddings, norms), this led to "argument pollution" in constructors.
239+
* **Change**: We refactored `LTX2VideoTransformerBlock`, `LTX2Attention`, and `LTX2AdaLayerNormSingle` to accept the `sharding_specs` object directly. This kept constructor signatures clean and maintainable.
240+
241+
### 2. Duck Typing for Shared Components
242+
For shared components like `NNXSimpleFeedForward` and `NNXPixArtAlphaTextProjection`, we wanted to avoid coupling them to LTX2-specific dataclasses.
243+
* **Change**: These shared components now accept a generic `sharding_specs` object (typed as `Any`) and use `getattr` with defaults to read specific attributes (e.g., `getattr(sharding_specs, "net_0_kernel", (None, "mlp"))`). This allows different models to pass their own specs objects without shared files importing model-specific code.
244+
245+
### 3. Full Coverage and Component-Specific Specs
246+
We expanded the centralization to cover all identified hardcoded sharding specs in the LTX2 model.
247+
* **Change**: We added specific specs for the text connector (`TextConnectorShardingSpecs`) and updated the VAE and Connectors to also receive their respective specs from the registry via the pipeline.
248+
249+
### 4. Configuration-Driven VAE Replication
250+
The pipeline previously hardcoded the replication of VAE weights during decoding. We made this configuration-driven by adding a `force_replication` flag to `VAEShardingSpecs` and using it in `ltx2_pipeline.py`.
251+
252+
### 5. Integration with LTX-2.3 Model Support
253+
When LTX-2.3 support was introduced to the codebase, we seamlessly integrated it into the centralized sharding registry:
254+
* **Gated Attention**: LTX-2.3 added a gated attention mechanism (`gated_attn`) in the self-attention and cross-attention blocks. We parameterized the gate projection sharding specs (`gate_logits_kernel` and `gate_logits_bias`) in both the DIT and Text Connector specs dataclasses, replacing the newly introduced hardcoded `nnx.with_partitioning` calls in `attention_ltx2.py`.
255+
* **Unified Configuration Support**: Added default sharding specs structures to the new LTX-2.3 configuration profile (`ltx2_3_video.yml`), ensuring users can cleanly configure and override sharding settings via CLI flags for LTX-2.3.

‎src/maxdiffusion/aot_cache.py‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ def transformer_forward_pass(...):
5858
import hashlib
5959
import inspect
6060
import os
61+
import json
6162
import pickle
6263
import re
6364
import threading
@@ -72,6 +73,12 @@ def transformer_forward_pass(...):
7273
_FORMAT_VERSION = 1
7374

7475

76+
def _metadata_fingerprint(meta: dict[str, Any]) -> str:
77+
"""Returns the stable filename fingerprint for install-time metadata."""
78+
serialized = json.dumps(meta, sort_keys=True, default=str)
79+
return hashlib.sha256(serialized.encode("utf-8")).hexdigest()[:12]
80+
81+
7582
def _dynamic_signature(args: tuple, kwargs: dict) -> str:
7683
"""Deterministic digest of everything that selects an executable.
7784
@@ -372,21 +379,30 @@ def install(cache_dir: str, meta: dict[str, Any], mesh: Any) -> None:
372379
mesh: The pipeline mesh; pins device order for deserialization and
373380
provides the context for re-lowering at save time.
374381
"""
375-
if not cache_dir:
376-
return
377-
os.makedirs(cache_dir, exist_ok=True)
378-
_STATE.cache_dir = cache_dir
379-
_STATE.fingerprint = hashlib.sha256(repr(sorted(meta.items())).encode()).hexdigest()[:12]
380-
_STATE.mesh = mesh
381-
_STATE.enabled = True
382+
# A process may construct multiple pipelines with different cache settings.
383+
# Finish any previous loads, then reset all install-scoped state even when
384+
# the new cache directory is empty.
385+
wait_for_loads()
386+
_STATE.cache_dir = ""
387+
_STATE.fingerprint = ""
388+
_STATE.mesh = None
389+
_STATE.enabled = False
382390
for entry in _REGISTRY:
383391
with entry._lock:
384-
# Cached state belongs to the previous install's dir/fingerprint.
385392
entry._compiled.clear()
386393
entry._out_specs.clear()
387394
entry._pending.clear()
388395
entry._adapters.clear()
389396
entry._on_disk.clear()
397+
398+
if not cache_dir:
399+
return
400+
os.makedirs(cache_dir, exist_ok=True)
401+
_STATE.cache_dir = cache_dir
402+
_STATE.fingerprint = _metadata_fingerprint(meta)
403+
_STATE.mesh = mesh
404+
_STATE.enabled = True
405+
for entry in _REGISTRY:
390406
thread = threading.Thread(target=entry.load_from_disk, name=f"aot-load-{entry.name}", daemon=True)
391407
thread.start()
392408
_LOAD_THREADS.append(thread)

‎src/maxdiffusion/configs/ltx2_3_video.yml‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,18 @@
11
#hardware
22
hardware: 'tpu'
33
skip_jax_distributed_system: False
4+
# Supported attention kernels:
5+
# dot_product, flash, tokamax_flash, tokamax_ring, tokamax_ring_custom,
6+
# ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_ring,
7+
# ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_bidir,
8+
# and cudnn_flash_te (GPU only).
49
attention: 'flash'
10+
use_base2_exp: False
11+
use_experimental_scheduler: False
12+
# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
13+
ulysses_shards: -1
14+
# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
15+
ulysses_attention_chunks: 1
516
a2v_attention_kernel: 'flash'
617
v2a_attention_kernel: 'dot_product'
718
attention_sharding_uniform: True
@@ -12,6 +23,8 @@ names_which_can_be_offloaded: []
1223
remat_policy: "NONE"
1324

1425
jax_cache_dir: ''
26+
# Local/shared filesystem directory for content-addressed PyTorch->Flax transformer weights ('' = disabled).
27+
converted_weights_dir: ''
1528
weights_dtype: 'bfloat16'
1629
activations_dtype: 'bfloat16'
1730
text_encoder_dtype: 'bfloat16'
@@ -94,7 +107,7 @@ flash_min_seq_length: 4096
94107
dcn_context_parallelism: 1
95108
dcn_tensor_parallelism: 1
96109
ici_data_parallelism: 1
97-
ici_fsdp_parallelism: 1
110+
ici_fsdp_parallelism: 1
98111
ici_context_parallelism: -1 # recommended ICI axis to be auto-sharded
99112
ici_tensor_parallelism: 1
100113
enable_profiler: False
@@ -106,6 +119,10 @@ enable_ondemand_xprof: True
106119
skip_first_n_steps_for_profiler: 0
107120
profiler_steps: 5
108121

122+
# Enable JAX named scopes for detailed profiling and debugging
123+
# When enabled, adds named scopes around key operations in transformer and attention layers
124+
enable_jax_named_scopes: False
125+
109126
replicate_vae: False
110127

111128
run_text_encoder_on_tpu: False
@@ -173,4 +190,17 @@ upsampler_temporal_patch_size: 1
173190
upsampler_adain_factor: 0.0
174191
upsampler_tone_map_compression_ratio: 0.0
175192
upsampler_rational_spatial_scale: 2.0
176-
upsampler_output_type: "pil"
193+
upsampler_output_type: "np_uint8"
194+
195+
aot_cache_dir: ''
196+
# Immutable package/build revision used when Git metadata is unavailable.
197+
aot_build_revision: ''
198+
# Tile-size auto-tuning. When enable_tile_search: True, generate_ltx2 runs a fast one-DiT-block
199+
# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and updates every
200+
# effective flash_block_sizes field with the winner. Default off (no-op).
201+
enable_tile_search: False
202+
tile_search_mode: 'smart' # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep)
203+
tile_search_iters: 10
204+
tile_search_out: '' # dir for the results CSV; '' -> print only
205+
tile_search_vmem_limit_bytes: 67108864 # 64 MiB; shared by benchmark and production custom kernels
206+
use_kv_cache: False

0 commit comments

Comments
 (0)