ista-daslab-optimizers 0.0.1__cp39-cp39-manylinux_2_34_x86_64.whl

ista_daslab_optimizers/__init__.py

@@ -0,0 +1,4 @@
+from .acdc import *
+from .micro_adam import *
+from .sparse_mfac import *
+from .dense_mfac import *

ista_daslab_optimizers/acdc/__init__.py

@@ -0,0 +1,5 @@
+from .acdc import ACDC
+
+__all__ = [
+    'ACDC'
+]

ista_daslab_optimizers/acdc/acdc.py

@@ -0,0 +1,387 @@
+import math
+from enum import Enum
+import torch
+import wandb
+from .wd_scheduler import WeightDecayScheduler
+
+class ACDC_Action(Enum):
+    FILL_MASK_WITH_ONES = 1
+    SPARSIFY_MODEL_AND_UPDATE_MASK = 2
+    KEEP_MASK_FIXED = 3
+
+class ACDC_Phase(Enum):
+    DENSE = 1
+    SPARSE = 2
+
+class ACDC_Scheduler:
+    """
+        This class will hold a list of epochs where the sparse training is performed
+    """
+    def __init__(self, warmup_epochs, epochs, phase_length_epochs, finetuning_epochs, zero_based=False):
+        """
+        Builds an AC/DC Scheduler
+        :param warmup_epochs: the warm-up length (dense training)
+        :param epochs: total number of epochs for training
+        :param phase_length_epochs: the length of dense and sparse phases (both are equal)
+        :param finetuning_epochs: the epoch when the finetuning_epochs starts and is considered sparse training
+        :param zero_based: True if epochs start from zero, False if epochs start from one
+        """
+        # print(f'AC/DC Scheduler: {warmup_epochs=}, {phase_length_epochs=}, {finetuning_epochs=}, {epochs=}')
+        self.warmup_epochs = warmup_epochs
+        self.epochs = epochs
+        self.phase_length_epochs = phase_length_epochs
+        self.finetuning_epochs = finetuning_epochs
+        self.zero_based = zero_based
+
+        self.sparse_epochs = []
+        self._build_sparse_epochs()
+
+    def _build_sparse_epochs(self):
+        is_sparse = True
+        for i, e in enumerate(range(self.warmup_epochs, self.epochs)):
+            if is_sparse or e >= self.finetuning_epochs:
+                self.sparse_epochs.append(e)
+
+            if (e-self.warmup_epochs) % self.phase_length_epochs == self.phase_length_epochs - 1:
+            # if e % self.phase_length_epochs == self.phase_length_epochs - 1:
+                is_sparse = not is_sparse
+
+        if not self.zero_based:
+            for i in range(len(self.sparse_epochs)):
+                self.sparse_epochs[i] += 1
+
+    def is_sparse_epoch(self, epoch):
+        """
+            Returns True if sparse training should performed, otherwise returns False
+            :param epoch: a zero-based epoch number
+        """
+        return epoch in self.sparse_epochs
+
+    def is_finetuning_epoch(self, epoch):
+        return epoch >= self.finetuning_epochs
+
+    def get_action(self, epoch):
+        is_crt_epoch_sparse = self.is_sparse_epoch(epoch)
+        is_prev_epoch_sparse = self.is_sparse_epoch(epoch-1)
+        if is_crt_epoch_sparse:
+            if is_prev_epoch_sparse:
+                return ACDC_Action.KEEP_MASK_FIXED # mask was updated with top-k at the first dense epoch and now do nothing
+            return ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK # first dense epoch, update the mask using topk
+        else: # dense epoch
+            if epoch == int(not self.zero_based) or is_prev_epoch_sparse:
+                return ACDC_Action.FILL_MASK_WITH_ONES  # fill mask with ones
+            else:
+                return ACDC_Action.KEEP_MASK_FIXED # do not change mask
+
+    def get_phase(self, epoch):
+        if self.is_sparse_epoch(epoch):
+            return ACDC_Phase.SPARSE
+        return ACDC_Phase.DENSE
+
+class ACDC(torch.optim.Optimizer):
+    """
+        This class implements the Default AC/DC schedule from the original paper https://arxiv.org/abs/2106.12379.pdf:
+        We use the model only to add names for the parameters in wandb logs and print the sparsity, norm and weight decay
+        For example, fc.bias (1D, that requires weight decay) and ln/bn.weight/bias (1D, that do not require weight decay).
+            * LN = Layer Normalization
+            * BN = Batch Normalization
+        Example for 100 epochs (from original AC/DC paper https://arxiv.org/pdf/2106.12379.pdf)
+            - first 10 epochs warmup (dense training)
+            - alternate sparse/dense training phases once at 5 epochs
+            - last 10 epochs finetuning (sparse training)
+            !!!!! SEE THE HUGE COMMENT AT THE END OF THIS FILE FOR A MORE DETAILED EXAMPLE !!!!!
+
+        To use this class, make sure you call method `update_acdc_state` at the beginning of each epoch.
+
+        The following information is required:
+            - params
+            - model
+            - momentum
+            - weight_decay
+            - wd_type
+            - k
+            - total_epochs
+            - warmup_epochs
+            - phase_length_epochs
+            - finetuning_epochs
+    """
+    def __init__(self,
+                 params, model, # optimization set/model
+                 lr, momentum, weight_decay, wd_type, k, # hyper-parameters
+                 total_epochs, warmup_epochs, phase_length_epochs, finetuning_epochs): # acdc schedulers
+        super(ACDC, self).__init__(params, defaults=dict(lr=lr, weight_decay=weight_decay, momentum=momentum, k=k))
+
+        self.model = model
+        self.lr = lr
+        self.momentum = momentum
+        self.weight_decay = weight_decay
+        self.wd_type = wd_type
+        self.k = k
+
+        self.acdc_scheduler = ACDC_Scheduler(
+            warmup_epochs=warmup_epochs,
+            epochs=total_epochs,
+            phase_length_epochs=phase_length_epochs,
+            finetuning_epochs=finetuning_epochs,
+            zero_based=False)
+
+        self.phase = None
+        self.is_finetuning_epoch = False
+        self.update_mask_flag = None
+
+        self.current_epoch = 0
+        self.steps = 0
+        self.log_interval = 250
+
+        self._initialize_param_states()
+
+    def _initialize_param_states(self):
+        for group in self.param_groups:
+            for p in group['params']:
+
+                # this method is called before the first forward pass and all gradients will be None
+                # if p.grad is None:
+                #     continue
+
+                state = self.state[p]
+
+                # initialize the state for each individual parameter p
+                if len(state) == 0:
+                    # v is the momentum buffer
+                    state['v'] = torch.zeros_like(p)
+
+                    # set density to be used in top-k call (only for multi-dim tensors)
+                    state['density'] = int(self.k * p.numel())
+
+                    # 1D tensors, like:
+                    # - batch/layer normalization layers
+                    # - biases for other layers
+                    # will always have mask=1 because they will never be pruned
+                    state['mask'] = torch.ones_like(p)
+
+                    # set the weight decay scheduler for each parameter individually
+                    # all biases and batch/layer norm layers are not decayed
+                    if len(p.shape) == 1:
+                        state['wd_scheduler'] = WeightDecayScheduler(weight_decay=0, wd_type='const')
+                    else:
+                        state['wd_scheduler'] = WeightDecayScheduler(weight_decay=self.weight_decay, wd_type=self.wd_type)
+
+    @torch.no_grad()
+    def update_acdc_state(self, epoch):
+        self.current_epoch = epoch
+        phase = self.acdc_scheduler.get_phase(self.current_epoch)
+        action = self.acdc_scheduler.get_action(self.current_epoch)
+
+        print(f'{epoch=}, {phase=}')
+
+        self._set_phase(phase)
+
+        if action == ACDC_Action.FILL_MASK_WITH_ONES:
+            for group in self.param_groups:
+                for p in group['params']:
+                    self.state[p]['mask'].fill_(1)
+        elif action == ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK:
+            # update mask and sparsify model
+            for group in self.param_groups:
+                for p in group['params']:
+                    is_multi_dim = (len(p.shape) > 1)
+                    if is_multi_dim:
+                        state = self.state[p]
+                        # original_shape = p.shape.clone()
+                        indices = torch.topk(p.reshape(-1).abs(), k=state['density']).indices
+
+                        # zerorize, view mask as 1D and then set 1 to specific indices. The result will have p.shape
+                        state['mask'].zero_().reshape(-1)[indices] = 1.
+
+                        # apply the mask to the parameters
+                        p.mul_(state['mask'])
+        elif action == ACDC_Action.KEEP_MASK_FIXED:
+            pass  # do nothing
+
+    def _zerorize_momentum_buffer(self):
+        # zerorize momentum buffer only for the multi-dimensional parameters
+        for group in self.param_groups:
+            for p in group['params']:
+                self.state[p]['v'].zero_()
+
+    def _set_phase(self, phase):
+        if phase != self.phase:
+            self.phase = phase
+            if self.phase == ACDC_Phase.DENSE:
+                # AC/DC: zerorize momentum buffer at the transition SPARSE => DENSE
+                # The following quote is copy-pasted from original ACDC paper, page 7:
+                # "We reset SGD momentum at the beginning of every decompression phase."
+                self._zerorize_momentum_buffer()
+
+    @torch.no_grad()
+    def _wandb_log(self):
+        if self.steps % self.log_interval == 0:
+            wandb_dict = dict()
+
+            total_params = 0
+            global_sparsity = 0
+            global_params_norm = 0
+            global_grad_norm = 0
+            for name, p in self.model.named_parameters():
+                total_params += p.numel()
+                crt_sparsity = (p == 0).sum().item()
+                norm_param = p.norm(p=2)
+                norm_grad = p.grad.norm(p=2)
+
+                wandb_dict[f'weight_sparsity_{name}'] = crt_sparsity / p.numel() * 100.
+                wandb_dict[f'mask_sparsity_{name}'] = (self.state[p]['mask'] == 0).sum().item() / p.numel() * 100.
+                wandb_dict[f'norm_param_{name}'] = norm_param
+                wandb_dict[f'norm_grad_{name}'] = norm_grad
+
+                if self.wd_type == 'awd':
+                    wandb_dict[f'awd_{name}'] = self.state[p]['wd_scheduler'].get_wd()
+
+                global_params_norm += norm_param ** 2
+                global_grad_norm += norm_grad ** 2
+                global_sparsity += crt_sparsity
+
+            wandb_dict[f'global_params_norm'] = math.sqrt(global_params_norm)
+            wandb_dict[f'global_grad_norm'] = math.sqrt(global_grad_norm)
+            wandb_dict[f'global_sparsity'] = global_sparsity / total_params * 100.
+
+            wandb_dict['optimizer_epoch'] = self.current_epoch
+            wandb_dict['optimizer_step'] = self.steps
+
+            wandb_dict['is_dense_phase'] = int(self.phase == ACDC_Phase.DENSE)
+            wandb_dict['is_sparse_phase'] = int(self.phase == ACDC_Phase.SPARSE)
+            wandb.log(wandb_dict)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        self.steps += 1
+        for group in self.param_groups:
+            lr = group['lr']
+            momentum = group['momentum']
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                # holds all buffers for the current parameter
+                state = self.state[p]
+
+                ### apply mask to the gradient on sparse phase
+                ### do not modify gradient in place via p.grad.mul_(state['mask'])
+                ### because this will affect the norm statistics in self._wandb_log
+                ### this will create intermediary tensors
+                # grad = p.grad
+                if self.phase == ACDC_Phase.SPARSE:
+                    # grad = grad * state['mask']
+                    p.grad.mul_(state['mask']) # sparsify gradient
+
+                state['v'].mul_(momentum).add_(p.grad)
+
+                wd = state['wd_scheduler'](w=p, g=p.grad) # use sparsified gradient
+                p.mul_(1 - lr * wd).sub_(other=state['v'], alpha=lr).mul_(state['mask'])
+
+        self._wandb_log()
+
+"""
+This is what ACDC_Scheduler outputs for the default ACDC schedule, presented at page 5, Figure 1 in the paper https://arxiv.org/pdf/2106.12379.pdf
+
+epoch	is_sparse_epoch	phase	            action
+1		False			ACDC_Phase.DENSE	ACDC_Action.FILL_MASK_WITH_ONES
+2		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+3		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+4		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+5		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+6		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+7		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+8		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+9		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+10		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+11		True			ACDC_Phase.SPARSE	ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK
+12		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+13		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+14		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+15		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+16		False			ACDC_Phase.DENSE	ACDC_Action.FILL_MASK_WITH_ONES
+17		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+18		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+19		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+20		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+21		True			ACDC_Phase.SPARSE	ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK
+22		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+23		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+24		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+25		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+26		False			ACDC_Phase.DENSE	ACDC_Action.FILL_MASK_WITH_ONES
+27		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+28		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+29		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+30		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+31		True			ACDC_Phase.SPARSE	ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK
+32		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+33		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+34		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+35		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+36		False			ACDC_Phase.DENSE	ACDC_Action.FILL_MASK_WITH_ONES
+37		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+38		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+39		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+40		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+41		True			ACDC_Phase.SPARSE	ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK
+42		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+43		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+44		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+45		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+46		False			ACDC_Phase.DENSE	ACDC_Action.FILL_MASK_WITH_ONES
+47		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+48		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+49		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+50		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+51		True			ACDC_Phase.SPARSE	ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK
+52		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+53		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+54		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+55		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+56		False			ACDC_Phase.DENSE	ACDC_Action.FILL_MASK_WITH_ONES
+57		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+58		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+59		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+60		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+61		True			ACDC_Phase.SPARSE	ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK
+62		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+63		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+64		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+65		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+66		False			ACDC_Phase.DENSE	ACDC_Action.FILL_MASK_WITH_ONES
+67		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+68		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+69		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+70		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+71		True			ACDC_Phase.SPARSE	ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK
+72		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+73		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+74		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+75		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+76		False			ACDC_Phase.DENSE	ACDC_Action.FILL_MASK_WITH_ONES
+77		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+78		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+79		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+80		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+81		True			ACDC_Phase.SPARSE	ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK
+82		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+83		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+84		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+85		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+86		False			ACDC_Phase.DENSE	ACDC_Action.FILL_MASK_WITH_ONES
+87		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+88		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+89		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+90		False			ACDC_Phase.DENSE	ACDC_Action.KEEP_MASK_FIXED
+91		True			ACDC_Phase.SPARSE	ACDC_Action.SPARSIFY_MODEL_AND_UPDATE_MASK
+92		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+93		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+94		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+95		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+96		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+97		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+98		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+99		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+100		True			ACDC_Phase.SPARSE	ACDC_Action.KEEP_MASK_FIXED
+"""

