ista-daslab-optimizers 1.1.8__py3-none-any.whl

ista_daslab_optimizers/sparse_mfac/sparse_mfac.py

@@ -0,0 +1,87 @@
+import wandb
+import torch
+from .sparse_core_mfac_w_ef import SparseCoreMFACwithEF
+from ..tools import get_first_device, get_gpus, get_weights_and_gradients, update_model, get_gpu_mem_usage
+
+
+class SparseMFAC(torch.optim.Optimizer):
+    def __init__(self, params, lr: float, damp: float, m: int, k_init: float, weight_decay: float, use_bf16: bool):
+        super(SparseMFAC, self).__init__(params, dict(lr=lr, weight_decay=weight_decay))
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.m = m
+        self.damp = damp
+        self.use_bf16 = use_bf16
+        self.k_init = k_init
+
+        self.device = get_first_device()
+        self.d = sum([p.numel() for group in self.param_groups for p in group['params']])
+
+        ##### Sparse M-FAC preconditioner
+        self.core_mfac = SparseCoreMFACwithEF(
+            m=self.m,
+            d=self.d,
+            k_init=self.k_init,
+            dev=self.device,
+            gpus=[self.device] if torch.distributed.is_initialized() else get_gpus(),
+            damp=damp,
+            use_bf16=use_bf16)
+
+        ##### scalar variables
+        self.steps = 0
+        self.log_interval = 100
+        self.grad_norms_sum = 0
+
+        self.wandb_data = dict()
+        self.cos = torch.nn.CosineSimilarity(dim=0, eps=1e-6)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        self.steps += 1
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        ##################################################
+        ########## [1] GET GRADIENT
+        ##################################################
+        g_dense = get_weights_and_gradients(self.param_groups, get_weights=False, get_grad=True, grad_bf16=self.use_bf16)
+        norm_g_dense = g_dense.norm(p=2)
+        self.grad_norms_sum += norm_g_dense
+
+        ##################################################
+        ########## [2] PRECONDITION
+        ##################################################
+        update = self.core_mfac.apply_ef_then_update_buffer_then_precondition(g_dense)
+
+        ##################################################
+        ########## [3] UPDATE THE MODEL
+        ##################################################
+        lr = self.param_groups[0]['lr']
+
+        update_model(
+            params=self.param_groups,
+            update=update,
+            weight_decay=self.weight_decay,
+            alpha=None,
+            multiply_wd_w_lr=True)
+
+        ##################################################
+        ########## LOGS
+        ##################################################
+        if self.log_interval > 0 and self.steps % self.log_interval == 0:
+            norm_error = self.core_mfac.error.norm(p=2)
+            self.wandb_data.update({
+                'epoch/step': self.steps,
+                'epoch/norm_g': norm_g_dense,
+                'epoch/norm_error': norm_error,
+                'epoch/ef_norm_div_grad_norm_sum': norm_error / self.grad_norms_sum,
+                'epoch/norm_u': update.norm(p=2),
+                'epoch/gpu_mem_usage': get_gpu_mem_usage(),
+            })
+            self.wandb_data.update(self.core_mfac.wandb_data)
+            wandb.log(self.wandb_data)
+
+        return loss

ista_daslab_optimizers/tools.py

