TharunSivamani commited on
Commit
fc45a7c
1 Parent(s): 4e54cb1

model's utils file

Browse files
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