k/llm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

set to gpt2 hyprs

Browse Source
This commit is contained in:
k
2026-01-09 12:45:01 -05:00
parent 496916f428
commit c78a31362a
1 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
train.py
View File

@@ -11,11 +11,11 @@ import log
import sys import sys
hypr = { hypr = {
"embed_size": 512, "embed_size": 768,
"n_heads": 8, "n_heads": 12,
"n_blocks": 6, "n_blocks": 12,
"block_size": 256, "block_size": 512,
"batch_size": 16, "batch_size": 8,
"starting_lr": 6e-4, "starting_lr": 6e-4,
"minimum_lr": 6e-5, "minimum_lr": 6e-5,
"warmup": 1_000, "warmup": 1_000,
@@ -25,6 +25,7 @@ hypr = {
"subset": "cosmopedia-v2", "subset": "cosmopedia-v2",
"chat_dataset": "yahma/alpaca-cleaned", "chat_dataset": "yahma/alpaca-cleaned",
"chat_subset": None, "chat_subset": None,
"half": False,
} }
print(Device.DEFAULT) print(Device.DEFAULT)
@@ -32,6 +33,8 @@ chat = len(sys.argv) > 1
if(chat): if(chat):
hypr["dataset"] = hypr["chat_dataset"] hypr["dataset"] = hypr["chat_dataset"]
hypr["subset"] = hypr["chat_subset"] hypr["subset"] = hypr["chat_subset"]
hypr["starting_lr"] *= 0.1
hypr["minimum_lr"] *= 0.1
#for loging #for loging
loger = ThreadPoolExecutor(max_workers=2) loger = ThreadPoolExecutor(max_workers=2)
@@ -49,6 +52,9 @@ batch = data.startDataWorker(dataset,encoding,hypr["batch_size"],hypr["block_siz
model = Transformer(hypr["vocab_size"],hypr["embed_size"],hypr["n_heads"],hypr["n_blocks"],hypr["block_size"]) model = Transformer(hypr["vocab_size"],hypr["embed_size"],hypr["n_heads"],hypr["n_blocks"],hypr["block_size"])
if (chat): if (chat):
load_state_dict(model,safe_load(sys.argv[1])) load_state_dict(model,safe_load(sys.argv[1]))
if hypr["half"]:
from tinygrad import dtypes
model = model.cast(dtypes.float16)
params = nn.state.get_parameters(model) params = nn.state.get_parameters(model)
optimizer = optm.llmOptimizer(params,hypr["steps"],hypr["starting_lr"],hypr["minimum_lr"]) optimizer = optm.llmOptimizer(params,hypr["steps"],hypr["starting_lr"],hypr["minimum_lr"])