125m.yaml

data_local: &data_local ./my-copy-c4
data_remote: &data_remote null
tokenizer_name: &tokenizer_name gpt2
max_seq_len: &max_seq_len 2048
global_seed: &global_seed 17

# Run Name
run_name: gpt-125m

# Model
model:
  name: mosaic_gpt
  device: meta
  tokenizer_name: *tokenizer_name
  d_model: 768
  n_heads: 12
  n_layers: 12
  mlp_ratio: 4
  max_seq_len: *max_seq_len
  vocab_size: 50257
  init_std: 0.02
  attn_pdrop: 0.0
  resid_pdrop: 0.0
  emb_pdrop: 0.0
  attn_impl: flash

# Tokenizer
tokenizer:
  type: hftokenizer
  args:
    tokenizer_name: *tokenizer_name
    max_seq_len: *max_seq_len

# Dataloaders
train_loader:
  name: text
  dataset:
    local: *data_local
    remote: *data_remote
    split: train
    shuffle: true
    tokenizer_name: *tokenizer_name
    max_seq_len: *max_seq_len
    group_method: concat
    shuffle_seed: *global_seed
  drop_last: true
  num_workers: 8

eval_loader:
  name: text
  dataset:
    local: *data_local
    remote: *data_remote
    split: val
    shuffle: false
    tokenizer_name: *tokenizer_name
    max_seq_len: *max_seq_len
    group_method: truncate
    shuffle_seed: *global_seed
  drop_last: false
  num_workers: 8

# Optimization
scheduler:
  name: cosine_with_warmup
  t_warmup: 100ba
  alpha_f: 0.1

optimizer:
  name: decoupled_adamw
  lr: 6.0e-4
  betas:
  - 0.9
  - 0.95
  eps: 1.0e-08
  weight_decay: 0.0

algorithms:
  gradient_clipping:
    clipping_type: norm
    clipping_threshold: 1.0

max_duration: 4800ba # ~ 2.5B tokens
eval_interval: 500ba
eval_subset_num_batches: -1
global_train_batch_size: 256

# System
seed: *global_seed
device_eval_batch_size: 16
device_train_microbatch_size: 16
# device_train_microbatch_size: auto
precision: amp_bf16

# FSDP
fsdp_config:
  sharding_strategy: FULL_SHARD
  min_params: 1e8
  mixed_precision: DEFAULT
  activation_checkpointing: false
  activation_cpu_offload: false
  verbose: true

# Logging
progress_bar: false
log_to_console: true
console_log_interval: 1ba

callbacks:
  speed_monitor:
    window_size: 10
  lr_monitor: {}
  memory_monitor: {}

# loggers:
#   wandb: {}

# Checkpoint to local filesystem or remote object store
# save_interval: 500ba
# save_num_checkpoints_to_keep: 1  # Important, this cleans up checkpoints saved to DISK
# save_folder: ./{run_name}/checkpoints
# save_folder: s3://my-bucket/my-folder/{run_name}/checkpoints

# Load from local filesystem or remote object store
# load_path: ./gpt-125m/checkpoints/latest-rank{rank}.pt
# load_path: s3://my-bucket/my-folder/gpt-125m/checkpoints/latest-rank{rank}.pt