h3-blackwell-runtime/DISTRIBUTED.md

135 lines
4.8 KiB
Markdown
Raw Normal View History

# Distributed H3 Execution
The runtime supports two single-node distributed denoiser modes:
- `ulysses`: token-sharded blocks with QKV sequence-to-head all-to-all and the
inverse head-to-sequence all-to-all around attention.
- `tensor`: token-sharded blocks plus true NVFP4 tensor parallelism across QKV,
attention output, MLP FC1, and MLP FC2.
Both modes preserve the native packed H3 sequence and reject empty token or head
partitions. They do not add semantic padding. Temporary collective padding used
by ragged row gathering is removed before model operations.
## Partition Contract
H3 has 56 attention heads of width 128. Head ownership is balanced and can be
ragged:
| GPUs | Heads per rank |
| --- | --- |
| 2 | 28, 28 |
| 4 | 14, 14, 14, 14 |
| 6 | 10, 10, 9, 9, 9, 9 |
| 8 | 7, 7, 7, 7, 7, 7, 7, 7 |
Token ranges use the same quotient/remainder partitioning. Segment boundaries
are clipped to each rank's token interval and rebased before AdaLN modulation
and residual gating.
In tensor mode, QKV output rows are selected by local head ownership. Attention
output and MLP FC2 input columns are sliced on NVFP4 alignment boundaries; each
rank computes a partial output and a ragged reduce-scatter sums and assigns token
rows. MLP FC1 selects corresponding local ranges from both the gate and value
halves. Projection bias is omitted from rank partials and added exactly once
after reduction.
## Launch
The launcher uses every visible GPU when the world size is omitted:
```bash
tools/run_distributed_t2va.sh ulysses
tools/run_distributed_t2va.sh tensor
```
Pass an explicit count and attention backend when needed:
```bash
tools/run_distributed_t2va.sh tensor 8 sdpa
```
Relevant environment variables:
- `H3_WORLD_SIZE`: fallback world size when no positional count is supplied.
- `H3_DISTRIBUTED_BENCHMARK`: benchmark JSON path.
- `H3_DISTRIBUTED_OUTPUT`: report and optional latent output directory.
- `H3_MODEL_PATH` and `H3_TEXT_ENCODER_PATH`: checkpoint paths.
- `H3_SAVE_LATENTS=0`: write reports without retaining large latent files.
Run every feasible target count on the current machine with:
```bash
H3_SAVE_LATENTS=0 tools/run_distributed_matrix.sh
```
The default matrix runs Ulysses and tensor modes at 1, 2, 4, 6, and 8 GPUs and
skips counts larger than the visible device count. Override the lists with
`H3_GPU_COUNTS` and `H3_DISTRIBUTED_MODES`.
The lower-level transport benchmark does not load H3 weights:
```bash
tools/run_ulysses_benchmark.sh
tools/run_ulysses_benchmark.sh 6 sdpa
```
## Validation
Automated Gloo tests cover:
- Ulysses transport identity at 2, 4, 6, and 8 ranks.
- Two-rank distributed SDPA parity.
- Ragged gather and reduce-scatter behavior.
- Distributed final-projection parity.
- TP attention and MLP parity at 2 and 6 ranks.
- Packed NVFP4 column- and row-shard layout preservation.
On GB10, the real 50-block NVFP4 checkpoint completed both one-rank distributed
paths at 864x480, 141 frames, 12 steps, seed 440420. Ulysses and tensor modes
produced identical video and audio tensors with zero maximum absolute error.
This validates integration and the world-size-one identity path, but it does not
replace multi-GPU NCCL parity testing.
## RunPod
The target is one eight-GPU RTX PRO 6000 Blackwell machine. Query current stock
using the guarded API v2 client:
```powershell
$env:RUNPOD_API_KEY = "..."
python .\tools\runpod_api.py catalog --count 8
```
After choosing an available data center, create the pod explicitly:
```powershell
python .\tools\runpod_api.py create --count 8 --datacenter US-XX-N --yes
```
The client defaults to the server-edition RTX PRO 6000 Blackwell and RunPod's
x86_64 CUDA 13.0, Torch 2.9.1 image. It creates persistent workspace storage and
enables SSH. Creation and termination require `--yes` to avoid accidental spend
or data loss. Use `get` to poll status and `terminate POD_ID --yes` when finished.
Transfer this checkout plus these two existing checkpoints to the pod workspace:
- `minimax_h3_fl2va_pruned_nvfp4.safetensors` (12,528,636,800 bytes)
- `qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors` (15,687,142,551 bytes)
Install the project dependencies in the RunPod image, set the checkpoint path
variables, run `test_distributed.py`, then run the transport and generation
matrices. `Dockerfile.runpod` provides the equivalent reproducible custom image
once an x86_64 image builder and registry are available.
SageAttention2 is not currently packaged in the generic x86 image, so cloud
correctness and scaling start with SDPA. Sage2 can be measured after an x86
wheel is added without changing the distributed layout.
## Remaining GPU Gates
- Real NCCL identity at 2, 4, 6, and 8 GPUs.
- Distributed-versus-single latent parity above one rank.
- Full 1/2/4/6/8 timing, transport, and memory reports.
- Quality comparison after selecting an x86 attention backend.