Add formal YaRN RoPE config fields for hybrid models - #5465
Draft
guihong-nv wants to merge 7 commits into
Draft
Conversation
HybridModel already builds YarnRotaryEmbedding from getattr(self.config, "yarn_*") (NVIDIA#4244), but those attributes were never declared on TransformerConfig, so --position-embedding-type yarn raised AttributeError and YaRN could not be configured. Declaring the 7 yarn_* fields on TransformerConfig fixes this and also makes the --yarn-* CLI flags appear automatically: _add_network_size_args builds a TransformerConfig argument group via ArgumentGroupFactory(TransformerConfig), which derives one CLI flag per dataclass field. No manual argument plumbing is needed -- adding the flags by hand duplicates the auto-generated ones and raises argparse 'conflicting option string'. - transformer_config.py: declare the 7 yarn_* fields with safe defaults (yarn_rotary_scaling_factor default 1.0 is a no-op). - test_hybrid_model.py: drive the yarn test config via the constructor and add tests for the dataclass fields and config->embedding plumbing. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Guihong Li <guihongl@nvidia.com>
Contributor
Author
|
/ok to test ad69cde |
The hybrid MoE golden-config drift test (test_hybrid_moe_model.py) pins the exact set of TransformerConfig fields. Adding the yarn_* fields trips its [ADDED ARGS] check, so register them in GOLDEN_CONFIG with their default values. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Guihong Li <guihongl@nvidia.com>
Contributor
Author
|
/ok to test c6fd678 |
main adds the --yarn-* CLI flags manually in _add_network_size_args. Now that the yarn_* fields are declared on TransformerConfig, ArgumentGroupFactory would also auto-generate those flags, producing duplicate argparse option strings and breaking every test that builds the parser. Exclude the yarn_* fields from auto-generation so the manual flags remain the single source; the fields still provide config storage + defaults consumed by gpt_model/hybrid_model. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Guihong Li <guihongl@nvidia.com>
Contributor
Author
|
/ok to test 1636571 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Declares the YaRN RoPE parameters as formal
TransformerConfigfields so that--position-embedding-type yarnis actually usable for hybrid models.HybridModelalready instantiatesYarnRotaryEmbeddingfromgetattr(self.config, "yarn_*")(added in #4244), but thoseyarn_*attributes were never declared onTransformerConfig— they were only set dynamically in tests. As a result, launching with--position-embedding-type yarnraisedAttributeError.megatron/core/transformer/transformer_config.py— declares the 7 YaRN fields onTransformerConfigwith safe defaults (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). The defaultyarn_rotary_scaling_factor=1.0is a no-op, so existing configs are unaffected.tests/unit_tests/models/test_hybrid_model.py— drives the YaRN test config via the constructor (the formal-field path) and adds two tests: one asserting the fields exist as dataclass fields with the expected defaults, and one asserting the config values flow into theYarnRotaryEmbeddinginstance.The matching
--yarn-*CLI flags are produced automatically:_add_network_size_argsbuilds aTransformerConfigargument group viaArgumentGroupFactory(TransformerConfig), which derives one CLI flag (and the args→config plumbing) per dataclass field. No manual argument plumbing is needed — adding the flags by hand would duplicate the auto-generated ones and raiseargparse: conflicting option string.No behavior change for non-YaRN paths. Validated on GPU: the full arg parser builds with the
--yarn-*flags auto-generated (no conflict), andtests/unit_tests/models/test_hybrid_model.pyplustests/unit_tests/transformer/test_multi_latent_attention.pypass (one pre-existing, environment-related context-parallel SIGABRT reproduces on unmodifiedmainand is unrelated to this change).Issue tracking
Related to #4244 (which added the YaRN wiring in
HybridModel).Contribution process
Pre-checks