Skip to content

Add YaRN RoPE scaling and correct dynamic scaling - #474

Closed
dkuku wants to merge 1 commit into
elixir-nx:mainfrom
dkuku:feat/rope-yarn-dynamic
Closed

dkuku wants to merge 1 commit into
elixir-nx:mainfrom
dkuku:feat/rope-yarn-dynamic

Conversation

@dkuku

@dkuku dkuku commented Oct 2, 2026

Copy link
Copy Markdown
Contributor

YaRN checkpoints currently cannot be loaded, and dynamic RoPE raises the entire base to the scaling exponent instead of applying the exponent only to its multiplier. This adds YaRN and corrects dynamic scaling to use the highest position ID rather than the padded/cache allocation length.

  • Support YaRN frequency blending, attention scaling, and truncation, with documented defaults for direct configuration.
  • Validate RoPE configuration and type aliases while preserving legacy Phi-3 LongRoPE aliases.
  • Add regression tests for the formulas, malformed options, and the documented YaRN configuration.

Validation: 38 targeted tests passed on this branch independently, covering layer/configuration tests and existing Llama, Mistral, Qwen3, Gemma3Text, and Phi3 tests. Formatting, compilation with --warnings-as-errors, and unused-dependency checks pass on the complete stack. A dependent embedding PR adds end-to-end Python Transformers parity fixtures for the RoPE variants.

The full test suite was not run. No associated issue is being closed.

AI disclosure: this contribution was developed with substantial LLM/coding-agent assistance, including implementation, tests, and review. Validation above was executed by the coding agent.

@dkuku
dkuku marked this pull request as draft October 2, 2026 20:25
@dkuku dkuku closed this Oct 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant