Fix type for model_init_kwargs when passed as CLI JSON string (#5230)

Author: albertvillanova

trl/experimental/bco/bco_config.py

@@ -142,7 +142,7 @@ class BCOConfig(_BaseConfig):
             "needed."
         },
     )
-    model_init_kwargs: dict[str, Any] | None = field(
+    model_init_kwargs: dict[str, Any] | str | None = field(
         default=None,
         metadata={
             "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the "

trl/experimental/cpo/cpo_config.py

@@ -163,7 +163,7 @@ class CPOConfig(_BaseConfig):
         default=None,
         metadata={"help": "Whether the model is an encoder-decoder model."},
     )
-    model_init_kwargs: dict[str, Any] | None = field(
+    model_init_kwargs: dict[str, Any] | str | None = field(
         default=None,
         metadata={
             "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model "

trl/experimental/kto/kto_config.py

@@ -132,7 +132,7 @@ class KTOConfig(_BaseConfig):
             "This is useful when training without the reference model to reduce the total GPU memory needed."
         },
     )
-    model_init_kwargs: dict[str, Any] | None = field(
+    model_init_kwargs: dict[str, Any] | str | None = field(
         default=None,
         metadata={
             "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model "

trl/experimental/online_dpo/online_dpo_config.py

@@ -16,6 +16,8 @@
 from dataclasses import dataclass, field
 from typing import Any
 
+from transformers import TrainingArguments
+
 from ...trainer.base_config import _BaseConfig
 
 
@@ -159,6 +161,8 @@ class may differ from those in [`~transformers.TrainingArguments`].
     > - `learning_rate`: Defaults to `5e-7` instead of `5e-5`.
     """
 
+    _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"]
+
     # Parameters whose default values are overridden from TrainingArguments
     learning_rate: float = field(
         default=5e-7,
@@ -361,7 +365,7 @@ class may differ from those in [`~transformers.TrainingArguments`].
             "is not compatible with vLLM generation."
         },
     )
-    model_init_kwargs: dict[str, Any] | None = field(
+    model_init_kwargs: dict[str, Any] | str | None = field(
         default=None,
         metadata={
             "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model "

trl/experimental/orpo/orpo_config.py

@@ -123,7 +123,7 @@ class ORPOConfig(_BaseConfig):
             "argument, you need to specify if the model returned by the callable is an encoder-decoder model."
         },
     )
-    model_init_kwargs: dict[str, Any] | None = field(
+    model_init_kwargs: dict[str, Any] | str | None = field(
         default=None,
         metadata={
             "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model "

trl/trainer/dpo_config.py

@@ -138,7 +138,7 @@ class DPOConfig(_BaseConfig):
     )
 
     # Parameters that control the model
-    model_init_kwargs: dict[str, Any] | None = field(
+    model_init_kwargs: dict[str, Any] | str | None = field(
         default=None,
         metadata={
             "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of "

trl/trainer/grpo_config.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass, field
+from typing import Any
 
 from transformers import TrainingArguments
 
@@ -332,7 +333,7 @@ class GRPOConfig(_BaseConfig):
     )
 
     # Parameters that control the model and reference model
-    model_init_kwargs: dict | str | None = field(
+    model_init_kwargs: dict[str, Any] | str | None = field(
         default=None,
         metadata={
             "help": "Keyword arguments for `transformers.AutoModelForCausalLM.from_pretrained`, used when the `model` "

trl/trainer/reward_config.py

@@ -91,7 +91,7 @@ class may differ from those in [`~transformers.TrainingArguments`].
     )
 
     # Parameters that control the model
-    model_init_kwargs: dict[str, Any] | None = field(
+    model_init_kwargs: dict[str, Any] | str | None = field(
         default=None,
         metadata={
             "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of "

trl/trainer/rloo_config.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass, field
+from typing import Any
 
 from transformers import TrainingArguments
 
@@ -224,7 +225,7 @@ class RLOOConfig(_BaseConfig):
     )
 
     # Parameters that control the model and reference model
-    model_init_kwargs: dict | str | None = field(
+    model_init_kwargs: dict[str, Any] | str | None = field(
         default=None,
         metadata={
             "help": "Keyword arguments for `transformers.AutoModelForCausalLM.from_pretrained`, used when the `model` "

trl/trainer/sft_config.py

@@ -121,7 +121,7 @@ class SFTConfig(_BaseConfig):
     )
 
     # Parameters that control the model
-    model_init_kwargs: dict[str, Any] | None = field(
+    model_init_kwargs: dict[str, Any] | str | None = field(
         default=None,
         metadata={
             "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of "