1-bit LAMB (Error-Compensated LAMB-Preconditioned Large-Batch Compressed Optimization)

Because LAMB's per-layer variance norms v_t^{(l)} continue decreasing by up to two orders of magnitude throughout large-batch training (Figure 1b), simply freezing both v_{T_w}^{(l)} and LAMB's scaling coefficient c_t^{(l)} ('LAMB + basic 1-bit') degrades BERT validation loss (1.494 vs 1.451) and SQuAD/GLUE accuracy (Tables 2-3); 1-bit LAMB compensates for the frozen denominator sqrt(v_{T_w}^{(l)}) by reconstructing the global layer gradient g_t^{(l)} = (m_t^{(l)} - beta_1 m_{t-1}^{(l)}) / (1 - beta_1) from consecutive compressed momentums, updating a fresh variance v_t^{(l)} = beta_2 v_{t-1}^{(l)} + (1 - beta_2)(g_t^{(l)})^2, and multiplying the warmup moving average c_avg^{(l)} by the soft- and hard-clipped infinity-norm ratio r_t^{(l)} = clip(clip(||v_{T_w}^{(l)} / v_t^{(l)}||_inf, (1 - r_threshold) r_{t-1}^{(l)}, (1 + r_threshold) r_{t-1}^{(l)}), r_min, r_max) (Algorithm 1 Lines 15-21).

To avoid issuing 302 separate per-layer compressed collectives in BERT-Large while preventing layers with small or large momentum magnitudes from suffering severe 1-bit quantization error when fused into a single 1D buffer, 1-bit LAMB computes a fixed per-layer momentum scale coefficient from the average momentum norm at the end of warmup, multiplies each layer's local momentum by its scale coefficient prior to fused 1-bit NCCL Alltoall/AllGather communication, and divides each layer's global compressed momentum by the same fixed coefficient afterward without perturbing error-compensation cancellation.

With time-varying effective layerwise learning rate gamma_t = gamma c_{t-1} in [gamma_min, gamma_max] and step-ratio-scaled error feedback m_{t+1/2}^{(i)} = C_omega[m_t^{(i)} + (c_{t-2} / c_{t-1}) delta_t^{(i)}], the global momentum satisfies m_t = beta_1 m_{t-1} + (1 - beta_1) g_bar_t + (gamma_{t-1} / gamma_t) delta_tilde_{t-1} - delta_tilde_t, and the virtual sequence y_t = x_t - (gamma_t m_t + gamma_{t-1} delta_tilde_t) / (1 - beta_1) satisfies y_{t+1} - y_t = -gamma_t g_bar_{t+1}; under Assumption 1 (L-Lipschitz gradients, variance sigma^2, and bounded compression error eps), Theorem 1 and Corollary 1 establish (V_max / T) sum_{t=0}^{T-1} E||grad f(x_t)||_V^2 = O(sigma / sqrt(n T) + eps^{2/3} / T^{2/3} + 1 / T).

Across BERT-Large pre-training with batch sizes from 8K to 64K on up to 256 V100 GPUs using the 3-phase NCCL Alltoall/AllGather compressed backend, 1-bit LAMB achieves 4.6x and 4.1x end-to-end communication volume reduction for sequence lengths 128 (16.7% warmup) and 512 (19.3% warmup), up to 4.5x compression-stage throughput speedup and 2.8x end-to-end speedup (cutting 256-GPU Ethernet total runtime from 731 minutes to 351 minutes), while matching uncompressed LAMB validation loss (1.443 vs 1.451), SQuAD 1.1 F1 (90.524 avg / 90.788 max), and GLUE average score (83.9 vs 83.9).