# 多节点分布式训练脚本

# 生成唯一的训练ID，基于时间戳
uid="$(date +%Y%m%d_%H%M%S)"

# 基础模型和训练超参数设置
base_model="Qwen/Qwen2.5-32B-Instruct" # 也可以使用 meta-llama/Llama-3.1-70B-Instruct
lr=1e-5                                # 学习率
min_lr=0                               # 最小学习率
epochs=5                               # 训练轮数
micro_batch_size=1                     # 如果有2个节点每个8个GPU，总批次大小将是16
push_to_hub=true                       # 是否推送到HuggingFace Hub
gradient_accumulation_steps=1          # 梯度累积步数
max_steps=-1                           # 最大训练步数，-1表示不限制
gpu_count=$(nvidia-smi -L | wc -l)     # 获取每个节点可用的GPU数量

# 使用torchrun启动多节点分布式训练
torchrun \
--nnodes ${NUM_NODES}:${NUM_NODES} \        # 总节点数和最小节点数
--node_rank=$REPLICA_RANK \                  # 当前节点的排名
--nproc-per-node ${gpu_count} \             # 每个节点的进程数（GPU数）
--rdzv_id=12347 \                           # 集合点ID
--rdzv_backend=c10d \                       # 集合点后端
--rdzv_conf='read_timeout=420' \            # 集合点配置：读取超时时间
--rdzv_endpoint=$REPLICA_HOSTNAME:29401 \    # 集合点终端节点
train/sft.py \
--per_device_train_batch_size=${micro_batch_size} \    # 每个设备的训练批次大小
--per_device_eval_batch_size=${micro_batch_size} \     # 每个设备的评估批次大小
--gradient_accumulation_steps=${gradient_accumulation_steps} \  # 梯度累积步数
--train_file_path="simplescaling/s1K_tokenized" \      # 训练数据路径
--block_size=32768 \                                   # 序列最大长度
--model_name=${base_model} \                           # 模型名称
--warmup_ratio=0.05 \                                  # 预热比例
--fsdp="full_shard auto_wrap" \                        # FSDP配置：完全分片和自动包装
--fsdp_config="train/scripts/fsdp_config_qwen.json" \  # FSDP具体配置文件
--bf16=True \                                          # 使用BF16精度
--eval_strategy="steps" \                              # 按步数进行评估
--eval_steps=50 \                                      # 每50步评估一次
--logging_steps=1 \                                    # 每步都记录日志
--save_steps=100 \                                     # 每100步保存一次模型
--lr_scheduler_type cosine \                           # 使用余弦学习率调度器
--learning_rate ${lr} \                                # 学习率
--weight_decay 1e-4 \                                  # 权重衰减
--adam_beta1 0.9 \                                     # Adam优化器参数beta1
--adam_beta2 0.95 \                                    # Adam优化器参数beta2
--output_dir="ckpts/s1_${uid}" \                      # 输出目录
--hub_model_id="simplescaling/s1-${uid}" \            # Hub上的模型ID
--push_to_hub=True \                                   # 推送到Hub
--hub_always_push=True \                               # 始终推送到Hub
--num_train_epochs ${epochs} \                         # 训练轮数
--save_only_model=True \                               # 仅保存模型权重
--gradient_checkpointing=True                          # 启用梯度检查点
