#!/bin/bash
export MIOPEN_FIND_MODE=3
export GPU_MAX_HW_QUEUES=16
lrank=$OMPI_COMM_WORLD_LOCAL_RANK
comm_rank=$OMPI_COMM_WORLD_RANK
comm_size=$OMPI_COMM_WORLD_SIZE
export LOCAL_RANK=$OMPI_COMM_WORLD_LOCAL_RANK
export RANK=$comm_rank
export WORLD_SIZE=$comm_size
export MASTER_ADDR=$1
export MASTER_PORT=29500

# export NCCL_DEBUG=info   #打印nccl通信的日志
export NCCL_IB_HCA=mlx5
export NCCL_SOCKET_IFNAME=ib0 
export HIP_DIRECT_DISPATCH=0

    #   "offload_optimizer": {
    #     "device": "cpu",
    #     "pin_memory": true
    #   },
    #   "offload_param": {
    #     "device": "cpu",
    #     "pin_memory": true
    #   }

APP="python ./src/train_bash.py \
    --deepspeed deepspeed.json  --model_name_or_path /work/home/hepj/model/Qwen-7B-Chat \
    --do_train --dataset alpaca_gpt4_zh --template chatml --finetuning_type lora --lora_target c_attn  \
    --output_dir ./output/ft_qwen \
    --per_device_train_batch_size 1 --gradient_accumulation_steps 1  --lr_scheduler_type cosine \
    --logging_steps 10 --save_steps 1000 --learning_rate 5e-5 --num_train_epochs 3.0 --fp16"
#--overwrite_cache

case ${lrank} in
[0])
  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
  export UCX_NET_DEVICES=mlx5_0:1
  export UCX_IB_PCI_BW=mlx5_0:50Gbs
  NCCL_SOCKET_IFNAME=ib0 numactl --cpunodebind=0 --membind=0 ${APP}
  ;;
[1])
  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
  export UCX_NET_DEVICES=mlx5_1:1
  export UCX_IB_PCI_BW=mlx5_1:50Gbs
  NCCL_SOCKET_IFNAME=ib0 numactl --cpunodebind=0 --membind=0 ${APP}
  ;;
[2])
  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
  export UCX_NET_DEVICES=mlx5_2:1
  export UCX_IB_PCI_BW=mlx5_2:50Gbs
  NCCL_SOCKET_IFNAME=ib0 numactl --cpunodebind=0 --membind=0 ${APP}
  ;;
[3])
  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
  export UCX_NET_DEVICES=mlx5_3:1
  export UCX_IB_PCI_BW=mlx5_3:50Gbs
  NCCL_SOCKET_IFNAME=ib0 numactl --cpunodebind=0 --membind=0 ${APP}
  ;;
[4])
  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
  export UCX_NET_DEVICES=mlx5_4:1
  export UCX_IB_PCI_BW=mlx5_4:50Gbs
  NCCL_SOCKET_IFNAME=ib0 numactl --cpunodebind=3 --membind=3 ${APP}
  ;;
[5])
  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
  export UCX_NET_DEVICES=mlx5_5:1
  export UCX_IB_PCI_BW=mlx5_5:50Gbs
  NCCL_SOCKET_IFNAME=ib0 numactl --cpunodebind=3 --membind=3 ${APP}
  ;;
[6])
  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
  export UCX_NET_DEVICES=mlx5_6:1
  export UCX_IB_PCI_BW=mlx5_6:50Gbs
  NCCL_SOCKET_IFNAME=ib0 numactl --cpunodebind=3 --membind=3 ${APP}
  ;;
[7])
  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
  export UCX_NET_DEVICES=mlx5_7:1
  export UCX_IB_PCI_BW=mlx5_7:50Gbs
  NCCL_SOCKET_IFNAME=ib0 numactl --cpunodebind=3 --membind=3 ${APP}
  ;;  
esac