microsoft / tutel

Tutel MoE: An Optimized Mixture-of-Experts Implementation

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Tutel with pytorch automatic mixed precision package

MiZhenxing opened this issue · comments

Hi, thank you for your excellent package. I wonder whether tutel could be used seamless together with the automatic mixed precision package of PyTorch. If so, could you provide some hints on how to use it or provide an example python script?
Thank you very much!

Yes, Tutel is able to support native Pytorch AMP. Please follow this example: https://github.com/microsoft/tutel/blob/main/tutel/examples/helloworld_amp.py#L76

Note that you need to use @autocast() and with autocast: properly according to pytorch's doc.
Welcome to report problems if you encounter any issue.

Thanks for your quick response. Sorry I didn't notice that example 🤣. I will do some testing on my code.