Speed decrease over time

[GreenLandisaLie]

I don't know why but sometimes the training speed decreases considerably out of no where. It starts fine for the first steps then takes a nose dive.

INFO:musubi_tuner.dataset.image_video_dataset:epoch is incremented. current_epoch: 0, epoch: 1
steps:   7%|███▎                                             | 250/3646 [2:01:03<27:24:21, 29.05s/it, avr_loss=0.00565]
saving checkpoint: O:/SD/DATASETS/ANIFY/musubi_output\qwen_edit_anify_1024x1024-step00000250.safetensors
steps:  14%|██████▋                                          | 500/3646 [6:03:45<38:08:48, 43.65s/it, avr_loss=0.00564]
saving checkpoint: O:/SD/DATASETS/ANIFY/musubi_output\qwen_edit_anify_1024x1024-step00000500.safetensors
steps:  21%|█████████▊                                      | 750/3646 [15:59:42<61:45:46, 76.78s/it, avr_loss=0.00589]
saving checkpoint: O:/SD/DATASETS/ANIFY/musubi_output\qwen_edit_anify_1024x1024-step00000750.safetensors
steps:  22%|██████████▌                                     | 805/3646 [18:10:12<64:07:33, 81.26s/it, avr_loss=0.00583]

The log above is for a Qwen-Edit lora being trained at 1024x1024 - which I've done many times (without the speed loss) with pretty much the exact same config.

But every so often the speed decrease you see here happens for no reason. VRAM usage suddenly increases and overflows to RAM - causing this issue. But why does it happen sometimes and not every time? How can I avoid it?

Its not as though I launched a game or opened another AI model that consumed VRAM while it was training - or something like that. The log above was from over-night training so no external variables took place.

EDIT:
I should probably show you my config:

accelerate launch ^
--num_cpu_threads_per_process 1 ^
--mixed_precision bf16 src/musubi_tuner/qwen_image_train_network.py ^
--dit "F:/SD/Qwen-Image-Edit/qwen_image_edit_bf16.safetensors" ^
--vae "F:/SD/Qwen-Image-Edit/qwen_image_vae.safetensors" ^
--text_encoder "F:/SD/Qwen-Image-Edit/qwen_2.5_vl_7b.safetensors" ^
--dataset_config "O:/SD/DATASETS/ANIFY/config_1024.toml" ^
--edit ^
--fp8_base ^
--fp8_scaled ^
--fp8_vl ^
--blocks_to_swap 55 ^
--sdpa ^
--mixed_precision bf16 ^
--timestep_sampling qwen_shift ^
--weighting_scheme none ^
--optimizer_type adamw8bit ^
--learning_rate 0.0002 ^
--gradient_checkpointing ^
--max_data_loader_n_workers 2 ^
--persistent_data_loader_workers ^
--network_module networks.lora_qwen_image ^
--network_dim 64 ^
--max_train_epochs 2 ^
--save_every_n_epochs 1 ^
--save_every_n_steps 250 ^
--seed 42 ^
--logging_dir=logs ^
--network_weights "O:/SD/DATASETS/ANIFY/musubi_output/qwen_edit_anify_640x640-step00006000.safetensors" ^
--output_dir "O:/SD/DATASETS/ANIFY/musubi_output" ^
--output_name "qwen_edit_anify_1024x1024"

I used to train at rank 16 or 32 with 45 blocks swapped at 1024x1024 - successfully. I can't remember if I ever had this issue with those settings.

At some point I decided to increase the number of swapping blocks to 55 to have enough spare VRAM for safety and I'm not sure if that was the moment I started having this issue or if it was due to updating the program but from there on out I was already facing this issue - even when training at rank 16. Right now I'm training at rank 64 but I don't think that's related.

EDIT2: Windows 10, RTX 5060Ti 16Gb, 64Gb RAM

[kohya-ss]

I can't think of any reason why this is happening, but does it occur because of checkpoint saving? Or is it unrelated?

Also, if the batch size during training is different, VRAM usage can increase when large batches are trained. (The number of elements in a batch is variable in Musubi Tuner.)

Could you please share the image resolution, number of images, and number of batches logged at the start of training? Something like this:

INFO:musubi_tuner.dataset.image_video_dataset:bucket: (768, 1360), count: 4
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (848, 1232), count: 9
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (896, 1168), count: 10
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1024, 1024), count: 4
INFO:musubi_tuner.dataset.image_video_dataset:total batches: 8

[GreenLandisaLie]

