search

NVIDIA / modulus

Open-source deep-learning framework for building, training, and fine-tuning deep learning models using state-of-the-art Physics-ML methods
https://developer.nvidia.com/modulus
Apache License 2.0
941 stars 220 forks source link

GraphCast: Add the option to use transformer_engine.LayerNorm #516

Closed mnabian closed 4 months ago

mnabian commented 4 months ago

Note from Max: Can we add the option to use transformer_engine.LayerNorm? In AIFS benchmarks, we just could get a 1.3x end-to-end improvement from doing so since the PyTorch implementation is rather bad for the sizes we encounter in these workloads.

mnabian commented 4 months ago

https://github.com/NVIDIA/modulus/pull/510