File size: 4,421 Bytes
ef503c4 | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | #!/bin/bash
# Multi-Language GRPO Training Script
# This script trains a model on multiple programming languages (Python, C, C++, Java, JavaScript)
# Set CUDA device (modify this to select which GPU to use)
# Examples:
# export CUDA_VISIBLE_DEVICES=0 # Use only GPU 0
# export CUDA_VISIBLE_DEVICES=1 # Use only GPU 1
# export CUDA_VISIBLE_DEVICES=0,1 # Use GPU 0 and 1
# export CUDA_VISIBLE_DEVICES=0,1,2,3 # Use all 4 GPUs
# export CUDA_VISIBLE_DEVICES=0 # Default: use GPU 0
# Set project directory
export PROJECT_DIR="/ocean/projects/cis240137p/yfei4/secureCodeGen"
export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH}"
# Set environment variables for API access
export AI_GATEWAY_API_KEY="${AI_GATEWAY_API_KEY}"
# Training configuration
MODEL_ABBREVIATION="gemma-7B"
MODEL_NAME="/ocean/projects/cis240137p/yfei4/secureCodeGen/models/merged_PSC-2_codegemma-7b-it_sft_2-epochs"
MAX_SEQ_LENGTH=1280
LOAD_IN_4BIT="true"
# Training hyperparameters
LEARNING_RATE=1e-5
NUM_EPOCHS=10
BATCH_SIZE=1
GRAD_ACCUM_STEPS=4 # Increased to maintain effective batch size
MAX_STEPS=-1 # -1 means use num_epochs
# SAVE_STRATEGY="epoch"
# SAVE_STEPS=100
SAVE_STRATEGY="steps"
SAVE_STEPS=100
SAVE_TOTAL_LIMIT=3
# GRPO configuration (Optimized for 79GB GPU to avoid OOM)
NUM_GENERATIONS=4 # Reduced from 4 to 2 - most critical for memory!
MAX_PROMPT_LENGTH=512 # Reduced from 512
MAX_COMPLETION_LENGTH=768 # Reduced from 1024
TEMPERATURE=0.9
TOP_P=1.0
# Data paths
TRAIN_DATA="${PROJECT_DIR}/data/rl_training_data_complete.json"
EVAL_DATA="" # Optional
# Reward configuration
SECURITY_WEIGHT=0.6
CAPABILITY_WEIGHT=0.4
TIMEOUT=10
MAX_TEST_CASES=25
# Reward component weights
REWARD_WEIGHT_FORMAT=0.2
REWARD_WEIGHT_FUNCTION=0.2
REWARD_WEIGHT_RUNNABLE=0.2
REWARD_WEIGHT_CAPABILITY=0.2
REWARD_WEIGHT_SECURITY=0.2
# Format pattern for code extraction
FORMAT_PATTERN='<think>(.+?)</think>\n<code>(.+?)</code>'
# Output configuration
TIMESTAMP=$(date +%Y%m%d-%H%M%S)
EXPERIMENT_NAME="${MODEL_ABBREVIATION}-multilang-grpo-${TIMESTAMP}"
OUTPUT_DIR="${PROJECT_DIR}/outputs/${MODEL_ABBREVIATION}_multilang_grpo_${TIMESTAMP}"
LOGGING_DIR="${OUTPUT_DIR}/logs"
# Wandb configuration
USE_WANDB="--use_wandb" # Remove this flag to disable wandb
WANDB_PROJECT="llm-multilang-code-gen-security"
WANDB_ENTITY="arihants-carnegie-mellon-university"
WANDB_RUN_NAME="${EXPERIMENT_NAME}"
# Resume from checkpoint (optional)
# RESUME_FROM_CHECKPOINT="true" # Auto-find latest
# RESUME_FROM_CHECKPOINT="${OUTPUT_DIR}/checkpoint-100" # Specific checkpoint
RESUME_FROM_CHECKPOINT="" # Don't resume
# Create output directory
mkdir -p ${OUTPUT_DIR}
mkdir -p ${LOGGING_DIR}
# Save this script to output directory
cp $0 ${OUTPUT_DIR}/training_script.sh
# Run training
python ${PROJECT_DIR}/rl/train_multilang_grpo.py \
--model_name ${MODEL_NAME} \
--max_seq_length ${MAX_SEQ_LENGTH} \
--load_in_4bit ${LOAD_IN_4BIT} \
--learning_rate ${LEARNING_RATE} \
--num_train_epochs ${NUM_EPOCHS} \
--per_device_train_batch_size ${BATCH_SIZE} \
--gradient_accumulation_steps ${GRAD_ACCUM_STEPS} \
--max_steps ${MAX_STEPS} \
--save_strategy ${SAVE_STRATEGY} \
--save_steps ${SAVE_STEPS} \
--save_total_limit ${SAVE_TOTAL_LIMIT} \
--num_generations ${NUM_GENERATIONS} \
--max_prompt_length ${MAX_PROMPT_LENGTH} \
--max_completion_length ${MAX_COMPLETION_LENGTH} \
--temperature ${TEMPERATURE} \
--top_p ${TOP_P} \
--train_dataset_path ${TRAIN_DATA} \
--security_weight ${SECURITY_WEIGHT} \
--capability_weight ${CAPABILITY_WEIGHT} \
--timeout ${TIMEOUT} \
--max_test_cases ${MAX_TEST_CASES} \
--reward_weight_format ${REWARD_WEIGHT_FORMAT} \
--reward_weight_function ${REWARD_WEIGHT_FUNCTION} \
--reward_weight_runnable ${REWARD_WEIGHT_RUNNABLE} \
--reward_weight_capability ${REWARD_WEIGHT_CAPABILITY} \
--reward_weight_security ${REWARD_WEIGHT_SECURITY} \
--format_pattern "${FORMAT_PATTERN}" \
--output_dir ${OUTPUT_DIR} \
--logging_dir ${LOGGING_DIR} \
--experiment_name ${EXPERIMENT_NAME} \
${USE_WANDB} \
--wandb_project ${WANDB_PROJECT} \
--wandb_entity ${WANDB_ENTITY} \
--wandb_run_name ${WANDB_RUN_NAME} \
--wandb_bash_script $0 \
${RESUME_FROM_CHECKPOINT:+--resume_from_checkpoint ${RESUME_FROM_CHECKPOINT}}
echo "Training completed! Model saved to: ${OUTPUT_DIR}"
|