tests/functional_tests/data/test_samplers.py (1)

620-622: LGTM!

Migration is correct.

Optional: The _DummyBridge class with GPTModelProvider is duplicated across 5 test methods. Consider extracting to a module-level helper or pytest fixture if this pattern expands further.

,

♻️ Optional refactor using pytest fixture

# At module level or in conftest.py
`@pytest.fixture`
def mock_llama3_bridge():
    """Mock AutoBridge to avoid HF downloads, returning GPTModelProvider."""
    from unittest import mock as _mock

    class _DummyBridge:
        def to_megatron_provider(self, load_weights=False):
            from megatron.bridge.models.gpt_provider import GPTModelProvider
            return GPTModelProvider()

    with _mock.patch(
        "megatron.bridge.recipes.llama.llama3.AutoBridge.from_hf_pretrained",
        return_value=_DummyBridge()
    ):
        yield

Then tests can simply use def test_...(self, mock_llama3_bridge): without the repeated boilerplate.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/functional_tests/data/test_samplers.py` around lines 620 - 622, Extract
the repeated _DummyBridge + GPTModelProvider pattern into a reusable pytest
fixture or module-level helper to avoid duplication: create a fixture (e.g.,
mock_llama3_bridge) that defines class _DummyBridge with method
to_megatron_provider returning GPTModelProvider, and use unittest.mock.patch to
patch "megatron.bridge.recipes.llama.llama3.AutoBridge.from_hf_pretrained" to
return that _DummyBridge, then update the five tests to accept the fixture
(e.g., def test_...(mock_llama3_bridge)) instead of repeating the inline class
and patch.

tests/unit_tests/recipes/test_moonlight_recipes.py (1)

136-136: Consolidate repeated MLAModelProvider monkeypatching into a pytest fixture.

The identical monkeypatch.setattr(mod, "MLAModelProvider", _FakeMoonlightModelProvider16B) setup is repeated in many tests. Please move this into a shared fixture (module-level or reusable helper fixture) to reduce duplication and keep setup consistent.

As per coding guidelines, "tests/**/*.py: Use pytest fixtures for common setup in unit tests".

Also applies to: 171-171, 196-196, 222-222, 238-238, 261-261, 284-284, 307-307, 329-329

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit_tests/recipes/test_moonlight_recipes.py` at line 136, Create a
reusable pytest fixture that performs the repeated monkeypatching of
MLAModelProvider to _FakeMoonlightModelProvider16B and replace the duplicated
monkeypatch.setattr(mod, "MLAModelProvider", _FakeMoonlightModelProvider16B)
lines in tests with that fixture; implement a fixture (e.g.,
fake_mlamodel_provider or patch_mla_provider) that accepts monkeypatch, calls
monkeypatch.setattr(mod, "MLAModelProvider", _FakeMoonlightModelProvider16B),
mark it module- or session-scoped as appropriate, and add the fixture name to
the tests that previously had the inline monkeypatch to eliminate duplication
while preserving behavior.

src/megatron/bridge/recipes/moonlight/moonlight_16b.py (2)

413-468: Optional: Extract shared architecture parameters to reduce duplication.

The MLAModelProvider instantiation in moonlight_16b_peft_config duplicates ~40 architecture parameters that are identical to moonlight_16b_sft_config (lines 200-241 vs 414-454). Only the parallelism settings differ.

This is acceptable for clarity and explicit configuration, but if maintainability becomes a concern, consider extracting the shared architecture parameters into a module-level constant.

♻️ Example: Extract shared architecture parameters

# Module-level constant for shared architecture
_MOONLIGHT_16B_ARCH = dict(
    num_layers=27,
    hidden_size=2048,
    ffn_hidden_size=11264,
    # ... remaining shared params
)

