[WIP] Large scale rejection sampling with VLLM

[msaroufim]

NOt intended to merge

This is a POC using Deepseek R (70B distlillation) 1, for some reason I'm running into connectivity issues for Qwen for large scale rejection sampling

If you squint it's a form of RL where you have a policy (generated kernels) with a tradeoff between explore (generate completely random samples) and exploit (generate using the best existing samples). It's not really gradient based RL but infra for that tends to be more complex whereas with this approach you can just purely use inference engines

python scripts/run_vllm_prototype.py --operations relu

deepseek-ai/DeepSeek-R1-Distill-Llama-70B

The core idea is that we have

1. Workers that generate samples running VLLM
2. Workers that execute those samples
3. A coordinator running on CPU and collecting metrics
4. Queue requests and store metrics using redis

An obvious problem is that utilization is insanely low for executors but that might be unavoidable considering noisy neighbor problems

Prototyping on 8 GPUs but the goal is that this should run on 1,000