Remove redundant get_kbit_device_map()

examples/scripts/distillation.py

@@ -82,12 +82,7 @@ def main(script_args, training_args, model_args):
     from datasets import load_dataset
     from transformers import GenerationConfig
 
-    from trl import (
-        LogCompletionsCallback,
-        get_kbit_device_map,
-        get_peft_config,
-        get_quantization_config,
-    )
+    from trl import LogCompletionsCallback, get_peft_config, get_quantization_config
     from trl.experimental.distillation import DistillationTrainer
 
     ################
@@ -99,7 +94,6 @@ def main(script_args, training_args, model_args):
         attn_implementation=model_args.attn_implementation,
         dtype=model_args.dtype,
         use_cache=False if training_args.gradient_checkpointing else True,
-        device_map=get_kbit_device_map() if quantization_config is not None else None,
         quantization_config=quantization_config,
     )
     training_args.model_init_kwargs = model_kwargs
@@ -109,7 +103,6 @@ def main(script_args, training_args, model_args):
         attn_implementation=model_args.attn_implementation,
         dtype=model_args.dtype,
         use_cache=True,
-        device_map=get_kbit_device_map() if quantization_config is not None else None,
         quantization_config=quantization_config,
     )
     if training_args.teacher_model_init_kwargs is not None:

examples/scripts/dpo_vlm.py

@@ -66,7 +66,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -89,7 +88,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     model = AutoModelForImageTextToText.from_pretrained(

examples/scripts/gkd.py

@@ -57,7 +57,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -80,7 +79,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     training_args.model_init_kwargs = model_kwargs
@@ -93,7 +91,6 @@
     )
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     training_args.teacher_model_init_kwargs = teacher_model_kwargs

examples/scripts/gold.py

@@ -60,7 +60,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -84,7 +83,6 @@
         attn_implementation=model_args.attn_implementation,
         dtype=model_args.dtype,
         use_cache=False if training_args.gradient_checkpointing else True,
-        device_map=get_kbit_device_map() if quantization_config is not None else None,
         quantization_config=quantization_config,
     )
     training_args.model_init_kwargs = model_kwargs
@@ -96,7 +94,6 @@
         attn_implementation=model_args.attn_implementation,
         dtype=model_args.dtype,
         use_cache=True,
-        device_map=get_kbit_device_map() if quantization_config is not None else None,
         quantization_config=quantization_config,
     )
     if training_args.teacher_model_init_kwargs is not None:

examples/scripts/grpo_vlm.py

@@ -71,7 +71,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -93,7 +92,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        training_args.model_init_kwargs["device_map"] = get_kbit_device_map()
         training_args.model_init_kwargs["quantization_config"] = quantization_config
 
     ################

examples/scripts/gspo.py

@@ -60,7 +60,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -82,7 +81,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        training_args.model_init_kwargs["device_map"] = get_kbit_device_map()
         training_args.model_init_kwargs["quantization_config"] = quantization_config
 
     ################

examples/scripts/gspo_vlm.py

@@ -60,7 +60,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -82,7 +81,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        training_args.model_init_kwargs["device_map"] = get_kbit_device_map()
         training_args.model_init_kwargs["quantization_config"] = quantization_config
 
     ################

examples/scripts/mpo_vlm.py

@@ -50,7 +50,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -73,7 +72,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     model = AutoModelForImageTextToText.from_pretrained(

examples/scripts/nash_md.py

@@ -58,14 +58,7 @@
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, GenerationConfig
 
-from trl import (
-    LogCompletionsCallback,
-    ModelConfig,
-    ScriptArguments,
-    TrlParser,
-    get_kbit_device_map,
-    get_quantization_config,
-)
+from trl import LogCompletionsCallback, ModelConfig, ScriptArguments, TrlParser, get_quantization_config
 from trl.experimental.nash_md import NashMDConfig, NashMDTrainer
 
 
@@ -84,7 +77,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, **model_kwargs)

examples/scripts/online_dpo.py

@@ -57,7 +57,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -79,7 +78,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, **model_kwargs)

examples/scripts/online_dpo_vlm.py

@@ -90,7 +90,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -113,7 +112,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     # Load the VLM model using correct architecture (from GRPO pattern)

examples/scripts/ppo/ppo.py

@@ -32,7 +32,7 @@
     HfArgumentParser,
 )
 
