transformer_engine/jax/cpp_extensions.py · 387397a2afeabfe52de019ea9f03f46be3f47f2d · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Support layernorm/rmsnorm sm_margin control through environment variable (#520) · 753eed31
zlsh80826 authored Nov 30, 2023
```
Support layernorm sm_margin through environment variables
Signed-off-by: Reese Wang <rewang@nvidia.com>
```
753eed31

cpp_extensions.py 146 KB