Fp8 all gather hack #1136
Open
Fp8 all gather hack #1136
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
jspark1105
force-pushed
the
fp8_all_gather
branch
from
dfe122b
to
fe3a0d6
Compare
jspark1105
force-pushed
the
fp8_all_gather
branch
from
fe3a0d6
to
7dd000e
Compare
jspark1105
force-pushed
the
fp8_all_gather
branch
from
7dd000e
to
0224797
Compare
jspark1105
changed the base branch from
main
to
ngoyal_added_zero2_shard_modelparams_multiple_gpus
|
Will merge main_grad related changes with #1142 |
jspark1105
force-pushed
the
fp8_all_gather
branch
2 times, most recently
from
bd70153
to
af3d2d7
Compare
jianyuh
reviewed
Oct 6, 2023
| # Cast grad to FP32. | ||
| grad_or_main_grad.data = grad_or_main_grad.to(param.dtype) | ||
| elif self._is_fp8_dtype(): | ||
| # Use bf16 wgrad for fp8 weights (TODO: handle fp8 wgrad) |
There was a problem hiding this comment.
Currently this is not working with the latest FP8 wgrad ?
There was a problem hiding this comment.
This meant to be for future work when we have fp8 reduce-scatter. I'll update the comment.
| @@ -1393,7 +1447,11 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: | |||
|
|
|||
| # For root and mixed precision, we convert the input to FP16 (no_grad is needed for | |||
| # the conversion). | |||
| is_bf16 = self.compute_dtype == torch.bfloat16 | |||
| is_bf16 = self.compute_dtype in [ | |||
jspark1105
force-pushed
the
fp8_all_gather
branch
2 times, most recently
from
b9b093b
to
a2b49d1
Compare
jspark1105
commented
Oct 7, 2023
| @@ -2265,8 +2361,7 @@ def local_metadata_dict(self) -> Dict[str, Any]: | |||
| backing_param_name = m.module.flat_param_names[i] | |||
| names, shapes, numels = m.module.metadata(i) | |||
| else: | |||
| assert len(m._param_name_groups[i]) == 1 | |||
| backing_param_name = m._param_name_groups[i][0] | |||
| backing_param_name = m._param_name_groups[m._num_flatten_params][i - m._num_flatten_params] | |||
There was a problem hiding this comment.
Need to make sure checkpointing works properly with this.
jspark1105
force-pushed
the
fp8_all_gather
branch
2 times, most recently
from
d92dc0f
to
6a4d7f4
Compare
… during post_backward_hook
jspark1105
force-pushed
the
fp8_all_gather
branch
from
6a4d7f4
to
c0f4b97
Compare
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 is based on ngoyal_added_zero2_shard_modelparams_multiple_gpus and adding hacks to use fp8 all-gather with Nvidia's transformer engine (see the latest commit for the changes on top of ngoyal_added_zero2_shard_modelparams_multiple_gpus branch).
This depends on transformer engine changes in https://github.com/facebookresearch/TransformerEngine/pull/20
See https://github.com/fairinternal/xlformers/pull/1403 for an example how to use.
Also depends on PyTorch changes in pytorch/pytorch#109654
To use fp8 allgather, set compute_dtype=torch.float8_e4m3fn and mixed_precision=True
We separate out precision critical parameters like affine weights for norm as non flattened params and hard-code to use bf16.
We update scale/scale_inv inside forward before _rebuild_full_params that calls _cast_fp32_param_shards_to_fp16 vs. TE baseline that updates scale/scale_inv in prepare_forward. This because we need fp8 quantization of weights earlier before allgather. (One can consider doing this in post backward but this has a problem since updating bwd amax update is done after bwd of all layers are finished which can be later than post backward so we won't be using the latest bwd amax info for scale/scale_inv update).
We hard-code special handling for a couple of TransformerEngine layers like Linear, LayerNormLinear, and LayerNormMLP in _cast_fp32_param_shards_to_fp16 to access their fp8 meta data to quantize with right scales (TODO: we may want to extract this as a user customizable call back functions?)
CC @awgu @ngoyal2707 @vedanuj @jiecaoyu @yf225 @GD06