Skip to content

Conversation

@awaelchli
Copy link
Contributor

When using ddp spawn, moving the model to gpu before processes are spawn creates extra memory usage on the root gpu.
We already move the model to device in the plugins, so this shouldn't be necessary in setup.

torch.cuda.set_device(self.root_device)
model.to(self.root_device)
return super().setup(trainer, model)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just wanted to double check here with @tchaton and @justusschock
is it good to remove this? any implications for rpc/sequential?
For the other plugins I have not seen anything breaking.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is fine

@justusschock
Copy link
Member

I'll close this in favor of #5866

@tchaton tchaton deleted the refactor/memory branch February 8, 2021 17:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants