LAMB / LARS never update a parameter whose norm is zero (zero-initialised biases, LoRA B matrices, zero-init residual gains) #2116

Activity

1. SalyyS1 commented on Nov 12, 2026

   

   SalyyS1

   on Nov 12, 2026

   More actions

   For LoRA specifically, this bug is even worse than freezing lora_B alone—it causes a permanent coupled deadlock of both lora_A and lora_B:

   1. Under standard LoRA initialization ($A \sim \text{Kaiming}, B = \mathbf{0}$), the gradients at step $t=0$ are:
      $$\nabla_B \mathcal{L} = (\nabla_{\Delta W} \mathcal{L}) A^T \ne \mathbf{0}, \qquad \nabla_A \mathcal{L} = B^T (\nabla_{\Delta W} \mathcal{L}) = \mathbf{0}$$
   2. At step 1, lora_A has zero gradient (so lora_A cannot move until lora_B becomes non-zero), while lora_B has a non-zero gradient but $|B|_2 = 0$.
   3. Because the trust-ratio clip checks if (unorm > max_unorm * pnorm) with pnorm == 0.0f and sets update_scale = max_unorm * pnorm / unorm = 0.0f, lora_B remains identically $\mathbf{0}$ after step 1.
   4. Consequently, $\nabla_A \mathcal{L} = B^T (\nabla_{\Delta W} \mathcal{L})$ remains $\mathbf{0}$ at step 2, 3, ..., $\infty$ as well, so neither lora_A nor lora_B ever moves for the entire training run (and if lora_B.numel() < min_8bit_size = 4096, e.g. d_out = 512, r = 4, even LAMB8bit / LARS8bit hits this 32-bit path).

   In the original LAMB paper (You et al., ICLR 2020) and apex.optimizers.FusedLAMB, the trust ratio is explicitly defined to fall back to 1.0 whenever either $|w_t|_2 == 0$ or $|u_t|_2 == 0$:

   float update_scale = 1.0f;
   if (pnorm > 0.0f && unorm > max_unorm * pnorm) {
       update_scale = (max_unorm * pnorm) / unorm;
   }

   Guarding with pnorm > 0.0f across the CPU, CUDA, and Triton kernels matches FusedLAMB and unblocks step-1 escape from zero initialization.

   (We actually had to add a dedicated DEAD_LAYER check for $|\Delta W|_F = |B A|_F = 0$ in loraspectra loraspectra audit because users training LoRAs with trust-ratio optimizers or rank-expansion schedules kept ending up with .safetensors checkpoints where lora_B was all zeros despite normal-looking training loss logs.)
