[WIP] Adding support for FP8 training

[shahromil16]

• Changes made to train.py, model.py, and norms.py to check if Transformer Engine can be imported for P5 instances or H100s and use FP8 for Linear and LayerNorm layers.
• Minor modifications to main.py for FP8 support