🙋 Add Optional Eager Execution Mode for vLLM Serving

[ucalyptus]

Overview

This pull request introduces a new optional flag, enforce_eager, to the VLLM serving script. This enhancement gives users control over the execution mode of the model, allowing them to choose between eager execution or the default hybrid CUDA graph approach based on their specific needs. See https://docs.vllm.ai/en/stable/api/offline_inference/llm.html#vllm.LLM

Changes Made

• Added a new boolean parameter enforce_eager_flag to the ScriptArguments class in trl/scripts/vllm_serve.py
• Updated the main function to pass this flag as enforce_eager to the model configuration
• Set the default value to False to maintain backward compatibility with existing behavior

Benefits

• Memory Optimization: Enforcing eager execution can reduce memory usage in certain scenarios
• Performance Flexibility: Allows users to choose between memory efficiency and potential performance gains
• Better Control: Provides explicit control over execution strategy without modifying source code

Technical Details

The changes primarily affect:

• trl/scripts/vllm_serve.py - Added new flag to the argument parser and passed it to the model configuration

Testing

Verified the flag works correctly in both states:

• When enforce_eager_flag=True: Model uses pure eager execution
• When enforce_eager_flag=False (default): Model uses hybrid CUDA graph execution

Related Issues

Fixes #XXXX

Checklist Before Submitting

• Read the [contributor guidelines](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request)
• Updated documentation to reflect the new parameter
• Tested the feature with different flag values
• Maintained backward compatibility

Who Can Review?

Anyone in the community is welcome to review once tests have passed. Team members with expertise in model serving or CUDA optimization would be particularly valuable reviewers.

[ucalyptus]

@qgallouedec

[qgallouedec]

LGTM! Thanks

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.