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

ista_daslab_optimizers/dense_mfac/dense_mfac.py

@@ -0,0 +1,93 @@
+import wandb
+import torch
+from ..tools import get_weights_and_gradients, update_model, get_first_device, get_gpus
+from .dense_core_mfac import DenseCoreMFAC
+
+# Disable tensor cores as they can mess with precision
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cudnn.allow_tf32 = False
+
+class DenseMFAC(torch.optim.Optimizer):
+    def __init__(self,
+                 params,
+                 lr: float,
+                 weight_decay: float,
+                 ngrads: int,
+                 damp: float,
+                 create_G=False):
+
+        super(DenseMFAC, self).__init__(params, dict(lr=lr, weight_decay=weight_decay))
+
+        self.m = ngrads
+        self.lr = lr
+        self.damp = damp
+        self.weight_decay = weight_decay
+        self.device = get_first_device()
+        self.create_G = create_G
+
+        self.model_size = None
+        self.steps = 0
+        self.wandb_data = dict()
+
+        self.model_size = sum([p.numel() for group in self.param_groups for p in group['params']])
+
+        self.hinv = None
+
+    def _create_hinv(self):
+        self.hinv = DenseCoreMFAC(
+            grads=torch.zeros((self.m, self.model_size), dtype=torch.float),
+            dev=self.device,
+            gpus=get_gpus(),
+            damp=self.damp,
+            create_G=self.create_G)
+
+    @torch.no_grad()
+    def empty_buffer(self):
+        self.hinv.empty_buffer()
+
+
+    @torch.no_grad()
+    def integrate_gradient(self, g):
+        _ = self.hinv.integrate_gradient(g)
+
+    @torch.no_grad()
+    def compute_update(self, g, x):
+        update_method = self.hinv.integrate_gradient_and_precondition
+        # if self.use_sq_newton:
+        #     update_method = self.hinv.integrate_gradient_and_precondition_twice
+
+        update = update_method(g, x).to(self.device)
+        return update
+
+    @torch.no_grad()
+    def log_data(self, update, g):
+        lr = self.param_groups[0]['lr']
+        self.wandb_data.update(dict(norm_upd_w_lr=lr * update.norm(p=2), norm_g=g.norm(p=2)))
+        self.wandb_data.update(self.hinv.wandb_data)
+        # self.wandb_data.update(quantify_preconditioning(g=g, u=update.to(g.device), return_distribution=False, use_abs=True, optim_name=self.optim_name))
+        wandb.log(self.wandb_data)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        self.steps += 1
+
+        if self.hinv is None:
+            self._create_hinv()
+
+        g = get_weights_and_gradients(self.param_groups, get_weights=False)
+        update = self.compute_update(g, g)
+
+        update_model(params=self.param_groups, update=update, alpha=None)
+
+        if self.steps % self.m == 0:
+            self.log_data(update, g)
+
+
+# def kfac_update_rescaling(self, g, u):
+#     # rescaling_kfac64,  # use the alpha rescaling from K-FAC paper, section 6.4
+#     T1 = torch.dot(g, u)
+#     T2 = (self.hinv.grads_matmul(u).norm(p=2) ** 2).mean()
+#     T3 = self.damp * (u.norm(p=2) ** 2)
+#     alpha = -T1 / (T2 + T3)
+#     self.wandb_data.update(dict(kfac_T1=T1, kfac_T2=T2, kfac_T3=T3, kfac_alpha=alpha))
+#     return alpha

ista_daslab_optimizers/fft_low_rank/dct_adamw.py

