Weights dont sychronise in multi GPU training

distributed
1

Hi all,

I have been trying to train a GNN in multiple gpus. I use the standard python libraries, PyG and DDP. I can see the data distributed in the GPUs, and each of them training a model. However the weights are not sychronized. Before the training all the weights are exactly the same, but after one epoch they are different. I feel that the problem comes from :

“[rank0]: RuntimeError: Expected to mark a variable ready only once. This error is caused by one of the following reasons: 1) Use of a module parameter outside the forward function. Please make sure model parameters are not shared across multiple concurrent forward-backward passes. or try to use _set_static_graph() as a workaround if this module graph does not change during training loop.2) Reused parameters in multiple reentrant backward passes. For example, if you use multiple checkpoint functions to wrap the same part of your model, it would result in the same set of parameters been used by different reentrant backward passes multiple times, and hence marking a variable ready multiple times. DDP does not support such use cases in default. You can try to use _set_static_graph() as a workaround if your module graph does not change over iterations.”

which i masked/bypassed (?) by setting set_static_graph(). Does anyone have any ideas what is the problem and what I should do?

thank you, Panagiotis

2

Very interesting, that just as i experience something very similar, i hit this post. Indeed, i see a similar phenomenon here, but without this RuntimeError. It is the same with pytorch 2.8.0 and 2.11.0.

We use a model implementation that extends and modifies TCN/TCN/tcn.py at master · locuslab/TCN · GitHub and define our own loss functions.

Before attempting to use pytorch.distributed it just worked either on GPUs or distributed on up to 96 CPUs with horovod. Results are not identical, but very similar and of comparable quality, no matter if i train distributed or not or on CPU or on GPU.

Now, that horovod is considered insecure, cause its communication uses pickle to (de)serialize data, and pytorch > 2.8 does not support horovod anymore, we are moving away. It is basically running with pytorch.distributed, but the results are significantly worse. Tracking down the issue i found, that the weights are not averaged across the ranks during training. According to explanations in the pytorch documentation (if i understand them correctly) this is intended. Initially they are the same in all ranks.

In contrast, horovod does this averaging with each optimizer.step() with good training results as outlined above. torch.distributed does not and has bad results. So i added my own averaging of the weights after the optimizer.step() and found, that training convergence and the results are much better and in the expected range.

There is another “but”: I have only good results with torch.distributed, when i use RelU in the network. When i switch to Softsign, what we found to yield a somewhat more stable training with our data, the prediction results are worse.

Can someone explain, why torch.distributed does not synchronize the weights or what might prevent it from doing so ? And why might only this partially linear activation function RelU be working well with pytorch.distributed ? Might it have to do with the way, that torch.distributed handles the gradients ?

Thank you !

3

It would be wasteful to sync the parameters in every step. Instead the initial parameters are synchronized before the training starts. In each iteration the gradients are then synchronized and the optimizer.step() method updates the parameters on each rank individually. Since the initial weights were identical and the gradients are synchronized, the updated weights are equal on each rank again. This tutorial explains this in mode details.

@pbellos are you using any module outside of the forward method? If so, the gradient sync might be skipped resulting in this issue.

4

Thank you, i went through this tutorial and other documentation several times and don’t find a clue, why what you state is obviously not the case with our program and when i synchronize the weights myself, i get different and especially better results. I printed out a checksum over all weights in every step in all ranks and could see, that they differ. With horovod, they don’t

5

The linked doc describes a general overview of DDP and also links to the internal design which exactly matches my description explaining why a weight synchronization is not needed and in fact wasteful since the parameters will be updated from the same initial state using the same gradients and optimizer. Let me know if you have any questions after reading the internal design and which part is unclear.

With that said, your issue could be caused by manually calling modules outside the forward as mentioned before.

6

Hmm I think I missed the last line from the error message, it looks quite suspicious

“[rank1]: Parameter at index 116 with name 19.fn.linear.bias has been marked as ready twice. This means that multiple autograd engine hooks have fired for this particular parameter during this iteration.”

This shouldnt happen if the parameter is called only inside forward, right? Could it cause the weight sychronization problem, if covered by set_static_graph() ?

7

Hello,

so at least for our case i have tracked down the issue and think, this should be documented somewhere. At least i haven’t found a hint to this:

Someone of earlier developers of our entire package had decided to add a method, that computes at least a part of a loss function, to the model, that is derived from nn.Module . This method calls self.forward(…) . What does NOT happen during the run of this method is - whyever - that grad is applied, as if requires_grad was set to False during methods of the model itself. This way, the gradients are not in sync in the different ranks. I’ve put this function out of the class and wrote and call it just as a function. Now the grads are equal in the different ranks and backward behaves as expected and torch.distributed can be applied.

Please mention somewhere: Do not implement the loss function or subordinate functionality as method of the model.

8

It was the unused parameters that were breaking the synchronization. I slightly changed the architecture, so that all of them are now used, and everything works fine