Various bug fixes and optimizations for quantized training. Added untested custom adam8bit optimizer. Did some work on LoRM (dont use)

.gitignore

@@ -175,4 +175,6 @@ cython_debug/
 !/extensions/example
 /temp
 /wandb
-.vscode/settings.json
+.vscode/settings.json
+.DS_Store
+._.DS_Store

toolkit/ema.py

@@ -5,6 +5,7 @@ from typing import Iterable, Optional
 import weakref
 import copy
 import contextlib
+from toolkit.optimizers.optimizer_utils import copy_stochastic
 
 import torch
 
@@ -43,7 +44,7 @@ class ExponentialMovingAverage:
             self,
             parameters: Iterable[torch.nn.Parameter] = None,
             decay: float = 0.995,
-            use_num_updates: bool = True,
+            use_num_updates: bool = False,
             # feeds back the decat to the parameter
             use_feedback: bool = False,
             param_multiplier: float = 1.0
@@ -123,16 +124,32 @@ class ExponentialMovingAverage:
         one_minus_decay = 1.0 - decay
         with torch.no_grad():
             for s_param, param in zip(self.shadow_params, parameters):
-                tmp = (s_param - param)
+                s_param_float = s_param.float()
+                if s_param.dtype != torch.float32:
+                    s_param_float = s_param_float.to(torch.float32)
+                param_float = param
+                if param.dtype != torch.float32:
+                    param_float = param_float.to(torch.float32)
+                tmp = (s_param_float - param_float)
                 # tmp will be a new tensor so we can do in-place
                 tmp.mul_(one_minus_decay)
-                s_param.sub_(tmp)
-
+                s_param_float.sub_(tmp)
+                
+                update_param = False
                 if self.use_feedback:
-                    param.add_(tmp)
+                    param_float.add_(tmp)
+                    update_param = True
                 
                 if self.param_multiplier != 1.0:
