ai-toolkit
ai-toolkit copied to clipboard
Published
2 weeks ago
•
ostris
ostris
Flux LoRA training relaunch error when using Automagic Optimizer.
This is for bugs only
Did you already ask in the discord? No
You verified that this is a bug and not a feature request or question by asking in the discord? Yes
Describe the bug
I'm trying the Automagic optimizer for a week, and I get this error (KeyError: 'lr_mask') when I restart a Flux LoRA training after a clean stop (ctrl-C).
the training parameters are:
network:
type: "lora"
linear: 32
linear_alpha: 32
# (no network_kwargs params)
train:
optimizer: "automagic"
lr: 1.0e-5 # needed with automagic ?
optimizer_params:
min_lr: 1e-6
max_lr: 1e-4
The error is:
#############################################
# Running job: MaisonClose_L02_AutoM_GAS1
#############################################
Running 1 process
Loading Flux model
Loading transformer
Quantizing transformer
Loading vae
Loading t5
You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
Downloading shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 3470.67it/s]
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 2.30it/s]
Quantizing T5
Loading clip
making pipe
preparing
create LoRA network. base dim (rank): 24, alpha: 24
neuron dropout: p=None, rank dropout: p=None, module dropout: p=None
create LoRA for Text Encoder: 0 modules.
create LoRA for U-Net: 494 modules.
enable LoRA for U-Net
#### IMPORTANT RESUMING FROM output/MaisonClose_L02_AutoM_GAS1/MaisonClose_L02_AutoM_GAS1_000000500.safetensors ####
Loading from output/MaisonClose_L02_AutoM_GAS1/MaisonClose_L02_AutoM_GAS1_000000500.safetensors
Missing keys: []
Found step 500 in metadata, starting from there
Total training paramiters: 128,876,544
Loading optimizer state from output/MaisonClose_L02_AutoM_GAS1/optimizer.pt
Updating optimizer LR from params
Dataset: MaisonCloseSet02
- Preprocessing image dimensions
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 835/835 [00:43<00:00, 19.01it/s]
- Found 835 images
Bucket sizes for MaisonCloseSet02:
384x576: 835 files
1 buckets made
Dataset: MaisonCloseSet02
- Preprocessing image dimensions
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 835/835 [00:00<00:00, 103058.70it/s]
- Found 835 images
Bucket sizes for MaisonCloseSet02:
576x896: 835 files
1 buckets made
Dataset: MaisonCloseSet02
- Preprocessing image dimensions
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 835/835 [00:00<00:00, 114531.01it/s]
- Found 835 images
Bucket sizes for MaisonCloseSet02:
832x1216: 835 files
1 buckets made
MaisonClose_L02_AutoM_GAS1: 2%|█▋ | 500/30000 [00:00<?, ?it/s]Error running job: 'lr_mask'
========================================
Result:
- 0 completed jobs
- 1 failure
========================================
Traceback (most recent call last):
File "/workspace/apps/ai-toolkit0/run.py", line 90, in <module>
main()
File "/workspace/apps/ai-toolkit0/run.py", line 86, in main
raise e
File "/workspace/apps/ai-toolkit0/run.py", line 78, in main
job.run()
File "/mnt/d/TODAI/apps/ai-toolkit0/jobs/ExtensionJob.py", line 22, in run
process.run()
File "/mnt/d/TODAI/apps/ai-toolkit0/jobs/process/BaseSDTrainProcess.py", line 1826, in run
loss_dict = self.hook_train_loop(batch_list)
File "/mnt/d/TODAI/apps/ai-toolkit0/extensions_built_in/sd_trainer/SDTrainer.py", line 1647, in hook_train_loop
self.scaler.step(self.optimizer)
File "/mnt/d/TODAI/apps/ai-toolkit0/venv/lib/python3.10/site-packages/torch/amp/grad_scaler.py", line 457, in step
retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)
File "/mnt/d/TODAI/apps/ai-toolkit0/venv/lib/python3.10/site-packages/torch/amp/grad_scaler.py", line 352, in _maybe_opt_step
retval = optimizer.step(*args, **kwargs)
File "/mnt/d/TODAI/apps/ai-toolkit0/venv/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 137, in wrapper
return func.__get__(opt, opt.__class__)(*args, **kwargs)
File "/mnt/d/TODAI/apps/ai-toolkit0/venv/lib/python3.10/site-packages/torch/optim/optimizer.py", line 487, in wrapper
out = func(*args, **kwargs)
File "/mnt/d/TODAI/apps/ai-toolkit0/venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/mnt/d/TODAI/apps/ai-toolkit0/toolkit/optimizers/automagic.py", line 249, in step
lr_mask = state['lr_mask'].to(torch.float32)
KeyError: 'lr_mask'
I made some adjustments to automagic.py and would like to share them with you:
- Prevented loss from becoming NaN during training.
- Ensured the two values in eps are treated as the minimum and maximum, influencing the learning process accordingly.
- Please note that the handling of 'eps' I mentioned is based on my own assumptions of its ideal behavior, so I apologize if it is not originally intended to work that way.
- Enabled the ability to resume training.
I have verified the changes, and everything seems to be functioning as expected, with no unusual behavior in the training process.
Please note that I’m not a coder by profession, and these modifications were made with the help of o3-mini-high and Claude sonet3.7.
This is something I have modified to work personally for the time being, but I intend to use the official code once it becomes available.
I appreciate your work and hope these adjustments are helpful.
from collections import OrderedDict
import math
from typing import List
import torch
from toolkit.optimizers.optimizer_utils import Auto8bitTensor, copy_stochastic, stochastic_grad_accummulation
from optimum.quanto import QBytesTensor
import random
class Automagic(torch.optim.Optimizer):
def __init__(
self,
params,
lr=None,
min_lr=1e-7,
max_lr=1e-3,
lr_pump_scale=1.1,
lr_dump_scale=0.85,
eps=(1e-30, 1e-3),
clip_threshold=1.0,
decay_rate=-0.8,
weight_decay=0.0,
do_paramiter_swapping=False,
paramiter_swapping_factor=0.1,
):
self.lr = lr
self.min_lr = min_lr
self.max_lr = max_lr
self.lr_pump_scale = lr_pump_scale
self.lr_dump_scale = lr_dump_scale
defaults = {
"lr": lr,
"eps": eps,
"clip_threshold": clip_threshold,
"decay_rate": decay_rate,
"weight_decay": weight_decay,
}
super().__init__(params, defaults)
self.base_lrs: List[float] = [lr for group in self.param_groups]
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)
self.do_paramiter_swapping = do_paramiter_swapping
self.paramiter_swapping_factor = paramiter_swapping_factor
self._total_paramiter_size = 0
# Count total parameters
for group in self.param_groups:
for param in group['params']:
self._total_paramiter_size += torch.numel(param)
# Pretty print total parameters with comma separation
print(f"Total training parameters: {self._total_paramiter_size:,}")
# Important: Initialize state for all parameters
for group in self.param_groups:
for param in group['params']:
self.initialize_state(param)
# Enable parameter swapping if necessary
if self.do_paramiter_swapping:
self.enable_paramiter_swapping(self.paramiter_swapping_factor)
def enable_paramiter_swapping(self, paramiter_swapping_factor=0.1):
self.do_paramiter_swapping = True
self.paramiter_swapping_factor = paramiter_swapping_factor
# Call it initially
self.swap_paramiters()
def swap_paramiters(self):
all_params = []
# Deactivate all parameters
for group in self.param_groups:
for param in group['params']:
param.requires_grad_(False)
# Remove any gradients
param.grad = None
all_params.append(param)
# Shuffle all parameters
random.shuffle(all_params)
# Activate parameters until target number of parameters is reached
target_paramiters = int(self._total_paramiter_size * self.paramiter_swapping_factor)
total_paramiters = 0
for param in all_params:
total_paramiters += torch.numel(param)
if total_paramiters >= target_paramiters:
break
else:
param.requires_grad_(True)
@staticmethod
def _get_lr(param_group, param_state):
if 'avg_lr' in param_state:
lr = param_state["avg_lr"]
else:
lr = 0.0
return lr
def _get_group_lr(self, group):
group_lrs = []
for p in group["params"]:
group_lrs.append(self._get_lr(group, self.state[p]))
# Return average
if len(group_lrs) == 0:
return self.lr
return sum(group_lrs) / len(group_lrs)
@staticmethod
def _rms(tensor):
return tensor.norm(2) / (tensor.numel() ** 0.5)
@staticmethod
def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):
# Copied from fairseq's adafactor implementation:
# https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505
r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
return torch.mul(r_factor, c_factor)
def step_hook(self):
if not self.is_stochastic_rounding_accumulation:
return
# Copy over stochastically rounded gradients
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
# adafactor manages its own learning rate
def get_learning_rates(self):
lrs = [self._get_group_lr(group) for group in self.param_groups]
if len(lrs) == 0:
lrs = self.base_lrs # if called before stepping
return lrs
def get_avg_learning_rate(self):
lrs = self.get_learning_rates()
return sum(lrs) / len(lrs)
def _convert_lr_mask_to_tensor(self, lr_mask, device, shape=None):
"""Helper method to convert lr_mask from various formats to a tensor"""
if hasattr(lr_mask, 'to'): # Already a tensor or tensor-like object
return lr_mask
try:
if isinstance(lr_mask, dict) and 'quantized' in lr_mask and 'scale' in lr_mask and 'orig_dtype' in lr_mask:
# Auto8bitTensor format
return lr_mask['quantized'].to(lr_mask['orig_dtype']) * lr_mask['scale']
elif isinstance(lr_mask, dict):
# Some other dictionary format - try to extract values
values = list(lr_mask.values())
if isinstance(values[0], (int, float)):
return torch.tensor(values, device=device)
# Default fallback - try to convert to tensor directly
return torch.tensor(lr_mask, device=device)
except Exception as e:
print(f"ERROR: Failed to convert lr_mask to tensor: {e}")
# Return a default tensor as fallback
if shape is not None:
return torch.ones(shape, device=device) * self.lr
else:
return torch.ones(1, device=device) * self.lr
def initialize_state(self, p):
state = self.state[p]
# Basic state initialization
if "step" not in state:
state["step"] = 0
# Initialize lr_mask
if 'lr_mask' not in state:
state['lr_mask'] = Auto8bitTensor(torch.ones(p.shape).to(p.device, dtype=torch.float32) * self.lr)
elif isinstance(state['lr_mask'], dict):
# Convert from dictionary format to tensor
try:
if 'quantized' in state['lr_mask'] and 'scale' in state['lr_mask']:
tensor = state['lr_mask']['quantized'].to(
state['lr_mask'].get('orig_dtype', torch.float32)
) * state['lr_mask']['scale']
state['lr_mask'] = Auto8bitTensor(tensor)
else:
# Fallback
state['lr_mask'] = Auto8bitTensor(torch.ones(p.shape).to(p.device, dtype=torch.float32) * self.lr)
except Exception as e:
print(f"Error converting lr_mask: {e}")
state['lr_mask'] = Auto8bitTensor(torch.ones(p.shape).to(p.device, dtype=torch.float32) * self.lr)
# Other states
if 'avg_lr' not in state:
state['avg_lr'] = self.lr
if 'last_polarity' not in state:
state['last_polarity'] = torch.zeros(p.shape, dtype=torch.bool, device=p.device)
factored = len(p.shape) >= 2
if factored:
if "exp_avg_sq_row" not in state:
state["exp_avg_sq_row"] = torch.zeros(p.shape[:-1]).to(p)
if "exp_avg_sq_col" not in state:
state["exp_avg_sq_col"] = torch.zeros(p.shape[:-2] + p.shape[-1:]).to(p)
else:
if "exp_avg_sq" not in state:
state["exp_avg_sq"] = torch.zeros_like(p)
if "RMS" not in state:
state["RMS"] = 0
@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.
"""
self.step_hook()
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
for p in group["params"]:
if p.grad is None or not p.requires_grad:
continue
grad = p.grad
if grad.dtype != torch.float32:
grad = grad.to(torch.float32)
if grad.is_sparse:
raise RuntimeError("Automagic does not support sparse gradients.")
state = self.state[p]
# Initialize state if not already initialized
if len(state) == 0 or 'lr_mask' not in state:
self.initialize_state(p)
grad_shape = grad.shape
factored = len(grad_shape) >= 2
if factored:
state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)
state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)
else:
state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)
p_data_fp32 = p
if isinstance(p_data_fp32, QBytesTensor):
p_data_fp32 = p_data_fp32.dequantize()
if p.dtype != torch.float32:
p_data_fp32 = p_data_fp32.clone().float()
state["step"] += 1
state["RMS"] = self._rms(p_data_fp32)
beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])
eps = group["eps"]
# Dynamic setting of eps: if eps is a tuple/list, interpolate based on average lr_mask
if isinstance(eps, (tuple, list)):
eps_min, eps_max = eps
# Handle case when lr_mask doesn't exist
if 'lr_mask' not in state:
state['lr_mask'] = Auto8bitTensor(torch.ones(p.shape).to(p.device, dtype=torch.float32) * self.lr)
# Safe conversion of lr_mask to tensor
if not isinstance(state['lr_mask'], torch.Tensor) or not hasattr(state['lr_mask'], 'to'):
state['lr_mask'] = self._convert_lr_mask_to_tensor(
state['lr_mask'], p.device, shape=p.shape
)
lr_mask = state['lr_mask'].to(torch.float32)
avg_lr_mask = torch.mean(lr_mask)
norm = ((avg_lr_mask - self.min_lr) / (self.max_lr - self.min_lr)).clamp(0, 1)
eps = eps_min + norm * (eps_max - eps_min)
update = (grad**2) + eps
if factored:
exp_avg_sq_row = state["exp_avg_sq_row"]
exp_avg_sq_col = state["exp_avg_sq_col"]
exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t))
exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t))
# Approximate exponential moving average of squared gradient
update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
update.mul_(grad)
else:
exp_avg_sq = state["exp_avg_sq"]
exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t))
update = exp_avg_sq.rsqrt().mul_(grad)
update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
# Initialize last_polarity if it doesn't exist
if 'last_polarity' not in state:
state['last_polarity'] = torch.zeros(p.shape, dtype=torch.bool, device=p.device)
# Get signs of previous and current updates
last_polarity = state['last_polarity']
current_polarity = (update > 0).to(torch.bool)
sign_agreement = torch.where(last_polarity == current_polarity, 1, -1)
state['last_polarity'] = current_polarity
# Initialize lr_mask if it doesn't exist
if 'lr_mask' not in state:
state['lr_mask'] = Auto8bitTensor(torch.ones(p.shape).to(p.device, dtype=torch.float32) * self.lr)
# Safe handling of lr_mask - ensure it's a tensor
if not isinstance(state['lr_mask'], torch.Tensor) or not hasattr(state['lr_mask'], 'to'):
# print(f"Converting lr_mask: {type(state['lr_mask'])}")
state['lr_mask'] = self._convert_lr_mask_to_tensor(
state['lr_mask'], p.device, shape=p.shape
)
# Convert lr_mask to float32
lr_mask = state['lr_mask'].to(torch.float32)
# Update learning rate mask based on sign agreement
new_lr = torch.where(
sign_agreement > 0,
lr_mask * self.lr_pump_scale, # Increase lr
lr_mask * self.lr_dump_scale # Decrease lr
)
# Clip new learning rates to specified bounds
new_lr = torch.clamp(new_lr, min=self.min_lr, max=self.max_lr)
# Apply new learning rate mask to update
update.mul_(new_lr)
state['lr_mask'] = Auto8bitTensor(new_lr)
state['avg_lr'] = torch.mean(new_lr)
if group["weight_decay"] != 0:
p_data_fp32.add_(p_data_fp32, alpha=(-group["weight_decay"] * new_lr))
p_data_fp32.add_(-update)
if p.dtype != torch.float32:
# Apply stochastic rounding if necessary
copy_stochastic(p, p_data_fp32)
return loss
# Override the state_dict to save the lr_mask
def state_dict(self):
orig_state_dict = super().state_dict()
new_state = {}
for p_key, state in orig_state_dict['state'].items():
save_state = {k: v for k, v in state.items() if k != 'lr_mask'}
if 'lr_mask' in state:
if hasattr(state['lr_mask'], 'state_dict'):
save_state['lr_mask'] = state['lr_mask'].state_dict()
elif isinstance(state['lr_mask'], dict):
save_state['lr_mask'] = state['lr_mask']
elif isinstance(state['lr_mask'], torch.Tensor):
# Create a dict representation for the tensor
save_state['lr_mask'] = {
'quantized': state['lr_mask'],
'orig_dtype': state['lr_mask'].dtype,
'scale': 1.0
}
new_state[p_key] = save_state
orig_state_dict['state'] = new_state
return orig_state_dict
def load_state_dict(self, state_dict):
# Make a copy to avoid modifying the input
state_dict_copy = {'param_groups': state_dict['param_groups'], 'state': {}}
# Process state entries
for p_key, saved_state in state_dict['state'].items():
state_dict_copy['state'][p_key] = {k: v for k, v in saved_state.items() if k != 'lr_mask'}
# Find matching parameter in current optimizer
param_found = False
for group in self.param_groups:
for p in group['params']:
if str(id(p)) == str(p_key) or id(p) == p_key:
param_found = True
# Handle lr_mask separately
if 'lr_mask' in saved_state:
try:
if isinstance(saved_state['lr_mask'], dict) and 'quantized' in saved_state['lr_mask']:
tensor = saved_state['lr_mask']['quantized'].to(
saved_state['lr_mask'].get('orig_dtype', torch.float32)
) * saved_state['lr_mask'].get('scale', 1.0)
self.state[p]['lr_mask'] = Auto8bitTensor(tensor)
elif isinstance(saved_state['lr_mask'], torch.Tensor):
self.state[p]['lr_mask'] = Auto8bitTensor(saved_state['lr_mask'])
except Exception as e:
print(f"Error loading lr_mask: {e}")
# Initialize with default
self.state[p]['lr_mask'] = Auto8bitTensor(
torch.ones(p.shape).to(p.device, dtype=torch.float32) * self.lr
)
break
if param_found:
break
# Load the processed state dict
super().load_state_dict(state_dict_copy)