seed: 1234
output_dir: './output'

# checkpoints setting
load_checkpoint: '' # 'path/to/unified_safe'
load_ckpt_format: 'safetensors'
auto_trans_ckpt: False  # If true, auto transform load_checkpoint to load in distributed model
only_save_strategy: False
resume_training: False

run_mode: 'train'
use_parallel: True
use_legacy: False
print_separate_loss: True

# trainer config
trainer:
  type: CausalLanguageModelingTrainer
  model_name: 'qwen3_moe'

# runner config
runner_config:
  epochs: 1
  batch_size: 1
  sink_mode: True
  sink_size: 1
  gradient_accumulation_steps: 1

# optimizer
optimizer:
  type: AdamW
  betas: [0.9, 0.95]
  eps: 1.e-8
  weight_decay: 0.0

lr_schedule:
  type: ConstantWarmUpLR
  learning_rate: 1.e-5
  warmup_steps: 0
  total_steps: -1 # -1 means it will load the total steps of the dataset

# dataset
train_dataset: &train_dataset
  data_loader:
    type: BlendedMegatronDatasetDataLoader
    datasets_type: "GPTDataset"
    sizes:
      - 8000 # 训练集数据样本数
      - 0    # 测试集数据样本数，当前不支持配置
      - 0    # 评测集数据样本数，当前不支持配置
    config:  # GPTDataset配置项
      seed: 1234                         # 数据采样随机种子
      split: "1, 0, 0"                   # 训练、测试、评测集使用比例，当前不支持配置
      seq_length: 4096                   # 数据集返回数据的序列长度
      eod_mask_loss: False               # 是否在eod处计算loss
      reset_position_ids: False          # 是否在eod处重置position_ids
      create_attention_mask: False       # 是否返回attention_mask
      reset_attention_mask: False        # 是否在eod处重置attention_mask，返回阶梯状attention_mask
      create_compressed_eod_mask: False  # 是否返回压缩后的attention_mask
      eod_pad_length: 128                # 设置压缩后attention_mask的长度
      eod: 1                             # 数据集中eod的token id
      pad: -1                            # 数据集中pad的token id

      data_path:  # Megatron数据集采样比例以及路径
        - '1'
        - "/path/to/wiki103-megatron_text_document"

  input_columns: ["input_ids", "labels", "loss_mask", "position_ids"]
  construct_args_key: ["input_ids", "labels", "loss_mask", "position_ids"]

  num_parallel_workers: 8
  python_multiprocessing: False
  drop_remainder: True
  numa_enable: False
  prefetch_size: 1
  seed: 1234

train_dataset_task:
  type: CausalLanguageModelDataset
  dataset_config: *train_dataset

# mindspore context init config
context:
  mode: 0 # 0--Graph Mode; 1--Pynative Mode
  device_target: "Ascend"
  max_call_depth: 10000
  max_device_memory: "58GB"
  save_graphs: False
  save_graphs_path: "./graph"
  device_id: 0
  memory_optimize_level: "O0"
  jit_config:
    jit_level: "O0"
  ascend_config:
    precision_mode: "must_keep_origin_dtype"
    parallel_speed_up_json_path: "configs/qwen3_moe/parallel_speed_up.json"

parallel_config:
  data_parallel: &dp 1
  model_parallel: 4
  pipeline_stage: 4
  micro_batch_num: &micro_batch_num 16
  expert_parallel: 4
  use_seq_parallel: True
  gradient_aggregation_group: 1
# when model parallel is greater than 1, we can set micro_batch_interleave_num=2, that may accelerate the train process.
micro_batch_interleave_num: 1

# parallel context config
parallel:
  parallel_mode: 1 # 0-data parallel, 1-semi-auto parallel, 2-auto parallel, 3-hybrid parallel
  gradients_mean: False
  enable_alltoall: True
  full_batch: False
  dataset_strategy: [
    [*dp, 1],
    [*dp, 1],
    [*dp, 1],
    [*dp, 1]
  ]
  search_mode: "sharding_propagation"
  enable_parallel_optimizer: False
  strategy_ckpt_config:
    save_file: "./ckpt_strategy.ckpt"
    only_trainable_params: False
  parallel_optimizer_config:
    gradient_accumulation_shard: False
    parallel_optimizer_threshold: 8

# recompute config
recompute_config:
  recompute: True
  select_recompute: False
  parallel_optimizer_comm_recompute: True
  mp_comm_recompute: True

# model config
model:
  model_config:
    model_type: "qwen3_moe"
    architectures: "Qwen3MoeForCausalLM"

    # model args
    offset: 0
    vocab_size: 151936
    hidden_size: 2048
    intermediate_size: 6144
    num_hidden_layers: 48
    max_position_embeddings: 40960
    hidden_act: 'silu'
    num_attention_heads: 32
    num_key_value_heads: 4
    head_dim: 128
    rms_norm_eps: 1.e-6
    add_bias_linear: False
    tie_word_embeddings: False
    use_flash_attention: True
    use_contiguous_weight_layout_attention: False
    input_sliced_sig: True

    # dropout args
    hidden_dropout: 0.0
    attention_dropout: 0.0

    # rope args
    position_embedding_type: "rope"
    rope_theta: 1000000

    # dtype args
    params_dtype: "float32"
    compute_dtype: "bfloat16"
    layernorm_compute_dtype: "float32"
    softmax_compute_dtype: "float32"
    rotary_dtype: "float32"
    initializer_range: 0.02

    # moe args
    moe_grouped_gemm: True
    num_experts: 128
    num_experts_per_tok: 8
    moe_intermediate_size: 768

    router_dense_type: "float32"
    gated_linear_unit: True
    norm_topk_prob: False
    moe_router_pre_softmax: True
    moe_token_drop_policy: probs
    routed_scaling_factor: 1.5

    router_aux_loss_coef: 0.001
    moe_aux_loss_coeff: 0.001
    moe_router_load_balancing_type: "gbs_aux_loss"

# callbacks
callbacks:
  - type: MFLossMonitor

# wrapper cell config
runner_wrapper:
  type: MFTrainOneStepCell
  scale_sense: 1.0
  use_clip_grad: True

profile: False
profile_start_step: 1
profile_stop_step: 10
init_start_profile: False
profile_communication: False
profile_memory: True