@@ -0,0 +1,218 @@
+import os
+import gpustat
+import torch
+from enum import Enum
+from importlib import import_module
+import ista_daslab_cuda_tools
+
+def get_cuda_capability(device=0):
+    cc = torch.cuda.get_device_capability(device) # tuple, for example (8, 6) for CUDA Capability 8.6
+    return f'{cc[0]}{cc[1]}'
+
+class CopyDirection(Enum):
+    k2d = 0
+    d2k = 1
+
+class Strategy(Enum):
+    """Apply Top-K globally"""
+    GLOBAL = 1
+
+    """Apply Top-K in blocks of specific size"""
+    BLOCK = 2
+
+    @staticmethod
+    def factory(name: str):
+        if name == 'gl': return Strategy.GLOBAL
+        if name == 'bl': return Strategy.BLOCK
+        raise RuntimeError('Invalid strategy name')
+
+def get_first_device():
+    if not torch.cuda.is_available():
+        return torch.device('cpu')
+    if torch.distributed.is_initialized():
+        return torch.device(f'cuda:{torch.distributed.get_rank()}')
+    return torch.device('cuda:0')
+
+
+def get_gpus():
+    if not torch.cuda.is_available():
+        return ['cpu']
+    device = get_first_device()
+    if torch.cuda.device_count() == 1:
+        return [device]
+
+    return [
+        torch.device(f'cuda:{i}')
+        for i in range(len(os.environ['CUDA_VISIBLE_DEVICES'].split(',')))
+    ]
+
+def get_gpu_mem_usage():
+    """
+        This method returns the GPU memory usage for the current process.
+        It uses gpustat to query the GPU used by the current process (using CUDA_VISIBLE_DEVICES)
+
+        GPUSTAT usage:
+        stat = gpustat.new_query().gpus # this is a list containing information about each GPU indexed from 0 to 7
+        stat[i] (GPU #i) has the following keys:
+            - 'index'
+            - 'uuid'
+            - 'name'
+            - 'temperature.gpu'
+            - 'fan.speed'
+            - 'utilization.gpu'
+            - 'utilization.enc'
+            - 'utilization.dec'
+            - 'power.draw'
+            - 'enforced.power.limit'
+            - 'memory.used'
+            - 'memory.total'
+            - 'processes'
+        Among these keys, only the key 'processes' is used here.
+        stat[i].processes is a list of dicts, where each dict contains information about each process currently running on the GPU #i
+            - 'username'
+            - 'command'
+            - 'full_command'
+            - 'gpu_memory_usage'
+            - 'cpu_percent'
+            - 'cpu_memory_usage'
+            - 'pid'
+    """
+    gpus = gpustat.new_query().gpus
+    gids = list(map(int, os.environ['CUDA_VISIBLE_DEVICES'].split(',')))
+    gpu_mem = sum([int(proc['gpu_memory_usage']) for gid in gids for proc in gpus[gid]['processes'] if int(proc['pid']) == os.getpid()])
+    return gpu_mem
+
+def block_split(model_size, block_size):
+    if model_size < block_size:
+        return 1, model_size
+    ### this is the shorter version that only returns the number of full blocks of size "block_size"
+    ### and the starting position of the last and smallest block
+    blocks_count = int(model_size / block_size)
+    start_index_last_block = model_size - model_size % block_size
+    return blocks_count, start_index_last_block
+
+def get_weights_and_gradients(params, get_weights, get_grad=True, grad_bf16=False):
+    """
+        This method returns:
+        - w: the raw weights collected from the model if get_weights=True
+        - g: the gradients (without WD added)
+    """
+    w, g = [], []
+    for group in params:
+        for p in group['params']:
+            if p.grad is None or not p.requires_grad:
+                continue
+
+            if get_weights:
+                w.append(p.reshape(-1))
+            if get_grad:
+                if grad_bf16:
+                    if p.grad.dtype != torch.bfloat16:
+                        g.append(p.grad.reshape(-1).to(dtype=torch.bfloat16))
+                    else:
+                        g.append(p.grad.reshape(-1))
+                else:
+                    g.append(p.grad.reshape(-1))
+
+    if get_weights and get_grad:
+        return torch.cat(w), torch.cat(g)
+    if get_weights:
+        return torch.cat(w)
+    if get_grad:
+        return torch.cat(g)
+    raise RuntimeError(f'invalid combination of parameters: {get_weights=}, {get_grad=}')
+
+
+def update_model(params, update, weight_decay=0, alpha=None, multiply_wd_w_lr=False):
+    """
+        Applies the `update` to the model
+        When alpha=None, alpha is set to lr in the group
+        Returns the shrinking factor for the weights
+    """
+    count = 0
+    for group in params:
+        lr = group['lr']
+        wd = group.get('weight_decay', weight_decay) # if the param groups do not have weight decay, then use the externally provided one
+        for p in group['params']:
+            if p.grad is None:
+                continue
+            u = update[count:(count + p.numel())].reshape(p.shape).to(p.device)
+            if wd > 0:
+                if multiply_wd_w_lr:
+                    p.mul_(1 - lr * wd)
+                else:
+                    p.mul_(1 - wd)
+            p.add_(u, alpha=-lr if alpha is None else alpha)
+            count += p.numel()
+
+class KernelVersionsManager:
+    def __init__(self, version_SP, version_LCG, m, d, d_block_size):
+        self.version_SP = version_SP
+        self.version_LCG = version_LCG
+        self.m = m
+        self.d = d
+        self.d_block_size = d_block_size
+
+        self.BLOCK_INDEX = 0
+        self.THREAD_INDEX = 1
+
+        # set number of blocks (initially None) based on the number of threads (see page 80 in the PhD #8)
+        # if self.d > 300_000_000:
+        #     print(f'Model size is larger than 300M. Switching SP version from {self.version_SP} to 252')
+        #     self.version_SP = 252
+
+        self.SP_BLOCKS_THREADS = {
+            23: [self.m, self.m],
+            # 24: [1024, 1024],
+            # 251: [None, 1024],
+            # 252: [None, self.m],
+            # 261: [None, 128],
+            # 262: [None, 128],
+            # 272: [None, 1024],
+        }
+
+        self.LCG_BLOCKS_THREADS = {
+            # 42: [68, 256],
+            # 43: [117, 32],
+            51: [None, 1024],
+            # 524: [None, 128],
+            # 53: [None, 128],
+            # 54: [None, 128],
+        }
+
+        self.set_blocks_count(self.SP_BLOCKS_THREADS, self.version_SP, op='SP')
+        self.set_blocks_count(self.LCG_BLOCKS_THREADS, self.version_LCG, op='LCG')
+        # self.SP_BLOCKS_THREADS[self.version_SP][0] = 10
+
+    def set_blocks_count(self, op_blocks_threads, op_version, op):
+        """
+        Safety measure: for small models, there might be too many thread blocks launched and most of them will process data out of bounds of arrays out, indices and values
+        """
+        def div_inc(a, b):
+            r = a // b
+            return (r + 1) if (a % b > 0) else r
+
+        if op_blocks_threads[op_version][self.BLOCK_INDEX] is None:
+            blocks_count = div_inc(self.d, self.d_block_size)
+            op_max_blocks = ista_daslab_cuda_tools.get_sm_count()
+            op_required_blocks = min(blocks_count, op_max_blocks)
+            if op_required_blocks < op_max_blocks:
+                print(f'Maximum number of blocks for {op} is {op_max_blocks}, but this model requires only {op_required_blocks}')
+                # return op_required_blocks
+                op_blocks_threads[op_version][self.BLOCK_INDEX] = op_required_blocks
+            op_blocks_threads[op_version][self.BLOCK_INDEX] = op_max_blocks
+
+        print(f'{op_blocks_threads=}, {op_version=}, {op=}, {op_blocks_threads[op_version][self.BLOCK_INDEX]=}')
+
+    def get_SP_blocks(self):
+        return self.SP_BLOCKS_THREADS[self.version_SP][self.BLOCK_INDEX]
+
+    def get_SP_threads(self):
+        return self.SP_BLOCKS_THREADS[self.version_SP][self.THREAD_INDEX]
+
+    def get_LCG_blocks(self):
+        return self.LCG_BLOCKS_THREADS[self.version_LCG][self.BLOCK_INDEX]
+
+    def get_LCG_threads(self):
+        return self.LCG_BLOCKS_THREADS[self.version_LCG][self.THREAD_INDEX]
+

