31 lines
614 B
YAML
31 lines
614 B
YAML
|
# Small LLM ~125M parameters (GPT-2 small equivalent)
|
||
|
|
model:
|
||
|
|
vocab_size: 32000
|
||
|
|
d_model: 768
|
||
|
|
n_layers: 12
|
||
|
|
n_heads: 12
|
||
|
|
n_kv_heads: 12 # MHA (same as n_heads)
|
||
|
|
max_seq_len: 2048
|
||
|
|
rope_theta: 10000.0
|
||
|
|
dropout: 0.0
|
||
|
|
bias: false
|
||
|
|
use_flash_attn: true
|
||
|
|
|
||
|
|
train:
|
||
|
|
max_steps: 100000
|
||
|
|
batch_size: 8 # per GPU
|
||
|
|
grad_accum_steps: 4 # effective batch = 8 * 8 GPUs * 4 = 256
|
||
|
|
lr: 3.0e-4
|
||
|
|
weight_decay: 0.1
|
||
|
|
warmup_steps: 2000
|
||
|
|
max_grad_norm: 1.0
|
||
|
|
log_interval: 10
|
||
|
|
save_interval: 1000
|
||
|
|
eval_interval: 500
|
||
|
|
use_amp: true
|
||
|
|
compile_model: false
|
||
|
|
|
||
|
|
tokenizer:
|
||
|
|
vocab_size: 32000
|
||
|
|
type: bpe
|