[MAST] Commit qwen3_30b_a3b_mast.yaml

.meta/mast/qwen3_30b_a3b_mast.yaml

@@ -0,0 +1,165 @@
+# Grouped Relative Policy Optimization (GRPO)
+# >>> ./.meta/mast/launch.sh .meta/mast/qwen3_30b_a3b_mast.yaml
+# WARNING: ========================================================================
+# WARNING: This setup is only tested for successfully kicking off the training loop.
+# WARNING: ========================================================================
+
+# Global configuration
+group_size: 4
+local_batch_size: 1 # per-device batch size
+max_req_tokens: 1024
+max_res_tokens: 1024
+model: /mnt/wsfuse/teamforge/hf/qwen3_30b_a3b_2
+off_by_n: 1 # Off by one by default
+launcher: mast
+
+# Main loop configuration
+rollout_threads: 32 # make this 4x the number of policy replicas seems to work well
+
+# Observability configuration
+metric_logging:
+  wandb:
+    entity: torchforge
+    project: grpo-training
+    group: grpo_exp_${oc.env:USER}
+    logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce
+  console:
+    logging_mode: global_reduce
+
+# Dataset configuration
+dataset:
+  path: /mnt/wsfuse/teamforge/hf/gsm8k
+  revision: "main"
+  data_split: "train"
+  streaming: true
+  model: ${model}
+
+# Policy configuration
+policy:
+  engine_args:  # https://docs.vllm.ai/en/v0.10.0/api/vllm/engine/arg_utils.html#vllm.engine.arg_utils.EngineArgs
+    model: /mnt/wsfuse/teamforge/hf/qwen3_30b_a3b_2
+    tensor_parallel_size: 4
+    pipeline_parallel_size: 1
+    enforce_eager: false
+  sampling_params:  # https://docs.vllm.ai/en/v0.10.0/api/vllm/sampling_params.html#vllm.sampling_params.SamplingParams
+    n: ${group_size}
+    max_tokens: ${max_res_tokens}
+    temperature: 1.0
+    top_p: 1.0
+
+# Trainer configuration
+trainer:
+  model:
+    name: qwen3
+    flavor: 30B-A3B
+    hf_assets_path: /mnt/wsfuse/teamforge/hf/qwen3_30b_a3b_2
+  optimizer:
+    name: AdamW
+    lr: 1e-5
+    eps: 1e-8
+  lr_scheduler:
+    warmup_steps: 1
+  training:
+    local_batch_size: ${local_batch_size}
+    seq_len: ${sum:${max_req_tokens},${max_res_tokens}}  # seq_len >= max_req_tokens + max_res_tokens
+    max_norm: 1.0
+    steps: 1000000
+    dtype: bfloat16
+    gc_freq: 1
+  compile:
+    enable: false
+  parallelism:
+    data_parallel_replicate_degree: 1
+    data_parallel_shard_degree: -1
+    tensor_parallel_degree: 1
+    pipeline_parallel_degree: 1
+    context_parallel_degree: 1
+    expert_parallel_degree: 1
+    expert_tensor_parallel_degree: 1
+    disable_loss_parallel: true
+  checkpoint:
+    enable: true
+    initial_load_path: /mnt/wsfuse/teamforge/hf/qwen3_30b_a3b_2
+    initial_load_in_hf: true
+    folder: ${checkpoint_folder}
+    last_save_in_hf: true
+    interval: 500
+    async_mode: "disabled"
+  activation_checkpoint:
+    mode: full
+  comm:
+    # 30B MoE model can require more time to load checkpoint than 1.7B's 1200s
+    init_timeout_seconds: 1800
+  dcp_path: ${checkpoint_folder}
+
+# Replay buffer configuration
+replay_buffer:
+  batch_size: ${local_batch_size}
+  max_policy_age: ${off_by_n}
+  dp_size: ${actors.trainer.procs}
+
+# Reference model configuration
+ref_model:
+  model:
+    name: qwen3
+    flavor: 30B-A3B
+    hf_assets_path: /mnt/wsfuse/teamforge/hf/qwen3_30b_a3b_2
+  training:
+    seq_len: ${trainer.training.seq_len}
+    dtype: bfloat16
+    gc_freq: 1
+  compile:
+    enable: false
+  parallelism:
+    data_parallel_replicate_degree: 1
+    data_parallel_shard_degree: -1
+    tensor_parallel_degree: 1
+    pipeline_parallel_degree: 1
+    context_parallel_degree: 1
+    expert_parallel_degree: 1
+  checkpoint:
+    enable: true
+    initial_load_path: /mnt/wsfuse/teamforge/hf/qwen3_30b_a3b_2
+    initial_load_in_hf: true
+  comm:
+    # Reference model also loads from checkpoint; so setting a higher timeout
+    init_timeout_seconds: 1800
+
+# All resource allocations
+services:
+  policy:
+    procs: ${policy.engine_args.tensor_parallel_size}
+    num_replicas: 1
+    hosts: 1
+    with_gpus: true
+    mesh_name: policy
+  ref_model:
+    procs: 4
+    num_replicas: 1
+    hosts: 1
+    with_gpus: true
+    mesh_name: ref_model
+  reward_actor:
+    procs: 1
+    num_replicas: 1
+    with_gpus: false
+    mesh_name: reward_actor
+
+actors:
+  dataset:
+    procs: 1
+    with_gpus: false
+    mesh_name: dataset
+  trainer:
+    procs: 8
+    hosts: 1
+    with_gpus: true
+    mesh_name: trainer
+  replay_buffer:
+    procs: 1
+    with_gpus: false
+    mesh_name: replay_buffer
+  compute_advantages:
+    procs: 1
+    with_gpus: false
+    mesh_name: compute_advantages