ista_daslab_optimizers/utils/dct.py

@@ -0,0 +1,45 @@
+import torch
+import math
+
+from ista_daslab_optimizers.utils.global_cache import GlobalCache
+
+def dct3_matrix(n, dtype, device):
+    """
+        This function returns the orthogonal transformation for Discrete Cosine Transform (DCT-3).
+    """
+    lin = torch.arange(n)
+    I = lin.repeat(n, 1).to(torch.float)
+    Q = math.sqrt(2 / n) * torch.cos(torch.pi * (I.t() * (2. * I + 1.)) / (2. * n))
+    del lin, I
+    Q[0, :] *= math.sqrt(0.5)
+    return Q.to(device=device, dtype=dtype)
+
+def dct_type2_makhoul(X):
+    N = X.shape[1]
+
+    if GlobalCache.contains(category='perm', key=N):
+        perm = GlobalCache.get(category='perm', key=N)
+    else:
+        even_idx = torch.arange(0, N, 2)  # 0, 2, 4, ...
+        odd_idx = torch.arange(1, N, 2).flip(0)  # last odd → first odd
+        perm = torch.cat([even_idx, odd_idx]).to(X.device)
+
+        GlobalCache.add(category='perm', key=N, item=perm)
+    #
+    # X_input = X[:, perm]
+    # if X_input.dtype != torch.float:
+    #     X_input = X_input.to(torch.float)
+    # X_fft = torch.fft.fft(X_input, dim=1)
+
+    X_fft = torch.fft.fft(X[:, perm].contiguous(), dim=1)
+
+    if GlobalCache.contains(category='twiddle', key=N):
+        W = GlobalCache.get(category='twiddle', key=N)
+    else:
+        W = 2 * torch.exp((-1j * torch.pi * torch.arange(N, device=X.device) / (2 * N)))
+        W[0] /= math.sqrt(4 * N)
+        W[1:] /= math.sqrt(2 * N)
+
+        GlobalCache.add(category='twiddle', key=N, item=W.reshape(1, N))
+
+    return (X_fft * W).real

