Merge pull request #145 from georgian-io/remove-accelerate

Remove unused `accelerate` code

Author: benjaminye

llmtune/finetune/lora.py

@@ -1,7 +1,5 @@
 from os.path import join
 
-import bitsandbytes as bnb
-import torch
 from datasets import Dataset
 from peft import (
     LoraConfig,
@@ -38,21 +36,10 @@ def __init__(self, config: Config, directory_helper: DirectoryHelper):
         self.model = None
         self.tokenizer = None
 
-        """ TODO: Figure out how to handle multi-gpu
-        if config.accelerate:
-            self.accelerator = Accelerator()
-            self.accelerator.state.deepspeed_plugin.deepspeed_config[
-                "train_micro_batch_size_per_gpu"
-            ] = self.config.training.training_args.per_device_train_batch_size
-
-        if config.accelerate:
-            # device_index = Accelerator().process_index
-            self.device_map = None #{"": device_index}
-        else:
-        """
         self.device_map = self._model_config.device_map
 
         self._load_model_and_tokenizer()
+        self._inject_lora()
 
     def _load_model_and_tokenizer(self):
         ckpt = self._model_config.hf_model_ckpt
@@ -67,11 +54,7 @@ def _load_model_and_tokenizer(self):
     def _get_model(self):
         model = AutoModelForCausalLM.from_pretrained(
             self._model_config.hf_model_ckpt,
-            quantization_config=(
-                BitsAndBytesConfig(**self._model_config.bitsandbytes.model_dump())
-                if not self.config.accelerate
-                else None
-            ),
+            quantization_config=BitsAndBytesConfig(**self._model_config.bitsandbytes.model_dump()),
             use_cache=False,
             device_map=self.device_map,
             torch_dtype=self._model_config.casted_torch_dtype,
@@ -90,19 +73,10 @@ def _get_tokenizer(self):
         return tokenizer
 
     def _inject_lora(self):
-        if not self.config.accelerate:
-            self.model.gradient_checkpointing_enable()
-            self.model = prepare_model_for_kbit_training(self.model)
+        self.model.gradient_checkpointing_enable()
+        self.model = prepare_model_for_kbit_training(self.model)
         self.model = get_peft_model(self.model, self._lora_config)
 
-        if not self.config.accelerate:
-            self.optimizer = bnb.optim.Adam8bit(self.model.parameters(), lr=self._training_args.learning_rate)
-            self.lr_scheduler = torch.optim.lr_scheduler.ConstantLR(self.optimizer)
-        if self.config.accelerate:
-            self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
-                self.model, self.optimizer, self.lr_scheduler
-            )
-
     def finetune(self, train_dataset: Dataset):
         logging_dir = join(self._weights_path, "/logs")
         training_args = TrainingArguments(
@@ -123,7 +97,6 @@ def finetune(self, train_dataset: Dataset):
             args=training_args,
             dataset_text_field="formatted_prompt",  # TODO: maybe move consts to a dedicated folder
             callbacks=[progress_callback],
-            # optimizers=[self.optimizer, self.lr_scheduler],
             **self._sft_args.model_dump(),
         )
 

llmtune/inference/lora.py

@@ -44,16 +44,11 @@ def _get_merged_model(self, weights_path: str):
         self.model = AutoPeftModelForCausalLM.from_pretrained(
             weights_path,
             torch_dtype=self.config.model.casted_torch_dtype,
-            quantization_config=(BitsAndBytesConfig(**self.config.model.bitsandbytes.model_dump())),
+            quantization_config=BitsAndBytesConfig(**self.config.model.bitsandbytes.model_dump()),
             device_map=self.device_map,
             attn_implementation=self.config.model.attn_implementation,
         )
 
-        """TODO: figure out multi-gpu
-        if self.config.accelerate:
-            self.model = self.accelerator.prepare(self.model)
-        """
-
         model = self.model.merge_and_unload()
 
         tokenizer = AutoTokenizer.from_pretrained(self._weights_path, device_map=self.device_map)

llmtune/pydantic_models/config_model.py

@@ -237,10 +237,6 @@ class AblationConfig(BaseModel):
 class Config(BaseModel):
     save_dir: Optional[str] = Field("./experiments", description="Folder to save to")
     ablation: AblationConfig
-    accelerate: Optional[bool] = Field(
-        False,
-        description="set to True if you want to use multi-gpu training; then launch with `accelerate launch --config_file ./accelerate_config toolkit.py`",
-    )
     data: DataConfig
     model: ModelConfig
     lora: LoraConfig