SemRep: Generative Code Representation Learning with Code Transformations

🏗️ Project Organization

Core Components

📁 src/ — Core Source Code

• eval/ — Evaluation scripts
  • kernelbench/ - KernelBench evaluation and generation scripts
  • editbench/ - EditBench evaluation and generation scripts
• preprocess/ — Data preprocessing and synthesis
• utils/ — Utility functions
  • llm_client.py - LLM API client wrapper
  • prompter.py - Prompt generation utilities
  • kernel_prompter.py - Kernel-specific prompt generation
  • sand_box.py - Code execution sandbox
• services/ - Services API
  • kernel-api.py - KernelBench evaluation API service
  • editbench-api.py - EditBench evaluation API service

📁 benchmark/ — Evaluation Benchmarks

• KernelBench/ — Comprehensive GPU kernel optimization benchmark suite
• editbench/ — Code editing benchmark

📁 external/ — External Dependencies

• verl/ — Volcengine RL framework for scalable reinforcement learning training
  • Custom reward logic for KernelBench in verl/utils/reward_score/semrep_editing_rewards/kernelbench_reward.py
  • Custom reward logic for EditBench in verl/utils/reward_score/semrep_editing_rewards/editbench_reward.py
  • Custom reward logic for CodeContest in verl/utils/reward_score/semrep_editing_rewards/execution_time.py
  • Parallel reward computation for both benchmarks

📁 scripts/ — Training and Evaluation Scripts

• Training Scripts:
  • run_on_cloud.sh - Cloud-based GRPO training
• Evaluation Scripts:
  • kernel_eval.sh - KernelBench evaluation (generation + evaluation)
  • kernel_sem_eval.sh - KernelBench semantic evaluation
  • edit_eval.sh - EditBench evaluation (generation + evaluation)
  • eval_gen.sh - General generation evaluation
• Generation Scripts:
  • one_step_gen.sh - Single-step code generation
  • two_step_gen.sh - Two-step code generation
  • synthesize.sh - Code synthesis

📁 prompts/ — Prompt Templates

• kernelbench/ - KernelBench-specific prompts (ours, kevin, semrep_first)
• editbench/ - EditBench-specific prompts (semrep_first, ours)

🚀 Quick Start

Installation

1. Clone the repository:

git clone <repository-url>
cd CodeRL

2. Set up the environment:

conda create --name coderl python=3.10
conda activate coderl
bash cloud_setup.sh

Basic Usage

1. KernelBench: Generate and Evaluate GPU Kernels

# Generate kernels
python src/eval/kernelbench/generation.py \
    run_name=my_run \
    dataset_src=local \
    model_name=model_name \
    temperature=0.7 \
    num_completions=16

# Evaluate generated kernels
python src/eval/kernelbench/eval_from_generations.py \
    run_name=my_run \
    dataset_src=local \
    level=1 \
    num_gpu_devices=2 \
    timeout=500

2. EditBench: Generate and Evaluate Code Edits

# Generate code edits 
python src/eval/editbench/generate.py \
    --model model_path \
    --generation_path output_path \
    --template_path template_path \
    --temperature 0.7 \
    --top_p 0.9 \
    --split test \
    --num_generations 16 \

# Evaluate generated edits using Docker
cd benchmark/editbench/
bash run_experiment.sh examples/batch_eval.py --gen_path generation_dir \
    --split test \
    --output_file output_file\
    --baseline_file baseline_results.json \
    --max_workers 4

3. Train with GRPO

python -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files=$TRAIN_DATA_FILE \
    data.val_files=$VAL_DATA_FILE \
    data.shuffle=True \
    data.train_batch_size=$ROLLOUT_N_QUERY \
    data.max_prompt_length=$MAX_PROMPT_LENGTH \
    data.max_response_length=$MAX_RESPONSE_LENGTH \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
    actor_rollout_ref.rollout.max_num_batched_tokens=$MAX_NUM_BATCHED_TOKENS \
    actor_rollout_ref.rollout.enable_chunked_prefill=True \
    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \
    actor_rollout_ref.rollout.temperature=0.7 \
    actor_rollout_ref.rollout.top_p=0.9 \
    actor_rollout_ref.actor.strategy=fsdp2 \
    actor_rollout_ref.ref.strategy=fsdp2 \
    actor_rollout_ref.actor.use_dynamic_bsz=True \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.actor.ulysses_sequence_parallel_size=4 \
    actor_rollout_ref.actor.fsdp_config.offload_policy=True \
    actor_rollout_ref.model.enable_activation_offload=True \
    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$MAX_NUM_BATCHED_TOKENS \
    actor_rollout_ref.actor.ppo_epochs=2 \
    actor_rollout_ref.actor.optim.lr=2e-6 \
    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0 \
    actor_rollout_ref.actor.grad_clip=0.5 \
    actor_rollout_ref.actor.clip_ratio_high=0.28 \
    algorithm.kl_ctrl.kl_coef=0.0  \
    actor_rollout_ref.actor.loss_agg_mode="seq-mean-token-sum-norm" \
    actor_rollout_ref.actor.use_kl_loss=False  \
    algorithm.norm_adv_by_std_in_grpo=False \
    actor_rollout_ref.actor.kl_loss_coef=0.0 \
    actor_rollout_ref.model.path=$BASE_MODEL \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.actor.ppo_mini_batch_size=$GLOBAL_BATCH_SIZE \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$MICRO_BATCH_PER_GPU \
    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
    actor_rollout_ref.model.enable_gradient_checkpointing=True \
    actor_rollout_ref.rollout.log_prob_micro_batch_size=256 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.n=$ROLLOUT_N_SAMPLE \
    actor_rollout_ref.ref.log_prob_micro_batch_size=256 \
    actor_rollout_ref.ref.fsdp_config.param_offload=True \
    trainer.logger=['wandb'] \
    trainer.project_name='CodeRL' \
    trainer.experiment_name=$EXPERIMENT_NAME \
    trainer.rollout_data_dir=$SAVE_DIR/rollout_data \
    trainer.val_before_train=False \
    trainer.nnodes=1 \
    trainer.default_local_dir=$SAVE_DIR \
    trainer.n_gpus_per_node=$TRAINER_GPUS_PER_NODE \
    trainer.save_freq=5 \
    trainer.test_freq=5 \
    trainer.total_epochs=$MAX_EPOCHS \
    reward_model.reward_manager=$REWARD_MANAGER \
    +reward_model.reward_kwargs.reward_model_configuration.pretrain_steps=$PRETRAIN_STEPS \
    +reward_model.reward_kwargs.reward_model_configuration.reject_exact_duplicates=True \
    +reward_model.reward_kwargs.reward_model_configuration.duplicate_penalty=0.0