Skip to content

Commit 2fe99b1

Browse files
committed
src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py
Signed-off-by: Li Ding <liding@nvidia.com>
1 parent b30a725 commit 2fe99b1

8 files changed

Lines changed: 15 additions & 145 deletions

File tree

‎src/megatron/bridge/models/conversion/param_mapping.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -531,7 +531,8 @@ def gather_from_tp_ranks(self, tensor: torch.Tensor) -> List[torch.Tensor]:
531531
torch.distributed.all_gather(gathered, tensor, group=self.tp_group)
532532
return gathered
533533

534-
def _count_wildcard_groups(self, pattern: str) -> int:
534+
@staticmethod
535+
def _count_wildcard_groups(pattern: str) -> int:
535536
"""Count the number of wildcard capture groups in a pattern.
536537
537538
Args:

‎src/megatron/bridge/models/mamba/mamba_provider.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -166,9 +166,6 @@ class MambaModelProvider(TransformerConfig, ModelProviderMixin[MCoreMambaModel])
166166
mamba_stack_spec: Union[ModuleSpec, Callable[[], ModuleSpec], Callable[["MambaModelProvider"], ModuleSpec]] = (
167167
get_default_mamba_stack_spec
168168
)
169-
mtp_mamba_stack_spec: Union[ModuleSpec, Callable[[], ModuleSpec], Callable[["MambaModelProvider"], ModuleSpec]] = (
170-
get_default_mamba_stack_spec
171-
)
172169
vocab_size: Optional[int] = None
173170
should_pad_vocab: bool = False
174171
hf_model_id: Optional[str] = None

‎src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py‎

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -37,17 +37,6 @@
3737
logger = logging.getLogger(__name__)
3838

3939

40-
def _count_wildcards(pattern: str) -> int:
41-
# Mirrors MegatronParamMapping._count_wildcard_groups (private there).
42-
count = 0
43-
remaining = pattern
44-
while "**" in remaining:
45-
count += 1
46-
remaining = remaining.replace("**", "", 1)
47-
count += remaining.count("*")
48-
return count
49-
50-
5140
def _replace_wildcards(pattern: str, captures: Tuple[str, ...]) -> str:
5241
"""Replace ** then * sequentially with captures."""
5342
out = pattern
@@ -108,8 +97,10 @@ def __init__(
10897

10998
# Precompute wildcard counts to disambiguate resolve() calls coming from
11099
# megatron_to_hf_lookup vs hf_to_megatron_lookup.
111-
self._megatron_wc = _count_wildcards(self.megatron_param)
112-
self._hf_wc = _count_wildcards(self.hf_param) if isinstance(self.hf_param, str) else 0
100+
self._megatron_wc = MegatronParamMapping._count_wildcard_groups(self.megatron_param)
101+
self._hf_wc = (
102+
MegatronParamMapping._count_wildcard_groups(self.hf_param) if isinstance(self.hf_param, str) else 0
103+
)
113104

114105
def resolve(self, captures: Tuple[str, ...]) -> MegatronParamMapping:
115106
# We primarily expect captures from Megatron pattern matching.
@@ -161,7 +152,9 @@ def resolve(self, captures: Tuple[str, ...]) -> MegatronParamMapping:
161152

162153
return AutoMapping(megatron_param=resolved_megatron, hf_param=resolved_hf)
163154

164-
# These are never called: this mapping always resolves into an AutoMapping.
155+
# Required by ABC but never called at runtime: resolve() always returns an
156+
# AutoMapping, so hf_to_megatron/megatron_to_hf are invoked on that resolved
157+
# instance, not on this class.
165158
def hf_to_megatron(self, hf_weights: torch.Tensor, megatron_module: torch.nn.Module) -> torch.Tensor:
166159
raise NotImplementedError
167160

@@ -189,8 +182,8 @@ def __init__(
189182
self._mtp_layers_per_block = int(mtp_layers_per_block)
190183
if self._mtp_layers_per_block <= 0:
191184
raise ValueError("mtp_layers_per_block must be > 0 for MTP flattening QKV mapping.")
192-
self._megatron_wc = _count_wildcards(self.megatron_param)
193-
self._hf_wc = _count_wildcards(q) # q/k/v share pattern structure here
185+
self._megatron_wc = MegatronParamMapping._count_wildcard_groups(self.megatron_param)
186+
self._hf_wc = MegatronParamMapping._count_wildcard_groups(q) # q/k/v share pattern structure here
194187

195188
def resolve(self, captures: Tuple[str, ...]) -> MegatronParamMapping:
196189
# Expect captures from Megatron lookup: (outer, inner)
@@ -207,6 +200,9 @@ def resolve(self, captures: Tuple[str, ...]) -> MegatronParamMapping:
207200

208201
return QKVMapping(megatron_param=resolved_megatron, q=resolved_q, k=resolved_k, v=resolved_v)
209202

203+
# Required by ABC but never called at runtime: resolve() always returns a
204+
# QKVMapping, so hf_to_megatron/megatron_to_hf are invoked on that resolved
205+
# instance, not on this class.
210206
def hf_to_megatron(self, hf_weights: Dict[str, torch.Tensor], megatron_module: torch.nn.Module) -> torch.Tensor:
211207
raise NotImplementedError
212208

‎src/megatron/bridge/models/nemotronh/nemotron_h_provider.py‎

Lines changed: 0 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -167,36 +167,6 @@ class Nemotron3NanoProvider(NemotronHModelProvider):
167167
moe_router_group_topk: int = 1
168168

169169

170-
@dataclass
171-
class Nemotron3SuperProvider(NemotronHModelProvider):
172-
"""Configuration for a 3B parameter Nemotron 3 Super model."""
173-
174-
seq_length: int = 262144
175-
num_query_groups: int = 2
176-
hybrid_layer_pattern: str = (
177-
"MEMEMEM*EMEMEMEM*EMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEM*EMEMEMEME"
178-
)
179-
hidden_size: int = 4096
180-
mamba_num_heads: int = 128
181-
kv_channels: int = 128
182-
mamba_state_dim: int = 128
183-
ffn_hidden_size: int = 2688
184-
num_attention_heads: int = 32
185-
mamba_head_dim: int = 64
186-
num_moe_experts: int = 512
187-
moe_ffn_hidden_size: int = 2688
188-
moe_shared_expert_intermediate_size: int = 5376
189-
moe_router_topk: int = 22
190-
moe_router_topk_scaling_factor: float = 5.0
191-
moe_router_num_groups: int = 1
192-
moe_router_group_topk: int = 1
193-
moe_latent_size: int = 1024
194-
moe_shared_expert_overlap: bool = False
195-
mtp_num_layers: int = 2
196-
mtp_hybrid_override_pattern: str = "*E"
197-
mtp_use_repeated_layer: bool = False
198-
199-
200170
# -----------------------------------------------------------------------------
201171
# Deprecated aliases (to be removed in a future release)
202172
# -----------------------------------------------------------------------------

tests/functional_tests/launch_scripts/active/L1_Launch_models_nemotronh.sh renamed to tests/functional_tests/launch_scripts/active/L0_Launch_models_nemotronh.sh

File renamed without changes.

tests/functional_tests/launch_scripts/active/L1_Launch_recipes_nemotronh.sh renamed to tests/functional_tests/launch_scripts/active/L0_Launch_recipes_nemotronh.sh

File renamed without changes.

‎tests/unit_tests/models/nemotronh/test_nemotron_h_provider.py‎

Lines changed: 0 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717

1818
from megatron.bridge.models.nemotronh.nemotron_h_provider import (
1919
Nemotron3NanoProvider,
20-
Nemotron3SuperProvider,
2120
NemotronHModel4BProvider,
2221
NemotronHModel8BProvider,
2322
NemotronHModel47BProvider,
@@ -284,10 +283,6 @@ def test_nemotron_nano_12b_v2_inherits_from_base(self):
284283
"""Test Nemotron Nano v2 12B provider inherits from NemotronHModelProvider."""
285284
assert issubclass(NemotronNano12Bv2Provider, NemotronHModelProvider)
286285

287-
def test_nemotron_3_super_inherits_from_base(self):
288-
"""Test Nemotron 3 Super provider inherits from NemotronHModelProvider."""
289-
assert issubclass(Nemotron3SuperProvider, NemotronHModelProvider)
290-
291286
def test_provide_method_inherited(self):
292287
"""Test that provide method works correctly in inherited classes."""
293288
# Test with Nemotron-H 4B
@@ -299,7 +294,6 @@ def test_provide_method_inherited(self):
299294
NemotronNano9Bv2Provider(),
300295
NemotronNano12Bv2Provider(),
301296
Nemotron3NanoProvider(),
302-
Nemotron3SuperProvider(),
303297
]
304298

305299
for provider in providers:
@@ -463,93 +457,6 @@ def test_nemotron_3_nano_inherits_moe_defaults(self):
463457
assert provider.moe_shared_expert_overlap is True
464458

465459

466-
class TestNemotron3SuperProvider:
467-
"""Test cases for Nemotron3SuperProvider class."""
468-
469-
def test_nemotron_3_super_default_configuration(self):
470-
"""Test Nemotron 3 Super model has correct default configuration."""
471-
provider = Nemotron3SuperProvider()
472-
473-
# Check Nemotron 3 Super specific configuration
474-
assert provider.seq_length == 262144
475-
assert provider.hidden_size == 4096
476-
assert provider.num_attention_heads == 32
477-
assert provider.num_query_groups == 2
478-
assert provider.mamba_num_heads == 128
479-
assert provider.kv_channels == 128
480-
assert provider.mamba_state_dim == 128
481-
assert provider.ffn_hidden_size == 2688
482-
assert provider.mamba_head_dim == 64
483-
assert (
484-
provider.hybrid_layer_pattern
485-
== "MEMEMEM*EMEMEMEM*EMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEM*EMEMEMEME"
486-
)
487-
# num_layers is derived from hybrid_layer_pattern in finalize()
488-
assert provider.num_layers is None
489-
assert len(provider.hybrid_layer_pattern) == 88
490-
491-
def test_nemotron_3_super_moe_configuration(self):
492-
"""Test Nemotron 3 Super model MoE-specific configuration."""
493-
provider = Nemotron3SuperProvider()
494-
495-
# Check MoE-specific configuration
496-
assert provider.num_moe_experts == 512
497-
assert provider.moe_ffn_hidden_size == 2688
498-
assert provider.moe_shared_expert_intermediate_size == 5376
499-
assert provider.moe_router_topk == 22
500-
assert provider.moe_router_topk_scaling_factor == 5.0
501-
assert provider.moe_router_num_groups == 1
502-
assert provider.moe_router_group_topk == 1
503-
assert provider.moe_latent_size == 1024
504-
assert provider.moe_shared_expert_overlap is False
505-
506-
def test_nemotron_3_super_mtp_configuration(self):
507-
"""Test Nemotron 3 Super model MTP-specific configuration."""
508-
provider = Nemotron3SuperProvider()
509-
510-
# Check MTP-specific configuration
511-
assert provider.mtp_num_layers == 2
512-
assert provider.mtp_hybrid_override_pattern == "*E"
513-
assert provider.mtp_use_repeated_layer is False
514-
515-
def test_nemotron_3_super_override_configuration(self):
516-
"""Test Nemotron 3 Super model with overridden configuration."""
517-
provider = Nemotron3SuperProvider(
518-
seq_length=16384,
519-
hidden_dropout=0.1,
520-
num_moe_experts=64,
521-
)
522-
523-
# Check overridden values
524-
assert provider.seq_length == 16384
525-
assert provider.hidden_dropout == 0.1
526-
assert provider.num_moe_experts == 64
527-
528-
# Check critical defaults remain (num_layers derived from pattern in finalize())
529-
assert provider.num_layers is None
530-
assert len(provider.hybrid_layer_pattern) == 88
531-
assert provider.hidden_size == 4096
532-
assert provider.mamba_num_heads == 128
533-
534-
def test_nemotron_3_super_inherits_from_base(self):
535-
"""Test Nemotron 3 Super provider inherits from NemotronHModelProvider."""
536-
assert issubclass(Nemotron3SuperProvider, NemotronHModelProvider)
537-
538-
def test_nemotron_3_super_inherits_moe_defaults(self):
539-
"""Test Nemotron 3 Super inherits MoE defaults from base class."""
540-
provider = Nemotron3SuperProvider()
541-
542-
# Check inherited MoE defaults from NemotronHModelProvider
543-
assert provider.moe_aux_loss_coeff == 0.0001
544-
assert provider.moe_router_score_function == "sigmoid"
545-
assert provider.moe_router_enable_expert_bias is True
546-
assert provider.moe_router_load_balancing_type == "seq_aux_loss"
547-
assert provider.moe_router_dtype == "fp32"
548-
assert provider.moe_grouped_gemm is True
549-
assert provider.moe_token_dispatcher_type == "alltoall"
550-
assert provider.moe_permute_fusion is True
551-
552-
553460
class TestHybridPatterns:
554461
"""Test hybrid override patterns of Nemotron-H providers."""
555462

@@ -563,7 +470,6 @@ def test_hybrid_patterns_contain_mamba_and_attention(self):
563470
NemotronNano9Bv2Provider(),
564471
NemotronNano12Bv2Provider(),
565472
Nemotron3NanoProvider(),
566-
Nemotron3SuperProvider(),
567473
]
568474

569475
for provider in providers:

‎tests/unit_tests/recipes/nemotronh/test_nemotron_3_super.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,7 @@ def test_moe_model_configuration(self, recipe_fn):
311311
"""Test MoE-specific model configuration from provider."""
312312
config = recipe_fn()
313313

314-
# Check MoE settings from Nemotron3SuperProvider
314+
# Check MoE settings from AutoBridge provider
315315
assert config.model.num_moe_experts == 512
316316
assert config.model.moe_ffn_hidden_size == 2688
317317
assert config.model.moe_shared_expert_intermediate_size == 5376

0 commit comments

Comments
 (0)