mpt-1b-gaudi2.yaml

data_local: ./my-copy-c4
data_remote: # If blank, files must be present in data_local
max_seq_len: 2048
global_seed: 17
device: hpu

# Run Name
run_name: # If left blank, will be read from env var $RUN_NAME

# Model
model:
  name: mpt_causal_lm
  init_device: hpu
  d_model: 2048
  n_heads: 16 # Modified 24->16 so that d_head == 128 to statisfy FlashAttention
  n_layers: 24
  expansion_ratio: 4
  max_seq_len: ${max_seq_len}
  vocab_size: 50368
  attn_config:
    attn_impl: torch
  loss_fn: torch_crossentropy

# Tokenizer
tokenizer:
  name: EleutherAI/gpt-neox-20b
  kwargs:
    model_max_length: ${max_seq_len}

# Dataloaders
train_loader:
  name: text
  dataset:
    local: ${data_local}
    remote: ${data_remote}
    split: train
    shuffle: true
    max_seq_len: ${max_seq_len}
    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
    max_seq_len: ${max_seq_len}
    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: 2.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: 24800ba # ~ 26B tokens
eval_interval: 2000ba
eval_first: false
eval_subset_num_batches: -1
global_train_batch_size: 512

# System
seed: ${global_seed}
device_eval_batch_size: 4
device_train_microbatch_size: 8
# device_train_microbatch_size: auto
precision: amp_bf16

# FSDP
# fsdp_config:
#  sharding_strategy: FULL_SHARD
#  mixed_precision: PURE
#  activation_checkpointing: false
#  activation_checkpointing_reentrant: false
#  activation_cpu_offload: false
#  limit_all_gathers: true
#  verbose: false
deepspeed_config:
  bf16:
    enabled: true
  gradient_clipping: 1.0
  zero_optimization:
    stage: 2
  zero_allow_untested_optimizer: true

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

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

# loggers:
#   wandb: {}

# Checkpoint to local filesystem or remote object store
# save_interval: 2000ba
# 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-1b/checkpoints/latest-rank{rank}.pt
# load_path: s3://my-bucket/my-folder/gpt-1b/checkpoints/latest-rank{rank}.pt