metaseq
metaseq copied to clipboard
Add GPU timers
The Megatron codebase has timers scattered all over portions of their code (i.e. https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/optimizer/optimizer.py#L412). We should add similar timers to see if we can find areas of improvement.
- Per @anj-s: these should be gated behind a debug flag in case there's a perf hit for measurement / to reduce log spam.