Sorry, should have given that info as well.
The log above does make it seem like its related to saving checkpoints but I could almost swear it was just a coincidence in this case (I think in some cases I noticed a drop before it ever saved but I'm not so sure anymore).
I'll double check on that and in the meantime I'll give you the info you asked.

This is what the .toml looks like:

# general configurations
[general]
resolution = [512, 512] # optional, [W, H], default is [960, 544]. This is the default resolution for all datasets
caption_extension = ".txt" # optional, default is None. This is the default caption extension for all datasets
batch_size = 1 # optional, default is 1. This is the default batch size for all datasets
num_repeats = 1 # optional, default is 1. Number of times to repeat the dataset. Useful to balance the multiple datasets with different sizes.
enable_bucket = true # optional, default is false. Enable bucketing for datasets
bucket_no_upscale = false # optional, default is false. Disable upscaling for bucketing. Ignored if enable_bucket is false

[[datasets]]
image_directory = "O:/SD/DATASETS/ANIFY/img"
control_directory = "O:/SD/DATASETS/ANIFY/ctrl"
cache_directory = "O:/SD/DATASETS/ANIFY/musubi_cache_1024"
resolution = [1024, 1024]

The dataset contains 1823 image pairs (big I know). Their cache has been pre-generated.

Actually I just found something wrong after checking the log:

INFO:musubi_tuner.dataset.image_video_dataset:  1820 images have 1 control images
INFO:musubi_tuner.dataset.image_video_dataset:  3 images have 2 control images

This is for the original Qwen-Edit model - I'm not even sure it supports training on multiple controls and my dataset is supposed to be just single pairs.
I'm constantly tweaking the dataset so there would be times when this might have happened and times when it might not (only noticed now tho) - unlikely but maybe related?
After checking the code I found its because I have 'imgname_1.ext' and 'imgname_11.ext' files in my dataset (both on the image and control directories) and you added multiple controls support using '_0000' suffix - this was kind of a bad idea, no?
I mean, it kinda broke compatibility in my case. Maybe demanding multiple control dirs would have been better but whatever.

This is the part of the log you asked:

INFO:musubi_tuner.dataset.image_video_dataset:bucket: (576, 1808), count: 2
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (624, 1680), count: 2
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (640, 1632), count: 1
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (656, 1584), count: 1
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (672, 1552), count: 1
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (688, 1520), count: 2
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (704, 1488), count: 5
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (720, 1456), count: 8
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (736, 1424), count: 9
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (752, 1392), count: 10
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (768, 1360), count: 44
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (784, 1328), count: 27
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (800, 1296), count: 25
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (816, 1280), count: 52
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (832, 1248), count: 192
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (848, 1232), count: 58
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (864, 1200), count: 240
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (880, 1184), count: 241
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (896, 1168), count: 60
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (912, 1136), count: 43
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (928, 1120), count: 30
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (944, 1104), count: 31
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (960, 1088), count: 17
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (976, 1072), count: 25
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (992, 1056), count: 13
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1008, 1040), count: 2
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1024, 1024), count: 65
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1040, 1008), count: 7
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1056, 992), count: 13
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1072, 976), count: 12
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1088, 960), count: 10
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1104, 944), count: 7
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1120, 928), count: 10
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1136, 912), count: 29
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1168, 896), count: 12
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1184, 880), count: 142
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1200, 864), count: 53
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1232, 848), count: 17
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1248, 832), count: 170
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1280, 816), count: 25
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1296, 800), count: 18
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1328, 784), count: 6
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1360, 768), count: 74
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1392, 752), count: 2
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1424, 736), count: 6
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1488, 704), count: 1
INFO:musubi_tuner.dataset.image_video_dataset:bucket: (1520, 688), count: 3
INFO:musubi_tuner.dataset.image_video_dataset:total batches: 1823

I'll do some test runs now and see if the issue above is related or not and if saving the checkpoint is related as well.

[GreenLandisaLie]

So I'm re-running the same training as before but this time saving every 100 steps (and with the extra controls fixed).

VRAM consumption is about 14.7Gb (with 1.1Gb shared for some reason).
It remained like that from step 15 all the way to 105+ without ever changing even after saving on step 100.

it/s also remained consistent even after saving.

Its still too soon to say and I could be wrong but I think the issue might have really been related to the unintended multiple controls - they should increase VRAM by quite a bit and at 14.7Gb there's not enough room on 16Gb for an extra control.

The reason why sometimes this happened and sometimes it didn't - was probably because the code had a flaw that would cause some images to be skipped and others to be re-used on the same epoch - I believe you already fixed that bug.

