kohya-ss / kohya-ss/sd-scripts

Training with Adafactor raises KeyError: 'step'

Open
#342 3 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
7.2k
Forks
1.2k
Avg merge
11m
Merged PRs (30d)
2

Description

When trying to train with Adafactor as the optimiser, it gives the following error:

```
import network module: networks.lora
create LoRA network. base dim (rank): 8, alpha: 1.0
create LoRA for Text Encoder: 72 modules.
create LoRA for U-Net: 192 modules.
enable LoRA for text encoder
enable LoRA for U-Net
prepare optimizer, data loader etc.
use Adafactor optimizer | {'relative_step': True}
relative_step is true / relative_stepがtrueです
learning rate is used as initial_lr / 指定したlearning rateはinitial_lrとして使用されます
unet_lr and text_encoder_lr are ignored / unet_lrとtext_encoder_lrは無視されます
use adafactor_scheduler / スケジューラにadafactor_schedulerを使用します
running training / 学習開始
num train images * repeats / 学習画像の数×繰り返し回数: 400
num reg images / 正則化画像の数: 0
num batches per epoch / 1epochのバッチ数: 50
num epochs / epoch数: 60
batch size per device / バッチサイズ: 8
gradient accumulation steps / 勾配を合計するステップ数 = 2
total optimization steps / 学習ステップ数: 1500
steps: 0%| | 0/1500 [00:00
train(args)
File "D:\stable-diffusion\kohya_ss\train_network.py", line 571, in train
lr_scheduler.step()
File "D:\stable-diffusion\kohya_ss\venv\lib\site-packages\accelerate\scheduler.py", line 76, in step
self.scheduler.step(*args, **kwargs)
File "D:\stable-diffusion\kohya_ss\venv\lib\site-packages\torch\optim\lr_scheduler.py", line 154, in step
values = self.get_lr()
File "D:\stable-diffusion\kohya_ss\venv\lib\site-packages\transformers\optimization.py", line 646, in get_lr
lrs = [
File "D:\stable-diffusion\kohya_ss\venv\lib\site-packages\transformers\optimization.py", line 647, in
opt._get_lr(group, opt.state[group["params"][0]])
File "D:\stable-diffusion\kohya_ss\venv\lib\site-packages\transformers\optimization.py", line 508, in _get_lr
rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"]))
KeyError: 'step'
steps: 0%| | 0/1500 [00:21
File "D:\stable-diffusion\kohya_ss\venv\lib\site-packages\accelerate\commands\accelerate_cli.py", line 45, in main
args.func(args)
File "D:\stable-diffusion\kohya_ss\venv\lib\site-packages\accelerate\commands\launch.py", line 1104, in launch_command
simple_launcher(args)
File "D:\stable-diffusion\kohya_ss\venv\lib\site-packages\accelerate\commands\launch.py", line 567, in simple_launcher
raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
```

For reference, this is the command I used for training:

```
accelerate launch
--num_cpu_threads_per_process=2
"train_network.py"
--enable_bucket
--pretrained_model_name_or_path="D:/models/anythingV5Anything_anythingV5PrtRE.safetensors"
--train_data_dir="D:\2023-03-26\src3"
--resolution=512,512
--output_dir="D:\2023-03-26\trn_v3"
--logging_dir=""
--network_alpha="1"
--save_model_as=safetensors
--network_module=networks.lora
--text_encoder_lr=5e-5
--unet_lr=0.0001
--network_dim=8
--gradient_accumulation_steps=2
--output_name="my_loha_v3.0"
--lr_scheduler_num_cycles="30"
--learning_rate="0.0001"
--lr_scheduler="cosine"
--lr_warmup_steps="150"
--train_batch_size="8"
--max_train_steps="1500"
--save_every_n_epochs="1"
--mixed_precision="bf16"
--save_precision="fp16"
--caption_extension=".txt"
--cache_latents
--optimizer_type="Adafactor"
--clip_skip=2
--vae_batch_size="3"
--bucket_reso_steps=64
--xformers
--bucket_no_upscale
--min_snr_gamma=5
--sample_sampler=euler_a
--sample_prompts="D:\2023-03-26\trn_v3\sample\prompt.txt"
--sample_every_n_epochs="1"
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with train_network.py around line 571, where lr_scheduler.step() is called, and trace the Adafactor setup and scheduler interaction shown in the traceback. Reproduce with the supplied accelerate launch command and inspect the installed PyTorch, Transformers, and Accelerate versions. Done means Adafactor training passes the initial scheduler step without the KeyError: 'step'.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.