[BUG] train_text_to_image_lora.py not support Multi-nodes or Multi-gpus training.

[WindVChen]

In train_text_to_image_lora.py, I notice that the LORA parameters are extracted into an AttnProcsLayers class:

518    lora_layers = AttnProcsLayers(unet.attn_processors)

And it is only the lora_layers that is wrapped by DistributedDataParallel in the following code:

670    lora_layers, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
            lora_layers, optimizer, train_dataloader, lr_scheduler
          )

In the training process, it seems that the lora_layers are not explicitly used but only the unet is used:

776    model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample

My question is that when using Multi-GPUs or Multi-Machines, will the gradients be successfully averaged across all processes in the above way?

It is true that in each process, the gradients will be backward to unet.attn_processors, and these gradients will be shared by lora_layers, so we can use optimizer to update the weights. However, since we actually use unet.attn_processors to do the forward operation, but not the wrapped lora_layers, can the gradients be correctly averaged? From here, it seems that a wrapped module will have a different forward compared to its original forward operation.

I am not quite familiar with torch.nn.parallel.DistributedDataParallel wrapper, and I do worry about whether the current code in train_text_to_image_lora.py will lead to different LORA weights in different processes (if the gradients failed to broadcast among processes).

Hope to find some help here, thank you.

[WindVChen]

After carefully printing out the gradients and weights in different processes, it seems quite sure that the current LORA training script fails to be applied to Multi-nodes or Multi-GPUs training: The gradients failed to broadcast among the processes, which then lead to different LORA weights in different processes.

In the following, I give a printout when I start a 2-GPU task:

The first 3 number denotes process/epoch/step, and the following two tensors are the gradients and weights of lora_layers.layers[31].to_out_lora.up (A linear layer). It is obvious that the gradients are different in different processes. Thus, after three steps, it will turn to:

Obviously, the weights of LORA are different in different processes.

I think it is a bug since there is no description of that in the script. People who use the script to train on multi-processes will finally and actually get the result on a single process.

[patrickvonplaten]

cc @williamberman @sayakpaul

[sayakpaul]

@WindVChen

Thanks for elaborating on this!

From gauging this briefly, it seems like passing unet to the prepare step rather than just the LoRA layers might just fix this.

Could you confirm this once?

[WindVChen]

Hi @sayakpaul ,

Yes. Passing unet to the prepare step can quickly fix it. However, it also can bring some inconvenience: 1) The batch size has to be reduced by half, or the training will lead to out-of-memory; 2) The checkpoint stored intermediate is much larger (~3G compared to the original ~3M), as it needs to store the whole unet structure.

I wonder if there is a more elegant way to solve it without sacrificing memory and storage. I also tried wrapping each "loraAttnprocessor" manually, but found it required a lot of source code modification, so I gave up 😢.

[sayakpaul]

Thanks for sharing!

I also tried wrapping each "loraAttnprocessor" manually, but found it required a lot of source code modification, so I gave up

Could you expand a bit more on this?

[WindVChen]

Yes.

It seems that the problem above happened just because the training script used the unwrapped unet to do the forward, not the wrapped lora_layers, so my previous idea is to wrap every loraAttnprocessor in the unet one by one (directly use accelerate.prepare) and then replace the original unwrapped loraAttnprocessors in unet with these wrapped loraAttnprocessors. In this way, maybe we can expect that in the training phase, the gradients can broadcast among different processes.

However, since I'm not quite familiar with torch.nn.parallel.DistributedDataParallel, I'm not sure whether this solution can work.

[sayakpaul]

Hmm. I will defer to what @williamberman has to point out here. Also ccing @muellerzr from the accelerate team here.

[muellerzr]

When using Accelerate any models that get gradient/weight updates should be passed to .prepare. Do note that gradient accumulation currently won't work with multiple models, we're adding that very soon, but hope that helps.

[WindVChen]

Hi @muellerzr ,

Could you expand a bit more about "gradient accumulation currently won't work with multiple models"? Or are there any blogs or issues related to this? Because I am going to use gradient accumulation for multiple models.:sweat_smile:

[muellerzr]

No one has quite pointed out this "issue" yet actually, once the PR is opened (probably today or tomorrow) I'll link to it but basically multiple forward passes followed by multiple backward passes (for each model's loss, for instance) leads to some headaches in torch distributed.

[WindVChen]

OK, thanks.

Hope to be sure. So does that mean that a GAN training script like the one below will cause some BUGs?

"Suppose gradient_accumulation is set to 2"
optimizer_gen = optim(generator.parameters())
optimizer_disc = optim(discriminator.parameters())
with accelerator.accumulate(generator):
        outputs = optimizer_gen(input)
        loss = loss_func(outputs)
        loss.backward()
        optimizer_gen.step()
        optimizer_gen.zero_grad()

        outputs = optimizer_disc(input)
        loss = loss_func(outputs)
        loss.backward()
        optimizer_disc.step()
        optimizer_disc.zero_grad()

[muellerzr]

Yes because technically the discriminator should also be under accumulate. So the discriminator might still be getting updated each step. (Though again, working on support there)

[WindVChen]

Ah, I see. Maybe I can double-wrap the pipeline for a temporary fix before the PR? Like this:

with accelerator.accumulate(generator):
    with accelerator.accumulate(discriminator):
        outputs = optimizer_gen(input)
        ...