So, for as long as those with multiple controls were being skipped - it would be fine but the moment they were not (on 1024x1024), it would cause memory to overflow during that particular step but not crash due to NVIDIA's default memory offload. Once that happened the block swapping mechanics freaked out and never really fixed themselves.

At least this is my hope. I'll let it continue training and see if it happens again at later steps.

[kohya-ss]

Thank you for sharing the information.

Since the process is common between Qwen-Image-Edit and Edit-2509, training on multiple control images can also be done with Edit (although this may or may not work).

It seems likely that multiple control images are part of the problem. The script does not explicitly free memory during training.
The timing of when those 3 images out of 1800 are trained is random, so I think the slowing down may be caused by when they are trained.

[GreenLandisaLie]

I'm now 100% positive that was indeed the issue.

Did a test run with just 10 images and 1 of them had 2 controls.
Also injected some code into the program so it would log which image(s) were being used on each training step.
As soon as it reached the image with 2 controls -> VRAM increased from 14.7Gb +1.1Gb shared to 15.5Gb +3.0Gb shared and that step took significantly more time to complete than the previous ones.

Additionally VRAM usage did not decrease after that step and every consequent step was slower than it should be. Maybe this is something that could be improved/fixed if you consider it a 'bug'. If not, feel free to close this issue as it can be narrowed down to low VRAM.

EDIT: just in case I did not phrase myself properly - the reason why I think this could be considered a 'bug' its because its natural for VRAM consumption to increase (and potentially overflow) when going from 1 control to 2 (or more) - since its a huge bump in resolution due to image concatenation. The problem, is that the program does not handle 'memory' properly after overflowing and going back to inputs that should originally fit in the user's VRAM.

[FurkanGozukara]

@GreenLandisaLie i found error it is this litearlly : #678

#678 (comment)

[GreenLandisaLie]

@FurkanGozukara Yes that was literally what caused me this problem - it definitely needs to be fixed.
But that aside, there is also a slight/major inefficiency issue in regards to temporary memory overflows. When 1 particular step (or even a dataset when training on multiple ones) causes memory overflow - the program should fix that once it starts training again on inputs that should fit in VRAM. I know this might be difficult and maybe even not possible, but I had to mention it.

[FurkanGozukara]

@FurkanGozukara Yes that was literally what caused me this problem - it definitely needs to be fixed. But that aside, there is also a slight/major inefficiency issue in regards to temporary memory overflows. When 1 particular step (or even a dataset when training on multiple ones) causes memory overflow - the program should fix that once it starts training again on inputs that should fit in VRAM. I know this might be difficult and maybe even not possible, but I had to mention it.

after this fix i didnt see vram increase over time

but just 1 control image reduces speed like 100% you noticed same? any ideas? i wouldnt expect such dramatic speed drop

[GreenLandisaLie]

Controls are simply concatenated to the target image, right? If there is no internal rescaling of any sort in the training code then at 1328x1328:

1 control: (1328+1328)*1328 = 3,527,168 pixels
2 controls: (1328+1328+1328)*1328 = 5,290,752 pixels

Ofc this is not factors VAE encoding but still shows why you could very easily see huge speed gains by just reducing the number of controls at high res buckets.

It is precisely because of this that if you have a particular target image (or a few) with more controls than the others in your dataset - and that single image causes a memory overflow: then after training on that step the program should 'refresh the model+optimizer memory allocation' and resume training at optimal speed instead of keeping being chaotic all the time. As it stands, after a memory overflow it keeps being slow until the end of training - no matter what.

[FurkanGozukara]

@GreenLandisaLie i am training 1 fine tuning 1 lora with 1 train image + 1 control image atm

all images are 1328x1328

it has been like 7 hours here

0 vram leak however speed is slow

[GreenLandisaLie]

@FurkanGozukara Bro you telling me that's slow? Try bucket 1328x1328 with 1 control on a 5060Ti and enjoy near 1min/it speeds. This ofc with all speed and memory optimizations available - otherwise it would probably be 3min or something.

[FurkanGozukara]

@FurkanGozukara Bro you telling me that's slow? Try bucket 1328x1328 with 1 control on a 5060Ti and enjoy near 1min/it speeds. This ofc with all speed and memory optimizations available - otherwise it would probably be 3min or something.

i mean it is slow compared to Qwen Image base model :)

[GreenLandisaLie]

Did 2 test runs to demonstrate the overflow problem.

The combination of all of the used settings will trigger a memory overflow on 16Gb GPUs when training on a 2 control image but not on images with 1 control.
The test dataset contains 6 images, 1 of them has 2 controls and the others have 1.

