Enable custom model flops calculation #3552
Description
Is your feature request related to a problem? Please describe.
Need to support a custom model flops calculation interface if the generic calculation is not accurate for some custom model. Currently, only way is to build your special cases into the generic calculation. See Megatron-Bridge here.
Describe the solution you'd like
Some interface through the ModelConfig or ModelBuilder introduced in NVIDIA-NeMo/Megatron-Bridge#2241 that provides user with all necessary config and state to calculate model flops for their model on their own. Ideally, all models supported in Megatron-Bridge can use this interface instead of adding complexity to flop_utils.py (linked above).