Class Optimizer8bit

bitsandbytes/optim/optimizer.py:117–400

class Optimizer8bit(torch.optim.Optimizer):
    _FSDP_WRAPPED_QUANT_STATE_KEY = "__bnb_optimizer_quant_state__"

    def __init__(self, params, defaults, optim_bits=32, is_paged=False):
        """
        Base 8-bit optimizer class.

        Arguments:
            params (`torch.Tensor`):
                The input parameters to optimize.
            optim_bits (`int`, defaults to 32):
                The number of bits of the optimizer state.
            is_paged (`bool`, defaults to `False`):
                Whether the optimizer is a paged optimizer or not.
        """
        super().__init__(params, defaults)
        self.initialized = False
        self.name2qmap = {}
        self.is_paged = is_paged
        self.page_mng = F.GlobalPageManager.get_instance()

        self.mng = GlobalOptimManager.get_instance()
        self.non_castable_tensor_keys = {
            "qmap1",
            "qmap2",
            "max1",
            "max2",
            "new_max1",
            "new_max2",
            "state1",
            "state2",
            "gnorm_vec",
            "absmax1",
            "absmax2",
            "unorm_vec",
        }

        if optim_bits == 8:
            self.fill_qmap()

    def fill_qmap(self):
        self.name2qmap["dynamic"] = F.create_dynamic_map(signed=True)
        self.name2qmap["udynamic"] = F.create_dynamic_map(signed=False)

    def state_dict(self):
        """Return optimizer state, wrapping quantization tensors for FSDP compatibility.

        FSDP's full_optim_state_dict gathers all tensor states across ranks.
        Quantization states (state1, state2, absmax, etc.) have different shapes
        than model parameters, causing gather operations to fail. By wrapping
        these tensors in a nested dict, FSDP skips them during gathering.
        """
        state_dict = super().state_dict()

        # Deep copy the state to avoid modifying the original optimizer state
        # PyTorch's state_dict() only does a shallow copy
        state_dict["state"] = {
            k: {kk: vv for kk, vv in v.items()} if isinstance(v, dict) else v for k, v in state_dict["state"].items()