Add FSDP integration test for 4-bit state_dict save (#1405)

Add a subprocess-based pytest test that launches a single-GPU FSDP
process via torchrun to verify get_model_state_dict with cpu_offload=True
works for QLoRA-style models with frozen Params4bit base weights.

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

Author: TimDettmers

tests/fsdp_state_dict_save.py

@@ -0,0 +1,80 @@
+"""FSDP state_dict save integration test for 4-bit quantized models (#1405).
+
+This script must be launched via torchrun (not directly):
+    torchrun --nproc_per_node=1 tests/fsdp_state_dict_save.py
+
+It wraps a QLoRA-style model (frozen 4-bit base + trainable adapter) in FSDP
+and calls get_model_state_dict with cpu_offload=True, which exercises the
+_get_fqns() getattr traversal that previously crashed with:
+    AttributeError: 'Params4bit' object has no attribute 'absmax'
+"""
+
+import sys
+
+import torch
+import torch.distributed as dist
+from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+import torch.nn as nn
+
+import bitsandbytes as bnb
+
+
+class SimpleQLoRAModel(nn.Module):
+    """Minimal model with a frozen 4-bit base layer and a trainable adapter."""
+
+    def __init__(self, quant_type="nf4"):
+        super().__init__()
+        self.base = bnb.nn.Linear4bit(64, 64, bias=False, quant_type=quant_type)
+        self.adapter = nn.Linear(64, 64, bias=False)
+
+    def forward(self, x):
+        return self.base(x) + self.adapter(x)
+
+
+def main():
+    dist.init_process_group(backend="nccl")
+    rank = dist.get_rank()
+    torch.cuda.set_device(rank)
+
+    errors = []
+
+    for quant_type in ("nf4", "fp4"):
+        model = SimpleQLoRAModel(quant_type=quant_type)
+        model = model.to("cuda")
+
+        # Freeze quantized base weights (as in real QLoRA)
+        for p in model.base.parameters():
+            p.requires_grad = False
+
+        # Tell FSDP to ignore the frozen quantized params (can't flatten int dtypes)
+        ignored = list(model.base.parameters())
+        fsdp_model = FSDP(model, device_id=rank, ignored_states=ignored, use_orig_params=True)
+
+        options = StateDictOptions(full_state_dict=True, cpu_offload=True)
+        try:
+            state_dict = get_model_state_dict(fsdp_model, options=options)
+
+            # Verify expected keys are present
+            expected_substrings = ["base.weight", "absmax", "quant_map", "adapter.weight"]
+            for substr in expected_substrings:
+                if not any(substr in k for k in state_dict.keys()):
+                    errors.append(f"{quant_type}: missing key containing '{substr}' in {list(state_dict.keys())}")
+
+            print(f"{quant_type}: SUCCESS ({len(state_dict)} keys)", flush=True)
+        except Exception as e:
+            errors.append(f"{quant_type}: {type(e).__name__}: {e}")
+            print(f"{quant_type}: FAILED: {e}", flush=True)
+
+    dist.destroy_process_group()
+
+    if errors:
+        print("\nFAILURES:\n" + "\n".join(errors), file=sys.stderr, flush=True)
+        sys.exit(1)
+    else:
+        print("\nAll FSDP state_dict tests passed.", flush=True)
+        sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()

tests/test_linear4bit.py

@@ -1,7 +1,9 @@
 import copy
 import os
+import pathlib
 import pickle
 import platform
+import subprocess
 import sys
 from tempfile import TemporaryDirectory
 
@@ -497,3 +499,30 @@ def test_params4bit_quant_state_attr_access(device, quant_type, compress_statist
     assert isinstance(w.quant_state, bnb.functional.QuantState)
     assert isinstance(w.bnb_quantized, bool)
     assert w.bnb_quantized is True
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="FSDP requires CUDA")
+@pytest.mark.skipif(
+    not torch.distributed.is_nccl_available(),
+    reason="FSDP test requires NCCL backend",
+)
+def test_fsdp_state_dict_save_4bit():
+    """Integration test: FSDP get_model_state_dict with cpu_offload on a 4-bit model (#1405).
+
+    Launches a single-GPU FSDP process via torchrun to exercise the real
+    _get_fqns() code path that previously crashed with:
+        AttributeError: 'Params4bit' object has no attribute 'absmax'
+    """
+    script = pathlib.Path(__file__).with_name("fsdp_state_dict_save.py")
+    result = subprocess.run(
+        ["torchrun", "--nproc_per_node=1", str(script)],
+        capture_output=True,
+        text=True,
+        timeout=120,
+    )
+    if result.returncode != 0:
+        pytest.fail(
+            f"FSDP state_dict test failed (exit {result.returncode}):\n"
+            f"stdout: {result.stdout}\n"
+            f"stderr: {result.stderr}"
+        )