Memory optimization in async tp-linear #208
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
…rentiable distributed operations
3outeille
reviewed
Jul 22, 2024
| @@ -141,22 +142,27 @@ def forward(ctx, tensor, weight, bias, group, tp_mode): | |||
| # `tensor` can sometimes not be contiguous | |||
| # https://cs.github.com/pytorch/pytorch/blob/2b267fa7f28e18ca6ea1de4201d2541a40411457/torch/distributed/nn/functional.py#L317 | |||
| tensor = tensor.contiguous() | |||
| ctx.save_for_backward(tensor, weight) | |||
| # ctx.save_for_backward(tensor, weight) | |||
xrsrke
reviewed
Aug 1, 2024
|
lgtm as well |
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 introduces the memory optimization methods implemented in #203, but now allows them to be used in the async comm regime. This PR also includes commits related to fixing the row-parallel tp-linear in #172, so both of those PRs should be merged first to make reviewing this one easier. This PR also includes two modes: recomputing or not the all_gather. Recomputing is more-memory efficient, but slightly slower.
Here is a table that summarizes my observations on a tp4 llama8b model on A100 GPUs:
These changes should prove to be very useful for more efficient training. I recommend using the async-recompute setting, but using async-norecompute might make more sense for the extra throughput when memory is not a concern. In addition, as dp and pp increase and optimizer states and parameters become more sharded, the memory savings this PR brings should only increase as these affect memory activations. Very useful for scaling to larger models.
I attach the wandb logs for llama8b (top) experiments and a tiny 152M model (bottom) to study the effects on smaller models.
(blue = sync baseline, green = sync no-recompute, purple = sync recompute, yellow = async baseline, gray = async recompute, red = async no-recompute).