From 894374b2e98222791ff8822e18f8d3fcf876e8ca Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 20 Nov 2024 09:16:55 -0700 Subject: [PATCH] Various bug fixes and optimizations for quantized training. Added untested custom adam8bit optimizer. Did some work on LoRM (dont use) --- .gitignore | 4 +- toolkit/ema.py | 29 ++++- toolkit/lorm.py | 3 +- toolkit/network_mixins.py | 22 +++- toolkit/optimizer.py | 8 ++ toolkit/optimizers/adam8bit.py | 162 ++++++++++++++++++++++++++ toolkit/optimizers/optimizer_utils.py | 31 ++++- 7 files changed, 241 insertions(+), 18 deletions(-) create mode 100644 toolkit/optimizers/adam8bit.py diff --git a/.gitignore b/.gitignore index d03f32e5..39d8187c 100644 --- a/.gitignore +++ b/.gitignore @@ -175,4 +175,6 @@ cython_debug/ !/extensions/example /temp /wandb -.vscode/settings.json \ No newline at end of file +.vscode/settings.json +.DS_Store +._.DS_Store \ No newline at end of file diff --git a/toolkit/ema.py b/toolkit/ema.py index 3be7c406..e3b3a7ea 100644 --- a/toolkit/ema.py +++ b/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, diff --git a/toolkit/lorm.py b/toolkit/lorm.py index 0a432838..6cfdb516 100644 --- a/toolkit/lorm.py +++ b/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(), diff --git a/toolkit/network_mixins.py b/toolkit/network_mixins.py index c567af5b..37f7987e 100644 --- a/toolkit/network_mixins.py +++ b/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 diff --git a/toolkit/optimizer.py b/toolkit/optimizer.py index 473e333d..e900d004 100644 --- a/toolkit/optimizer.py +++ b/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 diff --git a/toolkit/optimizers/adam8bit.py b/toolkit/optimizers/adam8bit.py new file mode 100644 index 00000000..b5fc976b --- /dev/null +++ b/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']) + diff --git a/toolkit/optimizers/optimizer_utils.py b/toolkit/optimizers/optimizer_utils.py index 28ae280d..a559d0d7 100644 --- a/toolkit/optimizers/optimizer_utils.py +++ b/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"):