NCCL-BOT: Closing the Loop on NCCL's Cost Model with Online Bandits

원문 초록

NCCL picks an algorithm-protocol pair for every collective from an analytic cost model (AUTO) that never sees a measured latency. The model is a strong prior but a brittle final decision: priced on forced-arm runs, its choice is exact on most contexts and badly wrong in a few. We built NCCL-BOT, a tuner plugin that closes the loop: a per-context bandit, warm-started from the model, learns from measured latencies. We replay a forced-arm H100 sweep, tuned at 4-8 nodes and scored at 16-32, and price the model's choice on the same forced arms; that sweep has no stock-NCCL run. The bandit cuts the contexts where the model's choice loses more than 10% from 19 to 1 (p99 103% to 9.9%, mean 6.1% to 1.6%; p=0.047 over nine scenarios after Holm correction). The prior and continued learning both matter: cold-started bandits, and committing after one probe of each arm, lose more on the mean in every scenario. Supervised selectors fit once on the same data and tuned with the same budget match the mean, and their heavier tail is not significant (7 against 1 contexts above 10%, p=0.06); at seen scales they match or beat online selection. With hyperparameters frozen from H100, UCB1 cuts AUTO's p99 from 60% to 23% and from 56% to 15% on two GB200 fabrics, but its contexts above 10% go from 16 to 8 on one and from 10 to 13 on the other, and scenario-level tests do not separate the two. Live, on two latency-emulated node sets of one H100 system with one run per policy, the learners repair most of a gross AUTO error on one (selection loss 31% to 6-8% for deterministic Greedy and UCB1; average all-reduce bandwidth 1.18-1.24x stock) but average 0.94-0.99x stock on the other, and on both lose to stock below 1 MiB (on the first through plugin overhead); on H100 the cost model is exact there offline, so small sizes can be left to it. We do not show a net end-to-end win.