# Reference Running: bash train/sft.sh
# {'train_runtime': 5268.8407, 'train_samples_per_second': 0.949, 'train_steps_per_second': 0.119, 'train_loss': 0.1172730620391667, 'epoch': 5.0}

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

# 基础模型和训练超参数设置
base_model="Qwen/Qwen2.5-32B-Instruct"  # 使用的基础模型
lr=1e-5                                  # 学习率
min_lr=0                                 # 最小学习率
epochs=5                                 # 训练轮数
weight_decay=1e-4                        # 权重衰减，与slurm训练使用相同的训练管道
micro_batch_size=1                       # 每个GPU的批次大小，如果有16个GPU则总批次大小为16
gradient_accumulation_steps=1            # 梯度累积步数，需要更多GPU内存
max_steps=-1                             # 最大训练步数，-1表示不限制
gpu_count=$(nvidia-smi -L | wc -l)       # 获取可用的GPU数量
push_to_hub=false                        # 是否推送到HuggingFace Hub

# 使用torchrun启动分布式训练
torchrun --nproc-per-node ${gpu_count} --master_port 12345 \
    train/sft.py \
    --block_size=32768 \                          # 序列最大长度
    --per_device_train_batch_size=${micro_batch_size} \  # 每个设备的训练批次大小
    --per_device_eval_batch_size=${micro_batch_size} \   # 每个设备的评估批次大小
    --gradient_accumulation_steps=${gradient_accumulation_steps} \  # 梯度累积步数
    --num_train_epochs=${epochs} \                # 训练轮数
    --train_file_path="simplescaling/s1K_tokenized" \  # 训练数据路径
    --model_name=${base_model} \                  # 模型名称
    --warmup_ratio=0.05 \                        # 预热比例
    --fsdp="full_shard auto_wrap" \              # FSDP配置：完全分片和自动包装
    --fsdp_config="train/fsdp_config_qwen.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/s1-${uid}" \            # 输出目录
    --push_to_hub=${push_to_hub} \              # 是否推送到Hub
    --save_only_model=True                       # 仅保存模型权重
    # --gradient_checkpointing=True \            # 启用梯度检查点以在8个H100 GPU上实现高效内存使用
    # --accelerator_config='{"gradient_accumulation_kwargs": {"sync_each_batch": true}}'  # 加速器配置

