Fix multi-gpu bug with central_value
#191
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Denys88
reviewed
Jul 23, 2022
| if param.grad is not None: | ||
| all_grads_list.append(param.grad.view(-1)) | ||
| all_grads = torch.cat(all_grads_list) | ||
| dist.all_reduce(all_grads, op=dist.ReduceOp.SUM) |
There was a problem hiding this comment.
why do we need to work with gradients manually?
Are there any prallel optimizer like we had in horovod?
There was a problem hiding this comment.
I will merge it now but wanna look into the torch distributed module later.
There was a problem hiding this comment.
See vwxyzjn/cleanrl#162 (comment). Option 2 (implemented here) is faster than the pytorch recommended gradient averaging -
def average_gradients(model):
size = float(dist.get_world_size())
for param in model.parameters():
dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
param.grad.data /= size
There was a problem hiding this comment.
The trouble with cat is that for a bigger network, we will have OOM because it's concatenating all the grads on a single GPU i.e. we can only roughly train a model of size 2GB on 8 GPU V100 with 16GBs each.
Denys88
approved these changes
Jul 23, 2022
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR fixes multi-gpu bug w/
central_value.Tested w/
torchrun --standalone --nnodes=1 --nproc_per_node=2 train.py multi_gpu=True task=Ant task=ShadowHandOpenAI_LSTM num_envs=2048 wandb_activate=True wandb_entity=costa-huang train.params.config.minibatch_size=8192Which yields this run