Method step

bitsandbytes/optim/optimizer.py:306–341

Perform a single optimization step. Arguments: closure (`Callable`, *optional*, defaults to `None`): A closure that reevaluates the model and returns the loss.

(self, closure=None)

    @torch.no_grad()
    def step(self, closure=None):
        """Perform a single optimization step.

        Arguments:
            closure (`Callable`, *optional*, defaults to `None`):
                A closure that reevaluates the model and returns the loss.
        """
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        if not self.initialized:
            self.check_overrides()
            self.to_gpu()  # needed for fairseq pure fp16 training
            self.initialized = True

        # if self.is_paged: self.page_mng.prefetch_all()
        p = None
        for gindex, group in enumerate(self.param_groups):
            for pindex, p in enumerate(group["params"]):
                if p.grad is None:
                    continue
                state = self.state[p]
                if len(state) == 0:
                    self.init_state(group, p, gindex, pindex)

                self.prefetch_state(p)
                self.update_step(group, p, gindex, pindex)
                sync_gpu(p)
        if self.is_paged and p is not None:
            # all paged operations are asynchronous, we need
            # to sync to make sure all tensors are in the right state
            sync_gpu(p)

        return loss

    def get_config(self, gindex, pindex, group):
        config = {}