Closed yaox12 closed 1 month ago
/te-ci pytorch
/te-ci pytorch
@timmoon10 Can I get this PR merged if it looks good to you? It's sufficient to be a drop-in replacement of Apex.FusedAdam in MCore. Other features to be migrated or implemented will come in separate PRs.
Description
This PR moves
FusedAdam
/FusedSGD
and necessarymulti_tensor_*
kernels from Apex to TE, as well as the tests.Type of change
Changes
Please list the changes introduced in this PR:
multi_tensor_*
kernelsFusedAdam
FusedSGD
Checklist: