Choice A code framework: SFT + GRPO + vLLM rollouts

This repository is a code framework for the course project Choice A: Post-training Implementation, covering:

• Part 1: SFT on Qwen/Qwen2.5-1.5B using the AI-MO/NuminaMath-CoT subset with source == "gsm8k".
• Part 2: GRPO on Qwen/Qwen2.5-1.5B using GSM8K.
• Part 3: vLLM rollout acceleration with HF baseline, colocate mode, and HTTP server mode.

Repository layout

.
├── configs/
│   ├── README.md                 # config layout notes
│   ├── templates/                 # canonical hand-edited YAML templates
│   │   ├── sft.yaml
│   │   ├── grpo.yaml
│   │   └── eval.yaml
│   ├── generated/                 # generated SFT, GRPO, and eval YAML variants
│   ├── deepspeed_zero2.json      # recommended default for 8 × RTX 4090
│   └── deepspeed_zero3.json      # optional memory-saving config
├── scripts/
│   ├── doctor.py                 # environment sanity checks
│   ├── resolve_flash_attn_wheel.py
│   ├── run_sft_8gpu.sh
│   ├── run_grpo_8gpu.sh
│   ├── eval_base.sh
│   ├── eval_sft.sh
│   └── eval_grpo.sh
├── src/llm_project/
│   ├── cli/                      # train/eval entrypoints
│   ├── data/                     # dataset loading and prompt formatting
│   ├── evaluation/               # GSM8K and MMLU evaluators
│   ├── training/                 # losses, rewards, generation, checkpoint helpers
│   ├── rollout/                  # HF/vLLM rollout backends and weight sync helpers
│   ├── config.py
│   ├── distributed.py
│   ├── math_utils.py
│   └── models.py
├── tests/
├── report_templates/
├── env_setup.md
└── reference.md

1. Environment setup

Follow env_setup.md first. The key steps are:

uv python install 3.10
uv venv .venv --python 3.10
source .venv/bin/activate

# Install PyTorch first, then project dependencies and flash-attn.
python scripts/doctor.py

# Training and evaluation log metrics directly to Weights & Biases.
wandb login

The provided server has 8 × RTX 4090 GPUs, so the default scripts use --num_gpus 8 and configs/deepspeed_zero2.json.

2. Where to modify configs

Most changes should be made in these files:

• configs/templates/sft.yaml: model path, NuminaMath filtering, max sequence length, SFT learning rate, epoch count, output directory.
• configs/templates/grpo.yaml: initial policy path, reference model path, rollout group size, max generation length, reward weights, KL coefficient, GRPO learning rate.
• configs/templates/eval.yaml: evaluation model path defaults, GSM8K/MMLU sample limits, output paths.
• configs/deepspeed_zero2.json: micro-batch size, gradient accumulation, bf16, ZeRO stage.

For a quick smoke test, set these values before a full run:

# configs/templates/sft.yaml
dataset:
  max_train_samples: 64
  max_eval_samples: 32
train:
  num_train_epochs: 1
  eval_steps: 5

# configs/templates/grpo.yaml
dataset:
  max_train_samples: 32
rollout:
  group_size: 2
  max_new_tokens: 128

3. Baseline evaluation before post-training

Run baseline GSM8K and MMLU evaluation on the base model with vLLM:

source .venv/bin/activate
bash scripts/eval_base.sh

Outputs:

outputs/eval/base_gsm8k.json
outputs/eval/base_mmlu.json

For a fast debug run:

python -m llm_project.cli.eval_gsm8k --model Qwen/Qwen2.5-1.5B --max_samples 16 --output outputs/eval/debug_gsm8k.json
python -m llm_project.cli.eval_mmlu --model Qwen/Qwen2.5-1.5B --subjects abstract_algebra --max_samples_per_subject 16 --output outputs/eval/debug_mmlu.json

4. Part 1: Run SFT

Launch SFT on all eight GPUs:

bash scripts/run_sft_8gpu.sh

Equivalent explicit command:

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 deepspeed --num_gpus 8 -m llm_project.cli.train_sft \
  --config configs/templates/sft.yaml \
  --ds_config configs/deepspeed_zero2.json

Main outputs:

outputs/sft_qwen25_1p5b_numina_gsm8k/
├── hf/                            # Hugging Face checkpoint for evaluation
├── resolved_sft_config.yaml
└── deepspeed_config.json

SFT train loss, validation loss, and validation perplexity are logged to the llm-course-project wandb project.

Evaluate the SFT checkpoint:

bash scripts/eval_sft.sh outputs/sft_qwen25_1p5b_numina_gsm8k/hf

For the report, export or screenshot the wandb chart containing sft/train/loss, sft/eval/val_loss, and sft/eval/val_ppl.

5. Part 2: Run GRPO

By default, GRPO starts from Qwen/Qwen2.5-1.5B. To run RL after the SFT checkpoint, edit configs/templates/grpo.yaml:

model:
  init_from_sft_checkpoint: outputs/sft_qwen25_1p5b_numina_gsm8k/hf
  reference_model_name_or_path: Qwen/Qwen2.5-1.5B

Launch GRPO:

bash scripts/run_grpo_8gpu.sh

Equivalent explicit command:

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 deepspeed --num_gpus 8 -m llm_project.cli.train_grpo \
  --config configs/templates/grpo.yaml \
  --ds_config configs/deepspeed_zero2.json

Main outputs:

outputs/grpo_qwen25_1p5b_gsm8k/
├── hf/                            # Hugging Face checkpoint for evaluation
├── resolved_grpo_config.yaml
└── deepspeed_config.json

GRPO reward, KL diagnostics, and rollout accuracy are logged to the llm-course-project wandb project.

Evaluate the final GRPO checkpoint:

bash scripts/eval_grpo.sh outputs/grpo_qwen25_1p5b_gsm8k/hf

For Part 3 rollout benchmarking, use the dedicated configs and scripts:

bash scripts/run_grpo_hf_baseline.sh
bash scripts/run_grpo_vllm_colocate_8gpu.sh
bash scripts/run_grpo_vllm_server_pynccl.sh
bash scripts/run_grpo_vllm_server_tmp.sh

The benchmark configs keep the same short GRPO setup while switching only the rollout backend. vLLM modes request sampled-token logprobs and apply token-level truncated importance sampling in the GRPO policy loss.

For the report, export or screenshot the wandb chart containing grpo/train/reward, then run final vLLM evaluation with scripts/eval_grpo.sh.

6. Submit-ready result files to collect

For Part 1 SFT:

wandb chart: sft/train/loss and sft/eval/val_loss
outputs/eval/base_gsm8k.json
outputs/eval/base_mmlu.json
outputs/eval/sft_gsm8k.json
outputs/eval/sft_mmlu.json

For Part 2 GRPO:

wandb chart: grpo/train/reward and eval/grpo/gsm8k_accuracy
outputs/eval/grpo_gsm8k.json
outputs/eval/grpo_mmlu.json

Use report_templates/results_tables.md to copy results into the final report.

7. Implementation notes

SFT uses standard next-token prediction with labels masked over the prompt tokens. The data collator pads input_ids, attention_mask, and labels, using -100 for label padding.

GRPO uses grouped rollouts per prompt. Original GRPO z-normalizes rewards within each group to produce advantages, while Dr. GRPO uses mean-centered rewards without reward-std normalization. The loss uses a clipped policy-gradient term plus a reference-model KL penalty. GSM8K rewards are based on exact normalized numeric answer matching, with an optional small formatting reward.

Standalone MMLU and GSM8K evaluation use vLLM generation. MMLU compares the extracted final boxed answer against A/B/C/D, accepting either boxed letters or boxed 1-through-4 aliases. GSM8K uses numeric answer extraction.

8. Practical memory settings for 8 × RTX 4090

The default config is conservative:

• train_micro_batch_size_per_gpu = 1
• gradient_accumulation_steps = 8
• bf16 enabled
• ZeRO-2
• gradient checkpointing enabled

For faster training, increase train_micro_batch_size_per_gpu and the matching per_device_*_batch_size in the YAML files after confirming memory headroom with nvidia-smi.