# NeMo/Megatron target config for the rebalanced Laoyao 2B MoE experiment.
model:
  architecture: decoder_only_moe
  library_target: nemo_megatron
  tokenizer:
    library: huggingface_tokenizers
    compatibility_target: GLM-5.2
    path: /mnt/beegfs/yi/laoyao_2b_moe/tokenizer/glm5.2/tokenizer.json
    vocab_size: 154820
    design_note: use ModelScope zai-org/GLM-5.2 tokenizer.json, sha256 19e773648cb4e65de8660ea6365e10acca112d42a854923df93db4a6f333a82d, for serving/SFT compatibility

  # Derived from laoyaomodel preset: 2b-moe-active0.9b-balanced.
  num_layers: 12
  hidden_size: 1536
  num_attention_heads: 24
  num_query_groups: 4
  ffn_hidden_size: 4608
  max_position_embeddings: 16384
  seq_length: 16384
  normalization: rmsnorm
  activation: swiglu
  position_embedding_type: rope
  share_embeddings_and_output_weights: false

  moe:
    enabled: true
    layer_indices: [2, 4, 6, 8, 10]
    num_moe_experts: 12
    moe_router_topk: 4
    moe_ffn_hidden_size: 6144
    moe_aux_loss_coeff: 0.02
    moe_z_loss_coeff: 0.001
    moe_token_dispatcher_type: alltoall
    moe_expert_capacity_factor: 1.25
    router_load_balancing: aux_loss_plus_bias_update
    router_bias_update_rate: 0.02
    router_bias_final_update_rate: 0.002
    router_bias_decay_steps: 2000

parallelism:
  tensor_model_parallel_size: 1
  pipeline_model_parallel_size: 1
  expert_model_parallel_size: 1
  context_parallel_size: 1
  data_parallel_size: auto

precision:
  params_dtype: bf16
  autocast_dtype: bf16
  attention_backend: flash
  fp8: false
