remove torch install for now

install torch for tests
fix cuda version for autogptq index
set torch in requirements so that it installs properly
move gptq install around to work with github cicd

.github/workflows/main.yml

@@ -23,11 +23,6 @@ jobs:
             python_version: "3.10"
             pytorch: 2.0.1
             axolotl_extras:
-          - cuda: 118
-            cuda_version: 11.8.0
-            python_version: "3.9"
-            pytorch: 2.0.1
-            axolotl_extras: gptq
     runs-on: self-hosted
     steps:
       - name: Checkout
@@ -73,11 +68,6 @@ jobs:
             pytorch: 2.0.1
             axolotl_extras:
             is_latest: true
-          - cuda: 118
-            cuda_version: 11.8.0
-            python_version: "3.9"
-            pytorch: 2.0.1
-            axolotl_extras: gptq
     runs-on: self-hosted
     steps:
       - name: Checkout

.github/workflows/tests.yml

@@ -24,7 +24,7 @@ jobs:
 
       - name: Install dependencies
         run: |
-          pip install -e .[peft]
+          pip install -e .
           pip install -r requirements-tests.txt
 
       - name: Run tests

README.md

@@ -665,7 +665,6 @@ fsdp:
 fsdp_config:
   fsdp_offload_params: true
   fsdp_state_dict_type: FULL_STATE_DICT