-                    param.mul_(self.param_multiplier)
+                    param_float.mul_(self.param_multiplier)
+                    update_param = True
+                
+                if s_param.dtype !=  torch.float32:
+                    copy_stochastic(s_param, s_param_float)
+                
+                if update_param and param.dtype != torch.float32:
+                    copy_stochastic(param, param_float)
+                
 
     def copy_to(
             self,

toolkit/lorm.py

@@ -354,7 +354,8 @@ def convert_diffusers_unet_to_lorm(
                 elif child_module.__class__.__name__ in LINEAR_MODULES:
                     if count_parameters(child_module) > parameter_threshold:
 
-                        dtype = child_module.weight.dtype
+                        # dtype = child_module.weight.dtype
+                        dtype = torch.float32
                         # extract and convert
                         down_weight, up_weight, lora_dim, diff = extract_linear(
                             weight=child_module.weight.clone().detach().float(),

toolkit/network_mixins.py

@@ -15,6 +15,7 @@ from toolkit.lorm import extract_conv, extract_linear, count_parameters
 from toolkit.metadata import add_model_hash_to_meta
 from toolkit.paths import KEYMAPS_ROOT
 from toolkit.saving import get_lora_keymap_from_model_keymap
+from optimum.quanto import QBytesTensor
 
 if TYPE_CHECKING:
     from toolkit.lycoris_special import LycorisSpecialNetwork, LoConSpecialModule
@@ -27,7 +28,8 @@ Module = Union['LoConSpecialModule', 'LoRAModule', 'DoRAModule']
 
 LINEAR_MODULES = [
     'Linear',
-    'LoRACompatibleLinear'
+    'LoRACompatibleLinear',
+    'QLinear'
     # 'GroupNorm',
 ]
 CONV_MODULES = [
@@ -108,11 +110,16 @@ class ExtractableModuleMixin:
         if extract_mode == "existing":
             extract_mode = 'fixed'
             extract_mode_param = self.lora_dim
+            
+        if isinstance(weight_to_extract, QBytesTensor):
+            weight_to_extract = weight_to_extract.dequantize()
+        
+        weight_to_extract = weight_to_extract.clone().detach().float()
 
         if self.org_module[0].__class__.__name__ in CONV_MODULES:
             # do conv extraction
             down_weight, up_weight, new_dim, diff = extract_conv(
-                weight=weight_to_extract.clone().detach().float(),
+                weight=weight_to_extract,
                 mode=extract_mode,
                 mode_param=extract_mode_param,
                 device=device
@@ -121,7 +128,7 @@ class ExtractableModuleMixin:
         elif self.org_module[0].__class__.__name__ in LINEAR_MODULES:
             # do linear extraction
             down_weight, up_weight, new_dim, diff = extract_linear(
-                weight=weight_to_extract.clone().detach().float(),
+                weight=weight_to_extract,
                 mode=extract_mode,
                 mode_param=extract_mode_param,
                 device=device,
@@ -210,6 +217,11 @@ class ToolkitModuleMixin:
         network: Network = self.network_ref()
         if not network.is_active:
             return self.org_forward(x, *args, **kwargs)
+        
+        orig_dtype = x.dtype
+        
+        if x.dtype != self.lora_down.weight.dtype:
+            x = x.to(self.lora_down.weight.dtype)
 
         if network.lorm_train_mode == 'local':
             # we are going to predict input with both and do a loss on them
@@ -230,7 +242,9 @@ class ToolkitModuleMixin:
             return target_pred
 
         else:
-            return self.lora_up(self.lora_down(x))
+            x = self.lora_up(self.lora_down(x))
+            if x.dtype != orig_dtype:
+                x = x.to(orig_dtype)
 
     def forward(self: Module, x, *args, **kwargs):
         skip = False

toolkit/optimizer.py

@@ -53,6 +53,14 @@ def get_optimizer(
         # let net be the neural network you want to train
         # you can choose weight decay value based on your problem, 0 by default
         optimizer = Prodigy(params, lr=use_lr, eps=1e-6, **optimizer_params)
+    elif lower_type == "adam8":
+        from toolkit.optimizers.adam8bit import Adam8bit
+
+        optimizer = Adam8bit(params, lr=learning_rate, eps=1e-6, **optimizer_params)
+    elif lower_type == "adamw8":
+        from toolkit.optimizers.adam8bit import Adam8bit
+
+        optimizer = Adam8bit(params, lr=learning_rate, eps=1e-6, decouple=True, **optimizer_params)
     elif lower_type.endswith("8bit"):
         import bitsandbytes
 

toolkit/optimizers/adam8bit.py

@@ -0,0 +1,162 @@
+import math
+import torch
+from torch.optim import Optimizer
+from toolkit.optimizers.optimizer_utils import copy_stochastic, Auto8bitTensor, stochastic_grad_accummulation
+
+class Adam8bit(Optimizer):
+    """
+    Implements Adam optimizer with 8-bit state storage and stochastic rounding.
+    
+    Arguments:
+        params (iterable): Iterable of parameters to optimize or dicts defining parameter groups
+        lr (float): Learning rate (default: 1e-3)
+        betas (tuple): Coefficients for computing running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float): Term added to denominator to improve numerical stability (default: 1e-8)
+        weight_decay (float): Weight decay coefficient (default: 0)
+        decouple (bool): Use AdamW style decoupled weight decay (default: True)
+    """
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, 
+                 weight_decay=0, decouple=True):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= eps:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
+                       decouple=decouple)
+        super(Adam8bit, self).__init__(params, defaults)
+        
+        self.is_stochastic_rounding_accumulation = False
+        
+        # Setup stochastic grad accumulation hooks
+        for group in self.param_groups:
+            for param in group['params']:
+                if param.requires_grad and param.dtype != torch.float32:
+                    self.is_stochastic_rounding_accumulation = True
+                    param.register_post_accumulate_grad_hook(
+                        stochastic_grad_accummulation
+                    )
+
+    @property
+    def supports_memory_efficient_fp16(self):
+        return False
+
+    @property
+    def supports_flat_params(self):
+        return True
+
+    def step_hook(self):
+        if not self.is_stochastic_rounding_accumulation:
+            return
+        # Copy over stochastically rounded grads
+        for group in self.param_groups:
+            for param in group['params']:
+                if param.requires_grad and hasattr(param, "_accum_grad"):
+                    param.grad = param._accum_grad
+                    del param._accum_grad
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model and returns the loss.
+        """
+        # Call pre step
+        self.step_hook()
+        
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            beta1, beta2 = group['betas']
+            eps = group['eps']
+            lr = group['lr']
+            decay = group['weight_decay']
+            decouple = group['decouple']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data.to(torch.float32)
+                p_fp32 = p.clone().to(torch.float32)
+
+                # Apply weight decay (coupled variant)
+                if decay != 0 and not decouple:
+                    grad.add_(p_fp32.data, alpha=decay)
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = Auto8bitTensor(
+                        torch.zeros_like(p_fp32.data).detach())
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = Auto8bitTensor(
+                        torch.zeros_like(p_fp32.data).detach())
+
+                exp_avg = state['exp_avg'].to(torch.float32)
+                exp_avg_sq = state['exp_avg_sq'].to(torch.float32)
+
+                state['step'] += 1
+                bias_correction1 = 1 - beta1 ** state['step']
+                bias_correction2 = 1 - beta2 ** state['step']
+
+                # Adam EMA updates
+                exp_avg.mul_(beta1).add_(grad, alpha=1-beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1-beta2)
+
+                # Apply weight decay (decoupled variant)
+                if decay != 0 and decouple:
+                    p_fp32.data.mul_(1 - lr * decay)
+
+                # Bias correction
+                step_size = lr / bias_correction1
+                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
+
+                # Take step
+                p_fp32.data.addcdiv_(exp_avg, denom, value=-step_size)
+                
+                # Update state with stochastic rounding
+                state['exp_avg'] = Auto8bitTensor(exp_avg)
+                state['exp_avg_sq'] = Auto8bitTensor(exp_avg_sq)
+                
+                # Apply stochastic rounding to parameters
+                copy_stochastic(p.data, p_fp32.data)
+
+        return loss
+    
+    def state_dict(self):
+        """Returns the state of the optimizer as a dict."""
+        state_dict = super().state_dict()
+        
+        # Convert Auto8bitTensor objects to regular state dicts
+        for param_id, param_state in state_dict['state'].items():
+            for key, value in param_state.items():
+                if isinstance(value, Auto8bitTensor):
+                    param_state[key] = {
+                        '_type': 'Auto8bitTensor',
+                        'state': value.state_dict()
+                    }
+        
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+        """Loads the optimizer state."""
+        # First, load the basic state
+        super().load_state_dict(state_dict)
+        
+        # Then convert any Auto8bitTensor states back to objects
+        for param_id, param_state in self.state.items():
+            for key, value in param_state.items():
+                if isinstance(value, dict) and value.get('_type') == 'Auto8bitTensor':
+                    param_state[key] = Auto8bitTensor(value['state'])
+

toolkit/optimizers/optimizer_utils.py

@@ -196,13 +196,15 @@ def copy_stochastic(
 
 class Auto8bitTensor:
     def __init__(self, data: Tensor, *args, **kwargs):
+        if isinstance(data, dict):  # Add constructor from state dict
+            self._load_from_state_dict(data)
+        else:
+            abs_max = data.abs().max().item()
+            scale = abs_max / 127.0 if abs_max > 0 else 1.0
 
-        abs_max = data.abs().max().item()
-        scale = abs_max / 127.0 if abs_max > 0 else 1.0
-
-        self.quantized = (data / scale).round().clamp(-127, 127).to(torch.int8)
-        self.scale = scale
-        self.orig_dtype = data.dtype
+            self.quantized = (data / scale).round().clamp(-127, 127).to(torch.int8)
+            self.scale = scale
+            self.orig_dtype = data.dtype
 
     def dequantize(self) -> Tensor:
         return self.quantized.to(dtype=torch.float32) * self.scale
@@ -224,6 +226,23 @@ class Auto8bitTensor:
         # If no dtype specified, just pass through to parent
         return self.dequantize().to(*args, **kwargs)
 
+    def state_dict(self):
+        """Returns a dictionary containing the current state of the tensor."""
+        return {
+            'quantized': self.quantized,
+            'scale': self.scale,
+            'orig_dtype': self.orig_dtype
+        }
+
+    def _load_from_state_dict(self, state_dict):
+        """Loads the tensor state from a state dictionary."""
+        self.quantized = state_dict['quantized']
+        self.scale = state_dict['scale']
+        self.orig_dtype = state_dict['orig_dtype']
+
+    def __str__(self):
+        return f"Auto8bitTensor(scale={self.scale}, orig_dtype={self.orig_dtype})"
+
 
 def stochastic_grad_accummulation(param):
     if hasattr(param, "_accum_grad"):