transformer_engine/jax/checkpoint_policies.py · 858755c0c3898d5bc971fcec031b26ff26e7e4fb · OpenDAS / TransformerEngine

[JAX] TE GEMM checkpointing policies (#2003) · 858755c0

jberchtold-nvidia authored Jul 30, 2025



* TE primitive checkpointing policies
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Remove batched gemm policy
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

---------
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

858755c0

checkpoint_policies.py 1.23 KB

Replace checkpoint_policies.py