INFO:musubi_tuner.dataset.image_video_dataset:  5 images have 1 control images
INFO:musubi_tuner.dataset.image_video_dataset:  1 images have 2 control images

The name of the image with 2 controls is: 'img4_1'
Note: this is with Musubi-Tuner before the control filename logic was fixed.

I injected some code in the for step training loop to create a simple log.

This is what happens right now when a overflow occurs:

epoch 1 - step 1 - time(s): 29.10 - filenames: img4_11
epoch 1 - step 2 - time(s): 28.34 - filenames: img2
epoch 1 - step 3 - time(s): 28.51 - filenames: img1
epoch 1 - step 4 - time(s): 28.87 - filenames: img0
epoch 1 - step 5 - time(s): 28.34 - filenames: img3
epoch 1 - step 6 - time(s): 313.44 - filenames: img4_1
epoch 2 - step 7 - time(s): 117.89 - filenames: img2
epoch 2 - step 8 - time(s): 117.61 - filenames: img4_11
epoch 2 - step 9 - time(s): 146.57 - filenames: img1
epoch 2 - step 10 - time(s): 141.05 - filenames: img3
epoch 2 - step 11 - time(s): 88.39 - filenames: img0
epoch 2 - step 12 - time(s): 311.05 - filenames: img4_1
epoch 3 - step 13 - time(s): 129.61 - filenames: img2
epoch 3 - step 14 - time(s): 87.06 - filenames: img0
epoch 3 - step 15 - time(s): 115.09 - filenames: img4_11
epoch 3 - step 16 - time(s): 158.36 - filenames: img1
epoch 3 - step 17 - time(s): 120.77 - filenames: img3

Notice how even after the 'img4_1' was used - the steps for the next ones (which should fit in memory), are still very slow.
This is because the input tensor memory allocation has expanded to the size of 'img4_1' and it never shrinks so when populating it with the other (smaller) images - its still overflowing.

I tried running this at the end of the step just to see if it would fix it:

if files_str == "img4_1": # don't mind the 'files_str' here - its part of my own code for a custom logger similar to the one I'm showing
    del latents, noise, noisy_model_input, weighting, model_pred, target, loss
    clean_memory_on_device(accelerator.device)

And this is the result:

epoch 1 - step 1 - time(s): 29.20 - filenames: img4_11
epoch 1 - step 2 - time(s): 28.32 - filenames: img2
epoch 1 - step 3 - time(s): 28.49 - filenames: img1
epoch 1 - step 4 - time(s): 29.05 - filenames: img0
epoch 1 - step 5 - time(s): 28.73 - filenames: img3
epoch 1 - step 6 - time(s): 315.91 - filenames: img4_1
epoch 2 - step 7 - time(s): 28.22 - filenames: img2
epoch 2 - step 8 - time(s): 28.47 - filenames: img4_11
epoch 2 - step 9 - time(s): 28.45 - filenames: img1

Problem fixed. Input tensors are cleared, their size now fits again in the GPU and speed falls back to what is supposed to be.

The problem is we obviously cannot run 'clean_memory_on_device(accelerator.device)' on every step because it would decrease training speed.
It should run in only cases like this with a well defined logic.
Implementing that would be harsh for me since I'm not a programmer but you can probably do so without breaking a sweat.
Please consider it.

[kohya-ss]

Thanks for reporting this. clean_memory_on_device seems to solve the problem, but it's a bit ad-hoc. Regardless of this, batches with insufficient VRAM will be very slow (if there are 3 control images in the batch it will take a very very long time).

The solution might be to configure your training settings so that your most memory-intensive batch doesn't run out of VRAM.

[GreenLandisaLie]

Ofc but an user might decide to accept that compromise and tolerate low speeds for a particular small part of a dataset (or a particular dataset when training on multiple ones).

The solution might be to configure your training settings so that your most memory-intensive batch doesn't run out of VRAM.

One should assume the user is already using all available optimizations and if a part of their dataset still does not fit in VRAM - no matter what - then we shouldn't allow the whole training to become extremely slow if we can fix it with a few lines of code. Specially when this happens for a really small part of their dataset.

clean_memory_on_device took 0.28s to run in my case - its obviously not something we can or should run every time.
To put it simply, it should only run when:

• the current step is overflowing
• the next step should fit

The question is, how do we determine that? Is there an easy, fast and reliable way to do so?
Maybe checking how much VRAM is available before the main training loop and trying to calculate the size of the input tensors based on their resolution and number of controls, etc? Adding a key with info about the latent size to batch so we can predict the VRAM consumption of the next step?
Dunno - this is the part I'd rather leave to an actual developer.