transformer_engine/pytorch/distributed.py · fe80ca0615e2b629b914f0d7dc34342ac2aa4ea7 · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[PyTorch] Fixed assert on primary Fp8 weights in `prepare_te_modules_for_fsdp()` (#916) · fe80ca06
Alp Dener authored Jun 12, 2024
```
restricted fsdp asserts on primary fp8 weights to TE modules
Signed-off-by: Alp Dener <adener@nvidia.com>
```
fe80ca06

distributed.py 36 KB