-from trl import ModelConfig, ScriptArguments, get_kbit_device_map, get_peft_config, get_quantization_config
+from trl import ModelConfig, ScriptArguments, get_peft_config, get_quantization_config
 from trl.experimental.ppo import PPOConfig, PPOTrainer
 
 
@@ -83,7 +83,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, padding_side="left")

examples/scripts/ppo/ppo_tldr.py

@@ -32,7 +32,7 @@
     HfArgumentParser,
 )
 
-from trl import ModelConfig, ScriptArguments, get_kbit_device_map, get_peft_config, get_quantization_config
+from trl import ModelConfig, ScriptArguments, get_peft_config, get_quantization_config
 from trl.experimental.ppo import PPOConfig, PPOTrainer
 
 
@@ -90,7 +90,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, padding_side="left")

examples/scripts/prm.py

@@ -52,13 +52,7 @@
 from datasets import load_dataset
 from transformers import AutoModelForTokenClassification, AutoTokenizer, HfArgumentParser
 
-from trl import (
-    ModelConfig,
-    ScriptArguments,
-    get_kbit_device_map,
-    get_peft_config,
-    get_quantization_config,
-)
+from trl import ModelConfig, ScriptArguments, get_peft_config, get_quantization_config
 from trl.experimental.prm import PRMConfig, PRMTrainer
 
 
@@ -80,7 +74,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, use_fast=True)

examples/scripts/reward_modeling.py

@@ -55,15 +55,7 @@
 from datasets import load_dataset
 from transformers import AutoModelForSequenceClassification, HfArgumentParser
 
-from trl import (
-    ModelConfig,
-    RewardConfig,
-    RewardTrainer,
-    ScriptArguments,
-    get_kbit_device_map,
-    get_peft_config,
-    get_quantization_config,
-)
+from trl import ModelConfig, RewardConfig, RewardTrainer, ScriptArguments, get_peft_config, get_quantization_config
 
 
 logger = logging.get_logger(__name__)
@@ -85,7 +77,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     model = AutoModelForSequenceClassification.from_pretrained(

examples/scripts/rloo_vlm.py

@@ -71,7 +71,6 @@
     RLOOTrainer,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -93,7 +92,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        training_args.model_init_kwargs["device_map"] = get_kbit_device_map()
         training_args.model_init_kwargs["quantization_config"] = quantization_config
 
     ################

examples/scripts/sdft.py

@@ -71,7 +71,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -331,7 +330,6 @@ def _run_tooluse_eval(
     )
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     training_args.model_init_kwargs = model_kwargs

examples/scripts/sdpo.py

@@ -73,7 +73,6 @@
     ModelConfig,
     ScriptArguments,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -298,7 +297,6 @@ def _run_accuracy_eval(
     )
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
-        training_args.model_init_kwargs["device_map"] = get_kbit_device_map()
         training_args.model_init_kwargs["quantization_config"] = quantization_config
 
     if script_args.dataset_path is not None:

examples/scripts/sft_video_llm.py

@@ -62,7 +62,7 @@
 from qwen_vl_utils import process_vision_info
 from transformers import AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig, Qwen2VLProcessor
 
-from trl import ModelConfig, ScriptArguments, SFTConfig, SFTTrainer, TrlParser, get_kbit_device_map
+from trl import ModelConfig, ScriptArguments, SFTConfig, SFTTrainer, TrlParser
 
 
 def download_video(url: str, cache_dir: str) -> str:
@@ -195,7 +195,6 @@ class CustomScriptArguments(ScriptArguments):
     model_kwargs = dict(
         revision=model_args.model_revision,
         dtype=dtype,
-        device_map=get_kbit_device_map(),
         quantization_config=bnb_config,
     )
 

examples/scripts/sft_vlm.py

@@ -63,7 +63,6 @@
     SFTConfig,
     SFTTrainer,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -86,7 +85,6 @@
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     model = AutoModelForImageTextToText.from_pretrained(model_args.model_name_or_path, **model_kwargs)

examples/scripts/sft_vlm_gemma3.py

@@ -68,7 +68,6 @@
     SFTConfig,
     SFTTrainer,
     TrlParser,
-    get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
 )
@@ -152,7 +151,6 @@ def main():
     quantization_config = get_quantization_config(model_args)
     if quantization_config is not None:
         # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
-        model_kwargs["device_map"] = get_kbit_device_map()
         model_kwargs["quantization_config"] = quantization_config
 
     model = AutoModelForImageTextToText.from_pretrained(model_args.model_name_or_path, **model_kwargs)