transformers
84ee0694 - tp loading: batch coalesce scatter/broadcast via _coalescing_manager

Commit
80 days ago
tp loading: batch coalesce scatter/broadcast via _coalescing_manager Wrap each batch of 64 mappings' redistribute calls in torch.distributed.distributed_c10d._coalescing_manager with device=local_device, which lowers per-call launch overhead via ncclGroupStart/ncclGroupEnd. Falls back to the plain loop on gloo (CPU synthetic test) since gloo has no coalescing primitive. Also threads mapping.convert() across a small thread pool (up to 4 workers) so the batch's converts run concurrently with each other while the previous batch's scatter is still in-flight. Added VIZTRACER_OUTPUT env var to the benchmark script for rank-0 profiling of the loading path. Results (8×B200, tp_plan=auto): Qwen2.5-7B (ws=4): 13.21s (main: 10.42s) Qwen2.5-14B (ws=8): 22.94s (main: 12.54s) Qwen3-30B-A3B MoE (ws=4): 14.97s (main: 13.56s) viztracer shows the per-collective launch overhead is gone; the remaining gap vs main is the coalesced NCCL wait itself (~8.5s on 14B). NVLink is not saturated — further gains likely need grouping multiple tensors into a single scatter payload.
Author
Parents
Loading