Contrib/ltx 2.3

[jimburtoft]

Description

Adds LTX-2.3, a 22B parameter DiT (Diffusion Transformer) for joint audio-video generation, running on Trainium 2 via NxD Inference. The model generates synchronized video + audio from text prompts using bidirectional audio-video cross-attention with flow matching.
Both the 22B DiT backbone and Gemma 3 12B text encoder run on Neuron at TP=4, sharing 4 NeuronCores and executing sequentially. CPU handles VAE decode and optional latent upscaling. Uses the native ltx-core framework (not Diffusers).
Key results:

• Single forward pass cosine similarity: 0.9999 (video), 0.9999 (audio) vs CPU reference
• Neuron Gemma3 text encoding: 6.6s (83x faster than CPU fallback at ~162s)
• Warm denoising step: ~0.3s per step (22B DiT, 8 Euler steps)
• E2E output: 25 frames @ 384x512 video + stereo 48kHz audio

Model Information

Model Name: LTX-2.3
Model Architecture: DiT (Diffusion Transformer) — 48 transformer blocks, 32 attention heads, 4096 video dim, 2048 audio dim. Bidirectional audio-video cross-attention, gated attention, QK-RMSNorm, split RoPE, flow matching.
Purpose: Joint text-to-video and text-to-audio generation with optional spatial/temporal latent upscaling.

Checklist

Please ensure your PR includes the following items. Refer to the contrib/CONTRIBUTING.md for detailed guidelines.

Required Components

• Accuracy Test (ex. test/integration/test_model.py)
  • At least one integration test that validates model accuracy
  • Uses logit validation or equivalent accuracy verification
  • Test can compile and run the model on Neuron
• README.md with the following sections:
  • Usage Example: Clear code example showing how to use the model
  • Compatibility Matrix: Table showing tested Neuron SDK versions and instance types (Trn1/Trn2/Inf2)
  • Example Checkpoints: Links to compatible model checkpoints (e.g., HuggingFace Hub)
  • Testing Instructions: Command to run the test suite for the model
• Source Code (src/)
  • Modeling code following NxD Inference patterns
  • Properly structured in the contrib folder hierarchy

Optional Components

• Unit Tests (CPU or Neuron-based)
  • Tests for individual modeling components
  • Located in test/unit/ directory

Folder Structure

Confirm your contribution follows this structure:
/contrib/models/LTX-2.3/
README.md
/src
init.py
modeling_ltx23.py # Core DiT backbone (TP sharding, DistributedRMSNorm, SDPA replacement)
modeling_gemma3_encoder.py # Custom Gemma3 encoder-only model (returns all 49 hidden states)
pipeline.py # NeuronTransformerWrapper (CPU preprocessing, backbone routing)
compile_transformer.py # DiT backbone compilation script
compile_gemma3.py # Gemma3 encoder compilation script
shard_gemma3_weights.py # Pre-shard Gemma3 weights to per-rank files
load_with_weights.py # DiT backbone weight sharding and injection
generate_ltx23.py # E2E generation pipeline
/test
init.py
/unit
init.py
/integration
init.py
test_model.py

Testing

How did you test this change?
All testing performed on trn2.3xlarge (TP=4, LNC=2) with Neuron SDK 2.28, PyTorch 2.9, Deep Learning AMI Neuron (Ubuntu 24.04) 20260227.
Tests include:

1. Integration tests (4 tests): model loading, forward pass validity, accuracy vs CPU reference, latency measurement
2. E2E generation: Full pipeline with Neuron Gemma3 text encoding → Neuron DiT denoising (8 steps) → CPU VAE decode, producing 25 video frames @ 384x512 + stereo 48kHz audio
   Test Results:
   test/integration/test_model.py::test_model_loads PASSED
   test/integration/test_model.py::test_forward_pass_no_nan PASSED
   test/integration/test_model.py::test_accuracy_vs_cpu PASSED
   test/integration/test_model.py::test_performance_latency PASSED
   ========================= 4 passed in 456.89s =========================
   Accuracy validation (measured against unsharded BF16 CPU reference):
   | Component | Metric | Value |
   |-----------|--------|-------|
   | Single forward pass (video) | Cosine similarity | 0.999947 |
   | Single forward pass (audio) | Cosine similarity | 0.999867 |
   | 8-step denoised latent (real text) | Cosine similarity | 0.972 |

Compatibility

Tested with:

• Neuron SDK Version(s): 2.28
• Instance Type(s): trn2.3xlarge (TP=4, LNC=2)
• PyTorch Version: 2.9
• Python Version: 3.12

Additional Information

• Two Neuron models share NeuronCores: Both the 22B DiT backbone and 12B Gemma3 text encoder are compiled for TP=4 and share the same 4 NeuronCores, executing sequentially (text encoding runs once, then denoising loop runs 8 steps).
• Different compiler flags per model: DiT uses --auto-cast matmult for throughput; Gemma3 uses --auto-cast=none --enable-saturate-infinity --enable-mixed-precision-accumulation for text encoder precision.
• Cold start: First two denoising steps are slow (~144s + ~177s) due to Neuron device initialization. Warm steps run at ~0.3s each.
• CPU fallback available: Gemma3 can run on CPU (~162s) without Neuron compilation, using --gemma-path without --neuron-gemma.
• Stage 1 only: This submission covers Stage 1 generation with optional latent upscaling. Stage 2 refinement denoising (recompiling at larger latent shapes) is planned as a follow-up.
• Uses native ltx-core framework, not Diffusers.

Related Issues

N/A — first contribution for this model.

vLLM Integration

• This model/feature is intended for use with vLLM
• Documentation includes vLLM registration instructions

---

By submitting this PR, I confirm that:

• I have read and followed the contributing guidelines
• This is a community contribution and may have limited testing compared to officially-supported models
• The code follows best practices and is well-documented
• All required components listed above are included