Enable Training on 6GB Cards... with DeepSpeed?

[brucethemoose]

I am trying to squeeze training onto my 6GB laptop RTX 2060, and cant quite manage it with the "low memory" config:

accelerate launch --num_cpu_threads_per_process 8 train_db.py \
--pretrained_model_name_or_path="/home/alpha/Storage/AIModels/Stable-diffusion/panatomy05full_0.7-AIModels_Anything-V3.0-pruned-fp16_0.3-Weighted_sum-merged.ckpt" \
--train_data_dir="/home/alpha/Storage/TrainingData/test/training_data" \
--output_dir="/home/alpha/Storage/TrainingOutput/test/" \
--prior_loss_weight=1.0 \
--resolution=512 \
--train_batch_size=1 \
--learning_rate=1e-6 \
--max_train_steps=1600 \
--use_8bit_adam \
--xformers \
--mixed_precision="fp16" \
--cache_latents \
--gradient_checkpointing \
--save_precision="fp16" \
--full_fp16 \
--save_model_as="safetensors" \

So, I figured I would investigate Deepspeed cpu offloading with the accelerate config... but I keep running into errors on both the git version and the 0.7.7 release from pypi. Here is an error from the pypi release:

Traceback (most recent call last):
  File "/home/alpha/clone/sd-scripts/train_db.py", line 332, in <module>
    train(args)
  File "/home/alpha/clone/sd-scripts/train_db.py", line 154, in train
    unet, text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
  File "/home/alpha/.local/lib/python3.10/site-packages/accelerate/accelerator.py", line 619, in prepare
    result = self._prepare_deepspeed(*args)
  File "/home/alpha/.local/lib/python3.10/site-packages/accelerate/accelerator.py", line 805, in _prepare_deepspeed
    engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
  File "/home/alpha/.local/lib/python3.10/site-packages/deepspeed/__init__.py", line 125, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/home/alpha/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 330, in __init__
    self._configure_optimizer(optimizer, model_parameters)
  File "/home/alpha/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1210, in _configure_optimizer
    self.optimizer = self._configure_zero_optimizer(basic_optimizer)
  File "/home/alpha/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1455, in _configure_zero_optimizer
    optimizer = DeepSpeedZeroOptimizer(
  File "/home/alpha/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 532, in __init__
    self._param_slice_mappings = self._create_param_mapping()
  File "/home/alpha/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 544, in _create_param_mapping
    lp_name = self.param_names[lp]
KeyError: <exception str() failed>
[2023-01-12 13:13:52,241] [INFO] [launch.py:318:sigkill_handler] Killing subprocess 5398
[2023-01-12 13:13:52,244] [ERROR] [launch.py:324:sigkill_handler] ['/usr/bin/python', '-u', 'train_db.py', '--pretrained_model_name_or_path=/home/alpha/Storage/AIModels/Stable-diffusion/panatomy05full_0.7-AIModels_Anything-V3.0-pruned-fp16_0.3-Weighted_sum-merged.ckpt', '--train_data_dir=/home/alpha/Storage/TrainingData/test/training_data', '--output_dir=/home/alpha/Storage/TrainingOutput/test/', '--prior_loss_weight=1.0', '--resolution=512', '--train_batch_size=1', '--learning_rate=1e-6', '--max_train_steps=1600', '--use_8bit_adam', '--xformers', '--mixed_precision=fp16', '--cache_latents', '--gradient_checkpointing', '--save_precision=fp16', '--full_fp16', '--save_model_as=safetensors'] exits with return code = 1
Traceback (most recent call last):
  File "/home/alpha/.local/bin/accelerate", line 8, in <module>
    sys.exit(main())
  File "/home/alpha/.local/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py", line 43, in main
    args.func(args)
  File "/home/alpha/.local/lib/python3.10/site-packages/accelerate/commands/launch.py", line 827, in launch_command
    deepspeed_launcher(args)
  File "/home/alpha/.local/lib/python3.10/site-packages/accelerate/commands/launch.py", line 540, in deepspeed_launcher
    raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['deepspeed', '--no_local_rank', '--num_gpus', '1', 'train_db.py', '--pretrained_model_name_or_path=/home/alpha/Storage/AIModels/Stable-diffusion/panatomy05full_0.7-AIModels_Anything-V3.0-pruned-fp16_0.3-Weighted_sum-merged.ckpt', '--train_data_dir=/home/alpha/Storage/TrainingData/test/training_data', '--output_dir=/home/alpha/Storage/TrainingOutput/test/', '--prior_loss_weight=1.0', '--resolution=512', '--train_batch_size=1', '--learning_rate=1e-6', '--max_train_steps=1600', '--use_8bit_adam', '--xformers', '--mixed_precision=fp16', '--cache_latents', '--gradient_checkpointing', '--save_precision=fp16', '--full_fp16', '--save_model_as=safetensors']' returned non-zero exit status 1.

Is there anything in particular that needs to be changed for this repo to support deepspeed? Or maybe there is some other tweak to squeeze LORA onto 6GB?

[brucethemoose]

Tried upgrading accelerate and downgrading pytorch to 1.13.1, still no dice:

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /home/alpha/clone/sd-scripts/train_db.py:332 in <module>                                         │
│                                                                                                  │
│   329 │   │   │   │   │     help="steps to stop text encoder training, -1 for no training / Te   │
│   330                                                                                            │
│   331   args = parser.parse_args()                                                               │
│ ❱ 332   train(args)                                                                              │
│   333                                                                                            │
│                                                                                                  │
│ /home/alpha/clone/sd-scripts/train_db.py:154 in train                                            │
│                                                                                                  │
│   151                                                                                            │
│   152   # acceleratorがなんかよろしくやってくれるらしい                                          │
│   153   if train_text_encoder:                                                                   │
│ ❱ 154 │   unet, text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(   │
│   155 │   │   unet, text_encoder, optimizer, train_dataloader, lr_scheduler)                     │
│   156   else:                                                                                    │
│   157 │   unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimize   │
│                                                                                                  │
│ /home/alpha/.local/lib/python3.10/site-packages/accelerate/accelerator.py:872 in prepare         │
│                                                                                                  │
│    869 │   │   │   old_named_params = self._get_named_parameters(*args)                          │
│    870 │   │                                                                                     │
│    871 │   │   if self.distributed_type == DistributedType.DEEPSPEED:                            │
│ ❱  872 │   │   │   result = self._prepare_deepspeed(*args)                                       │
│    873 │   │   elif self.distributed_type == DistributedType.MEGATRON_LM:                        │
│    874 │   │   │   result = self._prepare_megatron_lm(*args)                                     │
│    875 │   │   else:                                                                             │
│                                                                                                  │
│ /home/alpha/.local/lib/python3.10/site-packages/accelerate/accelerator.py:1093 in                │
│ _prepare_deepspeed                                                                               │
│                                                                                                  │
│   1090 │   │   │   │   │   │   if type(scheduler).__name__ in deepspeed.runtime.lr_schedules.VA  │
│   1091 │   │   │   │   │   │   │   kwargs["lr_scheduler"] = scheduler                            │
│   1092 │   │   │                                                                                 │
│ ❱ 1093 │   │   │   engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)           │
│   1094 │   │   │   if optimizer is not None:                                                     │
│   1095 │   │   │   │   optimizer = DeepSpeedOptimizerWrapper(optimizer)                          │
│   1096 │   │   │   if scheduler is not None:                                                     │
│                                                                                                  │
│ /home/alpha/.local/lib/python3.10/site-packages/deepspeed/__init__.py:125 in initialize          │
│                                                                                                  │
│   122 │   assert model is not None, "deepspeed.initialize requires a model"                      │
│   123 │                                                                                          │
│   124 │   if not isinstance(model, PipelineModule):                                              │
│ ❱ 125 │   │   engine = DeepSpeedEngine(args=args,                                                │
│   126 │   │   │   │   │   │   │   │    model=model,                                              │
│   127 │   │   │   │   │   │   │   │    optimizer=optimizer,                                      │
│   128 │   │   │   │   │   │   │   │    model_parameters=model_parameters,                        │
│                                                                                                  │
│ /home/alpha/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py:330 in __init__      │
│                                                                                                  │
│    327 │   │   │   model_parameters = self.module.parameters()                                   │
│    328 │   │                                                                                     │
│    329 │   │   if has_optimizer:                                                                 │
│ ❱  330 │   │   │   self._configure_optimizer(optimizer, model_parameters)                        │
│    331 │   │   │   self._configure_lr_scheduler(lr_scheduler)                                    │
│    332 │   │   │   self._report_progress(0)                                                      │
│    333 │   │   elif self.zero_optimization():                                                    │
│                                                                                                  │
│ /home/alpha/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py:1210 in              │
│ _configure_optimizer                                                                             │
│                                                                                                  │
│   1207 │   │   optimizer_wrapper = self._do_optimizer_sanity_check(basic_optimizer)              │
│   1208 │   │                                                                                     │
│   1209 │   │   if optimizer_wrapper == ZERO_OPTIMIZATION:                                        │
│ ❱ 1210 │   │   │   self.optimizer = self._configure_zero_optimizer(basic_optimizer)              │
│   1211 │   │   elif optimizer_wrapper == AMP:                                                    │
│   1212 │   │   │   amp_params = self.amp_params()                                                │
│   1213 │   │   │   log_dist(f"Initializing AMP with these params: {amp_params}", ranks=[0])      │
│                                                                                                  │
│ /home/alpha/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py:1455 in              │
│ _configure_zero_optimizer                                                                        │
│                                                                                                  │
│   1452 │   │   │   │   │   │   "Pipeline parallelism does not support overlapped communication,  │
│   1453 │   │   │   │   │   )                                                                     │
│   1454 │   │   │   │   │   overlap_comm = False                                                  │
│ ❱ 1455 │   │   │   optimizer = DeepSpeedZeroOptimizer(                                           │
│   1456 │   │   │   │   optimizer,                                                                │
│   1457 │   │   │   │   self.param_names,                                                         │
│   1458 │   │   │   │   timers=timers,                                                            │
│                                                                                                  │
│ /home/alpha/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py:532 in   │
│ __init__                                                                                         │
│                                                                                                  │
│    529 │   │                                                                                     │
│    530 │   │   self._link_all_hp_params()                                                        │
│    531 │   │   self._enable_universal_checkpoint()                                               │
│ ❱  532 │   │   self._param_slice_mappings = self._create_param_mapping()                         │
│    533 │                                                                                         │
│    534 │   def _enable_universal_checkpoint(self):                                               │
│    535 │   │   for lp_param_group in self.bit16_groups:                                          │
│                                                                                                  │
│ /home/alpha/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py:544 in   │
│ _create_param_mapping                                                                            │
│                                                                                                  │
│    541 │   │   │   param_mapping_per_group = OrderedDict()                                       │
│    542 │   │   │   for lp in self.bit16_groups[i]:                                               │
│    543 │   │   │   │   if lp._hp_mapping is not None:                                            │
│ ❱  544 │   │   │   │   │   lp_name = self.param_names[lp]                                        │
│    545 │   │   │   │   │   param_mapping_per_group[                                              │
│    546 │   │   │   │   │   │   lp_name] = lp._hp_mapping.get_hp_fragment_address()               │
│    547 │   │   │   param_mapping.append(param_mapping_per_group)                                 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
KeyError: <exception str() failed>

[kohya-ss]

Hi,
Unfortunately I have no experience about DeepSpeed.

However, you seem to run train_db.py. It is intended to train DreamBooth. If you want to train LoRA, please use train_network.py (some argments should be modified.) LoRA training might be work with 6GB VRAM without full_fp16 option.

[brucethemoose]

Hi, Unfortunately I have no experience about DeepSpeed.

However, you seem to run train_db.py. It is intended to train DreamBooth. If you want to train LoRA, please use train_network.py (some argments should be modified.) LoRA training might be work with 6GB VRAM without full_fp16 option.

🤦 I knew train_network.py was LORA, but I apparently copy pasted the wrong file in. Thanks 👍

I figured full-fp16 would reduce VRAM usage, but I will leave it out if thats not the case.

[brucethemoose]

Yeah, looks like LORA training uses about 5.1GB without deepspeed, torch inductor or anything weird like that.

But I am leaving this open because I can't get either to work yet.

[healthyfat]

I am trying to squeeze training onto my 6GB laptop RTX 2060

Same here. Looking to run any training on my 2060 RTX 6GB VRAM. So far I had most success with LORA training on Automatic1111 after adjusting some settings, but even though I manage to complete training it didn't do any actual training for some reason. Even after 5000 epochs on 7 images. So I am still looking for the solution.

[lolxdmainkaisemaanlu]

I am able to do excellent LORA training on my 6GB GTX 1060. Using xformers, 8bit adam, cache latents, aspect ratio bucket, fp16 save precision. I don't have to do full fp16 experimental. No deepspeed ( I doubt this even works rn ). Gradient checkpointing also not needed.

[healthyfat]

That is great news. Are you using latest Dreambooth version in Auto1111? Cuz I have to use older one, otherwise I am getting OOM errors.

The only settings configuration that let me even start training without OOM error is:

adam
fp16
xformers
gradient checkpointing

without Train Text Encoder

I don't have to do full fp16 experimental.

What is that?

As well what model you use?

Do you use Train Text Encoder?

[lolxdmainkaisemaanlu]

https://github.com/bmaltais/kohya_ss

I am using bmaltais GUI for training. I don't train in automatic's GUI.

the full fp16 experimental is less VRAM consuming than mixed precision training. It's an option in the GUI i'm using.

I use Dreamlike Photoreal v2 model.

Yes, I train the text encoder too.

[brucethemoose]

I would still like to get Accelerate's Deepspeed stage 1 or 2 config working if only so I can train at faster settings or a higher batch size.

I'm trying to investigate whats going wrong, but I can't even get the git version of DeepSpeed to work (see that linked issue above) :/

[brucethemoose]

Also, --full_fp16 doesn't seem to change VRAM consumption or training speed for me.

[healthyfat]

I am using bmaltais GUI for training.

@lolxdmainkaisemaanlu

Are you running this on Windows?

How much RAM do you have?

[lolxdmainkaisemaanlu]

I am using bmaltais GUI for training.

@lolxdmainkaisemaanlu

Are you running this on Windows?

How much RAM do you have?

Yes, I am running it on Windows 11. But I use a custom Windows 11.

I have 16 GB RAM.

Recently I tried LoRA training with batch size = 2 and only some 48 mb could not be allocated! I think with some tweaks and maybe deepspeed implementation it would be possible to even do batch size = 2 on 6 GB VRAM.

I think if deepspeed stage 2 is implemented LoRA training would be possible even on 4 GB VRAM cards!

[brucethemoose]

@lolxdmainkaisemaanlu

I can already manage a batch size of 2 on my 2060, even without gradient checkpointing enabled or any stage of deepspeed enabled. Says it sits at about 5700mb last I checked.

I am running my display through a laptop with an IGP though, so its VRAM is totally empty.

[healthyfat]

Yes, I am running it on Windows 11. But I use a custom Windows 11.

@lolxdmainkaisemaanlu

Ok. So bmaltais GUI is compatible with Windows then. Cuz I am getting some errors during training. I was thinking its because
bitsandbytes is not supported on Windows.

Are you able to train in Dreambooth or only LORA?