|
| 1 | +# LTX-2 Distillation Training with ModelOpt |
| 2 | + |
| 3 | +Knowledge distillation for LTX-2 DiT models using NVIDIA ModelOpt. A frozen **teacher** guides a trainable **student** through a combined loss: |
| 4 | + |
| 5 | +```text |
| 6 | +L_total = α × L_task + (1-α) × L_distill |
| 7 | +``` |
| 8 | + |
| 9 | +Currently supported: |
| 10 | + |
| 11 | +- **Quantization-Aware Distillation (QAD)** — student uses ModelOpt fake quantization |
| 12 | + |
| 13 | +Planned: |
| 14 | + |
| 15 | +- **Sparsity-Aware Distillation (SAD)** — student uses ModelOpt sparsity |
| 16 | + |
| 17 | +## Installation |
| 18 | + |
| 19 | +```bash |
| 20 | +# From the distillation example directory |
| 21 | +cd examples/diffusers/distillation |
| 22 | + |
| 23 | +# Install Model-Optimizer (from repo root) |
| 24 | +pip install -e ../../.. |
| 25 | + |
| 26 | +# Install all dependencies (ltx-trainer, ltx-core, ltx-pipelines, omegaconf) |
| 27 | +pip install -r requirements.txt |
| 28 | +``` |
| 29 | + |
| 30 | +## Quick Start |
| 31 | + |
| 32 | +### 1. Prepare Your Dataset |
| 33 | + |
| 34 | +Use the ltx-trainer preprocessing to extract latents and text embeddings: |
| 35 | + |
| 36 | +```bash |
| 37 | +python -m ltx_trainer.preprocess \ |
| 38 | + --input_dir /path/to/videos \ |
| 39 | + --output_dir /path/to/preprocessed \ |
| 40 | + --model_path /path/to/ltx2/checkpoint.safetensors |
| 41 | +``` |
| 42 | + |
| 43 | +### 2. Configure |
| 44 | + |
| 45 | +Copy and edit the example config: |
| 46 | + |
| 47 | +```bash |
| 48 | +cp configs/distillation_example.yaml configs/my_experiment.yaml |
| 49 | +``` |
| 50 | + |
| 51 | +Key settings to update: |
| 52 | + |
| 53 | +```yaml |
| 54 | +model: |
| 55 | + model_path: "/path/to/ltx2/checkpoint.safetensors" |
| 56 | + text_encoder_path: "/path/to/gemma/model" |
| 57 | + |
| 58 | +data: |
| 59 | + preprocessed_data_root: "/path/to/preprocessed/data" |
| 60 | + |
| 61 | +distillation: |
| 62 | + distillation_alpha: 0.5 # 1.0 = pure task loss, 0.0 = pure distillation |
| 63 | + quant_cfg: "FP8_DEFAULT_CFG" # or INT8_DEFAULT_CFG, NVFP4_DEFAULT_CFG, null |
| 64 | + |
| 65 | +# IMPORTANT: disable ltx-trainer's built-in quantization |
| 66 | +acceleration: |
| 67 | + quantization: null |
| 68 | +``` |
| 69 | +
|
| 70 | +### 3. Run Training |
| 71 | +
|
| 72 | +#### Single GPU |
| 73 | +
|
| 74 | +```bash |
| 75 | +python distillation_trainer.py --config configs/my_experiment.yaml |
| 76 | +``` |
| 77 | + |
| 78 | +#### Multi-GPU (Single Node) with Accelerate |
| 79 | + |
| 80 | +```bash |
| 81 | +accelerate launch \ |
| 82 | + --config_file configs/accelerate/fsdp.yaml \ |
| 83 | + --num_processes 8 \ |
| 84 | + distillation_trainer.py --config configs/my_experiment.yaml |
| 85 | +``` |
| 86 | + |
| 87 | +#### Multi-node Training with Accelerate |
| 88 | + |
| 89 | +To launch on multiple nodes, make sure to set the following environment variables on each node: |
| 90 | + |
| 91 | +- `NUM_NODES`: Total number of nodes |
| 92 | +- `GPUS_PER_NODE`: Number of GPUs per node |
| 93 | +- `NODE_RANK`: Unique rank/index of this node (0-based) |
| 94 | +- `MASTER_ADDR`: IP address of the master node (rank 0) |
| 95 | +- `MASTER_PORT`: Communication port (e.g., 29500) |
| 96 | + |
| 97 | +Then run this (on every node): |
| 98 | + |
| 99 | +```bash |
| 100 | +accelerate launch \ |
| 101 | + --config_file configs/accelerate/fsdp.yaml \ |
| 102 | + --num_machines $NUM_NODES \ |
| 103 | + --num_processes $((NUM_NODES * GPUS_PER_NODE)) \ |
| 104 | + --machine_rank $NODE_RANK \ |
| 105 | + --main_process_ip $MASTER_ADDR \ |
| 106 | + --main_process_port $MASTER_PORT \ |
| 107 | + distillation_trainer.py --config configs/my_experiment.yaml |
| 108 | +``` |
| 109 | + |
| 110 | +**Config overrides** can be passed via CLI using dotted notation: |
| 111 | + |
| 112 | +```bash |
| 113 | +accelerate launch ... distillation_trainer.py \ |
| 114 | + --config configs/my_experiment.yaml \ |
| 115 | + ++distillation.distillation_alpha=0.6 \ |
| 116 | + ++distillation.quant_cfg=INT8_DEFAULT_CFG \ |
| 117 | + ++optimization.learning_rate=1e-5 |
| 118 | +``` |
| 119 | + |
| 120 | +## Configuration Reference |
| 121 | + |
| 122 | +### Calibration |
| 123 | + |
| 124 | +Before training begins, calibration runs full denoising inference to collect activation statistics for accurate quantizer scales. This is cached as a step-0 checkpoint and reused on subsequent runs. |
| 125 | + |
| 126 | +| Parameter | Default | Description | |
| 127 | +|-----------|---------|-------------| |
| 128 | +| `calibration_prompts_file` | null | Text file with one prompt per line. Use the HuggingFace dataset 'Gustavosta/Stable-Diffusion-Prompts' if null. | |
| 129 | +| `calibration_size` | 128 | Number of prompts (each runs a full denoising loop) | |
| 130 | +| `calibration_n_steps` | 30 | Denoising steps per prompt | |
| 131 | +| `calibration_guidance_scale` | 4.0 | CFG scale (should match inference-time) | |
| 132 | + |
| 133 | +### Checkpoint Resume |
| 134 | + |
| 135 | +| Parameter | Default | Description | |
| 136 | +|-----------|---------|-------------| |
| 137 | +| `resume_from_checkpoint` | null | `"latest"` to auto-detect, or explicit path | |
| 138 | +| `must_save_by` | null | Minutes after which to save and exit (for Slurm time limits) | |
| 139 | +| `restore_quantized_checkpoint` | null | Restore a pre-quantized model (skips calibration) | |
| 140 | +| `save_quantized_checkpoint` | null | Path to save the final quantized model | |
| 141 | + |
| 142 | +### Custom Quantization Configs |
| 143 | + |
| 144 | +To define custom quantization configs, add entries to `CUSTOM_QUANT_CONFIGS` in `distillation_trainer.py`: |
| 145 | + |
| 146 | +```python |
| 147 | +CUSTOM_QUANT_CONFIGS["MY_FP8_CFG"] = { |
| 148 | + "quant_cfg": mtq.FP8_DEFAULT_CFG["quant_cfg"], |
| 149 | + "algorithm": "max", |
| 150 | +} |
| 151 | +``` |
| 152 | + |
| 153 | +Then reference it in your YAML: `quant_cfg: MY_FP8_CFG`. |
0 commit comments