{
  "model_config": {
    "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_config": {
    "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
  },
  "training_config": {
    "total_tokens": 1000000000,
    "batch_size": 32,
    "dataset_name": "jacobcd52/simplestories-tokenized",
    "tokenizer_name": "SimpleStories/SimpleStories-1.25M"
  }
}