add PrepareFloat8ModuleInput for sequence parallel #275
Conversation
when applying Sequence Parallel to a module with more than 2 linear layers for input proj, we often want to transform from Shard to Replicate once (allgather once) and then reuse the allgathered result, for fp8 we would need to do the casting before the shard -> replicate so that we can perform the fp8 allgather. This PR subclasses the PrepareModuleInput to add the fp8 casting logic to make sure we run the fp8 allgather instead of bf16 allgather then do the casting for computation. Also adjust the test cases to test the real ffn case for sequence parallel
facebook-github-bot
added
the
CLA Signed
|
@wanchaol has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
| @@ -109,3 +114,94 @@ def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: | |||
| ) | |||
|
|
|||
| return super()._apply(module, device_mesh) | |||
|
|
|||
|
|
|||
| class PrepareFloat8ModuleInput(PrepareModuleInput): | |||
There was a problem hiding this comment.
nit: maybe we can have e4m3 in the name, and maybe add a TODO to support the AMD version of e4m3 eventually?
There was a problem hiding this comment.
maybe a quick docblock to explain that this is ensuring the float8 cast happens before the all-gather if there are multiple float8 users of the input activation?
There was a problem hiding this comment.
maybe we can have e4m3 in the name, and maybe add a TODO to support the AMD version of e4m3 eventually
I wonder what's your thought on these two choice: 1. make e4m3 appears in the name of this class 2. make this class constructor take an additional argument of fp8 dtype, i.e. float8_dtype=torch.float8_e4m3fn, and we default to this e4m3fn dtype, and then later we can add on the AMD version of e4m3 by passing a different float8_dtype` arg?
There was a problem hiding this comment.
make this class constructor take an additional argument of fp8 dtype
sgtm
|
|
||
| # search for ScaledMM configs for all the submodules and make sure they are the same | ||
| fwd_linear_config = None | ||
| for mod in module.modules(): |
There was a problem hiding this comment.
WDYT something like the following to avoid the logic below?
PrepareFloat8ModuleInputtakes a ScaledMMConfig constructor argumentFloat8DynamicLinearhas logic where if the input is already aFloat8Tensor, there is a check to verify the config matches
?
There was a problem hiding this comment.
I thought about this option too, the concern I have on this is that, this would make the API diverges from the TP API offered in core, so making the switch between fp8 and bf16 be harder.
Also I think user would need to know how to construct the ScaledMMConfig, this basically make ScaledMMConfig be a public facing API. I wasn't sure this is sth we want or not?
There was a problem hiding this comment.
this basically make ScaledMMConfig be a public facing API
Yeah, good point, that's not intended to be a user facing thing. How about something like requiring a name of the module to get the config from?
I think the user API of the current code is great (no extra args), but the restriction that all configs in the module need the same config is not ideal. If we are ok with changing that later, current API sgtm.
There was a problem hiding this comment.
I think this make sense! Let me draft up the changes for accept the module_fqn to get the scaled mm config from. My current thinking on how we could approach this:
- We add a
fwd_config_module_fqnarg to the constructor so that user can specify which module config to take from - This arg could be optional, where if user don't pass it in, we still do the search and restrict all configs in this specific module should all be the same.
|
@wanchaol has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
This PR is a follow up PR to enable fp8 allgather in TP after these PR landed: * pytorch/pytorch#128431 * pytorch-labs/float8_experimental#275 One need to update their pytorch/float8_experimental to have those changes in to train with fp8 changes. Since fp8 is not enabled as part of our integration tests yet, there should be no issues on CI
This PR is a follow up PR to enable fp8 allgather in TP after these PR landed: * pytorch/pytorch#128431 * pytorch-labs/float8_experimental#275 One need to update their pytorch/float8_experimental to have those changes in to train with fp8 changes. Since fp8 is not enabled as part of our integration tests yet, there should be no issues on CI or trains that does not use fp8
when applying Sequence Parallel to a module with more than 2 linear layers for input proj, we often want to transform from Shard to Replicate once (allgather once) and then reuse the allgathered result, for fp8 we would need to do the casting before the shard -> replicate so that we can perform the fp8 allgather.
This PR subclasses the PrepareModuleInput to add the fp8 casting logic to make sure we run the fp8 allgather instead of bf16 allgather then do the casting for computation.
Also adjust the test cases to test the real ffn case for sequence parallel
torchtitan perf benchmarks (8 H100 devgpu, Llama3 8b, 2-way DP, 4-way TP):
So even in eager we got around 20% perf improvement with every allgather runs in fp8, and compiled fp8 allgather perf is more than doubled (102% more WPS) :)