byte_level_ffn_sharing.yaml

model:
  core_model:
    core_model_type: generic_ffn_sharing
    num_layers: 16
    ffn:
      ffn_type: swiglu
      ffn_dim: 1536
      normalization: rms_norm
      bias: false
    attn:
      attn_type: generic
      num_heads: 16
      normalization: rms_norm
      group_size: 4
      bias: false
      is_causal: true
  embedder:
    tokenizer_type: "gpt2"
    byte_tokenizer_type: "bpe"
    embedding_model_type: byte_level
    byte_context_window: 12
    dataset_name: simple_en_wiki
  lm_head:
    normalization: rms_norm
    bias: false
    lm_head_type: byte_level

  hidden_dim: 512
  context_window: 512
  vocab_size: 50257
  byte_vocab_size: 258
  byte_context_window: 12
  byte_embedding_dim: 128
  model_shell_type: standard
  embedding_weight_tying: false
  positional_encoding_type: rope
trainer:
  dropout_scheduler:
    dropout_type: linear
    start_dropout_p: 0.0
    end_dropout_p: 0.1
    start_iter: 0
    end_iter: 10000
  dataset: simple_en_wiki
  training:
    trainer_type: base_trainer
    batch_size: 24
    gradient_accumulation_steps: 20
    max_iters: 25000
    lr_decay_iters: 25000
    warmup_iters: 5000
    eval_interval: 5000
    log_interval: 100
    eval_iters: 500
    checkpoint_interval: 1000000000.0
    run_profiler: false
  optimizer:
    name: nanoGPTadamW
    lr: 0.0006
    min_lr: 6.0e-05
    weight_decay: 0.1
    beta1: 0.9
    beta2: 0.95
    grad_clip: 1.0
    decay_lr: true
    warmup_iters: 5000
  lr_scheduler:
    name: cosine
  dataloader:
    name: byte_pooling
  loss_fn:
    name: cross_entropy
  eval:
    benchmarks:
      - "winograd"
      - "hellaswag"
      - "arc"
      - "mmlu"
      - "blimp"
    num_samples: 5000
    evaluator: "mcq"

general:
  logging:
    wandb_log: true
    wandb_project: SuperTinyLanguageModels
  paths:
    output_dir: outputs
    data_dir: data
    checkpoint_dir: checkpoints
  seed: 489
  device: cuda