ista_daslab_optimizers/acdc/wd_scheduler.py

@@ -0,0 +1,31 @@
+class WeightDecayScheduler:
+    def __init__(self, weight_decay: float, wd_type: str):
+        # assert wd_type in ['const', 'balanced', 'opt-a', 'opt-b']
+        self.weight_decay = weight_decay
+        self.wd_type = wd_type
+        self.awd = 0 # for AWD from apple
+
+    def get_wd(self):
+        if self.wd_type == 'const':
+            return self.weight_decay
+        if self.wd_type == 'awd':
+            return self.awd
+
+    def __call__(self, w=None, g=None):
+        """
+        :param w: tensor that contains weights
+        :param g: tensor that contains gradients
+        :return: the value for the weight decay
+        """
+        if self.wd_type == 'const':
+            return self.weight_decay
+
+        if self.wd_type == 'awd': # AWD from the Apple paper: https://openreview.net/pdf?id=ajnThDhuq6
+            assert (w is not None) and (g is not None),\
+                'The balanced weight decay scheduler requires valid for w and g, but at least one is None!'
+
+            # in the paper, lambda_awd is set by the user and they return a lambda_bar
+            # here, self.awd will be the moving average from the line 8 in their algorithm and
+            # the input to our algorithm is self.weight_decay for all wd_types!
+            self.awd = 0.1 * self.awd + 0.9 * self.weight_decay * g.norm(p=2) / w.norm(p=2)
+            return self.awd