ista_daslab_optimizers/utils/global_cache.py

@@ -0,0 +1,45 @@
+class GlobalCache:
+    _instance = None
+
+    @staticmethod
+    def init():
+        if GlobalCache._instance is None:
+            GlobalCache._instance = GlobalCache()
+
+    @staticmethod
+    def contains(category, key):
+        return GlobalCache.get_instance()._contains(category, key)
+
+    @staticmethod
+    def get_instance():
+        if GlobalCache._instance is None:
+            GlobalCache.init()
+        return GlobalCache._instance
+
+    @staticmethod
+    def get(category, key):
+        return GlobalCache.get_instance()._get(category, key)
+
+    @staticmethod
+    def add(category, key, item):
+        return GlobalCache.get_instance()._add(category, key, item)
+
+    def __init__(self):
+        self.storage = dict()
+
+    def _contains(self, category, key):
+        if category not in self.storage:
+            return False
+        return key in self.storage[category]
+
+    def _add(self, category, key, item):
+        if category not in self.storage:
+            self.storage[category] = { key: item }
+        elif key not in self.storage[category]:
+            self.storage[category][key] = item
+
+    def _get(self, category, key):
+        # print(self.storage)
+        if self._contains(category, key):
+            return self.storage[category][key]
+        raise ValueError(f'GlobalCache does not contain category {category} and/or key {key}')

ista_daslab_optimizers/utils/matrix_storage.py

@@ -0,0 +1,58 @@
+import torch
+import torch.distributed as dist
+from memory_efficient_optimizers.utils.dct import dct3_matrix
+
+PROJ_DCT = 'dct'
+PROJ_HDM = 'hdm'
+
+ALL_PROJ = [
+    PROJ_DCT, # DCT projection
+    PROJ_HDM, # Hadamard projection
+]
+
+class MatrixStorage:
+    """
+        This singleton class stores a dictionary where:
+        - keys = the matrix size
+        - values = the corresponding orthogonal matrix of DCT-3 or Hadamard transforms of size stored in the key
+    """
+    _instance = None
+
+    @staticmethod
+    def init():
+        if MatrixStorage._instance is None:
+            MatrixStorage._instance = MatrixStorage()
+
+    @staticmethod
+    def get_instance():
+        if MatrixStorage._instance is None:
+            MatrixStorage.init()
+        return MatrixStorage._instance
+
+    @staticmethod
+    def get_matrix(size, proj, dtype):
+        return MatrixStorage.get_instance()._get_matrix(size, proj, dtype)
+
+    @staticmethod
+    def add_matrix(size, proj, dtype):
+        return MatrixStorage.get_instance()._add_matrix(size, proj, dtype)
+
+    def __init__(self):
+        self.storage = dict()
+        self.dtype = None
+        self.device = f'cuda:{dist.get_rank()}' if dist.is_initialized() else 'cuda:0'
+
+    def _add_matrix(self, size, proj, dtype):
+        if size not in self.storage:
+            if proj == PROJ_DCT:
+                self.storage[size] = dct3_matrix(size).to(device=self.device, dtype=dtype) # first row is zero
+            elif proj == PROJ_HDM:
+                self.storage[size] = hadamard_transform(torch.eye(size).to(device=self.device, dtype=dtype), scale=1./math.sqrt(size))
+            else:
+                raise RuntimeError(f'Projection {proj} is currently not supported!')
+
+    def _get_matrix(self, size, proj, dtype):
+        if size not in self.storage:
+            assert dtype is not None
+            self._add_matrix(size, proj, dtype)
+        return self.storage[size]