[xuantengh]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

This PR fixes the CI import checking error caused by the DSA spec introduced in #7156.

[github-actions]

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:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[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.

[svcnvidia-nemo-ci]

The path-copying change looks sound in this light static review. One non-blocking coverage suggestion:

Add a regression test for non-deepcopyable builders

The new derivation logic deliberately avoids copying optional-dependency builders. The existing DSA spec tests check module selection, upstream nonmutation, and mixer differences, but do not exercise a builder that rejects deepcopy. With dependencies installed, those assertions also pass under the old implementation, leaving the import-failure mechanism unpinned by a unit test.

Suggested addition to that test file, using its existing imports:

from dataclasses import replace

from megatron.core.transformer.experimental_attention_variant.dsa_layer_specs import (
    _with_dsa_gqa_attention,
)


class _UncopyableBuilder:
    def __call__(self, *args, **kwargs):
        raise AssertionError("Spec derivation must not instantiate builders")

    def __deepcopy__(self, memo):
        raise TypeError("cannot pickle 'cell' object")


@pytest.mark.parametrize("upstream", [hybrid_stack_spec, gated_delta_product_stack_spec])
def test_dsa_derivation_does_not_copy_optional_builders(upstream):
    mixer = replace(upstream.submodules.mamba_layer, module=_UncopyableBuilder())
    source = replace(upstream, submodules=replace(upstream.submodules, mamba_layer=mixer))
    derived = _with_dsa_gqa_attention(source)

    assert derived.submodules.mamba_layer is mixer
    assert derived.submodules.attention_layer is not source.submodules.attention_layer

Preservation and acceptance: This proposed test should return normally for both upstream stacks; reverting to whole-stack deepcopy should raise the synthetic TypeError. Retain the existing assertions that both derived stacks use DSGroupedSelfAttention, neither upstream attention module changes, and the two mixers remain different. Also check that separate derivations receive distinct indexer specs and retain the upstream dense-attention builder for warmup. The invariant is selective copying without instantiating builders or mutating upstream specs—not restoring whole-stack deepcopy.

Suggested code requires author verification. Tests were not executed.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

[xuantengh]

/ok to test 41855ab

[xuantengh]

/ok to test 57f1ed7

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/39796311996