CS336A5 小显存适配

用小显卡直接训练：CS336A5 的 OLMo-2 全参数强化学习，已在单张 RTX 5070 Ti 16GB 上跑通。

这个项目为 CS336A5 的推理强化学习训练做了低显存适配。按照提供的配置，你可以直接使用自己的 RTX 5070 Ti 16GB 在本地训练 OLMo-2，无需租用云 GPU。通过 BF16 训练状态、梯度重计算、logits 分块，以及 vLLM 批量生成与训练交替使用显存，把全参数训练放到一张消费级显卡上。

This project runs full-parameter reasoning RL for allenai/OLMo-2-0425-1B on a single RTX 5070 Ti. The checkpoint actually contains 1,484,916,736 parameters (~1.485B). The implementation uses native BF16 parameters and AdamW states, checkpointed/chunked scoring, and a colocated vLLM backend with explicit CPU offload and sleep.

Scope of the headline: this specific model and the documented sequence lengths, batches, and software stack. It is not a promise that arbitrary models or contexts fit in 16GB. This is full-parameter training, not LoRA, QLoRA, or 8-bit Adam.

Start here

• Full implementation — this branch: training code, memory optimizations, tests, and experiment notes.
• No-solution starter — upstream assignment skeleton with unimplemented adapters, for working through the assignment yourself. Both branches are public.
• Reproduce locally — environment setup, a capacity check, and the single-GPU training command.
• Evidence and limitations — measured memory, diagnostic speed, and the completed local experiments.

What was demonstrated

• Real full-policy updates on 16GB: a two-update GSPO capacity probe initialized both AdamW moments for all 179 parameter tensors, checked finite/nonzero gradients, and verified parameter changes. At the second update, PyTorch peak allocated memory was 12.802 GiB, with 13.033 GiB reserved. This was a synthetic capacity probe, not a whole-run GPU peak or an accuracy benchmark. Probe evidence.
• Completed local experiments: seven fresh-base 100-rollout runs plus a MaxRL run split into 50 + 50 rollouts. MaxRL's second half was a weights-only warm start with a new seed and optimizer, not exact resumption. Run evidence.
• 497/1024 correct (48.54%) at the final checkpoint of offpolicy_naive on the fixed evaluation subset. This is a single-seed, format-sensitive automated score, not a multi-seed leaderboard result. Other methods did not all improve. Interpretation.
• 8.33x faster diagnostic cycles: colocated vLLM reduced the measured steady-state generation-plus-one-B8-update cycle from 171.875s to 20.643s, including transfers and weight synchronization. This is not an 8.33x end-to-end training speedup. Diagnostic evidence.

How the memory fits

1. BF16 training state: parameters, gradients, and the native AdamW first/second moments use BF16, with no separate FP32 master weights. Probability, entropy, advantages, ratios, and clipping arithmetic remain FP32. This differs from standard FP32-master-weight AMP and has numerical trade-offs.
2. Microbatch 1 + gradient accumulation: keep the intended effective batch and complete reward groups while limiting activation memory.
3. Recompute activations: enable decoder gradient checkpointing. Compute the vocabulary output head in checkpointed chunks of 32 sequence positions, avoiding retention of full-sequence vocabulary logits. This optimized path is OLMo2-specific.
4. Avoid optimizer temporaries: use AdamW(foreach=False).
5. Share the GPU over time: move training weights and optimizer state to CPU before vLLM generation, then discard vLLM weights/KV cache with level-2 sleep before training resumes.

The colocated cycle is:

Trainer -> CPU
    -> wake vLLM weights -> load and verify current weights -> wake KV cache
    -> generate the whole rollout batch -> vLLM sleep(level=2)
    -> trainer + optimizer back to GPU -> score -> backward -> update
    -> synchronize the newly updated weights on the next generation

Sleep does not release every CUDA context allocation. CPU RAM and /dev/shm are part of the memory budget. This is explicit offload, not Windows shared-GPU-memory spillover. See the implementation notes and vLLM's sleep-mode documentation.

Run one local experiment

After following environment setup, run this from the repository root in Linux/WSL. Replace both example paths with your own. This starts a real, potentially long GPU training job; it does not rent cloud hardware.

