test: Add e2e test for fast_inference vLLM path#7136
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new test file tests/fast_inference/test_fast_inference.py to verify GRPO training with the fast_inference=True vLLM rollout path. The review feedback suggests dynamically selecting the tensor data type (falling back to torch.float16 if bfloat16 is not supported) to prevent runtime errors on older GPUs, and correcting misleading type hints in the helper function get_unsloth_peft_model.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| model, tokenizer = FastLanguageModel.from_pretrained( | ||
| model_name=MODEL_NAME, | ||
| max_seq_length=MAX_SEQ_LENGTH, | ||
| dtype=torch.bfloat16, |
There was a problem hiding this comment.
Dynamically select the tensor data type (e.g., using torch.bfloat16 only if supported, falling back to torch.float16) to prevent runtime errors on devices that do not natively support bfloat16.
| dtype=torch.bfloat16, | |
| dtype=torch.bfloat16 if is_bfloat16_supported() else torch.float16, |
References
- When probing GPU capabilities or running test kernels on older GPUs (such as Turing/Volta, sm < 80), dynamically select the tensor data type (e.g., using
torch.bfloat16only if supported, falling back totorch.float16) to prevent runtime errors on devices that do not natively support bfloat16.
for more information, see https://pre-commit.ci
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4b46889498
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| ) | ||
|
|
||
|
|
||
| if __name__ == "__main__": |
There was a problem hiding this comment.
In the repo CPU workflow I checked, .github/workflows/studio-backend-ci.yml runs python -m pytest tests/ and does not ignore tests/fast_inference, but pytest will only import this file because the actual model load/train path is guarded by if __name__ == "__main__". That means CI, and anyone invoking pytest tests/fast_inference, gets no executable assertions for this new test, so regressions in the vLLM fast_inference path this file is meant to catch won't fail. Please wrap the body in a test_* function with appropriate GPU/vLLM marks/skips, or add the script to an explicit runner.
Useful? React with 👍 / 👎.
| trainer = GRPOTrainer( | ||
| model=model, | ||
| processing_class=tokenizer, | ||
| reward_funcs=[length_reward_func], | ||
| args=training_args, | ||
| train_dataset=dataset, | ||
| ) |
There was a problem hiding this comment.
This instantiation never verifies that GRPOTrainer actually switches to the colocated vLLM rollout backed by model.vllm_engine; if the Unsloth patch that sets args.use_vllm or assigns self.llm = model.vllm_engine regresses, training can still succeed via the normal generation path and this test would pass despite not exercising the path it claims to cover. After constructing the trainer, please assert the vLLM mode/engine binding (for example that the trainer's llm is the model's vllm_engine) before calling train().
Useful? React with 👍 / 👎.
| def length_reward_func(completions, **kwargs) -> list[float]: | ||
| """Simple reward: prefer non-empty completions.""" | ||
| responses = [completion[0]["content"] for completion in completions] | ||
| return [1.0 if len(r) > 0 else 0.0 for r in responses] |
There was a problem hiding this comment.
When both generations for a prompt are non-empty — the normal case for this model — this reward gives every completion in the GRPO group the same score, so the group-normalized advantages are zero and the LoRA optimizer path can do no meaningful update while trainer.train() still returns successfully. If this is meant to be an end-to-end GRPO training check, use a reward that creates variation within each prompt group or assert that adapter weights changed.
Useful? React with 👍 / 👎.
1 participant
What This PR Does
This PR adds an end-to-end test script for the GRPO trainer when argument
fast_inference=TruefromFastLanguageModel.from_pretrained()is set for vLLM rollout. This add a single test filetests/fast_inference/test_fast_inference.py.The script can be run with
python tests/fast_inference/test_fast_inference.py. It is designed to match the end-to-end scripts undertests/qlora/.Please let me know if any improvements are required and if the file location is appropriate.
Why
To add coverage for vLLM rollout on different devices.
Validation
This test passes on CUDA, but fails on Intel XPU hardware with
AssertionError: Torch not compiled with CUDA enablederror, which is the aim this test captures.Traceback:
Environments
NVIDIA A100 (CUDA):
Intel Arc Pro B60 (XPU):