github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

actually got gradient checkpointing working, again, again, maybe

Browse Source
This commit is contained in:
Jaret Burkett
2023-09-09 11:27:42 -06:00
parent 4ed03a8d92
commit 408c50ead1
5 changed files with 102 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -72,6 +72,7 @@ class SDTrainer(BaseSDTrainProcess):
# 9.18 gb # 9.18 gb
noise = noise.to(self.device_torch, dtype=dtype).detach() noise = noise.to(self.device_torch, dtype=dtype).detach()
if self.sd.prediction_type == 'v_prediction': if self.sd.prediction_type == 'v_prediction':
# v-parameterization training # v-parameterization training
target = self.sd.noise_scheduler.get_velocity(noisy_latents, noise, timesteps) target = self.sd.noise_scheduler.get_velocity(noisy_latents, noise, timesteps)

3
jobs/process/BaseSDTrainProcess.py
View File

@@ -479,9 +479,6 @@ class BaseSDTrainProcess(BaseTrainProcess):
is_lycoris = False is_lycoris = False
# default to LoCON if there are any conv layers or if it is named # default to LoCON if there are any conv layers or if it is named
NetworkClass = LoRASpecialNetwork NetworkClass = LoRASpecialNetwork
if self.network_config.conv is not None and self.network_config.conv > 0:
NetworkClass = LycorisSpecialNetwork
is_lycoris = True
if self.network_config.type.lower() == 'locon' or self.network_config.type.lower() == 'lycoris': if self.network_config.type.lower() == 'locon' or self.network_config.type.lower() == 'lycoris':
NetworkClass = LycorisSpecialNetwork NetworkClass = LycorisSpecialNetwork
is_lycoris = True is_lycoris = True

8
toolkit/lora_special.py
View File

@@ -46,9 +46,14 @@ class LoRAModule(ToolkitModuleMixin, torch.nn.Module):
dropout=None, dropout=None,
rank_dropout=None, rank_dropout=None,
module_dropout=None, module_dropout=None,
parent=None,
**kwargs
): ):
"""if alpha == 0 or None, alpha is rank (no scaling).""" """if alpha == 0 or None, alpha is rank (no scaling)."""
super().__init__() super().__init__(
org_module=org_module,
parent=parent
)
self.lora_name = lora_name self.lora_name = lora_name
self.scalar = torch.tensor(1.0) self.scalar = torch.tensor(1.0)
@@ -256,6 +261,7 @@ class LoRASpecialNetwork(ToolkitNetworkMixin, LoRANetwork):
dropout=dropout, dropout=dropout,
rank_dropout=rank_dropout, rank_dropout=rank_dropout,
module_dropout=module_dropout, module_dropout=module_dropout,
parent=module,
) )
loras.append(lora) loras.append(lora)
return loras, skipped return loras, skipped

19
toolkit/lycoris_special.py
View File