ista_daslab_optimizers/dense_mfac/__init__.py

@@ -0,0 +1,5 @@
+from .dense_mfac import DenseMFAC
+
+__all__ = [
+    'DenseMFAC'
+]

ista_daslab_optimizers/dense_mfac/dense_core_mfac.py

@@ -0,0 +1,164 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+USE_CUDA = True
+try:
+    import ista_daslab_dense_mfac
+except Exception as e:
+    USE_CUDA = False
+    print('\n\t[WARNING] The module "ista_daslab_dense_mfac" is not installed, using slower PyTorch implementation!\n')
+
+class DenseCoreMFAC:
+    def __init__(self, grads, dev, gpus, damp=1e-5, create_G=False):
+        self.m, self.d = grads.shape
+        self.dev = dev
+        self.gpus = gpus
+        self.dtype = grads.dtype
+        self.gpus = gpus
+        self.grads_count = 0
+        self.wandb_data = dict()
+        self.damp = None
+        self.lambd = None
+        self.set_damp(damp)
+        self.create_G = create_G
+        if self.create_G:
+            self.G = grads
+
+        if USE_CUDA and self.m % 32 != 0 or self.m > 1024:
+            raise ValueError('CUDA implementation currently on supports $m$ < 1024 and divisible by 32.')
+
+        self.dper = self.d // len(gpus) + 1
+        self.grads = []  # matrix $G$ in the paper
+        for idx in range(len(gpus)):
+            start, end = idx * self.dper, (idx + 1) * self.dper
+            self.grads.append(grads[:, start:end].to(gpus[idx]))
+        self.dots = torch.zeros((self.m, self.m), device=self.dev, dtype=self.dtype)  # matrix $GG^T$
+        for idx in range(len(gpus)):
+            self.dots += self.grads[idx].matmul(self.grads[idx].t()).to(self.dev)
+
+        self.last = 0  # ringbuffer index
+        self.giHig = self.lambd * self.dots  # matrix $D$
+        self.denom = torch.zeros(self.m, device=self.dev, dtype=self.dtype)  # $D_ii + m$
+        self.coef = self.lambd * torch.eye(self.m, device=self.dev, dtype=self.dtype)  # matrix $B$
+
+        self.setup()
+
+    def empty_buffer(self):
+        for g in self.grads:
+            g.zero_()
+
+    def set_damp(self, new_damp):
+        self.damp = new_damp
+        self.lambd = 1. / new_damp
+
+    def reset_optimizer(self):
+        self.grads_count = 0
+        for idx in range(len(self.gpus)):
+            self.grads[idx].zero_()
+        self.dots.zero_()
+        for idx in range(len(self.gpus)):
+            self.dots += self.grads[idx].matmul(self.grads[idx].t()).to(self.dev)
+        self.last = 0
+        self.giHig = self.lambd * self.dots  # matrix $D$
+        self.denom = torch.zeros(self.m, device=self.dev, dtype=self.dtype)  # $D_ii + m$
+        self.coef = self.lambd * torch.eye(self.m, device=self.dev, dtype=self.dtype)  # matrix $B$
+        self.setup()
+
+    # Calculate $D$ / `giHig` and $B$ / `coef`
+    def setup(self):
+        self.giHig = self.lambd * self.dots
+        diag = torch.diag(torch.full(size=[self.m], fill_value=self.m, device=self.dev, dtype=self.dtype))
+        self.giHig = torch.lu(self.giHig + diag, pivot=False)[0]
+        self.giHig = torch.triu(self.giHig - diag)
+        self.denom = self.m + torch.diagonal(self.giHig) # here we should use min(grads_count, m)
+        tmp = -self.giHig.t().contiguous() / self.denom.reshape((1, -1))
+
+        if USE_CUDA:
+            diag = torch.diag(torch.full(size=[self.m], fill_value=self.lambd, device=self.dev, dtype=self.dtype))
+            self.coef = ista_daslab_dense_mfac.hinv_setup(tmp, diag)
+        else:
+            for i in range(max(self.last, 1), self.m):
+                self.coef[i, :i] = tmp[i, :i].matmul(self.coef[:i, :i])
+
+    # Replace oldest gradient with `g` and then calculate the IHVP with `g`
+    def integrate_gradient_and_precondition(self, g, x):
+        tmp = self.integrate_gradient(g)
+        p = self.precondition(x, tmp)
+        return p
+
+    # Replace oldest gradient with `g`
+    def integrate_gradient(self, g):
+        self.set_grad(self.last, g)
+        tmp = self.compute_scalar_products(g)
+        self.dots[self.last, :] = tmp
+        self.dots[:, self.last] = tmp
+        self.setup()
+        self.last = (self.last + 1) % self.m
+        return tmp
+
+    # Distributed `grads[j, :] = g`
+    def set_grad(self, j, g):
+        # for the eigenvalue experiment:
+        if self.create_G:
+            self.G[j, :] = g
+
+        self.grads_count += 1
+        def f(i):
+            start, end = i * self.dper, (i + 1) * self.dper
+            self.grads[i][j, :] = g[start:end]
+
+        nn.parallel.parallel_apply(
+            [f] * len(self.grads), list(range(len(self.gpus)))
+        )
+
+    # Distributed `grads.matmul(x)`
+    def compute_scalar_products(self, x):
+        def f(i):
+            start, end = i * self.dper, (i + 1) * self.dper
+            G = self.grads[i]
+            return G.matmul(x[start:end].to(G.device)).to(self.dev)
+
+        outputs = nn.parallel.parallel_apply(
+            [f] * len(self.gpus), list(range(len(self.gpus)))
+        )
+        return sum(outputs)
+
+    # Product with inverse of dampened empirical Fisher
+    def precondition(self, x, dots=None):
+        if dots is None:
+            dots = self.compute_scalar_products(x)
+        giHix = self.lambd * dots
+        if USE_CUDA:
+            giHix = ista_daslab_dense_mfac.hinv_mul(self.m, self.giHig, giHix)
+        else:
+            for i in range(1, self.m):
+                giHix[i:].sub_(self.giHig[i - 1, i:], alpha=giHix[i - 1] / self.denom[i - 1])
+        """
+            giHix size: 1024
+            denom size: 1024
+            coef size: 1024x1024
+            M size: 1024
+            x size: d
+        """
+        M = (giHix / self.denom).matmul(self.coef)
+        partA = self.lambd * x
+        partB = self.compute_linear_combination(M)
+        self.wandb_data.update({f'norm_partA': partA.norm(p=2), f'norm_partB': partB.norm(p=2)})
+        return partA.to(self.dev) - partB.to(self.dev)
+
+    # Distributed `x.matmul(grads)`
+    def compute_linear_combination(self, x):
+        def f(G):
+            return (x.to(G.device).matmul(G)).to(self.dev)
+        outputs = nn.parallel.parallel_apply(
+            [f] * len(self.grads), self.grads
+        )
+        """
+            x size: 1024
+            grads: 1024 x d
+        """
+        x = x.detach().cpu().numpy()
+        norm = np.linalg.norm(x)
+        self.wandb_data.update({f'lin_comb_coef_norm': norm})
+        return torch.cat(outputs)