Fix hybrid stack spec serialization in Megatron-Bridge checkpoints #2452
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,46 @@ | ||
| # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| """Megatron-Core layer specs used to build models for ModelOpt workflows.""" | ||
|
|
||
| import copy | ||
|
|
||
| from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec | ||
| from megatron.core.models.hybrid.hybrid_layer_specs import ( | ||
| hybrid_stack_spec as _te_hybrid_stack_spec, | ||
| ) | ||
| from megatron.core.transformer.spec_utils import ModuleSpec | ||
|
|
||
| __all__ = ["te_hybrid_stack_spec_sequential_mlp"] | ||
|
|
||
|
|
||
| def te_hybrid_stack_spec_sequential_mlp() -> ModuleSpec: | ||
| """Return the TE Hybrid stack spec with SequentialMLP MoE experts. | ||
|
|
||
| Named and zero-argument so a provider can store this function instead of the ModuleSpec it | ||
| builds; see ``set_moe_expert_layout`` for why a built spec cannot be serialized. | ||
|
|
||
| Its module path and name are written into ``run_config.yaml`` as a ``_target_``, so moving or | ||
| renaming it breaks every SequentialMLP hybrid checkpoint already saved. | ||
| """ | ||
| # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE. | ||
| # Replace it with SequentialMLP (TE linear layers, no grouped gemm dependency). | ||
| # num_experts only has to be non-zero to select the MoE branch; the real count comes from the | ||
| # model config at build time. | ||
| te_hybrid_stack_spec = copy.deepcopy(_te_hybrid_stack_spec) | ||
| te_hybrid_stack_spec.submodules.moe_layer.submodules.mlp = get_moe_module_spec( | ||
| use_te=True, num_experts=8, moe_grouped_gemm=False | ||
| ) | ||
| return te_hybrid_stack_spec | ||
|
kevalmorabia97 marked this conversation as resolved.
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,74 @@ | ||
| # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| import pytest | ||
| import yaml | ||
| from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider | ||
| from megatron.bridge.utils.instantiate_utils import instantiate | ||
| from megatron.bridge.utils.yaml_utils import dump_dataclass_to_yaml | ||
|
|
||
| from modelopt.torch.utils.plugins.mbridge import set_moe_expert_layout | ||
|
|
||
|
|
||
| def _round_trip(value): | ||
| """Serialize through the writer used for run_config.yaml, then reload.""" | ||
| node = yaml.safe_load(dump_dataclass_to_yaml({"spec": value}))["spec"] | ||
| return node["_target_"], instantiate(node) | ||
|
|
||
|
|
||
| @pytest.mark.parametrize( | ||
| ("moe_grouped_gemm", "expected_experts", "expected_target"), | ||
| [ | ||
| ( | ||
| True, | ||
| "TEGroupedMLP", | ||
| "megatron.bridge.models.hybrid.hybrid_provider.transformer_engine_hybrid_stack_spec", | ||
| ), | ||
| ( | ||
| False, | ||
| "SequentialMLP", | ||
| "modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp", | ||
| ), | ||
| ], | ||
| ) | ||
| def test_set_moe_expert_layout_survives_run_config_round_trip( | ||
| moe_grouped_gemm, expected_experts, expected_target | ||
| ): | ||
| """A provider's stack spec must still build real submodules after a run_config round trip. | ||
|
|
||
| A built ``ModuleSpec`` loses its ``MLPSubmodules`` / ``MoESubmodules`` when written to | ||
| ``run_config.yaml``, so ``set_moe_expert_layout`` stores a factory function instead. | ||
| """ | ||
| provider = HybridModelProvider(num_layers=2, hidden_size=64, num_attention_heads=4) | ||
| set_moe_expert_layout(provider, moe_grouped_gemm=moe_grouped_gemm) | ||
| assert provider.moe_grouped_gemm == moe_grouped_gemm | ||
|
|
||
| assert callable(provider.hybrid_stack_spec) | ||
|
|
||
| target, factory = _round_trip(provider.hybrid_stack_spec) | ||
| # The target is an on-disk contract: renaming or moving the factory breaks saved checkpoints. | ||
| assert target == expected_target | ||
|
|
||
| provider.hybrid_stack_spec = factory | ||
| spec = provider._resolve_hybrid_stack_spec() | ||
|
|
||
| mlp = spec.submodules.mlp_layer.submodules.mlp.keywords["submodules"] | ||
| assert mlp.linear_fc1 is not None | ||
| assert mlp.linear_fc2 is not None | ||
|
|
||
| moe = spec.submodules.moe_layer.submodules.mlp.keywords["submodules"] | ||
| assert moe.experts is not None | ||
| # Experts are built through a partial for the grouped-GEMM layout. | ||
| assert getattr(moe.experts, "func", moe.experts).__name__ == expected_experts | ||
|
kevalmorabia97 marked this conversation as resolved.
kevalmorabia97 marked this conversation as resolved.
Comment on lines
+71
to
+74
There was a problem hiding this comment. [SUGGESTION] This assertion path looks like it only holds for the
Why it matters: this test is the PR's only automated coverage for the fix, and the PR notes it has not been executed — if it errors on the SequentialMLP case, the half that exercises the new ModelOpt factory (and the Suggested fix: normalize the node shape before asserting, e.g. def _submodules(node):
"""MoE/MLP nodes are a functools.partial in the upstream spec and a ModuleSpec once replaced."""
return node.keywords["submodules"] if hasattr(node, "keywords") else node.submodules
def _name(node):
"""experts is a class, a partial of one, or a ModuleSpec wrapping one."""
node = getattr(node, "module", node)
return getattr(node, "func", node).__name__and use
Collaborator
Author
There was a problem hiding this comment. Both parametrizations have now been run, and both pass. The premise does not hold in megatron-core 0.19: The |
||
Uh oh!
There was an error while loading. Please reload this page.