#!/bin/bash

# 获取系统中可用的GPU数量
# nvidia-smi --query-gpu=gpu_name：查询所有GPU的名称
# format=csv,noheader：以CSV格式输出，不显示表头
# wc -l：统计行数，即GPU数量
NUM_GPUS=$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l)

# 使用torchrun启动分布式训练
# --nproc_per_node：每个节点(机器)上启动的进程数，这里设置为GPU数量
# 每个进程会被分配到一个GPU上
torchrun --nproc_per_node=$NUM_GPUS \
    train.py \
    --model_name_or_path "/root/models/DeepSeek-R1-Distill-Qwen-1.5B" \
    --train_file "./data/train.jsonl" \
    --eval_file "./data/eval.jsonl" \
    --epochs 3 \
    --train_batch_size 4 \
    --grad_accum_steps 2 \
    --model_dir "./output_model" 

# 重要说明：
# 1. 每个GPU会运行一个独立的进程，进程间通过NCCL后端进行通信
# 2. 模型会自动进行FSDP（完全分片数据并行）配置
# 3. 如果要使用特定的GPU，可以设置CUDA_VISIBLE_DEVICES环境变量：
#    CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 train.py ...
# 4. 训练数据会自动在多个GPU之间平均分配
# 5. 主进程（rank 0）负责保存模型和记录日志 