#!/bin/bash

# SLURM集群训练脚本

# 设置训练参数
lr=1e-5                        # 学习率
epochs=5                       # 训练轮数
batch_size=16                  # 总批次大小
weight_decay=1e-4              # 权重衰减
train_dataset_name="s1K_tokenized"  # 训练数据集名称
uid="$(date +%Y%m%d_%H%M%S)"   # 生成唯一的训练ID

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        --lr) lr="$2"; shift 2 ;;                                 # 设置学习率
        --epochs) epochs="$2"; shift 2 ;;                         # 设置训练轮数
        --batch_size) batch_size="$2"; shift 2 ;;                 # 设置批次大小
        --weight_decay) weight_decay="$2"; shift 2 ;;             # 设置权重衰减
        --train_dataset_name) train_dataset_name="$2"; shift 2 ;; # 设置数据集名称
        --uid) uid="$2"; shift 2 ;;                              # 设置训练ID
        *) echo "Unknown parameter: $1"; exit 1 ;;                # 未知参数处理
    esac
done

# 获取节点信息
node_array=$(scontrol show hostnames $SLURM_JOB_NODELIST)  # 获取所有节点主机名
nnodes=$(echo $node_array | wc -w)                         # 计算节点数量
head_node=($node_array)                                    # 获取头节点
head_node_ip=$(ssh $head_node hostname --ip-address)       # 获取头节点IP地址
export NCCL_DEBUG=INFO                                     # 设置NCCL调试级别
export NCCL_TIMEOUT=3600                                   # 设置NCCL超时时间


# 计算梯度累积步数
gpu_count=$(nvidia-smi -L | wc -l)                        # 获取每个节点的GPU数量
grad_acc=$((batch_size/(gpu_count * nnodes)))             # 计算需要的梯度累积步数

echo "Number of nodes: $nnodes"                           # 打印节点数量
echo "Number of GPUs per node: $gpu_count"                # 打印每个节点的GPU数量
echo "Head node IP: $head_node_ip"                        # 打印头节点IP

# 使用torchrun启动分布式训练
run_name="qwen_${train_dataset_name}_bs${batch_size}_lr${lr}_epoch${epochs}_wd${weight_decay}_${uid}"  # 生成运行名称
torchrun \
    --nnodes=$nnodes \                                    # 设置节点数量
    --nproc_per_node=$gpu_count \                        # 设置每个节点的进程数
    --rdzv_id=$SLURM_JOB_ID \                           # 设置集合点ID为SLURM作业ID
    --rdzv_backend=c10d \                                # 设置集合点后端
    --rdzv_endpoint=$head_node_ip:29500 \                # 设置集合点终端节点
    train/sft.py \
    --block_size=32768 \                                 # 序列最大长度
    --per_device_train_batch_size=1 \                    # 每个设备的训练批次大小
    --per_device_eval_batch_size=1 \                     # 每个设备的评估批次大小
    --gradient_accumulation_steps=$grad_acc \            # 梯度累积步数
    --num_train_epochs=${epochs} \                       # 训练轮数
    --train_file_path="simplescaling/${train_dataset_name}" \  # 训练数据路径
    --model_name="Qwen/Qwen2.5-32B-Instruct" \          # 模型名称
    --warmup_ratio=0.05 \                               # 预热比例
    --report_to="none" \                                # 不报告训练进度
    --fsdp="full_shard auto_wrap" \                     # FSDP配置：完全分片和自动包装
    --fsdp_config="train/fsdp_config_qwen_cpu.json" \   # FSDP具体配置文件
    --bf16=True \                                       # 使用BF16精度
    --eval_strategy="no" \                              # 不进行评估
    --logging_steps=1 \                                 # 每步都记录日志
    --save_strategy="no" \                              # 不保存中间检查点
    --lr_scheduler_type="cosine" \                      # 使用余弦学习率调度器
    --learning_rate=${lr} \                             # 学习率
    --weight_decay=${weight_decay} \                    # 权重衰减
    --adam_beta1=0.9 \                                 # Adam优化器参数beta1
    --adam_beta2=0.95 \                                # Adam优化器参数beta2
    --output_dir="ckpts/${run_name}" \                 # 输出目录
    --push_to_hub=false \                              # 不推送到Hub
    --save_only_model=True \                           # 仅保存模型权重
    --gradient_checkpointing=True \                    # 启用梯度检查点
    --accelerator_config='{"gradient_accumulation_kwargs": {"sync_each_batch": true}}'  # 加速器配置