Commit
•
4379f92
1
Parent(s):
be43a05
added model util files
Browse files- tsai_gpt/__init__.py +15 -0
- tsai_gpt/config.py +1181 -0
- tsai_gpt/model.py +342 -0
- tsai_gpt/packed_dataset.py +235 -0
- tsai_gpt/rmsnorm.py +26 -0
- tsai_gpt/speed_monitor.py +425 -0
- tsai_gpt/tokenizer.py +103 -0
- tsai_gpt/utils.py +347 -0
tsai_gpt/__init__.py
ADDED
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from tsai_gpt.model import GPT
|
2 |
+
from tsai_gpt.config import Config
|
3 |
+
from tsai_gpt.tokenizer import Tokenizer
|
4 |
+
|
5 |
+
from lightning_utilities.core.imports import RequirementCache
|
6 |
+
|
7 |
+
_LIGHTNING_AVAILABLE = RequirementCache("lightning>=2.1.0.dev0")
|
8 |
+
if not bool(_LIGHTNING_AVAILABLE):
|
9 |
+
raise ImportError(
|
10 |
+
"Lit-GPT requires lightning==2.1. Please run:\n"
|
11 |
+
f" pip uninstall -y lightning; pip install -r requirements.txt\n{str(_LIGHTNING_AVAILABLE)}"
|
12 |
+
)
|
13 |
+
|
14 |
+
|
15 |
+
__all__ = ["GPT", "Config", "Tokenizer"]
|
tsai_gpt/config.py
ADDED
@@ -0,0 +1,1181 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import json
|
2 |
+
from copy import deepcopy
|
3 |
+
from dataclasses import dataclass, field
|
4 |
+
from pathlib import Path
|
5 |
+
from typing import Any, Literal, Optional, Type, Union
|
6 |
+
|
7 |
+
import torch
|
8 |
+
from typing_extensions import Self
|
9 |
+
|
10 |
+
import tsai_gpt.model
|
11 |
+
from tsai_gpt.utils import find_multiple
|
12 |
+
|
13 |
+
|
14 |
+
@dataclass
|
15 |
+
class Config:
|
16 |
+
name: str = ""
|
17 |
+
hf_config: dict = field(default_factory=dict)
|
18 |
+
block_size: int = 4096
|
19 |
+
vocab_size: int = 50254
|
20 |
+
padding_multiple: int = 512
|
21 |
+
padded_vocab_size: Optional[int] = None
|
22 |
+
n_layer: int = 16
|
23 |
+
n_head: int = 32
|
24 |
+
n_embd: int = 4096
|
25 |
+
rotary_percentage: float = 0.25
|
26 |
+
parallel_residual: bool = True
|
27 |
+
bias: bool = True
|
28 |
+
lm_head_bias: bool = False
|
29 |
+
# to use multi-head attention (MHA), set this to `n_head` (default)
|
30 |
+
# to use multi-query attention (MQA), set this to 1
|
31 |
+
# to use grouped-query attention (GQA), set this to a value in between
|
32 |
+
# Example with `n_head=4`
|
33 |
+
# ┌───┐┌───┐┌───┐┌───┐ ┌───┐ ┌───┐ ┌───┐
|
34 |
+
# │ v ││ v ││ v ││ v │ │ v │ │ v │ │ v │
|
35 |
+
# └───┘└───┘└───┘└───┘ └───┘ └───┘ └───┘
|
36 |
+
# │ │ │ │ │ │ │
|
37 |
+
# ┌───┐┌───┐┌───┐┌───┐ ┌───┐ ┌───┐ ┌───┐
|
38 |
+
# │ k ││ k ││ k ││ k │ │ k │ │ k │ │ k │
|
39 |
+
# └───┘└───┘└───┘└───┘ └───┘ └───┘ └───┘
|
40 |
+
# │ │ │ │ ┌──┴──┐ ┌──┴──┐ ┌────┬──┴─┬────┐
|
41 |
+
# ┌───┐┌───┐┌───┐┌───┐ ┌───┐┌───┐┌───┐┌───┐ ┌───┐┌───┐┌───┐┌───┐
|
42 |
+
# │ q ││ q ││ q ││ q │ │ q ││ q ││ q ││ q │ │ q ││ q ││ q ││ q │
|
43 |
+
# └───┘└───┘└───┘└───┘ └───┘└───┘└───┘└───┘ └───┘└───┘└───┘└───┘
|
44 |
+
# ◀──────────────────▶ ◀──────────────────▶ ◀──────────────────▶
|
45 |
+
# MHA GQA MQA
|
46 |
+
# n_query_groups=4 n_query_groups=2 n_query_groups=1
|
47 |
+
#
|
48 |
+
# credit https://arxiv.org/pdf/2305.13245.pdf
|
49 |
+
n_query_groups: Optional[int] = None
|
50 |
+
shared_attention_norm: bool = False
|
51 |
+
_norm_class: Literal["LayerNorm", "RMSNorm"] = "LayerNorm"
|
52 |
+
norm_eps: float = 1e-5
|
53 |
+
_mlp_class: Literal["GptNeoxMLP", "LLaMAMLP"] = "GptNeoxMLP"
|
54 |
+
gelu_approximate: str = "none"
|
55 |
+
intermediate_size: Optional[int] = None
|
56 |
+
rope_condense_ratio: int = 1
|
57 |
+
rope_base: int = 10000
|
58 |
+
|
59 |
+
def __post_init__(self):
|
60 |
+
if not self.name:
|
61 |
+
self.name = self.hf_config.get("name", self.name)
|
62 |
+
|
63 |
+
assert self.n_embd % self.n_head == 0
|
64 |
+
self.head_size = self.n_embd // self.n_head
|
65 |
+
|
66 |
+
# vocab size should be a power of 2 to be optimal on hardware. compute the closest value
|
67 |
+
if self.padded_vocab_size is None:
|
68 |
+
self.padded_vocab_size = find_multiple(self.vocab_size, self.padding_multiple)
|
69 |
+
else:
|
70 |
+
# vocab size shouldn't be larger than padded vocab size
|
71 |
+
self.vocab_size = min(self.vocab_size, self.padded_vocab_size)
|
72 |
+
|
73 |
+
# compute the number of query groups
|
74 |
+
if self.n_query_groups is not None:
|
75 |
+
assert self.n_head % self.n_query_groups == 0
|
76 |
+
else:
|
77 |
+
self.n_query_groups = self.n_head
|
78 |
+
|
79 |
+
# compute the intermediate size for MLP if not set
|
80 |
+
if self.intermediate_size is None:
|
81 |
+
if self._mlp_class == "LLaMAMLP":
|
82 |
+
raise ValueError("The config needs to set the `intermediate_size`")
|
83 |
+
self.intermediate_size = 4 * self.n_embd
|
84 |
+
|
85 |
+
self.rope_n_elem = int(self.rotary_percentage * self.head_size)
|
86 |
+
|
87 |
+
@classmethod
|
88 |
+
def from_name(cls, name: str, **kwargs: Any) -> Self:
|
89 |
+
if name not in name_to_config:
|
90 |
+
# search through all `config['hf_config']['name']`
|
91 |
+
conf_dict = next(config for config in configs if name == config["hf_config"]["name"])
|
92 |
+
else:
|
93 |
+
conf_dict = name_to_config[name]
|
94 |
+
|
95 |
+
conf_dict = conf_dict.copy()
|
96 |
+
if "condense_ratio" in kwargs: # legacy name
|
97 |
+
kwargs["rope_condense_ratio"] = kwargs.pop("condense_ratio")
|
98 |
+
conf_dict.update(kwargs)
|
99 |
+
return cls(**conf_dict)
|
100 |
+
|
101 |
+
@classmethod
|
102 |
+
def from_json(cls, path: Union[str, Path], **kwargs: Any) -> Self:
|
103 |
+
with open(path, encoding="utf-8") as fp:
|
104 |
+
json_kwargs = json.load(fp)
|
105 |
+
if "condense_ratio" in json_kwargs: # legacy name
|
106 |
+
json_kwargs["rope_condense_ratio"] = json_kwargs.pop("condense_ratio")
|
107 |
+
if "condense_ratio" in kwargs: # legacy name
|
108 |
+
kwargs["rope_condense_ratio"] = kwargs.pop("condense_ratio")
|
109 |
+
if "org" in json_kwargs: # legacy name
|
110 |
+
json_kwargs["hf_config"] = {"name": json_kwargs["name"], "org": json_kwargs.pop("org")}
|
111 |
+
if "org" in kwargs: # legacy name
|
112 |
+
kwargs["hf_config"] = {"name": kwargs.get("name", json_kwargs["name"]), "org": kwargs.pop("org")}
|
113 |
+
json_kwargs.update(kwargs)
|
114 |
+
return cls(**json_kwargs)
|
115 |
+
|
116 |
+
@property
|
117 |
+
def mlp_class(self) -> Type:
|
118 |
+
# `self._mlp_class` cannot be the type to keep the config json serializable
|
119 |
+
return getattr(tsai_gpt.model, self._mlp_class)
|
120 |
+
|
121 |
+
@property
|
122 |
+
def norm_class(self) -> Type:
|
123 |
+
# `self._norm_class` cannot be the type to keep the config json serializable
|
124 |
+
if self._norm_class == "RMSNorm":
|
125 |
+
from tsai_gpt.rmsnorm import RMSNorm
|
126 |
+
|
127 |
+
return RMSNorm
|
128 |
+
return getattr(torch.nn, self._norm_class)
|
129 |
+
|
130 |
+
|
131 |
+
########################
|
132 |
+
# Stability AI StableLM
|
133 |
+
########################
|
134 |
+
configs = [
|
135 |
+
# https://huggingface.co/stabilityai/stablelm-base-alpha-3b/blob/main/config.json
|
136 |
+
dict(name="stablelm-base-alpha-3b", hf_config=dict(org="stabilityai", name="stablelm-base-alpha-3b")),
|
137 |
+
# https://huggingface.co/stabilityai/stablelm-base-alpha-7b/blob/main/config.json
|
138 |
+
dict(
|
139 |
+
name="stablelm-base-alpha-7b",
|
140 |
+
hf_config=dict(org="stabilityai", name="stablelm-base-alpha-7b"),
|
141 |
+
n_head=48,
|
142 |
+
n_embd=6144,
|
143 |
+
padding_multiple=256,
|
144 |
+
),
|
145 |
+
# https://huggingface.co/stabilityai/stablelm-tuned-alpha-3b/blob/main/config.json
|
146 |
+
dict(name="stablelm-tuned-alpha-3b", hf_config=dict(org="stabilityai", name="stablelm-tuned-alpha-3b"), n_head=32),
|
147 |
+
# https://huggingface.co/stabilityai/stablelm-tuned-alpha-7b/blob/main/config.json
|
148 |
+
dict(
|
149 |
+
name="stablelm-tuned-alpha-7b",
|
150 |
+
hf_config=dict(org="stabilityai", name="stablelm-tuned-alpha-7b"),
|
151 |
+
n_head=48,
|
152 |
+
n_embd=6144,
|
153 |
+
padding_multiple=256,
|
154 |
+
),
|
155 |
+
]
|
156 |
+
|
157 |
+
####################
|
158 |
+
# EleutherAI Pythia
|
159 |
+
####################
|
160 |
+
pythia = [
|
161 |
+
# https://huggingface.co/EleutherAI/pythia-70m/blob/main/config.json
|
162 |
+
dict(
|
163 |
+
name="pythia-70m",
|
164 |
+
hf_config=dict(org="EleutherAI", name="pythia-70m"),
|
165 |
+
block_size=2048,
|
166 |
+
n_layer=6,
|
167 |
+
n_embd=512,
|
168 |
+
n_head=8,
|
169 |
+
padding_multiple=128,
|
170 |
+
),
|
171 |
+
# https://huggingface.co/EleutherAI/pythia-160m/blob/main/config.json
|
172 |
+
dict(
|
173 |
+
name="pythia-160m",
|
174 |
+
hf_config=dict(org="EleutherAI", name="pythia-160m"),
|
175 |
+
block_size=2048,
|
176 |
+
n_layer=12,
|
177 |
+
n_embd=768,
|
178 |
+
n_head=12,
|
179 |
+
padding_multiple=128,
|
180 |
+
),
|
181 |
+
# https://huggingface.co/EleutherAI/pythia-410m/blob/main/config.json
|
182 |
+
dict(
|
183 |
+
name="pythia-410m",
|
184 |
+
hf_config=dict(org="EleutherAI", name="pythia-410m"),
|
185 |
+
block_size=2048,
|
186 |
+
n_layer=24,
|
187 |
+
n_embd=1024,
|
188 |
+
n_head=16,
|
189 |
+
padding_multiple=128,
|
190 |
+
),
|
191 |
+
# https://huggingface.co/EleutherAI/pythia-1b/blob/main/config.json
|
192 |
+
dict(
|
193 |
+
name="pythia-1b",
|
194 |
+
hf_config=dict(org="EleutherAI", name="pythia-1b"),
|
195 |
+
block_size=2048,
|
196 |
+
n_embd=2048,
|
197 |
+
n_head=8,
|
198 |
+
padding_multiple=128,
|
199 |
+
),
|
200 |
+
# https://huggingface.co/EleutherAI/pythia-1.4b/blob/main/config.json
|
201 |
+
dict(
|
202 |
+
name="pythia-1.4b",
|
203 |
+
hf_config=dict(org="EleutherAI", name="pythia-1.4b"),
|
204 |
+
block_size=2048,
|
205 |
+
n_layer=24,
|
206 |
+
n_embd=2048,
|
207 |
+
n_head=16,
|
208 |
+
padding_multiple=128,
|
209 |
+
),
|
210 |
+
# https://huggingface.co/EleutherAI/pythia-2.8b/blob/main/config.json
|
211 |
+
dict(
|
212 |
+
name="pythia-2.8b",
|
213 |
+
hf_config=dict(org="EleutherAI", name="pythia-2.8b"),
|
214 |
+
block_size=2048,
|
215 |
+
n_layer=32,
|
216 |
+
n_embd=2560,
|
217 |
+
padding_multiple=128,
|
218 |
+
),
|
219 |
+
# https://huggingface.co/EleutherAI/pythia-6.9b/blob/main/config.json
|
220 |
+
dict(
|
221 |
+
name="pythia-6.9b",
|
222 |
+
hf_config=dict(org="EleutherAI", name="pythia-6.9b"),
|
223 |
+
block_size=2048,
|
224 |
+
n_layer=32,
|
225 |
+
padding_multiple=256,
|
226 |
+
),
|
227 |
+
# https://huggingface.co/EleutherAI/pythia-12b/blob/main/config.json
|
228 |
+
dict(
|
229 |
+
name="pythia-12b",
|
230 |
+
hf_config=dict(org="EleutherAI", name="pythia-12b"),
|
231 |
+
block_size=2048,
|
232 |
+
n_layer=36,
|
233 |
+
n_embd=5120,
|
234 |
+
n_head=40,
|
235 |
+
),
|
236 |
+
]
|
237 |
+
configs.extend(pythia)
|
238 |
+
for c in pythia:
|
239 |
+
copy = c.copy()
|
240 |
+
copy["name"] = f"{c['name']}-deduped"
|
241 |
+
copy["hf_config"]["name"] = f"{c['hf_config']['name']}-deduped"
|
242 |
+
configs.append(copy)
|
243 |
+
|
244 |
+
|
245 |
+
####################################
|
246 |
+
# togethercomputer RedPajama INCITE
|
247 |
+
####################################
|
248 |
+
redpajama_incite = [
|
249 |
+
# https://huggingface.co/togethercomputer/RedPajama-INCITE-Base-3B-v1/blob/main/config.json
|
250 |
+
dict(
|
251 |
+
name="RedPajama-INCITE-{}-3B-v1",
|
252 |
+
hf_config=dict(org="togethercomputer", name="RedPajama-INCITE-{}-3B-v1"),
|
253 |
+
block_size=2048,
|
254 |
+
n_layer=32,
|
255 |
+
n_embd=2560,
|
256 |
+
padding_multiple=256,
|
257 |
+
rotary_percentage=1.0,
|
258 |
+
parallel_residual=False,
|
259 |
+
),
|
260 |
+
# https://huggingface.co/togethercomputer/RedPajama-INCITE-7B-Base/blob/main/config.json
|
261 |
+
dict(
|
262 |
+
name="RedPajama-INCITE-7B-{}",
|
263 |
+
hf_config=dict(org="togethercomputer", name="RedPajama-INCITE-7B-{}"),
|
264 |
+
block_size=2048,
|
265 |
+
n_layer=32,
|
266 |
+
padding_multiple=256,
|
267 |
+
rotary_percentage=1.0,
|
268 |
+
parallel_residual=False,
|
269 |
+
),
|
270 |
+
# this redirects to the checkpoint above. kept for those who had the old weights already downloaded
|
271 |
+
dict(
|
272 |
+
name="RedPajama-INCITE-{}-7B-v0.1",
|
273 |
+
hf_config=dict(org="togethercomputer", name="RedPajama-INCITE-{}-7B-v0.1"),
|
274 |
+
block_size=2048,
|
275 |
+
n_layer=32,
|
276 |
+
padding_multiple=256,
|
277 |
+
rotary_percentage=1.0,
|
278 |
+
parallel_residual=False,
|
279 |
+
),
|
280 |
+
]
|
281 |
+
for c in redpajama_incite:
|
282 |
+
for kind in ("Base", "Chat", "Instruct"):
|
283 |
+
copy = c.copy()
|
284 |
+
copy["name"] = c["name"].format(kind)
|
285 |
+
copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
|
286 |
+
configs.append(copy)
|
287 |
+
|
288 |
+
|
289 |
+
#################
|
290 |
+
# TII UAE Falcon
|
291 |
+
#################
|
292 |
+
falcon = [
|
293 |
+
# https://huggingface.co/tiiuae/falcon-7b/blob/main/config.json
|
294 |
+
dict(
|
295 |
+
name="falcon-7b{}",
|
296 |
+
hf_config=dict(org="tiiuae", name="falcon-7b{}"),
|
297 |
+
block_size=2048,
|
298 |
+
vocab_size=65024,
|
299 |
+
padded_vocab_size=65024,
|
300 |
+
n_layer=32,
|
301 |
+
n_head=71,
|
302 |
+
n_embd=4544,
|
303 |
+
rotary_percentage=1.0,
|
304 |
+
n_query_groups=1,
|
305 |
+
bias=False,
|
306 |
+
# this is not in the config, but in the original model implementation, only for this config
|
307 |
+
shared_attention_norm=True,
|
308 |
+
),
|
309 |
+
# https://huggingface.co/tiiuae/falcon-40b/blob/main/config.json
|
310 |
+
dict(
|
311 |
+
name="falcon-40b{}",
|
312 |
+
hf_config=dict(org="tiiuae", name="falcon-40b{}"),
|
313 |
+
block_size=2048,
|
314 |
+
vocab_size=65024,
|
315 |
+
padded_vocab_size=65024,
|
316 |
+
n_layer=60,
|
317 |
+
n_head=128,
|
318 |
+
n_embd=8192,
|
319 |
+
rotary_percentage=1.0,
|
320 |
+
n_query_groups=8,
|
321 |
+
bias=False,
|
322 |
+
),
|
323 |
+
]
|
324 |
+
for c in falcon:
|
325 |
+
for kind in ("", "-instruct"):
|
326 |
+
copy = c.copy()
|
327 |
+
copy["name"] = c["name"].format(kind)
|
328 |
+
copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
|
329 |
+
configs.append(copy)
|
330 |
+
|
331 |
+
# https://huggingface.co/tiiuae/falcon-180b/blob/main/config.json
|
332 |
+
falcon180b = dict(
|
333 |
+
name="falcon-180B{}",
|
334 |
+
hf_config=dict(org="tiiuae", name="falcon-180B{}"),
|
335 |
+
block_size=2048,
|
336 |
+
vocab_size=65024,
|
337 |
+
padded_vocab_size=65024,
|
338 |
+
n_layer=80,
|
339 |
+
n_head=232,
|
340 |
+
n_embd=14848,
|
341 |
+
rotary_percentage=1.0,
|
342 |
+
n_query_groups=8,
|
343 |
+
bias=False,
|
344 |
+
)
|
345 |
+
|
346 |
+
for kind in ("", "-chat"):
|
347 |
+
copy = falcon180b.copy()
|
348 |
+
copy["name"] = falcon180b["name"].format(kind)
|
349 |
+
copy["hf_config"]["name"] = falcon180b["hf_config"]["name"].format(kind)
|
350 |
+
configs.append(copy)
|
351 |
+
|
352 |
+
|
353 |
+
#############################
|
354 |
+
# OpenLM Research Open LLaMA
|
355 |
+
#############################
|
356 |
+
open_LLaMA = [
|
357 |
+
# https://huggingface.co/openlm-research/open_llama_3b/blob/main/config.json
|
358 |
+
dict(
|
359 |
+
name="open_llama_3b",
|
360 |
+
hf_config=dict(org="openlm-research", name="open_llama_3b"),
|
361 |
+
block_size=2048,
|
362 |
+
vocab_size=32000,
|
363 |
+
padding_multiple=64,
|
364 |
+
n_layer=26,
|
365 |
+
n_embd=3200,
|
366 |
+
rotary_percentage=1.0,
|
367 |
+
parallel_residual=False,
|
368 |
+
bias=False,
|
369 |
+
_norm_class="RMSNorm",
|
370 |
+
norm_eps=1e-6,
|
371 |
+
_mlp_class="LLaMAMLP",
|
372 |
+
intermediate_size=8640,
|
373 |
+
),
|
374 |
+
# https://huggingface.co/openlm-research/open_llama_7b/blob/main/config.json
|
375 |
+
dict(
|
376 |
+
name="open_llama_7b",
|
377 |
+
hf_config=dict(org="openlm-research", name="open_llama_7b"),
|
378 |
+
block_size=2048,
|
379 |
+
vocab_size=32000,
|
380 |
+
padding_multiple=64,
|
381 |
+
n_layer=32,
|
382 |
+
rotary_percentage=1.0,
|
383 |
+
parallel_residual=False,
|
384 |
+
bias=False,
|
385 |
+
_norm_class="RMSNorm",
|
386 |
+
norm_eps=1e-6,
|
387 |
+
_mlp_class="LLaMAMLP",
|
388 |
+
intermediate_size=11008,
|
389 |
+
),
|
390 |
+
# https://huggingface.co/openlm-research/open_llama_13b/blob/main/config.json
|
391 |
+
dict(
|
392 |
+
name="open_llama_13b",
|
393 |
+
hf_config=dict(org="openlm-research", name="open_llama_13b"),
|
394 |
+
block_size=2048,
|
395 |
+
vocab_size=32000,
|
396 |
+
padding_multiple=64,
|
397 |
+
n_layer=40,
|
398 |
+
n_head=40,
|
399 |
+
n_embd=5120,
|
400 |
+
rotary_percentage=1.0,
|
401 |
+
parallel_residual=False,
|
402 |
+
bias=False,
|
403 |
+
_norm_class="RMSNorm",
|
404 |
+
norm_eps=1e-6,
|
405 |
+
_mlp_class="LLaMAMLP",
|
406 |
+
intermediate_size=13824,
|
407 |
+
),
|
408 |
+
]
|
409 |
+
configs.extend(open_LLaMA)
|
410 |
+
|
411 |
+
|
412 |
+
###############
|
413 |
+
# LMSYS Vicuna
|
414 |
+
###############
|
415 |
+
vicuna = [
|
416 |
+
# https://huggingface.co/lmsys/vicuna-7b-v1.3/blob/main/config.json
|
417 |
+
dict(
|
418 |
+
name="vicuna-7b-v1.3",
|
419 |
+
hf_config=dict(org="lmsys", name="vicuna-7b-v1.3"),
|
420 |
+
block_size=2048,
|
421 |
+
vocab_size=32000,
|
422 |
+
padding_multiple=64,
|
423 |
+
n_layer=32,
|
424 |
+
rotary_percentage=1.0,
|
425 |
+
parallel_residual=False,
|
426 |
+
bias=False,
|
427 |
+
_norm_class="RMSNorm",
|
428 |
+
norm_eps=1e-6,
|
429 |
+
_mlp_class="LLaMAMLP",
|
430 |
+
intermediate_size=11008,
|
431 |
+
),
|
432 |
+
# https://huggingface.co/lmsys/vicuna-13b-v1.3/blob/main/config.json
|
433 |
+
dict(
|
434 |
+
name="vicuna-13b-v1.3",
|
435 |
+
hf_config=dict(org="lmsys", name="vicuna-13b-v1.3"),
|
436 |
+
block_size=2048,
|
437 |
+
vocab_size=32000,
|
438 |
+
padding_multiple=64,
|
439 |
+
n_layer=40,
|
440 |
+
n_head=40,
|
441 |
+
n_embd=5120,
|
442 |
+
rotary_percentage=1.0,
|
443 |
+
parallel_residual=False,
|
444 |
+
bias=False,
|
445 |
+
_norm_class="RMSNorm",
|
446 |
+
norm_eps=1e-6,
|
447 |
+
_mlp_class="LLaMAMLP",
|
448 |
+
intermediate_size=13824,
|
449 |
+
),
|
450 |
+
# https://huggingface.co/lmsys/vicuna-33b-v1.3/blob/main/config.json
|
451 |
+
dict(
|
452 |
+
name="vicuna-33b-v1.3",
|
453 |
+
hf_config=dict(org="lmsys", name="vicuna-33b-v1.3"),
|
454 |
+
block_size=2048,
|
455 |
+
vocab_size=32000,
|
456 |
+
padding_multiple=64,
|
457 |
+
n_layer=60,
|
458 |
+
n_head=52,
|
459 |
+
n_embd=6656,
|
460 |
+
rotary_percentage=1.0,
|
461 |
+
parallel_residual=False,
|
462 |
+
bias=False,
|
463 |
+
_norm_class="RMSNorm",
|
464 |
+
norm_eps=1e-6,
|
465 |
+
_mlp_class="LLaMAMLP",
|
466 |
+
intermediate_size=17920,
|
467 |
+
),
|
468 |
+
# https://huggingface.co/lmsys/vicuna-7b-v1.5/blob/main/config.json
|
469 |
+
dict(
|
470 |
+
name="vicuna-7b-v1.5",
|
471 |
+
hf_config=dict(org="lmsys", name="vicuna-7b-v1.5"),
|
472 |
+
vocab_size=32000,
|
473 |
+
padding_multiple=64,
|
474 |
+
n_layer=32,
|
475 |
+
rotary_percentage=1.0,
|
476 |
+
parallel_residual=False,
|
477 |
+
bias=False,
|
478 |
+
_norm_class="RMSNorm",
|
479 |
+
_mlp_class="LLaMAMLP",
|
480 |
+
intermediate_size=11008,
|
481 |
+
),
|
482 |
+
# https://huggingface.co/lmsys/vicuna-7b-v1.5-16k/blob/main/config.json
|
483 |
+
dict(
|
484 |
+
name="vicuna-7b-v1.5-16k",
|
485 |
+
hf_config=dict(org="lmsys", name="vicuna-7b-v1.5-16k"),
|
486 |
+
block_size=16384,
|
487 |
+
vocab_size=32000,
|
488 |
+
padding_multiple=64,
|
489 |
+
n_layer=32,
|
490 |
+
rotary_percentage=1.0,
|
491 |
+
parallel_residual=False,
|
492 |
+
bias=False,
|
493 |
+
_norm_class="RMSNorm",
|
494 |
+
_mlp_class="LLaMAMLP",
|
495 |
+
intermediate_size=11008,
|
496 |
+
rope_condense_ratio=4,
|
497 |
+
),
|
498 |
+
# https://huggingface.co/lmsys/vicuna-13b-v1.5/blob/main/config.json
|
499 |
+
dict(
|
500 |
+
name="vicuna-13b-v1.5",
|
501 |
+
hf_config=dict(org="lmsys", name="vicuna-13b-v1.5"),
|
502 |
+
vocab_size=32000,
|
503 |
+
padding_multiple=64,
|
504 |
+
n_layer=40,
|
505 |
+
n_head=40,
|
506 |
+
n_embd=5120,
|
507 |
+
rotary_percentage=1.0,
|
508 |
+
parallel_residual=False,
|
509 |
+
bias=False,
|
510 |
+
_norm_class="RMSNorm",
|
511 |
+
_mlp_class="LLaMAMLP",
|
512 |
+
intermediate_size=13824,
|
513 |
+
),
|
514 |
+
# https://huggingface.co/lmsys/vicuna-13b-v1.5-16k/blob/main/config.json
|
515 |
+
dict(
|
516 |
+
name="vicuna-13b-v1.5-16k",
|
517 |
+
hf_config=dict(org="lmsys", name="vicuna-13b-v1.5-16k"),
|
518 |
+
block_size=16384,
|
519 |
+
vocab_size=32000,
|
520 |
+
padding_multiple=64,
|
521 |
+
n_layer=40,
|
522 |
+
n_head=40,
|
523 |
+
n_embd=5120,
|
524 |
+
rotary_percentage=1.0,
|
525 |
+
parallel_residual=False,
|
526 |
+
bias=False,
|
527 |
+
_norm_class="RMSNorm",
|
528 |
+
_mlp_class="LLaMAMLP",
|
529 |
+
intermediate_size=13824,
|
530 |
+
rope_condense_ratio=4,
|
531 |
+
),
|
532 |
+
]
|
533 |
+
configs.extend(vicuna)
|
534 |
+
|
535 |
+
|
536 |
+
#################
|
537 |
+
# LMSYS LongChat
|
538 |
+
#################
|
539 |
+
long_chat = [
|
540 |
+
# https://huggingface.co/lmsys/longchat-7b-16k/blob/main/config.json
|
541 |
+
dict(
|
542 |
+
name="longchat-7b-16k",
|
543 |
+
hf_config=dict(org="lmsys", name="longchat-7b-16k"),
|
544 |
+
block_size=16384,
|
545 |
+
vocab_size=32000,
|
546 |
+
padding_multiple=64,
|
547 |
+
n_layer=32,
|
548 |
+
rotary_percentage=1.0,
|
549 |
+
parallel_residual=False,
|
550 |
+
bias=False,
|
551 |
+
_norm_class="RMSNorm",
|
552 |
+
norm_eps=1e-6,
|
553 |
+
_mlp_class="LLaMAMLP",
|
554 |
+
intermediate_size=11008,
|
555 |
+
rope_condense_ratio=8,
|
556 |
+
),
|
557 |
+
# https://huggingface.co/lmsys/longchat-13b-16k/blob/main/config.json
|
558 |
+
dict(
|
559 |
+
name="longchat-13b-16k",
|
560 |
+
hf_config=dict(org="lmsys", name="longchat-13b-16k"),
|
561 |
+
block_size=16384,
|
562 |
+
vocab_size=32000,
|
563 |
+
padding_multiple=64,
|
564 |
+
n_layer=40,
|
565 |
+
n_head=40,
|
566 |
+
n_embd=5120,
|
567 |
+
rotary_percentage=1.0,
|
568 |
+
parallel_residual=False,
|
569 |
+
bias=False,
|
570 |
+
_norm_class="RMSNorm",
|
571 |
+
norm_eps=1e-6,
|
572 |
+
_mlp_class="LLaMAMLP",
|
573 |
+
intermediate_size=13824,
|
574 |
+
rope_condense_ratio=8,
|
575 |
+
),
|
576 |
+
]
|
577 |
+
configs.extend(long_chat)
|
578 |
+
|
579 |
+
|
580 |
+
######################
|
581 |
+
# NousResearch Hermes
|
582 |
+
######################
|
583 |
+
nous_research = [
|
584 |
+
# https://huggingface.co/NousResearch/Nous-Hermes-llama-2-7b/blob/main/config.json
|
585 |
+
dict(
|
586 |
+
name="Nous-Hermes-llama-2-7b",
|
587 |
+
hf_config=dict(org="NousResearch", name="Nous-Hermes-llama-2-7b"),
|
588 |
+
padded_vocab_size=32000,
|
589 |
+
n_layer=32,
|
590 |
+
rotary_percentage=1.0,
|
591 |
+
parallel_residual=False,
|
592 |
+
bias=False,
|
593 |
+
_norm_class="RMSNorm",
|
594 |
+
norm_eps=1e-05,
|
595 |
+
_mlp_class="LLaMAMLP",
|
596 |
+
intermediate_size=11008,
|
597 |
+
),
|
598 |
+
# https://huggingface.co/NousResearch/Nous-Hermes-13B/blob/main/config.json
|
599 |
+
dict(
|
600 |
+
name="Nous-Hermes-13b",
|
601 |
+
hf_config=dict(org="NousResearch", name="Nous-Hermes-13b"),
|
602 |
+
block_size=2048,
|
603 |
+
vocab_size=32000,
|
604 |
+
padded_vocab_size=32001,
|
605 |
+
n_layer=40,
|
606 |
+
n_head=40,
|
607 |
+
n_embd=5120,
|
608 |
+
rotary_percentage=1.0,
|
609 |
+
parallel_residual=False,
|
610 |
+
bias=False,
|
611 |
+
_norm_class="RMSNorm",
|
612 |
+
norm_eps=1e-6,
|
613 |
+
_mlp_class="LLaMAMLP",
|
614 |
+
intermediate_size=13824,
|
615 |
+
),
|
616 |
+
# https://huggingface.co/NousResearch/Nous-Hermes-Llama2-13b
|
617 |
+
dict(
|
618 |
+
name="Nous-Hermes-Llama2-13b",
|
619 |
+
hf_config=dict(org="NousResearch", name="Nous-Hermes-Llama2-13b"),
|
620 |
+
vocab_size=32000,
|
621 |
+
padded_vocab_size=32032,
|
622 |
+
n_layer=40,
|
623 |
+
n_head=40,
|
624 |
+
n_embd=5120,
|
625 |
+
rotary_percentage=1.0,
|
626 |
+
parallel_residual=False,
|
627 |
+
bias=False,
|
628 |
+
_norm_class="RMSNorm",
|
629 |
+
norm_eps=1e-05,
|
630 |
+
_mlp_class="LLaMAMLP",
|
631 |
+
intermediate_size=13824,
|
632 |
+
),
|
633 |
+
]
|
634 |
+
configs.extend(nous_research)
|
635 |
+
|
636 |
+
|
637 |
+
###############
|
638 |
+
# Meta LLaMA 2
|
639 |
+
###############
|
640 |
+
llama_2 = [
|
641 |
+
# https://huggingface.co/meta-llama/Llama-2-7b-hf/blob/main/config.json
|
642 |
+
dict(
|
643 |
+
name="Llama-2-7b{}-hf",
|
644 |
+
hf_config=dict(org="meta-llama", name="Llama-2-7b{}-hf"),
|
645 |
+
vocab_size=32000,
|
646 |
+
padding_multiple=64,
|
647 |
+
n_layer=32,
|
648 |
+
rotary_percentage=1.0,
|
649 |
+
parallel_residual=False,
|
650 |
+
bias=False,
|
651 |
+
_norm_class="RMSNorm",
|
652 |
+
_mlp_class="LLaMAMLP",
|
653 |
+
intermediate_size=11008,
|
654 |
+
),
|
655 |
+
# https://huggingface.co/meta-llama/Llama-2-13b-hf/blob/main/config.json
|
656 |
+
dict(
|
657 |
+
name="Llama-2-13b{}-hf",
|
658 |
+
hf_config=dict(org="meta-llama", name="Llama-2-13b{}-hf"),
|
659 |
+
vocab_size=32000,
|
660 |
+
padding_multiple=64,
|
661 |
+
n_layer=40,
|
662 |
+
n_head=40,
|
663 |
+
n_embd=5120,
|
664 |
+
rotary_percentage=1.0,
|
665 |
+
parallel_residual=False,
|
666 |
+
bias=False,
|
667 |
+
_norm_class="RMSNorm",
|
668 |
+
_mlp_class="LLaMAMLP",
|
669 |
+
intermediate_size=13824,
|
670 |
+
),
|
671 |
+
# https://huggingface.co/meta-llama/Llama-2-70b-hf/blob/main/config.json
|
672 |
+
dict(
|
673 |
+
name="Llama-2-70b{}-hf",
|
674 |
+
hf_config=dict(org="meta-llama", name="Llama-2-70b{}-hf"),
|
675 |
+
vocab_size=32000,
|
676 |
+
padding_multiple=64,
|
677 |
+
n_layer=80,
|
678 |
+
n_head=64,
|
679 |
+
n_embd=8192,
|
680 |
+
n_query_groups=8,
|
681 |
+
rotary_percentage=1.0,
|
682 |
+
parallel_residual=False,
|
683 |
+
bias=False,
|
684 |
+
_norm_class="RMSNorm",
|
685 |
+
_mlp_class="LLaMAMLP",
|
686 |
+
intermediate_size=28672,
|
687 |
+
),
|
688 |
+
]
|
689 |
+
for c in llama_2:
|
690 |
+
for kind in ("", "-chat"):
|
691 |
+
copy = c.copy()
|
692 |
+
copy["name"] = c["name"].format(kind)
|
693 |
+
copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
|
694 |
+
configs.append(copy)
|
695 |
+
|
696 |
+
|
697 |
+
##########################
|
698 |
+
# Stability AI FreeWilly2
|
699 |
+
##########################
|
700 |
+
freewilly_2 = [
|
701 |
+
# https://huggingface.co/stabilityai/FreeWilly2/blob/main/config.json
|
702 |
+
dict(
|
703 |
+
name="FreeWilly2",
|
704 |
+
hf_config=dict(org="stabilityai", name="FreeWilly2"),
|
705 |
+
vocab_size=32000,
|
706 |
+
padding_multiple=64,
|
707 |
+
n_layer=80,
|
708 |
+
n_head=64,
|
709 |
+
n_embd=8192,
|
710 |
+
n_query_groups=8,
|
711 |
+
rotary_percentage=1.0,
|
712 |
+
parallel_residual=False,
|
713 |
+
bias=False,
|
714 |
+
_norm_class="RMSNorm",
|
715 |
+
_mlp_class="LLaMAMLP",
|
716 |
+
intermediate_size=28672,
|
717 |
+
)
|
718 |
+
]
|
719 |
+
configs.extend(freewilly_2)
|
720 |
+
|
721 |
+
|
722 |
+
##################
|
723 |
+
# Meta Code Llama
|
724 |
+
##################
|
725 |
+
code_llama = [
|
726 |
+
# https://huggingface.co/codellama/CodeLlama-7b-hf/blob/main/config.json
|
727 |
+
dict(
|
728 |
+
name="CodeLlama-7b-hf",
|
729 |
+
hf_config=dict(org="codellama", name="CodeLlama-7b-hf"),
|
730 |
+
block_size=16384,
|
731 |
+
vocab_size=32016,
|
732 |
+
padding_multiple=16,
|
733 |
+
n_layer=32,
|
734 |
+
rotary_percentage=1.0,
|
735 |
+
parallel_residual=False,
|
736 |
+
bias=False,
|
737 |
+
_norm_class="RMSNorm",
|
738 |
+
norm_eps=1e-05,
|
739 |
+
_mlp_class="LLaMAMLP",
|
740 |
+
intermediate_size=11008,
|
741 |
+
rope_base=1000000,
|
742 |
+
),
|
743 |
+
# https://huggingface.co/codellama/CodeLlama-13b-hf/blob/main/config.json
|
744 |
+
dict(
|
745 |
+
name="CodeLlama-13b-hf",
|
746 |
+
hf_config=dict(org="codellama", name="CodeLlama-13b-hf"),
|
747 |
+
block_size=16384,
|
748 |
+
vocab_size=32016,
|
749 |
+
padding_multiple=16,
|
750 |
+
n_layer=40,
|
751 |
+
n_head=40,
|
752 |
+
n_embd=5120,
|
753 |
+
rotary_percentage=1.0,
|
754 |
+
parallel_residual=False,
|
755 |
+
bias=False,
|
756 |
+
_norm_class="RMSNorm",
|
757 |
+
norm_eps=1e-05,
|
758 |
+
_mlp_class="LLaMAMLP",
|
759 |
+
intermediate_size=13824,
|
760 |
+
rope_base=1000000,
|
761 |
+
),
|
762 |
+
# https://huggingface.co/codellama/CodeLlama-34b-hf/blob/main/config.json
|
763 |
+
dict(
|
764 |
+
name="CodeLlama-34b-hf",
|
765 |
+
hf_config=dict(org="codellama", name="CodeLlama-34b-hf"),
|
766 |
+
block_size=16384,
|
767 |
+
vocab_size=32000,
|
768 |
+
padding_multiple=64,
|
769 |
+
n_layer=48,
|
770 |
+
n_head=64,
|
771 |
+
n_embd=8192,
|
772 |
+
n_query_groups=8,
|
773 |
+
rotary_percentage=1.0,
|
774 |
+
parallel_residual=False,
|
775 |
+
bias=False,
|
776 |
+
_norm_class="RMSNorm",
|
777 |
+
norm_eps=1e-05,
|
778 |
+
_mlp_class="LLaMAMLP",
|
779 |
+
intermediate_size=22016,
|
780 |
+
rope_base=1000000,
|
781 |
+
),
|
782 |
+
# https://huggingface.co/codellama/CodeLlama-7b-Python-hf/blob/main/config.json
|
783 |
+
dict(
|
784 |
+
name="CodeLlama-7b-Python-hf",
|
785 |
+
hf_config=dict(org="codellama", name="CodeLlama-7b-Python-hf"),
|
786 |
+
block_size=16384,
|
787 |
+
vocab_size=32000,
|
788 |
+
padding_multiple=64,
|
789 |
+
n_layer=32,
|
790 |
+
rotary_percentage=1.0,
|
791 |
+
parallel_residual=False,
|
792 |
+
bias=False,
|
793 |
+
_norm_class="RMSNorm",
|
794 |
+
norm_eps=1e-05,
|
795 |
+
_mlp_class="LLaMAMLP",
|
796 |
+
intermediate_size=11008,
|
797 |
+
rope_base=1000000,
|
798 |
+
),
|
799 |
+
# https://huggingface.co/codellama/CodeLlama-13b-Python-hf/blob/main/config.json
|
800 |
+
dict(
|
801 |
+
name="CodeLlama-13b-Python-hf",
|
802 |
+
hf_config=dict(org="codellama", name="CodeLlama-13b-Python-hf"),
|
803 |
+
block_size=16384,
|
804 |
+
vocab_size=32000,
|
805 |
+
padding_multiple=64,
|
806 |
+
n_layer=40,
|
807 |
+
n_head=40,
|
808 |
+
n_embd=5120,
|
809 |
+
rotary_percentage=1.0,
|
810 |
+
parallel_residual=False,
|
811 |
+
bias=False,
|
812 |
+
_norm_class="RMSNorm",
|
813 |
+
norm_eps=1e-05,
|
814 |
+
_mlp_class="LLaMAMLP",
|
815 |
+
intermediate_size=13824,
|
816 |
+
rope_base=1000000,
|
817 |
+
),
|
818 |
+
# https://huggingface.co/codellama/CodeLlama-34b-Python-hf/blob/main/config.json
|
819 |
+
dict(
|
820 |
+
name="CodeLlama-34b-Python-hf",
|
821 |
+
hf_config=dict(org="codellama", name="CodeLlama-34b-Python-hf"),
|
822 |
+
block_size=16384,
|
823 |
+
vocab_size=32000,
|
824 |
+
padding_multiple=64,
|
825 |
+
n_layer=48,
|
826 |
+
n_head=64,
|
827 |
+
n_embd=8192,
|
828 |
+
n_query_groups=8,
|
829 |
+
rotary_percentage=1.0,
|
830 |
+
parallel_residual=False,
|
831 |
+
bias=False,
|
832 |
+
_norm_class="RMSNorm",
|
833 |
+
norm_eps=1e-05,
|
834 |
+
_mlp_class="LLaMAMLP",
|
835 |
+
intermediate_size=22016,
|
836 |
+
rope_base=1000000,
|
837 |
+
),
|
838 |
+
# https://huggingface.co/codellama/CodeLlama-7b-Instruct-hf/tree/main/config.json
|
839 |
+
dict(
|
840 |
+
name="CodeLlama-7b-Instruct-hf",
|
841 |
+
hf_config=dict(org="codellama", name="CodeLlama-7b-Instruct-hf"),
|
842 |
+
block_size=16384,
|
843 |
+
vocab_size=32016,
|
844 |
+
padding_multiple=16,
|
845 |
+
n_layer=32,
|
846 |
+
rotary_percentage=1.0,
|
847 |
+
parallel_residual=False,
|
848 |
+
bias=False,
|
849 |
+
_norm_class="RMSNorm",
|
850 |
+
norm_eps=1e-05,
|
851 |
+
_mlp_class="LLaMAMLP",
|
852 |
+
intermediate_size=11008,
|
853 |
+
rope_base=1000000,
|
854 |
+
),
|
855 |
+
# https://huggingface.co/codellama/CodeLlama-13b-Instruct-hf/blob/main/config.json
|
856 |
+
dict(
|
857 |
+
name="CodeLlama-13b-Instruct-hf",
|
858 |
+
hf_config=dict(org="codellama", name="CodeLlama-13b-Instruct-hf"),
|
859 |
+
block_size=2048,
|
860 |
+
vocab_size=32016,
|
861 |
+
padding_multiple=16,
|
862 |
+
n_layer=40,
|
863 |
+
n_head=40,
|
864 |
+
n_embd=5120,
|
865 |
+
rotary_percentage=1.0,
|
866 |
+
parallel_residual=False,
|
867 |
+
bias=False,
|
868 |
+
_norm_class="RMSNorm",
|
869 |
+
norm_eps=1e-05,
|
870 |
+
_mlp_class="LLaMAMLP",
|
871 |
+
intermediate_size=13824,
|
872 |
+
rope_base=1000000,
|
873 |
+
),
|
874 |
+
# https://huggingface.co/codellama/CodeLlama-34b-Instruct-hf/blob/main/config.json
|
875 |
+
dict(
|
876 |
+
name="CodeLlama-34b-Instruct-hf",
|
877 |
+
hf_config=dict(org="codellama", name="CodeLlama-34b-Instruct-hf"),
|
878 |
+
block_size=16384,
|
879 |
+
vocab_size=32000,
|
880 |
+
padding_multiple=64,
|
881 |
+
n_layer=48,
|
882 |
+
n_head=64,
|
883 |
+
n_embd=8192,
|
884 |
+
n_query_groups=8,
|
885 |
+
rotary_percentage=1.0,
|
886 |
+
parallel_residual=False,
|
887 |
+
bias=False,
|
888 |
+
_norm_class="RMSNorm",
|
889 |
+
norm_eps=1e-05,
|
890 |
+
_mlp_class="LLaMAMLP",
|
891 |
+
intermediate_size=22016,
|
892 |
+
rope_base=1000000,
|
893 |
+
),
|
894 |
+
]
|
895 |
+
configs.extend(code_llama)
|
896 |
+
|
897 |
+
|
898 |
+
########################
|
899 |
+
# garage-bAInd Platypus
|
900 |
+
########################
|
901 |
+
platypus = [
|
902 |
+
# https://huggingface.co/garage-bAInd/Platypus-30B/blob/main/config.json
|
903 |
+
dict(
|
904 |
+
name="Platypus-30B",
|
905 |
+
hf_config=dict(org="garage-bAInd", name="Platypus-30B"),
|
906 |
+
block_size=2048,
|
907 |
+
padded_vocab_size=32000,
|
908 |
+
n_layer=60,
|
909 |
+
n_head=52,
|
910 |
+
n_embd=6656,
|
911 |
+
rotary_percentage=1.0,
|
912 |
+
parallel_residual=False,
|
913 |
+
bias=False,
|
914 |
+
_norm_class="RMSNorm",
|
915 |
+
norm_eps=1e-06,
|
916 |
+
_mlp_class="LLaMAMLP",
|
917 |
+
intermediate_size=17920,
|
918 |
+
),
|
919 |
+
# https://huggingface.co/garage-bAInd/Platypus2-7B/blob/main/config.json
|
920 |
+
dict(
|
921 |
+
name="Platypus2-7B",
|
922 |
+
hf_config=dict(org="garage-bAInd", name="Platypus2-7B"),
|
923 |
+
padded_vocab_size=32000,
|
924 |
+
n_layer=32,
|
925 |
+
rotary_percentage=1.0,
|
926 |
+
parallel_residual=False,
|
927 |
+
bias=False,
|
928 |
+
_norm_class="RMSNorm",
|
929 |
+
norm_eps=1e-05,
|
930 |
+
_mlp_class="LLaMAMLP",
|
931 |
+
intermediate_size=11008,
|
932 |
+
),
|
933 |
+
# https://huggingface.co/garage-bAInd/Platypus2-13B/blob/main/config.json
|
934 |
+
dict(
|
935 |
+
name="Platypus2-13B",
|
936 |
+
hf_config=dict(org="garage-bAInd", name="Platypus2-13B"),
|
937 |
+
padded_vocab_size=32000,
|
938 |
+
n_layer=40,
|
939 |
+
n_head=40,
|
940 |
+
n_embd=5120,
|
941 |
+
rotary_percentage=1.0,
|
942 |
+
parallel_residual=False,
|
943 |
+
bias=False,
|
944 |
+
_norm_class="RMSNorm",
|
945 |
+
norm_eps=1e-05,
|
946 |
+
_mlp_class="LLaMAMLP",
|
947 |
+
intermediate_size=13824,
|
948 |
+
),
|
949 |
+
# https://huggingface.co/garage-bAInd/Platypus2-70B/blob/main/config.json
|
950 |
+
dict(
|
951 |
+
name="Platypus2-70B",
|
952 |
+
hf_config=dict(org="garage-bAInd", name="Platypus2-70B"),
|
953 |
+
padded_vocab_size=32000,
|
954 |
+
n_layer=80,
|
955 |
+
n_head=64,
|
956 |
+
n_embd=8192,
|
957 |
+
rotary_percentage=1.0,
|
958 |
+
parallel_residual=False,
|
959 |
+
bias=False,
|
960 |
+
_norm_class="RMSNorm",
|
961 |
+
_mlp_class="LLaMAMLP",
|
962 |
+
intermediate_size=28672,
|
963 |
+
),
|
964 |
+
# https://huggingface.co/garage-bAInd/Camel-Platypus2-13B/blob/main/config.json
|
965 |
+
dict(
|
966 |
+
name="Camel-Platypus2-13B",
|
967 |
+
hf_config=dict(org="garage-bAInd", name="Camel-Platypus2-13B"),
|
968 |
+
padded_vocab_size=32000,
|
969 |
+
n_layer=40,
|
970 |
+
n_head=40,
|
971 |
+
n_embd=5120,
|
972 |
+
rotary_percentage=1.0,
|
973 |
+
parallel_residual=False,
|
974 |
+
bias=False,
|
975 |
+
_norm_class="RMSNorm",
|
976 |
+
_mlp_class="LLaMAMLP",
|
977 |
+
intermediate_size=13824,
|
978 |
+
),
|
979 |
+
# https://huggingface.co/garage-bAInd/Camel-Platypus2-70B/blob/main/config.json
|
980 |
+
dict(
|
981 |
+
name="Camel-Platypus2-70B",
|
982 |
+
hf_config=dict(org="garage-bAInd", name="Camel-Platypus2-70B"),
|
983 |
+
padded_vocab_size=32000,
|
984 |
+
n_layer=80,
|
985 |
+
n_head=64,
|
986 |
+
n_embd=8192,
|
987 |
+
n_query_groups=8,
|
988 |
+
rotary_percentage=1.0,
|
989 |
+
parallel_residual=False,
|
990 |
+
bias=False,
|
991 |
+
_norm_class="RMSNorm",
|
992 |
+
_mlp_class="LLaMAMLP",
|
993 |
+
intermediate_size=28672,
|
994 |
+
),
|
995 |
+
# https://huggingface.co/garage-bAInd/Stable-Platypus2-13B/blob/main/config.json
|
996 |
+
dict(
|
997 |
+
name="Stable-Platypus2-13B",
|
998 |
+
hf_config=dict(org="garage-bAInd", name="Stable-Platypus2-13B"),
|
999 |
+
padded_vocab_size=32000,
|
1000 |
+
n_layer=40,
|
1001 |
+
n_head=40,
|
1002 |
+
n_embd=5120,
|
1003 |
+
rotary_percentage=1.0,
|
1004 |
+
parallel_residual=False,
|
1005 |
+
bias=False,
|
1006 |
+
_norm_class="RMSNorm",
|
1007 |
+
_mlp_class="LLaMAMLP",
|
1008 |
+
intermediate_size=13824,
|
1009 |
+
),
|
1010 |
+
# https://huggingface.co/garage-bAInd/Platypus2-70B-instruct/blob/main/config.json
|
1011 |
+
dict(
|
1012 |
+
name="Platypus2-70B-instruct",
|
1013 |
+
hf_config=dict(org="garage-bAInd", name="Platypus2-70B-instruct"),
|
1014 |
+
padded_vocab_size=32000,
|
1015 |
+
n_layer=80,
|
1016 |
+
n_head=64,
|
1017 |
+
n_embd=8192,
|
1018 |
+
n_query_groups=8,
|
1019 |
+
rotary_percentage=1.0,
|
1020 |
+
parallel_residual=False,
|
1021 |
+
bias=False,
|
1022 |
+
_norm_class="RMSNorm",
|
1023 |
+
_mlp_class="LLaMAMLP",
|
1024 |
+
intermediate_size=28672,
|
1025 |
+
),
|
1026 |
+
]
|
1027 |
+
configs.extend(platypus)
|
1028 |
+
|
1029 |
+
|
1030 |
+
##########################
|
1031 |
+
# Stability AI StableCode
|
1032 |
+
##########################
|
1033 |
+
stablecode = [
|
1034 |
+
# https://huggingface.co/stabilityai/stablecode-completion-alpha-3b/blob/main/config.json
|
1035 |
+
dict(
|
1036 |
+
name="stablecode-completion-alpha-3b",
|
1037 |
+
hf_config=dict(org="stabilityai", name="stablecode-completion-alpha-3b"),
|
1038 |
+
block_size=16384,
|
1039 |
+
vocab_size=49152,
|
1040 |
+
n_layer=32,
|
1041 |
+
n_embd=2560,
|
1042 |
+
),
|
1043 |
+
# https://huggingface.co/stabilityai/stablecode-completion-alpha-3b-4k/blob/main/config.json
|
1044 |
+
dict(
|
1045 |
+
name="stablecode-completion-alpha-3b-4k",
|
1046 |
+
hf_config=dict(org="stabilityai", name="stablecode-completion-alpha-3b-4k"),
|
1047 |
+
vocab_size=49152,
|
1048 |
+
n_layer=32,
|
1049 |
+
n_embd=2560,
|
1050 |
+
),
|
1051 |
+
# https://huggingface.co/stabilityai/stablecode-instruct-alpha-3b/blob/main/config.json
|
1052 |
+
dict(
|
1053 |
+
name="stablecode-instruct-alpha-3b",
|
1054 |
+
hf_config=dict(org="stabilityai", name="stablecode-instruct-alpha-3b"),
|
1055 |
+
vocab_size=49152,
|
1056 |
+
n_layer=32,
|
1057 |
+
n_embd=2560,
|
1058 |
+
),
|
1059 |
+
]
|
1060 |
+
configs.extend(stablecode)
|
1061 |
+
|
1062 |
+
|
1063 |
+
##################################
|
1064 |
+
# togethercomputer LLaMA-2-7B-32K
|
1065 |
+
##################################
|
1066 |
+
together_llama2_32k = [
|
1067 |
+
# https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/blob/main/config.json
|
1068 |
+
dict(
|
1069 |
+
name="LLaMA-2-7B-32K",
|
1070 |
+
hf_config=dict(org="togethercomputer", name="LLaMA-2-7B-32K"),
|
1071 |
+
vocab_size=32000,
|
1072 |
+
padding_multiple=64,
|
1073 |
+
n_layer=32,
|
1074 |
+
rotary_percentage=1.0,
|
1075 |
+
parallel_residual=False,
|
1076 |
+
bias=False,
|
1077 |
+
_norm_class="RMSNorm",
|
1078 |
+
_mlp_class="LLaMAMLP",
|
1079 |
+
intermediate_size=11008,
|
1080 |
+
rope_condense_ratio=8,
|
1081 |
+
)
|
1082 |
+
]
|
1083 |
+
configs.extend(together_llama2_32k)
|
1084 |
+
|
1085 |
+
|
1086 |
+
################
|
1087 |
+
# Microsoft Phi
|
1088 |
+
################
|
1089 |
+
phi = [
|
1090 |
+
# https://huggingface.co/microsoft/phi-1_5/blob/main/config.json
|
1091 |
+
dict(
|
1092 |
+
name="phi-1_5",
|
1093 |
+
hf_config=dict(org="microsoft", name="phi-1_5"),
|
1094 |
+
vocab_size=50257,
|
1095 |
+
padded_vocab_size=51200,
|
1096 |
+
block_size=2048,
|
1097 |
+
n_embd=2048,
|
1098 |
+
n_layer=24,
|
1099 |
+
rotary_percentage=0.5, # 32 / (n_embd / n_head) = 32 / 64
|
1100 |
+
shared_attention_norm=True,
|
1101 |
+
lm_head_bias=True,
|
1102 |
+
gelu_approximate="tanh",
|
1103 |
+
)
|
1104 |
+
]
|
1105 |
+
configs.extend(phi)
|
1106 |
+
|
1107 |
+
|
1108 |
+
#############
|
1109 |
+
# Mistral AI
|
1110 |
+
#############
|
1111 |
+
mistral = [
|
1112 |
+
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
|
1113 |
+
dict(
|
1114 |
+
name="Mistral-7B-{}v0.1",
|
1115 |
+
hf_config=dict(org="mistralai", name="Mistral-7B-{}v0.1"),
|
1116 |
+
padded_vocab_size=32000,
|
1117 |
+
block_size=4096, # should be 32768 but sliding window attention is not implemented
|
1118 |
+
n_layer=32,
|
1119 |
+
n_query_groups=8,
|
1120 |
+
rotary_percentage=1.0,
|
1121 |
+
parallel_residual=False,
|
1122 |
+
bias=False,
|
1123 |
+
_norm_class="RMSNorm",
|
1124 |
+
norm_eps=1e-05,
|
1125 |
+
_mlp_class="LLaMAMLP",
|
1126 |
+
intermediate_size=14336,
|
1127 |
+
)
|
1128 |
+
]
|
1129 |
+
for c in mistral:
|
1130 |
+
for kind in ("", "Instruct-"):
|
1131 |
+
copy = c.copy()
|
1132 |
+
copy["name"] = c["name"].format(kind)
|
1133 |
+
copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
|
1134 |
+
configs.append(copy)
|
1135 |
+
|
1136 |
+
|
1137 |
+
############
|
1138 |
+
# TinyLlama
|
1139 |
+
############
|
1140 |
+
tiny_llama = [
|
1141 |
+
dict(
|
1142 |
+
name="tiny-llama-1.1b",
|
1143 |
+
hf_config=dict(org="PY007", name="TinyLlama-1.1B-intermediate-step-480k-1T"),
|
1144 |
+
block_size=2048,
|
1145 |
+
vocab_size=32000,
|
1146 |
+
padding_multiple=64,
|
1147 |
+
n_layer=22,
|
1148 |
+
n_head=32,
|
1149 |
+
n_embd=2048,
|
1150 |
+
rotary_percentage=1.0,
|
1151 |
+
parallel_residual=False,
|
1152 |
+
bias=False,
|
1153 |
+
_norm_class="RMSNorm", # original TinyLlama uses FusedRMSNorm
|
1154 |
+
norm_eps=1e-5,
|
1155 |
+
_mlp_class="LLaMAMLP",
|
1156 |
+
intermediate_size=5632,
|
1157 |
+
n_query_groups=4,
|
1158 |
+
),
|
1159 |
+
dict(
|
1160 |
+
name="tiny-llama-new",
|
1161 |
+
hf_config=dict(org="PY007", name="TinyLlama-1.1B-intermediate-step-480k-1T"),
|
1162 |
+
block_size=768,
|
1163 |
+
vocab_size=32000,
|
1164 |
+
padding_multiple=64,
|
1165 |
+
n_layer=18,
|
1166 |
+
n_head=32,
|
1167 |
+
n_embd=1024,
|
1168 |
+
rotary_percentage=1.0,
|
1169 |
+
parallel_residual=False,
|
1170 |
+
bias=False,
|
1171 |
+
_norm_class="RMSNorm", # original TinyLlama uses FusedRMSNorm
|
1172 |
+
norm_eps=1e-5,
|
1173 |
+
_mlp_class="LLaMAMLP",
|
1174 |
+
intermediate_size=5632,
|
1175 |
+
n_query_groups=4,
|
1176 |
+
),
|
1177 |
+
]
|
1178 |
+
configs.extend(tiny_llama)
|
1179 |
+
|
1180 |
+
|
1181 |
+
name_to_config = {config["name"]: config for config in configs}
|
tsai_gpt/model.py
ADDED
@@ -0,0 +1,342 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""Full definition of a GPT NeoX Language Model, all of it in this single file.
|
2 |
+
|
3 |
+
Based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT and
|
4 |
+
https://github.com/EleutherAI/gpt-neox/tree/main/megatron/model.
|
5 |
+
"""
|
6 |
+
import math
|
7 |
+
from typing import Any, Optional, Tuple
|
8 |
+
|
9 |
+
import torch
|
10 |
+
import torch.nn as nn
|
11 |
+
from typing_extensions import Self
|
12 |
+
|
13 |
+
from tsai_gpt.config import Config
|
14 |
+
|
15 |
+
|
16 |
+
|
17 |
+
class GPT(nn.Module):
|
18 |
+
def __init__(self, config: Config) -> None:
|
19 |
+
super().__init__()
|
20 |
+
assert config.padded_vocab_size is not None
|
21 |
+
self.config = config
|
22 |
+
|
23 |
+
self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=config.lm_head_bias)
|
24 |
+
self.transformer = nn.ModuleDict(
|
25 |
+
dict(
|
26 |
+
wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
|
27 |
+
h=nn.ModuleList(Block(config) for _ in range(config.n_layer)),
|
28 |
+
ln_f=config.norm_class(config.n_embd, eps=config.norm_eps),
|
29 |
+
)
|
30 |
+
)
|
31 |
+
self.max_seq_length = self.config.block_size
|
32 |
+
self.mask_cache: Optional[torch.Tensor] = None
|
33 |
+
|
34 |
+
@property
|
35 |
+
def max_seq_length(self) -> int:
|
36 |
+
return self._max_seq_length
|
37 |
+
|
38 |
+
@max_seq_length.setter
|
39 |
+
def max_seq_length(self, value: int) -> None:
|
40 |
+
"""
|
41 |
+
When doing inference, the sequences used might be shorter than the model's context length.
|
42 |
+
This allows setting a smaller number to avoid allocating unused memory
|
43 |
+
"""
|
44 |
+
if value > self.config.block_size:
|
45 |
+
raise ValueError(f"Cannot attend to {value}, block size is only {self.config.block_size}")
|
46 |
+
self._max_seq_length = value
|
47 |
+
if not hasattr(self, "cos"):
|
48 |
+
# first call
|
49 |
+
cos, sin = self.rope_cache()
|
50 |
+
self.register_buffer("cos", cos, persistent=False)
|
51 |
+
self.register_buffer("sin", sin, persistent=False)
|
52 |
+
elif value != self.cos.size(0):
|
53 |
+
# override
|
54 |
+
self.cos, self.sin = self.rope_cache(device=self.cos.device)
|
55 |
+
# the mask and kv cache size will get updated on `set_kv_cache`. we cannot update it here because we don't know
|
56 |
+
# if the kv cache is expected
|
57 |
+
|
58 |
+
def reset_parameters(self) -> None:
|
59 |
+
# Trigger resetting the rope-cache
|
60 |
+
self.max_seq_length = self.config.block_size
|
61 |
+
|
62 |
+
def _init_weights(self, module: nn.Module) -> None:
|
63 |
+
"""Meant to be used with `gpt.apply(gpt._init_weights)`."""
|
64 |
+
if isinstance(module, nn.Linear):
|
65 |
+
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
66 |
+
if module.bias is not None:
|
67 |
+
torch.nn.init.zeros_(module.bias)
|
68 |
+
elif isinstance(module, nn.Embedding):
|
69 |
+
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
70 |
+
|
71 |
+
def forward(self, idx: torch.Tensor, input_pos: Optional[torch.Tensor] = None) -> torch.Tensor:
|
72 |
+
T = idx.size(1)
|
73 |
+
if self.max_seq_length < T:
|
74 |
+
raise ValueError(f"Cannot forward sequence of length {T}, max seq length is only {self.max_seq_length}.")
|
75 |
+
|
76 |
+
if input_pos is not None: # use the kv cache
|
77 |
+
cos = self.cos.index_select(0, input_pos)
|
78 |
+
sin = self.sin.index_select(0, input_pos)
|
79 |
+
if self.mask_cache is None:
|
80 |
+
raise TypeError("You need to call `gpt.set_kv_cache()`")
|
81 |
+
mask = self.mask_cache.index_select(2, input_pos)
|
82 |
+
else:
|
83 |
+
cos = self.cos[:T]
|
84 |
+
sin = self.sin[:T]
|
85 |
+
mask = None
|
86 |
+
|
87 |
+
x = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
|
88 |
+
for block in self.transformer.h:
|
89 |
+
x = block(x, cos, sin, mask, input_pos)
|
90 |
+
x = self.transformer.ln_f(x)
|
91 |
+
return self.lm_head(x) # (b, t, vocab_size)
|
92 |
+
|
93 |
+
@classmethod
|
94 |
+
def from_name(cls, name: str, **kwargs: Any) -> Self:
|
95 |
+
return cls(Config.from_name(name, **kwargs))
|
96 |
+
|
97 |
+
def rope_cache(self, device: Optional[torch.device] = None) -> Tuple[torch.Tensor, torch.Tensor]:
|
98 |
+
return build_rope_cache(
|
99 |
+
seq_len=self.max_seq_length,
|
100 |
+
n_elem=self.config.rope_n_elem,
|
101 |
+
device=device,
|
102 |
+
condense_ratio=self.config.rope_condense_ratio,
|
103 |
+
base=self.config.rope_base,
|
104 |
+
)
|
105 |
+
|
106 |
+
def set_kv_cache(
|
107 |
+
self,
|
108 |
+
batch_size: int,
|
109 |
+
rope_cache_length: Optional[int] = None,
|
110 |
+
device: Optional[torch.device] = None,
|
111 |
+
dtype: Optional[torch.dtype] = None,
|
112 |
+
) -> None:
|
113 |
+
if rope_cache_length is None:
|
114 |
+
rope_cache_length = self.cos.size(-1)
|
115 |
+
max_seq_length = self.max_seq_length
|
116 |
+
|
117 |
+
# initialize the kv cache for all blocks
|
118 |
+
for block in self.transformer.h:
|
119 |
+
block.attn.kv_cache = block.attn.build_kv_cache(
|
120 |
+
batch_size, max_seq_length, rope_cache_length, device, dtype
|
121 |
+
)
|
122 |
+
|
123 |
+
if self.mask_cache is None or self.mask_cache.size(3) != max_seq_length:
|
124 |
+
# passing `attn_mask` to SDPA downgrades it to use the inefficient implementation. since we only need the mask
|
125 |
+
# for the kv-cache support (only during inference), we only create it in that situation
|
126 |
+
# this will be resolved by https://github.com/pytorch/pytorch/issues/96099
|
127 |
+
ones = torch.ones((max_seq_length, max_seq_length), device=device, dtype=torch.bool)
|
128 |
+
self.mask_cache = torch.tril(ones).unsqueeze(0).unsqueeze(0)
|
129 |
+
|
130 |
+
def clear_kv_cache(self) -> None:
|
131 |
+
self.mask_cache = None
|
132 |
+
for block in self.transformer.h:
|
133 |
+
block.attn.kv_cache = None
|
134 |
+
|
135 |
+
|
136 |
+
class Block(nn.Module):
|
137 |
+
def __init__(self, config: Config) -> None:
|
138 |
+
super().__init__()
|
139 |
+
self.norm_1 = config.norm_class(config.n_embd, eps=config.norm_eps)
|
140 |
+
self.attn = CausalSelfAttention(config)
|
141 |
+
self.norm_2 = None if config.shared_attention_norm else config.norm_class(config.n_embd, eps=config.norm_eps)
|
142 |
+
self.mlp = config.mlp_class(config)
|
143 |
+
|
144 |
+
self.config = config
|
145 |
+
|
146 |
+
def forward(
|
147 |
+
self,
|
148 |
+
x: torch.Tensor,
|
149 |
+
cos: torch.Tensor,
|
150 |
+
sin: torch.Tensor,
|
151 |
+
mask: Optional[torch.Tensor] = None,
|
152 |
+
input_pos: Optional[torch.Tensor] = None,
|
153 |
+
) -> torch.Tensor:
|
154 |
+
n_1 = self.norm_1(x)
|
155 |
+
h = self.attn(n_1, cos, sin, mask, input_pos)
|
156 |
+
if self.config.parallel_residual:
|
157 |
+
n_2 = n_1 if self.config.shared_attention_norm else self.norm_2(x)
|
158 |
+
x = self.mlp(n_2) + h + x
|
159 |
+
else:
|
160 |
+
if self.config.shared_attention_norm:
|
161 |
+
raise NotImplementedError(
|
162 |
+
"No checkpoint amongst the ones we support uses this configuration"
|
163 |
+
" (non-parallel residual and shared attention norm)."
|
164 |
+
)
|
165 |
+
x = h + x
|
166 |
+
x = self.mlp(self.norm_2(x)) + x
|
167 |
+
return x
|
168 |
+
|
169 |
+
|
170 |
+
class CausalSelfAttention(nn.Module):
|
171 |
+
def __init__(self, config: Config) -> None:
|
172 |
+
super().__init__()
|
173 |
+
shape = (config.n_head + 2 * config.n_query_groups) * config.head_size
|
174 |
+
# key, query, value projections for all heads, but in a batch
|
175 |
+
self.attn = nn.Linear(config.n_embd, shape, bias=config.bias)
|
176 |
+
# output projection
|
177 |
+
self.proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
|
178 |
+
# disabled by default
|
179 |
+
self.kv_cache: Optional[KVCache] = None
|
180 |
+
|
181 |
+
self.config = config
|
182 |
+
|
183 |
+
def forward(
|
184 |
+
self,
|
185 |
+
x: torch.Tensor,
|
186 |
+
cos: torch.Tensor,
|
187 |
+
sin: torch.Tensor,
|
188 |
+
mask: Optional[torch.Tensor] = None,
|
189 |
+
input_pos: Optional[torch.Tensor] = None,
|
190 |
+
) -> torch.Tensor:
|
191 |
+
B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
|
192 |
+
|
193 |
+
qkv = self.attn(x)
|
194 |
+
|
195 |
+
# assemble into a number of query groups to support MHA, MQA and GQA together (see `config.n_query_groups`)
|
196 |
+
q_per_kv = self.config.n_head // self.config.n_query_groups
|
197 |
+
total_qkv = q_per_kv + 2 # each group has 1+ queries, 1 key, and 1 value
|
198 |
+
qkv = qkv.view(B, T, self.config.n_query_groups, total_qkv, self.config.head_size)
|
199 |
+
qkv = qkv.permute(0, 2, 3, 1, 4) # (B, n_query_groups, total_qkv, T, hs)
|
200 |
+
|
201 |
+
# split batched computation into three
|
202 |
+
q, k, v = qkv.split((q_per_kv, 1, 1), dim=2)
|
203 |
+
|
204 |
+
# maybe repeat k and v if for the non multi-head attention cases
|
205 |
+
# training: flash attention requires it
|
206 |
+
# inference: multi-query would require a full kv cache so avoid it to limit its memory usage
|
207 |
+
if self.config.n_query_groups != self.config.n_head and (input_pos is None or self.config.n_query_groups != 1):
|
208 |
+
k = k.expand(B, self.config.n_query_groups, q_per_kv, T, self.config.head_size)
|
209 |
+
v = v.expand(B, self.config.n_query_groups, q_per_kv, T, self.config.head_size)
|
210 |
+
|
211 |
+
q = q.reshape(B, -1, T, self.config.head_size) # (B, nh_q, T, hs)
|
212 |
+
k = k.reshape(B, -1, T, self.config.head_size) # (B, nh_k, T, hs)
|
213 |
+
v = v.reshape(B, -1, T, self.config.head_size) # (B, nh_v, T, hs)
|
214 |
+
|
215 |
+
q_roped = apply_rope(q[..., : self.config.rope_n_elem], cos, sin)
|
216 |
+
k_roped = apply_rope(k[..., : self.config.rope_n_elem], cos, sin)
|
217 |
+
q = torch.cat((q_roped, q[..., self.config.rope_n_elem :]), dim=-1)
|
218 |
+
k = torch.cat((k_roped, k[..., self.config.rope_n_elem :]), dim=-1)
|
219 |
+
|
220 |
+
if input_pos is not None:
|
221 |
+
if not isinstance(self.kv_cache, KVCache):
|
222 |
+
raise TypeError("You need to call `gpt.set_kv_cache()`")
|
223 |
+
k, v = self.kv_cache(input_pos, k, v)
|
224 |
+
|
225 |
+
y = self.scaled_dot_product_attention(q, k, v, mask)
|
226 |
+
|
227 |
+
y = y.reshape(B, T, C) # re-assemble all head outputs side by side
|
228 |
+
|
229 |
+
# output projection
|
230 |
+
return self.proj(y)
|
231 |
+
|
232 |
+
def scaled_dot_product_attention(
|
233 |
+
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor] = None
|
234 |
+
) -> torch.Tensor:
|
235 |
+
scale = 1.0 / math.sqrt(self.config.head_size)
|
236 |
+
y = torch.nn.functional.scaled_dot_product_attention(
|
237 |
+
q, k, v, attn_mask=mask, dropout_p=0.0, scale=scale, is_causal=mask is None
|
238 |
+
)
|
239 |
+
return y.transpose(1, 2)
|
240 |
+
|
241 |
+
def build_kv_cache(
|
242 |
+
self,
|
243 |
+
batch_size: int,
|
244 |
+
max_seq_length: int,
|
245 |
+
rope_cache_length: Optional[int] = None,
|
246 |
+
device: Optional[torch.device] = None,
|
247 |
+
dtype: Optional[torch.dtype] = None,
|
248 |
+
) -> "KVCache":
|
249 |
+
heads = 1 if self.config.n_query_groups == 1 else self.config.n_head
|
250 |
+
v_shape = (batch_size, heads, max_seq_length, self.config.head_size)
|
251 |
+
if rope_cache_length is None:
|
252 |
+
if self.config.rotary_percentage != 1.0:
|
253 |
+
raise TypeError("Please pass the `rope_cache_length=gpt.cos.size(-1)` value")
|
254 |
+
k_shape = v_shape
|
255 |
+
else:
|
256 |
+
k_shape = (
|
257 |
+
batch_size,
|
258 |
+
heads,
|
259 |
+
max_seq_length,
|
260 |
+
rope_cache_length + self.config.head_size - self.config.rope_n_elem,
|
261 |
+
)
|
262 |
+
return KVCache(k_shape, v_shape, device=device, dtype=dtype)
|
263 |
+
|
264 |
+
|
265 |
+
class GptNeoxMLP(nn.Module):
|
266 |
+
def __init__(self, config: Config) -> None:
|
267 |
+
super().__init__()
|
268 |
+
self.fc = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
|
269 |
+
self.proj = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias)
|
270 |
+
|
271 |
+
self.config = config
|
272 |
+
|
273 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
274 |
+
x = self.fc(x)
|
275 |
+
x = torch.nn.functional.gelu(x, approximate=self.config.gelu_approximate)
|
276 |
+
return self.proj(x)
|
277 |
+
|
278 |
+
|
279 |
+
class LLaMAMLP(nn.Module):
|
280 |
+
def __init__(self, config: Config) -> None:
|
281 |
+
super().__init__()
|
282 |
+
self.fc_1 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
|
283 |
+
self.fc_2 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
|
284 |
+
self.proj = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias)
|
285 |
+
|
286 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
287 |
+
x_fc_1 = self.fc_1(x)
|
288 |
+
x_fc_2 = self.fc_2(x)
|
289 |
+
x = torch.nn.functional.silu(x_fc_1) * x_fc_2
|
290 |
+
return self.proj(x)
|
291 |
+
|
292 |
+
|
293 |
+
def build_rope_cache(
|
294 |
+
seq_len: int, n_elem: int, device: Optional[torch.device] = None, base: int = 10000, condense_ratio: int = 1
|
295 |
+
) -> Tuple[torch.Tensor, torch.Tensor]:
|
296 |
+
"""Enhanced Transformer with Rotary Position Embedding.
|
297 |
+
|
298 |
+
Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
|
299 |
+
transformers/rope/__init__.py. MIT License:
|
300 |
+
https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
|
301 |
+
"""
|
302 |
+
# $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
|
303 |
+
theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, device=device).float() / n_elem))
|
304 |
+
|
305 |
+
# Create position indexes `[0, 1, ..., seq_len - 1]`
|
306 |
+
seq_idx = torch.arange(seq_len, device=device) / condense_ratio
|
307 |
+
|
308 |
+
# Calculate the product of position index and $\theta_i$
|
309 |
+
idx_theta = torch.outer(seq_idx, theta).repeat(1, 2)
|
310 |
+
|
311 |
+
return torch.cos(idx_theta), torch.sin(idx_theta)
|
312 |
+
|
313 |
+
|
314 |
+
def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
|
315 |
+
head_size = x.size(-1)
|
316 |
+
x1 = x[..., : head_size // 2] # (B, nh, T, hs/2)
|
317 |
+
x2 = x[..., head_size // 2 :] # (B, nh, T, hs/2)
|
318 |
+
rotated = torch.cat((-x2, x1), dim=-1) # (B, nh, T, hs)
|
319 |
+
roped = (x * cos) + (rotated * sin)
|
320 |
+
return roped.type_as(x)
|
321 |
+
|
322 |
+
|
323 |
+
class KVCache(nn.Module):
|
324 |
+
def __init__(
|
325 |
+
self,
|
326 |
+
k_shape: Tuple[int, int, int, int],
|
327 |
+
v_shape: Tuple[int, int, int, int],
|
328 |
+
device: Optional[torch.device] = None,
|
329 |
+
dtype: Optional[torch.dtype] = None,
|
330 |
+
) -> None:
|
331 |
+
super().__init__()
|
332 |
+
self.register_buffer("k", torch.zeros(k_shape, device=device, dtype=dtype), persistent=False)
|
333 |
+
self.register_buffer("v", torch.zeros(v_shape, device=device, dtype=dtype), persistent=False)
|
334 |
+
|
335 |
+
def forward(self, input_pos: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
336 |
+
# move the buffer to the activation dtype for when AMP is used
|
337 |
+
self.k = self.k.to(k.dtype)
|
338 |
+
self.v = self.v.to(v.dtype)
|
339 |
+
# update the cache
|
340 |
+
k = self.k.index_copy_(2, input_pos, k)
|
341 |
+
v = self.v.index_copy_(2, input_pos, v)
|
342 |
+
return k, v
|
tsai_gpt/packed_dataset.py
ADDED
@@ -0,0 +1,235 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Very loosely inspired by indexed_dataset in Fairseq, Megatron
|
2 |
+
# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/indexed_dataset.py
|
3 |
+
|
4 |
+
|
5 |
+
import os
|
6 |
+
import random
|
7 |
+
import struct
|
8 |
+
|
9 |
+
import numpy as np
|
10 |
+
import torch
|
11 |
+
from torch.utils.data import IterableDataset, get_worker_info
|
12 |
+
|
13 |
+
dtypes = {1: np.uint8, 2: np.int8, 3: np.int16, 4: np.int32, 5: np.int64, 6: np.float32, 7: np.float64, 8: np.uint16}
|
14 |
+
|
15 |
+
|
16 |
+
def code(dtype):
|
17 |
+
for k in dtypes:
|
18 |
+
if dtypes[k] == dtype:
|
19 |
+
return k
|
20 |
+
raise ValueError(dtype)
|
21 |
+
|
22 |
+
|
23 |
+
HDR_MAGIC = b"LITPKDS"
|
24 |
+
HDR_SIZE = 24 # bytes
|
25 |
+
|
26 |
+
|
27 |
+
class PackedDataset(IterableDataset):
|
28 |
+
def __init__(
|
29 |
+
self, filenames, n_chunks, block_size, seed=12345, shuffle=True, wrap=False, num_processes=1, process_rank=0
|
30 |
+
):
|
31 |
+
self._filenames = filenames
|
32 |
+
self._n_chunks = n_chunks
|
33 |
+
self._block_size = block_size
|
34 |
+
self._seed = seed
|
35 |
+
self._shuffle = shuffle
|
36 |
+
self._wrap = wrap
|
37 |
+
self._num_processes = num_processes
|
38 |
+
self._process_rank = process_rank
|
39 |
+
|
40 |
+
def __iter__(self):
|
41 |
+
worker_info = get_worker_info()
|
42 |
+
num_workers = worker_info.num_workers if worker_info is not None else 1
|
43 |
+
worker_id = worker_info.id if worker_info is not None else 0
|
44 |
+
num_shards = num_workers * self._num_processes
|
45 |
+
shard_id = self._process_rank * num_workers + worker_id
|
46 |
+
|
47 |
+
max_num_files = len(self._filenames) // num_shards * num_shards
|
48 |
+
filenames = self._filenames[shard_id:max_num_files:num_shards]
|
49 |
+
|
50 |
+
return PackedDatasetIterator(
|
51 |
+
filenames=filenames,
|
52 |
+
n_chunks=self._n_chunks,
|
53 |
+
block_size=self._block_size,
|
54 |
+
seed=self._seed,
|
55 |
+
shuffle=self._shuffle,
|
56 |
+
wrap=self._wrap,
|
57 |
+
)
|
58 |
+
|
59 |
+
|
60 |
+
class PackedDatasetBuilder(object):
|
61 |
+
def __init__(self, outdir, prefix, chunk_size, sep_token, dtype="auto", vocab_size=None):
|
62 |
+
if dtype == "auto":
|
63 |
+
if vocab_size is None:
|
64 |
+
raise ValueError("vocab_size cannot be None when dtype='auto'")
|
65 |
+
if vocab_size is not None and vocab_size < 65500:
|
66 |
+
self._dtype = np.uint16
|
67 |
+
else:
|
68 |
+
self._dtype = np.int32
|
69 |
+
else:
|
70 |
+
self._dtype = dtype
|
71 |
+
self._counter = 0
|
72 |
+
self._chunk_size = chunk_size
|
73 |
+
self._outdir = outdir
|
74 |
+
self._prefix = prefix
|
75 |
+
self._sep_token = sep_token
|
76 |
+
self._arr = np.zeros(self._chunk_size, dtype=self._dtype)
|
77 |
+
self._arr.fill(self._sep_token)
|
78 |
+
self._idx = 0
|
79 |
+
self._version = 1
|
80 |
+
self._filenames = []
|
81 |
+
|
82 |
+
def _write_chunk(self):
|
83 |
+
filename = f"{self._prefix}_{self._counter:010d}.bin"
|
84 |
+
filename = os.path.join(self._outdir, filename)
|
85 |
+
|
86 |
+
with open(filename, "wb") as f:
|
87 |
+
f.write(HDR_MAGIC)
|
88 |
+
f.write(struct.pack("<Q", self._version))
|
89 |
+
f.write(struct.pack("<B", code(self._dtype)))
|
90 |
+
f.write(struct.pack("<Q", self._chunk_size))
|
91 |
+
f.write(self._arr.tobytes(order="C"))
|
92 |
+
|
93 |
+
self._filenames.append(filename)
|
94 |
+
self._counter += 1
|
95 |
+
self._arr.fill(self._sep_token)
|
96 |
+
self._idx = 0
|
97 |
+
|
98 |
+
@property
|
99 |
+
def dtype(self):
|
100 |
+
return self._dtype
|
101 |
+
|
102 |
+
@property
|
103 |
+
def filenames(self):
|
104 |
+
return self._filenames.copy()
|
105 |
+
|
106 |
+
def add_array(self, arr):
|
107 |
+
while self._idx + arr.shape[0] > self._chunk_size:
|
108 |
+
part_len = self._chunk_size - self._idx
|
109 |
+
self._arr[self._idx : self._idx + part_len] = arr[:part_len]
|
110 |
+
self._write_chunk()
|
111 |
+
arr = arr[part_len:]
|
112 |
+
|
113 |
+
arr_len = arr.shape[0]
|
114 |
+
self._arr[self._idx : self._idx + arr_len] = arr
|
115 |
+
self._idx += arr_len
|
116 |
+
|
117 |
+
def write_reminder(self):
|
118 |
+
self._write_chunk()
|
119 |
+
|
120 |
+
|
121 |
+
class PackedDatasetIterator:
|
122 |
+
def __init__(self, filenames, n_chunks, block_size, seed, shuffle, wrap):
|
123 |
+
self._seed = seed
|
124 |
+
self._shuffle = shuffle
|
125 |
+
self._rng = np.random.default_rng(seed) if shuffle else None
|
126 |
+
self._block_idxs = None
|
127 |
+
|
128 |
+
self._wrap = wrap
|
129 |
+
|
130 |
+
# TODO: instead of filenames, we could have a single text stream
|
131 |
+
# (or text file) with the sequence of all files to be
|
132 |
+
# fetched/loaded.
|
133 |
+
self._filenames = filenames
|
134 |
+
self._file_idx = 0
|
135 |
+
|
136 |
+
self._n_chunks = n_chunks
|
137 |
+
|
138 |
+
self._dtype = None
|
139 |
+
self._block_size = block_size
|
140 |
+
self._n_blocks = None
|
141 |
+
|
142 |
+
self._mmaps = []
|
143 |
+
self._buffers = []
|
144 |
+
|
145 |
+
self._block_idxs = []
|
146 |
+
self._curr_idx = 0
|
147 |
+
|
148 |
+
self._load_n_chunks()
|
149 |
+
|
150 |
+
def _read_header(self, path):
|
151 |
+
with open(path, "rb") as f:
|
152 |
+
magic = f.read(len(HDR_MAGIC))
|
153 |
+
assert magic == HDR_MAGIC, "File doesn't match expected format."
|
154 |
+
version = struct.unpack("<Q", f.read(8))
|
155 |
+
assert version == (1,)
|
156 |
+
(dtype_code,) = struct.unpack("<B", f.read(1))
|
157 |
+
dtype = dtypes[dtype_code]
|
158 |
+
(chunk_size,) = struct.unpack("<Q", f.read(8))
|
159 |
+
return dtype, chunk_size
|
160 |
+
|
161 |
+
def _close_mmaps(self):
|
162 |
+
for mmap in self._mmaps:
|
163 |
+
mmap._mmap.close()
|
164 |
+
|
165 |
+
def _load_n_chunks(self):
|
166 |
+
self._close_mmaps()
|
167 |
+
self._mmaps = []
|
168 |
+
self._buffers = []
|
169 |
+
|
170 |
+
if self._n_chunks > len(self._filenames[self._file_idx :]):
|
171 |
+
if not self._wrap:
|
172 |
+
raise StopIteration
|
173 |
+
self._file_idx = 0
|
174 |
+
|
175 |
+
for i in range(self._n_chunks):
|
176 |
+
filename = self._filenames[self._file_idx + i]
|
177 |
+
if self._dtype is None:
|
178 |
+
self._dtype, self._chunk_size = self._read_header(filename)
|
179 |
+
self._n_blocks = self._chunk_size // self._block_size
|
180 |
+
# TODO: check header matches with previous files
|
181 |
+
mmap = np.memmap(filename, mode="r", order="C", offset=HDR_SIZE)
|
182 |
+
self._mmaps.append(mmap)
|
183 |
+
self._buffers.append(memoryview(mmap))
|
184 |
+
|
185 |
+
self._file_idx += self._n_chunks
|
186 |
+
n_all_blocks = self._n_chunks * self._n_blocks
|
187 |
+
|
188 |
+
self._block_idxs = self._rng.permutation(n_all_blocks) if self._shuffle else range(n_all_blocks)
|
189 |
+
|
190 |
+
self._curr_idx = 0
|
191 |
+
|
192 |
+
def __del__(self):
|
193 |
+
self._close_mmaps()
|
194 |
+
del self._mmaps
|
195 |
+
del self._buffers
|
196 |
+
|
197 |
+
def __iter__(self):
|
198 |
+
return self
|
199 |
+
|
200 |
+
def __next__(self):
|
201 |
+
if self._curr_idx >= len(self._block_idxs):
|
202 |
+
self._load_n_chunks()
|
203 |
+
# TODO: trigger fetching next next n_chunks if remote
|
204 |
+
block_idx = self._block_idxs[self._curr_idx]
|
205 |
+
chunk_id = block_idx // self._n_blocks
|
206 |
+
buffer = self._buffers[chunk_id]
|
207 |
+
elem_id = (block_idx % self._n_blocks) * self._block_size
|
208 |
+
offset = np.dtype(self._dtype).itemsize * elem_id
|
209 |
+
arr = np.frombuffer(buffer, dtype=self._dtype, count=self._block_size, offset=offset)
|
210 |
+
self._curr_idx += 1
|
211 |
+
return torch.from_numpy(arr.astype(np.int64))
|
212 |
+
|
213 |
+
|
214 |
+
class CombinedDataset(IterableDataset):
|
215 |
+
def __init__(self, datasets, seed, weights=None):
|
216 |
+
self._seed = seed
|
217 |
+
self._datasets = datasets
|
218 |
+
self._weights = weights
|
219 |
+
n_datasets = len(datasets)
|
220 |
+
if weights is None:
|
221 |
+
self._weights = [1 / n_datasets] * n_datasets
|
222 |
+
|
223 |
+
def __iter__(self):
|
224 |
+
return CombinedDatasetIterator(self._datasets, self._seed, self._weights)
|
225 |
+
|
226 |
+
|
227 |
+
class CombinedDatasetIterator:
|
228 |
+
def __init__(self, datasets, seed, weights):
|
229 |
+
self._datasets = [iter(el) for el in datasets]
|
230 |
+
self._weights = weights
|
231 |
+
self._rng = random.Random(seed)
|
232 |
+
|
233 |
+
def __next__(self):
|
234 |
+
(dataset,) = self._rng.choices(self._datasets, weights=self._weights, k=1)
|
235 |
+
return next(dataset)
|
tsai_gpt/rmsnorm.py
ADDED
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
|
3 |
+
|
4 |
+
class RMSNorm(torch.nn.Module):
|
5 |
+
"""Root Mean Square Layer Normalization.
|
6 |
+
|
7 |
+
Derived from https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py. BSD 3-Clause License:
|
8 |
+
https://github.com/bzhangGo/rmsnorm/blob/master/LICENSE.
|
9 |
+
"""
|
10 |
+
|
11 |
+
def __init__(self, size: int, dim: int = -1, eps: float = 1e-5) -> None:
|
12 |
+
super().__init__()
|
13 |
+
self.weight = torch.nn.Parameter(torch.ones(size))
|
14 |
+
self.eps = eps
|
15 |
+
self.dim = dim
|
16 |
+
|
17 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
18 |
+
dtype = x.dtype
|
19 |
+
x = x.float()
|
20 |
+
# NOTE: the original RMSNorm paper implementation is not equivalent
|
21 |
+
norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
|
22 |
+
x_normed = x * torch.rsqrt(norm_x + self.eps)
|
23 |
+
return (self.weight * x_normed).to(dtype=dtype)
|
24 |
+
|
25 |
+
def reset_parameters(self) -> None:
|
26 |
+
torch.nn.init.ones_(self.weight)
|
tsai_gpt/speed_monitor.py
ADDED
@@ -0,0 +1,425 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import time
|
2 |
+
from collections import deque
|
3 |
+
from contextlib import nullcontext
|
4 |
+
from typing import Any, Callable, Deque, Dict, Optional
|
5 |
+
|
6 |
+
import torch
|
7 |
+
from lightning import Callback, Fabric, LightningModule, Trainer
|
8 |
+
from lightning.fabric.accelerators.xla import _XLA_GREATER_EQUAL_2_1
|
9 |
+
from lightning.fabric.plugins import (
|
10 |
+
BitsandbytesPrecision,
|
11 |
+
DoublePrecision,
|
12 |
+
FSDPPrecision,
|
13 |
+
HalfPrecision,
|
14 |
+
MixedPrecision,
|
15 |
+
Precision,
|
16 |
+
TransformerEnginePrecision,
|
17 |
+
XLAPrecision,
|
18 |
+
)
|
19 |
+
from lightning.fabric.utilities.rank_zero import rank_zero_only as fabric_rank_zero_only
|
20 |
+
from lightning.pytorch.plugins import (
|
21 |
+
DoublePrecisionPlugin,
|
22 |
+
FSDPPrecisionPlugin,
|
23 |
+
HalfPrecisionPlugin,
|
24 |
+
MixedPrecisionPlugin,
|
25 |
+
XLAPrecisionPlugin,
|
26 |
+
)
|
27 |
+
from lightning.pytorch.utilities.rank_zero import rank_zero_only as trainer_rank_zero_only
|
28 |
+
from torch.utils.flop_counter import FlopCounterMode
|
29 |
+
|
30 |
+
from tsai_gpt import GPT
|
31 |
+
from tsai_gpt.utils import num_parameters
|
32 |
+
|
33 |
+
GPU_AVAILABLE_FLOPS = {
|
34 |
+
# source: https://resources.nvidia.com/en-us-tensor-core/nvidia-tensor-core-gpu-datasheet
|
35 |
+
# nvidia publishes spec sheet with a 2x sparsity factor
|
36 |
+
"h100-sxm": {
|
37 |
+
torch.float64: 67e12,
|
38 |
+
torch.float32: 67e12,
|
39 |
+
torch.bfloat16: 1.979e15 / 2,
|
40 |
+
torch.float16: 1.979e15 / 2,
|
41 |
+
torch.int8: 3.958e15 / 2,
|
42 |
+
},
|
43 |
+
"h100-pcie": {
|
44 |
+
torch.float64: 51e12,
|
45 |
+
torch.float32: 51e12,
|
46 |
+
torch.bfloat16: 1.513e15 / 2,
|
47 |
+
torch.float16: 1.513e15 / 2,
|
48 |
+
torch.int8: 3.026e15 / 2,
|
49 |
+
},
|
50 |
+
# source: https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/a100/pdf/nvidia-a100-datasheet-us-nvidia-1758950-r4-web.pdf
|
51 |
+
# sxm and pcie have same flop counts
|
52 |
+
"a100": {torch.float64: 19.5e12, torch.float32: 19.5e12, torch.bfloat16: 312e12, torch.float16: 312e12},
|
53 |
+
# source: https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/a10/pdf/a10-datasheet.pdf
|
54 |
+
"a10g": {torch.float32: 31.2e12, torch.bfloat16: 125e12, torch.float16: 125e12},
|
55 |
+
# source: https://images.nvidia.com/content/technologies/volta/pdf/volta-v100-datasheet-update-us-1165301-r5.pdf
|
56 |
+
"v100-sxm": {torch.float64: 7.8e12, torch.float32: 15.7e12, torch.float16: 125e12},
|
57 |
+
"v100-pcie": {torch.float64: 7e12, torch.float32: 14e12, torch.float16: 112e12},
|
58 |
+
"v100s-pcie": {torch.float64: 8.2e12, torch.float32: 16.4e12, torch.float16: 130e12},
|
59 |
+
# source: https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/tesla-t4/t4-tensor-core-datasheet-951643.pdf
|
60 |
+
# sxm and pcie have same flop counts
|
61 |
+
"t4": {torch.float32: 8.1e12, torch.float16: 65e12, torch.int8: 130e12},
|
62 |
+
# https://www.nvidia.com/content/dam/en-zz/Solutions/design-visualization/quadro-product-literature/quadro-rtx-5000-data-sheet-us-nvidia-704120-r4-web.pdf
|
63 |
+
"quadro rtx 5000": {torch.float32: 11.2e12, torch.float16: 89.2e12},
|
64 |
+
}
|
65 |
+
|
66 |
+
TPU_AVAILABLE_FLOPS = {
|
67 |
+
# flop count for each TPU generation is the same for all precisions
|
68 |
+
# since bfloat16 precision is always used for performing matrix operations
|
69 |
+
# for more info: https://cloud.google.com/tpu/docs/bfloat16#choosing_bfloat16
|
70 |
+
# source: https://arxiv.org/pdf/1907.10701.pdf
|
71 |
+
"v2": 45e12,
|
72 |
+
# source: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v3
|
73 |
+
"v3": 123e12,
|
74 |
+
# source: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v4
|
75 |
+
"v4": 275e12,
|
76 |
+
# source: https://cloud.google.com/tpu/docs/v5e-training
|
77 |
+
"v5litepod": 197e12,
|
78 |
+
}
|
79 |
+
|
80 |
+
|
81 |
+
def get_flops_available(device: torch.device, dtype: torch.dtype) -> Optional[float]:
|
82 |
+
if device.type == "cuda":
|
83 |
+
device_name = torch.cuda.get_device_name(device).lower()
|
84 |
+
if "h100" in device_name and "hbm3" in device_name:
|
85 |
+
device_name = "h100-sxm"
|
86 |
+
elif "h100" in device_name and ("pcie" in device_name or "hbm2e" in device_name):
|
87 |
+
device_name = "h100-pcie"
|
88 |
+
elif "a100" in device_name:
|
89 |
+
device_name = "a100"
|
90 |
+
elif "a10g" in device_name:
|
91 |
+
device_name = "a10g"
|
92 |
+
elif "v100-sxm" in device_name:
|
93 |
+
device_name = "v100-sxm"
|
94 |
+
elif "v100-pcie" in device_name:
|
95 |
+
device_name = "v100-pcie"
|
96 |
+
elif "t4" in device_name:
|
97 |
+
device_name = "t4"
|
98 |
+
elif "quadro rtx 5000" in device_name:
|
99 |
+
device_name = "quadro rtx 5000"
|
100 |
+
else:
|
101 |
+
device_name = None
|
102 |
+
|
103 |
+
if device_name is not None:
|
104 |
+
try:
|
105 |
+
return int(GPU_AVAILABLE_FLOPS[device_name][dtype])
|
106 |
+
except KeyError:
|
107 |
+
raise KeyError(
|
108 |
+
f"flop count not found for {device_name} with dtype: {dtype}; "
|
109 |
+
"MFU cannot be calculated and reported."
|
110 |
+
)
|
111 |
+
elif device.type == "xla":
|
112 |
+
if _XLA_GREATER_EQUAL_2_1:
|
113 |
+
from torch_xla._internal import tpu
|
114 |
+
else:
|
115 |
+
from torch_xla.experimental import tpu
|
116 |
+
|
117 |
+
device_name = tpu.get_tpu_env()["TYPE"].lower()
|
118 |
+
try:
|
119 |
+
return int(TPU_AVAILABLE_FLOPS[device_name])
|
120 |
+
except KeyError:
|
121 |
+
raise KeyError(
|
122 |
+
f"flop count not found for {device_name} with dtype: {dtype}; MFU cannot be calculated and reported."
|
123 |
+
)
|
124 |
+
|
125 |
+
return None
|
126 |
+
|
127 |
+
|
128 |
+
# Adapted from https://github.com/mosaicml/composer/blob/f2a2dc820cb75023b9eb7c46fdfd25273712abd0/composer/callbacks/speed_monitor.py
|
129 |
+
|
130 |
+
|
131 |
+
class SpeedMonitorBase:
|
132 |
+
"""Logs the training throughput and utilization.
|
133 |
+
|
134 |
+
+-------------------------------------+-----------------------------------------------------------+
|
135 |
+
| Key | Logged data |
|
136 |
+
+=====================================+===========================================================+
|
137 |
+
| | Rolling average (over `window_size` most recent |
|
138 |
+
| `throughput/batches_per_sec` | batches) of the number of batches processed per second |
|
139 |
+
| | |
|
140 |
+
+-------------------------------------+-----------------------------------------------------------+
|
141 |
+
| | Rolling average (over `window_size` most recent |
|
142 |
+
| `throughput/samples_per_sec` | batches) of the number of samples processed per second |
|
143 |
+
| | |
|
144 |
+
+-------------------------------------+-----------------------------------------------------------+
|
145 |
+
| | Rolling average (over `window_size` most recent |
|
146 |
+
| `throughput/tokens_per_sec` | batches) of the number of tokens processed per second. |
|
147 |
+
| | This may include padding depending on dataset |
|
148 |
+
+-------------------------------------+-----------------------------------------------------------+
|
149 |
+
| | Estimates flops by `flops_per_batch * batches_per_sec` |
|
150 |
+
| `throughput/flops_per_sec` | |
|
151 |
+
| | |
|
152 |
+
+-------------------------------------+-----------------------------------------------------------+
|
153 |
+
| `throughput/device/batches_per_sec` | `throughput/batches_per_sec` divided by world size |
|
154 |
+
+-------------------------------------+-----------------------------------------------------------+
|
155 |
+
| `throughput/device/samples_per_sec` | `throughput/samples_per_sec` divided by world size |
|
156 |
+
+-------------------------------------+-----------------------------------------------------------+
|
157 |
+
| | `throughput/tokens_per_sec` divided by world size. This |
|
158 |
+
| `throughput/device/tokens_per_sec` | may include pad tokens depending on dataset |
|
159 |
+
| | |
|
160 |
+
+-------------------------------------+-----------------------------------------------------------+
|
161 |
+
| | `throughput/flops_per_sec` divided by world size. Only |
|
162 |
+
| `throughput/device/flops_per_sec` | logged when model has attribute `flops_per_batch` |
|
163 |
+
| | |
|
164 |
+
+-------------------------------------+-----------------------------------------------------------+
|
165 |
+
| | `throughput/device/flops_per_sec` divided by world size. |
|
166 |
+
| `throughput/device/mfu` | |
|
167 |
+
| | |
|
168 |
+
+-------------------------------------+-----------------------------------------------------------+
|
169 |
+
| `time/train` | Total elapsed training time |
|
170 |
+
+-------------------------------------+-----------------------------------------------------------+
|
171 |
+
| `time/val` | Total elapsed validation time |
|
172 |
+
+-------------------------------------+-----------------------------------------------------------+
|
173 |
+
| `time/total` | Total elapsed time (time/train + time/val) |
|
174 |
+
+-------------------------------------+-----------------------------------------------------------+
|
175 |
+
|
176 |
+
Notes:
|
177 |
+
- The implementation assumes that devices are homogeneous as it normalizes by the world size.
|
178 |
+
- Tokens/sec, flops/sec and MFU do not account for padding tokens if present. We suggest using samples/sec or
|
179 |
+
batches/sec to measure throughput under this circumstance.
|
180 |
+
- Be careful when comparing MFU numbers across projects, as this will highly depend on the ``flops_per_batch``.
|
181 |
+
There is no widespread, realistic, and reliable implementation to compute them.
|
182 |
+
We suggest using our ``measure_flops`` function, but many other works will use ``estimated_flops`` which
|
183 |
+
will almost always be an overestimate when compared to the true value.
|
184 |
+
|
185 |
+
Args:
|
186 |
+
window_size (int, optional): Number of batches to use for a rolling average of throughput.
|
187 |
+
Defaults to 100.
|
188 |
+
time_unit (str, optional): Time unit to use for `time` logging. Can be one of
|
189 |
+
'seconds', 'minutes', 'hours', or 'days'. Defaults to 'hours'.
|
190 |
+
"""
|
191 |
+
|
192 |
+
def __init__(
|
193 |
+
self,
|
194 |
+
flops_available: float,
|
195 |
+
log_dict: Callable[[Dict, int], None],
|
196 |
+
window_size: int = 100,
|
197 |
+
time_unit: str = "hours",
|
198 |
+
):
|
199 |
+
self.flops_available = flops_available
|
200 |
+
self.log_dict = log_dict
|
201 |
+
|
202 |
+
# Track the batch num samples and wct to compute throughput over a window of batches
|
203 |
+
self.history_samples: Deque[int] = deque(maxlen=window_size + 1)
|
204 |
+
self.history_wct: Deque[float] = deque(maxlen=window_size + 1)
|
205 |
+
self.history_lengths: Deque[int] = deque(maxlen=window_size + 1)
|
206 |
+
self.history_flops: Deque[int] = deque(maxlen=window_size + 1)
|
207 |
+
|
208 |
+
self.divider = 1
|
209 |
+
if time_unit == "seconds":
|
210 |
+
self.divider = 1
|
211 |
+
elif time_unit == "minutes":
|
212 |
+
self.divider = 60
|
213 |
+
elif time_unit == "hours":
|
214 |
+
self.divider = 60 * 60
|
215 |
+
elif time_unit == "days":
|
216 |
+
self.divider = 60 * 60 * 24
|
217 |
+
else:
|
218 |
+
raise ValueError(
|
219 |
+
f'Invalid time_unit: {time_unit}. Must be one of "seconds", "minutes", "hours", or "days".'
|
220 |
+
)
|
221 |
+
|
222 |
+
# Keep track of time spent evaluating
|
223 |
+
self.total_eval_wct = 0.0
|
224 |
+
self.step = -1
|
225 |
+
|
226 |
+
def on_train_batch_end(
|
227 |
+
self,
|
228 |
+
samples: int, # total samples seen (per device)
|
229 |
+
train_elapsed: float, # total training time (seconds)
|
230 |
+
world_size: int,
|
231 |
+
flops_per_batch: Optional[int] = None, # (per device)
|
232 |
+
lengths: Optional[int] = None, # total length of the samples seen (per device)
|
233 |
+
) -> None:
|
234 |
+
self.step += 1
|
235 |
+
step = self.step
|
236 |
+
metrics = {}
|
237 |
+
|
238 |
+
self.history_samples.append(samples)
|
239 |
+
if lengths is not None:
|
240 |
+
self.history_lengths.append(lengths)
|
241 |
+
# if lengths are passed, there should be as many values as samples
|
242 |
+
assert len(self.history_samples) == len(self.history_lengths)
|
243 |
+
self.history_wct.append(train_elapsed)
|
244 |
+
if len(self.history_wct) == self.history_wct.maxlen:
|
245 |
+
elapsed_batches = len(self.history_samples) - 1
|
246 |
+
elapsed_samples = self.history_samples[-1] - self.history_samples[0]
|
247 |
+
elapsed_wct = self.history_wct[-1] - self.history_wct[0]
|
248 |
+
samples_per_sec = elapsed_samples * world_size / elapsed_wct
|
249 |
+
dev_samples_per_sec = elapsed_samples / elapsed_wct
|
250 |
+
metrics.update(
|
251 |
+
{
|
252 |
+
"throughput/batches_per_sec": elapsed_batches * world_size / elapsed_wct,
|
253 |
+
"throughput/samples_per_sec": samples_per_sec,
|
254 |
+
"throughput/device/batches_per_sec": elapsed_batches / elapsed_wct,
|
255 |
+
"throughput/device/samples_per_sec": dev_samples_per_sec,
|
256 |
+
}
|
257 |
+
)
|
258 |
+
if lengths is not None:
|
259 |
+
elapsed_lengths = int(self.history_lengths[-1]) - int(self.history_lengths[0])
|
260 |
+
avg_length = elapsed_lengths / elapsed_batches
|
261 |
+
metrics.update(
|
262 |
+
{
|
263 |
+
"throughput/tokens_per_sec": samples_per_sec * avg_length,
|
264 |
+
"throughput/device/tokens_per_sec": dev_samples_per_sec * avg_length,
|
265 |
+
}
|
266 |
+
)
|
267 |
+
|
268 |
+
if flops_per_batch is not None:
|
269 |
+
# sum of flops per batch across ranks
|
270 |
+
self.history_flops.append(flops_per_batch * world_size)
|
271 |
+
if len(self.history_flops) == self.history_flops.maxlen:
|
272 |
+
elapsed_flops = sum(self.history_flops) - self.history_flops[0]
|
273 |
+
elapsed_wct = self.history_wct[-1] - self.history_wct[0]
|
274 |
+
flops_per_sec = elapsed_flops / elapsed_wct
|
275 |
+
device_flops_per_sec = flops_per_sec / world_size
|
276 |
+
metrics.update(
|
277 |
+
{"throughput/flops_per_sec": flops_per_sec, "throughput/device/flops_per_sec": device_flops_per_sec}
|
278 |
+
)
|
279 |
+
if self.flops_available:
|
280 |
+
metrics["throughput/device/mfu"] = device_flops_per_sec / self.flops_available
|
281 |
+
|
282 |
+
metrics.update(
|
283 |
+
{
|
284 |
+
"time/train": train_elapsed / self.divider,
|
285 |
+
"time/val": self.total_eval_wct / self.divider,
|
286 |
+
"time/total": (train_elapsed + self.total_eval_wct) / self.divider,
|
287 |
+
"samples": samples,
|
288 |
+
}
|
289 |
+
)
|
290 |
+
|
291 |
+
self.log_dict(metrics, step)
|
292 |
+
|
293 |
+
def eval_end(self, eval_elapsed: float) -> None:
|
294 |
+
self.total_eval_wct += eval_elapsed # seconds
|
295 |
+
|
296 |
+
|
297 |
+
def plugin_to_compute_dtype(plugin: Precision) -> torch.dtype:
|
298 |
+
if isinstance(plugin, BitsandbytesPrecision):
|
299 |
+
return plugin.dtype
|
300 |
+
if isinstance(plugin, (HalfPrecision, MixedPrecision, HalfPrecisionPlugin)):
|
301 |
+
return plugin._desired_input_dtype
|
302 |
+
if isinstance(plugin, MixedPrecisionPlugin):
|
303 |
+
return torch.bfloat16 if plugin.precision == "bf16-mixed" else torch.half
|
304 |
+
if isinstance(plugin, (DoublePrecision, DoublePrecisionPlugin)):
|
305 |
+
return torch.double
|
306 |
+
if isinstance(plugin, (XLAPrecision, XLAPrecisionPlugin)):
|
307 |
+
return plugin._desired_dtype
|
308 |
+
if isinstance(plugin, TransformerEnginePrecision):
|
309 |
+
return torch.int8
|
310 |
+
if isinstance(plugin, (FSDPPrecision, FSDPPrecisionPlugin)):
|
311 |
+
return plugin.mixed_precision_config.reduce_dtype
|
312 |
+
if isinstance(plugin, Precision):
|
313 |
+
return torch.float32
|
314 |
+
raise NotImplementedError(plugin)
|
315 |
+
|
316 |
+
|
317 |
+
class SpeedMonitorFabric(SpeedMonitorBase):
|
318 |
+
def __init__(self, fabric: Fabric, *args: Any, **kwargs: Any) -> None:
|
319 |
+
dtype = plugin_to_compute_dtype(fabric.strategy.precision)
|
320 |
+
flops_available = get_flops_available(fabric.device, dtype)
|
321 |
+
super().__init__(flops_available, fabric.log_dict, *args, **kwargs)
|
322 |
+
|
323 |
+
@fabric_rank_zero_only
|
324 |
+
def on_train_batch_end(self, *args: Any, **kwargs: Any) -> None:
|
325 |
+
super().on_train_batch_end(*args, **kwargs)
|
326 |
+
|
327 |
+
|
328 |
+
class SpeedMonitorCallback(Callback):
|
329 |
+
def __init__(self, length_fn: Callable[[Any], int], batch_size: int, **kwargs: Any) -> None:
|
330 |
+
super().__init__()
|
331 |
+
self.speed_monitor: Optional[SpeedMonitorBase] = None
|
332 |
+
self.speed_monitor_kwargs = kwargs
|
333 |
+
self.length_fn = length_fn
|
334 |
+
self.batch_size = batch_size
|
335 |
+
self.eval_t0: int = 0
|
336 |
+
self.train_t0: int = 0
|
337 |
+
self.total_lengths: int = 0
|
338 |
+
|
339 |
+
def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
|
340 |
+
if self.speed_monitor is not None:
|
341 |
+
return # already setup
|
342 |
+
dtype = plugin_to_compute_dtype(trainer.precision_plugin)
|
343 |
+
flops_available = get_flops_available(trainer.strategy.root_device, dtype)
|
344 |
+
self.speed_monitor = SpeedMonitorBase(flops_available, trainer.logger.log_metrics, **self.speed_monitor_kwargs)
|
345 |
+
|
346 |
+
@trainer_rank_zero_only
|
347 |
+
def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
|
348 |
+
if trainer.fit_loop._should_accumulate():
|
349 |
+
return
|
350 |
+
|
351 |
+
self.train_t0 = time.perf_counter()
|
352 |
+
|
353 |
+
@trainer_rank_zero_only
|
354 |
+
def on_train_batch_end(
|
355 |
+
self, trainer: Trainer, pl_module: LightningModule, outputs: Any, batch: Any, batch_idx: int
|
356 |
+
) -> None:
|
357 |
+
self.total_lengths += self.length_fn(batch)
|
358 |
+
if trainer.fit_loop._should_accumulate():
|
359 |
+
return
|
360 |
+
train_elapsed = time.perf_counter() - self.train_t0
|
361 |
+
assert self.speed_monitor is not None
|
362 |
+
iter_num = trainer.fit_loop.total_batch_idx
|
363 |
+
assert (measured_flops := pl_module.measured_flops) is not None
|
364 |
+
self.speed_monitor.on_train_batch_end(
|
365 |
+
(iter_num + 1) * self.batch_size,
|
366 |
+
train_elapsed,
|
367 |
+
# this assumes that device FLOPs are the same and that all devices have the same batch size
|
368 |
+
trainer.world_size,
|
369 |
+
flops_per_batch=measured_flops,
|
370 |
+
lengths=self.total_lengths,
|
371 |
+
)
|
372 |
+
|
373 |
+
@trainer_rank_zero_only
|
374 |
+
def on_validation_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
|
375 |
+
self.eval_t0 = time.perf_counter()
|
376 |
+
|
377 |
+
@trainer_rank_zero_only
|
378 |
+
def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
|
379 |
+
eval_elapsed = time.perf_counter() - self.eval_t0
|
380 |
+
assert self.speed_monitor is not None
|
381 |
+
self.speed_monitor.eval_end(eval_elapsed)
|
382 |
+
|
383 |
+
|
384 |
+
def flops_per_param(max_seq_length: int, n_layer: int, n_embd: int, n_params: int) -> int:
|
385 |
+
flops_per_token = 2 * n_params # each parameter is used for a MAC (2 FLOPS) per network operation
|
386 |
+
# this assumes that all samples have a fixed length equal to the block size
|
387 |
+
# which is most likely false during finetuning
|
388 |
+
flops_per_seq = flops_per_token * max_seq_length
|
389 |
+
attn_flops_per_seq = n_layer * 2 * 2 * (n_embd * (max_seq_length**2))
|
390 |
+
return flops_per_seq + attn_flops_per_seq
|
391 |
+
|
392 |
+
|
393 |
+
def estimate_flops(model: GPT) -> int:
|
394 |
+
"""Measures estimated FLOPs for MFU.
|
395 |
+
|
396 |
+
Refs:
|
397 |
+
* https://ar5iv.labs.arxiv.org/html/2205.05198#A1
|
398 |
+
* https://ar5iv.labs.arxiv.org/html/2204.02311#A2
|
399 |
+
"""
|
400 |
+
# using all parameters for this is a naive over estimation because not all model parameters actually contribute to
|
401 |
+
# this FLOP computation (e.g. embedding, norm). For this reason, the result will be higher by a fixed percentage
|
402 |
+
# (~10%) compared to the measured FLOPs, making those lower but more realistic.
|
403 |
+
# For a proper estimate, this needs a more fine-grained calculation as in Appendix A of the paper.
|
404 |
+
n_trainable_params = num_parameters(model, requires_grad=True)
|
405 |
+
trainable_flops = flops_per_param(
|
406 |
+
model.max_seq_length, model.config.n_layer, model.config.n_embd, n_trainable_params
|
407 |
+
)
|
408 |
+
# forward + backward + gradients (assumes no gradient accumulation)
|
409 |
+
ops_per_step = 3 if model.training else 1
|
410 |
+
n_frozen_params = num_parameters(model, requires_grad=False)
|
411 |
+
frozen_flops = flops_per_param(model.max_seq_length, model.config.n_layer, model.config.n_embd, n_frozen_params)
|
412 |
+
# forward + backward
|
413 |
+
frozen_ops_per_step = 2 if model.training else 1
|
414 |
+
return ops_per_step * trainable_flops + frozen_ops_per_step * frozen_flops
|
415 |
+
|
416 |
+
|
417 |
+
def measure_flops(model: GPT, x: torch.Tensor) -> int:
|
418 |
+
"""Measures real FLOPs for HFU"""
|
419 |
+
flop_counter = FlopCounterMode(model, display=False)
|
420 |
+
ctx = nullcontext() if model.training else torch.no_grad()
|
421 |
+
with ctx, flop_counter:
|
422 |
+
y = model(x)
|
423 |
+
if model.training:
|
424 |
+
y.sum().backward()
|
425 |
+
return flop_counter.get_total_flops()
|
tsai_gpt/tokenizer.py
ADDED
@@ -0,0 +1,103 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import json
|
2 |
+
from pathlib import Path
|
3 |
+
from typing import Optional
|
4 |
+
|
5 |
+
import torch
|
6 |
+
|
7 |
+
|
8 |
+
class Tokenizer:
|
9 |
+
def __init__(self, checkpoint_dir: Path) -> None:
|
10 |
+
self.use_bos = self.check_if_bos_token_used(checkpoint_dir)
|
11 |
+
self.bos_id = None
|
12 |
+
self.eos_id = None
|
13 |
+
|
14 |
+
# some checkpoints have both files, `.model` takes precedence
|
15 |
+
if (vocabulary_path := checkpoint_dir / "tokenizer.model").is_file():
|
16 |
+
from sentencepiece import SentencePieceProcessor
|
17 |
+
|
18 |
+
self.processor = SentencePieceProcessor(model_file=str(vocabulary_path))
|
19 |
+
self.backend = "sentencepiece"
|
20 |
+
self.bos_id = self.processor.bos_id()
|
21 |
+
self.eos_id = self.processor.eos_id()
|
22 |
+
|
23 |
+
elif (vocabulary_path := checkpoint_dir / "tokenizer.json").is_file():
|
24 |
+
from tokenizers import Tokenizer as HFTokenizer
|
25 |
+
|
26 |
+
self.processor = HFTokenizer.from_file(str(vocabulary_path))
|
27 |
+
self.backend = "huggingface"
|
28 |
+
|
29 |
+
if (special_tokens_path := checkpoint_dir / "tokenizer_config.json").is_file():
|
30 |
+
with open(special_tokens_path) as fp:
|
31 |
+
config = json.load(fp)
|
32 |
+
bos_token = config.get("bos_token")
|
33 |
+
self.bos_id = self.token_to_id(bos_token) if bos_token is not None else None
|
34 |
+
eos_token = config.get("eos_token")
|
35 |
+
self.eos_id = self.token_to_id(eos_token) if eos_token is not None else None
|
36 |
+
if (special_tokens_path := checkpoint_dir / "generation_config.json").is_file():
|
37 |
+
with open(special_tokens_path) as fp:
|
38 |
+
config = json.load(fp)
|
39 |
+
if self.bos_id is None:
|
40 |
+
self.bos_id = config.get("bos_token_id")
|
41 |
+
if self.eos_id is None:
|
42 |
+
self.eos_id = config.get("eos_token_id")
|
43 |
+
else:
|
44 |
+
raise NotImplementedError
|
45 |
+
|
46 |
+
@property
|
47 |
+
def vocab_size(self) -> int:
|
48 |
+
if self.backend == "huggingface":
|
49 |
+
return self.processor.get_vocab_size(with_added_tokens=False)
|
50 |
+
if self.backend == "sentencepiece":
|
51 |
+
return self.processor.vocab_size()
|
52 |
+
raise RuntimeError
|
53 |
+
|
54 |
+
def token_to_id(self, token: str) -> int:
|
55 |
+
if self.backend == "huggingface":
|
56 |
+
id_ = self.processor.token_to_id(token)
|
57 |
+
elif self.backend == "sentencepiece":
|
58 |
+
id_ = self.processor.piece_to_id(token)
|
59 |
+
else:
|
60 |
+
raise RuntimeError
|
61 |
+
if id_ is None:
|
62 |
+
raise ValueError(f"token {token!r} not found in the collection.")
|
63 |
+
return id_
|
64 |
+
|
65 |
+
def check_if_bos_token_used(self, checkpoint_dir: Path) -> bool:
|
66 |
+
if not (tokenizer_config_path := checkpoint_dir / "tokenizer_config.json").is_file():
|
67 |
+
return False
|
68 |
+
with open(tokenizer_config_path) as fp:
|
69 |
+
config = json.load(fp)
|
70 |
+
if any(config.get(check, False) for check in ("add_bos_token", "add_prefix_space")):
|
71 |
+
return True
|
72 |
+
# for examples that also use the Llama tokenizer, but do not have or set add_bos_token to True.
|
73 |
+
# ex: https://huggingface.co/stabilityai/StableBeluga2/blob/main/tokenizer_config.json#L2
|
74 |
+
return config.get("add_bos_token") is None and config.get("tokenizer_class") == "LlamaTokenizer"
|
75 |
+
|
76 |
+
def encode(
|
77 |
+
self,
|
78 |
+
string: str,
|
79 |
+
device: Optional[torch.device] = None,
|
80 |
+
bos: Optional[bool] = None,
|
81 |
+
eos: bool = False,
|
82 |
+
max_length: int = -1,
|
83 |
+
) -> torch.Tensor:
|
84 |
+
if self.backend == "huggingface":
|
85 |
+
tokens = self.processor.encode(string).ids
|
86 |
+
elif self.backend == "sentencepiece":
|
87 |
+
tokens = self.processor.encode(string)
|
88 |
+
else:
|
89 |
+
raise RuntimeError
|
90 |
+
if bos or (bos is None and self.use_bos):
|
91 |
+
bos_id = self.bos_id
|
92 |
+
if bos_id is None:
|
93 |
+
raise NotImplementedError("This tokenizer does not have a defined a bos token")
|
94 |
+
tokens = [bos_id] + tokens
|
95 |
+
if eos:
|
96 |
+
tokens = tokens + [self.eos_id]
|
97 |
+
if max_length > 0:
|
98 |
+
tokens = tokens[:max_length]
|
99 |
+
return torch.tensor(tokens, dtype=torch.int, device=device)
|
100 |
+
|
101 |
+
def decode(self, tensor: torch.Tensor) -> str:
|
102 |
+
tokens = [tensor.item()] if tensor.ndim == 0 else tensor.tolist()
|
103 |
+
return self.processor.decode(tokens)
|
tsai_gpt/utils.py
ADDED
@@ -0,0 +1,347 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""Utility functions for training and inference."""
|
2 |
+
import math
|
3 |
+
import pickle
|
4 |
+
import sys
|
5 |
+
from contextlib import nullcontext
|
6 |
+
from io import BytesIO
|
7 |
+
from pathlib import Path
|
8 |
+
from typing import TYPE_CHECKING, ContextManager, Dict, List, Mapping, Optional, TypeVar, Union
|
9 |
+
|
10 |
+
import lightning as L
|
11 |
+
import torch
|
12 |
+
import torch.nn as nn
|
13 |
+
import torch.utils._device
|
14 |
+
from lightning.fabric.strategies import FSDPStrategy
|
15 |
+
from lightning.fabric.utilities.load import _lazy_load as lazy_load
|
16 |
+
from torch.serialization import normalize_storage_type
|
17 |
+
|
18 |
+
if TYPE_CHECKING:
|
19 |
+
from model import GPT
|
20 |
+
|
21 |
+
|
22 |
+
def find_multiple(n: int, k: int) -> int:
|
23 |
+
assert k > 0
|
24 |
+
if n % k == 0:
|
25 |
+
return n
|
26 |
+
return n + k - (n % k)
|
27 |
+
|
28 |
+
|
29 |
+
def num_parameters(module: nn.Module, requires_grad: Optional[bool] = None) -> int:
|
30 |
+
total = 0
|
31 |
+
for p in module.parameters():
|
32 |
+
if requires_grad is None or p.requires_grad == requires_grad:
|
33 |
+
if hasattr(p, "quant_state"):
|
34 |
+
# bitsandbytes 4bit layer support
|
35 |
+
total += math.prod(p.quant_state[1])
|
36 |
+
else:
|
37 |
+
total += p.numel()
|
38 |
+
return total
|
39 |
+
|
40 |
+
|
41 |
+
def gptq_quantization(enabled: bool = False) -> ContextManager:
|
42 |
+
if not enabled:
|
43 |
+
return nullcontext()
|
44 |
+
|
45 |
+
from lightning.fabric.plugins.precision.utils import _ClassReplacementContextManager
|
46 |
+
|
47 |
+
from quantize.gptq import ColBlockQuantizedLinear
|
48 |
+
|
49 |
+
class QuantizedLinear(ColBlockQuantizedLinear):
|
50 |
+
def __init__(self, *args, **kwargs):
|
51 |
+
super().__init__(*args, bits=4, tile_cols=-1, **kwargs)
|
52 |
+
|
53 |
+
return _ClassReplacementContextManager({"torch.nn.Linear": QuantizedLinear})
|
54 |
+
|
55 |
+
|
56 |
+
def check_valid_checkpoint_dir(checkpoint_dir: Path) -> None:
|
57 |
+
files = {
|
58 |
+
"lit_model.pth": (checkpoint_dir / "lit_model.pth").is_file(),
|
59 |
+
"lit_config.json": (checkpoint_dir / "lit_config.json").is_file(),
|
60 |
+
"tokenizer.json OR tokenizer.model": (checkpoint_dir / "tokenizer.json").is_file() or (
|
61 |
+
checkpoint_dir / "tokenizer.model"
|
62 |
+
).is_file(),
|
63 |
+
"tokenizer_config.json": (checkpoint_dir / "tokenizer_config.json").is_file(),
|
64 |
+
}
|
65 |
+
if checkpoint_dir.is_dir():
|
66 |
+
if all(files.values()):
|
67 |
+
# we're good
|
68 |
+
return
|
69 |
+
problem = f" is missing the files: {[f for f, exists in files.items() if not exists]!r}"
|
70 |
+
else:
|
71 |
+
problem = " is not a checkpoint directory"
|
72 |
+
|
73 |
+
# list locally available checkpoints
|
74 |
+
available = list(Path("checkpoints").glob("*/*"))
|
75 |
+
if available:
|
76 |
+
options = "\n --checkpoint_dir ".join([""] + [repr(str(p.resolve())) for p in available])
|
77 |
+
extra = f"\nYou have downloaded locally:{options}\n"
|
78 |
+
else:
|
79 |
+
extra = ""
|
80 |
+
|
81 |
+
error_message = (
|
82 |
+
f"--checkpoint_dir {str(checkpoint_dir.absolute())!r}{problem}."
|
83 |
+
"\nFind download instructions at https://github.com/Lightning-AI/lit-gpt/blob/main/tutorials\n"
|
84 |
+
f"{extra}\nSee all download options by running:\n python scripts/download.py"
|
85 |
+
)
|
86 |
+
print(error_message, file=sys.stderr)
|
87 |
+
raise SystemExit(1)
|
88 |
+
|
89 |
+
|
90 |
+
class SavingProxyForStorage:
|
91 |
+
def __init__(self, obj, saver, protocol_version=5):
|
92 |
+
self.protocol_version = protocol_version
|
93 |
+
self.saver = saver
|
94 |
+
if not (isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj)):
|
95 |
+
raise TypeError(f"expected storage, not {type(obj)}")
|
96 |
+
|
97 |
+
# this logic is taken from PyTorch 2.0+ torch/serialization.py
|
98 |
+
if isinstance(obj, torch.storage.TypedStorage):
|
99 |
+
# PT upstream wants to deprecate this eventually...
|
100 |
+
storage = obj._untyped_storage
|
101 |
+
storage_type_str = obj._pickle_storage_type()
|
102 |
+
storage_type = getattr(torch, storage_type_str)
|
103 |
+
storage_numel = obj._size()
|
104 |
+
else:
|
105 |
+
storage = obj
|
106 |
+
storage_type = normalize_storage_type(type(obj))
|
107 |
+
storage_numel = storage.nbytes()
|
108 |
+
|
109 |
+
storage_key = saver._write_storage_and_return_key(storage)
|
110 |
+
location = torch.serialization.location_tag(storage)
|
111 |
+
|
112 |
+
self.storage_info = ("storage", storage_type, storage_key, location, storage_numel)
|
113 |
+
|
114 |
+
def __reduce_ex__(self, protocol_version):
|
115 |
+
assert False, "this should be handled with out of band"
|
116 |
+
|
117 |
+
|
118 |
+
class SavingProxyForTensor:
|
119 |
+
def __init__(self, tensor, saver, protocol_version=5):
|
120 |
+
self.protocol_version = protocol_version
|
121 |
+
self.reduce_ret_fn, reduce_args = tensor.__reduce_ex__(protocol_version)
|
122 |
+
if reduce_args[0] == torch._utils._rebuild_tensor_v2:
|
123 |
+
# for Tensors with Python attributes
|
124 |
+
(a0, a1, (storage, *a2_other), *other_reduce_args) = reduce_args
|
125 |
+
assert isinstance(storage, torch.storage.TypedStorage), "Please check for updates"
|
126 |
+
storage_proxy = SavingProxyForStorage(storage, saver, protocol_version=protocol_version)
|
127 |
+
self.reduce_args = (a0, a1, (storage_proxy, *a2_other), *other_reduce_args)
|
128 |
+
else:
|
129 |
+
(storage, *other_reduce_args) = reduce_args
|
130 |
+
assert isinstance(storage, torch.storage.TypedStorage), "Please check for updates"
|
131 |
+
storage_proxy = SavingProxyForStorage(storage, saver, protocol_version=protocol_version)
|
132 |
+
self.reduce_args = (storage_proxy, *other_reduce_args)
|
133 |
+
|
134 |
+
def __reduce_ex__(self, protocol_version):
|
135 |
+
if protocol_version != self.protocol_version:
|
136 |
+
raise RuntimeError(f"Unexpected protocol version: expected {self.protocol_version}, got {protocol_version}")
|
137 |
+
return self.reduce_ret_fn, self.reduce_args
|
138 |
+
|
139 |
+
|
140 |
+
class IncrementalPyTorchPickler(pickle.Pickler):
|
141 |
+
def __init__(self, saver, *args, **kwargs):
|
142 |
+
super().__init__(*args, **kwargs)
|
143 |
+
self.storage_dtypes = {}
|
144 |
+
self.saver = saver
|
145 |
+
self.id_map = {}
|
146 |
+
|
147 |
+
# this logic is taken from PyTorch 2.0+ torch/serialization.py
|
148 |
+
def persistent_id(self, obj):
|
149 |
+
# FIXME: the docs say that persistent_id should only return a string
|
150 |
+
# but torch store returns tuples. This works only in the binary protocol
|
151 |
+
# see
|
152 |
+
# https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects
|
153 |
+
# https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537
|
154 |
+
if isinstance(obj, SavingProxyForStorage):
|
155 |
+
return obj.storage_info
|
156 |
+
|
157 |
+
if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj):
|
158 |
+
if isinstance(obj, torch.storage.TypedStorage):
|
159 |
+
# TODO: Once we decide to break serialization FC, this case
|
160 |
+
# can be deleted
|
161 |
+
storage = obj._untyped_storage
|
162 |
+
storage_dtype = obj.dtype
|
163 |
+
storage_type_str = obj._pickle_storage_type()
|
164 |
+
storage_type = getattr(torch, storage_type_str)
|
165 |
+
storage_numel = obj._size()
|
166 |
+
|
167 |
+
else:
|
168 |
+
storage = obj
|
169 |
+
storage_dtype = torch.uint8
|
170 |
+
storage_type = normalize_storage_type(type(obj))
|
171 |
+
storage_numel = storage.nbytes()
|
172 |
+
|
173 |
+
# If storage is allocated, ensure that any other saved storages
|
174 |
+
# pointing to the same data all have the same dtype. If storage is
|
175 |
+
# not allocated, don't perform this check
|
176 |
+
if storage.data_ptr() != 0:
|
177 |
+
if storage.data_ptr() in self.storage_dtypes:
|
178 |
+
if storage_dtype != self.storage_dtypes[storage.data_ptr()]:
|
179 |
+
raise RuntimeError(
|
180 |
+
"Cannot save multiple tensors or storages that view the same data as different types"
|
181 |
+
)
|
182 |
+
else:
|
183 |
+
self.storage_dtypes[storage.data_ptr()] = storage_dtype
|
184 |
+
|
185 |
+
storage_key = self.id_map.get(storage._cdata)
|
186 |
+
if storage_key is None:
|
187 |
+
storage_key = self.saver._write_storage_and_return_key(storage)
|
188 |
+
self.id_map[storage._cdata] = storage_key
|
189 |
+
location = torch.serialization.location_tag(storage)
|
190 |
+
|
191 |
+
return ("storage", storage_type, storage_key, location, storage_numel)
|
192 |
+
|
193 |
+
return None
|
194 |
+
|
195 |
+
|
196 |
+
class incremental_save:
|
197 |
+
def __init__(self, name):
|
198 |
+
self.name = name
|
199 |
+
self.zipfile = torch._C.PyTorchFileWriter(str(name))
|
200 |
+
self.has_saved = False
|
201 |
+
self.next_key = 0
|
202 |
+
|
203 |
+
def __enter__(self):
|
204 |
+
return self
|
205 |
+
|
206 |
+
def store_early(self, tensor):
|
207 |
+
if isinstance(tensor, torch.Tensor):
|
208 |
+
return SavingProxyForTensor(tensor, self)
|
209 |
+
raise TypeError(f"can only store tensors early, not {type(tensor)}")
|
210 |
+
|
211 |
+
def save(self, obj):
|
212 |
+
if self.has_saved:
|
213 |
+
raise RuntimeError("have already saved")
|
214 |
+
# Write the pickle data for `obj`
|
215 |
+
data_buf = BytesIO()
|
216 |
+
pickler = IncrementalPyTorchPickler(self, data_buf, protocol=5)
|
217 |
+
pickler.dump(obj)
|
218 |
+
data_value = data_buf.getvalue()
|
219 |
+
self.zipfile.write_record("data.pkl", data_value, len(data_value))
|
220 |
+
self.has_saved = True
|
221 |
+
|
222 |
+
def _write_storage_and_return_key(self, storage):
|
223 |
+
if self.has_saved:
|
224 |
+
raise RuntimeError("have already saved")
|
225 |
+
key = self.next_key
|
226 |
+
self.next_key += 1
|
227 |
+
name = f"data/{key}"
|
228 |
+
if storage.device.type != "cpu":
|
229 |
+
storage = storage.cpu()
|
230 |
+
num_bytes = storage.nbytes()
|
231 |
+
self.zipfile.write_record(name, storage.data_ptr(), num_bytes)
|
232 |
+
return key
|
233 |
+
|
234 |
+
def __exit__(self, type, value, traceback):
|
235 |
+
self.zipfile.write_end_of_file()
|
236 |
+
|
237 |
+
|
238 |
+
T = TypeVar("T")
|
239 |
+
|
240 |
+
|
241 |
+
def chunked_cross_entropy(
|
242 |
+
logits: Union[torch.Tensor, List[torch.Tensor]], targets: torch.Tensor, chunk_size: int = 128
|
243 |
+
) -> torch.Tensor:
|
244 |
+
# with large max_sequence_lengths, the beginning of `backward` allocates a large memory chunk which can dominate
|
245 |
+
# the memory usage in fine-tuning settings with low number of parameters.
|
246 |
+
# as a workaround hack, the cross entropy computation is chunked to force it to deallocate on the go, reducing
|
247 |
+
# the memory spike's magnitude
|
248 |
+
|
249 |
+
# lm_head was chunked (we are fine-tuning)
|
250 |
+
if isinstance(logits, list):
|
251 |
+
# don't want to chunk cross entropy
|
252 |
+
if chunk_size == 0:
|
253 |
+
logits = torch.cat(logits, dim=1)
|
254 |
+
logits = logits.reshape(-1, logits.size(-1))
|
255 |
+
targets = targets.reshape(-1)
|
256 |
+
return torch.nn.functional.cross_entropy(logits, targets, ignore_index=-1)
|
257 |
+
|
258 |
+
# chunk cross entropy
|
259 |
+
logit_chunks = [logit_chunk.reshape(-1, logit_chunk.size(-1)) for logit_chunk in logits]
|
260 |
+
target_chunks = [target_chunk.reshape(-1) for target_chunk in targets.split(logits[0].size(1), dim=1)]
|
261 |
+
loss_chunks = [
|
262 |
+
torch.nn.functional.cross_entropy(logit_chunk, target_chunk, ignore_index=-1, reduction="none")
|
263 |
+
for logit_chunk, target_chunk in zip(logit_chunks, target_chunks)
|
264 |
+
]
|
265 |
+
return torch.cat(loss_chunks).mean()
|
266 |
+
|
267 |
+
# no chunking at all
|
268 |
+
logits = logits.reshape(-1, logits.size(-1))
|
269 |
+
targets = targets.reshape(-1)
|
270 |
+
if chunk_size == 0:
|
271 |
+
return torch.nn.functional.cross_entropy(logits, targets, ignore_index=-1)
|
272 |
+
|
273 |
+
# lm_head wasn't chunked, chunk cross entropy
|
274 |
+
logit_chunks = logits.split(chunk_size)
|
275 |
+
target_chunks = targets.split(chunk_size)
|
276 |
+
loss_chunks = [
|
277 |
+
torch.nn.functional.cross_entropy(logit_chunk, target_chunk, ignore_index=-1, reduction="none")
|
278 |
+
for logit_chunk, target_chunk in zip(logit_chunks, target_chunks)
|
279 |
+
]
|
280 |
+
return torch.cat(loss_chunks).mean()
|
281 |
+
|
282 |
+
|
283 |
+
def map_old_state_dict_weights(state_dict: Dict, mapping: Mapping, prefix: str) -> Dict:
|
284 |
+
for checkpoint_name, attribute_name in mapping.items():
|
285 |
+
full_checkpoint_name = prefix + checkpoint_name
|
286 |
+
if full_checkpoint_name in state_dict:
|
287 |
+
full_attribute_name = prefix + attribute_name
|
288 |
+
state_dict[full_attribute_name] = state_dict.pop(full_checkpoint_name)
|
289 |
+
return state_dict
|
290 |
+
|
291 |
+
|
292 |
+
def get_default_supported_precision(training: bool) -> str:
|
293 |
+
"""Return default precision that is supported by the hardware: either `bf16` or `16`.
|
294 |
+
|
295 |
+
Args:
|
296 |
+
training: `-mixed` or `-true` version of the precision to use
|
297 |
+
|
298 |
+
Returns:
|
299 |
+
default precision that is suitable for the task and is supported by the hardware
|
300 |
+
"""
|
301 |
+
from lightning.fabric.accelerators import MPSAccelerator
|
302 |
+
|
303 |
+
if MPSAccelerator.is_available() or (torch.cuda.is_available() and not torch.cuda.is_bf16_supported()):
|
304 |
+
return "16-mixed" if training else "16-true"
|
305 |
+
return "bf16-mixed" if training else "bf16-true"
|
306 |
+
|
307 |
+
|
308 |
+
def load_checkpoint(fabric: L.Fabric, model: nn.Module, checkpoint_path: Path, strict: bool = True) -> None:
|
309 |
+
if isinstance(fabric.strategy, FSDPStrategy):
|
310 |
+
fabric.load_raw(checkpoint_path, model, strict=strict)
|
311 |
+
else:
|
312 |
+
state_dict = lazy_load(checkpoint_path)
|
313 |
+
state_dict = state_dict.get("model", state_dict)
|
314 |
+
model.load_state_dict(state_dict, strict=strict)
|
315 |
+
|
316 |
+
|
317 |
+
def flops_per_param(max_seq_length: int, n_layer: int, n_embd: int, n_params: int) -> int:
|
318 |
+
flops_per_token = 2 * n_params # each parameter is used for a MAC (2 FLOPS) per network operation
|
319 |
+
# this assumes that all samples have a fixed length equal to the block size
|
320 |
+
# which is most likely false during finetuning
|
321 |
+
flops_per_seq = flops_per_token * max_seq_length
|
322 |
+
attn_flops_per_seq = n_layer * 2 * 2 * (n_embd * (max_seq_length**2))
|
323 |
+
return flops_per_seq + attn_flops_per_seq
|
324 |
+
|
325 |
+
|
326 |
+
def estimate_flops(model: "GPT", training: bool) -> int:
|
327 |
+
"""Measures estimated FLOPs for MFU.
|
328 |
+
|
329 |
+
Refs:
|
330 |
+
* https://ar5iv.labs.arxiv.org/html/2205.05198#A1
|
331 |
+
* https://ar5iv.labs.arxiv.org/html/2204.02311#A2
|
332 |
+
"""
|
333 |
+
# using all parameters for this is a naive over estimation because not all model parameters actually contribute to
|
334 |
+
# this FLOP computation (e.g. embedding, norm). For this reason, the result will be higher by a fixed percentage
|
335 |
+
# (~10%) compared to the measured FLOPs, making those lower but more realistic.
|
336 |
+
# For a proper estimate, this needs a more fine-grained calculation as in Appendix A of the paper.
|
337 |
+
n_trainable_params = num_parameters(model, requires_grad=True)
|
338 |
+
trainable_flops = flops_per_param(
|
339 |
+
model.max_seq_length, model.config.n_layer, model.config.n_embd, n_trainable_params
|
340 |
+
)
|
341 |
+
# forward + backward + gradients (assumes no gradient accumulation)
|
342 |
+
ops_per_step = 3 if training else 1
|
343 |
+
n_frozen_params = num_parameters(model, requires_grad=False)
|
344 |
+
frozen_flops = flops_per_param(model.max_seq_length, model.config.n_layer, model.config.n_embd, n_frozen_params)
|
345 |
+
# forward + backward
|
346 |
+
frozen_ops_per_step = 2 if training else 1
|
347 |
+
return ops_per_step * trainable_flops + frozen_ops_per_step * frozen_flops
|