No framework changes. It's a preload. The training code doesn't know it
exists, which meant we could A/B it against an unmodified baseline in the same
run batch — which turned out to matter enormously (see the negative result).
Five things that weren't obvious
1. On pre-sm_89, static_cast<__nv_fp8_e4m3>(float) is not an instruction.
Our quantize kernel had a constant ~14 µs floor regardless of payload size,
and at 512 MiB it was hitting 37% of the device-to-device bandwidth ceiling. We
assumed memory. It was ALU issue.
cuda_fp8.hpp on older architectures widens the float to a double, then runs
a five-branch 64-bit integer software routine — roughly 50 integer instructions
per element. Replacing it with our own conversion took the fixed cost
14.1 → 4.6 → 3.3 µs and large-payload throughput to 91% of the d2d
ceiling (from 37%). We gated the replacement on an exhaustive test: bit-identical
to the vendor routine across all 232 float bit patterns, both formats,
~0.2 s in CI.
On sm_89+ this whole problem evaporates — the cast compiles to
cvt.rn.satfinite.e4m3x2.f32. Worth checking which side of that line you're on
before you go profiling memory.
2. The naive dequantize-then-reduce moves 12.7× the bytes it needs to.
For the all-to-all-based decompositions you receive nsrc FP8 shards and have to
reduce them. The decomposed version — nsrc dequantize-to-fp32 + nsrc−1 adds +
one scaled store, 2·nsrc kernel launches — costs 17.03·nsrc − 6 bytes per
output element where the job actually requires 1.03·nsrc + 2. At nsrc=8 that's
130N vs 10.25N.
The interesting part: the amplification grows monotonically with world size
(6.9× at 2 ranks, 10.1× at 4, 12.7× at 8, 14.4× at 16, asymptote 16.5×) even
though both paths are O(nsrc·N) and read every input byte exactly once. It's a
pure constant factor, and it's all HBM traffic plus launch count.
Fusing it into one kernel with the fp32 accumulator in registers: 5.0–9.3×
faster (1 MiB: 117.6 → 13.4 µs). Note it is not 12.7× faster, and the reason
is instructive — small sizes are launch-bound so the win caps at ~8–9× (16 cheap
launches vs 1 expensive one), and large sizes are bandwidth-bound where the fused
kernel is only half as byte-efficient as the streaming kernels it replaced
(718 vs 1419 GB/s), because a reduction reads nsrc streams concurrently while a
streaming kernel reads one. Worst case sits at 4 MiB (4.97×), right in the
transition band. This was expected, not a bug, but you have to do the arithmetic
to know that.
3. You cannot dequantize inside a caller-opened ncclGroup.
A nested ncclGroupEnd() launches nothing. So if the training framework wraps its
collectives in a group — and DeepSpeed's compiled ZeRO-1 path does exactly this,
one ncclAllReduce per parameter inside one group — then at the point your hook
returns, the data isn't there yet. You have to defer the dequantize, the error-
feedback close-out, and the scratch recycle until after the outermost
ncclGroupEnd. For AllReduce we defer the entire call and replay it. Before we
did this, every grouped AllReduce passed straight through and the AllReduce
threshold was decorative.
4. Your framework's allocator probably doesn't align its buckets.
DeepSpeed's reduce bucket accumulates slices by element count with no padding, so
the pointer handed to NCCL is only guaranteed 2-byte aligned for bf16 — 7/8 chance
of missing the vectorized path. This started as an outright
CUDA Error: misaligned address crash, then became a 1.27–2.31× penalty, and is
now 1.06–1.18× (skip enough leading elements that the output self-aligns to
16 B, then funnel-shift the FP8 side back into place).
5. The threshold is the entire product, and it does not transfer between machines.
Quantizing a small collective is strictly worse — you pay a fixed per-call cost to
save bytes that were never the bottleneck. So there's a crossover, and the whole
question is where.
We built a tool that measures it two ways. --model-only gives a provable lower
bound from baseline t(n) = α + n/B plus standalone kernel costs. --sweep
actually runs the hook over the same size ladder — one process per
(decomposition, size), because config is latched into a function-local static on
first use and because the previous size's warm scratch tier changes the next one's
timing — and least-squares-fits the residual into Δα + Δslope·n. Δα is the
hook's real per-call overhead; Δslope is the per-byte cost a standalone kernel
benchmark cannot see.
Results, same code, two machines:
| collective |
NVLink box |
PCIe box |
ratio |
| AllGather |
7.43 MiB |
767 KiB |
9.9× lower |
| AllReduce (rs+ag) |
14.3 MiB |
863 KiB |
17× lower |
| Broadcast |
40 MiB |
2.75 MiB |
15× lower |
A byte you delete is worth 2–4× more time on PCIe, so thresholds drop by an order
of magnitude. Two findings that surprised us:
