Validation Loss Stagnate on BLT Arch

[CharlesDDDD]

Hi Team,

Thank you for your great work! I notice a wierd training phenomenon when training 0.8B BLT model on fineweb-edu-10bt, that training loss would go down smoothly (even around 0.2) but val loss (bpb) stagnate at a around ~1.4-1.8 for all patching mode, I wonder if I mistake anything or any specific hyperparams I should be aware to set.

# Template config, need to change dump_dir, data.root_dir and tokenizer.path
# Evals can be activated by uncommenting its config
# python -m launchers.stool config=apps/main/configs/debug.yaml nodes=8 account=fair_amaia_cw_codegen qos=lowest

dump_dir: xxx
name: "debug"
steps: 250_000
probe_freq: null
seed: 777
optim:
  lr: 4e-04
  warmup: 10000
  lr_min_ratio: 0.1
  clip: 1.0
distributed:
  fsdp_type: full_shard
  model_dtype: bf16
  matmul_allow_tf32: false
  selective_activation_checkpointing: false
  tp_size: 1

model:
  n_heads: 8
  dim: 256
  n_heads_global: 10
  n_heads_local_decoder: 8
  n_heads_local_encoder: 8
  dim_global: 1280
  dim_local_decoder: 512
  dim_local_encoder: 512
  vocab_size: 260
  dim_token: 512
  patch_size: 6
  patching_mode: "bpe"
  tie_local_encoder_decoder_logits: false
  patch_in_forward: false
  max_encoder_seq_length: 12288
  pad_to_max_length: true
  patching_threshold: 0.5439168453216553
  encoder_hash_byte_group_size: [3,4,5,6,7,8]
  encoder_hash_byte_group_vocab: 100002
  encoder_hash_byte_group_nb_functions: 1
  encoder_enable_byte_ngrams: false
  cross_attn_encoder: true # assuming cross_attention is true
  cross_attn_decoder: true # assuming cross_attention is true
  cross_attn_window_encoder: 512
  cross_attn_window_decoder: 512
  cross_attn_k: 8
  cross_attn_nheads: 4
  cross_attn_all_layers_decoder: true
  cross_attn_all_layers_encoder: true
  cross_attn_use_flex_attention: true
  cross_attn_init_by_pooling: true
  log_patch_lengths: true
  non_linearity: "swiglu"
  use_rope: true
  recompute_fc1_out: false
  recompute_fc3_out: false
  recompute_attn: false
  custom_bwd: false
  layer_ckpt: "none"
  use_local_encoder_transformer: true
  init_use_gaussian: true
  init_use_depth: "current"
  attn_impl: "xformers"
  attn_bias_type: "block_causal"
  alpha_depth: "disabled"
  max_length: 256
  local_attention_window_len: 512
  max_seqlen: 12288
  downsampling_by_pooling: "max"
  n_layers_global: 24
  n_layers_local_encoder: 1
  n_layers_local_decoder: 7

data:
  root_dir: xxx
  sources:
    fineweb_edu_10bt_shuffled: 1.0
  batch_size: 4
  prefetch_size: 64
  seq_len: 4096
  load_async: true
  preprocess_dir: xxx
  add_patches: true
  patcher_args:
    patching_mode: bpe
    patching_device: cuda
    entropy_model_checkpoint_dir: null
    realtime_patching: false
    threshold: 0.5439168453216553
    threshold_add: null
    max_patch_length: null
    patch_size: 6
    patching_batch_size: 1
    device: cuda
    monotonicity: true
    log_time: false
  tokenizer_args:
    name: blt
    init_kwargs:
      bpe_tokenizer_path: xxx
      bpe_delim: true

profiling:
  run: false

checkpoint:
  dump:
    every: 500
    keep: 1
  eval:
    every: 1000
    keep: -1

logging:
  freq: 10

eval_on_gpus: 8
eval:
  run_ppl: true
  run_tasks: false
  consolidate_if_needed: true
  validation:
    root_dir: xxx
    sources:
      - fineweb_edu_10bt.val.arrow
    batch_size: 4
    max_n_docs: null
    max_n_batches: null

Thank you in advance!