issues
search
xrsrke
/
pipegoose
Large scale 4D parallelism pre-training for 🤗 transformers in Mixture of Experts *(still work in progress)*
MIT License
76
stars
17
forks
source link
End-to-end FP8 training
#45
Open
xrsrke
opened
9 months ago
xrsrke
commented
9 months ago
Notes
Write an
FP8Tensor
that inherits from torch.Tensor (just support type hints).
Write an
FP8Linear
that binds to TransformerEngine's FP8 kernel in the forward pass
TODO
[ ]
nn.Linear
but in FP8
[ ] Recursively convert all
nn.Linear
to FP8 Linear
[ ]
nn.Embedding
in FP8
[ ] DP in FP8
[ ] TP in FP8
[ ] MoE in FP8
[ ] ZeRO-1
[x] PP in FP8 (you get it for free)
3outeille
commented
9 months ago
@xrsrke On it
Notes
FP8Tensor
that inherits from torch.Tensor (just support type hints).FP8Linear
that binds to TransformerEngine's FP8 kernel in the forward passTODO
nn.Linear
but in FP8nn.Linear
to FP8 Linearnn.Embedding
in FP8