@@ -0,0 +1,351 @@
+import torch
+import torch.distributed as dist
+
+import math
+import numpy as np
+
+from fast_hadamard_transform import hadamard_transform
+from ista_daslab_optimizers.utils.dct import dct3_matrix
+from ista_daslab_optimizers.utils.quantizers import Quantizer4bit, Quantizer8bit
+from ista_daslab_optimizers.fft_low_rank.fft_projector import FFTLowRankProjector
+
+PROJ_DCT = 'dct'
+PROJ_HDM = 'hdm'
+PROJ_RAND_QR = 'rqr'
+
+ALL_PROJ = [
+    PROJ_DCT, # DCT projection
+    PROJ_HDM, # Hadamard projection
+    PROJ_RAND_QR, # Random-QR projection
+]
+
+STATE_M = 'm'
+STATE_V = 'v'
+STATE_Q = 'Q'
+STATE_ID = 'param-id'
+STATE_EF = 'ef'
+# STATE_EF_MIN = 'ef-min-vals'
+# STATE_EF_MAX = 'ef-max-vals'
+STATE_FFT_LRP = 'fft-low-rank-projector'
+STATE_BROADCAST_SOURCE = 'broadcast-src' # the process rank that computes the update for a parameter p will broadcast the parameter p to other workers
+
+
+class DCTAdamW(torch.optim.Optimizer):
+    def __init__(self,
+                 params,
+                 lr,
+                 weight_decay,
+                 rank,
+                 proj,
+                 use_ef=False,
+                 q_ef=False,
+                 distributed=False,
+                 update_proj_gap=1,
+                 rotate_subspace=False,
+                 sim_type='matmul',
+                 ell_norm=1,
+                 max_shape=32_000,
+                 betas=(0.9, 0.999),
+                 eps=1e-8):
+        assert proj in ALL_PROJ
+
+        super().__init__(params, dict(lr=lr, weight_decay=weight_decay))
+
+        self.rank = rank
+        self.proj = proj
+        self.use_ef = use_ef
+        self.q_ef = q_ef
+        self.distributed = distributed
+        self.update_proj_gap = update_proj_gap
+        self.rotate_subspace = rotate_subspace
+        self.sim_type = sim_type
+        self.ell_norm = ell_norm
+        self.max_shape = max_shape # apply low-rank to 2D parameters that have both dimensions smaller than max_shape
+        self.betas = betas
+        self.eps = eps
+
+        self.steps = 0
+        self.is_state_initialized = False
+        self.Q = None # the full transformation matrix (non-truncated, all rows and columns)
+        self.Q_cols_norm = None
+        self.use_theoretical_similarity = (self.ell_norm < 0)
+        self.ell_norm = abs(self.ell_norm)
+
+        if proj == PROJ_DCT:
+            assert sim_type in ['matmul', 'makhoul']
+        else:
+            assert sim_type == 'matmul'
+
+    def setup_Q(self, p):
+        if self.Q is None:
+            size = min(p.shape)
+            if self.proj == PROJ_DCT:
+                Qdct3 = dct3_matrix(size, p.dtype, p.device)  # first row is zero
+                if self.sim_type == 'makhoul':
+                    self.Q = Qdct3.t()
+                    print(f'\n\t!!!!! Initialized DCT-2 matrix of size {size} !!!!!\n')
+                elif self.sim_type == 'matmul':
+                    self.Q = Qdct3
+                    print(f'\n\t!!!!! Initialized DCT-3 matrix of size {size} !!!!!\n')
+                else:
+                    raise RuntimeError(f'Unknown sim_type: {self.sim_type}')
+            elif self.proj == PROJ_HDM:
+                self.Q = hadamard_transform(torch.eye(size).to(device=p.device, dtype=p.dtype), scale=1. / math.sqrt(size))
+                print(f'\n\t!!!!! Initialized Hadamard matrix of size {size} !!!!!\n')
+            elif self.proj == PROJ_RAND_QR:
+                random = torch.randn(size, size, dtype=p.dtype, device=p.device)
+                self.Q, _ = torch.linalg.qr(random)
+                del random
+            else:
+                raise RuntimeError(f'Projection {self.proj} is currently not supported!')
+
+            if self.use_theoretical_similarity:
+                self.Q_cols_norm = self.Q.norm(p=self.ell_norm, dim=0)
+
+    def should_compute_update(self, p):
+        """
+            This function returns a boolean indicating whether the update for the parameter p should be computed on the current GPU
+        """
+        state = self.state[p]
+        param_id = state[STATE_ID]
+        return param_id % dist.get_world_size() == dist.get_rank()
+
+    def should_update_projection(self):
+        return self.steps == 1 or self.steps % self.update_proj_gap == 0
+
+    def init_state(self, p):
+        state = self.state[p]
+        if p.ndim == 1: # adam update
+            print(f'Parameter of size {tuple(p.shape)} will receive original AdamW update with state shape {tuple(p.shape)}')
+            state[STATE_M] = torch.zeros_like(p)
+            state[STATE_V] = torch.zeros_like(p)
+        elif p.ndim == 2: # low-rank adam update
+            n, m = p.shape
+            if n >= self.max_shape or m >= self.max_shape:  # apply full-rank
+                print(f'Parameter of size {tuple(p.shape)} will receive original AdamW update with state shape {tuple(p.shape)}')
+                state[STATE_M] = torch.zeros_like(p)
+                state[STATE_V] = torch.zeros_like(p)
+            else: # apply low-rank using the DCT transform as orthogonal matrix
+                if n >= m:
+                    low_rank_shape = (n, self.rank)
+                else:
+                    # fix for Llama-3-8B that has a layer of size (1024, 4096)
+                    # fix for Qwen2.5-7B that has a layer of size (512, 3584)
+                    if n in [512, 1024] and m in [3584, 4096]:
+                        low_rank_shape = (n, self.rank)
+                    else:
+                        low_rank_shape = (self.rank, m)
+                # low_rank_shape = (n, self.rank) if n >= m else (self.rank, m)
+                print(f'Parameter of size {tuple(p.shape)} will receive low-rank update with state shape {low_rank_shape}')
+                state[STATE_M] = torch.zeros(*low_rank_shape, dtype=p.dtype, device=p.device)
+                state[STATE_V] = torch.zeros(*low_rank_shape, dtype=p.dtype, device=p.device)
+                state[STATE_FFT_LRP] = FFTLowRankProjector(p,
+                                                          rank=self.rank,
+                                                          proj=self.proj,
+                                                          rotate_subspace=self.rotate_subspace,
+                                                          sim_type=self.sim_type,
+                                                          ell_norm=self.ell_norm,
+                                                          use_th_sim=self.use_theoretical_similarity)
+                if self.use_ef:
+                    if self.q_ef > 0:
+                        # state[STATE_EF] = torch.zeros(p.numel() // 2, dtype=torch.uint8, device=p.device)
+                        # state[STATE_EF_MIN] = torch.zeros(p.shape[0], dtype=torch.bfloat16, device=p.device)
+                        # state[STATE_EF_MAX] = torch.zeros(p.shape[0], dtype=torch.bfloat16, device=p.device)
+                        quantClass = {4: Quantizer4bit, 8: Quantizer8bit}[self.q_ef]
+                        if self.q_ef == 4:
+                            quantClass = Quantizer4bit
+                            print(f'\n\t!!!!! Quantizing EF to 4 bits !!!!!\n')
+                        elif self.q_ef == 8:
+                            quantClass = Quantizer8bit
+                            print(f'\n\t!!!!! Quantizing EF to 8 bits !!!!!\n')
+                        else:
+                            raise RuntimeError(f'Quantization on {self.q_ef} bits is currently not supported!')
+                        state[STATE_EF] = quantClass(shape=p.shape, device=p.device, dtype=p.dtype, bucket_size=p.shape[1])
+                    else:
+                        state[STATE_EF] = torch.zeros_like(p)
+
+                ### initialize Q
+                print('calling setup_Q')
+                self.setup_Q(p)
+        # end if
+
+    def init(self):
+        # init broadcast info
+        self.is_state_initialized = True
+        bcast_src_list = []
+        param_id = 0 # parameter id
+        for group in self.param_groups:
+            for p in group['params']:
+                if p is None: continue
+                if p.grad is None: continue
+
+                state = self.state[p]
+                if len(state) == 0:
+                    if self.distributed:
+                        state[STATE_ID] = param_id
+                        param_id += 1
+                        if self.should_compute_update(p):
+                            # if the current process computes the update, then it will also broadcast the parameters to all other workers
+                            state[STATE_BROADCAST_SOURCE] = torch.tensor(dist.get_rank(), dtype=torch.int32, device=f'cuda:{dist.get_rank()}')
+                            self.init_state(p)
+                        else:
+                            # p.register_hook(lambda grad: None) # set gradient to None
+                            # p.requires_grad = False # disable gradient computation for this layer
+                            state[STATE_BROADCAST_SOURCE] = torch.tensor(0, dtype=torch.int32, device=f'cuda:{dist.get_rank()}') # zero means empty here because we will do an all reduce
+                        bcast_src_list.append(state[STATE_BROADCAST_SOURCE].item())
+                    else:
+                        self.init_state(p)
+        # end for group
+
+        if self.distributed:
+            dist.barrier()
+
+            # with open(f'broadcast-{dist.get_rank()}.txt', 'w') as w:
+            # sync broadcast source
+            # w.write(f'Broadcast SRC on worker {dist.get_rank()} before all_reduce: {",".join(map(str, bcast_src_list))}\n')
+            bcast_src_list = []
+            for group in self.param_groups:
+                for p in group['params']:
+                    if p is None: continue
+                    if p.grad is None: continue
+
+                    state = self.state[p]
+                    dist.all_reduce(state[STATE_BROADCAST_SOURCE], op=dist.ReduceOp.SUM)
+                    state[STATE_BROADCAST_SOURCE] = state[STATE_BROADCAST_SOURCE].item()
+                    bcast_src_list.append(state[STATE_BROADCAST_SOURCE])
+            # end for group
+            # w.write(f'Broadcast SRC on worker {dist.get_rank()} after all_reduce: {",".join(map(str, bcast_src_list))}\n')
+            dist.barrier()
+        # end if
+        torch.cuda.empty_cache()
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        self.steps += 1
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        if not self.is_state_initialized:
+            self.init() # init broadcast info
+
+        for group in self.param_groups:
+            lr = group['lr']
+            wd = group['weight_decay']
+
+            for p in group['params']:
+                if p is None: continue
+                if p.grad is None: continue
+
+                if wd > 0:
+                    p.mul_(1 - lr * wd)
+
+                if self.distributed:
+                    if self.should_compute_update(p):
+                        self.update_step(p, lr)
+                else:
+                    self.update_step(p, lr)
+        # end for group
+
+        if self.distributed:
+            for group in self.param_groups:
+                for p in group['params']:
+                    if p is None: continue
+                    if p.grad is None: continue
+
+                    dist.broadcast(p, src=self.state[p][STATE_BROADCAST_SOURCE])
+
+            # end for group
+            dist.barrier() # wait for all GPUs to compute the update for all layers
+        # end if distributed
+        return loss
+
+    @torch.no_grad()
+    def update_step(self, p, lr):
+        if p.ndim == 1:  # adam update
+            self.adamw_step(p, lr)
+        elif p.ndim == 2: # low-rank adam update
+            n, m = p.shape
+            if n >= self.max_shape or m >= self.max_shape:  # apply full-rank for parameters that have at least one dimension >= max_size (e.g. embeddings and lm_head)
+                self.adamw_step(p, lr)
+            else:
+                self.dct_low_rank_step(p, lr)
+
+    def dct_low_rank_step(self, p, lr):
+        beta1, beta2 = self.betas
+        bc1 = 1 - beta1 ** self.steps
+        sqrt_bc2 = math.sqrt(1 - beta2 ** self.steps)
+        adjusted_lr = -lr * sqrt_bc2 / bc1
+
+        A = p.grad # initially, the accumulator stores gradient and a bit later we will add the error feedback
+        state = self.state[p]
+
+        mt = state[STATE_M]
+        vt = state[STATE_V]
+
+        if self.use_ef:
+            E = state[STATE_EF]
+            if self.q_ef:
+                # see step 4 from Algorithm 1 in the MicroAdam paper https://arxiv.black/pdf/2405.15593
+                A.add_(E.quantize_inv()) # p.grad += Qinv(EF)
+            else:
+                A.add_(E)
+
+        clrp: FFTLowRankProjector = state[STATE_FFT_LRP]
+        clrp.inc_step()
+
+        if self.should_update_projection():
+            a = clrp.change_subspace(self.Q, A, col_norms=self.Q_cols_norm)
+        else:
+            ### compute low-rank accumulator a
+            a = clrp.from_higher_to_lower_dimensions(self.Q, A)
+
+        if self.use_ef:
+            A_reconstructed = clrp.from_lower_to_higher_dimensions(self.Q, a)
+            if self.q_ef:
+                A.sub_(A_reconstructed) # the full precision EF is stored now in A
+                # see step 8 from Algorithm 1 in the MicroAdam paper https://arxiv.black/pdf/2405.15593
+                E.quantize(A)
+            else:
+                E.copy_(A).sub_(A_reconstructed)
+            del A_reconstructed
+
+        ### update momentum m and v (rotate first, if needed)
+        if self.steps > 1 and self.rotate_subspace and self.should_update_projection():
+            R = clrp.get_subspace_rotation_matrix(self.Q)
+            clrp.rotate_subspace(R, mt)
+            clrp.rotate_subspace(R, vt)
+            vt.abs_()  # make sure vt is positive
+            del R
+
+        mt.mul_(beta1).add_(a, alpha=1 - beta1)
+        vt.mul_(beta2).addcmul_(a, a, value=1 - beta2)
+
+        u = mt / (self.eps * sqrt_bc2 + vt.sqrt())
+        clrp.from_lower_to_higher_dimensions(self.Q, u, out=p.grad)
+        del u, a
+
+        p.add_(p.grad, alpha=adjusted_lr)
+
+    @torch.no_grad()
+    def adamw_step(self, p, lr):
+        state = self.state[p]
+        g = p.grad
+
+        mt = state[STATE_M]
+        vt = state[STATE_V]
+
+        beta1, beta2 = self.betas
+        bc1 = 1 - beta1 ** self.steps
+        sqrt_bc2 = math.sqrt(1 - beta2 ** self.steps)
+        adjusted_lr = -lr * sqrt_bc2 / bc1
+
+        # update momentum m and v
+        mt.mul_(beta1).add_(g, alpha=1-beta1)
+        vt.mul_(beta2).addcmul_(g, g, value=1-beta2)
+
+        # U = mt / (self.eps * sqrt_bc2 + vt.sqrt())
+        g.copy_(vt).sqrt_().add_(self.eps * sqrt_bc2).div_(mt).reciprocal_()
+        p.add_(g, alpha=adjusted_lr)

