[Refactor] Refactor model nparams calculation into single static function #1767
Conversation
meta-cla
bot
added
the
CLA Signed
| from torchtitan.tools.logging import logger | ||
|
|
||
|
|
||
| def get_dense_model_nparams_and_flops( |
There was a problem hiding this comment.
I feel this function for dense can be entirely covered by the MoE one. Should we keep it for readability? I'm OK with that.
There was a problem hiding this comment.
I made separate functions because in moe model nparams function, there are several places accessing model_args.moe_args. For our dense model definition, we don't have this field. So if we want to merge them, we need to add several checks like if hasattr(model_args, "xxx"). Technically we can merge them, but for readability I make the separate for now
| # If weight tying is enabled, subtract embedding parameters from total count | ||
| if hasattr(model_args, "enable_weight_tying") and model_args.enable_weight_tying: | ||
| nparams = nparams - nparams_embedding |
There was a problem hiding this comment.
May not be needed by real models in torchtitan, but I feel this should be added to the dense part as well, because later on new dense + weight tying model could be added so if they reuse the dense fn it's still correct. An example would be Llama 3.2 1B / 3B.
3 participants
As titled, remove repeated copy