Monitoring Parallelism in Large Training Runs

When a large training run is slow, the usual dashboard shows GPU utilization, memory used and NVLink traffic per GPU. None of these tells a researcher which parallelism dimension is costing time.

A GPU waiting on a pipeline bubble, a GPU waiting on a tensor parallel group that crossed a node boundary, and a GPU waiting on a video batch with four times more tokens than its neighbors all look the same.

On NVIDIA B200 this matters even more. NVLink gives each GPU about 900 GB/s in each direction inside a node, while the network gives each GPU about 50 GB/s. This is critical to manage as parallelism designs follow from this 18x drop.

A monitoring framework for large training runs needs to have these 5 essential components.

a) Every sample from DCGM, NCCL and the training loop needs to be tagged with its data parallel, tensor parallel, pipeline, context parallel and expert parallel rank, plus the node, NVLink domain and network it runs on.

b) Every step is split into compute time, exposed communication for each parallelism dimension, pipeline bubble, data loading wait, optimizer time and checkpoint time. Exposed communication is the time a collective (all-gather etc) runs while the GPU idles.

c) An agentic component - planner computes expected memory per pipeline stage and expected bytes per collective from the model config. The dashboard shows expected and measured values side by side, so a large gap points to a bug or a bad placement.

d) A simple ratio of FLOPs per byte moved compared with what the link can feed can give significant insights. For example - on a B200 in BF16 the limit is about 2,500 FLOPs per byte on NVLink and about 45,000 over the network. Tensor parallelism produces bytes roughly 3 times the hidden size divided by the tensor parallel count so I would definitely put it inside the NVLink domain.

e) The data loader monitors attention work per rank. For every step and every rank, it monitors the number of tokens, the longest sequence, the attention kernel time and the peak memory. When step time mostly consists of busiest rank’s attention work, the cause is likely data imbalance.

The output for researchers could be summarized as a cost table which lets researchers see what a design choice costs on the hardware before they scale it.

Discover more from AI & Data Notes for Leaders

Subscribe now to keep reading and get access to the full archive.

Continue reading