#!/usr/bin/env bash
set -euo pipefail

export CUDA_VISIBLE_DEVICES=4,5,6,7
export NPROC_PER_NODE=4
export NNODES=1
export NODE_RANK=0
export MASTER_ADDR=127.0.0.1
export MASTER_PORT=29500

export USE_HF=1
export HF_HOME=/raid/data/weifeng/model_cache/huggingface
export HUGGINGFACE_HUB_CACHE=/raid/data/weifeng/model_cache/huggingface/hub
export TRANSFORMERS_CACHE=/raid/data/weifeng/model_cache/huggingface/transformers
export MODELSCOPE_CACHE=/raid/data/weifeng/model_cache/modelscope

mkdir -p "${HF_HOME}" "${HUGGINGFACE_HUB_CACHE}" "${TRANSFORMERS_CACHE}" "${MODELSCOPE_CACHE}"

swift pt \
  --model Qwen/Qwen2.5-Coder-7B \
  --tuner_type full \
  --dataset /raid/data/weifeng/Datasets/code_messages_pre_train.jsonl \
  --torch_dtype bfloat16 \
  --per_device_train_batch_size 1 \
  --learning_rate 1e-5 \
  --gradient_accumulation_steps 16 \
  --packing true \
  --save_steps 500 \
  --save_total_limit 5 \
  --logging_steps 10 \
  --max_length 8192 \
  --num_train_epochs 3 \
  --warmup_ratio 0.03 \
  --dataloader_num_workers 8 \
  --dataset_num_proc 8 \
  --use_hf true \
  --save_only_model false \
  --output_dir /raid/data/weifeng/Pre_train/output/qwen25_coder_7b_full \
  --deepspeed zero3 \
  --attn_impl flash_attn