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.0
  use_bias: true
  use_flash_attention: true
sparsity:
  enable_weight_sparsity: true
  target_l0_fraction: 0.05
  sparsity_anneal_start_fraction: 0.01
  sparsity_anneal_end_fraction: 0.5
  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.5
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.0
  warmup_fraction: 0.01
  enable_lr_decay: true
  min_lr_fraction: 0.0
  lr_decay_end_fraction: 1.0
  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/pgd_d3072
  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: ci_thresholding_l0_0.05
  wandb_entity: null
  use_wandb: true
  seed: 0
  hf_repo: asher577/ci_thresholding
  num_workers: 8
  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
  ablation_mode: pgd
  stochastic_n_bins: 10
  pgd_init: random
  pgd_lr: 0.2
  pgd_n_steps: 2
  pgd_use_sign_grad: true
  pgd_outer_lr: 0.0
  minimality_coeff: 1.0e-05
  minimality_coeff_final_frac: 1.0
  minimality_coeff_annealing: true
  minimality_coeff_anneal_start_frac: 0.01
  minimality_coeff_anneal_end_frac: 0.7
  minimality_coeff_anneal_type: linear
  minimality_pnorm: 2.0
  minimality_eps: 5.0e-07
  minimality_beta: 0.5
  minimality_p_anneal_start_frac: 0.01
  minimality_p_anneal_final_p: 0.5
  minimality_p_anneal_end_frac: 0.5
  initial_ci: 1.0
  minimality_warmup_frac: 0.02
  ci_thresholding: 0.1
  ci_learning_rate: null