Enable DDP for Off-Policy training branch #1099
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
emailweixu
reviewed
Dec 2, 2021
Base automatically changed from
PR/breakds/data_distributed_decorator
to
pytorch
December 2, 2021 22:08
breakds
force-pushed
the
PR/breakds/data_distrbuted_off_policy
branch
from
eecfb67
to
58496fa
Compare
|
Addressed comments. PTAL. |
emailweixu
reviewed
Dec 3, 2021
breakds
force-pushed
the
PR/breakds/data_distrbuted_off_policy
branch
from
ef9064d
to
e943e5c
Compare
emailweixu
approved these changes
Dec 4, 2021
pd-perry
pushed a commit
to pd-perry/alf
that referenced
this pull request
Dec 11, 2021
* DDP for off-policy training branch * More reliable way to ignore replay buffer's buffers * Use set instead of list
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.
Motivation
Previously DDP (data distributed parallel) is only enabled for the on-policy training branch (i.e.
train_from_unroll()). However, most of the algorithms actually runs on the off-policy training branch (i.e.train_from_replay_buffer()) which does not enjoy DDP yet.Solution
After #1098 we are ready to make this work for non-composite off-policy branch training, (e.g. standard PPO). According to this, we need wrap the computation that involves all trainable parameters. Therefore, the change is mainly:
train_step()(i.e. thecollection_train_info...()) and computation of loss into a standalone method ofAlgorithm, called_compute_train_info_and_loss_info._compute_train_info_and_loss_infowith@data_distributed.Also, a special treatment is done to ignore replay buffer from participating DDP's synchronization.
Testing
Tested with
(Note that
mini_batch_sizeis halved)Still under training but the throughput is almost doubled (13000 vs 7000) on a machine with two 3080s.
Also verified on
ppo_cart_poleas a small experiment, see here.