model_config.json

{
    "router_type": "top_k",
    "use_load_balancer": false,
    "num_epochs": 5,
    "batch_size": 16,
    "learning_rate": 0.0001,
    "debug_subset": 0,
    "output_dir": "moe_model_top_k",
    "push_to_hub": true,
    "hf_repo_id": "Sky6573/moe-topk",
    "use_wandb": true,
    "wandb_project": "moe-summarizer",
    "wandb_entity": null,
    "model": {
        "d_model": 256,
        "num_layers": 4,
        "num_heads": 8,
        "d_ffn": 1024,
        "num_experts": 8,
        "k": 2,
        "dropout": 0.1,
        "vocab_size": 30000
    }
}