File size: 2,045 Bytes
9f66630 | 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 | """
Nano-GPT Configuration
L4-SAFE: Reduced memory usage
"""
import torch
from dataclasses import dataclass
@dataclass
class NanoGPTConfig:
# Model Architecture
vocab_size: int = 32000
n_layers: int = 8 # REDUCED from 12
n_heads: int = 8 # REDUCED from 12
n_embd: int = 512 # REDUCED from 768
block_size: int = 512 # REDUCED from 1024 (KEY!)
dropout: float = 0.1
bias: bool = True
# Training Hyperparameters
batch_size: int = 16 # REDUCED from 32 (KEY!)
gradient_accumulation_steps: int = 8 # INCREASED from 4
learning_rate: float = 3e-4
max_iters: int = 100000
weight_decay: float = 0.1
beta1: float = 0.9
beta2: float = 0.95
grad_clip: float = 1.0
# Learning Rate Scheduling
decay_lr: bool = True
warmup_iters: int = 2000
lr_decay_iters: int = 100000
min_lr: float = 3e-5
# Evaluation & Logging
eval_interval: int = 1000
eval_iters: int = 100 # REDUCED from 200
log_interval: int = 100
# Checkpointing
save_interval: int = 5000
checkpoint_dir: str = "checkpoints"
# Data
dataset_mix: dict = None
# Hardware
device: str = 'cuda' if torch.cuda.is_available() else 'cpu'
dtype: str = 'bfloat16'
compile: bool = False # DISABLED torch.compile (uses more memory!)
# Reproducibility
seed: int = 42
def __post_init__(self):
if self.dataset_mix is None:
self.dataset_mix = {
'fineweb': 1.0
}
@property
def n_params(self):
return (2 * self.vocab_size * self.n_embd +
12 * self.n_layers * self.n_embd * self.n_embd) / 1e6
config = NanoGPTConfig()
if __name__ == "__main__":
print(f"Model size: ~{config.n_params:.1f}M parameters")
print(f"Sequence length: {config.block_size}")
print(f"Batch size: {config.batch_size}")
print(f"Effective batch: {config.batch_size * config.gradient_accumulation_steps}") |