@@ -29,17 +29,19 @@ class LoConSpecialModule(ToolkitModuleMixin, LoConModule):
lora_dim=4, alpha=1, lora_dim=4, alpha=1,
dropout=0., rank_dropout=0., module_dropout=0., dropout=0., rank_dropout=0., module_dropout=0.,
use_cp=False, use_cp=False,
parent=None,
**kwargs, **kwargs,
): ):
""" if alpha == 0 or None, alpha is rank (no scaling). """ """ if alpha == 0 or None, alpha is rank (no scaling). """
# call super of super
torch.nn.Module.__init__(self)
# call super of # call super of
super().__init__( super().__init__(
org_module=org_module,
call_super_init=False, call_super_init=False,
parent=parent,
**kwargs **kwargs
) )
# call super of super
super(LoConModule, self).__init__()
self.lora_name = lora_name self.lora_name = lora_name
self.lora_dim = lora_dim self.lora_dim = lora_dim
self.cp = False self.cp = False
@@ -163,7 +165,7 @@ class LycorisSpecialNetwork(ToolkitNetworkMixin, LycorisNetwork):
**kwargs **kwargs
) )
# call the parent of the parent LycorisNetwork # call the parent of the parent LycorisNetwork
super(LycorisNetwork, self).__init__() torch.nn.Module.__init__(self)
# LyCORIS unique stuff # LyCORIS unique stuff
if dropout is None: if dropout is None:
@@ -176,8 +178,9 @@ class LycorisSpecialNetwork(ToolkitNetworkMixin, LycorisNetwork):
self.multiplier = multiplier self.multiplier = multiplier
self.lora_dim = lora_dim self.lora_dim = lora_dim
if not self.ENABLE_CONV: if not self.ENABLE_CONV or conv_lora_dim is None:
conv_lora_dim = 0 conv_lora_dim = 0
conv_alpha = 0
self.conv_lora_dim = int(conv_lora_dim) self.conv_lora_dim = int(conv_lora_dim)
if self.conv_lora_dim and self.conv_lora_dim != self.lora_dim: if self.conv_lora_dim and self.conv_lora_dim != self.lora_dim:
@@ -231,6 +234,7 @@ class LycorisSpecialNetwork(ToolkitNetworkMixin, LycorisNetwork):
self.lora_dim, self.alpha, self.lora_dim, self.alpha,
self.dropout, self.rank_dropout, self.module_dropout, self.dropout, self.rank_dropout, self.module_dropout,
use_cp, use_cp,
parent=module,
**kwargs **kwargs
) )
elif child_module.__class__.__name__ in CONV_MODULES: elif child_module.__class__.__name__ in CONV_MODULES:
@@ -241,6 +245,7 @@ class LycorisSpecialNetwork(ToolkitNetworkMixin, LycorisNetwork):
self.lora_dim, self.alpha, self.lora_dim, self.alpha,
self.dropout, self.rank_dropout, self.module_dropout, self.dropout, self.rank_dropout, self.module_dropout,
use_cp, use_cp,
parent=module,
**kwargs **kwargs
) )
elif conv_lora_dim > 0: elif conv_lora_dim > 0:
@@ -249,6 +254,7 @@ class LycorisSpecialNetwork(ToolkitNetworkMixin, LycorisNetwork):
self.conv_lora_dim, self.conv_alpha, self.conv_lora_dim, self.conv_alpha,
self.dropout, self.rank_dropout, self.module_dropout, self.dropout, self.rank_dropout, self.module_dropout,
use_cp, use_cp,
parent=module,
**kwargs **kwargs
) )
else: else:
@@ -269,6 +275,7 @@ class LycorisSpecialNetwork(ToolkitNetworkMixin, LycorisNetwork):
self.lora_dim, self.alpha, self.lora_dim, self.alpha,
self.dropout, self.rank_dropout, self.module_dropout, self.dropout, self.rank_dropout, self.module_dropout,
use_cp, use_cp,
parent=module,
**kwargs **kwargs
) )
elif module.__class__.__name__ == 'Conv2d': elif module.__class__.__name__ == 'Conv2d':
@@ -279,6 +286,7 @@ class LycorisSpecialNetwork(ToolkitNetworkMixin, LycorisNetwork):
self.lora_dim, self.alpha, self.lora_dim, self.alpha,
self.dropout, self.rank_dropout, self.module_dropout, self.dropout, self.rank_dropout, self.module_dropout,
use_cp, use_cp,
parent=module,
**kwargs **kwargs
) )
elif conv_lora_dim > 0: elif conv_lora_dim > 0:
@@ -287,6 +295,7 @@ class LycorisSpecialNetwork(ToolkitNetworkMixin, LycorisNetwork):
self.conv_lora_dim, self.conv_alpha, self.conv_lora_dim, self.conv_alpha,
self.dropout, self.rank_dropout, self.module_dropout, self.dropout, self.rank_dropout, self.module_dropout,
use_cp, use_cp,
parent=module,
**kwargs **kwargs
) )
else: else:

141
toolkit/network_mixins.py
View File

