Updates for trl 0.16.0 - mostly for GRPO (#2437) [skip ci]

* add grpo scale_rewards config for trl#3135

* options to connect to vllm server directly w grpo trl#3094

* temperature support trl#3029

* sampling/generation kwargs for grpo trl#2989

* make vllm_enable_prefix_caching a config param trl#2900

* grpo multi-step optimizeations trl#2899

* remove overrides for grpo trainer

* bump trl to 0.16.0

* add cli  to start vllm-serve via trl

* call the python module directly

* update to use vllm with 2.6.0 too now and call trl vllm serve from module

* vllm 0.8.1

* use python3

* use sys.executable

* remove context and wait for start

* fixes to make it actually work

* fixes so the grpo tests pass with new vllm paradigm

* explicit host/port and check in start vllm

* make sure that vllm doesn't hang by setting quiet so outouts go to dev null

* also bump bnb to latest release

* add option for wait from cli and nccl debugging for ci

* grpo + vllm test on separate devices for now

* make sure grpo + vllm tests runs single worker since pynccl comms would conflict

* fix cli

* remove wait and add caching for argilla dataset

* refactoring configs

* chore: lint

* add vllm config

* fixup vllm grpo args

* fix one more incorrect schema/config path

* fix another vlllm reference and increase timeout

* make the tests run a bit faster

* change mbsz back so it is correct for grpo

* another change mbsz back so it is correct for grpo

* fixing cli args

* nits

* adding docs

* docs

* include tensor parallel size for vllm in pydantic schema

* moving start_vllm, more docs

* limit output len for grpo vllm

* vllm enable_prefix_caching isn't a bool cli arg

* fix env ordering in tests and also use pid check when looking for vllm

---------

Co-authored-by: Salman Mohammadi <salman.mohammadi@outlook.com>

src/axolotl/utils/schemas/trl.py

@@ -20,27 +20,30 @@ class TRLConfig(BaseModel):
     )
 
     # GRPO specific args
-    # Ref: https://github.com/huggingface/trl/blob/e3244d2d096ff1e2e248c931d06d39e165e20623/trl/trainer/grpo_config.py#L22
-    use_vllm: bool | None = Field(
+    # Ref: https://github.com/huggingface/trl/blob/26d86757a7c7e24e397ea44f57ecce6031dfac01/trl/trainer/grpo_config.py#L23
+    use_vllm: bool = Field(
         default=False,
         json_schema_extra={"description": "Whether to use VLLM for RL training"},
     )
-    vllm_device: str | None = Field(
-        default="auto",
-        json_schema_extra={"description": "Device to use for VLLM"},
+    vllm_server_host: str | None = Field(
+        default="0.0.0.0",  # nosec B104
+        json_schema_extra={"description": "Host of the vLLM server to connect to"},
     )
-    vllm_gpu_memory_utilization: float | None = Field(
-        default=0.9,
-        json_schema_extra={"description": "GPU memory utilization for VLLM"},
+    vllm_server_port: int | None = Field(
+        default=8000,
+        json_schema_extra={"description": "Port of the vLLM server to connect to"},
     )
-    vllm_dtype: str | None = Field(
-        default="auto",
-        json_schema_extra={"description": "Data type for VLLM"},
-    )
-    vllm_max_model_len: int | None = Field(
+    vllm_server_timeout: int | None = Field(
         default=None,
         json_schema_extra={
-            "description": "Maximum length of the model context for VLLM"
+            "description": "Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up "
+            "after the timeout, a `ConnectionError` is raised."
+        },
+    )
+    vllm_guided_decoding_regex: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."
         },
     )
 
@@ -85,3 +88,48 @@ class TRLConfig(BaseModel):
             "description": "Sync steps for the reference model. Requires `sync_ref_model=True`."
         },
     )
+    scale_rewards: bool = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Whether to scale the rewards for GRPO by dividing them by their standard deviation."
+        },
+    )
+
+    temperature: float | None = Field(
+        default=None,
+        json_schema_extra={"description": "Sampling temperature for the GRPO policy."},
+    )
+    top_p: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Top-p sampling probability for the generation policy."
+        },
+    )
+    top_k: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Top-k sampling for the generation policy."},
+    )
+    min_p: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Minimum probability for the generation policy."
+        },
+    )
+    repetition_penalty: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far."
+        },
+    )
+    num_iterations: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of iterations per batch (denoted as μ in the algorithm) for GRPO."
+        },
+    )
+    epsilon: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Epsilon value for clipping in the GRPO algorithm."
+        },
+    )