Skip to content

Support llama3.1-8b generation #947

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jul 24, 2024
Merged

Support llama3.1-8b generation #947

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
416c4f4
add llama 3.1 8b support
Gasoonjia Jul 24, 2024
abb8e95
Merge branch 'main' of https://github.com/pytorch/torchchat
Gasoonjia Jul 24, 2024
070ea67
replace math.pi with torch.pi
Gasoonjia Jul 24, 2024
1f309dd
add 3.1 8b base and 70b
Gasoonjia Jul 24, 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
1 change: 1 addition & 0 deletions build/known_model_params/Meta-Llama-3.1-70B.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
{"dim": 8192, "ffn_dim_multiplier": 1.3, "multiple_of": 4096, "n_heads": 64, "n_local_heads": 8, "n_layers": 80, "rope_base": 500000.0, "vocab_size": 128256, "use_tiktoken": true, "norm_eps": 1e-05, "use_scaled_rope": true}
1 change: 1 addition & 0 deletions build/known_model_params/Meta-Llama-3.1-8B.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
{"dim": 4096, "ffn_dim_multiplier": 1.3, "multiple_of": 1024, "n_heads": 32, "n_local_heads": 8, "n_layers": 32, "rope_base": 500000.0, "vocab_size": 128256, "use_tiktoken": true, "norm_eps": 1e-05, "use_scaled_rope": true}
31 changes: 30 additions & 1 deletion build/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
# LICENSE file in the root directory of this source tree.
import json
import os

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional
Expand Down Expand Up @@ -38,6 +39,7 @@ class ModelArgs:
ffn_dim_multiplier: Optional[int] = None
use_tiktoken: bool = False
max_seq_length: int = 8192
use_scaled_rope: bool = False

def __post_init__(self):
if self.n_local_heads == -1:
Expand Down Expand Up @@ -178,6 +180,7 @@ def setup_caches(self, max_batch_size, max_seq_length):
self.config.dim // self.config.n_heads,
self.config.block_size * 2,
self.config.rope_base,
use_scaled = self.config.use_scaled_rope,
)
self.register_buffer("freqs_cis", freqs_cis, persistent=True)
causal_mask = torch.tril(
Expand Down Expand Up @@ -361,15 +364,41 @@ def forward(self, x: Tensor) -> Tensor:
return output * self.weight


def apply_scaling(freqs: torch.Tensor):
# Values obtained from grid search
scale_factor = 8
low_freq_factor = 1
high_freq_factor = 4
old_context_len = 8192 # original llama3 length

low_freq_wavelen = old_context_len / low_freq_factor
high_freq_wavelen = old_context_len / high_freq_factor
new_freqs = []
Copy link
Contributor

@malfet malfet Jul 24, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

Hmm, it feels like one can write this logic in a much more pytorchy-style:

new_freqs = 360 / freq.rad2deg()
new_freqs[new_freqs > high_freq_wavelen ] /= scale_factor
new_freqs[new_freqs < low_freq_wavelen * scale_factor] = apply_smooth_here

for freq in freqs:
wavelen = 2 * torch.pi / freq
if wavelen < high_freq_wavelen:
new_freqs.append(freq)
elif wavelen > low_freq_wavelen:
new_freqs.append(freq / scale_factor)
else:
assert low_freq_wavelen != high_freq_wavelen
smooth = (old_context_len / wavelen - low_freq_factor) / (
high_freq_factor - low_freq_factor
)
new_freqs.append((1 - smooth) * freq / scale_factor + smooth * freq)
return torch.tensor(new_freqs, dtype=freqs.dtype, device=freqs.device)

def precompute_freqs_cis(
n_elem: int, seq_len: int, base: int = 10000, dtype=None
n_elem: int, seq_len: int, base: int = 10000, dtype=None, use_scaled: bool = False
) -> Tensor:
if not dtype:
dtype = get_precision()
freqs = 1.0 / (
base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem)
)
t = torch.arange(seq_len, device=freqs.device)
if use_scaled:
freqs = apply_scaling(freqs)
freqs = torch.outer(t, freqs)
freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
Expand Down
17 changes: 17 additions & 0 deletions config/data/models.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,23 @@
"distribution_path": "meta-llama/Meta-Llama-3-70B-Instruct",
"transformer_params_key": "Meta-Llama-3-70B"
},
"meta-llama/Meta-Llama-3.1-8B": {
"aliases": ["llama3.1-base"],
"distribution_channel": "HuggingFaceSnapshot",
"distribution_path": "meta-llama/Meta-Llama-3.1-8B"
},
"meta-llama/Meta-Llama-3.1-8B-Instruct": {
Copy link
Contributor

Choose a reason for hiding this comment

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

nit: technically you can also add support for the pre-trained model (meta-llama/Meta-Llama-3.1-8B) and the llama guard (meta-llama/Llama-Guard-3-8B). Not a requirement though

Copy link
Contributor

Choose a reason for hiding this comment

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

Probably want support for the pre-trained (non-Instruct) version to match with our support for Llama 3.0?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

oh yeah of course. This PR is focused on enabling llama3.1 in torchchat, so didn't cover all possible models. Will have another PR to handle that.

"aliases": ["llama3.1", "llama3.1-chat", "llama3.1-instruct"],
"distribution_channel": "HuggingFaceSnapshot",
"distribution_path": "meta-llama/Meta-Llama-3.1-8B-Instruct",
"transformer_params_key": "Meta-Llama-3.1-8B"
},
"meta-llama/Meta-Llama-3.1-70B-Instruct": {
"aliases": ["llama3.1-70b"],
"distribution_channel": "HuggingFaceSnapshot",
"distribution_path": "meta-llama/Meta-Llama-3.1-70B-Instruct",
"transformer_params_key": "Meta-Llama-3.1-70B"
},
"meta-llama/CodeLlama-7b-Python-hf": {
"aliases": ["codellama", "codellama-7b"],
"distribution_channel": "HuggingFaceSnapshot",
Expand Down
Loading