{
  "vocab_size": 206464,
  "d_model": 256,
  "n_layers": 4,
  "max_seq_len": 512,
  "mha_layer_indices": [
    1,
    3
  ],
  "dense_ffn_layer_indices": [
    0
  ],
  "hse_ffn_layer_indices": [
    2
  ],
  "arg_d_inner": 256,
  "arg_d_state": 8,
  "arg_d_rank": 8,
  "arg_conv_kernel": 4,
  "arg_local_window": 32,
  "arg_local_n_heads": 4,
  "arg_local_n_kv_heads": 1,
  "arg_head_dim": 32,
  "arg_scalar_gate": true,
  "mha_n_heads": 4,
  "mha_n_kv_heads": 1,
  "mha_head_dim": 64,
  "rope_theta": 500000.0,
  "rope_scaling_type": null,
  "rope_scaling_factor": 1.0,
  "dense_d_ff": 512,
  "hse_n_domains": 2,
  "hse_n_experts_per_domain": 2,
  "hse_top_k": 1,
  "hse_d_ff_expert": 128,
  "hse_capacity_factor": 1.25,
  "hse_expert_dropout": 0.1,
  "hse_aux_loss_alpha": 0.01,
  "hse_gumbel_tau_init": 1.0,
  "hse_gumbel_tau_final": 0.1,
  "hse_gumbel_tau_warmup_tokens": 20000000000,
  "norm_type": "dgn",
  "dgn_n_groups": 4,
  "norm_eps": 1e-06,
  "init_std": 0.02,
  "init_std_residual_scale": true,
  "dropout": 0.0,
  "tie_word_embeddings": false,
  "use_cache": true,
  "pad_token_id": 0,
  "bos_token_id": 1,
  "eos_token_id": 2,
  "use_flash_attention": false,
  "use_triton_kernels": false,
  "use_torch_compile": false,
  "compile_mode": "max-autotune",
  "use_gradient_checkpointing": true,
  "gradient_checkpointing_ratio": 0.5
}