[training migration] Finish ModelBuilder integration#5516
Conversation
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
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. |
maanug-nv
changed the title
|
|
||
| cfg = cfg_container | ||
| model_config = cfg.model | ||
| builder_cls = model_config.get_builder_cls() |
There was a problem hiding this comment.
model_provider_func is now unused, but the callers migrated in this PR (BERT, MIMO, T5, and VLM) still call pretrain_cfg_container_from_args(args) without a model config and pass their provider here. That helper deliberately leaves cfg.model as None, so each of those entry points now fails at first model setup with AttributeError: NoneType has no attribute get_builder_cls. Please either retain the provider fallback for configs without a builder or migrate those models to concrete ModelConfigs/builders before ignoring the provider.
| model_provider, | ||
| ModelType.encoder_or_decoder, | ||
| forward_step, | ||
| model_provider, |
There was a problem hiding this comment.
This entry point still cannot call the current pretrain API: the required first argument is a PretrainConfigContainer, and args_defaults/extra_args_provider are no longer accepted keyword arguments. As written, multimodal training raises TypeError before initialization. Please parse/validate the multimodal args, construct an appropriate model/config container, and then call pretrain with that container first.
|
|
||
| # get dense model | ||
| dense_model_for_upcycling = get_model(model_provider_func, model_type) | ||
| dense_model_for_upcycling = _build_model_wrapper(wrap_with_ddp=True) |
There was a problem hiding this comment.
The temporary mutations above no longer affect model construction: _build_model_wrapper always reuses cfg_container.model, whose transformer config was created earlier with the MoE num_experts and FFN size. This therefore builds a second MoE model instead of the dense source model that load_and_upcycle_model expects. Please derive/clone a dense ModelConfig after applying these overrides (and build from that), rather than reusing the original MoE config.
| data_parallel_random_init=cfg.rng.data_parallel_random_init, | ||
| ) | ||
|
|
||
| model = _build_model_wrapper(wrap_with_ddp) |
There was a problem hiding this comment.
The old get_model path detected ModelOpt state in args.load (or export_kd_teacher_load) and set args.modelopt_enabled before selecting the provider/spec. This direct builder path skips that detection, and gpt_config_from_args/hybrid_config_from_args never set restore_modelopt_state, so loading a ModelOpt checkpoint now constructs the ordinary model/spec and also takes the ordinary loss path. Please preserve the detection and propagate it into the ModelConfig before resolving the builder.
| seed_etp_group=getattr(init_pg_collection, "expt_tp", None), | ||
| ) | ||
| # TODO (@maanug): temporary until initialize.py is refactored to build pgcollection as bridge does | ||
| pg_collection = ProcessGroupCollection.use_mpu_process_groups() |
There was a problem hiding this comment.
get_model used to honor --create-all-gather-group here by creating the separate DP/EP all-gather communicators and assigning dp_cp_ag/expt_dp_ag on this collection. use_mpu_process_groups() initializes both fields to None, and the new builder path never creates them, so Megatron FSDP silently loses the requested AG/RS overlap communicators. Please preserve that setup when materializing this collection.
| """ | ||
| if profiling is None or not profiling.record_memory_history: | ||
| return | ||
| if safe_get_rank() not in profiling.profile_ranks: |
There was a problem hiding this comment.
An empty profile_ranks is the default and means “all ranks” in the existing profiling checks (len(args.profile_ranks) == 0 or rank in ...). This condition instead returns on every rank, so --record-memory-history with default rank settings never enables history and the later snapshot contains no allocation timeline. Please preserve the empty-list semantics (and align the selected rank with the existing last-rank snapshot dump).
| use_torch_fsdp2=cfg.dist.use_torch_fsdp2, | ||
| wrap_with_ddp=wrap_with_ddp, | ||
| data_parallel_random_init=cfg.rng.data_parallel_random_init, | ||
| ) |
There was a problem hiding this comment.
The builder path also drops the get_model handling for --freeze-all-layers. That option was added for logits-saving workflows and must call requires_grad_(False) (plus set router frozen_expert_bias) before DDP/optimizer construction. With this return path every parameter remains trainable, so a supposedly frozen logits-generation run can update the model. Please carry the freeze operation into a pre-wrap hook or otherwise apply it before wrapping.
| model_provider, | ||
| model_type, | ||
| forward_step_func, | ||
| model_provider=None, |
There was a problem hiding this comment.
There are still callers using the old positional order. In particular, examples/post_training/modelopt/finetune.py and megatron/elastification/pretrain_hybrid_flex.py pass model_provider as argument 3, so after this signature change it is interpreted as model_type (and the enum is interpreted as forward_step_func). Please migrate all repository pretrain(...) callers in the same change; these paths are not covered by the current green CI selection.
| checkpointing_context=None, | ||
| pg_collection=None, | ||
| *, | ||
| cfg_container: PretrainConfigContainer, |
There was a problem hiding this comment.
These new required keyword-only arguments also leave direct callers unmigrated. tasks/finetune_utils.py and tools/bert_embedding/embed.py still call this function with only provider/type, as do test_fp4_param.py, test_cuda_graphs.py, and several MTP tests. They now raise TypeError before model construction. Please update all direct callers or retain a compatible fallback/default while the migration is incomplete.
3 participants
What does this PR do ?
Replaces
model_builderfunction pointer with ModelBuilder and calls the correct methods in place ofget_model. Fixes necessary unit tests.Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.