GRPO Single GPU / colocate with TRL==0.18.x

⚠️ Please check that this feature request hasn't been suggested before.

• I searched previous Ideas in Discussions didn't find any similar feature requests.
• I searched previous Issues didn't find any similar feature requests.

Feature description

With TRL 0.18.0 release, TRL introduces colocating vLLM on the same GPU. This would first of all improve GPU utilisation and second it would enable single GPU trainig with vLLM and GRPO.

Currently users are forced to have a second GPU to serve vLLM. In past releases, running vLLM on the same GPU was still possible.

✔️ Solution

In order to include the colocate feature it's needed to adjust the yaml file, the schema and core/trainers/grpo/init.py to accept TRL's new config arguments:

• vllm_mode
• vllm_gpu_memory_utilization
• vllm_tensor_parallel_size

On Top, axolotl needs to change its logic since it currently tries to serve vLLM on a different GPU. I can image that the code in file src/axolotl/core/trainers/grpo/trainer.py needs to be adjusted.

In commit c907ac173e586e759da9c71aa2a51e294648d1f0 the function def _move_model_to_vllm(self): was still used before vLLM was served on a different GPU

❓ Alternatives

No response

Additional Context

No response

Acknowledgements

• My issue title is concise, descriptive, and in title casing.
• I have searched the existing issues to make sure this feature has not been requested yet.
• I have provided enough information for the maintainers to understand and evaluate this request.