[FSDP] use all_gather for 10X OSD consolidation speedup #595
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
facebook-github-bot
added
the
CLA Signed
sshleifer
changed the title
sshleifer
commented
Apr 8, 2021
|
|
||
| return non_tensor_state, tensor_state | ||
|
|
||
| def _gather_optim_state(self, sd_state: Dict[int, Dict[str, Any]]) -> Dict[int, Dict[str, List]]: |
There was a problem hiding this comment.
This is the new _all_gather logic
sshleifer
changed the title
sshleifer
commented
Apr 9, 2021
| @@ -627,15 +627,18 @@ def __init__(self, group, wrapper_config, checkpoint_act=False, delay_before_fre | |||
|
|
|||
| # "expert" params are different on each rank | |||
| torch.manual_seed(42 + group.rank()) | |||
| d_expert = 16 | |||
| expert = nn.Linear(d_expert, 4) | |||
| d_expert = 23 | |||
There was a problem hiding this comment.
make sure we unpad expert params correctly.
|
Planning to merge 10am PT tomorrow, barring further comments. |
min-xu-ai
approved these changes
Apr 13, 2021
| """Return the last known global optimizer state. The returned state is compatible with Pytorch, in that the | ||
| sharded properties are not exposed. Multiple parameter groups are not yet supported. | ||
|
|
||
| This should be called only on the root FSDP instance. | ||
| Nested FSDP instances are supported as long as they have the same world_size as the parent or world_size=1. |
|
cc @QuentinDuval @prigoyal FYI |
sorry for the delay, I was out for most of the last week. |
myleott
reviewed
Apr 13, 2021
myleott
reviewed
Apr 13, 2021
| new_sd = {"state": new_state, "param_groups": copy.deepcopy(sd["param_groups"])} | ||
| for k in sd.keys(): # if there are extra keys, like loss_scale, don't delete them | ||
| if k not in {"state", "param_groups", "uncollected_local_ids", "param_id_map"}: |
There was a problem hiding this comment.
I'm slightly uneasy about this falling out of sync with line 160. Thoughts on some way to enforce parity between them?
There was a problem hiding this comment.
module level constant + comment + assert
# These return keys are used by fairseq. To change, add @sshleifer as a reviewer.
UNFLAT_RETURN_KEYS = {"state", "param_groups", "uncollected_local_ids", "param_id_map"}
...
assert set(unflat_optim_state_dict.keys()) == UNFLAT_RETURN_KEYS
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.
TLDR: Using
all_gatherinstead ofbroadcastfor optimizer state consolidation appears to be a speed win without a memory cost.Approach:
all_gatherloss_scale,param_groups,num_padded) we usebroadcast.recipient_rank=0, since there are no other callers.Evidence of Win: