Use transformers utilities when possible (#2064)

* use transformers' availability functions

* require from transformers

* rm file

* fix no peft

* fix import

* don't alter  _peft_available

* fix require_diffusers

* style

* transformers>=4.40 and add back `is_liger_kernel_available`

examples/research_projects/stack_llama_2/scripts/sft_llama2.py

@@ -27,11 +27,12 @@
     AutoTokenizer,
     BitsAndBytesConfig,
     HfArgumentParser,
+    is_torch_npu_available,
+    is_torch_xpu_available,
     set_seed,
 )
 
 from trl import SFTConfig, SFTTrainer
-from trl.import_utils import is_npu_available, is_xpu_available
 from trl.trainer import ConstantLengthDataset
 
 
@@ -197,9 +198,9 @@ def create_datasets(tokenizer, args, seed=None):
 
 # Free memory for merging weights
 del base_model
-if is_xpu_available():
+if is_torch_xpu_available():
     torch.xpu.empty_cache()
-elif is_npu_available():
+elif is_torch_npu_available():
     torch.npu.empty_cache()
 else:
     torch.cuda.empty_cache()

examples/research_projects/toxicity/scripts/evaluate-toxicity.py

@@ -20,9 +20,7 @@
 import torch
 from datasets import load_dataset
 from tqdm import tqdm
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-from trl.import_utils import is_npu_available, is_xpu_available
+from transformers import AutoModelForCausalLM, AutoTokenizer, is_torch_npu_available, is_torch_xpu_available
 
 
 toxicity = evaluate.load("ybelkada/toxicity", "DaNLP/da-electra-hatespeech-detection", module_type="measurement")
@@ -66,9 +64,9 @@
 output_file = args.output_file
 max_new_tokens = args.max_new_tokens
 context_length = args.context_length
-if is_xpu_available():
+if is_torch_xpu_available():
     device = torch.xpu.current_device()
-elif is_npu_available():
+elif is_torch_npu_available():
     device = torch.npu.current_device()
 else:
     device = torch.cuda.current_device() if torch.cuda.is_available() else "cpu"
@@ -137,9 +135,9 @@
     print(f"Model: {model_id} - Mean: {mean} - Std: {std}")
 
     model = None
-    if is_xpu_available():
+    if is_torch_xpu_available():
         torch.xpu.empty_cache()
-    elif is_npu_available():
+    elif is_torch_npu_available():
         torch.npu.empty_cache()
     else:
         torch.cuda.empty_cache()

examples/scripts/ddpo.py

@@ -33,10 +33,9 @@
 import torch.nn as nn
 from huggingface_hub import hf_hub_download
 from huggingface_hub.utils import EntryNotFoundError
-from transformers import CLIPModel, CLIPProcessor, HfArgumentParser
+from transformers import CLIPModel, CLIPProcessor, HfArgumentParser, is_torch_npu_available, is_torch_xpu_available
 
 from trl import DDPOConfig, DDPOTrainer, DefaultDDPOStableDiffusionPipeline
-from trl.import_utils import is_npu_available, is_xpu_available
 
 
 @dataclass
@@ -116,9 +115,9 @@ def aesthetic_scorer(hub_model_id, model_filename):
         model_filename=model_filename,
         dtype=torch.float32,
     )
-    if is_npu_available():
+    if is_torch_npu_available():
         scorer = scorer.npu()
-    elif is_xpu_available():
+    elif is_torch_xpu_available():
         scorer = scorer.xpu()
     else:
         scorer = scorer.cuda()

examples/scripts/ppo.py

@@ -24,11 +24,10 @@
 from datasets import load_dataset
 from peft import LoraConfig
 from tqdm import tqdm
-from transformers import AutoTokenizer, HfArgumentParser, pipeline
+from transformers import AutoTokenizer, HfArgumentParser, is_torch_npu_available, is_torch_xpu_available, pipeline
 
 from trl import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead, PPOConfig, PPOTrainer, set_seed
 from trl.core import LengthSampler
-from trl.import_utils import is_npu_available, is_xpu_available
 
 
 tqdm.pandas()
@@ -142,9 +141,9 @@ def collator(data):
 # to the same device as the PPOTrainer.
 device = ppo_trainer.accelerator.device
 if ppo_trainer.accelerator.num_processes == 1:
-    if is_xpu_available():
+    if is_torch_xpu_available():
         device = "xpu:0"
-    elif is_npu_available():
+    elif is_torch_npu_available():
         device = "npu:0"
     else:
         device = 0 if torch.cuda.is_available() else "cpu"  # to avoid a `pipeline` bug

examples/scripts/ppo_multi_adapter.py

@@ -19,11 +19,16 @@
 from datasets import load_dataset
 from peft import LoraConfig
 from tqdm import tqdm
