16b_a3b.yaml

# Copyright 2026 FlagOS Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

system:
  no_shared_fs: ${experiment.runner.no_shared_fs}
  num_workers: 2
  tensor_model_parallel_size: 1
  pipeline_model_parallel_size: 2
  decoder_first_pipeline_num_layers: 13
  expert_model_parallel_size: 4
  context_parallel_size: 1
  sequence_parallel: true
  use_distributed_optimizer: true
  overlap_grad_reduce: true
  overlap_param_gather: true
  precision:
    bf16: true
    attention_softmax_in_fp32: true
    accumulate_allreduce_grads_in_fp32: true
  logging:
    log_interval: 1
    tensorboard_log_interval: 1
    wandb_project: ${experiment.exp_name}
    wandb_exp_name: ${experiment.exp_name}
    log_timers_to_tensorboard: true
    log_validation_ppl_to_tensorboard: true
    log_throughput: true
    log_params_norm: true
    log_num_zeros_in_grad: true
    log_memory_to_tensorboard: true
  checkpoint:
    save_interval: ${experiment.save_steps}
    load: ${experiment.load}
    ckpt_format: ${experiment.ckpt_format}

model:
  transformer_impl: transformer_engine
  num_layers: 27
  hidden_size: 2048
  num_attention_heads: 16
  num_query_groups: 16 # num_key_value_heads
  seq_length: 4096
  max_position_embeddings: 4096
  norm_epsilon: 1e-6
  use_rotary_position_embeddings: true
  rotary_base: 1000000
  swiglu: true
  normalization: RMSNorm
  qk_layernorm: true
  init_method_std: 0.02
  attention_dropout: 0.0
  hidden_dropout: 0.0
  position_embedding_type: rope
  untie_embeddings_and_output_weights: true
  no_position_embedding: true
  no_rope_fusion: true
  disable_bias_linear: true

  # mla args ==================
  multi_latent_attention: true
  kv_lora_rank: 512
  qk_head_dim: 128
  qk_pos_emb_head_dim: 64
  v_head_dim: 128

  # moe args ===================
  ffn_hidden_size: 11264
  moe_ffn_hidden_size: 1408
  moe_grouped_gemm: true
  moe_shared_expert_intermediate_size: 2816
  num_experts: 64
  moe_router_load_balancing_type: "seq_aux_loss"
  moe_router_score_function: sigmoid
  moe_router_enable_expert_bias: true
  moe_router_bias_update_rate: 0.001
  moe_aux_loss_coeff: 0.02
  moe_layer_freq: "[0]+[1]*26"
  # node limited routing
  moe_router_num_groups: 1
  moe_router_group_topk: 1
  moe_router_topk: 6
  moe_router_topk_scaling_factor: 2.446
  moe_token_dispatcher_type: "alltoall"

  # mtp args ====================
  mtp_num_layers: 1
  mtp_loss_scaling_factor: 0.3

  # training
  seed: ${experiment.seed}
  finetune: false
  micro_batch_size: 1
  global_batch_size: 128 #2048
  eval_iters: 0
  eval_interval: 1000
  train_iters: 102400

  optimizer:
    clip_grad: 1.0
    weight_decay: 0.1
    adam_beta1: 0.9
    adam_beta2: 0.95
    lr_scheduler:
      lr: 3.0e-3
      min_lr: 3.0e-4
      lr_warmup_fraction: 0.01
      lr_decay_style: WSD
      lr_wsd_decay_style: cosine
      lr_wsd_decay_iters: 10

data:
  reset_position_ids: True
  reset_attention_mask: True
  data_path: /path
  split: 1
  no_mmap_bin_files: true
  tokenizer:
    tokenizer_type: QwenTokenizerFS
    tokenizer_path: examples/aquila/qwentokenizer
    vocab_size: 151851
    make_vocab_size_divisible_by: 64