-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtrain_llama2.py
executable file
·126 lines (106 loc) · 3.69 KB
/
train_llama2.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import torch
import os
import urllib
from arch.llama2 import Llama2Model
from data import InMemoryDataset
from modules.trainer import TrainingModule
from torch.utils.data import DataLoader
import tiktoken
from lightning.pytorch.callbacks import ModelCheckpoint
import lightning as L
def main():
# download training data
file_path = "data/the-verdict.txt"
url = "https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/ch02/01_main-chapter-code/the-verdict.txt"
text_data = ''
if not os.path.exists(file_path):
with urllib.request.urlopen(url) as response:
text_data = response.read().decode('utf-8')
with open(file_path, "w", encoding="utf-8") as file:
file.write(text_data)
else:
with open(file_path, "r", encoding="utf-8") as file:
text_data = file.read()
train_ratio = 0.90
split_idx = int(train_ratio * len(text_data))
train_text = text_data[:split_idx]
val_text = text_data[split_idx:]
# configure for training
LLAMA2_CONFIG_7B = {
"vocab_size": 50257, # Vocabulary size
"context_length": 1024, # Context length
"emb_dim": 768, # Embedding dimension
"n_heads": 12, # Number of attention heads
"n_layers": 12, # Number of layers
"hidden_dim": 2752, # Different from GPT2: Size of the intermediate dimension in FeedForward
"dtype": torch.bfloat16 # Different from GPT2: Lower-precision dtype to save memory
}
TRAINING_CONFIG = {
"learning_rate": 5e-4,
"weight_decay": 0.1,
"batch_size": 2,
"num_epochs": 25,
}
DATA_LOADER_CONFIG = {
'num_workers': 0,
'drop_last_train': True,
'shuffle_train': True,
'drop_last_val': False,
"shuffle_val": False,
}
torch.manual_seed(123)
llm_model = Llama2Model(LLAMA2_CONFIG_7B)
tokenizer = tiktoken.get_encoding("gpt2")
training_module = TrainingModule(
model=llm_model,
learning_rate=TRAINING_CONFIG["learning_rate"],
weight_decay=TRAINING_CONFIG["weight_decay"],
tokenizer=tokenizer,
test_string="Every effort moves you",
temperature=0.1,
)
training_loader = DataLoader(
InMemoryDataset(
txt=train_text,
tokenizer=tokenizer,
max_length=LLAMA2_CONFIG_7B["context_length"],
stride=LLAMA2_CONFIG_7B["context_length"],
),
batch_size=TRAINING_CONFIG["batch_size"],
shuffle=DATA_LOADER_CONFIG["shuffle_train"],
drop_last=DATA_LOADER_CONFIG["drop_last_train"],
num_workers=DATA_LOADER_CONFIG["num_workers"],
)
val_loader = DataLoader(
InMemoryDataset(
txt=val_text,
tokenizer=tokenizer,
max_length=LLAMA2_CONFIG_7B["context_length"],
stride=LLAMA2_CONFIG_7B["context_length"],
),
batch_size=TRAINING_CONFIG["batch_size"],
shuffle=DATA_LOADER_CONFIG["shuffle_val"],
drop_last=DATA_LOADER_CONFIG["drop_last_val"],
num_workers=DATA_LOADER_CONFIG["num_workers"],
)
checkpoint_callback = ModelCheckpoint(
save_top_k=1,
monitor="val_loss",
)
trainer = L.Trainer(
max_epochs=TRAINING_CONFIG["num_epochs"],
callbacks=[checkpoint_callback],
accelerator="auto",
devices=[1],
num_sanity_val_steps=TRAINING_CONFIG["batch_size"],
val_check_interval=1,
enable_progress_bar=True,
)
# Train the model
trainer.fit(
model=training_module,
train_dataloaders=training_loader,
val_dataloaders=val_loader,
)
if __name__ == "__main__":
main()