OTIL for SDXL on 4x L40S

Implementation of OTIL (Only Transmit Informative Latents) from OTIL_IMPLEMENTATION_GUIDE.md, layered on top of MIT-Han-Lab DistriFusion.

Everything here is additive. distrifuser is pip-installed from git and never edited; the behaviour changes the guide asks for are applied at runtime from otil/patches.py. That keeps the upstream checkout clean and re-pullable.

---

Quick start

# 1. environment (once)
bash env/setup_env.sh
conda activate otil                 # or: source .venv/bin/activate

# 2. weights (once, single process -- 4 ranks racing the HF cache is the classic
#    way one of these jobs dies on a shared filesystem)
export HF_HOME=/scratch/$USER/hf
python env/fetch_models.py

# 3. sanity checks (seconds, CPU only, no GPU needed)
python scripts/selftest_otil.py
python scripts/selftest_schedule.py

# 4. the sweep
sbatch env/otil.sbatch

Edit env/otil.sbatch before submitting: uncomment and fill in --partition / --account, and point HF_HOME at real scratch space.

Results land in results/:

---

Presets

guide reproduces guide sections 4.1-4.3 exactly. otil is the same thing with the Pitfall 3 remedies actually applied. Every field is overridable:

torchrun --standalone --nproc_per_node=4 scripts/run_otil.py \
    --mode generation --preset guide --otil_gn_sync --otil_k 32 \
    --output_dir results/experiment

Run python scripts/run_otil.py --help for the full flag list.

---

Read this before interpreting the numbers

Three things about the guide are worth knowing up front. All three were found by reading the upstream sources, and two of them are handled in code here.

1. GPU count does not mean what you would expect

DistriConfig computes n_device_per_batch = world_size // 2 whenever classifier-free guidance and split_batch are both on, and patch parallelism only engages when that value exceeds 1:

So on 2 GPUs there is no patch parallelism at all and no image degradation is even possible -- which is why the sweep runs only one 2-GPU row instead of four. Passing --no_split_batch trades the batch split for more spatial splitting, giving a 4-way split on 4 GPUs (three slice seams instead of one). The sweep covers both.

2. The guide breaks classifier-free guidance on 4 GPUs

With split_batch on, output_buffer holds both CFG halves, filled by two separate torch.cat calls. Guide 4.2 writes only the local spatial slice of the local batch half, so after warmup ranks 0 and 1 keep a noise_pred_text frozen at warmup step 4, and the pipeline then computes uncond + scale * (text - uncond) against stale data.

otil/patches.py::install_cfg_repair_hook fixes it with one all_gather over the already-existing distri_config.split_group -- the two ranks holding the same spatial slice in opposite batch halves -- costing about 64 KB per step.

To see the failure for yourself:

torchrun --standalone --nproc_per_node=4 scripts/run_otil.py --mode generation \
    --preset guide --no_cfg_repair --output_dir results/cfg_bug

3. Skipping the noise-prediction gather saves ~5 ms, not more

The tensor guide 4.2 stops communicating is the noise prediction at latent resolution: 2 x 4 x 128 x 128 fp16 = 256 KB, roughly 100 us per step, about 5 ms across 50 steps.

DistriFusion's expensive traffic is somewhere else entirely -- the attention KV and conv halo buffers, on the order of 200 MB per step -- and that is removed by mode="no_sync" alone, not by OTIL. On an L40S node (PCIe, no NVLink) that is where the speedup comes from.

The consequence: OTIL's selective transmit replaces a nearly-free collective with a slightly cheaper one, while leaving remote latent regions stale. That is why balanced keeps the gather and spends the budget on quality instead, and why the sweep measures all of them rather than assuming. Decide from results/, not from the guide's table.

---

What is implemented where

Differences from the guide's reference code

Vectorised selection. The guide's 4.1 listing calls .item() inside three Python loops (K scoring, k packing, (N-1)*k scatter) -- roughly 128 device synchronisations per step, 6400 across a 50-step generation. Here scoring is one topk, packing is one index_select, and reconstruction is one index_copy_ per remote device. A step costs zero synchronisations and exactly two collectives.

