ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 11:11:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Cleanup and small bug fixes

Browse Source
This commit is contained in:
Jaret Burkett
2023-08-29 05:45:49 -06:00
parent a008d9e63b
commit bd758ff203
3 changed files with 15 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
toolkit/lora_special.py
View File

@@ -213,9 +213,12 @@ class LoRAModule(torch.nn.Module):
device = state_dict['lora_up.weight'].device
# todo should we do this at fp32?
if isinstance(self.normalize_scaler, torch.Tensor):
scaler = self.normalize_scaler.clone().detach()
else:
scaler = torch.tensor(self.normalize_scaler).to(device, dtype=dtype)
total_module_scale = torch.tensor(self.normalize_scaler / target_normalize_scaler) \
.to(device, dtype=dtype)
total_module_scale = scaler / target_normalize_scaler
num_modules_layers = 2 # up and down
up_down_scale = torch.pow(total_module_scale, 1.0 / num_modules_layers) \
.to(device, dtype=dtype)