BAdam (Block Coordinate Descent with Adam's Update Rule)

For an M-billion-parameter model partitioned into D equal-sized blocks, storing up-to-date parameters in FP16 (2M GB) and allocating FP32 master parameters, FP32 gradients, FP32 momentum, and FP32 second moments only for the single active block theta_{pi_i} reduces parameter/gradient/optimizer memory from 18M GB in mixed-precision Adam to 2M + 16M/D GB (Eq. 5, Table 1); furthermore, when blocks correspond to D consecutive transformer modules, updating only one active block at depth i requires only i unit-backward-passes per batch and avoids computing parameter gradients d z_l / d theta_l for intermediate modules, reducing total unit-backward-passes over KD batches from K D^2 in Adam/LoRA to K D(D + 1) / 2 (Eq. 6).

Under global L-smoothness, block L_i-smoothness (L_bar = max_i L_i), and bounded block partial derivatives ||g_i^{t,k}|| <= G (Assumptions D.1-D.2), the coordinate adaptive step matrix H_i^{t,k} = diag(1 / (sqrt(v_hat_i^{t,k}) + lambda)) satisfies (1 / (2G)) I <= H_i^{t,k} <= (1 / lambda) I (Corollary D.3), the inner Adam-vs-GD deviation is bounded by ||e_i^t|| <= (2 L_i alpha K / lambda^2) ||g_i^{t,1}|| (Lemma D.7), each block update satisfies L(theta_i^t) - L(theta_{i-1}^t) <= -(alpha K / (8G)) ||grad_i L(theta_i^t)||^2 (Corollary D.8), and with step size alpha = (lambda / (4 L_bar K D^{1/4})) min{1 / (K D^{1/4}), lambda / (6G)} one block-epoch achieves L(theta^{t+1}) - L(theta^t) <= -(alpha K / (32G)) ||grad L(theta^t)||^2, reaching a delta-stationary point in at most O(delta^{-2}) gradient evaluations (Theorem D.4, Corollary D.5).

Finetuning Llama 3-8B (D = 32) on a single RTX3090-24GB GPU (consuming 23.5 GB total memory vs 144.8 GB+ for Adam and cutting per-epoch backward time from 3.20 hours in LoRA and 3.70 hours in LOMO to 1.74 hours, Tables 2-3) and Llama 3-70B (D = 80) on 4xA100-80GB GPUs, BAdam learns full-parameter weight perturbations Delta W = W_K - W_0 whose effective rank min{r : cvar(r) >= 0.9} matches full Adam (up to ~2000-2500 across layers, Figure 2), outperforming LoRA, LOMO, and GaLore on Alpaca-GPT4 MT-bench (6.65 vs 6.27 LoRA and 6.23 Adam on Llama 3-8B, Table 5) and MathInstruct zero-shot math benchmarks (44.4 vs 43.3 LoRA and 44.1 Adam on Llama 3-8B; 62.7 vs 59.0 LoRA on Llama 3-70B, Table 6).

BAdam requires selecting the inner Adam step count K per block (guided by K = min{max{n / (B D), 50}, 100} in Eq. 7, since too small K yields insufficient Adam adaptation after zeroing states and too large K over-optimizes a single block before switching), and the formal O(delta^{-2}) convergence proof in Theorem D.4 is established for deterministic block gradients rather than stochastic mini-batch gradients.