config.json

{
  "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
  },
  "sparsity_config": {
    "enable_weight_sparsity": true,
    "target_l0_fraction": 0.01,
    "enable_activation_sparsity": false,
    "activation_topk_fraction": 0.25
  },
  "training_config": {
    "total_tokens": 1500000000,
    "batch_size": 32,
    "dataset_name": "jacobcd52/simplestories-tokenized",
    "tokenizer_name": "SimpleStories/SimpleStories-1.25M"
  }
}