File size: 4,755 Bytes
78c54ec | 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 | diff --git a/train_gpt.py b/train_gpt.py
index a9b3f62..9ae8e7a 100644
--- a/train_gpt.py
+++ b/train_gpt.py
@@ -1494,8 +1494,8 @@ def _load_data_shard(file: Path):
assert nbytes == 2 * num_tokens, "number of tokens read does not match header"
return tokens
-BOS_ID = 50256
-TRAIN_MAX_NUM_DOCS = {16384: 64, 32768: 96, 49152: 128}
+BOS_ID = 0 # Polish BPE <|endoftext|>
+TRAIN_MAX_NUM_DOCS = {16384: 384, 32768: 768, 49152: 1152} # bumped: dense short Polish docs
class Shard:
def __init__(self, tokens: Tensor, world_size: int = 1):
@@ -1602,7 +1602,7 @@ def distributed_data_generator(filename_pattern: str, num_tokens: int, max_seq_l
while True:
num_tokens_local = num_tokens // world_size
- max_num_docs = TRAIN_MAX_NUM_DOCS.get(num_tokens_local, next_multiple_of_n(num_tokens_local // 300, n=128))
+ max_num_docs = TRAIN_MAX_NUM_DOCS.get(num_tokens_local, next_multiple_of_n(num_tokens_local // 48, n=128))
if align_to_bos:
try:
@@ -1669,13 +1669,13 @@ def distributed_data_generator(filename_pattern: str, num_tokens: int, max_seq_l
class Hyperparameters:
# data
data_path = os.environ.get("DATA_PATH", ".")
- train_files: str = os.path.join(data_path, "data/fineweb10B/fineweb_train_*.bin") # input .bin to train on
- val_files: str = os.path.join(data_path, "data/fineweb10B/fineweb_val_*.bin") # input .bin to eval validation loss on
+ train_files: str = os.path.expanduser("~/dynaword/shards/polish_train_*.bin") # input .bin to train on
+ val_files: str = os.path.expanduser("~/dynaword/shards/polish_val_*.bin") # input .bin to eval validation loss on
val_tokens: int = 10485760 # how many tokens of validation data? it's important to keep this fixed for consistent comparisons
# batch sizes
val_batch_size: int = 4 * 64 * 1024 * 8
# schedule
- num_scheduled_iterations: int = 1375 # number of steps to complete lr and ws schedule
+ num_scheduled_iterations: int = 13200 # ~1 epoch of 3.47B Polish tokens # number of steps to complete lr and ws schedule
num_extension_iterations: int = 10 # number of steps to continue training at final lr and ws
# evaluation and logging
run_id: str = f"{uuid.uuid4()}"
@@ -1684,7 +1684,8 @@ class Hyperparameters:
# - (1 + m_r9) * x self-reference fuse on layer 9
# - backout_lambda fully removed (slot dropped from self.scalars; absorbed into MUDD bias init)
val_loss_every: int = 250 # every how many steps to evaluate val loss? 0 for only at the end
- save_checkpoint: bool = False
+ save_checkpoint: bool = True
+ checkpoint_every: int = 500 # save every N steps for crash-resume
run_evals: bool = False # run additional evaluations after training is completed
# bigram hash embedding
bigram_vocab_size: int = 50304 * 15
@@ -2014,7 +2015,7 @@ print0(nvidia_smi())
print0("="*100)
model: nn.Module = GPT(
- vocab_size=50257,
+ vocab_size=32896, # mult of 128, not power-of-2 (Karpathy); tokenizer stays 32768
num_layers=11,
num_heads=6,
head_dim=128,
@@ -2118,12 +2119,11 @@ for step in range(train_steps + 1):
torch.cuda.synchronize()
t0 = time.perf_counter()
+ if master_process and args.save_checkpoint and (last_step or (step > 0 and step % args.checkpoint_every == 0)):
+ log = dict(step=step, code=code, model=model.state_dict(), optimizer=training_manager.get_state())
+ os.makedirs(f"logs/{run_id}", exist_ok=True)
+ torch.save(log, f"logs/{run_id}/state_step{step:06d}.pt")
if last_step:
- if master_process and args.save_checkpoint:
- log = dict(step=step, code=code, model=model.state_dict(), optimizer=training_manager.get_state())
- os.makedirs(f"logs/{run_id}", exist_ok=True)
- torch.save(log, f"logs/{run_id}/state_step{step:06d}.pt")
- # the last step only has the validation loop, so break to avoid training
break
# --------------- TRAINING SECTION -----------------
diff --git a/triton_kernels.py b/triton_kernels.py
index 4f377ce..6b40884 100644
--- a/triton_kernels.py
+++ b/triton_kernels.py
@@ -898,7 +898,7 @@ ce_fwd_bwd_kernel = torch.cuda._compile_kernel(
CE_KERNEL_DECLS + CE_KERNEL_SOURCE,
"ce_fwd_bwd_kernel",
compute_capability="90",
- cuda_include_dirs=["/usr/local/cuda/include/"],
+ cuda_include_dirs=['/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/triton/backends/nvidia/include', '/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/nvidia/cuda_runtime/include'],
nvcc_options=["-lineinfo", "--use_fast_math"],
)
ce_fwd_bwd_kernel.set_shared_memory_config(CE_KERNEL_VOCAB_SIZE * 2)
|