#!/bin/bash
# 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="/capacity/userdata/qwen"  # 使用的基础模型
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 sft.py --block_size=18364  --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="../data/dataset/tokenizer"  --model_name=${base_model}  --warmup_ratio=0.05  --fsdp="full_shard auto_wrap"  --fsdp_config="fsdp_config_qwen.json"  --bf16=True   --eval_strategy="no"  --logging_steps=1  --save_strategy="no"  --save_only_model=True  --lr_scheduler_type="cosine"   --learning_rate=${lr}  --weight_decay=${weight_decay}  --adam_beta1=0.9  --output_dir="../ckpt/s1-${uid}"      --adam_beta2=0.95       --gradient_checkpointing True        --gradient_checkpointing_kwargs '{"use_reentrant": false}'

             
