Fix OPD teacher endpoint for vLLM test - #99
Conversation
There was a problem hiding this comment.
Code Review
This pull request updates the on-policy distillation (OPD) implementation to use the vLLM disaggregated /inference/v1/generate endpoint instead of the OpenAI-compatible /v1/completions endpoint. It also updates the test configuration to allow GPU settings and memory utilization to be configured via environment variables. Feedback includes ensuring that the teacher server is launched using vllm serve rather than vllm.entrypoints.openai.api_server to avoid 404 errors on the new endpoint, and adding an assertion to prevent an IndexError if the configured training GPUs exceed the total available GPUs.
| "1", | ||
| "--gpu-memory-utilization", | ||
| "0.6", | ||
| str(TEACHER_GPU_MEMORY_UTILIZATION), |
There was a problem hiding this comment.
The teacher server is launched using vllm.entrypoints.openai.api_server (line 46), which typically only exposes OpenAI-compatible endpoints (like /v1/completions). However, the test has been updated to use the disaggregated /inference/v1/generate endpoint (line 133). To ensure /inference/v1/generate is available and prevent 404 Not Found errors, consider launching the server using vllm serve instead of vllm.entrypoints.openai.api_server.
| NUM_GPUS = int(os.environ.get("SLIME_TEST_NUM_GPUS", "8")) | ||
| NUM_TRAIN_GPUS = int(os.environ.get("SLIME_TEST_NUM_TRAIN_GPUS", "4")) |
There was a problem hiding this comment.
If SLIME_TEST_NUM_TRAIN_GPUS is configured to be greater than or equal to SLIME_TEST_NUM_GPUS, _get_gpu_split() will raise an IndexError when attempting to access all_gpus[NUM_TRAIN_GPUS]. Adding an assertion here ensures that at least one GPU is reserved for the teacher server and prevents runtime errors.
| NUM_GPUS = int(os.environ.get("SLIME_TEST_NUM_GPUS", "8")) | |
| NUM_TRAIN_GPUS = int(os.environ.get("SLIME_TEST_NUM_TRAIN_GPUS", "4")) | |
| NUM_GPUS = int(os.environ.get("SLIME_TEST_NUM_GPUS", "8")) | |
| NUM_TRAIN_GPUS = int(os.environ.get("SLIME_TEST_NUM_TRAIN_GPUS", "4")) | |
| assert NUM_GPUS > NUM_TRAIN_GPUS, f"NUM_GPUS ({NUM_GPUS}) must be greater than NUM_TRAIN_GPUS ({NUM_TRAIN_GPUS}) to reserve at least one GPU for the teacher server." |
|
Superseded by #141. Root cause is the built-in OPD reward_func sending vime's disaggregated /inference/v1/generate payload, while the test launches a stock vllm openai api_server that only serves /v1/completions (the contract documented by merged #62). This PR flipped the test's --rm-url to /inference/v1/generate, but the stock teacher doesn't serve that route (would 404). #141 instead fixes the code to speak /v1/completions + echo + prompt_logprobs=1 and read choices[0].prompt_logprobs; validated e2e on gb200 (EXIT_RC=0, teacher_log_probs flow into opd_reverse_kl). Closing in favor of #141. |
|
Correction to my earlier close comment: I was wrong that flipping the URL to /inference/v1/generate would 404. That route IS mounted on the stock openai api_server in this vendored vLLM (build_app attaches the disagg router for any generate-capable model), so #99's approach was actually viable. The real reason #141 prefers /v1/completions is portability (it exists on any vLLM incl. upstream and any OpenAI-compatible server; /inference/v1/generate is a fork-only endpoint under entrypoints/serve/disagg/) plus matching the contract already documented by merged #62. The 400 itself came from the test pointing --rm-url at /v1/completions while the code POSTed an /inference/v1/generate-shaped body. Still closing in favor of #141 (validated e2e), just correcting the record. |
No description provided.