[ENHANCEMENT] MoE support in report_theoretical_memory
Open
community-request
enhancement
module: moe
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
Are there any plans to add support for theoretical memory calculations for MoE in
`compute_activation_memory` (located in `Megatron-LM/megatron/training/theoretical_memory_usage.py`)?
If not, I’d be happy to work on an implementation and open a PR.
Contributor guide
Assessment
This issue has not been assessed yet.