This research investigates the revitalization of layer dropout (stochastic depth) in Large Language Model (LLM) training, addressing its recent omission from standard pre-training recipes. Through over 2,400 experiments conducted on Cerebras CS-3 systems, the authors analyzed models ranging from 271M to 8.2B parameters across datasets up to 160B tokens. The study establishes that by optimizing layer distribution, time scheduling, and optimizer hyperparameters, layer dropout can significantly enhance efficiency. Specifically, at equivalent training FLOPs, models utilizing layer dropout achieve lower loss compared to those without it. For a fixed number of training steps, the method allows models to maintain or improve validation loss while reducing total training FLOPs by up to 25%. Beyond training, the researchers demonstrate that layer dropout facilitates superior post-training optimizations, including intermediate-layer skipping, early exit strategies, and self-speculative decoding. These techniques can yield up to a 1.5x speedup during inference with negligible impacts on accuracy. Ultimately, the findings suggest that integrating layer dropout into large-scale training regimes provides a robust pathway for improving both computational efficiency and model performance during both the training and deployment phases of LLM development.

Read original