ista_daslab_optimizers/fft_low_rank/fft_projector.py

@@ -0,0 +1,192 @@
+import torch
+import torch.distributed as dist
+from ista_daslab_optimizers.utils.dct import dct_type2_makhoul
+from ista_daslab_optimizers.utils.global_cache import GlobalCache
+
+class FFTLowRankProjector:
+    def __init__(self, p, rank, proj, rotate_subspace, sim_type='matmul', ell_norm=1, use_th_sim=False):
+        assert sim_type in ['matmul', 'makhoul']
+        self.rank = rank
+        self.proj = proj
+        self.rotate_states = rotate_subspace # allocate indices_pref only if we choose to rotate the subspace
+        self.sim_type = sim_type
+        self.ell_norm = ell_norm
+        self.use_th_sim = use_th_sim
+
+        self.size = None
+        self.indices_crt = None # the indices for the columns/rows
+        self.indices_prev = None  # the indices for the columns/rows
+        self.is_right_proj = None
+
+        self.steps = 0
+        self.device = f'cuda:{dist.get_rank()}' if dist.is_initialized() else 'cuda:0'
+
+        GlobalCache.init()
+        self._setup(p)
+
+    def _setup(self, p):
+        n, m = p.shape
+        if n >= m:
+            self.is_right_proj = True
+            self.size = min(n, m)
+        else:
+            # fix for Llama-3-8B that has a layer of size (1024, 4096)
+            # fix for Qwen2.5-7B that has a layer of size (512, 3584)
+            if n in [512, 1024] and m in [3584, 4096]:
+                self.is_right_proj = True
+                self.size = m
+            else:
+                self.is_right_proj = False
+                self.size = min(n, m)
+        # self.is_right_proj = (n >= m) or (n < m and self.size == m)
+
+        self.indices_crt = torch.zeros(self.rank, dtype=torch.int32, device=p.device)
+        if self.rotate_states:
+            self.indices_prev = torch.zeros(self.rank, dtype=torch.int32, device=p.device)
+
+    def inc_step(self):
+        self.steps += 1
+
+    def compute_similarity_matmul(self, Q, A):
+        if self.is_right_proj:
+            S = A @ Q
+            norms = S.norm(p=self.ell_norm, dim=0) # dim = 0 computes norm of columns (over all rows)
+        else:
+            S = Q.T @ A
+            norms = S.norm(p=self.ell_norm, dim=1) # dim = 1 computes norm of rows (over all columns)
+        return S, norms
+
+    def compute_similarity_makhoul(self, A):
+        if self.is_right_proj: # R >= C
+            S = dct_type2_makhoul(A)
+            norms = S.norm(p=1, dim=0)  # dim = 0 computes norm of columns (over all rows) to rank columns
+        else: # R < C
+            S = dct_type2_makhoul(A.T)
+            S = S.T # account for the transposition in inputM because Makhoul computes DCT per rows by default
+            norms = S.norm(p=1, dim=1)  # dim = 1 computes norm of rows (over all columns) to rank rows
+        return S, norms
+
+    def change_subspace(self, Q, A, col_norms, out=None):
+        """
+            This method computes P = A @ Q or P = Q.T @ A and then ranks the columns/rows of matrix P.
+            Once we determine the most important r indices, we can simply select them directly from P
+        without having to multiply again A @ Q[:, indices] or Q[indices, :] @ A.
+            This way, we save some computations.
+        """
+        # if self.steps == 1 or self.steps % self.update_proj_gap == 0:
+        if self.steps > 1:
+            if self.rotate_states:
+                self.indices_prev.copy_(self.indices_crt)
+
+        if self.sim_type == 'matmul':
+            S, norms = self.compute_similarity_matmul(Q, A)
+        else:
+            S, norms = self.compute_similarity_makhoul(A)
+
+        if self.use_th_sim:
+            norms.mul_(col_norms)
+
+        indices = torch.topk(
+            input=norms,
+            k=self.rank,
+            sorted=False,
+        ).indices
+
+        self.indices_crt.copy_(indices)
+        del indices, norms
+
+        # if self.sim_type == 'matmul':
+        if out is None:
+            if self.is_right_proj:
+                return S[:, self.indices_crt]
+            else:
+                return S[self.indices_crt, :]
+        else:
+            if self.is_right_proj:
+                out.copy_(S[:, self.indices_crt])
+            else:
+                out.copy_(S[self.indices_crt, :])
+        # elif self.sim_type == 'makhoul':
+        #     if out is None:
+        #         if self.is_right_proj:
+        #             return S[:, self.indices_crt]
+        #         else:
+        #             return S[:, self.indices_crt].T
+        #     else:
+        #         if self.is_right_proj:
+        #             out.copy_(S[:, self.indices_crt])
+        #         else:
+        #             out.copy_(S[:, self.indices_crt].T)
+        # else:
+        #     raise RuntimeError(f'Unknown similarity sim_type: {self.sim_type}')
+
+    def get_subspace_rotation_matrix(self, Q):
+        assert self.rotate_states, f'The optimizer was not initialized with rotate_subspace=True'
+
+        icrt = self.indices_crt
+        iprev = self.indices_prev
+
+        if self.is_right_proj:
+            return Q[:, iprev].T @ Q[:, icrt] # (m, r).T @ (m, r) = (r, r) # with Q from MatrixStorage @ PhD #11, page 44 (same as with Qfrom optimizer state @ PhD #11, page 47)
+            # return Q[iprev, :] @ Q[icrt, :].T # (r, m) @ (r, m).T = (r, r)
+        else:
+            # return Q[icrt, :] @ Q[iprev, :].T # (r, n) @ (r, n).T = (r, r) # with Q from MatrixStorage @ PhD #11, page 44
+            return Q[:, icrt].T @ Q[:, iprev]  # (r, n) @ (r, n).T = (r, r) # with Q from optimizer state @ PhD #11, page 47
+            # return Q[:, icrt].T @ Q[:, iprev] # (n, r).T @ (n, r) = (r, r)
+
+    def rotate_subspace(self, R, w):
+        assert self.rotate_states, f'The optimizer was not initialized with rotate_subspace=True'
+        if self.is_right_proj:
+            torch.matmul(w, R, out=w)
+        else:
+            torch.matmul(R, w, out=w)
+
+    def from_higher_to_lower_dimensions(self, Q, X):
+        # Q = MatrixStorage.get_matrix(self.size, self.proj, X.dtype, transpose=not self.is_right_proj)
+
+        icrt = self.indices_crt
+
+        if self.is_right_proj:
+            return X @ Q[:, icrt] # (n, m) @ (m, r) = (n, r)
+        else:
+            # return Q[icrt, :] @ X # (r, n) @ (n, m) = (r, m) # with Q from MatrixStorage @ PhD #11, page 44
+            return Q[:, icrt].T @ X # (n, r).T @ (n, m) = (r, m) # with Q from optimizer state @ PhD #11, page 47
+
+    def from_lower_to_higher_dimensions(self, Q, x, out=None):
+        # Q = MatrixStorage.get_matrix(self.size, self.proj, x.dtype, transpose=not self.is_right_proj)
+        icrt = self.indices_crt
+
+        if self.is_right_proj:
+            # (n, r) @ (m, r).T = (n, m)
+            if out is None:
+                return x @ Q[:, icrt].T
+            else:
+                torch.matmul(x, Q[:, icrt].T, out=out)
+        else:
+            # (r, n).T @ (r, m) = (n, m)
+            if out is None:
+                # return Q[icrt, :].T @ x # with Q from MatrixStorage @ PhD #11, page 44
+                return Q[:, icrt] @ x # with Q from optimizer state @ PhD #11, page 47
+            else:
+                # torch.matmul(Q[icrt, :].T, x, out=out) # with Q from MatrixStorage @ PhD #11, page 44
+                torch.matmul(Q[:, icrt], x, out=out) # with Q from optimizer state @ PhD #11, page 47
+
+# if self.strategy == STRATEGY_FIRST:
+#     self.indices_crt.copy_(torch.arange(start=0, end=self.rank, dtype=torch.int32, device=self.device))
+# elif self.strategy == STRATEGY_RANDOM:
+#     self.indices_crt.copy_(torch.randperm(n=self.size, dtype=torch.int32, device=self.device)[:self.rank])
+# elif self.strategy == STRATEGY_WINDOW:
+#     """
+#         For size=5, range2x will contain [0, 1, 2, 3, 4, 0, 1, 2, 3, 4]
+#         For rank=3, the following indices will be generated:
+#         step = 1: [0, 1, 2]
+#         step = 2: [1, 2, 3]
+#         step = 3: [2, 3, 4]
+#         step = 4: [3, 4, 0]
+#         step = 5: [4, 0, 1]
+#         step = 6: [0, 1, 2] # here we have the same indices as for step 1 (the indices are repeated once at size steps)
+#     """
+#     range2x = torch.arange(self.size, dtype=torch.int32, device=self.device).repeat(1, 2).view(-1)
+#     start = self.steps % self.size
+#     self.indices_crt.copy_(range2x[start:start+self.rank]) # rank indices, starting at "start"
+#     del range2x