File size: 1,769 Bytes
8a2fd95 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | #!/bin/bash
# Distributed training configuration
MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
MASTER_PORT=${MASTER_PORT:-$(shuf -i 20001-29999 -n 1)}
NNODES=${WORLD_SIZE:-1}
NODE_RANK=0
NPROC_PER_NODE=8
echo "MASTER_ADDR=$MASTER_ADDR"
echo "MASTER_PORT=$MASTER_PORT"
echo "NNODES=$NNODES"
echo "NPROC_PER_NODE=$NPROC_PER_NODE"
# DeepSpeed configuration
deepspeed=./scripts/zero3.json
# Model configuration
llm=/your/cold/start/ckpt # Using HuggingFace model ID
# Training hyperparameters
lr=2e-7
batch_size=1
grad_accum_steps=1
# Training entry point
entry_file=thinkstream/train.py
# Dataset configuration (replace with public dataset names)
datasets=stream_rlvr
# Output configuration
run_name="thinkstream-rl"
output_dir=./output/${run_name}
# Training arguments
args="
grpo \
--args.train.num_train_epochs 1 \
--args.train.output_dir ${output_dir} \
--args.train.deepspeed ${deepspeed} \
--args.train.per_device_train_batch_size ${batch_size} \
--args.train.warmup_ratio 0.03 \
--args.train.learning_rate ${lr} \
--args.train.gradient_accumulation_steps ${grad_accum_steps} \
--args.model.name_or_path ${llm} \
--args.train.save_steps 300 \
--args.data.dataset_use ${datasets} \
--args.model.model_type qwen2.5vl \
--args.train.group_size 8 \
--args.train.micro_batch_size 8 \
--args.train.rollout_max_pixels $((192*28*28)) \
--args.train.rollout_min_pixels $((128*28*28))"
# Launch training
TOKENIZERS_PARALLELISM=false \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
torchrun --nproc_per_node=${NPROC_PER_NODE} \
--master_addr=${MASTER_ADDR} \
--master_port=${MASTER_PORT} \
--node-rank=${NODE_RANK} \
--nnodes=${NNODES} \
${entry_file} ${args}
|