def moonlight_16b_sft_config() -> ConfigContainer:
    cfg.model = MLAModelProvider(
        **_MOONLIGHT_16B_ARCH,
        # SFT-specific parallelism
        tensor_model_parallel_size=2,
        expert_model_parallel_size=8,
        sequence_parallel=True,
        # ...
    )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/recipes/moonlight/moonlight_16b.py` around lines 413 -
468, The MLAModelProvider block in moonlight_16b_peft_config duplicates many
architecture params already used in moonlight_16b_sft_config; extract the shared
architecture into a module-level constant (e.g., _MOONLIGHT_16B_ARCH) containing
the common keys (num_layers, hidden_size, ffn_hidden_size, num_attention_heads,
kv_channels, vocab_size, rotary_*, layernorm_epsilon, init_method_std,
activation_func, gated_linear_unit, etc.), then update both
moonlight_16b_sft_config and moonlight_16b_peft_config to call
MLAModelProvider(**_MOONLIGHT_16B_ARCH, <their specific overrides>) so only
parallelism and config-specific overrides remain in each function.

---

199-255: Consider removing redundant parameter assignments.

The MLAModelProvider constructor sets moe_token_dispatcher_type="alltoall" (line 241) and moe_grouped_gemm=True (line 240), but these are reassigned later in the function:

• cfg.model.moe_token_dispatcher_type at lines 284
• cfg.model.moe_grouped_gemm at line 308
• cfg.model.moe_permute_fusion at lines 274 and 307

While functionally correct, the duplication could lead to drift if only one location is updated.

♻️ Suggested cleanup: remove redundant assignments from constructor or later code

Either remove from the constructor (if they're always overridden by the DeePEP logic):

         moe_grouped_gemm=True,
-        moe_token_dispatcher_type="alltoall",
         # Parallelism

Or consolidate all MoE settings in one place and remove the later redundant assignments.

Note: The Ruff S106 warnings about moe_token_dispatcher_type are false positives—this is a Mixture-of-Experts configuration parameter, not a password.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/recipes/moonlight/moonlight_16b.py` around lines 199 -
255, The constructor for MLAModelProvider sets MoE-related defaults
(moe_grouped_gemm, moe_token_dispatcher_type, moe_permute_fusion) that are later
reassigned on cfg.model (e.g., cfg.model.moe_token_dispatcher_type,
cfg.model.moe_grouped_gemm, cfg.model.moe_permute_fusion); remove the redundant
assignments either from the MLAModelProvider(...) argument list or from the
later reassignment block so all MoE settings are defined in one place—locate
MLAModelProvider(...) in moonlight_16b.py and the subsequent lines where
cfg.model.moe_permute_fusion, cfg.model.moe_token_dispatcher_type, and
cfg.model.moe_grouped_gemm are set, then keep only the preferred location and
delete the duplicates.

tests/unit_tests/training/test_config.py (1)

79-90: Consider renaming the helper to reflect the new abstraction.

create_test_deepseek_config now returns MLAModelProvider. A neutral helper name (e.g., create_test_mla_config) would reduce ambiguity as MLA coverage expands.

♻️ Optional clarity refactor

-def create_test_deepseek_config(**kwargs: Any) -> MLAModelProvider:
-    """Creates an instance of MLAModelProvider for testing."""
+def create_test_mla_config(**kwargs: Any) -> MLAModelProvider:
+    """Creates an instance of MLAModelProvider for testing."""
     defaults = {
         "num_layers": 1,
         "hidden_size": 128,
         "num_attention_heads": 4,
         "seq_length": 512,
         "apply_rope_fusion": False,
     }
     defaults.update(kwargs)
     return MLAModelProvider(**defaults)
+
+# Backward-compatible alias for existing tests.
+create_test_deepseek_config = create_test_mla_config

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit_tests/training/test_config.py` around lines 79 - 90, Rename the
helper function create_test_deepseek_config to a neutral name like
create_test_mla_config since it now returns MLAModelProvider; update the
function definition and all call sites to the new name (search for
create_test_deepseek_config and replace with create_test_mla_config) and keep
the same parameters/defaults and return MLAModelProvider to preserve behavior.

tests/unit_tests/recipes/nemotronh/test_nemotronh.py (1)

163-177: Consider simplifying the parameterized test.

Since all variants now use MambaModelProvider, the provider_cls parameter is redundant—every tuple maps to the same class. You could simplify by removing the parameter and using MambaModelProvider directly in the assertion.

That said, keeping the current structure is fine if you prefer consistency with other parameterized tests or want flexibility for future divergence.

♻️ Optional simplification

     `@pytest.mark.parametrize`(
-        "recipe_fn,provider_cls",
+        "recipe_fn",
         [
-            (nemotronh_4b_pretrain_config, MambaModelProvider),
-            (nemotronh_8b_pretrain_config, MambaModelProvider),
-            (nemotronh_47b_pretrain_config, MambaModelProvider),
-            (nemotronh_56b_pretrain_config, MambaModelProvider),
+            nemotronh_4b_pretrain_config,
+            nemotronh_8b_pretrain_config,
+            nemotronh_47b_pretrain_config,
+            nemotronh_56b_pretrain_config,
         ],
     )
-    def test_config_container_structure(self, recipe_fn, provider_cls):
+    def test_config_container_structure(self, recipe_fn):
         """Test that all configs return proper ConfigContainer with correct model provider."""
         config = recipe_fn()
 
         assert isinstance(config, ConfigContainer)
-        assert isinstance(config.model, provider_cls)
+        assert isinstance(config.model, MambaModelProvider)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit_tests/recipes/nemotronh/test_nemotronh.py` around lines 163 - 177,
The parameterized test test_config_container_structure currently includes a
redundant provider_cls parameter because all tuples use MambaModelProvider;
simplify by parameterizing only the recipe functions
(nemotronh_4b_pretrain_config, nemotronh_8b_pretrain_config,
nemotronh_47b_pretrain_config, nemotronh_56b_pretrain_config) and replace the
second assertion to directly assert isinstance(config.model, MambaModelProvider)
while keeping the initial assert isinstance(config, ConfigContainer); update the
`@pytest.mark.parametrize` decorator and the test signature accordingly to remove
provider_cls.