The dynamic-polling rule is unchanged, just expressed differently: "unvisited blocks first, then highest dissimilarity" is exactly topk(dissimilarity + 10 * unvisited, k), because cosine dissimilarity is bounded by [0, 2]. scripts/selftest_otil.py asserts the resulting selection covers every sub-block within ceil(K / k) steps.

Periodic layer-context refresh. Pitfall 3 notes that conv halos and cross-slice attention KV stay frozen at warmup step 4 forever. The fix needs no new communication code at all: unet.set_counter(0) already drives every DistriConv2dPP, DistriSelfAttentionPP and DistriGroupNorm onto its warmup branch for one step, and also selects CUDA graph index 0, so it works with graphs enabled. The counter is restored afterwards, which is why otil/schedule.py tracks the logical step index rather than reading unet.counter the way the guide's snippet does.

GroupNorm sync without the rest. Upstream's sync_gn mode also re-enables the conv and attention enqueue traffic, which is the exact thing OTIL removes. The patched DistriGroupNorm.forward adds one more way into the synchronised branch so a run can keep no_sync layer communication and still all-reduce the two scalars per group that decide slice brightness and contrast. Every other branch is byte-identical to upstream.

---

Self-tests

Both run on CPU in seconds and need neither a GPU nor distrifuser. Run them after any edit to otil/comm.py or otil/schedule.py.

python scripts/selftest_otil.py                      # 2 ranks, gloo
OTIL_SELFTEST_WORLD=4 python scripts/selftest_otil.py # 4-way spatial split
python scripts/selftest_schedule.py

selftest_otil.py checks the reconstruction rule from guide 3.5: the local slice comes out bit-exact fresh, every remote position is either a received sub-block or stale memory and never a blend, exactly k_transmit blocks arrive from each peer, and finalize() recovers the exact latent on every rank.

That last one matters. Ranks are expected to disagree in remote regions -- each device keeps its own patch in full but receives only the top-k from everyone else -- so the guarantee that makes it safe is that the pre-decode gather is exact, not that the intermediate latents match.

---

Manual runs

# single-GPU reference
torchrun --standalone --nproc_per_node=1 scripts/run_otil.py \
    --mode generation --output_dir results/ref

# 4 GPUs, guide-exact
torchrun --standalone --nproc_per_node=4 scripts/run_otil.py \
    --mode generation --preset guide --output_dir results/guide

# latency only, VAE decode excluded
torchrun --standalone --nproc_per_node=4 scripts/run_otil.py \
    --mode benchmark --preset balanced --repeats 3 --output_dir results/bench

# a subset of the sweep
python scripts/otil_sweep.py --plan 4:guide,4:balanced,4:balanced:nosplit \
    --output_dir results/

The seam check worth doing by eye: prompt 05 (a matte ceramic vase on a seamless pale grey background) is a flat field, which is where independent per-slice GroupNorm shows up as a horizontal contrast step at H/2. Compare results/4gpu_guide/04.png against results/4gpu_otil/04.png -- otil differs from guide only by GroupNorm sync and the periodic refresh.

---

Environment notes

Versions are pinned, and not cosmetically. distrifuser requires diffusers==0.24.0, which is also the version whose SDXL pipeline exposes callback_on_step_end -- the hook OTIL's step-level sync is built on. diffusers 0.24 imports cached_download from huggingface_hub, which was removed in hub 0.26, so the hub is pinned to 0.20.3.

setup_env.sh installs distrifuser with --no-deps on purpose: its install_requires lists unpinned transformers and torch>=2.2 and would otherwise replace the versions installed just above it. The resolved upstream commit is recorded in env/distrifuser.lock.

The UNet is loaded with variant="fp16", which upstream's from_pretrained does not do -- without it you download the ~10 GB fp32 checkpoint and then never use it.

Troubleshooting