36 lines
1.0 KiB
YAML
36 lines
1.0 KiB
YAML
# Small LLM ~125M parameters — FP8 variant (B200 TransformerEngine)
|
|
# Based on small.yaml; only changed fields are listed explicitly.
|
|
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
|
|
use_fp8: true # Enable TransformerEngine FP8 kernels
|
|
|
|
train:
|
|
max_steps: 100000
|
|
batch_size: 8 # per GPU; 8 * 2048 = 16384 tokens → divisible by 8 ✓
|
|
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: false # fp8_autocast replaces torch.autocast
|
|
compile_model: false # torch.compile + TE 2.10 stability not verified
|
|
fp8_amax_history_len: 16
|
|
fp8_amax_compute_algo: "max"
|
|
fp8_format: "MXFP8" # B200 native block scaling (better than HYBRID on Blackwell)
|
|
|
|
tokenizer:
|
|
vocab_size: 32000
|
|
type: bpe
|