src/megatron/bridge/recipes/nemotronh/nemotron_3_nano.py (1)

211-268: Configuration duplication across pretrain/SFT/PEFT is notable but acceptable.

The SFT and PEFT configs duplicate the MambaModelProvider instantiation with ~40 identical architecture parameters. While this could be extracted into a helper function, this pattern appears consistent with other recipes in the codebase.

♻️ Optional: Extract common architecture parameters

def _nemotron_3_nano_model_kwargs(seq_length: int = 8192) -> dict:
    """Common Nemotron 3 Nano architecture parameters."""
    return dict(
        hybrid_override_pattern="MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME",
        num_layers=52,
        hidden_size=2688,
        mamba_num_heads=64,
        kv_channels=128,
        mamba_state_dim=128,
        # ... remaining common params
        seq_length=seq_length,
    )

# Usage in configs:
cfg.model = MambaModelProvider(
    **_nemotron_3_nano_model_kwargs(seq_length=8192),
    tensor_model_parallel_size=4,
    # ... parallelism params
)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/recipes/nemotronh/nemotron_3_nano.py` around lines 211 -
268, Multiple recipes duplicate the same MambaModelProvider architecture block;
extract those shared parameters into a helper function (e.g.,
_nemotron_3_nano_model_kwargs) that returns a dict of the common keys
(hybrid_override_pattern, num_layers, hidden_size, mamba_num_heads, kv_channels,
mamba_state_dim, ffn_hidden_size, seq_length, etc.), then update cfg.model
assignments in pretrain/SFT/PEFT to call
MambaModelProvider(**_nemotron_3_nano_model_kwargs(seq_length=...), <override or
additional parallelism args>) so only recipe-specific fields (like
tensor_model_parallel_size, pipeline_model_parallel_size, pipeline_dtype,
expert_model_parallel_size) remain inline.

tests/functional_tests/training/test_pretrain_resume.py (1)

47-72: LGTM with optional consolidation opportunity.

The Llama3ModelProvider145M class correctly inherits from GPTModelProvider with Llama-3-style defaults. Note that this class definition is nearly identical to the one in test_sft.py (with minor differences in seq_length and num_layers defaults). Consider extracting a shared test utility class in a future refactor to reduce duplication across functional tests.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/functional_tests/training/test_pretrain_resume.py` around lines 47 -
72, The Llama3ModelProvider145M test class duplicates a nearly identical
definition found in test_sft.py; refactor by extracting the shared defaults into
a single reusable test utility (e.g., a base class or factory function) and have
both Llama3ModelProvider145M and the counterpart in test_sft.py inherit from or
invoke that utility; update the unique overrides (seq_length, num_layers) in
each test-specific subclass or factory call and remove the duplicated field
definitions to centralize maintenance.

tests/functional_tests/training/test_finetune_lora.py (1)

53-80: LGTM!

The Llama3ModelProvider145M class correctly inherits from GPTModelProvider with Llama-3-style defaults. The explicit vocab_size=128256 is appropriate for LoRA finetuning tests that use HuggingFace tokenizers.

