Repository navigation
Add formal YaRN RoPE config fields and CLI plumbing for hybrid models - #5461
Closed
guihong-nv wants to merge 1 commit into
Closed
guihong-nv wants to merge 1 commit into
guihong-nv wants to merge 1 commit into
Conversation
HybridModel already builds YarnRotaryEmbedding from getattr(self.config, "yarn_*") (NVIDIA#4244), but those attributes were never declared on TransformerConfig and had no CLI plumbing, so --position-embedding-type yarn raised AttributeError and YaRN could not be configured from the command line. - transformer_config.py: declare the 7 yarn_* fields with safe defaults (yarn_rotary_scaling_factor default 1.0 is a no-op). - arguments.py: add matching --yarn-* args; dest names match the field names so core_transformer_config_from_args plumbs them through automatically. - 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 8ec9a40 |
This branch was previously deployed
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
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