• 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
context.py 2.94 KB