🎀 New defaults: bf16=True

trl/trainer/bco_config.py

@@ -89,6 +89,15 @@ class BCOConfig(TrainingArguments):
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
 
     max_length: Optional[int] = field(
         default=1024,

trl/trainer/cpo_config.py

@@ -95,6 +95,15 @@ class CPOConfig(TrainingArguments):
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
 
     max_length: Optional[int] = field(
         default=1024,

trl/trainer/dpo_config.py

@@ -190,6 +190,15 @@ class DPOConfig(TrainingArguments):
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
 
     # Parameters that control the model and reference model
     model_init_kwargs: Optional[dict[str, Any]] = field(

trl/trainer/grpo_config.py

@@ -213,6 +213,15 @@ class GRPOConfig(TrainingArguments):
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
 
     # Parameters that control the model and reference model
     model_init_kwargs: Optional[Union[dict, str]] = field(

trl/trainer/iterative_sft_config.py

@@ -60,6 +60,15 @@ class may differ from those in [`~transformers.TrainingArguments`].
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
 
     # Parameters that control the model
     model_init_kwargs: Optional[dict[str, Any]] = field(

trl/trainer/kto_config.py

@@ -102,6 +102,15 @@ class KTOConfig(TrainingArguments):
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
 
     max_length: Optional[int] = field(
         default=1024,

trl/trainer/online_dpo_config.py

@@ -87,6 +87,15 @@ class may differ from those in [`~transformers.TrainingArguments`].
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
 
     reward_model_path: Optional[str] = field(
         default=None,

trl/trainer/orpo_config.py

@@ -80,6 +80,15 @@ class ORPOConfig(TrainingArguments):
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
 
     max_length: Optional[int] = field(
         default=1024,

trl/trainer/prm_config.py

@@ -62,6 +62,15 @@ class PRMConfig(TrainingArguments):
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
     average_tokens_across_devices: bool = field(
         default=True,
         metadata={

trl/trainer/reward_config.py

@@ -57,6 +57,15 @@ class may differ from those in [`~transformers.TrainingArguments`].
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
     average_tokens_across_devices: bool = field(
         default=True,
         metadata={

trl/trainer/sft_config.py

@@ -100,6 +100,15 @@ class SFTConfig(TrainingArguments):
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
     average_tokens_across_devices: bool = field(
         default=True,
         metadata={

trl/trainer/utils.py

@@ -1075,6 +1075,15 @@ class may differ from those in [`~transformers.TrainingArguments`].
             )
         },
     )
+    bf16: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA "
+                "architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
+            )
+        },
+    )
 
     run_name: Optional[str] = field(
         default=None,