fix moe freq config (#12)

Author: wenxie-amd

examples/deepseek/exp_pretrain.yaml

@@ -17,20 +17,33 @@ modules:
       # log
       wandb_project: "Primus_DeepSeek_Pretrain"
       # disable_wandb: false
+      # disable_tensorboard: false
       stderr_sink_level: DEBUG
 
       # debug
-      num_layers: 5
-      optimizer: adam
+      # num_layers: 5
+      # optimizer: adam
       moe_router_force_load_balancing: true
       moe_router_dtype: fp32
       log_avg_skip_iterations: 2
       log_avg_reset_interval: 5
+      # ddp_bucket_size: 629145600
+
+      # recompute
+      # recompute_granularity: full # full, selective
+      # recompute_method: block # uniform, block
+      # recompute_num_layers: 1 # int
+
+      # profile
+      # profile: true
+      # use_pytorch_profiler: true
+      # profile_step_end: 7
+      # profile_step_start: 6
 
       # hyber parameters
       train_iters: 10
       micro_batch_size: 1
-      global_batch_size: 16
+      global_batch_size: 64
       seq_length: 4096
       max_position_embeddings: 4096
       lr: 1.0e-5

examples/deepseek/run_pretrain.sh

@@ -42,8 +42,10 @@ export NCCL_IB_HCA=rdma0:1,rdma1:1,rdma2:1,rdma3:1,rdma4:1,rdma5:1,rdma6:1,rdma7
 export NCCL_IB_GID_INDEX=3
 export NCCL_CROSS_NIC=0
 export HSA_ENABLE_SDMA=0
-export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-ens51f0}
-export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-ens51f0}
+IP_INTERFACE=$(ip -o -4 addr show | awk -v ip="$(hostname -I | awk '{print $1}')" '$4 ~ ip {print $2}')
+export IP_INTERFACE
+export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-${IP_INTERFACE}}
+export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-${IP_INTERFACE}}
 export CUDA_DEVICE_MAX_CONNECTIONS=1 # Reducing to 1 ensures no PCIE traffic (even on single node)
 export NCCL_PROTO=Simple
 export RCCL_MSCCL_ENABLE=0
@@ -108,6 +110,7 @@ export HIP_VISIBLE_DEVICES=$gpus
 
 echo "[NODE-$NODE_RANK] MASTER_ADDR: $MASTER_ADDR"
 echo "[NODE-$NODE_RANK] MASTER_PORT: $MASTER_PORT"
+echo "[NODE-$NODE_RANK] IP_INTERFACE: $IP_INTERFACE"
 echo "[NODE-$NODE_RANK] NNODES: $NNODES"
 echo "[NODE-$NODE_RANK] NODE_RANK: $NODE_RANK"
 echo "[NODE-$NODE_RANK] GPUS_PER_NODE: $GPUS_PER_NODE"

examples/deepseek/run_slurm_pretrain.sh

@@ -8,8 +8,7 @@ SCRIPT_DIR=$(dirname "$(realpath "${BASH_SOURCE[0]}")")
 echo "Current script path: $SCRIPT_DIR"
 
 export RUN_ENV=slurm
-export NCCL_SOCKET_IFNAME=bond0
-export GLOO_SOCKET_IFNAME=bond0
+export MODEL_CONFIG=deepseek_v2_lite
 
 srun -N 2 \
      --gres=gpu:8 \

primus/configs/models/megatron/deepseek_v3.yaml

@@ -24,7 +24,7 @@ kv_channels: 128
 mtp_num_layers: 1
 mtp_loss_scaling_factor: 0.1
 # moe
-moe_layer_freq: 3
+moe_layer_freq: "([0]*3+[1]*58)"
 num_experts: 256
 moe_router_topk: 8
 # num_shared_experts: 1

primus/configs/models/megatron/deepseek_v3_17B.yaml

@@ -20,7 +20,7 @@ qk_pos_emb_head_dim: 0
 v_head_dim: 128
 kv_channels: 128
 # moe
-moe_layer_freq: 2
+moe_layer_freq: "([0]*2+[1]*26)"
 num_experts: 96
 moe_router_topk: 6
 # num_shared_experts: 1

primus/configs/models/megatron/deepseek_v3_393B.yaml

@@ -20,7 +20,7 @@ qk_pos_emb_head_dim: 0
 v_head_dim: 128
 kv_channels: 128
 # moe
-moe_layer_freq: 2
+moe_layer_freq: "([0]*2+[1]*62)"
 num_experts: 256
 moe_router_topk: 8
 # num_shared_experts: 1

primus/configs/models/megatron/deepseek_v3_45B.yaml

@@ -20,7 +20,7 @@ qk_pos_emb_head_dim: 64
 v_head_dim: 128
 kv_channels: 128
 # moe
-moe_layer_freq: 2
+moe_layer_freq: "([0]*2+[1]*30)"
 num_experts: 96
 moe_router_topk: 4
 # num_shared_experts: 1

primus/configs/modules/megatron/trainer_base.yaml

@@ -184,9 +184,9 @@ adlr_autoresume_interval: 1000
 
 # activation recomputation
 recompute_activations: false
-recompute_granularity: null
-recompute_method: null
-recompute_num_layers: null
+recompute_granularity: null # full, selective
+recompute_method: null # uniform, block
+recompute_num_layers: null # int
 distribute_saved_activations: false
 checkpoint_activations: false # deprecated
 
@@ -224,6 +224,7 @@ mmap_bin_files: true
 
 #profile:
 profile: false
+use_pytorch_profiler: false
 profile_ranks: [0]
 profile_step_end: 12
 profile_step_start: 10
@@ -329,7 +330,7 @@ straggler_minmax_count: 1
 inference_batch_times_seqlen_threshold: -1
 inference_dynamic_batching: false
 inference_dynamic_batching_buffer_size_gb: 40.0 # float
-inference_dynamic_batching_buffer_guaranteed_fraction: 0.2 # float   
+inference_dynamic_batching_buffer_guaranteed_fraction: 0.2 # float
 inference_dynamic_batching_buffer_overflow_factor: null # float
 inference_dynamic_batching_max_requests_override: null # int
 inference_dynamic_batching_max_tokens_override: null # int

primus/modules/trainer/megatron/trainer.py

@@ -77,7 +77,11 @@
     global_vars,
     one_logger_utils,
 )
-from megatron.training.arguments import core_transformer_config_from_args, validate_args
+from megatron.training.arguments import (
+    core_transformer_config_from_args,
+    moe_freq_type,
+    validate_args,
+)
 from megatron.training.async_utils import (
     init_persistent_async_worker,
     maybe_finalize_async_save,
@@ -404,6 +408,9 @@ def update_primus_config(
         if args.iterations_to_skip is None:
             args.iterations_to_skip = []
 
+        # support moe_freq_type
+        args.moe_layer_freq = moe_freq_type(args.moe_layer_freq)
+
     def vocab_size_with_padding(self, orig_vocab_size, args):
         """Pad vocab size so it is divisible by model parallel size and
         still having GPU friendly size."""