Skip to content

reproduce-local.md

Latest commit

 

History

History
117 lines (79 loc) · 6.41 KB

reproduce-local.md

File metadata and controls

117 lines (79 loc) · 6.41 KB
Raw

Reproduce the single-GPU run

The measured machine used an RTX 5070 Ti 16GB in WSL Ubuntu 24.04, Python 3.12, PyTorch 2.10.0+cu129, Transformers 5.7.0, and vLLM 0.19.1. The training and vLLM interpreters live in separate environments. Linux/WSL is required for the private Unix-socket worker and /dev/shm staging path.

These are setup instructions for a fresh checkout, not an instruction to rebuild a working training environment. The publication checks rerun CPU regressions in the existing validated environment; a fresh-machine installation and the long GPU experiment were not rerun for publication.

1. Clone the implementation

git clone --branch codex/full-code https://github.com/hahaws777/cs336_A5.git
cd cs336_A5
nvidia-smi

Make sure the NVIDIA GPU is visible inside WSL, other GPU-heavy tasks are closed, and the installed driver supports the CUDA wheels. Reserve sufficient host RAM for trainer/optimizer offload, vLLM, and temporary weight copies. /dev/shm must accommodate a roughly 3GB weight payload plus headroom. There is no measured universal minimum host-RAM claim here.

free -h
df -h /dev/shm .

The complete eight-method suite writes approximately 96GB of model snapshots in addition to the base model, environments, and logs. Do not start it on a nearly full disk.

2. Create fresh, separate environments

Install uv first. From the repository root:

uv venv --python 3.12 .venv
uv pip install --python .venv/bin/python --torch-backend=cu129 \
  -e . 'torch==2.10.0' 'transformers==5.7.0'

uv venv --python 3.12 .venv-vllm
uv pip install --python .venv-vllm/bin/python --torch-backend=cu129 \
  'vllm==0.19.1' 'torch==2.10.0' 'transformers==5.7.0'

.venv/bin/python -c 'import torch, transformers; print(torch.__version__, transformers.__version__, torch.cuda.is_available())'
.venv-vllm/bin/python -c 'import torch, transformers, vllm; print(torch.__version__, transformers.__version__, vllm.__version__)'

The local training path uses SDPA and does not require the optional FlashAttention wheel from the gpu extra. Do not install the entire cloud-oriented extra into an already-working local environment just to obtain vLLM. Pinning the key packages above is not a complete transitive lockfile for both environments.

uv's backend selection and the recommendation for separate environments are described in the official vLLM installation guide. If the resolver reports a conflict, investigate it rather than silently removing the version pins.

3. Prepare the base model

Use allenai/OLMo-2-0425-1B, revision a1847dff35000b4271fa70afc5db10fd29fedbdf. Weights are not in this Git repository. This explicit command downloads the pinned public checkpoint; it is not run automatically by the suite:

.venv/bin/python -c 'from huggingface_hub import snapshot_download; snapshot_download("allenai/OLMo-2-0425-1B", revision="a1847dff35000b4271fa70afc5db10fd29fedbdf", local_dir="models/OLMo-2-0425-1B")'
export MODEL_DIR="$(pwd)/models/OLMo-2-0425-1B"
export VLLM_PYTHON="$(pwd)/.venv-vllm/bin/python"

The local suite requires an existing model directory and runs in offline mode. You can instead point MODEL_DIR at a model already downloaded elsewhere. Keep the same model revision when comparing with the published evidence.

4. Check the implementation without GPU use

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

The Modal naming variable permits import-time launcher tests; no cloud workload is launched. Optional safety/SFT/DPO exercises are outside this targeted test command and are not all implemented.

5. Run a capacity probe before long training

This performs two real optimizer updates on synthetic responses, without saving or overwriting the base model:

PYTORCH_ALLOC_CONF=expandable_segments:True OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \
.venv/bin/python -m scripts.grpo_local_memory_probe \
  --model "$MODEL_DIR" --steps 2 --batch-size 8 --group-size 8 \
  --response-tokens 512 --generation-tokens 512 --logits-chunk-size 32 \
  --output-file results/capacity-new.json

The default diagnostic prompt is short. This command checks your setup but does not exactly reproduce the archived 305-token prompt probe; prompt/context length changes memory use. A passing short probe does not establish that every longer context or batch will fit. The published capacity evidence records the original shape and boundaries.

6. Train

For one method, use the command in the README. For all eight methods:

# Inspect the commands first. This does not launch GPU training.
.venv/bin/python -m scripts.grpo_experiments_local \
  --model "$MODEL_DIR" --output-root results/local-eight-new \
  --num-steps 100 --seed 0 \
  --generation-backend vllm-colocated --vllm-python "$VLLM_PYTHON" \
  --dry-run

# Removing --dry-run launches the long, serial experiment suite.

Use a new directory every time. The launcher rejects existing output roots, freezes source/data, records command/status/checkpoint evidence, and stops on the first failed run. It does not automatically resume or rent a GPU.

Monitor and interpret

tail -f results/local-eight-new/offpolicy_gspo.log
tail -n 3 results/local-eight-new/offpolicy_gspo/metrics.jsonl

step counts rollout iterations. optimizer_updates counts actual optimizer updates. A rollout is 32 questions with 8 answers each; zero-advantage groups may skip updates. The on-policy and off-policy methods make different numbers of updates per rollout.

Look at validation correctness and format rate, not just loss or training reward. Saved checkpoints contain the model/tokenizer but no AdamW/RNG state. Keep the machine awake during training. If switching back to gaming, stop your own training process after a verified checkpoint; do not terminate arbitrary GPU processes by name.

Optional cloud scripts

Modal is optional and not needed for the 16GB result. Set your own naming prefix before importing/running its launchers:

export SUNET_ID="your-unique-prefix"

Authenticate separately with Modal. Only explicit cloud launch commands create workloads, and they may incur charges. The local suite does not call them.