Fix AdEMAMix scheduler guard and add state_dict round-trip test (#1861)

* Fix AdEMAMix scheduler guard and add state_dict round-trip test (#1382)

Fix potential division-by-zero in AdEMAMix update_step when t_alpha or
t_beta3 is 0 (e.g. from get_config defaults). Change scheduler guards
from `is None` to falsy checks so that 0, None, and False all correctly
skip the scheduler path. Also change get_config defaults for t_alpha and
t_beta3 from 0 to None to match the intended semantics.

Add test_ademamix_state_dict_no_nan which saves and loads AdEMAMix
state_dict (8-bit, 32-bit, with and without schedulers) and verifies:
- loaded state matches original byte-for-byte
- training resumes without NaN or Inf
- two optimizers loaded from the same checkpoint produce identical updates

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* style: Fix ruff format violation in test_linear4bit.py

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

bitsandbytes/optim/ademamix.py

@@ -180,7 +180,7 @@ def init_state(self, group, p, gindex, pindex):
     def update_step(self, group, p, gindex, pindex):
         config = self.get_config(gindex, pindex, group)
 
-        if config["t_alpha"] is None and config["t_beta3"] is None:
+        if not config["t_alpha"] and not config["t_beta3"]:
             # Not using alpha/beta3 scheduler; we can fall through.
             super().update_step(group, p, gindex, pindex)
             return
@@ -201,13 +201,13 @@ def update_step(self, group, p, gindex, pindex):
         t_beta3 = config["t_beta3"]
 
         # Apply scheduler for alpha
-        if t_alpha is not None:
+        if t_alpha:
             alpha_t = min(step * alpha / t_alpha, alpha)
         else:
             alpha_t = alpha
 
         # Apply scheduler for beta3
-        if t_beta3 is not None:
+        if t_beta3:
             ln_beta1 = math.log(beta1)
             ln_beta3 = math.log(beta3)
             step_scale = step / t_beta3

bitsandbytes/optim/optimizer.py

@@ -341,8 +341,8 @@ def get_config(self, gindex, pindex, group):
         config["weight_decay"] = group["weight_decay"]
         config["lr"] = group["lr"]
         config["alpha"] = group.get("alpha", 0.0)
-        config["t_alpha"] = group.get("t_alpha", 0)
-        config["t_beta3"] = group.get("t_beta3", 0)
+        config["t_alpha"] = group.get("t_alpha", None)
+        config["t_beta3"] = group.get("t_beta3", None)
         config["optim_bits"] = self.args.optim_bits
         config["min_8bit_size"] = self.args.min_8bit_size
         config["percentile_clipping"] = self.args.percentile_clipping

tests/test_optim.py

@@ -592,3 +592,111 @@ def test_benchmark_blockwise(dim1, dim2, gtype, optim_name, device):
     params = (total_steps - total_steps // 5) * dim1 * dim2
     print(optim_name, gtype, s, params, s / params)
     # assert s < 3.9
+
+
+ademamix_state_dict_opts = [
+    ("AdEMAMix8bit", lambda p: bnb.optim.AdEMAMix8bit(p, lr=1e-3)),
+    ("AdEMAMix32bit", lambda p: bnb.optim.AdEMAMix(p, lr=1e-3)),
+    ("AdEMAMix8bit_scheduled", lambda p: bnb.optim.AdEMAMix8bit(p, lr=1e-3, t_alpha=100, t_beta3=100)),
+    ("AdEMAMix32bit_scheduled", lambda p: bnb.optim.AdEMAMix(p, lr=1e-3, t_alpha=100, t_beta3=100)),
+]
+
+
+@pytest.mark.parametrize(
+    "optim_name,optim_factory",
+    ademamix_state_dict_opts,
+    ids=[x[0] for x in ademamix_state_dict_opts],
+)
+@pytest.mark.parametrize("device", get_available_devices(no_cpu=True))
+@pytest.mark.skipif(not get_available_devices(no_cpu=True), reason="No device")
+def test_ademamix_state_dict_no_nan(optim_name, optim_factory, device):
+    """Test that AdEMAMix can save/load state_dict and continue training without NaN.
+
+    Regression test for https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1382
+    """
+    if device not in ["cuda", "xpu"]:
+        pytest.skip("Optimizers are only supported on CUDA and XPU")
+
+    import torch.nn as nn
+
+    torch.manual_seed(42)
+    model = nn.Linear(256, 64).to(device)
+    opt = optim_factory(model.parameters())
+
+    # Train a few steps to populate optimizer state
+    for _ in range(10):
+        x = torch.randn(8, 256, device=device)
+        loss = model(x).sum()
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+
+    # Save state
+    model_sd = {k: v.clone() for k, v in model.state_dict().items()}
+    opt_sd = opt.state_dict()
+    path = get_temp_dir()
+    torch.save(opt_sd, join(path, "opt.pt"))
+    torch.save(model_sd, join(path, "model.pt"))
+
+    # Create fresh model and optimizer, load state
+    model2 = nn.Linear(256, 64).to(device)
+    model2.load_state_dict(torch.load(join(path, "model.pt")))
+    opt2 = optim_factory(model2.parameters())
+    opt2.load_state_dict(torch.load(join(path, "opt.pt")))
+    rm_path(path)
+
+    # Verify loaded state matches original byte-for-byte
+    orig_params = list(model.parameters())
+    loaded_params = list(model2.parameters())
+    for p_idx in range(len(orig_params)):
+        s1 = opt.state[orig_params[p_idx]]
+        s2 = opt2.state[loaded_params[p_idx]]
+        for k in s1:
+            if isinstance(s1[k], torch.Tensor):
+                assert s1[k].shape == s2[k].shape, f"Shape mismatch for param {p_idx} {k}"
+                assert s1[k].dtype == s2[k].dtype, f"Dtype mismatch for param {p_idx} {k}"
+                torch.testing.assert_close(s1[k], s2[k])
+
+    # Resume training and verify no NaN
+    for i in range(10):
+        x = torch.randn(8, 256, device=device)
+        loss = model2(x).sum()
+        assert not torch.isnan(loss), f"NaN loss at step {i} after loading state_dict"
+        assert not torch.isinf(loss), f"Inf loss at step {i} after loading state_dict"
+        loss.backward()
+        opt2.step()
+        opt2.zero_grad()
+
+        # Check parameters for NaN/Inf after each step
+        for p in model2.parameters():
+            assert not p.isnan().any(), f"NaN in parameters at step {i} after loading state_dict"
+            assert not p.isinf().any(), f"Inf in parameters at step {i} after loading state_dict"
+
+    # Verify the original and loaded optimizers produce identical updates
+    # from the same starting point (immediately after loading, before any divergence)
+    torch.manual_seed(999)
+    x_orig = torch.randn(8, 256, device=device)
+    x_loaded = x_orig.clone()
+
+    # Reset models to the saved checkpoint weights
+    model.load_state_dict(model_sd)
+    model2.load_state_dict(model_sd)
+
+    # Reload optimizer states from the same checkpoint into two fresh optimizers
+    opt_fresh = optim_factory(model.parameters())
+    opt_fresh.load_state_dict(opt_sd)
+    opt_fresh2 = optim_factory(model2.parameters())
+    opt_fresh2.load_state_dict(opt_sd)
+
+    loss_a = model(x_orig).sum()
+    loss_a.backward()
+    opt_fresh.step()
+    opt_fresh.zero_grad()
+
+    loss_b = model2(x_loaded).sum()
+    loss_b.backward()
+    opt_fresh2.step()
+    opt_fresh2.zero_grad()
+
+    for p_a, p_b in zip(model.parameters(), model2.parameters()):
+        torch.testing.assert_close(p_a, p_b)