export MODEL_DIR="/absolute/path/to/OLMo-2-0425-1B"
export VLLM_PYTHON="/absolute/path/to/vllm-env/bin/python"

PYTORCH_ALLOC_CONF=expandable_segments:True \
OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \
.venv/bin/python -m scripts.grpo \
  --model "$MODEL_DIR" --no-trust-remote-code \
  --algorithm offpolicy_gspo --num-steps 100 --seed 0 \
  --learning-rate 1e-5 --rollout-batch-size 256 --group-size 8 \
  --train-batch-size 8 --gradient-accumulation-steps 8 \
  --cliprange 0.0003 --train-examples 6400 --val-examples 1024 \
  --max-tokens 512 --temperature 1 --top-p 1 \
  --low-memory --attention-implementation sdpa \
  --generation-backend vllm-colocated --vllm-python "$VLLM_PYTHON" \
  --eval-every 10 --save-every 25 --no-save-final \
  --output-dir results/gspo-100-new

Use a new output directory. To use Transformers generation instead, set --generation-backend transformers and omit --vllm-python; retain --low-memory.

For the serial eight-method suite, see the suite guide. Checkpoints are model/tokenizer snapshots only: AdamW and RNG state are not saved, so loading one is not exact training resumption.

Implementation map

• Core math and training step: tokenization/masks, rewards and advantages, log-probabilities, policy-gradient losses, accumulation.
• Training loop: rollout collection, fixed old-policy scoring, validation, checkpointing, and algorithm presets.
• Single-GPU vLLM backend: isolated worker, full-weight verification, CPU offload, sleep/wake, and cleanup.
• Capacity probe and backend diagnostic.
• Local suite launcher, with frozen source/data provenance and fail-stop execution.
• Prompting-baseline writeup and training-loop explanation.

Available presets: grpo, grpo_constant, dr_grpo, rft, maxrl, offpolicy_naive, offpolicy_noclip, offpolicy_clip, and offpolicy_gspo. The earlier standard-GRPO cloud baseline is separate from the eight-method local suite.

Tests

Run the GRPO implementation and infrastructure regressions without using the GPU:

CUDA_VISIBLE_DEVICES="" SUNET_ID=cpu-tests OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \
  .venv/bin/python -m pytest -q tests/test_grpo*.py tests/test_vllm_colocated.py \
  tests/test_publication_config.py

SUNET_ID is a configurable Modal resource-name prefix, not a credential. Setting it for import-time tests does not launch Modal. Cloud examples require your own Modal login and an explicit launch command, and can incur charges.

"Full code" means all current implementation work, not every optional assignment exercise: some safety/SFT/DPO adapters remain unimplemented. Do not interpret the targeted GRPO test command as claiming that every optional assignment test passes.

Honest limits

• No controlled before/after VRAM measurement is available, so no percentage-memory-reduction claim is made. Allocated memory is included in reserved memory; do not add the two.
• Training-stage whole-card headroom was narrow. Free memory during vLLM generation is not safe capacity for a larger training microbatch. Long contexts, bigger batches, or other GPU applications may OOM.
• The fixed 1024-example evaluation subset comes from GSM8K test, and was inspected repeatedly. It is not an independent held-out confirmation after model selection; sampling uses temperature 1.
• Strict answer parsing can mark mathematically correct but malformed answers incorrect. One clipped off-policy run degraded severely; these results are retained rather than presented as universal algorithm wins.
• The local 100-rollout experiments do not substitute for the assignment's complete 200-step/multiple-seed requirements. A rollout iteration is not necessarily one optimizer update.
• Raw logs, model checkpoints, credentials, caches, and temporary operator files are not shipped. Small sanitized evidence summaries are included; historical local paths in their provenance are identifiers, not downloadable artifacts.

Attribution

Based on Stanford CS336 Spring 2026 Assignment 5: Alignment, starter commit c2734a2. The original assignment material, datasets, graders, prompt templates, and starter tests are credited to their upstream authors; this is an independent student implementation and is not an official Stanford release.

See the assignment handout and optional safety/RLHF supplement. Existing upstream material and model/data terms still apply; no blanket license is assigned here to third-party content. Follow your course's rules when using solution branches.