Precision-Aware Learning Rate Transfer Across Block Count in LLM Post-Training

The post-training of pretrained large language models, including supervised finetuning (SFT) and direct preference optimization (DPO), is very sensitive to the learning rate, and tuning it through repeated runs is costly in both computation and memory. In this work, we propose TipSweep, a method to transfer the learning rate across the number of trainable transformer blocks. To do so efficiently, we search for the best learning rate by training only the last $k_0$ blocks, and show how to scale the selected rate for a subsequent set of $k$ blocks. Across SFT and DPO experiments with models ranging from 1 to 7 billion parameters, we observe that the optimal learning rate follows precision-dependent scaling rules: the optimal rate scales approximately as $k^{-1}$ with float32 (FP32) master weights, but as $k^{-1/2}$ when updates are applied directly to bfloat16 (BF16) weights without a master copy. We explain both rules using an analytical model that accounts for update rounding, under assumptions supported by profiling experiments. Finally, we develop a prototype that demonstrates the feasibility of learning rate search for a 6.7-billion-parameter model on a single consumer GPU with 16 GiB of memory.