Notes on Pretraining Parallelization and Failed Training Runs

Pretraining easily collapses the moment it touches causality or accumulates bias.

• token routing sends each token to its top-k experts, but the load per expert can become severely imbalanced.
• Using expert choice only during training to correct for this makes token placement affected by future information, breaking the train-deploy causality.
• token dropping belongs to the same category of problem.
• Rumor has it that this approach is linked to Llama 4's underwhelming performance and some Gemini 2 Pro issues.

Numerical error is more subtle. Early in GPT-4 training, FP16 all-reduce amplified rounding errors in cumulative summation, and this was far more dangerous because while variance can average out, bias accumulates.

Horace He's lecture explains pretraining parallelization as a chain of problems and solutions. Treating pretraining FLOPs as 6ND, you start with data parallel, but hit limits like B300's 288GB HBM. So each GPU keeping only 1/N of the parameters and gathering them per layer via all-gather — FSDP — becomes the default. This approach makes it easy to overlap communication and computation, and adding reduce-scatter can lower communication volume to about 3x relative to parameters. In multi-domain settings, a hierarchical collective chain of intra-domain reduce-scatter, inter-domain all-reduce, and then all-gather again further reduces the bottleneck. For the same reason, TPUs with more accelerators per domain are advantageous for FSDP. As batch size grows, the crossover shifts to the right, and the more sparse the model, the more it pulls to the left.

However, as GPU count increases, computation time shrinks but communication time does not, leading to a comms crossover. At this point pipeline parallelism needs to be mixed in, but this creates bubbles: GPUs handling later layers sit idle early in the batch, and GPUs handling earlier layers sit idle later. Since gradients must be aggregated before updating, overlapping with the next batch isn't easy, and long sequences plus large critical batch sizes also create a batch floor for FSDP. Kimi-style attention-to-residuals or mixed attention structures further tighten stage constraints, slowing down research iteration speed.

Ultimately, the problem isn't a single bottleneck but a systemic issue where numerics, routing, parallelism, and architecture all shake together. In inference for RL, numerical drift relative to the training engine can create subtle off-policy bias, and even Blackwell-level kernel optimization took Nvidia itself a long time. So automating kernel writing is hard to view as something that will be solved anytime soon, and since large-scale training keeps producing new types of failures, the process of combining compute multipliers must be managed even more rigorously.