stable-baselines3-contrib icon indicating copy to clipboard operation
stable-baselines3-contrib copied to clipboard

Published 20 hours ago • verified publisher icon Stable-Baselines-Team

Implement PPO MPI (SB2 PPO1)

Open araffin opened this issue 3 years ago • 2 comments

MPI can be quite useful to use multiprocessing full potential but it is dependency that can be tricky to install.

araffin avatar Nov 26 '20 10:11 araffin

Hey @araffin I prototyped multi-GPU support with torch.distributed https://github.com/vwxyzjn/cleanrl/pull/162. Preliminary experiments seem successful when controlling torch thread number to 1 per process and use SyncVecEnv:

ppo_atari_multigpu_batch_reduce.py was able to obtain 20-30% speed up at no cost of sample efficiency by leveraging data parallelism given by torch.distributed.

image

vwxyzjn avatar Apr 19 '22 20:04 vwxyzjn

@vwxyzjn ooh nice =) I didn't know you could do that with PyTorch (it looks like they included MPI but for GPU). I will try to have a look later this week. Does this work also with cpu only?

araffin avatar Apr 20 '22 08:04 araffin