pg19_200bb_360m_3d_stt.yaml

io:
  name_model: 1m_200b_360m_3d_stt
  output_dir: /data/output < set this!
  dataset_dir: /data/datasets/pg19 < set this!
  dataset_id: pg19
  num_models_to_save: 5
  validate_amount: 100
  log_train_loss_amount: 1000
params:
  num_tokens: 256
  pad_token_id: 0
  input_seq_len: 1_048_576
  seq_lens: [8192, 16, 8]
  hidden_dims: [1024, 1024, 1024]
  num_layers: [9, 8, 8]
  train_checkpoint_chunks: [10, 20]
  block:
    - d_state: 128
      d_conv: 4
      expand: 2
      headdim: 64
      pos_emb_type: null
    - attn_head_dims: 64
      attn_num_heads: 16
      attn_use_rot_embs: true
      attn_dropout: 0
      use_flash_attn: true
      pos_emb_type: fixed
    - attn_head_dims: 64
      attn_num_heads: 16
      attn_use_rot_embs: true
      attn_dropout: 0
      use_flash_attn: true
      pos_emb_type: fixed
train:
  target_elements: 200_000_000_000
  target_elements_strategy: sequence
  batch_size: 1
  shuffle_train: false
  learning_rate: 0.001
  gradient_clipping: 1
  gradient_accumulate_every: 48