SDXL controlnet training, sudden convergence

[brycegoh]

Hello, I am trying to train a controlnet model using the provided script that is conditioned on semantic segmentations. I am looking for a general gauge for how many steps for sudden convergence on sdxl controlnet especially from the diffusers team who have already trained sdxl controlnets.

For the dataset, I am using the ADE20k dataset (20k image pairs). For captions, I am using blip.

I have completed around 15k steps with learning rate of 1e-5, constant learning rate scheduler, gradient accumulation of 1 and batch size of 6, but still have not noticed any sudden convergence.

If anyone from the diffuser team reads this, did you notice any patterns in terms of the step range in which convergence happens during the training of the diffuser's sdxl controlnet? Any advice on the hyper-parameters, mixed-precision config, etc?

Any advice would be appreciated!

[rphly]

This is an interesting discussion and will help us as well. Hoping the team can respond to this.

[sayakpaul]

Thanks for opening this discussion. Does increasing the learning rate help?

Additionally, this script might be helpful for training large-scale:
https://github.com/huggingface/diffusers/blob/main/examples/research_projects/controlnet/train_controlnet_webdataset.py

This is what we used for https://huggingface.co/collections/diffusers/sdxl-controlnets-64f9c35846f3f06f5abe351f

[brycegoh]

Thanks for opening this discussion. Does increasing the learning rate help?

Additionally, this script might be helpful for training large-scale: https://github.com/huggingface/diffusers/blob/main/examples/research_projects/controlnet/train_controlnet_webdataset.py

This is what we used for https://huggingface.co/collections/diffusers/sdxl-controlnets-64f9c35846f3f06f5abe351f

Thanks for the quick reply Sayak.

1. May I ask what's the difference between the script you linked and the one in the example folder that makes it more suitable?

2. I used 1e-5 as that's the learning rate used by your team for train the depth controlnet (https://huggingface.co/diffusers/controlnet-depth-sdxl-1.0#hyper-parameters). Therefore, would you advice me to increase the learning rate by either:
   2.1 Literally increasing the learning rate
   2.2 Scale the learning rate by batch number as your team did for the canny controlnet here
   2.3 Changing the learning rate scheduler (for example, change constant to cosine). Note that it seems your team used constant learning rate for both canny and depth controlnets.

[sayakpaul]

I would recommend taking a step back and just trying to step through the dataloader in a separate script with multiple print statements thrown in between. This way, you will be able to point out the exact codepath that needs fixing.

[brycegoh]

Thanks Sayak, I did try that before. Let me explain it

I firstly added print statements to the group_by_keys_nothrow function that yields samples as such:

The stdouts are like 0 jpg meaning prefix = 0 and suffix = jpg

The prints just keeps looping through even though for each tar file I set the maximum number of samples to be 1000 meaning a total of 4000 files given that each sample has image, control_image, original_size and text.

std out example:

My code to get the first batch:

dataset = Text2ImageDataset(
        train_shards_path_or_url="00000.tar",
        eval_shards_path_or_url=None,
        num_train_examples=1000,
        per_gpu_batch_size=1,
        global_batch_size=1,
        num_workers=1,
        resolution=1024,
        center_crop=False,
        random_flip=False,
        shuffle_buffer_size=1000,
        pin_memory=True,
        persistent_workers=True,
        control_type="canny",
        feature_extractor=None,
    )
train_dataloader = dataset.train_dataloader

print("start dataset debugging")
try:
    first_batch = next(iter(train_dataloader))
    print("Successfully fetched a batch:", first_batch)
except StopIteration:
    print("The dataloader did not yield any batches.")
except Exception as e:
    print(f"An error occurred: {e}")

I am unsure if this is expected

Alternative

Is it fair to just change the noise scheduler on the the provided script to the EulerDiscreteScheduler instead of DDPMScheduler? Or are you aware if there are more in-depth changes?

[sayakpaul]

Are you able to print the shapes of the images and other columns?

I will also recommend using a simpler data input pipeline to see if you are able to successfully load from the tars and can stream from it.

Is it fair to just change the noise scheduler on the the provided script to the EulerDiscreteScheduler instead of DDPMScheduler? Or are you aware if there are more in-depth changes?

The script uses EDM formulation. So, DDPM won't work.

[brycegoh]

Ok, we have 2 different scripts here:

A) The one provided with documentation (https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet_sdxl.py)
B) The one that your diffuser team used (https://github.com/huggingface/diffusers/blob/main/examples/research_projects/controlnet/train_controlnet_webdataset.py)

B uses the EulerDiscreteScheduler and also the EDM formulation while A uses DDPMScheduler.

Therefore, instead of adjusting my dataset to webdataset format, could I intergrate the EDM formulation portion of the script B into script A?

If that's possible, would appreciate if you could point me in the right direction such as

1. where are the essential portions of the EDM formulation in script B that differs from script A?
2. Is it purely just EDM formulation that differs?

---

I am asking point 2 as well because comparing the code, there are slight differences such as:

 ### Script A
latents = vae.encode(pixel_values).latent_dist.sample() 
latents = torch.cat(latents, dim=0)


### Script B
# encode pixel values with batch size of at most 8.                                            
latents = []                                                                                                                 
for i in range(0, pixel_values.shape[0], 8):                                                          
    latents.append(vae.encode(pixel_values[i : i + 8]).latent_dist.sample())     
latents = torch.cat(latents, dim=0)                                                                        

Apologies as I am not familiar with the theory behind this.

[brycegoh]

@sayakpaul I tried referencing your recent PR that was included in the latest v0.27.0 release, #7126

My modifications can be referenced here, brycegoh#1
I made a local PR because it is evidently not ready for a PR on the official diffusers repo without the additional test cases and README modifications.

Apart from EDM style training modifications, do we need to add the crop_top_lefts coordinates and original sizes to each batch. I saw that the dataset loader on your PR keeps track of the coordinates instead of relying on the command inputs of default 0. I am not sure if that is needed. However, it does seem useless as we already resize and then do a centre crop so top left is always 0,0?

Would appreciate some validation, checks and guidance on my modifications.