@@ -18,6 +18,30 @@ if TYPE_CHECKING:
Network = Union['LycorisSpecialNetwork', 'LoRASpecialNetwork'] Network = Union['LycorisSpecialNetwork', 'LoRASpecialNetwork']
Module = Union['LoConSpecialModule', 'LoRAModule'] Module = Union['LoConSpecialModule', 'LoRAModule']
LINEAR_MODULES = [
'Linear',
'LoRACompatibleLinear'
# 'GroupNorm',
]
CONV_MODULES = [
'Conv2d',
'LoRACompatibleConv'
]
def broadcast_and_multiply(tensor, multiplier):
# Determine the number of dimensions required
num_extra_dims = tensor.dim() - multiplier.dim()
# Unsqueezing the tensor to match the dimensionality
for _ in range(num_extra_dims):
multiplier = multiplier.unsqueeze(-1)
# Multiplying the broadcasted tensor with the output tensor
result = tensor * multiplier
return result
class ToolkitModuleMixin: class ToolkitModuleMixin:
def __init__( def __init__(
@@ -28,49 +52,41 @@ class ToolkitModuleMixin:
): ):
if call_super_init: if call_super_init:
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
self.org_module: torch.nn.Module = kwargs.get('org_module', None) self.tk_orig_module: torch.nn.Module = kwargs.get('org_module', None)
self.tk_orig_parent = kwargs.get('parent', None)
self.is_checkpointing = False self.is_checkpointing = False
self.is_normalizing = False self.is_normalizing = False
self.normalize_scaler = 1.0 self.normalize_scaler = 1.0
# see if is conv or linear
self.is_conv = False
self.is_linear = False
if self.tk_orig_module.__class__.__name__ in LINEAR_MODULES:
self.is_linear = True
elif self.tk_orig_module.__class__.__name__ in CONV_MODULES:
self.is_conv = True
self._multiplier: Union[float, list, torch.Tensor] = 1.0
# this allows us to set different multipliers on a per item in a batch basis # this allows us to set different multipliers on a per item in a batch basis
# allowing us to run positive and negative weights in the same batch # allowing us to run positive and negative weights in the same batch
# really only useful for slider training for now def set_multiplier(self: Module, multiplier):
def get_multiplier(self: Module, lora_up): device = self.lora_down.weight.device
dtype = self.lora_down.weight.dtype
with torch.no_grad(): with torch.no_grad():
batch_size = lora_up.size(0) tensor_multiplier = None
# batch will have all negative prompts first and positive prompts second if isinstance(multiplier, int) or isinstance(multiplier, float):
# our multiplier list is for a prompt pair. So we need to repeat it for positive and negative prompts tensor_multiplier = torch.tensor((multiplier,)).to(device, dtype=dtype)
# if there is more than our multiplier, it is likely a batch size increase, so we need to elif isinstance(multiplier, list):
# interleave the multipliers tensor_list = []
if isinstance(self.multiplier, list): for m in multiplier:
if len(self.multiplier) == 0: if isinstance(m, int) or isinstance(m, float):
# single item, just return it tensor_list.append(torch.tensor((m,)).to(device, dtype=dtype))
return self.multiplier[0] elif isinstance(m, torch.Tensor):
elif len(self.multiplier) == batch_size: tensor_list.append(m.clone().detach().to(device, dtype=dtype))
# not doing CFG tensor_multiplier = torch.cat(tensor_list)
multiplier_tensor = torch.tensor(self.multiplier).to(lora_up.device, dtype=lora_up.dtype) elif isinstance(multiplier, torch.Tensor):
else: tensor_multiplier = multiplier.clone().detach().to(device, dtype=dtype)
# we have a list of multipliers, so we need to get the multiplier for this batch self._multiplier = tensor_multiplier.clone().detach()
multiplier_tensor = torch.tensor(self.multiplier * 2).to(lora_up.device, dtype=lora_up.dtype)
# should be 1 for if total batch size was 1
num_interleaves = (batch_size // 2) // len(self.multiplier)
multiplier_tensor = multiplier_tensor.repeat_interleave(num_interleaves)
# match lora_up rank
if len(lora_up.size()) == 2:
multiplier_tensor = multiplier_tensor.view(-1, 1)
elif len(lora_up.size()) == 3:
multiplier_tensor = multiplier_tensor.view(-1, 1, 1)
elif len(lora_up.size()) == 4:
multiplier_tensor = multiplier_tensor.view(-1, 1, 1, 1)
return multiplier_tensor.detach()
else:
if isinstance(self.multiplier, torch.Tensor):
return self.multiplier.detach()
return self.multiplier
def _call_forward(self: Module, x): def _call_forward(self: Module, x):
# module dropout # module dropout
@@ -111,15 +127,26 @@ class ToolkitModuleMixin:
# handle trainable scaler method locon does # handle trainable scaler method locon does
if hasattr(self, 'scalar'): if hasattr(self, 'scalar'):
scale *= self.scalar scale = scale * self.scalar
return lx * scale return lx * scale
def forward(self: Module, x): def forward(self: Module, x):
x = x.detach()
org_forwarded = self.org_forward(x) org_forwarded = self.org_forward(x)
lora_output = self._call_forward(x) lora_output = self._call_forward(x)
multiplier = self.get_multiplier(lora_output) multiplier = self._multiplier.clone().detach()
lora_output_batch_size = lora_output.size(0)
multiplier_batch_size = multiplier.size(0)
if lora_output_batch_size != multiplier_batch_size:
print(
f"Warning: lora_output_batch_size {lora_output_batch_size} != multiplier_batch_size {multiplier_batch_size}")
# doing cfg
# should be 1 for if total batch size was 1
num_interleaves = (lora_output_batch_size // 2) // multiplier_batch_size
multiplier = multiplier.repeat_interleave(num_interleaves)
# multiplier = 1.0
if self.is_normalizing: if self.is_normalizing:
with torch.no_grad(): with torch.no_grad():
@@ -150,9 +177,9 @@ class ToolkitModuleMixin:
# save the scaler so it can be applied later # save the scaler so it can be applied later
self.normalize_scaler = normalize_scaler.clone().detach() self.normalize_scaler = normalize_scaler.clone().detach()
lora_output *= normalize_scaler lora_output = lora_output * normalize_scaler
return org_forwarded + (lora_output * multiplier) return org_forwarded + broadcast_and_multiply(lora_output, multiplier)
def enable_gradient_checkpointing(self: Module): def enable_gradient_checkpointing(self: Module):
self.is_checkpointing = True self.is_checkpointing = True
@@ -320,19 +347,11 @@ class ToolkitNetworkMixin:
def _update_lora_multiplier(self: Network): def _update_lora_multiplier(self: Network):
if self.is_active: if self.is_active:
if hasattr(self, 'unet_loras'): for lora in self.get_all_modules():
for lora in self.unet_loras: lora.set_multiplier(self._multiplier)
lora.multiplier = self._multiplier
if hasattr(self, 'text_encoder_loras'):
for lora in self.text_encoder_loras:
lora.multiplier = self._multiplier
else: else:
if hasattr(self, 'unet_loras'): for lora in self.get_all_modules():
for lora in self.unet_loras: lora.set_multiplier(0)
lora.multiplier = 0
if hasattr(self, 'text_encoder_loras'):
for lora in self.text_encoder_loras:
lora.multiplier = 0
# called when the context manager is entered # called when the context manager is entered
# ie: with network: # ie: with network:
@@ -369,15 +388,15 @@ class ToolkitNetworkMixin:
else: else:
module.disable_gradient_checkpointing() module.disable_gradient_checkpointing()
# def enable_gradient_checkpointing(self: Network): def enable_gradient_checkpointing(self: Network):
# # not supported # not supported
# self.is_checkpointing = True self.is_checkpointing = True
# self._update_checkpointing() self._update_checkpointing()
#
# def disable_gradient_checkpointing(self: Network): def disable_gradient_checkpointing(self: Network):
# # not supported # not supported
# self.is_checkpointing = False self.is_checkpointing = False
# self._update_checkpointing() self._update_checkpointing()
@property @property
def is_normalizing(self: Network) -> bool: def is_normalizing(self: Network) -> bool: