torch.cuda.OutOfMemoryError: CUDA out of memory on Linux Ubuntu 22.04

[NeoLoger]

I have formatted my PC from Windows11 to Ubuntu 22.04 to see if I can get a small boost in performants.
It is configured as a server in a different room and I'm remotely connecting to it and running SDXL lora training.

After switching to Ubuntu and running a fresh install of kohya_ss I'm getting torch.cuda.OutOfMemoryError: CUDA out of memory. errors.
This is happening on the exact same data sets and training configuration that was working fine on the Windows install without any issues.

Here is the error output:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 7.78 GiB total capacity; 7.41 GiB already allocated; 9.06 MiB free; 7.47 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
steps:   0%|      

And this is the traceback that I get containing the command for the training.

Traceback (most recent call last):
  File "/home/NeoLoger/ai/kohya/kohya_ss/venv/bin/accelerate", line 8, in <module>
    sys.exit(main())
  File "/home/NeoLoger/ai/kohya/kohya_ss/venv/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py", line 47, in main
    args.func(args)
  File "/home/NeoLoger/ai/kohya/kohya_ss/venv/lib/python3.10/site-packages/accelerate/commands/launch.py", line 1017, in launch_command
    simple_launcher(args)
  File "/home/NeoLoger/ai/kohya/kohya_ss/venv/lib/python3.10/site-packages/accelerate/commands/launch.py", line 637, in simple_launcher
    raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['/home/NeoLoger/ai/kohya/kohya_ss/venv/bin/python', './sdxl_train_network.py', '--enable_bucket', '--min_bucket_reso=256', '--max_bucket_reso=1024', '--pretrained_model_name_or_path=/home/NeoLoger/ai/models/sd_xl_base_1.0.safetensors', '--train_data_dir=/home/NeoLoger/ai/ImagesForSDTrining/Lora training/person/Train/images', '--resolution=768,768', '--output_dir=/home/NeoLoger/ai/ImagesForSDTrining/Lora training/person/Train/model', '--logging_dir=/home/NeoLoger/ai/ImagesForSDTrining/Lora training/person/Train/logs', '--network_alpha=12', '--training_comment=Extremely person, some times referred also as Wasp Waist ', '--save_model_as=safetensors', '--network_module=networks.lora', '--text_encoder_lr=5e-05', '--unet_lr=0.0001', '--network_dim=64', '--output_name=Narrow_Waist_v0.1_SDXL', '--lr_scheduler_num_cycles=1', '--no_half_vae', '--learning_rate=0.0001', '--lr_scheduler=cosine', '--train_batch_size=1', '--max_train_steps=2776', '--save_every_n_epochs=1', '--mixed_precision=bf16', '--save_precision=bf16', '--seed=1234', '--caption_extension=.txt', '--cache_latents', '--cache_latents_to_disk', '--optimizer_type=Adafactor', '--optimizer_args', 'scale_parameter=False', 'relative_step=False', 'warmup_init=False', '--max_grad_norm=1', '--max_data_loader_n_workers=1', '--clip_skip=2', '--bucket_reso_steps=64', '--flip_aug', '--shuffle_caption', '--bucket_no_upscale', '--noise_offset=0.0']' returned non-zero exit status 1.

I want to stress again, this is on the same PC and same training configuration.
The only difference is that I'm running Ubuntu now.

Am I missing something on the OS?
Maybe something is wrong with Gradio? (there was an issue with that since yesterday)
Is something missing in my torch install?

Pleas anyone have any suggestions?

[kferterb]

Stable diffusion is well supported for windows. I don't think there will be a performance gain, but errors often occur. So I would not agonize and go back to windows 10, and use fork https://github.com/serpotapov/Kohya_ss-GUI-LoRA-Portable?tab=readme-ov-file in conjunction with fork https://github.com/XpucT/stable-diffusion-webui there is a large community there.

[NeoLoger]

Thanks, I'm aware that the performance gains are minimal if any at all using Linux, but I want to still try it.
Going back to Windows will be my last resort if I wont manage to figurate a solution.

I still want to try and make this work, I kinda like the new setup, there must be something I'm missing to make it work.
I refuse to believe that Windows is a superior setup for AI training...

Any advise on what the issue may be with the CUDA out of memory ?

[rubikstriangle]

Did you install cuda toolkit 11.8?

[anae-git]

I have been using Debian for kohya_ss for months. I bumped some packages and my kernel recently and also pulled the latest kohya_ss today and now training configs that worked fine 2 months ago are running out of memory with this backtrace:

  File "/u/SD/kohya_ss/sd-scripts/sdxl_train.py", line 792, in <module>
    train(args)
  File "/u/SD/kohya_ss/sd-scripts/sdxl_train.py", line 597, in train
    accelerator.backward(loss)
  File "/u/SD/kohya_ss/venv/lib/python3.10/site-packages/accelerate/accelerator.py", line 1905, in backward
    loss.backward(**kwargs)
  File "/u/SD/kohya_ss/venv/lib/python3.10/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/u/SD/kohya_ss/venv/lib/python3.10/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/u/SD/kohya_ss/venv/lib/python3.10/site-packages/torch/autograd/function.py", line 288, in apply
    return user_fn(self, *args)
  File "/u/SD/kohya_ss/venv/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 288, in backward
    torch.autograd.backward(outputs_with_grad, args_with_grad)
  File "/u/SD/kohya_ss/venv/lib/python3.10/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 38.00 MiB. GPU 0 has a total capacty of 23.64 GiB of which 39.19 MiB is free. Including non-PyTorch memory, this process has 22.93 GiB memory in use. Of the allocated memory 21.54 GiB is allocated by PyTorch, and 928.73 MiB is reserved by PyTorch but unallocated.

[maybleMyers]

Sometimes windows will use some ram as vram and avoid the oom errors in linux.