Introduce refactored model builder abstractions

[maanug-nv]

What does this PR do ?

This PR refactors the ModelProviderMixin to split up responsibilities.

This refactor preserves the existing level of model customizability as well as conversion compatibility, while separating model configuration and initialization/building into separate components. This architecture was determined to be more comprehensible based on several sources of feedback.
In particular, a custom model type can define one or both of:

1. A serializable configuration object encapsulating all config settings for the model. This may be a shallow dataclass hierarchy.
2. A builder object that determines how the model is initialized for distributed training or inference using aforementioned config.

Base model architectures like GPT and Mamba will follow this design.

These new abstractions are all new classes rather than direct modifications to ModelProviderMixin to ease migration. Additionally, ModelProviderMixin relied on the inheritance from TransformerConfig, modified some config attributes during model building, and contained some unused code. Therefore, I think we should deprecate the ModelProviderMixin after integrating the new abstractions across the codebase.

The new classes are:

1. ModelConfig - Has some required attributes, including import path of builder. Supports serialization to and from dictionary.
2. ModelBuilder - Contains stub methods for implementation of (distributed) model init. Also maintains pre-wrap and post-wrap hook registration from ModelProviderMixin which is used by PEFT. Usage of the hooks is left up to the child class.

This PR also includes the necessary changes for the Mamba base model to support this new interface. The refactor was easier to demonstrate on Mamba than on GPT. 'mamba/mamba_builder.py' contains:

1. MambaModelConfig which encapsulates all config settings used to create an MCore Mamba model.
2. MambaModelBuilder which defines how a distributed Mamba model should be built.

Since distributed model initialization is identical for GPT and Mamba models, the logic is extracted into a separate function under 'unimodal.py', which can be called by GPTModelBuilder in the future. The logic is also split into helpers for readability.

I plan to upstream all new files under src/ in this PR to MegatronLM in the near future. It may also be worthwhile to push the base model configs, e.g. MambaModelConfig to MCore eventually use directly with the MCore model. Additionally, if TransformerConfig is broken up in the future, the model-specific config can remain at the top-level of the hierarchy.

Some additional notes:

• Advanced use cases may need to customize how the model wrapping/distributed model initialization is done, eg MIMO. Therefore, it makes more sense for that functionality to be defined in the ModelBuilder, so that both the single-stage building logic and distributed building logic can be overriden from the same place.
• With the functional helpers in unimodal.py, these abstractions should appear as a simple layer on top of existing MLM code like get_model() and mamba_builders.py.

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

Summary by CodeRabbit

Release Notes

• New Features

  • Introduced a model configuration and building framework supporting distributed training with DDP/FSDP options.
  • Added Mamba model builder with configurable stack specifications.
  • Added hook system for customizing model wrapping behavior.
  • Enabled virtual pipeline staging support for models.

• Refactor

  • Consolidated distributed model utilities into shared module.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[ananthsub]

i see that this is the translation for the existing setup, but how do these args generalize for the MIMO use case?

[yaoyu-33]

mimo can take in pg collection and ddp config as dict for submodules? I dont want to over-design to fit mimo

[maanug-nv]

i think signature can be a little flexible until it's integrated into training loop, which will be in follow-up PR(s).

@yashaswikarnati do you have any feedback on how this signature can be a bit better for MIMO without changing too much?

[yashaswikarnati]

Sorry just coming back to this. I see that in the signature we accept only one pg_collection. do we plan to have a different signature for multimodal.

I dont think thats overdesign - I would say thats the bare minimum to achieve the functionality.