ista-daslab-optimizers 1.1.2__tar.gz → 1.1.3__tar.gz

{ista_daslab_optimizers-1.1.2/ista_daslab_optimizers.egg-info → ista_daslab_optimizers-1.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ista_daslab_optimizers
-Version: 1.1.2
+Version: 1.1.3
 Summary: Deep Learning optimizers developed in the Distributed Algorithms and Systems group (DASLab) @ Institute of Science and Technology Austria (ISTA)
 Author-email: Ionut-Vlad Modoranu <ionut-vlad.modoranu@ist.ac.at>
 Maintainer-email: Ionut-Vlad Modoranu <ionut-vlad.modoranu@ist.ac.at>
@@ -289,6 +289,8 @@ optimizer = MicroAdam(
 # Versions summary:
 
 ---
+- **1.1.3** @ September 5th, 2024:
+  - allow using `SparseCoreMFACwithEF` separately by importing it in `sparse_mfac.__init__.py`
 - **1.1.2** @ August 1st, 2024:
   - ***[1.1.0]:*** added support to densify the final update: introduced parameter alpha that controls the fraction of error feedback
   (EF) to be integrated into the update to make it dense. Finally, the fraction alpha will be discarded from the EF at

{ista_daslab_optimizers-1.1.2 → ista_daslab_optimizers-1.1.3}/README.md

@@ -64,6 +64,8 @@ optimizer = MicroAdam(
 # Versions summary:
 
 ---
+- **1.1.3** @ September 5th, 2024:
+  - allow using `SparseCoreMFACwithEF` separately by importing it in `sparse_mfac.__init__.py`
 - **1.1.2** @ August 1st, 2024:
   - ***[1.1.0]:*** added support to densify the final update: introduced parameter alpha that controls the fraction of error feedback
   (EF) to be integrated into the update to make it dense. Finally, the fraction alpha will be discarded from the EF at

{ista_daslab_optimizers-1.1.2 → ista_daslab_optimizers-1.1.3}/ista_daslab_optimizers/micro_adam/micro_adam.py

@@ -15,7 +15,7 @@ class MicroAdam(torch.optim.Optimizer):
         defaults = dict(lr=lr, weight_decay=weight_decay, eps=eps, alpha=alpha)
         super(MicroAdam, self).__init__(params, defaults)
 
-        assert 0 <= alpha < 1, 'Alpha must be in the [0, 1) interval'
+        assert (0 <= alpha < 1) or alpha == -2, 'Alpha must be in the [0, 1) interval or -2'
 
         self.m = m
         self.lr = lr
@@ -27,7 +27,9 @@ class MicroAdam(torch.optim.Optimizer):
         self.beta2 = betas[1]
         self.eps = eps
 
-        self.densify_update = (self.alpha > 0)
+        self.densify_update_using_ef = (self.alpha > 0)
+        self.densify_update_using_quant_error = (self.alpha == -2)
+
         self.model_size = sum([p.numel() for group in self.param_groups for p in group['params']])
 
         self.steps = 0  # how many optimization steps were performed so far
@@ -41,8 +43,12 @@ class MicroAdam(torch.optim.Optimizer):
         self.max_floats = ista_daslab_tools.get_max_floats_for_shared_memory_per_thread_block()
         self.d_block_size = self.max_floats // 2 // int(100 / self.shared_memory_carveout)
 
-        self.fsdp_dict_size_count = [{} for _ in range(
-            torch.distributed.get_world_size())]  # key = layer size, value = how many layers of that size the model has (per worker)
+        if torch.distributed.is_initialized():
+            self.fsdp_dict_size_count = [{} for _ in range(
+                torch.distributed.get_world_size())]  # key = layer size, value = how many layers of that size the model has (per worker)
+        else:
+            self.fsdp_dict_size_count = [{}]
+
         self.dict_size_count = {}  # key = layer size, value = how many layers of that size the model has
         for param in self.param_groups:
             for p in param['params']:
@@ -56,7 +62,10 @@ class MicroAdam(torch.optim.Optimizer):
         layer_size = p.numel()
         st = self.state[p]
 
-        rank = torch.distributed.get_rank()
+        rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+
+        if self.densify_update_using_quant_error:
+            st['quant_err'] = torch.zeros_like(p)
 
         st['blocks'] = max(1, int(math.floor(self.blocks * layer_size * self.fsdp_dict_size_count[rank][layer_size] / self.model_size)))
 
@@ -96,7 +105,7 @@ class MicroAdam(torch.optim.Optimizer):
                 loss = closure()
 
         if self.steps == 1:
-            rank = torch.distributed.get_rank()
+            rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
             for param in self.param_groups:
                 for p in param['params']:
                     if p is not None:
@@ -106,7 +115,7 @@ class MicroAdam(torch.optim.Optimizer):
 
         time_start = time.time()
 
-        norm_g, norm_u, norm_e, sparsity_u = 0, 0, 0, 0
+        norm_qe, norm_g, norm_u, norm_e, sparsity_u, sparsity_qe = 0, 0, 0, 0, 0, 0
 
         for group in self.param_groups:
             lr = group['lr']
@@ -119,23 +128,28 @@ class MicroAdam(torch.optim.Optimizer):
                 if p is None:
                     continue
 
-                ng, nu, ne, sp_u = self.update_step(p, lr, wd)
+                nqe, ng, nu, ne, sp_u, sp_qe = self.update_step(p, lr, wd)
+                norm_qe += nqe
                 norm_g += ng
                 norm_u += nu
                 norm_e += ne
                 sparsity_u += sp_u
+                sparsity_qe += sp_qe
 
         # torch.cuda.synchronize()
         time_end = time.time()
         elapsed_step = time_end - time_start
-        self._log(norm_g, norm_u, norm_e, sparsity_u, elapsed_step)
+        self._log(norm_qe, norm_g, norm_u, norm_e, sparsity_u, sparsity_qe, elapsed_step)
 
         return loss
 
     @torch.no_grad()
     def update_step(self, p, lr, wd):
-        norm_g, norm_u, norm_e, sp_u = 0, 0, 0, 0
+        norm_qe, norm_g, norm_u, norm_e, sp_u, sp_qe = 0, 0, 0, 0, 0, 0
 
+        # if p.grad.dtype != torch.bfloat16:
+        #     grad = p.grad.to(dtype=torch.bfloat16).reshape(-1)
+        # else:
         grad = p.grad.view(-1)
 
         if self.steps % self.log_interval == 0:
@@ -216,6 +230,48 @@ class MicroAdam(torch.optim.Optimizer):
         ##### STEP 8
         ista_daslab_micro_adam.asymm_block_quant(d, self.quant_block_size, error, min_vals, max_vals, grad)  # error = Q(a, min, max)
 
+        # weight decay step
+        if wd > 0:
+            p.mul_(1 - lr * wd)
+
+        ##### NEW: densify using quant error
+        if self.densify_update_using_quant_error:
+            # When entering this if-statement, we have:
+            #     - p is theta_t
+            #     - p.grad is a_t (from step 6 in algorithm 1)
+            #     - error is e_t+1 (from step 8 in algorithm 1)
+            #
+            # Below we have the formula to update the model parameters:
+            # [a = -1] with lr
+            #     theta_t+1 = theta_t - lr * (a_t - Qinv(e_t+1)) - lr * u_t
+            #               = theta_t - lr * a_t + lr * Qinv(e_t+1) - lr * u_t
+            #               = theta_t - lr * a_t              # STEP A below, in this if statmenet
+            #                         + lr * Qinv(e_t+1)      # STEP B below, in this if statmenet
+            #                         - lr * u_t              # this is steps 10-11
+            #
+            # [a = -2] without lr
+            #     theta_t+1 = theta_t - (a_t - Qinv(e_t+1)) - lr * u_t
+            #               = theta_t - a_t + Qinv(e_t+1) - lr * u_t
+            #               = theta_t - a_t              # STEP A below, in this if statmenet
+            #                         + Qinv(e_t+1)      # STEP B below, in this if statmenet
+            #                         - lr * u_t              # this is steps 10-11
+            quant_err = st['quant_err']
+            quant_err.zero_()
+            quant_err.add_(p.grad)
+
+            ##### STEP A
+            p.add_(p.grad, alpha=-1)
+
+            ##### STEP B
+            p.grad.zero_() # zerorize to prepare the accumulator for Qinv
+            ista_daslab_micro_adam.asymm_block_quant_inv(d, self.quant_block_size, error, min_vals, max_vals, grad, 1)
+            p.add_(p.grad)
+
+            quant_err.sub_(p.grad)
+
+            norm_qe = quant_err.norm(p=2) ** 2
+            sp_qe = (quant_err == 0).sum()
+
         ##### STEPS 10-11
         grad.zero_()
         ista_daslab_micro_adam.compute_microadam_update(blocks,  # blocks
@@ -237,16 +293,22 @@ class MicroAdam(torch.optim.Optimizer):
         ##### STEP 12: # side idea: only decay the weights that are update
 
         ##### if PRETRAINING #1
-        if self.densify_update: # we add alpha * EF to update that is stored in grad buffer
+        if self.densify_update_using_ef: # we add alpha * EF to update that is stored in grad buffer
             # p.grad += alpha * Qinv(error), alpha=0.1
             ista_daslab_micro_adam.asymm_block_quant_inv(d, self.quant_block_size, error, min_vals, max_vals, grad, self.alpha)
         ##### END IF PRETRAINING #1
 
         # if alpha > 0, then the update u=p.grad is dense now
-        p.mul_(1 - lr * wd).add_(p.grad, alpha=-lr)
+
+        # update model using MicroAdam update stored in p.grad
+        p.add_(p.grad, alpha=-lr)
+
+        if self.steps % self.log_interval == 0:
+            norm_u = grad.norm(p=2) ** 2
+            sp_u = (grad == 0).sum()  # check sparsity before zerorizing
 
         ##### if PRETRAINING #2
-        if self.densify_update:
+        if self.densify_update_using_ef:
             grad.zero_()
             ista_daslab_micro_adam.asymm_block_quant_inv(d, self.quant_block_size, error, min_vals, max_vals, grad, 1-self.alpha)
 
@@ -256,27 +318,29 @@ class MicroAdam(torch.optim.Optimizer):
 
         # compute error norm
         if self.steps % self.log_interval == 0:
-            norm_u = grad.norm(p=2) ** 2
-            sp_u = (grad == 0).sum()  # check sparsity before zerorizing
-
             grad.zero_()
             ista_daslab_micro_adam.asymm_block_quant_inv(d, self.quant_block_size, error, min_vals, max_vals, grad, 1.0)
 
             norm_e = grad.norm(p=2) ** 2
 
-        return norm_g, norm_u, norm_e, sp_u
+        # p.grad = p.grad.to(dtype=original_grad_type)
+
+        return norm_qe, norm_g, norm_u, norm_e, sp_u, sp_qe
 
-    def _log(self, norm_g, norm_u, norm_e, sparsity_u, elapsed_step):
+    def _log(self, norm_qe, norm_g, norm_u, norm_e, sparsity_u, sparsity_qe, elapsed_step):
         if self.steps % self.log_interval == 0:
-            sync_data = torch.tensor([norm_g, norm_u, norm_e, sparsity_u, elapsed_step], dtype=torch.float,
-                                     requires_grad=False).cuda()  # correct, loss, size
-            all_reduce(sync_data, op=ReduceOp.SUM)
-            norm_g, norm_u, norm_e, sparsity_u, elapsed_step = sync_data
+            if is_initialized():
+                sync_data = torch.tensor([norm_qe, norm_g, norm_u, norm_e, sparsity_u, sparsity_qe, elapsed_step], dtype=torch.float,
+                                         requires_grad=False).cuda()  # correct, loss, size
+                all_reduce(sync_data, op=ReduceOp.SUM)
+                norm_qe, norm_g, norm_u, norm_e, sparsity_u, sparsity_qe, elapsed_step = sync_data
 
             if not is_initialized() or get_rank() == 0:
                 wandb_data = {
                     'step/optimizer_steps': self.steps,
                     'step/gpu_mem_usage': get_gpu_mem_usage(),
+                    'step/norm_quant_err': math.sqrt(norm_qe),
+                    'step/sparsity_quant_err': sparsity_qe / self.model_size * 100.,
                     'step/norm_g': math.sqrt(norm_g),
                     'step/norm_u': math.sqrt(norm_u),
                     'step/norm_error': math.sqrt(norm_e),
@@ -335,4 +399,4 @@ class MicroAdam(torch.optim.Optimizer):
     #             st['quant_full_blocks_count'], st['d_index_quant'] = block_split(st['d'], self.quant_block_size)
     #             st['error'] = torch.zeros(int(math.ceil(st['d'] / 2)), dtype=torch.uint8, device=self.device)  # ceil(d/2) bytes
     #             st['min_vals'] = torch.zeros(st['quant_full_blocks_count'] + 1, dtype=torch.bfloat16, device=self.device)  # ceil(d/q_bsz)*2 bytes
-    #             st['max_vals'] = torch.zeros(st['quant_full_blocks_count'] + 1, dtype=torch.bfloat16, device=self.device)  # ceil(d/q_bsz)*2 bytes
+    #             st['max_vals'] = torch.zeros(st['quant_full_blocks_count'] + 1, dtype=torch.bfloat16, device=self.device)  # ceil(d/q_bsz)*2 bytes

ista_daslab_optimizers-1.1.3/ista_daslab_optimizers/sparse_mfac/__init__.py

@@ -0,0 +1,7 @@
+from .sparse_mfac import SparseMFAC
+from .sparse_core_mfac_w_ef import SparseCoreMFACwithEF
+
+__all__ = [
+    'SparseMFAC',
+    'SparseCoreMFACwithEF'
+]

{ista_daslab_optimizers-1.1.2 → ista_daslab_optimizers-1.1.3/ista_daslab_optimizers.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ista_daslab_optimizers
-Version: 1.1.2
+Version: 1.1.3
 Summary: Deep Learning optimizers developed in the Distributed Algorithms and Systems group (DASLab) @ Institute of Science and Technology Austria (ISTA)
 Author-email: Ionut-Vlad Modoranu <ionut-vlad.modoranu@ist.ac.at>
 Maintainer-email: Ionut-Vlad Modoranu <ionut-vlad.modoranu@ist.ac.at>
@@ -289,6 +289,8 @@ optimizer = MicroAdam(
 # Versions summary:
 
 ---
+- **1.1.3** @ September 5th, 2024:
+  - allow using `SparseCoreMFACwithEF` separately by importing it in `sparse_mfac.__init__.py`
 - **1.1.2** @ August 1st, 2024:
   - ***[1.1.0]:*** added support to densify the final update: introduced parameter alpha that controls the fraction of error feedback
   (EF) to be integrated into the update to make it dense. Finally, the fraction alpha will be discarded from the EF at

{ista_daslab_optimizers-1.1.2 → ista_daslab_optimizers-1.1.3}/ista_daslab_optimizers.egg-info/SOURCES.txt

@@ -3,7 +3,6 @@ MANIFEST.in
 README.md
 pyproject.toml
 setup.py
-./kernels/utils.h
 ./kernels/dense_mfac/dense_mfac.cpp
 ./kernels/dense_mfac/dense_mfac_kernel.cu
 ./kernels/micro_adam/micro_adam.cpp
@@ -32,4 +31,16 @@ ista_daslab_optimizers/micro_adam/__init__.py
 ista_daslab_optimizers/micro_adam/micro_adam.py
 ista_daslab_optimizers/sparse_mfac/__init__.py
 ista_daslab_optimizers/sparse_mfac/sparse_core_mfac_w_ef.py
-ista_daslab_optimizers/sparse_mfac/sparse_mfac.py
+ista_daslab_optimizers/sparse_mfac/sparse_mfac.py
+kernels/utils.h
+kernels/dense_mfac/dense_mfac.cpp
+kernels/dense_mfac/dense_mfac_kernel.cu
+kernels/micro_adam/micro_adam.cpp
+kernels/micro_adam/micro_adam_asymm_block_quant.cu
+kernels/micro_adam/micro_adam_asymm_block_quant_inv.cu
+kernels/micro_adam/micro_adam_update.cu
+kernels/sparse_mfac/sparse_mfac.cpp
+kernels/sparse_mfac/sparse_mfac_LCG_kernel.cu
+kernels/sparse_mfac/sparse_mfac_SP_kernel.cu
+kernels/tools/tools.cpp
+kernels/tools/tools_kernel.cu

{ista_daslab_optimizers-1.1.2 → ista_daslab_optimizers-1.1.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name='ista_daslab_optimizers'
-version='1.1.2'
+version='1.1.3'
 dependencies = [
     "torch", # >=2.3.1",
     "torchaudio", # >=2.3.1",

ista_daslab_optimizers-1.1.2/ista_daslab_optimizers/sparse_mfac/__init__.py

@@ -1,5 +0,0 @@
-from .sparse_mfac import SparseMFAC
-
-__all__ = [
-    'SparseMFAC',
-]