Skip to content
This repository was archived by the owner on Sep 10, 2025. It is now read-only.

Update Flamingo Builder to use Llama3ScaledRoPE instead of RotaryPositionalEmbeddings #1202

Merged
merged 2 commits into from
Sep 25, 2024
Merged

Update Flamingo Builder to use Llama3ScaledRoPE instead of RotaryPositionalEmbeddings #1202

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
175c967
update rope class
Gasoonjia Sep 25, 2024
078ac23
remove old rope class
Gasoonjia Sep 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions torchchat/cli/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@

from torchchat.model import Model, ModelArgs, ModelType

from torchtune.modules.position_embeddings import RotaryPositionalEmbeddings
from torchtune.models.llama3_1._position_embeddings import Llama3ScaledRoPE

from torchchat.model_config.model_config import resolve_model_config
from torchchat.utils.build_utils import (
Expand Down Expand Up @@ -402,7 +402,7 @@ def _load_model_default(builder_args: BuilderArgs) -> Model:
max_seq_len = decoder_config['max_seq_len']
rope_base = decoder_config['rope_base']
for submodule in model.modules():
if isinstance(submodule, RotaryPositionalEmbeddings):
if isinstance(submodule, Llama3ScaledRoPE):
submodule.__init__(head_dim, max_seq_len, rope_base)
Copy link
Contributor

@Jack-Khuu Jack-Khuu Sep 25, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: We should be populating the rest of the RoPE fields based on the config. This can be a separate PR

From tune:

scale_factor: int = 8,
low_freq_factor: int = 1,
high_freq_factor: int = 4,
old_context_len: int = 8192,

state_dict = flamingo_meta_to_tune(checkpoint)
model.model.load_state_dict(state_dict, assign=True, strict=False)
Expand Down
Loading