OLMo-1B.yaml

run_name: OLMo-1B-dolma2-tokenizer-wiki_${oc.env:TIME}
seed: 6198
dry_run: false

wandb:
  name: ${run_name}
  project: olmo-small
  entity: jhchen-ai-4s

model:
  d_model: 2048
  n_heads: 16
  n_layers: 16
  mlp_ratio: 8
  weight_tying: true
  alibi: false
  rope: true
  flash_attention: true  # not available on AMD
  attention_dropout: 0.0
  attention_layer_norm: false
  multi_query_attention: false
  include_bias: false
  block_type: sequential
  layer_norm_type: default
  layer_norm_with_affine: false
  bias_for_layer_norm: false
  attention_layer_norm_with_affine: false
  activation_type: swiglu
  residual_dropout: 0.0
  embedding_dropout: 0.0
  max_sequence_length: 2048
  vocab_size: 100278
  embedding_size: 100352
  eos_token_id: 100257
  pad_token_id: 100277
  init_device: meta
  init_fn: mitchell

compile: null  # causes instability on AMD GPUs

optimizer:
  name: adamw
  learning_rate: 4.0e-4
  weight_decay: 0.1
  betas:
  - 0.9
  - 0.95
  metrics_log_interval: 10

scheduler:
  name: cosine_with_warmup
  t_warmup: 2000
  alpha_f: 0.1

tokenizer:
  identifier: tokenizers/allenai_dolma2.json
  truncate_direction: right

save_folder: workspace/${run_name}_${oc.env:SLURM_JOB_ID}
save_overwrite: false
# Sharded checkpoints (best for restarts)
save_interval: 50
save_num_checkpoints_to_keep: 9
# Unsharded checkpoints (for final storage)
save_interval_unsharded: 100
save_num_unsharded_checkpoints_to_keep: -1

load_path: null

max_duration: 739_328  # 3.1T tokens
global_train_batch_size: 2048
device_train_microbatch_size: 2

precision: amp_bf16

fsdp:
  wrapping_strategy: null
  precision: mixed

max_grad_norm: 1.0
max_grad_norm_ratio: null

speed_monitor:
  window_size: 20

eval_interval: ${save_interval}
eval_subset_num_batches: -1
device_eval_batch_size: ${device_train_microbatch_size}

evaluators:
  # - label: all-small-ppl-validation
  #   data:
  #     num_workers: 0
  #     drop_last: true
  #     # generate_doc_lengths: true
  #     memmap_dtype: uint32
  #     datasets:
  #       c4_en-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/c4_en/val/part-0-00000.npy
  #       dolma_books-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/dolma_books/val/part-0-00000.npy
  #       dolma_common-crawl-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/dolma_common-crawl/val/part-0-00000.npy
  #       dolma_pes2o-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/dolma_pes2o/val/part-0-00000.npy
  #       dolma_reddit-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/dolma_reddit/val/part-0-00000.npy
  #       dolma_stack-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/dolma_stack/val/part-0-00000.npy
  #       dolma_wiki-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/dolma_wiki/val/part-0-00000.npy
  #       ice-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/ice/val/part-0-00000.npy
  #       m2d2_s2orc-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/m2d2_s2orc/val/part-0-00000.npy
  #       pile-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/pile/val/part-0-00000.npy
  #       wikitext_103-validation:
  #         - http://olmo-data.org/eval-data/perplexity/v3_small_dolma2-tokenizer/wikitext_103/val/part-0-00000.npy

  ##########################
  # Downstream evaluations #
  ##########################
  - label: piqa
    type: downstream

  - label: hellaswag
    type: downstream

  - label: winogrande
    type: downstream

  - label: openbook_qa
    type: downstream

  - label: boolq
    type: downstream

  - label: sciq
    type: downstream

  - label: arc_easy
    type: downstream

  - label: arc_challenge
    type: downstream

  - label: copa
    type: downstream

  - label: commonsense_qa
    type: downstream

  - label: social_iqa
    type: downstream

  - label: mmlu_stem_var
    type: downstream

  - label: mmlu_humanities_var
    type: downstream

  - label: mmlu_social_sciences_var
    type: downstream

  - label: mmlu_other_var
    type: downstream

  - label: mmlu_stem_mc_5shot
    type: downstream

  - label: mmlu_humanities_mc_5shot
    type: downstream

  - label: mmlu_social_sciences_mc_5shot
    type: downstream

  - label: mmlu_other_mc_5shot
    type: downstream

  - label: mmlu_stem_mc_5shot_test
    type: downstream

  - label: mmlu_humanities_mc_5shot_test
    type: downstream

  - label: mmlu_social_sciences_mc_5shot_test
    type: downstream

  - label: mmlu_other_mc_5shot_test
    type: downstream

  - label: basic_arithmetic
    type: downstream

  - label: trivia_qa_wiki_ppl
    type: downstream

  - label: natural_qs_open_ppl
    type: downstream

  - label: arc_easy_ppl
    type: downstream

data:
  pad_direction: right

  num_workers: 4
  drop_last: true
  pin_memory: true
  prefetch_factor: 8
  persistent_workers: true
  # be careful of the dtype, verify by using inspect_train_data.py
  memmap_dtype: uint32
  timeout: 0
  instance_filter:
    repetition_max_period: 13
    repetition_min_period: 1
    repetition_max_count: 32
  paths:
    # Wikipedia
    - http://olmo-data.org/preprocessed/olmo-mix/danyh-compiled-v1_7/documents/wiki/allenai/dolma2-tokenizer/part-0-00000.npy
    - http://olmo-data.org/preprocessed/olmo-mix/danyh-compiled-v1_7/documents/wiki/allenai/dolma2-tokenizer/part-1-00000.npy