2. Andreas-Hum commented on Nov 12, 2026

   

   Andreas-Hum

   on Nov 12, 2026

   Author

   More actions

   Thanks, that's right, and it makes the bug worse for LoRA. I checked it on a single LoRA layer (d = 512, PEFT's default init: Kaiming lora_A, zero lora_B), training for 50 steps toward a low-rank target. With bitsandbytes 0.50.2:

   optimizer	r = 4	r = 16
   torch.optim.AdamW (reference)	trains, loss 0.0159 → 0.0090	trains, 0.0159 → 0.0011
   LAMB, LARS, PytorchLARS	neither lora_A nor lora_B changes	neither changes
   LAMB8bit	neither changes	trains

   LAMB8bit is only affected at r = 4, because lora_B then has 2048 < min_8bit_size = 4096 elements and takes the 32-bit path. At r = 16 it takes the 8-bit blockwise path, which doesn't implement the trust ratio at all.

   With #2117 applied, every one of them moves. ‖B‖ after 50 steps at r = 4 (r = 16 is similar, except that LAMB8bit takes the 8-bit path and trains either way):

   optimizer	0.50.2	with #2117
   LAMB, lr 1e-3	0	4.5e-2
   LAMB8bit, lr 1e-3	0	4.5e-2
   LARS / PytorchLARS, lr 1	0	1.7e-3

   After that first step, lora_B still grows slowly. The trust ratio limits each step to about $\text{lr} \cdot \lVert B \rVert$, so a tensor that starts at zero grows by roughly a factor of $(1 + \text{lr})$ per step. LAMB as written in the paper (Algorithm 2 in pure PyTorch, with trust ratio 1 when a norm is 0) does the same: it reaches 4.1e-2, against 4.5e-2 for the fixed bitsandbytes LAMB, so this slow start comes from the method itself, not from bitsandbytes. For LoRA it is probably still worth excluding lora_B from the layer adaptation, as many LAMB setups do for biases.

   Script:

   lora_lamb.py

   """Does LAMB/LARS freeze both LoRA factors when lora_B starts at zero? Pure LoRA layer, CPU."""
   import math
   import bitsandbytes as bnb
   import torch
   
   print("bitsandbytes", bnb.__version__, bnb.__file__)
   torch.manual_seed(0)
   d, n = 512, 1024
   W0 = torch.randn(d, d) / math.sqrt(d)
   dW = torch.randn(d, 8) @ torch.randn(8, d) / d  # low-rank target change
   X = torch.randn(n, d)
   Y = X @ (W0 + dW).T
   
   
   class PaperLAMB(torch.optim.Optimizer):
       """LAMB as in You et al. (2020), Algorithm 2, phi = identity, no weight decay; trust ratio 1 if a norm is 0."""
   
       def __init__(self, params, lr, betas=(0.9, 0.999), eps=1e-6):
           super().__init__(params, dict(lr=lr, betas=betas, eps=eps))
   
       @torch.no_grad()
       def step(self):
           for g in self.param_groups:
               for p in g["params"]:
                   st = self.state[p]
                   if not st:
                       st["t"], st["m"], st["v"] = 0, torch.zeros_like(p), torch.zeros_like(p)
                   st["t"] += 1
                   b1, b2 = g["betas"]
                   st["m"].mul_(b1).add_(p.grad, alpha=1 - b1)
                   st["v"].mul_(b2).addcmul_(p.grad, p.grad, value=1 - b2)
                   u = (st["m"] / (1 - b1 ** st["t"])) / ((st["v"] / (1 - b2 ** st["t"])).sqrt() + g["eps"])
                   wn, un = p.norm(), u.norm()
                   ratio = wn / un if wn > 0 and un > 0 else 1.0
                   p.sub_(g["lr"] * ratio * u)
   
   
   def run(name, make_opt, r, steps=50):
       torch.manual_seed(1)
       A = torch.nn.Parameter(torch.empty(r, d))
       torch.nn.init.kaiming_uniform_(A, a=math.sqrt(5))  # PEFT's default init
       B = torch.nn.Parameter(torch.zeros(d, r))
       A0 = A.detach().clone()
       opt = make_opt([A, B])
       loss0 = None
       for t in range(steps):
           loss = ((X @ W0.T + (X @ A.T) @ B.T - Y) ** 2).mean()
           opt.zero_grad()
           loss.backward()
           if t == 0:
               loss0 = loss.item()
               gA0 = A.grad.norm().item()
           opt.step()
       loss = ((X @ W0.T + (X @ A.T) @ B.T - Y) ** 2).mean().item()
       print(f"{name:34s} r={r:2d} |B|={B.detach().norm():.3e} |A-A0|={(A.detach() - A0).norm():.3e} "
             f"grad_A0={gA0:.0e} loss {loss0:.5f} -> {loss:.5f}")
   
   
   for r in (4, 16):
       run("torch.optim.AdamW (reference)", lambda p: torch.optim.AdamW(p, lr=1e-3, weight_decay=0), r)
       run("bnb LAMB", lambda p: bnb.optim.LAMB(p, lr=1e-3), r)
       run("bnb LAMB8bit", lambda p: bnb.optim.LAMB8bit(p, lr=1e-3), r)
       run("LAMB paper reference (pure torch)", lambda p: PaperLAMB(p, lr=1e-3), r)
       for lr in (1e-3, 1.0):
           run(f"bnb LARS lr={lr:g}", lambda p: bnb.optim.LARS(p, lr=lr, momentum=0.9), r)
           run(f"bnb PytorchLARS lr={lr:g}", lambda p: bnb.optim.PytorchLARS(p, lr=lr, momentum=0.9), r)
           run(f"torch.optim.SGD lr={lr:g} (reference)", lambda p: torch.optim.SGD(p, lr=lr, momentum=0.9), r)