-from transformers import AutoTokenizer, BitsAndBytesConfig, HfArgumentParser
+from transformers import (
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    HfArgumentParser,
+    is_torch_npu_available,
+    is_torch_xpu_available,
+)
 
 from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer
 from trl.core import LengthSampler
-from trl.import_utils import is_npu_available, is_xpu_available
 
 
 input_min_text_length = 6
@@ -86,7 +91,7 @@ def tokenize(example):
 )
 model = AutoModelForCausalLMWithValueHead.from_pretrained(
     script_args.model_name,
-    device_map={"": "xpu:0"} if is_xpu_available() else {"": "npu:0"} if is_npu_available else {"": 0},
+    device_map={"": "xpu:0"} if is_torch_xpu_available() else {"": "npu:0"} if is_torch_npu_available else {"": 0},
     peft_config=lora_config,
     quantization_config=nf4_config,
     reward_adapter=script_args.rm_adapter,

requirements.txt

@@ -1,7 +1,7 @@
 datasets>=1.17.0
 torch>=1.4.0
 tqdm
-transformers>=4.39.0
+transformers>=4.40.0
 accelerate
 peft>=0.3.0
 tyro>=0.5.7

setup.py

@@ -77,7 +77,7 @@
 
 REQUIRED_PKGS = [
     "torch>=1.4.0",
-    "transformers>=4.39.0",
+    "transformers>=4.40.0",
     "numpy>=1.18.2;platform_system!='Windows'",
     "numpy<2;platform_system=='Windows'",
     "accelerate",

tests/slow/test_dpo_slow.py

@@ -21,18 +21,19 @@
 from datasets import load_dataset
 from parameterized import parameterized
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers.testing_utils import require_bitsandbytes, require_peft, require_torch_accelerator, torch_device
+from transformers.utils import is_peft_available
 
-from trl import DPOConfig, DPOTrainer, is_peft_available
+from trl import DPOConfig, DPOTrainer
 
-from ..testing_utils import require_bitsandbytes, require_non_cpu, require_peft, torch_device
 from .testing_constants import DPO_LOSS_TYPES, DPO_PRECOMPUTE_LOGITS, GRADIENT_CHECKPOINTING_KWARGS, MODELS_TO_TEST
 
 
 if is_peft_available():
     from peft import LoraConfig, PeftModel
 
 
-@require_non_cpu
+@require_torch_accelerator
 class DPOTrainerSlowTester(unittest.TestCase):
     def setUp(self):
         self.dataset = load_dataset("trl-internal-testing/mlabonne-chatml-dpo-pairs-copy", split="train[:10%]")

tests/slow/test_sft_slow.py

@@ -21,25 +21,26 @@
 from datasets import load_dataset
 from parameterized import parameterized
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-
-from trl import SFTConfig, SFTTrainer, is_peft_available
-from trl.models.utils import setup_chat_format
-
-from ..testing_utils import (
+from transformers.testing_utils import (
     require_bitsandbytes,
-    require_liger_kernel,
-    require_multi_accelerator,
-    require_non_cpu,
     require_peft,
+    require_torch_accelerator,
+    require_torch_multi_accelerator,
 )
+from transformers.utils import is_peft_available
+
+from trl import SFTConfig, SFTTrainer
+from trl.models.utils import setup_chat_format
+
+from ..testing_utils import require_liger_kernel
 from .testing_constants import DEVICE_MAP_OPTIONS, GRADIENT_CHECKPOINTING_KWARGS, MODELS_TO_TEST, PACKING_OPTIONS
 
 
 if is_peft_available():
     from peft import LoraConfig, PeftModel
 
 
-@require_non_cpu
+@require_torch_accelerator
 class SFTTrainerSlowTester(unittest.TestCase):
     def setUp(self):
         self.train_dataset = load_dataset("stanfordnlp/imdb", split="train[:10%]")
@@ -270,7 +271,7 @@ def test_sft_trainer_transformers_mp_gc_peft(self, model_name, packing, gradient
     @parameterized.expand(
         list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS, GRADIENT_CHECKPOINTING_KWARGS, DEVICE_MAP_OPTIONS))
     )
-    @require_multi_accelerator
+    @require_torch_multi_accelerator
     def test_sft_trainer_transformers_mp_gc_device_map(
         self, model_name, packing, gradient_checkpointing_kwargs, device_map
     ):

tests/test_alignprop_trainer.py

@@ -16,8 +16,9 @@
 
 import torch
 from parameterized import parameterized
+from transformers.utils import is_peft_available
 
-from trl import is_diffusers_available, is_peft_available
+from trl import is_diffusers_available
 
 from .testing_utils import require_diffusers
 

tests/test_bco_trainer.py

@@ -20,11 +20,12 @@
 from datasets import load_dataset
 from parameterized import parameterized
 from transformers import AutoModel, AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers.testing_utils import require_peft
 
 from trl import BCOConfig, BCOTrainer
 from trl.trainer.bco_trainer import _process_tokens, _tokenize
 
-from .testing_utils import require_no_wandb, require_peft
+from .testing_utils import require_no_wandb
 
 
 class BCOTrainerTester(unittest.TestCase):

tests/test_cpo_trainer.py

@@ -18,11 +18,10 @@
 from datasets import load_dataset
 from parameterized import parameterized
 from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers.testing_utils import require_peft
 
 from trl import CPOConfig, CPOTrainer
 
-from .testing_utils import require_peft
-
 
 class CPOTrainerTester(unittest.TestCase):
     def setUp(self):

tests/test_ddpo_trainer.py

@@ -15,8 +15,9 @@
 import unittest
 
 import torch
+from transformers.utils import is_peft_available
 
-from trl import is_diffusers_available, is_peft_available
+from trl import is_diffusers_available
 
 from .testing_utils import require_diffusers
 

tests/test_dpo_trainer.py

@@ -28,11 +28,12 @@
     AutoProcessor,
     AutoTokenizer,
 )
+from transformers.testing_utils import require_bitsandbytes, require_peft
 
 from trl import DPOConfig, DPOTrainer, FDivergenceType
 from trl.trainer.dpo_trainer import _build_tokenized_answer, _truncate_tokens
 
-from .testing_utils import require_bitsandbytes, require_no_wandb, require_peft
+from .testing_utils import require_no_wandb
 
 
 class TestBuildTokenizedAnswer(unittest.TestCase):

tests/test_kto_trainer.py

@@ -18,11 +18,12 @@
 from datasets import load_dataset
 from parameterized import parameterized
 from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers.testing_utils import require_peft
 
 from trl import KTOConfig, KTOTrainer
 from trl.trainer.kto_trainer import _get_kl_dataset, _process_tokens, _tokenize
 
-from .testing_utils import require_no_wandb, require_peft
+from .testing_utils import require_no_wandb
 
 
 class KTOTrainerTester(unittest.TestCase):

tests/test_no_peft.py

@@ -18,9 +18,8 @@
 
 import pytest
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-from .testing_utils import is_peft_available, require_peft
+from transformers import AutoTokenizer
+from transformers.utils import import_utils
 
 
 class DummyDataset(torch.utils.data.Dataset):
@@ -70,27 +69,14 @@ def __getitem__(self, idx):
 ]
 
 
