# 基于LAMRA实现拍立淘同款embed模型训练
set -x 

## 分布式配置，一个节点一张卡
NUM_GPUS=1
DISTRIBUTED_ARGS="
    --nnodes=1 \
    --nproc_per_node ${NUM_GPUS} \
    --rdzv_backend c10d \
    --rdzv_endpoint localhost:0 \
"

## 模型设置
### backbone proj模型设置
MODEL_ID=qwen3vl
MODEL_LOCAL_PATH=/data/oss_bucket_0/peifeng/vlm2vec/models/qwen/Qwen3-VL-Embedding-2B/
EMBED_DIM=2048

## 数据集设置
### 训练数据 json（本地路径，支持单个 json / json 目录 / 逗号分隔的多个 json）
DATA_PATH=/data/oss_bucket_0/peifeng/multi_vlm2vec/dataset/tuwen_train.json
MAX_LINES=10000
IMAGE_TYPE="both"
TEXT_REQUIRE=True
### 图片所在的本地目录（不传则用代码里的默认值）
IMAGE_PATH_PREFIX=/data/oss_bucket_0/peifeng/multi_vlm2vec/dataset/merge/image

## InfoNCE 候选池设置
TEMPERATURE=0.02
### InfoNCE 负样本共享的卡数：1=只用本卡 batch；8=单机 8 卡共享；-1/world_size=全部卡共享
### 候选池 = 2 * NEG_GATHER_SIZE * PER_DEVICE_BATCH_SIZE，负样本数 = 候选池 - 1
### 本地单卡训练保持 1 即可，多卡时按需放大
NEG_GATHER_SIZE=1

## 训练设置
### 训练参数设置
TRAIN_VISION_ENCODER=True                              
USE_VISION_LORA=False                 
TRAIN_LLM=True   

USE_LORA=True                                           
Q_LORA=False                                           
LORA_R=128                                                
LORA_ALPHA=256 

### DS和基本训练参数设置
DS_STAGE=zero2                                         
PER_DEVICE_BATCH_SIZE=4                            
GRAD_ACCUM=1                                            
NUM_EPOCHS=5                                         
SAVE_STEPS=200
LR=1e-5                                               
MODEL_MAX_LEN=512

## 保存和读取设置
OUTPUT_DIR=/data/oss_bucket_0/peifeng/multi_vlm2vec/multi_vlm2vec_train/test_train_2607/
RUN_ID=tbstars_train_debug
RESUME=False  ## 除了可以给定False或者True，也可以直接给定读取的地址，其中当没有ckpt保存在其中的时候，就选择False

echo "begin training!"

torchrun $DISTRIBUTED_ARGS train/train_tuwen_qwen3vl_infonce.py \
    --model_id $MODEL_ID \
    --data_path ${DATA_PATH} \
    --image_path_prefix ${IMAGE_PATH_PREFIX} \
    --max_lines $MAX_LINES \
    --output_dir ${OUTPUT_DIR}/${RUN_ID} \
    --report_to tensorboard \
    --run_name $RUN_ID \
    --deepspeed ./ds_configs/${DS_STAGE}.json \
    --bf16 True \
    --num_train_epochs $NUM_EPOCHS \
    --per_device_train_batch_size $PER_DEVICE_BATCH_SIZE \
    --per_device_eval_batch_size 12 \
    --gradient_accumulation_steps $GRAD_ACCUM \
    --eval_strategy "no" \
    --save_strategy "steps" \
    --save_total_limit 200 \
    --learning_rate ${LR} \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --lr_scheduler_type "constant_with_warmup" \
    --logging_steps 10 \
    --tf32 True \
    --model_max_length $MODEL_MAX_LEN \
    --gradient_checkpointing False \
    --dataloader_num_workers 0 \
    --train_vision_encoder $TRAIN_VISION_ENCODER \
    --use_vision_lora $USE_VISION_LORA \
    --use_lora $USE_LORA \
    --q_lora $Q_LORA \
    --lora_r $LORA_R \
    --lora_alpha $LORA_ALPHA \
    --model_local_path $MODEL_LOCAL_PATH \
    --use_flash_attn True \
    --save_steps $SAVE_STEPS \
    --stage classify \
    --embed_dim ${EMBED_DIM} \
    --resume_from_checkpoint $RESUME \
    --image_type $IMAGE_TYPE \
    --text_require $TEXT_REQUIRE \
    --train_llm $TRAIN_LLM \
    --temperature $TEMPERATURE \
    --neg_gather_size $NEG_GATHER_SIZE