Fix logging warning suppression for transformers 4.56.2

trl/trainer/reward_trainer.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import contextlib
+import logging
 import os
+import re
 from collections import defaultdict
 from collections.abc import Callable
 from contextlib import contextmanager
@@ -61,6 +63,30 @@
 # AutoModelForSequenceClassification adds a new classification head when loading a CausalLM. That head is randomly
 # initialized and triggers a harmless warning about uninitialized weights. We suppress just that specific warning to
 # avoid confusing users.
+
+
+# Old approach using logging filter (for transformers < 4.57.0)
+@contextmanager
+def suppress_from_pretrained_warning(logger: logging.Logger):
+    pattern = re.compile(
+        r"^Some weights of \S+ were not initialized from the model checkpoint at \S+ and are newly initialized: "
+        r"\[.*\]\nYou should probably TRAIN this model on a down-stream task to be able to use it for predictions and "
+        r"inference\.$"
+    )
+
+    class _Filter(logging.Filter):
+        def filter(self, record: logging.LogRecord) -> bool:
+            return not pattern.search(record.getMessage())
+
+    f = _Filter()
+    logger.addFilter(f)
+    try:
+        yield
+    finally:
+        logger.removeFilter(f)
+
+
+# New approach using scoped override (for transformers >= 4.57.0)
 @contextmanager
 def ignore_seqcls_score_missing_key():
     # Scoped override: ignore only the expected seq-clf head key.
@@ -76,6 +102,21 @@ def ignore_seqcls_score_missing_key():
         GenericForSequenceClassification._keys_to_ignore_on_load_missing = old
 
 
+# Version-aware wrapper that chooses the appropriate approach
+@contextmanager
+def suppress_seqcls_warning():
+    # Use the new approach for transformers >= 4.57.0, old approach for earlier versions
+    # The old approach is needed for 4.56.2 to avoid meta tensor issues with device_map=None
+    if Version(transformers.__version__) >= Version("4.57.0"):
+        with ignore_seqcls_score_missing_key():
+            yield
+    else:
+        # Get the transformers logger
+        transformers_logger = logging.getLogger("transformers.modeling_utils")
+        with suppress_from_pretrained_warning(transformers_logger):
+            yield
+
+
 def get_dataset_column_names(dataset: Dataset | IterableDataset) -> list[str]:
     return list(next(iter(dataset)).keys()) if dataset.column_names is None else dataset.column_names
 
@@ -305,7 +346,7 @@ def __init__(
             if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]:
                 model_init_kwargs["device_map"] = None
             model_init_kwargs["num_labels"] = 1  # the only output of the model is the reward score
-            with ignore_seqcls_score_missing_key():
+            with suppress_seqcls_warning():
                 model = create_model_from_path(model, AutoModelForSequenceClassification, **model_init_kwargs)
         else:
             if args.model_init_kwargs is not None: