| |
| |
| |
| |
| @@ -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 ----------------- |
| |
| |
| |
| |
| @@ -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) |
|
|
