training_config.yaml

model:
  n_layer: 2
  d_model: 3072
  n_ctx: 512
  d_head: 16
  d_mlp: 12288
  vocab_size: 4096
  use_rms_norm: true
  tie_embeddings: false
  use_positional_embeddings: false
  use_bigram_table: false
  use_attention_sinks: true
  activation: gelu
  dropout: 0
  use_bias: true
  use_flash_attention: true
  gradient_checkpointing: false
sparsity:
  enable_weight_sparsity: true
  target_l0_fraction: 0.01
  sparsity_anneal_start_fraction: 0.5
  sparsity_anneal_end_fraction: 0.75
  anneal_type: linear
  min_weights_per_neuron: 4
  enable_activation_sparsity: true
  activation_topk_fraction: 0.2
  activation_sparsity_locations: attn_q,attn_k,attn_v,attn_out,mlp_neuron,mlp_out
  activation_topk_mode: per_token
  mlp_neuron_pre_activation: true
  activation_topk_annealing: linear
  activation_topk_anneal_start_fraction: 0.01
  activation_topk_anneal_end_fraction: 0.33
optimizer:
  optimizer_type: adamw
  learning_rate: 0.001
  beta1: 0.9
  beta2: 0.95
  weight_decay: 0.001
  eps: 0.1
  enable_grad_clip: true
  grad_clip_rms: 1
  warmup_fraction: 0.01
  enable_lr_decay: true
  min_lr_fraction: 0
  lr_decay_start_fraction: 0.8
  lr_decay_end_fraction: 1
  use_sharkfin_schedule: true
training:
  dataset_name: jacobcd52/simplestories-tokenized
  dataset_split: train
  text_column: story
  tokenizer_name: SimpleStories/SimpleStories-1.25M
  total_tokens: 1000000000
  batch_size: 32
  gradient_accumulation_steps: 1
  mixed_precision: bf16
  checkpoint_dir: checkpoints/ablation_pretraining/separate_normalization
  checkpoint_every_n_steps: 0
  keep_n_checkpoints: 0
  log_every_n_steps: 10
  log_gradients_every_n_steps: 100
  log_weights_every_n_steps: 100
  log_sparsity_every_n_steps: 100
  eval_every_n_steps: 100
  val_split: test
  val_holdout_fraction: 0.01
  val_max_batches: 20
  wandb_project: nodesparse_adv_pretraining
  wandb_run_name: separate_normalization
  wandb_entity: null
  use_wandb: true
  seed: 0
  hf_repo: asher577/separate_normalization
  num_workers: 2
  resume_from_checkpoint: null
  use_torch_compile: false
  torch_compile_mode: default
  torch_compile_backend: inductor
ablation:
  ci_fn_type: global_reverse_residual
  d_resid_ci_fn: 1024
  reader_hidden_dims: 1024,1024
  use_attn_transitions: false
  attn_n_heads: 4
  attn_max_len: 2048
  ci_targets: attn_q,attn_k,attn_v,attn_out,mlp_pre,mlp_out
  pgd_loss_frac: 0.5
  stochastic_loss_frac: 0.16
  ci_loss_frac: 0.17
  unmasked_loss_frac: 0.17
  adv_kl_coeff: 0.5
  pgd_init: random
  pgd_lr: 0.4
  pgd_n_steps: 1
  pgd_use_sign_grad: true
  pgd_outer_lr: 0
  pgd_mode: batch
  separate_grad_normalize: true
  separate_grad_min_frac: 0.25
  minimality_coeff: 6.0e-08
  minimality_coeff_final_frac: 1
  minimality_coeff_annealing: true
  minimality_coeff_anneal_start_frac: 0.25
  minimality_coeff_anneal_end_frac: 0.5
  minimality_coeff_anneal_type: exp
  minimality_pnorm: 2
  minimality_eps: 1.0e-12
  minimality_beta: 325
  minimality_p_anneal_start_frac: 0.02
  minimality_p_anneal_final_p: 0.4
  minimality_p_anneal_end_frac: 0.95
  minimality_beta_pnorm: 2
  minimality_beta_p_anneal_start_frac: 0.02
  minimality_beta_p_anneal_final_p: 0.4
  minimality_beta_p_anneal_end_frac: 0.95
  initial_ci: 1
  minimality_warmup_frac: 0.02
  ci_thresholding: 0
  binarization_coeff: 0
  binarization_anneal_start_frac: 0.8
  binarization_anneal_end_frac: 1
  ci_learning_rate: 0.001