-@require_peft
 class TestPeftDependancy(unittest.TestCase):
     def setUp(self):
         self.causal_lm_model_id = "trl-internal-testing/tiny-random-GPTNeoXForCausalLM"
         self.seq_to_seq_model_id = "trl-internal-testing/tiny-random-T5ForConditionalGeneration"
 
-        if is_peft_available():
-            from peft import LoraConfig, get_peft_model
-
-            lora_config = LoraConfig(
-                r=16,
-                lora_alpha=32,
-                lora_dropout=0.05,
-                bias="none",
-                task_type="CAUSAL_LM",
-            )
-
-            causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
-            self.peft_model = get_peft_model(causal_lm_model, lora_config)
-
     def test_no_peft(self):
+        _peft_available = import_utils._peft_available
+        import_utils._peft_available = False  # required so that is_peft_available() returns False
         with patch.dict(sys.modules, {"peft": None}):
             from trl import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead
 
@@ -100,8 +86,11 @@ def test_no_peft(self):
 
             _trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.causal_lm_model_id)
             _trl_seq2seq_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(self.seq_to_seq_model_id)
+        import_utils._peft_available = _peft_available
 
     def test_imports_no_peft(self):
+        _peft_available = import_utils._peft_available
+        import_utils._peft_available = False  # required so that is_peft_available() returns False
         with patch.dict(sys.modules, {"peft": None}):
             from trl import (  # noqa: F401
                 AutoModelForCausalLMWithValueHead,
@@ -110,8 +99,11 @@ def test_imports_no_peft(self):
                 PPOTrainer,
                 PreTrainedModelWrapper,
             )
+        import_utils._peft_available = _peft_available
 
     def test_ppo_trainer_no_peft(self):
+        _peft_available = import_utils._peft_available
+        import_utils._peft_available = False  # required so that is_peft_available() returns False
         with patch.dict(sys.modules, {"peft": None}):
             from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer
 
@@ -154,3 +146,4 @@ def test_ppo_trainer_no_peft(self):
             # check expected stats
             for stat in EXPECTED_STATS:
                 assert stat in train_stats
+        import_utils._peft_available = _peft_available

tests/test_orpo_trainer.py

@@ -18,11 +18,10 @@
 from datasets import load_dataset
 from parameterized import parameterized
 from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers.testing_utils import require_peft
 
 from trl import ORPOConfig, ORPOTrainer
 
-from .testing_utils import require_peft
-
 
 class ORPOTrainerTester(unittest.TestCase):
     def setUp(self):