Vote on new features in Discussions

[tianyu-l]

Hi torchtitanists,

Thank you for your interests in torchtitan!

We created #693 for the community to add feature requests and vote on them. We'll try to prioritize on the most requested features. Please share what you'd like to see next!

[tianyu-l]

cc: @samsja @qsh-zh @Amerehei @FilyaGeikyan @thepowerfuldeez @Borda @starstream @leeeizhang @abatilo @mayank31398 @vermouth1992 @SeunghyunSEO @aniltrkkn @jiagaoxiang @casper-hansen @carmocca @daviswer @zigzagcai @jeromeku @Yangyi-Chen @tangjiasheng @zijian-hu @crcrpar @Doraemonzzz @jvwilliams23 @152334H @yzs981130 @MinghaoYan @rlrs @philippguevorguian

[tianyu-l]

cc: @LeoXinhaoLee @mathmax12 @yundai424 @vedantroy @dz1iang @goldhuang @evkogs @0781532 @jaysonfrancis @nighting0le01 @vighneshbirodkar @zyushun @medivh-xp @Yuxin-CV @qiziAI @ad8e @prathameshtd @bkchang @arjvik @jaywonchung @bhuvan777 @ahatamiz @dmammfl @TJ-Solergibert @karan-dalal @gau-nernst @lucasjinreal @qiongerfei @rshah918 @eminorhan @thomasneff @viai957 @pansershrek

[zigzagcai]

@tianyu-l

Hi developers,

Firstly, thanks for the great work that can demonstrate the power of PyTorch newly released features!

I just have one confusion about the usage of FSDP2 fully_shard.
Does FSDP2 support mixed precision within one warpping module, such like torch.float32 and torch.bfloat16 within a FSDPParamGroup?

To put it more clear, in most use cases of training LLM such like Lllama2, the precision of RMSNorm is usually torch.float32, but other components within the DecoderLayer is usually torch.bfloat16. When we want to train the model with the help of FSDP2, we have to wrap RMSNorm seperately since it has a seperate dtype, which will introduce additional all-gather and reduce-scatter.

From the profiling results, we found this approach (warpping RMSNorm seperately) will lead to bad computation-communication overlapping, especially in the backward pass.

Apart from that, there are also some other use cases: dtype of MoE gating layers is required to be torch.float32, but other components in the DecoderLayer is torch.bfloat16. We can also found that seperately warpping MoE.GateLayer would cause bad overlapping of computation-communication.

So, does mixed precision within a FSDPParamGroup supported? or could this be a new feature in the future?

Thanks!

[mayank31398]

@zigzagcai RMSNorm only has activations in fp32, the weights are still bf16.
Also, FSDP2 is quite flexible in having different dtypes for different tensors I believe

[tianyu-l]

cc: @awgu

[aniltrkkn]

it should be simple but

Gradient Accumulation

it is very useful for sfting big models.