You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
I was wondering if PyTorch's FullyShardedDataParallel (FSDP) is supported by TransformerEngine , especially if FP8 can work with FSDP. Thank you in advance.
In the current scheme, Transformer Engine modules use standard parameter tensors in standard dtypes (FP32/BF16/FP16). Optimizers typically require higher precision than FP8 to get good learning behavior. I don't see anything that would disrupt FSDP and I've been able to get TE and FSDP working together in some quick experiments.
What is the recommendation for MixedPrecision when using FP8 with FSDP @timmoon10 ?
fromtorch.distributed.fsdpimportMixedPrecisionprecision= ?
MixedPrecision(
param_dtype=precision, # should we be forcing dtype here?reduce_dtype=torch.float32, # reduce in FP32 as with AMP?buffer_dtype=precision, # buffers in FP32?cast_forward_inputs=is_amp_enabled, # should we cast the forward samples?
)
Transformer Engine manages FP8 casting internally (see transformer_engine.pytorch.fp8_autocast) and it can run into problems when combined with other mixed precision tools like torch.autocast or torch.distributed.fsdp.MixedPrecision. For the moment, FSDP mixed precision should only be used to handle the case where param_dtype and reduce_dtype do not match, e.g.:
Edit: On second thought, I'm skeptical about this post. I don't remember my main concerns and I've been able to have fp8_autocast working with torch.autocast and torch.distributed.fsdp.MixedPrecision. See #438 (comment).
@timmoon10 I think the community would greatly benefit from an e2e example on how to train a 7B model on H100 node(s). This is the main purpose of this lib - to train large models quickly. Yet there is not a single example on how to do it - instead we have a MNIST example...
Ideally, a simple pure PyTorch+TE file without any "framework" dependencies like compose or fabric.
Reacted by Jason Ramapuram, Shuo Qiu, c0g, naveen marri, praveenkanithi, Aleksandr Dremov and Kion Fallah
@yongyanrao#596 recently added support for deferred initialization via device='meta' to improve FSDP support for large models. This feature delays memory allocation on device until the FSDP wrap internally calls reset_parameters() after sharding the model weights, which ensures that TE modules are not duplicated on every device upon initialization.
We also introduced a small FSDP example in the repo as part of this effort to demonstrate how to use FSDP with TE modules. It works out of the box with fp8_autocast(), with the requirement that fp8_autocast(...,fp8_group=...) gets the same process group as FSDP (both will default to the world group if none given).
Reacted by Jason Ramapuram, Gautam Kumar and Victor Zhu
@denera Thank you for the feature! This will be super helpful for us too. Just a question, are there any additional steps needed to take to make primary_weights_in_fp8 work with FSDP? Or should it work out of the box?
NCCL itself does not support 8-bit floats (see this discussion for more detail) and FSDP needs to upcast TE Fp8 weights to Fp16 for all-gathers and reduce-scatters, which it cannot do until PyTorch starts natively supporting Fp8 tensors.
The workaround for this limitation is what already happens in TE when primary_weights_in_fp8 = False. TE modules maintain their own Fp8 weight copies, update them with the primary Fp16/Bf16 weights during the forward pass, and stash the Fp8 transposed-weights into the PyTorch autograd context to re-use it during the backward pass.
Activity
timmoon10 commentedon Aug 25, 2023
In the current scheme, Transformer Engine modules use standard parameter tensors in standard dtypes (FP32/BF16/FP16). Optimizers typically require higher precision than FP8 to get good learning behavior. I don't see anything that would disrupt FSDP and I've been able to get TE and FSDP working together in some quick experiments.
yongyanrao commentedon Aug 29, 2023
Hi Tim, thank you for the response. Did you try FP8 + FSDP yet?
MatthieuToulemont commentedon Aug 30, 2023
Hello, @timmoon10 which FSDP did you use ? Fairscale's ?
yongyanrao commentedon Aug 30, 2023
Hi, I was referring to pytorch's FullyShardedDataParallel.
MatthieuToulemont commentedon Aug 30, 2023
Thank you !
timmoon10 commentedon Aug 30, 2023
Yep, I used PyTorch FSDP with TE FP8. Be advised I haven't done full convergence experiments, just some basic sanity checking.
jramapuram commentedon Sep 1, 2023
What is the recommendation for MixedPrecision when using FP8 with FSDP @timmoon10 ?
timmoon10 commentedon Sep 1, 2023
Transformer Engine manages FP8 casting internally (see
transformer_engine.pytorch.fp8_autocast) and it can run into problems when combined with other mixed precision tools liketorch.autocastortorch.distributed.fsdp.MixedPrecision. For the moment, FSDP mixed precision should only be used to handle the case whereparam_dtypeandreduce_dtypedo not match, e.g.:Edit: On second thought, I'm skeptical about this post. I don't remember my main concerns and I've been able to have
fp8_autocastworking withtorch.autocastandtorch.distributed.fsdp.MixedPrecision. See #438 (comment).PiotrDabkowski commentedon Sep 2, 2023
@timmoon10 I think the community would greatly benefit from an e2e example on how to train a 7B model on H100 node(s). This is the main purpose of this lib - to train large models quickly. Yet there is not a single example on how to do it - instead we have a MNIST example...
Ideally, a simple pure PyTorch+TE file without any "framework" dependencies like compose or fabric.
ksivaman commentedon Sep 5, 2023
@PiotrDabkowski, please take a look at NeMo, which provides the exact scripts and tools to be able to do this.
jramapuram commentedon Sep 16, 2023
@ksivaman : AFAIK NeMo uses tensor parallelism which requires manual changes to layers and arch whereas FSDP is pretty generic and more easy to use.
naveenkumarmarri commentedon Oct 19, 2023
@jramapuram were you able to train FSDP with FP8?
denera commentedon Jan 24, 2024
@yongyanrao #596 recently added support for deferred initialization via
device='meta'to improve FSDP support for large models. This feature delays memory allocation on device until the FSDP wrap internally callsreset_parameters()after sharding the model weights, which ensures that TE modules are not duplicated on every device upon initialization.We also introduced a small FSDP example in the repo as part of this effort to demonstrate how to use FSDP with TE modules. It works out of the box with
fp8_autocast(), with the requirement thatfp8_autocast(...,fp8_group=...)gets the same process group as FSDP (both will default to the world group if none given).denizokt commentedon Jan 30, 2024
@denera Thank you for the feature! This will be super helpful for us too. Just a question, are there any additional steps needed to take to make primary_weights_in_fp8 work with FSDP? Or should it work out of the box?
Thank you!
denera commentedon Jan 30, 2024
@denizokt
fp8_model_init()is not supported with FSDP at the moment.NCCL itself does not support 8-bit floats (see this discussion for more detail) and FSDP needs to upcast TE Fp8 weights to Fp16 for all-gathers and reduce-scatters, which it cannot do until PyTorch starts natively supporting Fp8 tensors.
The workaround for this limitation is what already happens in TE when
primary_weights_in_fp8 = False. TE modules maintain their own Fp8 weight copies, update them with the primary Fp16/Bf16 weights during the forward pass, and stash the Fp8 transposed-weights into the PyTorch autograd context to re-use it during the backward pass.denizokt commentedon Jan 30, 2024
@denera Thank you for the information, this makes a lot of sense.