Skip to content

GitLab

You need to sign in or sign up before continuing.
Switch branch/tag
Find file BlameHistoryPermalink
  • Yifan Xiong's avatar
    Benchmarks - Support FP8 in BERT models (#446) · 5197cdf5
    Yifan Xiong authored 
    Support FP8 in PyTorch BERT models:

* add fp8 hybrid/e4m3/e5m2 in precision arguments
* build BERT encoders with `te.TransformerLayer` to repalce
`transformers.BertModel`
* wrap forward steps with fp8 autocast
    5197cdf5
pytorch_bert.py 11.3 KB