-  fsdp_sync_module_states: true
   fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
 ```
 

docker/Dockerfile

@@ -11,14 +11,13 @@ RUN apt-get update && \
 
 WORKDIR /workspace
 
-RUN pip3 install "peft @ git+https://github.com/huggingface/peft.git@main"
 RUN git clone --depth=1 https://github.com/OpenAccess-AI-Collective/axolotl.git
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN cd axolotl && \
     if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[flash-attn,$AXOLOTL_EXTRAS]; \
+        pip install -e .[flash-attn,gptq,$AXOLOTL_EXTRAS]; \
     else \
-        pip install -e .[flash-attn]; \
+        pip install -e .[flash-attn,gptq]; \
     fi
 
 # fix so that git fetch/pull from remote works

examples/gptq-lora-7b/README.md

@@ -1,8 +0,0 @@
-# LLaMa 7B using LoRA
-
-This is a good place to start for beginners. This will run on an NVIDIA RTX4090 with no other changes needed.
-
-```shell
-accelerate launch scripts/finetune.py examples/gptq-lora-7b/config.yml
-
-```

examples/gptq-lora-7b/config.yml

@@ -1,63 +0,0 @@
-base_model: Neko-Institute-of-Science/LLaMA-7B-4bit-128g
-base_model_config: Neko-Institute-of-Science/LLaMA-7B-4bit-128g
-model_type: LlamaForCausalLM
-tokenizer_type: LlamaTokenizer
-trust_remote_code:
-load_in_8bit: true
-gptq: true
-datasets:
-  - path: vicgalle/alpaca-gpt4
-    type: alpaca
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.02
-adapter:
-lora_model_dir:
-sequence_len: 2048
-max_packed_sequence_len:
-lora_r: 8
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_modules:
-  - q_proj
-  - v_proj
-lora_fan_in_fan_out: false
-wandb_project: llama-7b-lora-int4
-wandb_entity:
-wandb_watch:
-wandb_run_id:
-wandb_log_model:
-output_dir: ./llama-7b-lora-int4
-gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 3
-optimizer: adamw_bnb_8bit
-torchdistx_path:
-lr_scheduler: cosine
-learning_rate: 0.0000002
-train_on_inputs: false
-group_by_length: false
-fp16: true
-bf16: false
-tf32: true
-early_stopping_patience:
-resume_from_checkpoint:
-local_rank:
-logging_steps: 5
-xformers_attention:
-flash_attention:
-gradient_checkpointing: true
-gptq_groupsize: 128
-gptq_model_v1: false
-warmup_steps: 20
-eval_steps: 110
-save_steps: 660
-debug:
-deepspeed:
-weight_decay: 0.0001
-fsdp:
-fsdp_config:
-tokens:
-  pad_token: "<pad>"
-  bos_token: "<s>"
-  eos_token: "</s>"
-  unk_token: "<unk>"

examples/llama-2/gptq-lora.yml

@@ -0,0 +1,76 @@
+base_model: TheBloke/Llama-2-7B-GPTQ
+base_model_config: TheBloke/Llama-2-7B-GPTQ
+is_llama_derived_model: false
+gptq: true
+gptq_bits: 4
+model_type: AutoModelForCausalLM
+tokenizer_type: LlamaTokenizer
+tokenizer_use_fast: true
+tokenizer_legacy: true
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+push_dataset_to_hub:
+hf_use_auth_token: true
+datasets:
+  - path: mhenrichsen/alpaca_2k_test
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+adapter: lora
+lora_model_dir:
+sequence_len: 4096
+sample_packing:
+lora_r: 8
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_modules:
+  - k_proj
+  - o_proj
+  - q_proj
+  - v_proj
+lora_target_linear:
+lora_fan_in_fan_out:
+wandb_project:
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./model-out
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 3
+optimizer: adamw_torch
+adam_beta2: 0.95
+adam_eps: 0.00001
+max_grad_norm: 1.0
+torchdistx_path:
+lr_scheduler: cosine
+lr_quadratic_warmup: true
+learning_rate: 0.000017
+train_on_inputs: false
+group_by_length: false
+bf16: false
+fp16: false
+float16: true
+tf32: true
+gradient_checkpointing: true
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:
+logging_steps: 1
+xformers_attention:
+flash_attention:
+sdp_attention:
+flash_optimum:
+gptq_groupsize:
+gptq_model_v1:
+warmup_steps: 100
+eval_steps:
+save_steps:
+debug:
+deepspeed:
+weight_decay: 0.1
+special_tokens:
+  bos_token: "<s>"
+  eos_token: "</s>"
+  unk_token: "<unk>"

requirements.txt

@@ -1,10 +1,13 @@
+--extra-index-url https://download.pytorch.org/whl/cu118
+--extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
+torch==2.0.1
+auto-gptq
 packaging
 peft @ git+https://github.com/huggingface/peft.git
 transformers @ git+https://github.com/huggingface/transformers.git
 bitsandbytes>=0.41.1
 accelerate @ git+https://github.com/huggingface/accelerate@2a289f6108e77a77a4efffb3f6316bc98538413b
 addict
-evaluate
 fire
 PyYAML>=6.0
 datasets

setup.py

@@ -2,15 +2,27 @@
 
 from setuptools import find_packages, setup
 
-install_requires = []
+
-with open("./requirements.txt", encoding="utf-8") as requirements_file:
+def parse_requirements():
-    # don't include peft yet until we check the int4
+    _install_requires = []
-    # need to manually install peft for now...
+    _dependency_links = []
-    reqs = [r.strip() for r in requirements_file.readlines() if "peft" not in r]
+    with open("./requirements.txt", encoding="utf-8") as requirements_file:
-    reqs = [r for r in reqs if "flash-attn" not in r]
+        lines = [
-    reqs = [r for r in reqs if r and r[0] != "#"]
+            r.strip() for r in requirements_file.readlines() if "auto-gptq" not in r
-    for r in reqs:
+        ]
-        install_requires.append(r)
+        for line in lines:
+            if line.startswith("--extra-index-url"):
+                # Handle custom index URLs
+                _, url = line.split()
+                _dependency_links.append(url)
+            elif "flash-attn" not in line and line and line[0] != "#":
+                # Handle standard packages
+                _install_requires.append(line)
+    return _install_requires, _dependency_links
+
+
+install_requires, dependency_links = parse_requirements()
+
 
 setup(
     name="axolotl",
@@ -19,12 +31,10 @@ setup(
     package_dir={"": "src"},
     packages=find_packages(),
     install_requires=install_requires,
+    dependency_links=dependency_links,
     extras_require={
         "gptq": [
-            "alpaca_lora_4bit @ git+https://github.com/winglian/alpaca_lora_4bit.git@setup_pip",
+            "auto-gptq",
-        ],
-        "gptq_triton": [
-            "alpaca_lora_4bit[triton] @ git+https://github.com/winglian/alpaca_lora_4bit.git@setup_pip",
         ],
         "flash-attn": [
             "flash-attn==2.0.8",
@@ -32,8 +42,5 @@ setup(
         "extras": [
             "deepspeed",
         ],
-        "peft": [
-            "peft @ git+https://github.com/huggingface/peft.git",
-        ],
     },
 )

src/axolotl/monkeypatch/fsdp.py

@@ -1,45 +0,0 @@
-"""
-Monkeypatch to fix fsdp set state when no previous state was set
-"""
-
-import contextlib
-from typing import Generator, Optional
-
-import torch
-from torch import nn
-from torch.distributed.fsdp.api import (
-    OptimStateDictConfig,
-    StateDictConfig,
-    StateDictType,
-)
-from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
-
-
-@staticmethod
-@contextlib.contextmanager
-def state_dict_type_patch(
-    module: nn.Module,
-    state_dict_type: StateDictType,
-    state_dict_config: Optional[StateDictConfig] = None,
-    optim_state_dict_config: Optional[OptimStateDictConfig] = None,
-) -> Generator:
-    prev_state_dict_settings = FullyShardedDataParallel.set_state_dict_type(
-        module,
-        state_dict_type,
-        state_dict_config,
-        optim_state_dict_config,
-    )
-    yield
-    if prev_state_dict_settings.state_dict_type:
-        FullyShardedDataParallel.set_state_dict_type(
-            module,
-            prev_state_dict_settings.state_dict_type,
-            prev_state_dict_settings.state_dict_config,
-            prev_state_dict_settings.optim_state_dict_config,
-        )
-
-
-def replace_fsdp_state_dict_type():
-    torch.distributed.fsdp.fully_sharded_data_parallel.FullyShardedDataParallel.state_dict_type = (
-        state_dict_type_patch
-    )

src/axolotl/utils/callbacks.py

@@ -1,19 +1,9 @@
 """Callbacks for Trainer class"""
 
-from __future__ import annotations
-
 import logging
 import os
-from typing import TYPE_CHECKING, Dict, List
 
-import evaluate
-import numpy as np
-import pandas as pd
-import torch
-import torch.distributed as dist
-from datasets import load_dataset
 from optimum.bettertransformer import BetterTransformer
-from tqdm import tqdm
 from transformers import (
     TrainerCallback,
     TrainerControl,
@@ -23,19 +13,8 @@ from transformers import (
 from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, IntervalStrategy
 
 from axolotl.utils.bench import log_gpu_memory_usage
-from axolotl.utils.distributed import (
-    barrier,
-    gather_scalar_from_all_ranks,
-    get_world_size,
-    is_main_process,
-    zero_first,
-)
-
-if TYPE_CHECKING:
-    from axolotl.utils.trainer import AxolotlTrainingArguments
 
 LOG = logging.getLogger("axolotl.callbacks")
-IGNORE_INDEX = -100
 
 
 class SavePeftModelCallback(TrainerCallback):  # pylint: disable=too-few-public-methods
@@ -117,192 +96,3 @@ class GPUStatsCallback(
             log_gpu_memory_usage(LOG, "while training", self.cfg.device)
             self.logged = True
         return control
-
-
-def bench_eval_callback_factory(trainer, tokenizer):
-    accuracy = evaluate.load("accuracy")
-    abcd_idx = [
-        tokenizer("A", add_special_tokens=False).input_ids[0],
-        tokenizer("B", add_special_tokens=False).input_ids[0],
-        tokenizer("C", add_special_tokens=False).input_ids[0],
-        tokenizer("D", add_special_tokens=False).input_ids[0],
-        tokenizer("E", add_special_tokens=False).input_ids[0],
-        tokenizer("F", add_special_tokens=False).input_ids[0],
-        tokenizer("G", add_special_tokens=False).input_ids[0],
-    ]
-    bench_split = "eval"
-
-    def transform_bench_subject(example):
-        # Split on ':' and trim whitespace
-        parts = example["subject"].split(":")
-        first_part = (
-            parts[0].strip().lower().replace("-", "_")
-        )  # Lowercase the first part
-        second_part = (
-            parts[1].strip().replace("-", "_") if len(parts) > 1 else "all"
-        )  # Replace hyphens with underscores
-
-        # Return the transformed values
-        return {"name": first_part, "subject": second_part}
-
-    if trainer.args.bench_dataset == "mmlu-zs":
-        bench_dataset = load_dataset(
-            "openaccess-ai-collective/mmlu-evals",
-            data_files={
-                "eval": "zero_shot_mmlu_val.json",
-                "test": "zero_shot_mmlu_test.json",
-            },
-        )
-        # bench_dataset = bench_dataset.remove_columns("subject")
-    # MMLU Five-shot (Eval/Test only)
-    elif trainer.args.bench_dataset in ["mmlu", "mmlu-fs"]:
-        bench_dataset = load_dataset(
-            "openaccess-ai-collective/mmlu-evals",
-            data_files={
-                "eval": "five_shot_mmlu_val.json",
-                "test": "five_shot_mmlu_test.json",
-            },
-        )
-        # bench_dataset = bench_dataset.remove_columns('subject')
-    elif "/" in trainer.args.bench_dataset:
-        bench_ds = trainer.args.bench_dataset
-        bench_ds_name = "/".join(bench_ds.split("/", 2)[:2])
-        bench_ds_data_file = "/".join(bench_ds.split("/", 2)[2:])
-        bench_dataset = load_dataset(
-            bench_ds_name,
-            data_files={
-                "eval": bench_ds_data_file,
-            },
-        )
-        bench_dataset["eval"] = bench_dataset["eval"].map(transform_bench_subject)
-    else:
-        raise ValueError(
-            f"unhandled value `{trainer.args.bench_dataset}` for bench_dataset training args"
-        )
-    bench_dataset = bench_dataset[trainer.args.bench_split]
-    if trainer.args.max_bench_samples is not None:
-        bench_dataset = bench_dataset.select(range(trainer.args.max_bench_samples))
-
-    def tokenize_evals(example):
-        source = f"{tokenizer.bos_token}{example['input']}"
-        target = f"{example['output']}{tokenizer.eos_token}"
-
-        tokenized_source = tokenizer(
-            source,
-            max_length=2048,
-            truncation=True,
-            add_special_tokens=False,
-        )
-        tokenized_target = tokenizer(
-            target,
-            max_length=2048,
-            truncation=True,
-            add_special_tokens=False,
-        )
-        input_ids = tokenized_source["input_ids"] + tokenized_target["input_ids"]
-        labels = [IGNORE_INDEX] * len(tokenized_source["input_ids"]) + tokenized_target[
-            "input_ids"
-        ]
-
-        return {
-            "input_ids": input_ids,
-            "labels": labels,
-            "subject": example["subject"],
-        }
-
-    with zero_first(is_main_process()):
-        bench_dataset = bench_dataset.map(tokenize_evals)
-        bench_dataset = bench_dataset.filter(lambda x: x["labels"][-2] in abcd_idx)
-
-    class BenchEvalCallback(TrainerCallback):
-        """
-        TrainerCallback that runs the MMLU evals
-        """
-
-        def on_evaluate(
-            self,
-            args: AxolotlTrainingArguments,
-            state: TrainerState,  # pylint: disable=unused-argument
-            control: TrainerControl,  # pylint: disable=unused-argument
-            metrics: Dict[str, float],  # pylint: disable=unused-argument
-            **kwargs,  # pylint: disable=unused-argument
-        ):
-            data_loader = trainer.get_bench_dataloader(
-                bench_dataset.remove_columns(["input", "subject", "output", "name"])
-            )
-            trainer.model.eval()
-            preds, refs = [], []
-            loss_bench = 0
-            for batch in tqdm(data_loader, total=len(data_loader)):
-                (loss, logits, labels) = trainer.prediction_step(
-                    trainer.model,
-                    batch,
-                    prediction_loss_only=False,
-                )
-                # There are two tokens, the output, and eos token.
-                for i, logit in enumerate(logits):
-                    label_non_zero_id = (batch["labels"][i] != IGNORE_INDEX).nonzero()[
-                        0
-                    ][0]
-                    logit_abcd = logit[label_non_zero_id - 1][abcd_idx]
-                    preds.append(torch.argmax(logit_abcd).item())
-                labels = labels[labels != IGNORE_INDEX].view(-1, 2)[:, 0]
-                refs += [
-                    abcd_idx.index(label) if label in abcd_idx else -1
-                    for label in labels.tolist()
-                ]
-                loss_bench += loss.item()
-            # Extract results by subject.
-            bench_name = bench_dataset["name"]
-            bench_names: dict = {s: {"refs": [], "preds": []} for s in set(bench_name)}
-            for s, p, r in zip(bench_name, preds, refs):  # pylint: disable=invalid-name
-                bench_names[s]["preds"].append(p)
-                bench_names[s]["refs"].append(r)
-            barrier()
-            local_bench_names = bench_names
-            gathered_bench_names: List[Dict] = [{} for _ in range(get_world_size())]
-            # Gather results from all GPUs to GPU 0
-
-            loss_bench_ranks = gather_scalar_from_all_ranks(
-                lambda: loss_bench, get_world_size()
-            )
-            len_data_loader_ranks = gather_scalar_from_all_ranks(
-                lambda: len(data_loader), get_world_size()
-            )
-
-            if not is_main_process():
-                dist.gather_object(local_bench_names, dst=0)
-            else:
-                dist.gather_object(local_bench_names, gathered_bench_names, dst=0)
-                bench_loss = sum(loss_bench_ranks) / sum(len_data_loader_ranks)
-                results = {"bench_loss": bench_loss}
-
-                # Combine results from all GPUs
-                combined_bench_names: Dict[str, Dict[str, List]] = {}
-                for bench_name in gathered_bench_names:
-                    for name, data in bench_name.items():
-                        if name not in combined_bench_names:
-                            combined_bench_names[name] = {"refs": [], "preds": []}
-                        combined_bench_names[name]["refs"].extend(data["refs"])
-                        combined_bench_names[name]["preds"].extend(data["preds"])
-
-                bench_scores = []
-                for (
-                    bench_name
-                ) in combined_bench_names:  # pylint: disable=consider-using-dict-items
-                    bench_score = accuracy.compute(
-                        references=combined_bench_names[bench_name]["refs"],
-                        predictions=combined_bench_names[bench_name]["preds"],
-                    )["accuracy"]
-                    if not pd.isna(bench_score):
-                        results[
-                            f"bench_{bench_split}_accuracy_{bench_name}"
-                        ] = bench_score
-                        bench_scores.append(bench_score)
-                    else:
-                        results[f"bench_{bench_split}_accuracy_{bench_name}"] = 0.0
-                        bench_scores.append(0.0)
-                results[f"bench_{bench_split}_accuracy"] = np.mean(bench_scores)
-                trainer.log(results)
-
-    return BenchEvalCallback

src/axolotl/utils/config.py

@@ -97,9 +97,7 @@ def validate_config(cfg):
             "To calculate the equivalent gradient_accumulation_steps, divide batch_size / micro_batch_size / number of gpus.",
         )
     if cfg.load_4bit:
-        raise ValueError(
+        raise ValueError("cfg.load_4bit parameter has been deprecated")
-            "cfg.load_4bit parameter has been deprecated and replaced by cfg.gptq"
-        )
 
     if cfg.adapter == "qlora":
         if cfg.merge_lora:
@@ -152,16 +150,6 @@ def validate_config(cfg):
     if (cfg.base_model and "falcon" in cfg.base_model.lower()) and cfg.fsdp:
         raise ValueError("FSDP is not supported for falcon models")
 
-    if (
-        cfg.fsdp
-        and cfg.fsdp_config
-        and cfg.fsdp_config.fsdp_state_dict_type
-        and not cfg.fsdp_config.fsdp_sync_module_states
-    ):
-        LOG.warning(
-            "We recommend setting fsdp_config.fsdp_sync_module_states to `true`"
-        )
-
     if (
         cfg.base_model and "mpt" in cfg.base_model.lower()
     ) and cfg.gradient_checkpointing:

src/axolotl/utils/distributed.py

@@ -1,10 +1,8 @@
 """
 utility helpers for distributed checks
 """
-import os
 from contextlib import contextmanager
 
-import torch
 import torch.distributed as dist
 from accelerate import Accelerator
 
@@ -45,10 +43,6 @@ def is_main_process():
     return dist.get_rank() == 0
 
 
-def get_world_size():
-    return int(os.getenv("WORLD_SIZE", "1"))
-
-
 @contextmanager
 def zero_first(is_main):
     """
@@ -59,35 +53,3 @@ def zero_first(is_main):
     yield
     if is_main:  # then rank 0 waits after it has run the context
         barrier()
-
-
-def gather_scalar_from_all_ranks(fn, world_size=1):  # pylint: disable=invalid-name
-    """
-    Run a callable 'fn' on all ranks and gather the results on the specified rank.
-
-    Args:
-    - fn (callable): A function that computes the value. This should not have any side effects.
-    - rank (int, optional): The rank that gathers the values. Default is 0.
-    - world_size (int, optional): Total number of processes in the current distributed setup.
-
-    Returns:
-    - A list of computed values from all ranks if on the gathering rank, otherwise None.
-    """
-    value_scalar = fn()
-    value_tensor = torch.tensor(value_scalar, device=dist.get_rank()).float()
-
-    if not is_main_process():
-        dist.gather(value_tensor, dst=0)
-    else:
-        gathered_tensors = [torch.zeros_like(value_tensor) for _ in range(world_size)]
-        dist.gather(value_tensor, gather_list=gathered_tensors, dst=0)
-
-        # Convert tensors back to their original type (int or float)
-        gathered_values = []
-        for tensor in gathered_tensors:
-            if tensor == tensor.int():
-                gathered_values.append(int(tensor.item()))
-            else:
-                gathered_values.append(float(tensor.item()))
-        return gathered_values
-    return None

src/axolotl/utils/models.py

@@ -4,19 +4,19 @@
 import logging
 import math
 import os
-from pathlib import Path
 from typing import Optional, Tuple  # noqa: F401
 
 import bitsandbytes as bnb
 import torch
 import transformers
 from optimum.bettertransformer import BetterTransformer
-from peft import PeftConfig
+from peft import PeftConfig, prepare_model_for_kbit_training
 from transformers import (  # noqa: F401
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
+    GPTQConfig,
     LlamaConfig,
     PreTrainedModel,
     PreTrainedTokenizerBase,
@@ -155,32 +155,15 @@ def load_model(
         LOG.info("patching _expand_mask")
         hijack_expand_mask()
 
-    try:
-        if cfg.gptq:
-            from alpaca_lora_4bit.monkeypatch.peft_tuners_lora_monkey_patch import (
-                replace_peft_model_with_int4_lora_model,
-            )
-
-            replace_peft_model_with_int4_lora_model()
-    except Exception as err:
-        LOG.exception(err)
-        raise err
-
-    if not cfg.gptq and (
-        (cfg.adapter == "lora" and load_in_8bit)
-        or (cfg.adapter == "qlora" and cfg.load_in_4bit)
-    ):
-        try:
-            from peft import prepare_model_for_kbit_training
-        except ImportError:
-            # For backward compatibility
-            from peft import (
-                prepare_model_for_int8_training as prepare_model_for_kbit_training,
-            )
-
     model_kwargs = {}
     if cfg.model_revision:
         model_kwargs["revision"] = cfg.model_revision
+    if cfg.gptq:
+        # TODO we should figure out how read the models config.json first
+        model_kwargs["quantization_config"] = GPTQConfig(
+            bits=cfg.gptq_bits,
+            disable_exllama=True,
+        )
     if cfg.adapter == "qlora" and cfg.load_in_4bit:
         model_kwargs["quantization_config"] = BitsAndBytesConfig(
             load_in_4bit=True,
@@ -191,45 +174,7 @@ def load_model(
             bnb_4bit_quant_type="nf4",
         )
     try:
-        if cfg.gptq and cfg.is_llama_derived_model:
+        if cfg.is_llama_derived_model and not cfg.trust_remote_code and not cfg.gptq:
-            from alpaca_lora_4bit.autograd_4bit import load_llama_model_4bit_low_ram
-            from huggingface_hub import snapshot_download
-
-            try:
-                snapshot_download_kwargs = {}
-                if cfg.base_model_ignore_patterns:
-                    snapshot_download_kwargs[
-                        "ignore_patterns"
-                    ] = cfg.base_model_ignore_patterns
-                cache_model_path = Path(
-                    snapshot_download(base_model, **snapshot_download_kwargs)
-                )
-                files = (
-                    list(cache_model_path.glob("*.pt"))
-                    + list(cache_model_path.glob("*.safetensors"))
-                    + list(cache_model_path.glob("*.bin"))
-                )
-                if len(files) > 0:
-                    model_path = str(files[0])
-                else:
-                    LOG.warning(
-                        "unable to find a cached model file, this will likely fail..."
-                    )
-                    model_path = str(cache_model_path)
-            except Exception:  # pylint: disable=broad-exception-caught
-                model_path = cfg.base_model
-            model, _ = load_llama_model_4bit_low_ram(
-                base_model_config if base_model_config else base_model,
-                model_path,
-                device_map=cfg.device_map,
-                half=cfg.fp16,
-                groupsize=cfg.gptq_groupsize if cfg.gptq_groupsize else -1,
-                is_v1_model=cfg.gptq_model_v1
-                if cfg.gptq_model_v1 is not None
-                else True,
-            )
-            load_in_8bit = False
-        elif cfg.is_llama_derived_model and not cfg.trust_remote_code:
             from transformers import LlamaForCausalLM
 
             config_kwargs = {}
@@ -275,15 +220,24 @@ def load_model(
         #     )
         #     model.train() # sets to train instead of eval mode
         elif model_type and not cfg.trust_remote_code:
-            model = getattr(transformers, model_type).from_pretrained(
+            if cfg.gptq:
-                base_model,
+                model = AutoModelForCausalLM.from_pretrained(
-                device_map=cfg.device_map,
+                    base_model,
-                load_in_8bit=cfg.load_in_8bit and cfg.adapter is not None,
+                    device_map=cfg.device_map,
-                load_in_4bit=cfg.load_in_4bit and cfg.adapter is not None,
+                    torch_dtype=cfg.torch_dtype,
-                torch_dtype=cfg.torch_dtype,
+                    trust_remote_code=cfg.trust_remote_code or False,
-                trust_remote_code=cfg.trust_remote_code or False,
+                    **model_kwargs,
-                **model_kwargs,
+                )
-            )
+            else:
+                model = getattr(transformers, model_type).from_pretrained(
+                    base_model,
+                    device_map=cfg.device_map,
+                    load_in_8bit=cfg.load_in_8bit and cfg.adapter is not None,
+                    load_in_4bit=cfg.load_in_4bit and cfg.adapter is not None,
+                    torch_dtype=cfg.torch_dtype,
+                    trust_remote_code=cfg.trust_remote_code or False,
+                    **model_kwargs,
+                )
         else:
             config = AutoConfig.from_pretrained(
                 base_model,
@@ -359,11 +313,12 @@ def load_model(
                 module.to(torch.float32)
 
     needs_fa2_dtype = cfg.adapter or cfg.fsdp
-    if not cfg.gptq and (
+    if (cfg.adapter == "lora" and load_in_8bit) or (
-        (cfg.adapter == "lora" and load_in_8bit)
+        cfg.adapter == "qlora" and cfg.load_in_4bit
-        or (cfg.adapter == "qlora" and cfg.load_in_4bit)
     ):
         LOG.info("converting PEFT model w/ prepare_model_for_kbit_training")
+        if cfg.gradient_checkpointing:
+            model.gradient_checkpointing_enable()
         model = prepare_model_for_kbit_training(
             model, use_gradient_checkpointing=cfg.gradient_checkpointing
         )
@@ -385,22 +340,10 @@ def load_model(
     if cfg.ddp and not load_in_8bit:
         model.to(f"cuda:{cfg.local_rank}")
 
-    if cfg.gptq:
-        # Scales to half
-        LOG.info("Fitting 4bit scales and zeros to half")
-        for _, module in model.named_modules():
-            if "Autograd4bitQuantLinear" in str(type(module)) or "Linear4bitLt" in str(
-                type(module)
-            ):
-                if hasattr(module, "is_v1_model") and module.is_v1_model:
-                    module.zeros = module.zeros.half()
-                module.scales = module.scales.half()
-                module.bias = module.bias.half()
-
     if (
         torch.cuda.device_count() > 1
         and int(os.getenv("WORLD_SIZE", "1")) > 1
-        and (cfg.gptq or cfg.load_in_4bit)
+        and (cfg.load_in_4bit)
     ):
         # llama is PROBABLY model parallelizable, but the default isn't that it is
         # so let's only set it for the 4bit, see

src/axolotl/utils/trainer.py

@@ -12,15 +12,9 @@ from typing import Optional, Union
 
 import numpy as np
 import torch.cuda
-import transformers
 from datasets import Dataset, set_caching_enabled
 from torch.optim.lr_scheduler import OneCycleLR
-from torch.utils.data import (
+from torch.utils.data import DataLoader, DistributedSampler, RandomSampler
-    DataLoader,
-    DistributedSampler,
-    RandomSampler,
-    SequentialSampler,
-)
 from transformers import EarlyStoppingCallback, Trainer, TrainingArguments
 from transformers.trainer_pt_utils import SequentialDistributedSampler
 
@@ -29,7 +23,6 @@ from axolotl.utils.callbacks import (
     GPUStatsCallback,
     SaveBetterTransformerModelCallback,
     SavePeftModelCallback,
-    bench_eval_callback_factory,
 )
 from axolotl.utils.collators import DataCollatorForSeq2Seq
 from axolotl.utils.dataloader import MultipackDistributedDataloader
@@ -134,27 +127,6 @@ class AxolotlTrainingArguments(TrainingArguments):
         default=None,
         metadata={"help": "how many warmup steps to take after reset for ReLoRA"},
     )
-    bench_split: Optional[str] = field(
-        default="eval", metadata={"help": "The benchmark split to run on"}
-    )
-    bench_dataset: Optional[str] = field(
-        default="pharaouk/dharma-1/dharma_1_mini.json",
-        metadata={
-            "help": "Benchmark dataset to use: options are `mmlu-zs`, `mmlu-fs`, or the full path to the dataset file"
-        },
-    )
-    do_bench_eval: Optional[bool] = field(
-        default=False, metadata={"help": "Whether to run the Benchmark evaluation."}
-    )
-    max_bench_samples: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": "If set, only evaluates on `max_bench_samples` of the benchmark dataset."
-        },
-    )
-    bench_source_max_len: int = field(
-        default=2048, metadata={"help": "Maximum source sequence length for bench."}
-    )
 
 
 class AxolotlTrainer(Trainer):
@@ -164,10 +136,6 @@ class AxolotlTrainer(Trainer):
 
     args = None  # type: AxolotlTrainingArguments
 
-    def __init__(self, *args, bench_data_collator=None, **kwargs):
-        self.bench_data_collator = bench_data_collator
-        super().__init__(*args, **kwargs)
-
     def create_scheduler(
         self, num_training_steps: int, optimizer: torch.optim.Optimizer = None
     ):
@@ -258,31 +226,6 @@ class AxolotlTrainer(Trainer):
             )
         return super().get_eval_dataloader(eval_dataset)
 
-    def _get_bench_sampler(
-        self, bench_dataset: Dataset
-    ) -> Optional[torch.utils.data.Sampler]:
-        if self.args.world_size <= 1:
-            return SequentialSampler(bench_dataset)
-        return None
-
-    def get_bench_dataloader(
-        self,
-        bench_dataset: Dataset,
-    ) -> Union[DataLoader, MultipackDistributedDataloader]:
-        dataloader_params = {
-            "batch_size": self.args.eval_batch_size,
-            "collate_fn": self.bench_data_collator,
-            "num_workers": self.args.dataloader_num_workers,
-            "pin_memory": self.args.dataloader_pin_memory,
-        }
-
-        if not isinstance(bench_dataset, torch.utils.data.IterableDataset):
-            dataloader_params["sampler"] = self._get_bench_sampler(bench_dataset)
-            dataloader_params["drop_last"] = self.args.dataloader_drop_last
-
-        return DataLoader(bench_dataset, **dataloader_params)
-        # return self.accelerator.prepare(DataLoader(bench_dataset, **dataloader_params))
-
     def compute_loss(self, model, inputs, return_outputs=False):
         # use one's weighted cross entropy loss calc
         # if self.args.sample_packing:
@@ -471,9 +414,6 @@ def setup_fsdp_envs(cfg):
         os.environ[
             "FSDP_TRANSFORMER_CLS_TO_WRAP"
         ] = cfg.fsdp_config.fsdp_transformer_layer_cls_to_wrap
-    from axolotl.monkeypatch.fsdp import replace_fsdp_state_dict_type
-
-    replace_fsdp_state_dict_type()
 
 
 def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps):
@@ -507,23 +447,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_
         training_arguments_kwargs["seed"] = cfg.seed
 
     if cfg.gradient_checkpointing:
-        if cfg.gptq:
+        training_arguments_kwargs["gradient_checkpointing"] = cfg.gradient_checkpointing
-            from alpaca_lora_4bit.gradient_checkpointing import (
-                apply_gradient_checkpointing,
-            )
-
-            gradient_checkpointing_ratio = (
-                cfg.gradient_checkpointing_ratio
-                if cfg.gradient_checkpointing_ratio
-                else 1.0
-            )
-            apply_gradient_checkpointing(
-                model, checkpoint_ratio=gradient_checkpointing_ratio
-            )
-        else:
-            training_arguments_kwargs[
-                "gradient_checkpointing"
-            ] = cfg.gradient_checkpointing
     if cfg.fsdp:
         training_arguments_kwargs["fsdp"] = cfg.fsdp
         if cfg.fsdp_config:
@@ -577,11 +501,6 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_
             "steps" if cfg.save_steps else "epoch"
         )
 
-    if cfg.do_bench_eval:
-        training_arguments_kwargs["do_bench_eval"] = cfg.do_bench_eval
-        if cfg.bench_dataset:
-            training_arguments_kwargs["bench_dataset"] = cfg.bench_dataset
-
     training_args = AxolotlTrainingArguments(  # pylint: disable=unexpected-keyword-arg
         max_steps=total_num_steps if cfg.max_steps else -1,
         max_seq_length=cfg.sequence_len,
@@ -694,16 +613,8 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_
             return_tensors="pt",
             **data_collator_kwargs,
         ),
-        bench_data_collator=transformers.DataCollatorForSeq2Seq(
-            tokenizer,
-            return_tensors="pt",
-            **data_collator_kwargs,
-        ),
         callbacks=callbacks,
         **trainer_kwargs,
     )
 
-    if cfg.do_bench_eval:
-        trainer.add_callback(bench_eval_callback_factory(trainer, tokenizer))
-
     return trainer