[model] feat: Add YARN support for mamba_model from MCORE

[guihong-nv]

What does this PR do ?

Add the YARN support for mamba_model from MCORE.

Changelog

src/megatron/bridge/models/mamba/mamba_builder.py

• Added "yarn" to the position_embedding_type Literal
• Added 7 YaRN config fields (yarn_rotary_scaling_factor, yarn_original_max_position_embeddings, yarn_beta_fast, yarn_beta_slow, yarn_mscale, yarn_mscale_all_dim, yarn_correction_range_round_to_int)
• In build_model(), injects YaRN fields onto the embedded TransformerConfig before constructing MCoreMambaModel

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

• Same 7 YaRN fields added as dataclass fields (directly accessible by MCore since the provider IS the TransformerConfig)
• In provide(), computes yarn_original_max_position_embeddings default from seq_length / yarn_rotary_scaling_factor when not explicitly set

tests/unit_tests/models/mamba/test_mamba_builder.py

• TestMambaModelConfigYarnDefaults: field defaults, type acceptance, mutability
• TestMambaModelBuilderBuildModelWithYarn: YaRN attr injection onto transformer config, default/explicit yarn_original_max_position_embeddings, no injection for non-YaRN types

tests/unit_tests/models/mamba/test_mamba_provider.py

• TestMambaModelProviderYarnDefaults: field defaults and type acceptance
• TestMambaModelProviderProvideWithYarn: default computation, explicit value preservation, no injection for non-YaRN types

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (PR)

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