This is the third definition of Llama3ModelProvider145M across functional test files (also in test_sft.py and test_pretrain_resume.py). Consider consolidating these into a shared test utility module (e.g., tests/functional_tests/utils.py or a dedicated tests/functional_tests/training/model_providers.py) in a future refactor to reduce maintenance burden.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/functional_tests/training/test_finetune_lora.py` around lines 53 - 80,
Duplicate Llama3ModelProvider145M definitions should be consolidated into a
single shared test utility: create a new test helper module (e.g.,
tests/functional_tests/training/model_providers.py or utils.py), move the class
Llama3ModelProvider145M into that module, update usages in tests that currently
define the class (references in test_sft.py and test_pretrain_resume.py and this
file) to import Llama3ModelProvider145M from the new module, and remove the
duplicate class definitions so tests import the single canonical provider.

tests/functional_tests/training/test_decentralized_pg.py (1)

104-140: High duplication in model config setup across decentralized-PG tests.

A shared builder (with per-test overrides for TP/PP/CP/share-embeddings) would make these tests easier to maintain and less brittle.

Also applies to: 255-290, 410-446, 566-602, 722-758, 878-914

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/functional_tests/training/test_decentralized_pg.py` around lines 104 -
140, The duplicated GPTModelProvider block should be extracted into a shared
builder function (e.g., make_decentralized_model_cfg or build_base_gpt_model)
that returns a GPTModelProvider pre-populated with the common args shown
(normalization, activation_func, gated_linear_unit, position_embedding_type, ...
hidden_size, ffn_hidden_size, num_attention_heads, init_method_std,
layernorm_epsilon, rotary/rope settings, seq_length,
make_vocab_size_divisible_by, vocab_size, num_layers, etc.), and each test
should call that builder and apply per-test overrides for
tensor_model_parallel_size, pipeline_model_parallel_size, context_parallel_size,
sequence_parallel, and share_embeddings_and_output_weights; update all
referenced tests (the blocks at ~104-140, 255-290, 410-446, 566-602, 722-758,
878-914) to use the builder and only change the TP/PP/CP/share-embeddings fields
to reduce duplication and improve maintainability.

src/megatron/bridge/recipes/nemotronh/nemotron_nano_v2.py (1)

38-77: Consider factoring shared MambaModelProvider kwargs into reusable builders/constants.

These six blocks duplicate a large common parameter surface; extracting shared defaults would reduce maintenance overhead and accidental divergence across pretrain/SFT/PEFT variants.

Also applies to: 173-212, 310-349, 436-475, 571-610, 718-757

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/recipes/nemotronh/nemotron_nano_v2.py` around lines 38 -
77, Multiple MambaModelProvider invocations (e.g., the cfg.model =
MambaModelProvider block) duplicate a large common parameter set across
variants; factor those shared kwargs into a reusable constant or builder (for
example a MAMBA_DEFAULTS dict or a build_mamba_kwargs(name, overrides) helper)
and then construct each provider via unpacking the shared defaults plus only the
variant-specific overrides (use the existing symbols like MambaModelProvider and
cfg.model to locate callers and merge with per-variant keys such as num_layers,
seq_length, pipeline_dtype, etc.). Ensure the new helper/constant is
imported/defined once and replace the repeated blocks at the other occurrences
noted so differences remain explicit and minimal.

tests/functional_tests/training/test_pretrain.py (1)

73-108: Consider extracting a shared GPT model-config factory.

These two constructor blocks are mostly duplicated; a small helper would reduce config drift across tests when provider defaults/fields change.

Also applies to: 219-255

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/functional_tests/training/test_pretrain.py` around lines 73 - 108, The
two nearly identical GPTModelProvider constructor blocks (creating model_cfg)
should be extracted into a shared factory to avoid duplication and drift; add a
helper function (e.g., make_gpt_model_cfg or build_gpt_model_config) that
returns a configured GPTModelProvider with the common parameters shown
(normalization, activation_func, gated_linear_unit, position_embedding_type,
add_bias_linear, attention_dropout, hidden_dropout, bias_activation_fusion,
masked_softmax_fusion, persist_layer_norm, bias_dropout_fusion,
apply_rope_fusion, num_query_groups, init_method_std, layernorm_epsilon,
rotary_percent, rope_scaling, rope_scaling_factor,
share_embeddings_and_output_weights, rotary_base, hidden_size, ffn_hidden_size,
num_attention_heads, tensor_model_parallel_size, pipeline_model_parallel_size,
context_parallel_size, sequence_parallel, attention_softmax_in_fp32,
pipeline_dtype, bf16, seq_length, make_vocab_size_divisible_by, vocab_size,
num_layers) and replace the two inline constructor blocks (the model_cfg
assignments at the shown location and the similar block around lines 219–255)
with calls to that helper, allowing optional overrides for seq_length,
num_layers, vocab_size, etc.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro