more gpt-neox long ctx fixes

.github/workflows/base.yml

@@ -12,7 +12,6 @@ jobs:
     # this job needs to be run on self-hosted GPU runners...
     runs-on: self-hosted
     strategy:
-      fail-fast: false
       matrix:
         include:
           - cuda: "118"
@@ -26,7 +25,7 @@ jobs:
             pytorch: 2.0.0
             axolotl_extras:
           - cuda: "117"
-            cuda_version: 11.7.1
+            cuda_version: 11.7.0
             python_version: "3.9"
             pytorch: 1.13.1
             axolotl_extras:

.github/workflows/main.yml

@@ -11,7 +11,6 @@ jobs:
     if: github.repository_owner == 'OpenAccess-AI-Collective'
     # this job needs to be run on self-hosted GPU runners...
     strategy:
-      fail-fast: false
       matrix:
         include:
           - cuda: cu118
@@ -30,7 +29,7 @@ jobs:
             pytorch: 2.0.0
             axolotl_extras: gptq
           - cuda: cu117
-            cuda_version: 11.7.1
+            cuda_version: 11.7.0
             python_version: "3.9"
             pytorch: 1.13.1
             axolotl_extras:
@@ -85,7 +84,7 @@ jobs:
             pytorch: 2.0.0
             axolotl_extras: gptq
           - cuda: cu117
-            cuda_version: 11.7.1
+            cuda_version: 11.7.0
             python_version: "3.9"
             pytorch: 1.13.1
             axolotl_extras:

.github/workflows/tests.yml

@@ -7,7 +7,6 @@ jobs:
   test:
     runs-on: ubuntu-latest
     strategy:
-      fail-fast: false
       matrix:
         python_version: ["3.9", "3.10"]
     timeout-minutes: 10

.pre-commit-config.yaml

@@ -1,5 +1,5 @@
 default_language_version:
-    python: python3
+    python: python3.9
 
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks

FAQS.md

@@ -2,6 +2,3 @@
 
 - Can you train StableLM with this? Yes, but only with a single GPU atm. Multi GPU support is coming soon! Just waiting on this [PR](https://github.com/huggingface/transformers/pull/22874)
 - Will this work with Deepspeed? That's still a WIP, but setting `export ACCELERATE_USE_DEEPSPEED=true` should work in some cases
-- `Error invalid argument at line 359 in file /workspace/bitsandbytes/csrc/pythonInterface.c`
-`/arrow/cpp/src/arrow/filesystem/s3fs.cc:2598:  arrow::fs::FinalizeS3 was not called even though S3 was initialized.`
-This could lead to a segmentation fault at exit. Try reinstalling bitsandbytes and transformers from source.

README.md

@@ -16,14 +16,13 @@
 
 ## Axolotl supports
 
-|          | fp16/fp32 | lora | qlora | gptq | gptq w/ lora | gptq w/flash attn | flash attn | xformers attn |
-|----------|:----------|:-----|-------|------|:-------------|-------------------|------------|---------------|
-| llama    | ✅         | ✅    | ✅     | ✅    | ✅             | ✅                 | ✅          | ✅             |
-| Pythia   | ✅         | ✅    | ✅     | ❌    | ❓            | ❌                 | ❌          | ❓             |
-| cerebras | ✅         | ✅    | ✅     | ❌    | ❓            | ❌                 | ❌          | ✅             |
-| mpt      | ✅         | ❌    | ❓     | ❌    | ❓            | ❌                 | ❌          | ❓             |
-| falcon   | ✅         | ✅    | ✅     | ❌    | ❓            | ❌                 | ❌          | ✅             |
-| gpt-j    | ✅         | ✅    | ✅     | ❌    | ❓            | ❌                 | ❓          | ✅             |
+|         | fp16/fp32 | fp16/fp32 w/ lora | qlora | 4bit-quant | 4bit-quant w/flash attention | flash attention | xformers attention |
+|---------|:----------|:------------------|------|------------|------------------------------|-----------------|--------------------|
+| llama   | ✅         | ✅                 | ✅  | ✅          | ✅                            | ✅               | ✅                  |
+| Pythia  | ✅         | ✅                 | ❓  | ❌          | ❌                            | ❌               | ❓                  |
+| cerebras | ✅         | ✅                 | ❓  | ❌          | ❌                            | ❌               | ❓                  |
+| mpt     | ✅         | ❌                 | ❓  | ❌          | ❌                            | ❌               | ❓                  |
+| falcon  | ✅         | ❌                 | ❌  | ❌          | ❌                            | ❌               | ❓                  |
 
 
 ## Quickstart ⚡
@@ -34,15 +33,14 @@
 git clone https://github.com/OpenAccess-AI-Collective/axolotl
 
 pip3 install -e .
-pip3 install -U git+https://github.com/huggingface/peft.git
 
 accelerate config
 
 # finetune lora
-accelerate launch scripts/finetune.py examples/openllama-3b/lora.yml
+accelerate launch scripts/finetune.py examples/lora-openllama-3b/config.yml
 
 # inference
-accelerate launch scripts/finetune.py examples/openllama-3b/lora.yml \
+accelerate launch scripts/finetune.py examples/lora-openllama-3b/config.yml \
     --inference --lora_model_dir="./lora-out"
 ```
 
@@ -52,17 +50,10 @@ accelerate launch scripts/finetune.py examples/openllama-3b/lora.yml \
 
 - Docker
   ```bash
-  docker run --gpus '"all"' --rm -it winglian/axolotl:main-py3.9-cu118-2.0.0
-  ```
-  - `winglian/axolotl-runpod:main-py3.9-cu118-2.0.0`: for runpod
-  - `winglian/axolotl-runpod:main-py3.9-cu118-2.0.0-gptq`: for gptq
-  - `winglian/axolotl:dev`: dev branch (not usually up to date)
-
-  Or run on the current files for development:
-
-  ```sh
-  docker compose up -d
+  docker run --gpus '"all"' --rm -it winglian/axolotl:main
   ```
+  - `winglian/axolotl:dev`: dev branch
+  - `winglian/axolotl-runpod:main`: for runpod
 
 - Conda/Pip venv
   1. Install python **3.9**
@@ -70,65 +61,9 @@ accelerate launch scripts/finetune.py examples/openllama-3b/lora.yml \
   2. Install pytorch stable https://pytorch.org/get-started/locally/
 
   3. Install python dependencies with ONE of the following:
-      - Recommended, supports QLoRA, NO gptq/int4 support
-        ```bash
-        pip3 install -e .
-        pip3 install -U git+https://github.com/huggingface/peft.git
-        ```
-      - gptq/int4 support, NO QLoRA
-        ```bash
-        pip3 install -e .[gptq]
-        ```
-      - same as above but not recommended
-        ```bash
-        pip3 install -e .[gptq_triton]
-        ```
-
-- LambdaLabs
-  <details>
-
-  <summary>Click to Expand</summary>
-
-  1. Install python
-  ```bash
-  sudo apt update
-  sudo apt install -y python3.9
-
-  sudo update-alternatives --install /usr/bin/python python /usr/bin/python3.9 1
-  sudo update-alternatives --config python # pick 3.9 if given option
-  python -V # should be 3.9
-
-  ```
-
-  2. Install pip
-  ```bash
-  wget https://bootstrap.pypa.io/get-pip.py
-  python get-pip.py
-  ```
-
-  3. Install torch
-  ```bash
-  pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
-  ```
-
-  4. Axolotl
-  ```bash
-  git clone https://github.com/OpenAccess-AI-Collective/axolotl
-  cd axolotl
-
-  pip3 install -e . # change depend on needs
-  pip3 install protobuf==3.20.3
-  pip3 install -U requests
-  pip3 install -U --ignore-installed psutil
-  pip3 install -U scipy
-  pip3 install git+https://github.com/huggingface/peft.git # not for gptq
-  ```
-
-  5. Set path
-  ```bash
-  export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH
-  ```
-  </details>
+      - `pip3 install -e .` (recommended, supports QLoRA, no gptq/int4 support)
+      - `pip3 install -e .[gptq]` (next best if you don't need QLoRA, but want to use gptq)
+      - `pip3 install -e .[gptq_triton]`
 
 ### Dataset
 
@@ -138,7 +73,7 @@ Have dataset(s) in one of the following format (JSONL recommended):
   ```json
   {"instruction": "...", "input": "...", "output": "..."}
   ```
-- `sharegpt:chat`: conversations
+- `sharegpt`: conversations
   ```json
   {"conversations": [{"from": "...", "value": "..."}]}
   ```
@@ -179,66 +114,13 @@ Have dataset(s) in one of the following format (JSONL recommended):
   ```json
   {"article": "...", "summary": "..."}
   ```
-- `alpaca_chat`: basic instruct for alpaca chat
-  ```json
-  {"instruction": "...", "input": "...", "response": "..."}
-  ```
-- `alpaca_chat.load_qa`: question and answer for alpaca chat
-  ```json
-  {"question": "...", "answer": "..."}
-  ```
-- `alpaca_chat.load_concise`: question and answer for alpaca chat, for concise answers
-  ```json
-  {"instruction": "...", "input": "...", "response": "..."}
-  ```
-- `alpaca_chat.load_camel_ai`: question and answer for alpaca chat, for load_camel_ai
-  ```json
-  {"message_1": "...", "message_2": "..."}
-  ```
-- `context_qa`: in context question answering from an article
-  ```json
-  {"article": "...", "question": "...", "answer": "..."}
-  ```
-- `context_qa.load_404`: in context question answering from an article, with default response for no answer from context
-  ```json
-  {"article": "...", "unanswerable_question": "..."}
-  ```
-- `creative_acr.load_answer`: instruction and revision
-  ```json
-  {"instruction": "...", "revision": "..."}
-  ```
-- `creative_acr.load_critique`: critique
-  ```json
-  {"scores": "...", "critiques": "...", "instruction": "...", "answer": "..."}
-  ```
-- `creative_acr.load_revise`: critique and revise
-  ```json
-  {"scores": "...", "critiques": "...", "instruction": "...", "answer": "...", "revision": "..."}
-  ```
-- `pygmalion`: pygmalion
-  ```json
-  {"conversations": [{"role": "...", "value": "..."}]}
-  ```
-- `sharegpt_simple.load_role`: conversations where `role` is used instead of `from`
-  ```json
-  {"conversations": [{"role": "...", "value": "..."}]}
-  ```
-- `sharegpt_jokes`: creates a chat where bot is asked to tell a joke, then explain why the joke is funny
-  ```json
-  {"conversations": [{"title": "...", "text": "...", "explanation": "..."}]}
-  ```
+
+> Have some new format to propose? Check if it's already defined in [data.py](src/axolotl/utils/data.py) in `dev` branch!
 
 </details>
 
-#### How to add custom prompts
-
-  1. Add your method to a file in [prompt_strategies](src/axolotl/prompt_strategies). Please see other files as example.
-  2. Use your custom file name as the dataset type.
-
 Optionally, download some datasets, see [data/README.md](data/README.md)
 
-
-
 ### Config
 
 See sample configs in [configs](configs) folder or [examples](examples) for quick start. It is recommended to duplicate and modify to your needs. The most important options are:
@@ -264,8 +146,6 @@ See sample configs in [configs](configs) folder or [examples](examples) for quic
   bf16: true # require >=ampere
   fp16: true
   tf32: true # require >=ampere
-  bfloat16: true # require >=ampere, use instead of bf16 when you don't want AMP (automatic mixed precision)
-  float16: true # use instead of fp16 when you don't want AMP
   ```
   Note: Repo does not do 4-bit quantization.
 
@@ -302,8 +182,6 @@ model_type: AutoModelForCausalLM
 tokenizer_type: AutoTokenizer
 # Trust remote code for untrusted source
 trust_remote_code:
-# use_fast option for tokenizer loading from_pretrained, default to True
-tokenizer_use_fast:
 
 # whether you are training a 4-bit GPTQ quantized model
 gptq: true
@@ -394,15 +272,13 @@ num_epochs: 3
 warmup_steps: 100
 learning_rate: 0.00003
 logging_steps:
-save_steps:
-eval_steps:
 
 # whether to mask out or include the human's prompt from the training labels
 train_on_inputs: false
 # don't use this, leads to wonky training (according to someone on the internet)
 group_by_length: false
 
-# Whether to use gradient checkpointing https://huggingface.co/docs/transformers/v4.18.0/en/performance#gradient-checkpointing
+# does not work with current implementation of 4-bit LoRA
 gradient_checkpointing: false
 
 # stop training after this many evaluation losses have increased in a row
@@ -424,15 +300,7 @@ log_sweep_max_lr:
 optimizer:
 # specify weight decay
 weight_decay:
-# adamw hyperparams
-adam_beta1:
-adam_beta2:
-adam_epsilon:
-# Gradient clipping max norm
-max_grad_norm:
 
-# whether to bettertransformers
-flash_optimum:
 # whether to use xformers attention patch https://github.com/facebookresearch/xformers:
 xformers_attention:
 # whether to use flash attention patch https://github.com/HazyResearch/flash-attention:
@@ -440,11 +308,6 @@ flash_attention:  # require a100 for llama
 # whether to use scaled-dot-product attention
 # https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
 sdp_attention:
-# Landmark attention (only llama)
-landmark_attention:
-# xpos RoPE see https://github.com/kaiokendev/cutoff-len-is-context-len/blob/main/util/xpos_rope_llama_monkey_patch.py
-# llama only
-xpos_rope:
 
 # resume from a specific checkpoint dir
 resume_from_checkpoint:
@@ -512,16 +375,11 @@ Pass the appropriate flag to the train command:
 
 - Pretrained LORA:
   ```bash
-  --inference --lora_model_dir="./lora-output-dir"
+  --inference --lora_model_dir ./completed-model
   ```
 - Full weights finetune:
   ```bash
-  --inference --base_model="./completed-model"
-  ```
-- Full weights finetune w/ a prompt from a text file:
-  ```bash
-  cat /tmp/prompt.txt | python scripts/finetune.py configs/your_config.yml \
-    --base_model="./completed-model" --inference --prompter=None --load_in_8bit=True
+  --inference --base_model ./completed-model
   ```
 
 ### Merge LORA to base
@@ -532,12 +390,6 @@ Add below flag to train command above
 --merge_lora --lora_model_dir="./completed-model" --load_in_8bit=False --load_in_4bit=False
 ```
 
-If you run out of CUDA memory, you can try to merge in system RAM with
-
-```bash
-CUDA_VISIBLE_DEVICES="" python3 scripts/finetune.py ...
-```
-
 ## Common Errors 🧰
 
 > Cuda out of memory
@@ -545,7 +397,6 @@ CUDA_VISIBLE_DEVICES="" python3 scripts/finetune.py ...
 Please reduce any below
   - `micro_batch_size`
   - `eval_batch_size`
-  - `gradient_accumulation_steps`
   - `sequence_len`
 
 > RuntimeError: expected scalar type Float but found Half
@@ -556,7 +407,7 @@ Try set `fp16: true`
 
 Try to turn off xformers.
 
-## Need help? 🙋♂️
+## Need help? 🙋‍♂️
 
 Join our [Discord server](https://discord.gg/HhrNrHJPRb) where we can help you
 
@@ -570,16 +421,6 @@ Building something cool with Axolotl? Consider adding a badge to your model card
 
 [<img src="https://raw.githubusercontent.com/OpenAccess-AI-Collective/axolotl/main/image/axolotl-badge-web.png" alt="Built with Axolotl" width="200" height="32"/>](https://github.com/OpenAccess-AI-Collective/axolotl)
 
-## Community Showcase
-
-Open Access AI Collective
-- [Minotaur 13b](https://huggingface.co/openaccess-ai-collective/minotaur-13b)
-- [Manticore 13b](https://huggingface.co/openaccess-ai-collective/manticore-13b)
-- [Hippogriff 30b](https://huggingface.co/openaccess-ai-collective/hippogriff-30b-chat)
-
-PocketDoc Labs
-- [Dan's PersonalityEngine 13b LoRA](https://huggingface.co/PocketDoc/Dans-PersonalityEngine-13b-LoRA)
-
 ## Contributing 🤝
 
 Bugs? Please check for open issue else create a new [Issue](https://github.com/OpenAccess-AI-Collective/axolotl/issues/new).

configs/accelerate/default_config.yaml

@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: 'NO'
+downcast_bf16: 'no'
+gpu_ids: all
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

configs/cerebras_1_3B_alpaca.yml

@@ -0,0 +1,40 @@
+base_model: cerebras/Cerebras-GPT-1.3B
+model_type: AutoModelForCausalLM
+tokenizer_type: AutoTokenizer
+load_in_8bit: true
+datasets:
+  - path: data/alpaca_data_gpt4.jsonl
+    type: alpaca
+  - path: data/vicuna_cleaned.jsonl
+    type: sharegpt
+  - path: data/gpt4-instruct-similarity-0.6-dataset.jsonl
+    type: gpteacher
+  - path: data/roleplay-similarity_0.6-instruct-dataset.jsonl
+    type: gpteacher
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.05
+adapter: lora
+sequence_len: 2048
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules:
+  - c_attn
+lora_fan_in_fan_out: false
+wandb_project: pythia-1.4b-lora
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./lora-alpaca
+gradient_accumulation_steps: 1
+micro_batch_size: 4
+num_epochs: 5
+learning_rate: 0.0003
+train_on_inputs: false
+group_by_length: false
+bf16: True
+tf32: True
+gradient_checkpointing:
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:

configs/galactica_1_3B.yml

@@ -0,0 +1,41 @@
+base_model: facebook/galactica-1.3b
+model_type: AutoModelForCausalLM
+tokenizer_type: AutoTokenizer
+load_in_8bit: false
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+adapter:
+lora_model_dir:
+sequence_len: 1024
+max_packed_sequence_len: 1024
+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:
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./lora-llama-alpaca
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+num_epochs: 3
+learning_rate: 0.00003
+train_on_inputs: false
+group_by_length: false
+bf16: false
+tf32: false
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:
+tokens:
+  pad_token: "[PAD]"
+  bos_token: "<s>"
+  eos_token: "</s>"
+  unk_token: "<unk>"

configs/llama_13B_alpaca.yml

@@ -0,0 +1,39 @@
+base_model: huggyllama/llama-13b
+model_type: LlamaForCausalLM
+tokenizer_type: LlamaTokenizer
+load_in_8bit: true
+datasets:
+  - path: anon8231489123/ShareGPT_Vicuna_unfiltered
+    data_files: ShareGPT_V3_unfiltered_cleaned_split_no_imsorry.json
+    type: sharegpt
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.002
+adapter:
+lora_model_dir:
+sequence_len: 2048
+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:
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./llama-13b-sharegpt
+gradient_accumulation_steps: 1
+micro_batch_size: 2
+warmup_steps: 1000
+save_steps:
+eval_steps:
+num_epochs: 5
+learning_rate: 0.00003
+train_on_inputs: false
+group_by_length: false
+bf16: true
+tf32: true
+early_stopping_patience: 5
+resume_from_checkpoint:
+local_rank:

configs/llama_65B_alpaca.yml

@@ -0,0 +1,44 @@
+base_model: huggyllama/llama-65b
+model_type: LlamaForCausalLM
+tokenizer_type: LlamaTokenizer
+load_in_8bit: true
+datasets:
+  - path: data/alpaca_data_gpt4.jsonl
+    type: alpaca
+  - path: anon8231489123/ShareGPT_Vicuna_unfiltered
+    data_files: ShareGPT_V3_unfiltered_cleaned_split_no_imsorry.json
+    type: sharegpt
+  - path: data/gpt4-instruct-similarity-0.6-dataset.jsonl
+    type: gpteacher
+  - path: data/roleplay-similarity_0.6-instruct-dataset.jsonl
+    type: gpteacher
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.04
+adapter: lora
+lora_model_dir:
+sequence_len: 2048
+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-65b-lora
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./lora-llama-alpaca
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+warmup_steps: 1000
+save_steps:
+num_epochs: 5
+learning_rate: 0.00003
+train_on_inputs: false
+group_by_length: false
+bf16: true
+tf32: true
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:

configs/llama_7B_4bit.yml

@@ -0,0 +1,45 @@
+base_model: decapoda-research/llama-7b-hf-int4
+base_model_config: decapoda-research/llama-7b-hf
+model_type: LlamaForCausalLM
+tokenizer_type: LlamaTokenizer
+load_in_8bit: true
+datasets:
+  - path: tatsu-lab/alpaca  # original alpaca dataset
+    type: alpaca
+dataset_prepared_path: data/last_run_prepared
+val_set_size: 0.04
+adapter: lora
+lora_model_dir:
+sequence_len: 2048
+max_packed_sequence_len: 1024
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules:
+  - q_proj
+  - v_proj
+#  - k_proj
+#  - o_proj
+lora_fan_in_fan_out: false
+wandb_project:
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./lora-test
+gradient_accumulation_steps: 1
+micro_batch_size: 2
+num_epochs: 3
+warmup_steps: 100
+learning_rate: 0.00003
+train_on_inputs: false
+group_by_length: false
+bf16: true
+tf32: true
+gradient_checkpointing: false
+early_stopping_patience: 3
+resume_from_checkpoint:
+auto_resume_from_checkpoints: true
+local_rank:
+load_4bit: true
+xformers_attention: true
+flash_attention:

configs/llama_7B_alpaca.yml

@@ -0,0 +1,41 @@
+base_model: huggyllama/llama-7b
+model_type: LlamaForCausalLM
+tokenizer_type: LlamaTokenizer
+load_in_8bit: true
+datasets:
+  - path: data/alpaca_data_gpt4.jsonl
+    type: alpaca
+  - path: data/vicuna_cleaned.jsonl
+    type: sharegpt
+  - path: data/gpt4-instruct-similarity-0.6-dataset.jsonl
+    type: gpteacher
+  - path: data/roleplay-similarity_0.6-instruct-dataset.jsonl
+    type: gpteacher
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.04
+adapter: lora
+lora_model_dir:
+sequence_len: 2048
+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
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./lora-llama-alpaca
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+num_epochs: 5
+learning_rate: 0.00003
+train_on_inputs: false
+group_by_length: false
+bf16: true
+tf32: true
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:

configs/llama_7B_jeopardy.yml

@@ -7,28 +7,30 @@ datasets:
   - path: openaccess-ai-collective/jeopardy
     type: jeopardy
 dataset_prepared_path: last_run_prepared
-val_set_size: 0.02
+val_set_size: 0.01
 adapter:
 lora_model_dir:
-sequence_len: 512
-max_packed_sequence_len:
-lora_r:
-lora_alpha:
-lora_dropout:
+sequence_len: 2048
+max_packed_sequence_len: 2048
+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:
+wandb_project: jeopardy-bot-7b
 wandb_watch:
 wandb_run_id:
 wandb_log_model:
 output_dir: ./jeopardy-bot-7b
-gradient_accumulation_steps: 1
+gradient_accumulation_steps: 2
 micro_batch_size: 1
-num_epochs: 3
+num_epochs: 2
 optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine
-learning_rate: 0.00003
+learning_rate: 0.0000002
 train_on_inputs: false
 group_by_length: false
 bf16: true
@@ -46,10 +48,11 @@ eval_steps: 110
 save_steps: 660
 debug:
 deepspeed:
-weight_decay: 0.1
+weight_decay: 0.0001
 fsdp:
 fsdp_config:
 tokens:
+  pad_token: "[PAD]"
   bos_token: "<s>"
   eos_token: "</s>"
   unk_token: "<unk>"

configs/pythia_1_2B_alpaca.yml

@@ -1,29 +1,36 @@
 base_model: EleutherAI/pythia-1.4b-deduped
-base_model_config: EleutherAI/pythia-1.4b-deduped
+model_type: GPTNeoXForCausalLM
+tokenizer_type: AutoTokenizer
 load_in_8bit: true
 datasets:
-  - path: teknium/GPT4-LLM-Cleaned
+  - path: data/alpaca_data_gpt4.jsonl
     type: alpaca
+  - path: data/vicuna_cleaned.jsonl
+    type: sharegpt
+  - path: data/gpt4-instruct-similarity-0.6-dataset.jsonl
+    type: gpteacher
+  - path: data/roleplay-similarity_0.6-instruct-dataset.jsonl
+    type: gpteacher
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
 adapter: lora
 lora_model_dir:
-sequence_len: 512
-lora_r: 16
+sequence_len: 2048
+lora_r: 8
 lora_alpha: 32
 lora_dropout: 0.05
 lora_target_modules:
   - query_key_value
-lora_target_linear:
+#  - xxx
 lora_fan_in_fan_out: true  # pythia/GPTNeoX lora specific
-wandb_project:
+wandb_project: pythia-1.4b-lora
 wandb_watch:
 wandb_run_id:
 wandb_log_model:
-output_dir: ./lora-alpaca-pythia
+output_dir: ./lora-alpaca
 gradient_accumulation_steps: 1
 micro_batch_size: 4
-num_epochs: 3
+num_epochs: 5
 learning_rate: 0.00001
 train_on_inputs: false
 group_by_length: false
@@ -32,6 +39,3 @@ tf32: True
 early_stopping_patience:
 resume_from_checkpoint:
 local_rank:
-weight_decay: 0.1
-eval_steps: 20
-logging_steps: 1

configs/quickstart.yml

@@ -0,0 +1,45 @@
+base_model: decapoda-research/llama-7b-hf-int4
+base_model_config: decapoda-research/llama-7b-hf
+model_type: LlamaForCausalLM
+tokenizer_type: LlamaTokenizer
+load_in_8bit: true
+datasets:
+  - path: tatsu-lab/alpaca  # original alpaca dataset
+    type: alpaca
+dataset_prepared_path: data/last_run_prepared
+val_set_size: 0.04
+adapter: lora
+lora_model_dir:
+sequence_len: 1024
+max_packed_sequence_len: 1024
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules:
+  - q_proj
+  - v_proj
+#  - k_proj
+#  - o_proj
+lora_fan_in_fan_out: false
+wandb_project:
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./lora-test
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 3
+warmup_steps: 100
+learning_rate: 0.00003
+train_on_inputs: false
+group_by_length: false
+bf16: true
+tf32: true
+gradient_checkpointing: false
+early_stopping_patience: 3
+resume_from_checkpoint:
+auto_resume_from_checkpoints: true
+local_rank:
+gptq: true
+xformers_attention: true
+flash_attention:

configs/sample.yml

@@ -0,0 +1,87 @@
+# this is the huggingface model that contains *.pt, *.safetensors, or *.bin files
+# this can also be a relative path to a model on disk
+base_model: decapoda-research/llama-7b-hf-int4
+# you can specify an ignore pattern if the model repo contains more than 1 model type (*.pt, etc)
+base_model_ignore_patterns:
+# if the base_model repo on hf hub doesn't include configuration .json files,
+# you can set that here, or leave this empty to default to base_model
+base_model_config: decapoda-research/llama-7b-hf
+# If you want to specify the type of model to load, AutoModelForCausalLM is a good choice too
+model_type: AutoModelForCausalLM
+# Corresponding tokenizer for the model AutoTokenizer is a good choice
+tokenizer_type: AutoTokenizer
+# whether you are training a 4-bit quantized model
+load_4bit: true
+# this will attempt to quantize the model down to 8 bits and use adam 8 bit optimizer
+load_in_8bit: true
+# a list of one or more datasets to finetune the model with
+datasets:
+  # this can be either a hf dataset, or relative path
+  - path: vicgalle/alpaca-gpt4
+  # The type of prompt to use for training. [alpaca, sharegpt, gpteacher, oasst, reflection]
+    type: alpaca
+# axolotl attempts to save the dataset as an arrow after packing the data together so
+# subsequent training attempts load faster, relative path
+dataset_prepared_path: data/last_run_prepared
+# How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc
+val_set_size: 0.04
+# if you want to use lora, leave blank to train all parameters in original model
+adapter: lora
+# if you already have a lora model trained that you want to load, put that here
+lora_model_dir:
+# the maximum length of an input to train with, this should typically be less than 2048
+# as most models have a token/context limit of 2048
+sequence_len: 2048
+# max sequence length to concatenate training samples together up to
+# inspired by StackLLaMA. see https://huggingface.co/blog/stackllama#supervised-fine-tuning
+max_packed_sequence_len: 1024
+# lora hyperparameters
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules:
+  - q_proj
+  - v_proj
+#  - k_proj
+#  - o_proj
+lora_fan_in_fan_out: false
+# wandb configuration if your're using it
+wandb_project:
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+# where to save the finsihed model to
+output_dir: ./completed-model
+# training hyperparameters
+gradient_accumulation_steps: 1
+batch_size:
+micro_batch_size: 2
+num_epochs: 3
+warmup_steps: 100
+learning_rate: 0.00003
+# whether to mask out or include the human's prompt from the training labels
+train_on_inputs: false
+# don't use this, leads to wonky training (according to someone on the internet)
+group_by_length: false
+# Use CUDA bf16
+bf16: true
+# Use CUDA tf32
+tf32: true
+# does not work with current implementation of 4-bit LoRA
+gradient_checkpointing: false
+# stop training after this many evaluation losses have increased in a row
+# https://huggingface.co/transformers/v4.2.2/_modules/transformers/trainer_callback.html#EarlyStoppingCallback
+early_stopping_patience: 3
+# specify a scheduler to use with the optimizer. only one_cycle is supported currently
+lr_scheduler:
+# whether to use xformers attention patch https://github.com/facebookresearch/xformers:
+xformers_attention:
+# whether to use flash attention patch https://github.com/HazyResearch/flash-attention:
+flash_attention:
+# resume from a specific checkpoint dir
+resume_from_checkpoint:
+# if resume_from_checkpoint isn't set and you simply want it to start where it left off
+# be careful with this being turned on between different models
+auto_resume_from_checkpoints: false
+# don't mess with this, it's here for accelerate and torchrun
+local_rank:

configs/stability_3b.yml

@@ -0,0 +1,56 @@
+base_model: stabilityai/stablelm-base-alpha-3b
+base_model_config: stabilityai/stablelm-base-alpha-3b
+load_in_8bit: false
+datasets:
+  - path: vicgalle/alpaca-gpt4
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.04
+adapter:
+lora_model_dir:
+sequence_len: 4096
+max_packed_sequence_len: 4096
+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: stable-alpaca-3b
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./stable-alpaca-3b
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+torchdistx_path:
+lr_scheduler: cosine
+learning_rate: 0.0000002
+train_on_inputs: false
+group_by_length: false
+bf16: true
+tf32: true
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:
+logging_steps: 1
+xformers_attention: true
+flash_attention:
+gptq_groupsize:
+gptq_model_v1:
+warmup_steps: 100
+eval_steps: 50
+save_steps: 200
+debug:
+deepspeed:
+weight_decay: 0.01
+fsdp:
+fsdp_config:
+#tokens:
+#  pad_token: "[PAD]"
+#  bos_token: "<s>"
+#  eos_token: "</s>"
+#  unk_token: "<unk>"

configs/vicuna_13B_4bit_reflect.yml

@@ -0,0 +1,45 @@
+base_model: anon8231489123/vicuna-13b-GPTQ-4bit-128g
+base_model_config: anon8231489123/vicuna-13b-GPTQ-4bit-128g
+model_type: LlamaForCausalLM
+tokenizer_type: LlamaTokenizer
+load_in_8bit: false
+load_4bit: true
+gptq_groupsize: 128
+gptq_model_v1: false
+datasets:
+# https://github.com/vaguenebula/AlpacaDataReflect/blob/main/alpaca_reflect_pruned.json
+  - path: data/alpaca_reflect_pruned.jsonl
+    type: reflection
+dataset_prepared_path: data/last_run_prepared
+val_set_size: 0.04
+adapter: lora
+lora_model_dir:
+sequence_len: 2048
+max_packed_sequence_len: 2048
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules:
+  - q_proj
+  - v_proj
+#  - k_proj
+#  - o_proj
+lora_fan_in_fan_out: false
+wandb_project:
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./lora-reflect
+gradient_accumulation_steps: 1
+micro_batch_size: 2
+num_epochs: 3
+learning_rate: 0.00003
+train_on_inputs: false
+group_by_length: false
+bf16: true
+tf32: true
+gradient_checkpointing: false
+early_stopping_patience: 3
+resume_from_checkpoint:
+local_rank:
+flash_attention: true

data/README.md

@@ -10,10 +10,10 @@ curl https://github.com/teknium1/GPTeacher/blob/main/Roleplay/roleplay-similarit
 ## Convert the JSON data files to JSONL.
 
 ```shell
-python3 ./scripts/alpaca_json_to_jsonl.py --file data/alpaca_data_gpt4.json --output data/alpaca_data_gpt4.jsonl
-python3 ./scripts/alpaca_json_to_jsonl.py --file data/raw/vicuna_cleaned.json --output data/vicuna_cleaned.jsonl
-python3 ./scripts/alpaca_json_to_jsonl.py --file data/raw/roleplay-similarity_0.6-instruct-dataset.json --output data/roleplay-similarity_0.6-instruct-dataset.jsonl
-python3 ./scripts/alpaca_json_to_jsonl.py --file data/raw/gpt4-instruct-similarity-0.6-dataset.json --output data/gpt4-instruct-similarity-0.6-dataset.jsonl
+python3 ./scripts/alpaca_json_to_jsonl.py --input data/alpaca_data_gpt4.json > data/alpaca_data_gpt4.jsonl
+python3 ./scripts/alpaca_json_to_jsonl.py --input data/raw/vicuna_cleaned.json > data/vicuna_cleaned.jsonl
+python3 ./scripts/alpaca_json_to_jsonl.py --input data/raw/roleplay-similarity_0.6-instruct-dataset.json > data/roleplay-similarity_0.6-instruct-dataset.jsonl
+python3 ./scripts/alpaca_json_to_jsonl.py --input data/raw/gpt4-instruct-similarity-0.6-dataset.json > data/gpt4-instruct-similarity-0.6-dataset.jsonl
 ```
 ---
 

docker-compose.yaml

@@ -1,20 +0,0 @@
-# version: '3.8'
-services:
-  axolotl:
-    build:
-      context: .
-      dockerfile: ./docker/Dockerfile
-    volumes:
-      - .:/workspace/axolotl
-      - ~/.cache/huggingface/:/root/.cache/huggingface/
-    # set environment variables
-    environment:
-      - WANDB_API_KEY=${WANDB_API_KEY}
-    deploy:
-      resources:
-        reservations:
-          devices:
-            - driver: nvidia
-              # count: 1
-              capabilities: [gpu]
-    command: tail -f /dev/null

docker/Dockerfile

@@ -13,7 +13,8 @@ RUN pip3 install --force-reinstall "peft @ git+https://github.com/huggingface/pe
             "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \
             "transformers @ git+https://github.com/huggingface/transformers.git@main"
 
-RUN git clone --depth=1 https://github.com/OpenAccess-AI-Collective/axolotl.git
+RUN mkdir axolotl
+COPY . axolotl/
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN cd axolotl && \
     if [ "$AXOLOTL_EXTRAS" != "" ] ; then \

docker/Dockerfile-base

@@ -77,7 +77,7 @@ FROM base-builder
 RUN python3 -m pip uninstall -y apex
 RUN git clone https://github.com/NVIDIA/apex
 #  `MAX_JOBS=1` disables parallel building to avoid cpu memory OOM when building image on GitHub Action (standard) runners
-RUN cd apex && MAX_JOBS=1 python3 -m pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
+RUN cd apex && MAX_JOBS=1 python3 -m pip install --global-option="--cpp_ext" --global-option="--cuda_ext" --no-cache -v --disable-pip-version-check .
 
 RUN mkdir -p /workspace/builds
 COPY --from=bnb-builder /workspace/bitsandbytes /workspace/builds/bitsandbytes

examples/cerebras/qlora.yml

@@ -1,60 +0,0 @@
-base_model: cerebras/Cerebras-GPT-1.3B
-base_model_config: cerebras/Cerebras-GPT-1.3B
-load_in_8bit: false
-load_in_4bit: true
-strict: false
-push_dataset_to_hub:
-datasets:
-  - path: teknium/GPT4-LLM-Cleaned
-    type: alpaca
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.01
-adapter: qlora
-lora_model_dir:
-sequence_len: 2048
-max_packed_sequence_len: 2048
-lora_r: 16
-lora_alpha: 32
-lora_dropout: 0.05
-lora_target_modules:
-  - c_fc
-  - c_attn
-  - c_proj
-lora_target_linear:
-lora_fan_in_fan_out:
-wandb_project:
-wandb_watch:
-wandb_run_id:
-wandb_log_model:
-output_dir: ./qlora-out
-batch_size: 4
-micro_batch_size: 4
-num_epochs: 2
-optimizer: paged_adamw_8bit
-torchdistx_path:
-lr_scheduler: cosine
-learning_rate: 0.0002
-train_on_inputs: false
-group_by_length: true
-bf16: true
-fp16: false
-tf32: true
-gradient_checkpointing: true
-early_stopping_patience:
-resume_from_checkpoint:
-local_rank:
-logging_steps: 1
-xformers_attention: true
-flash_attention:
-gptq_groupsize:
-gptq_model_v1:
-warmup_steps: 10
-eval_steps: 20
-save_steps:
-debug:
-deepspeed:
-weight_decay: 0.1
-fsdp:
-fsdp_config:
-special_tokens:
-  pad_token: "<|endoftext|>"

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

@@ -23,7 +23,7 @@ lora_dropout: 0.0
 lora_target_modules:
 lora_target_linear: true
 lora_fan_in_fan_out:
-wandb_project:
+wandb_project: falcon-7b
 wandb_watch:
 wandb_run_id:
 wandb_log_model:

examples/falcon/config-7b-qlora.yml

@@ -1,92 +0,0 @@
-# 1b: tiiuae/falcon-rw-1b
-# 40b: tiiuae/falcon-40b
-base_model: tiiuae/falcon-7b
-base_model_config: tiiuae/falcon-7b
-# required by falcon custom model code: https://huggingface.co/tiiuae/falcon-7b/tree/main
-trust_remote_code: true
-model_type: AutoModelForCausalLM
-tokenizer_type: AutoTokenizer
-load_in_8bit: false
-# enable 4bit for QLoRA
-load_in_4bit: true
-gptq: false
-strict: false
-push_dataset_to_hub:
-datasets:
-  - path: QingyiSi/Alpaca-CoT
-    data_files:
-      - Chain-of-Thought/formatted_cot_data/gsm8k_train.json
-    type: "alpaca:chat"
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.01
-# enable QLoRA
-adapter: qlora
-lora_model_dir:
-sequence_len: 2048
-max_packed_sequence_len:
-
-# hyperparameters from QLoRA paper Appendix B.2
-# "We find hyperparameters to be largely robust across datasets"
-lora_r: 64
-lora_alpha: 16
-# 0.1 for models up to 13B
-# 0.05 for 33B and 65B models
-lora_dropout: 0.05
-# add LoRA modules on all linear layers of the base model
-lora_target_modules:
-lora_target_linear: true
-lora_fan_in_fan_out:
-
-wandb_project:
-wandb_watch:
-wandb_run_id:
-wandb_log_model:
-output_dir: ./qlora-out
-
-# QLoRA paper Table 9
-# - 16 for 7b & 13b
-# - 32 for 33b, 64 for 64b
-# Max size tested on A6000
-# - 7b: 40
-# - 40b: 4
-# decrease if OOM, increase for max VRAM utilization
-micro_batch_size: 1
-gradient_accumulation_steps: 2
-num_epochs: 3
-# Optimizer for QLoRA
-optimizer: paged_adamw_32bit
-torchdistx_path:
-lr_scheduler: cosine
-# QLoRA paper Table 9
-# - 2e-4 for 7b & 13b
-# - 1e-4 for 33b & 64b
-learning_rate: 0.0002
-train_on_inputs: false
-group_by_length: false
-bf16: true
-fp16: false
-tf32: true
-gradient_checkpointing: true
-# stop training after this many evaluation losses have increased in a row
-# https://huggingface.co/transformers/v4.2.2/_modules/transformers/trainer_callback.html#EarlyStoppingCallback
-early_stopping_patience: 3
-resume_from_checkpoint:
-auto_resume_from_checkpoints: true
-local_rank:
-logging_steps: 1
-xformers_attention: true
-flash_attention:
-gptq_groupsize:
-gptq_model_v1:
-warmup_steps: 10
-eval_steps: 5
-save_steps: 10
-debug:
-deepspeed:
-weight_decay: 0.000001
-fsdp:
-fsdp_config:
-special_tokens:
-  pad_token: "<|endoftext|>"
-  bos_token: ">>ABSTRACT<<"
-  eos_token: "<|endoftext|>"

examples/falcon/config-7b.yml

@@ -23,7 +23,7 @@ lora_dropout: 0.0
 lora_target_modules:
 lora_target_linear: true
 lora_fan_in_fan_out:
-wandb_project:
+wandb_project: falcon-7b
 wandb_watch:
 wandb_run_id:
 wandb_log_model:

examples/gptj/qlora.yml

@@ -1,57 +0,0 @@
-base_model: EleutherAI/gpt-j-6b
-base_model_config: EleutherAI/gpt-j-6b
-load_in_8bit: false
-load_in_4bit: true
-strict: false
-push_dataset_to_hub:
-datasets:
-  - path: teknium/GPT4-LLM-Cleaned
-    type: alpaca
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.01
-adapter: qlora
-lora_model_dir:
-sequence_len: 2048
-max_packed_sequence_len:
-lora_r: 8
-lora_alpha: 32
-lora_dropout: 0.05
-lora_target_modules:
-lora_target_linear: true
-lora_fan_in_fan_out:
-wandb_project:
-wandb_watch:
-wandb_run_id:
-wandb_log_model:
-output_dir: ./qlora-out
-gradient_accumulation_steps: 2
-micro_batch_size: 2
-num_epochs: 2
-optimizer: paged_adamw_8bit
-torchdistx_path:
-lr_scheduler: cosine
-learning_rate: 0.0001
-train_on_inputs: false
-group_by_length: true
-bf16: true
-fp16: false
-tf32: true
-gradient_checkpointing: true
-early_stopping_patience:
-resume_from_checkpoint:
-local_rank:
-logging_steps: 1
-xformers_attention: true
-flash_attention:
-gptq_groupsize:
-gptq_model_v1:
-warmup_steps: 10
-eval_steps: 20
-save_steps:
-debug:
-deepspeed:
-weight_decay: 0.1
-fsdp:
-fsdp_config:
-special_tokens:
-  pad_token: "<|endoftext|>"

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

@@ -3,6 +3,6 @@
 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
+accelerate launch scripts/finetune.py examples/4bit-lora-7b/config.yml
 
 ```

examples/lora-openllama-3b/config.yml

@@ -1,5 +1,5 @@
-base_model: openlm-research/open_llama_3b
-base_model_config: openlm-research/open_llama_3b
+base_model: openlm-research/open_llama_3b_600bt_preview
+base_model_config: openlm-research/open_llama_3b_600bt_preview
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 load_in_8bit: true
@@ -49,7 +49,7 @@ early_stopping_patience:
 resume_from_checkpoint:
 local_rank:
 logging_steps: 1
-xformers_attention: true
+xformers_attention:
 flash_attention:
 gptq_groupsize:
 gptq_model_v1:

examples/openllama-3b/README.md

@@ -1,16 +0,0 @@
-# openllama-3b
-
-Basic full tune
-```shell
-accelerate launch scripts/finetune.py examples/openllama-3b/config.yml
-```
-
-LoRA
-```shell
-accelerate launch scripts/finetune.py examples/openllama-3b/lora.yml
-```
-
-QLoRA
-```shell
-accelerate launch scripts/finetune.py examples/openllama-3b/qlora.yml
-```

examples/openllama-3b/config.yml

@@ -1,62 +0,0 @@
-base_model: openlm-research/open_llama_3b
-base_model_config: openlm-research/open_llama_3b
-model_type: LlamaForCausalLM
-tokenizer_type: LlamaTokenizer
-load_in_8bit: false
-load_in_4bit: false
-strict: false
-push_dataset_to_hub:
-datasets:
-  - path: teknium/GPT4-LLM-Cleaned
-    type: alpaca
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.02
-adapter:
-lora_model_dir:
-sequence_len: 256
-max_packed_sequence_len:
-lora_r:
-lora_alpha:
-lora_dropout:
-lora_target_modules:
-lora_target_linear:
-lora_fan_in_fan_out:
-wandb_project:
-wandb_watch:
-wandb_run_id:
-wandb_log_model:
-output_dir: ./openllama-out
-gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 3
-optimizer: adamw_bnb_8bit
-torchdistx_path:
-lr_scheduler: cosine
-learning_rate: 0.00001
-train_on_inputs: false
-group_by_length: false
-float16: true
-bf16: false
-fp16: false
-tf32: false
-gradient_checkpointing: true
-early_stopping_patience:
-resume_from_checkpoint:
-local_rank:
-logging_steps: 1
-xformers_attention: true
-flash_attention:
-gptq_groupsize:
-gptq_model_v1:
-warmup_steps: 10
-eval_steps: 50
-save_steps:
-debug:
-deepspeed:
-weight_decay: 0.1
-fsdp:
-fsdp_config:
-special_tokens:
-  bos_token: "<s>"
-  eos_token: "</s>"
-  unk_token: "<unk>"

examples/pythia-12b/README.md

@@ -1,4 +1,4 @@
-# Pythia 12B
+# Python 12B
 
 - Single-GPU A100 only (?)
 
@@ -7,3 +7,4 @@ python scripts/finetune.py examples/pythia-12b/config.yml
 ```
 
 ⚠️ Multiple-GPU A100 - Doesn't seem to work with multi-gpu without causing OOM! ⚠️
+

examples/pythia-12b/config.yml

@@ -22,7 +22,7 @@ lora_dropout: 0.0
 lora_target_modules:
 lora_target_linear: true
 lora_fan_in_fan_out: true  # pythia/GPTNeoX lora specific
-wandb_project:
+wandb_project: pythia-12b
 wandb_watch:
 wandb_run_id:
 wandb_log_model:
@@ -45,5 +45,5 @@ resume_from_checkpoint:
 local_rank:
 gradient_checkpointing: true
 fsdp:
-fsdp_config:
+fsdp_transformer_layer_cls_to_wrap:
 collator_pad_to_longest: true

examples/qlora-openllama-3b/README.md

@@ -0,0 +1,6 @@
+# qlora-openllama-3b
+
+```shell
+accelerate launch scripts/finetune.py examples/qlora-openllama-3b/config.yml
+
+```

examples/qlora-openllama-3b/config.yml

@@ -1,5 +1,5 @@
-base_model: openlm-research/open_llama_3b
-base_model_config: openlm-research/open_llama_3b
+base_model: openlm-research/open_llama_3b_600bt_preview
+base_model_config: openlm-research/open_llama_3b_600bt_preview
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 load_in_8bit: false

examples/redpajama/config-3b.yml

@@ -1,7 +1,7 @@
 base_model: togethercomputer/RedPajama-INCITE-Chat-3B-v1
 base_model_config: togethercomputer/RedPajama-INCITE-Chat-3B-v1
 model_type: GPTNeoXForCausalLM
-tokenizer_type: AutoTokenizer
+tokenizer_type: GPTNeoXTokenizer
 trust_remote_code:
 load_in_8bit: false
 datasets:

scripts/finetune.py

@@ -14,8 +14,9 @@ import torch
 import yaml
 
 # add src to the pythonpath so we don't need to pip install this
+from datasets import Dataset
 from optimum.bettertransformer import BetterTransformer
-from transformers import GenerationConfig, TextStreamer
+from transformers import GenerationConfig
 
 from axolotl.utils.data import load_prepare_datasets, load_pretraining_dataset
 from axolotl.utils.dict import DictDefault
@@ -49,7 +50,7 @@ def choose_device(cfg):
 
     cfg.device = get_device()
     if cfg.device_map != "auto":
-        if cfg.device.startswith("cuda"):
+        if cfg.device == "cuda":
             cfg.device_map = {"": cfg.local_rank}
         else:
             cfg.device_map = {"": cfg.device}
@@ -64,43 +65,23 @@ def get_multi_line_input() -> Optional[str]:
     return instruction
 
 
-def do_inference(cfg, model, tokenizer, prompter: Optional[str]):
-    default_tokens = {"unk_token": "<unk>", "bos_token": "<s>", "eos_token": "</s>"}
+def do_inference(cfg, model, tokenizer, prompter="AlpacaPrompter"):
+    tokenizer.add_special_tokens({"unk_token": "<unk>"})
+    tokenizer.add_special_tokens({"bos_token": "<s>"})
+    tokenizer.add_special_tokens({"eos_token": "</s>"})
 
-    for token, symbol in default_tokens.items():
-        # If the token isn't already specified in the config, add it
-        if not (cfg.special_tokens and token in cfg.special_tokens):
-            tokenizer.add_special_tokens({token: symbol})
-
-    prompter_module = None
-    if prompter:
-        prompter_module = getattr(
-            importlib.import_module("axolotl.prompters"), prompter
-        )
-
-    if cfg.landmark_attention:
-        from axolotl.monkeypatch.llama_landmark_attn import set_model_mem_id
-
-        set_model_mem_id(model, tokenizer)
-        model.set_mem_cache_args(
-            max_seq_len=255, mem_freq=50, top_k=5, max_cache_size=None
-        )
+    prompter_module = getattr(importlib.import_module("axolotl.prompters"), prompter)
 
     while True:
-        print("=" * 80)
         # support for multiline inputs
         instruction = get_multi_line_input()
         if not instruction:
             return
-        if prompter_module:
-            prompt: str = next(
-                prompter_module().build_prompt(instruction=instruction.strip("\n"))
-            )
-        else:
-            prompt = instruction.strip()
+        prompt: str = next(
+            prompter_module().build_prompt(instruction=instruction.strip("\n"))
+        )
         batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
 
-        print("=" * 40)
         model.eval()
         with torch.no_grad():
             generation_config = GenerationConfig(
@@ -119,13 +100,10 @@ def do_inference(cfg, model, tokenizer, prompter: Optional[str]):
                 output_hidden_states=False,
                 output_scores=False,
             )
-            streamer = TextStreamer(tokenizer)
             generated = model.generate(
                 inputs=batch["input_ids"].to(cfg.device),
                 generation_config=generation_config,
-                streamer=streamer,
             )
-        print("=" * 40)
         print(tokenizer.decode(generated["sequences"].cpu().tolist()[0]))
 
 
@@ -175,7 +153,7 @@ def train(
     cfg_keys = cfg.keys()
     for k, _ in kwargs.items():
         # if not strict, allow writing to cfg even if it's not in the yml already
-        if k in cfg_keys or not cfg.strict:
+        if k in cfg_keys or cfg.strict is False:
             # handle booleans
             if isinstance(cfg[k], bool):
                 cfg[k] = bool(kwargs[k])
@@ -207,30 +185,27 @@ def train(
             cfg.fp16 = True
         cfg.bf16 = False
 
-    if cfg.tf32:
-        torch.backends.cuda.matmul.allow_tf32 = True
-
     # load the tokenizer first
     tokenizer_config = cfg.tokenizer_config or cfg.base_model_config
     logging.info(f"loading tokenizer... {tokenizer_config}")
     tokenizer = load_tokenizer(tokenizer_config, cfg.tokenizer_type, cfg)
 
-    if (
-        check_not_in(["shard", "merge_lora"], kwargs) and not cfg.inference
+    if check_not_in(
+        ["inference", "shard", "merge_lora"], kwargs
     ):  # don't need to load dataset for these
         if not cfg.pretraining_dataset:
             train_dataset, eval_dataset = load_prepare_datasets(
                 tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH
             )
         else:
+            if cfg.pretraining_dataset is True:
+                pretraining_dataset = "togethercomputer/RedPajama-Data-1T"
+            else:
+                pretraining_dataset = cfg.pretraining_dataset
             train_dataset = load_pretraining_dataset(
-                cfg.pretraining_dataset,
-                tokenizer,
-                max_tokens=cfg.sequence_len,
-                seed=cfg.seed,
+                pretraining_dataset, tokenizer, max_tokens=cfg.sequence_len
             )
-            # https://discuss.huggingface.co/t/how-to-use-huggingface-trainer-streaming-datasets-without-wrapping-it-with-torchdatas-iterablewrapper/25230
-            train_dataset = train_dataset.with_format("torch")
+            train_dataset = Dataset.from_list(list(train_dataset))
             eval_dataset = None
 
     if cfg.debug or "debug" in kwargs:
@@ -255,6 +230,7 @@ def train(
         tokenizer,
         cfg,
         adapter=cfg.adapter,
+        inference=("inference" in kwargs),
     )
 
     if "merge_lora" in kwargs and cfg.adapter is not None:
@@ -267,21 +243,30 @@ def train(
             model.save_pretrained(str(Path(cfg.output_dir) / "merged"))
         return
 
-    if cfg.inference:
+    if "inference" in kwargs:
         logging.info("calling do_inference function")
-        prompter: Optional[str] = "AlpacaPrompter"
-        if "prompter" in kwargs:
-            if kwargs["prompter"] == "None":
-                prompter = None
-            else:
-                prompter = kwargs["prompter"]
-        do_inference(cfg, model, tokenizer, prompter=prompter)
+        do_inference(cfg, model, tokenizer)
         return
 
     if "shard" in kwargs:
         model.save_pretrained(cfg.output_dir)
         return
 
+    if cfg.debug:
+        logging.info("check_dataset_labels...")
+        check_dataset_labels(
+            train_dataset.select(
+                [random.randrange(0, len(train_dataset) - 1) for i in range(5)]  # nosec
+            ),
+            tokenizer,
+        )
+
+    if prepare_ds_only:
+        logging.info("Finished preparing dataset. Exiting...")
+        return
+
+    model.train()
+
     trainer = setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer)
 
     model.config.use_cache = False
@@ -326,8 +311,6 @@ def train(
                 f"Using Auto-resume functionality to start with checkpoint at {resume_from_checkpoint}"
             )
 
-    if not Path(cfg.output_dir).is_dir():
-        os.makedirs(cfg.output_dir, exist_ok=True)
     if cfg.flash_optimum:
         with torch.backends.cuda.sdp_kernel(
             enable_flash=True, enable_math=True, enable_mem_efficient=True

src/axolotl/datasets.py

@@ -33,16 +33,12 @@ class TokenizedPromptDataset(IterableDataset):
 
     def __iter__(self):
         iterator = iter(self.dataset)
-        count = 0
         # Loop through the entire dataset
         for example in iterator:
             try:
                 yield self.prompt_tokenizer.tokenize_prompt(example)
-                count += 1
             except InvalidDataException:
                 pass
-        if count == 0:
-            raise RuntimeError("Expected at least one datapoint in dataset.")
 
 
 # TODO this isn't the best since it can't interleave datasets
@@ -126,7 +122,6 @@ class ConstantLengthDataset(IterableDataset):
                     buffer_len = 0
 
                 if example:
-                    # FIXME
                     # just going to drop data points that are too long
                     if len(example["input_ids"]) <= self.seq_length:
                         input_ids = example["input_ids"]

src/axolotl/monkeypatch/xpos_rope_llama_monkey_patch.py

@@ -1,94 +0,0 @@
-# pylint: skip-file
-"""
-Copied from https://github.com/kaiokendev/cutoff-len-is-context-len/blob/main/util/xpos_rope_llama_monkey_patch.py
-"""
-import torch
-import transformers
-import transformers.models.llama.modeling_llama
-from einops import rearrange
-
-
-class XposRotaryEmbedding(torch.nn.Module):
-    def __init__(
-        self,
-        dim,
-        max_position_embeddings=2048,
-        base=10000,
-        device=None,
-        scale_base=2048,
-        use_xpos=True,
-    ):
-        super().__init__()
-        self.max_seq_len_cached = max_position_embeddings
-        self.scale_base = scale_base
-
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
-        t = torch.arange(self.max_seq_len_cached, device=device).type_as(inv_freq)
-        freqs = torch.einsum("i , j -> i j", t, inv_freq)
-        freqs = torch.cat((freqs, freqs), dim=-1)
-
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.register_buffer("freqs_cached", freqs, persistent=False)
-
-        if not use_xpos:
-            self.register_buffer("scale", None)
-            self.register_buffer("scale_cached", torch.ones(1))
-            return
-
-        scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
-        power = (t - (self.max_seq_len_cached // 2)) / self.scale_base
-        scale_cached = scale ** rearrange(power, "n -> n 1")
-        scale_cached = torch.cat((scale_cached, scale_cached), dim=-1)
-
-        self.register_buffer("scale", scale, persistent=False)
-        self.register_buffer("scale_cached", scale_cached, persistent=False)
-
-    def forward(
-        self,
-        x,
-        seq_len,
-    ):
-        if seq_len > self.max_seq_len_cached:
-            self.max_seq_len_cached = seq_len
-            t = torch.arange(self.max_seq_len_cached, device=x.device).type_as(
-                self.inv_freq
-            )
-            freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
-            freqs = torch.cat((freqs, freqs), dim=-1).to(dtype=x.dtype)
-
-            self.register_buffer("freqs_cached", freqs)
-
-            if self.scale is None:
-                self.register_buffer(
-                    "scale_cached", torch.ones(1, device=x.device).to(dtype=x.dtype)
-                )
-
-                return self.freqs_cached.to(dtype=x.dtype), self.scale_cached
-
-            power = (t - (seq_len // 2)) / self.scale_base
-            scale = self.scale ** rearrange(power, "n -> n 1")
-            scale = torch.cat((scale, scale), dim=-1).to(dtype=x.dtype)
-            self.register_buffer("scale_cached", scale)
-
-        return self.freqs_cached.to(dtype=x.dtype), self.scale_cached.to(dtype=x.dtype)
-
-
-def rotate_half(x):
-    x1, x2 = x.chunk(2, dim=-1)
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rotary_pos_emb(q, k, freqs, scale=1, position_ids=None):
-    freqs = freqs[position_ids, :]
-    if scale.shape[-1] != 1:
-        scale = scale[position_ids, :]
-
-    q_embed = (q * freqs.cos() * scale) + (rotate_half(q) * freqs.sin() * scale)
-    k_embed = (k * freqs.cos() * 1 / scale) + (rotate_half(k) * freqs.sin() * 1 / scale)
-
-    return q_embed, k_embed
-
-
-def replace_llama_rope_with_xpos_rope():
-    transformers.models.llama.modeling_llama.LlamaRotaryEmbedding = XposRotaryEmbedding
-    transformers.models.llama.modeling_llama.apply_rotary_pos_emb = apply_rotary_pos_emb

src/axolotl/prompt_strategies/alpaca_chat.py

@@ -6,7 +6,7 @@ from axolotl.prompt_tokenizers import (
     AlpacaPromptTokenizingStrategy,
     InstructionPromptTokenizingStrategy,
 )
-from axolotl.prompters import AlpacaPrompter, PromptStyle, UnpromptedPrompter
+from axolotl.prompters import AlpacaPrompter, PromptStyle
 
 
 def load(tokenizer, cfg):
@@ -18,42 +18,6 @@ def load(tokenizer, cfg):
     )
 
 
-class AlpacaConcisePrompter(AlpacaPrompter):
-    """
-    Alpaca Prompter extending the system prompt to ask for concise chat-instruct answers
-    """
-
-    system_prompt = "Below is an instruction from a USER that describes a task, paired with an input that provides further context. The ASSISTANT writes a response that concisely and appropriately completes the request.\n\n"
-    system_no_input_prompt = "Below is an instruction from a USER that describes a task. The ASSISTANT writes a response that appropriately and concisely completes the request.\n\n"
-
-
-class AlpacaChatPrompter(AlpacaPrompter):
-    """
-    Alpaca Chat Prompter extending the system prompt to for chat-instruct answers
-    """
-
-    system_prompt = "Below is an instruction from a USER that describes a task, paired with an input that provides further context. The ASSISTANT writes a response that concisely and appropriately completes the request.\n\n"
-    system_no_input_prompt = "Below is an instruction from a USER that describes a task. The ASSISTANT writes a response that appropriately and concisely completes the request.\n\n"
-
-    def __init__(self):  # pylint: disable=super-init-not-called
-        self.prompt_style = PromptStyle.CHAT.value
-        self.match_prompt_style()
-
-
-class NoSystemPrompter(AlpacaPrompter):
-    """
-    Null Prompter with no system prompts
-    """
-
-    system_prompt = ""
-    system_no_input_prompt = ""
-    turn_format = "{instruction} {input} "
-    turn_no_input_format = "{instruction} "
-
-    def __init__(self):  # pylint: disable=super-init-not-called
-        pass
-
-
 class AlpacaQAPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
     """
     Tokenizing strategy for AlpacaQA
@@ -67,49 +31,9 @@ class AlpacaQAPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
         )
 
 
-class CamelAIPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
-    """
-    Tokenizing strategy for CamelAI datasets
-    """
-
-    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
-        return (
-            prompt["message_1"],
-            "",
-            prompt["message_2"],
-        )
-
-
-def load_concise(tokenizer, cfg):
-    return AlpacaPromptTokenizingStrategy(
-        AlpacaConcisePrompter(PromptStyle.CHAT.value),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-
-
 def load_qa(tokenizer, cfg):
     return AlpacaQAPromptTokenizingStrategy(
-        AlpacaChatPrompter(),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-
-
-def load_camel_ai(tokenizer, cfg):
-    return CamelAIPromptTokenizingStrategy(
-        AlpacaChatPrompter(),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-
-
-def load_no_prompt(tokenizer, cfg):
-    return AlpacaPromptTokenizingStrategy(
-        UnpromptedPrompter(PromptStyle.CHAT.value),
+        AlpacaPrompter(PromptStyle.CHAT.value),
         tokenizer,
         cfg.train_on_inputs,
         cfg.sequence_len,

src/axolotl/prompt_strategies/alpaca_instruct.py

@@ -1,7 +1,7 @@
 """Module loading the AlpacaInstructPromptTokenizingStrategy class"""
 
 from axolotl.prompt_tokenizers import AlpacaPromptTokenizingStrategy
-from axolotl.prompters import AlpacaPrompter, PromptStyle, UnpromptedPrompter
+from axolotl.prompters import AlpacaPrompter, PromptStyle
 
 
 def load(tokenizer, cfg):
@@ -11,12 +11,3 @@ def load(tokenizer, cfg):
         cfg.train_on_inputs,
         cfg.sequence_len,
     )
-
-
-def load_no_prompt(tokenizer, cfg):
-    return AlpacaPromptTokenizingStrategy(
-        UnpromptedPrompter(PromptStyle.INSTRUCT.value),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )

src/axolotl/prompt_strategies/alpaca_w_system.py

@@ -1,84 +0,0 @@
-"""
-Prompt strategies loader for alpaca instruction datasets with system prompts
-"""
-from typing import Generator, Tuple, Union
-
-from axolotl.prompt_tokenizers import PromptTokenizingStrategy
-from axolotl.prompters import AlpacaPrompter, PromptStyle
-
-
-class InstructionWSystemPromptTokenizingStrategy(PromptTokenizingStrategy):
-    """
-    Tokenizing strategy for instruction-based prompts.
-    """
-
-    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str, str]:
-        return (
-            prompt["instruction"],
-            prompt["input"] if "input" in prompt else "",
-            prompt["output"],
-            prompt["system"],
-        )
-
-    def tokenize_prompt(self, prompt):
-        # pylint: disable=duplicate-code
-        (
-            instruction,
-            input,  # pylint: disable=redefined-builtin
-            response,
-            system,
-        ) = self.parse_instruction_fields(prompt)
-        user_prompt = next(
-            iter(
-                self.prompter.build_prompt_w_system(
-                    system,
-                    instruction,
-                    input,
-                )
-            )
-        )
-        tokenized_prompt = self._tokenize(user_prompt, add_eos_token=False)
-        if not self.train_on_inputs:
-            user_prompt_len = len(tokenized_prompt["input_ids"])
-            # TODO this could be sped up using numpy array slicing
-            tokenized_prompt["labels"] = [-100] * user_prompt_len
-        tokenized_res_prompt = self._tokenize(
-            response, strip_bos_token=True, add_eos_token=True
-        )
-        tokenized_prompt["input_ids"] += tokenized_res_prompt["input_ids"]
-        tokenized_prompt["attention_mask"] += tokenized_res_prompt["attention_mask"]
-        tokenized_prompt["labels"] += tokenized_res_prompt["input_ids"]
-
-        return tokenized_prompt
-
-
-class SystemDataPrompter(AlpacaPrompter):
-    """
-    Alpaca Style Prompter that uses system prompts from the dataset
-    """
-
-    def build_prompt_w_system(
-        self,
-        system: str,
-        instruction: str,
-        input: Union[None, str] = None,  # pylint: disable=redefined-builtin
-        output: Union[None, str] = None,
-    ) -> Generator[str, None, None]:
-        # returns the full prompt from instruction and optional input
-        # if a label (=response, =output) is provided, it's also appended.
-        if input:
-            res = system + self.turn_format.format(instruction=instruction, input=input)
-        else:
-            res = system + self.turn_no_input_format.format(instruction=instruction)
-        if output:
-            res = f"{res}{output}"
-        yield res
-
-
-def load(tokenizer, cfg):
-    return InstructionWSystemPromptTokenizingStrategy(
-        SystemDataPrompter(PromptStyle.CHAT.value),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )

src/axolotl/prompt_strategies/context_qa.py

@@ -1,67 +0,0 @@
-"""Module containing the classes for Context QA Prompt Tokenization Strategies"""
-from typing import Tuple
-
-from axolotl.prompt_tokenizers import InstructionPromptTokenizingStrategy
-from axolotl.prompters import AlpacaPrompter, PromptStyle
-
-
-# article, unanswerable_question, question, answer
-def load_404(tokenizer, cfg):
-    return AlpacaMissingInfoContextPromptTokenizingStrategy(
-        AlpacaContextPrompter(PromptStyle.CHAT.value),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-
-
-def load(tokenizer, cfg):
-    return AlpacaContextPromptTokenizingStrategy(
-        AlpacaContextPrompter(PromptStyle.CHAT.value),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-
-
-class AlpacaContextPrompter(AlpacaPrompter):
-    """
-    Customized system prompted for concise QA
-    """
-
-    system_prompt = (
-        "Use the following contextual information to concisely answer the question.\n"
-    )
-    system_no_input_prompt = (
-        "Use the following contextual information to concisely answer the question.\n"
-    )
-
-
-class AlpacaContextPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
-    """
-    Tokenization Strategy to combine in-context article with a question and answer
-    """
-
-    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
-        return (
-            prompt["article"] + "\n===\n" + prompt["question"],
-            "",
-            prompt["answer"],
-        )
-
-
-class AlpacaMissingInfoContextPromptTokenizingStrategy(
-    InstructionPromptTokenizingStrategy
-):
-    """
-    Tokenization Strategy to combine in-context article with a question that can't be answered
-    from the context and a default response to that effect
-    """
-
-    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
-        return (
-            prompt["article"] + "\n===\n" + prompt["unanswerable_question"],
-            "",
-            "The context provided does not contain any information about your inquiry. "
-            "Therefore, I'm unable to answer your question based on the given context.",
-        )

src/axolotl/prompt_strategies/sharegpt_jokes.py

@@ -1,28 +0,0 @@
-"""Module for Jokes prompts using sharegpt style """
-from axolotl.prompt_tokenizers import ShareGPTPromptTokenizingStrategy
-from axolotl.prompters import PromptStyle, ShareGPTPrompter
-
-
-def load(tokenizer, cfg):
-    return SimpleJokesShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompter(PromptStyle.CHAT.value),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-
-
-class SimpleJokesShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
-    """
-    Tokenization strategy for asking bot to tell a joke and then explain why its funny
-    """
-
-    # title, text, explanation
-    def get_conversation_thread(self, prompt):
-        title = "" if not prompt["title"] else prompt["title"] + " "
-        return [
-            {"from": "human", "value": "Tell me a joke."},
-            {"from": "gpt", "value": title + prompt["text"]},
-            {"from": "human", "value": "Why is that joke funny?"},
-            {"from": "gpt", "value": prompt["explanation"]},
-        ]

src/axolotl/prompt_strategies/sharegpt_simple.py

@@ -1,67 +0,0 @@
-"""Module containing the SimpleShareGPTPromptTokenizingStrategy class"""
-
-from axolotl.prompt_tokenizers import ShareGPTPromptTokenizingStrategy
-from axolotl.prompters import PromptStyle, ShareGPTPrompter
-
-
-def load(tokenizer, cfg):
-    return SimpleShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompter(PromptStyle.CHAT.value),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-
-
-def load_role(tokenizer, cfg):
-    return SimpleRoleShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompter(PromptStyle.CHAT.value),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-
-
-def load_guanaco(tokenizer, cfg):
-    return GuanacoShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompter(PromptStyle.CHAT.value),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-
-
-class SimpleShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
-    """
-    basic sharegpt strategy to grab conversations from the sample row
-    """
-
-    def get_conversation_thread(self, prompt):
-        return prompt["conversations"]
-
-
-class SimpleRoleShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
-    """
-    basic sharegpt strategy to grab conversations from the sample row, but uses role instead of from
-    """
-
-    def get_conversation_thread(self, prompt):
-        conversations = prompt["conversations"]
-        # remap role: prompter/assistant, text: ... => from: human/gpt, value: ...
-        turns = [{"from": t["role"], "value": t["value"]} for t in conversations]
-        return turns
-
-
-class GuanacoShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
-    """
-    sharegpt strategy that remaps oasst data to sharegpt format
-    """
-
-    def get_conversation_thread(self, prompt):
-        conversations = prompt["conversations"]
-        # remap role: prompter/assistant, text: ... => from: human/gpt, value: ...
-        role_map = {"prompter": "human", "assistant": "gpt"}
-        turns = [
-            {"from": role_map[t["role"]], "value": t["text"]} for t in conversations
-        ]
-        return turns

src/axolotl/prompt_tokenizers.py

@@ -87,9 +87,7 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy):
     Tokenizing strategy for instruction-based prompts.
     """
 
-    def parse_instruction_fields(
-        self, prompt
-    ) -> Union[Tuple[str, str, str], Tuple[str, str, str, str]]:
+    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
         raise NotImplementedError
 
     def tokenize_prompt(self, prompt):
@@ -98,27 +96,25 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy):
             input,  # pylint: disable=redefined-builtin
             response,
         ) = self.parse_instruction_fields(prompt)
-        user_prompt = next(
-            iter(
-                self.prompter.build_prompt(
-                    instruction,
-                    input,
+        full_prompt = self._build_full_prompt(instruction, input, response)
+        tokenized_full_prompt = self._tokenize(full_prompt)
+        if not self.train_on_inputs:
+            user_prompt = next(
+                iter(
+                    self.prompter.build_prompt(
+                        instruction,
+                        input,
+                    )
                 )
             )
-        )
-        tokenized_prompt = self._tokenize(user_prompt, add_eos_token=False)
-        if not self.train_on_inputs:
-            user_prompt_len = len(tokenized_prompt["input_ids"])
+            tokenized_user_prompt = self._tokenize(user_prompt, add_eos_token=False)
+            user_prompt_len = len(tokenized_user_prompt["input_ids"])
             # TODO this could be sped up using numpy array slicing
-            tokenized_prompt["labels"] = [-100] * user_prompt_len
-        tokenized_res_prompt = self._tokenize(
-            response, strip_bos_token=True, add_eos_token=True
-        )
-        tokenized_prompt["input_ids"] += tokenized_res_prompt["input_ids"]
-        tokenized_prompt["attention_mask"] += tokenized_res_prompt["attention_mask"]
-        tokenized_prompt["labels"] += tokenized_res_prompt["input_ids"]
+            tokenized_full_prompt["labels"] = [
+                -100
+            ] * user_prompt_len + tokenized_full_prompt["labels"][user_prompt_len:]
 
-        return tokenized_prompt
+        return tokenized_full_prompt
 
     def _build_full_prompt(
         self, instruction, input, response  # pylint: disable=redefined-builtin

src/axolotl/prompters.py

@@ -24,8 +24,6 @@ class AlpacaPrompter:
 
     system_prompt = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n"
     system_no_input_prompt = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n"
-    turn_format: str
-    turn_no_input_format: str
     prompt_style: Optional[PromptStyle] = None
 
     def __init__(self, prompt_style=PromptStyle.INSTRUCT.value):
@@ -34,13 +32,23 @@ class AlpacaPrompter:
 
     def match_prompt_style(self):
         if self.prompt_style == PromptStyle.INSTRUCT.value:
-            self.turn_format = "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
-            self.turn_no_input_format = (
-                "### Instruction:\n{instruction}\n\n### Response:\n"
+            self.prompt_input = (
+                self.system_prompt
+                + "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
             )
+            self.prompt_no_input = (
+                self.system_no_input_prompt
+                + "### Instruction:\n{instruction}\n\n### Response:\n"
+            )
+            self.response_split = "### Response:"
         if self.prompt_style == PromptStyle.CHAT.value:
-            self.turn_format = "USER: {instruction}\n{input}\nASSISTANT:"
-            self.turn_no_input_format = "USER: {instruction}\nASSISTANT:"
+            self.prompt_input = (
+                self.system_prompt + "USER: {instruction}\n{input}\nASSISTANT:"
+            )
+            self.prompt_no_input = (
+                self.system_no_input_prompt + "USER: {instruction}\nASSISTANT:"
+            )
+            self.response_split = "ASSISTANT:"
 
     def build_prompt(
         self,
@@ -51,17 +59,16 @@ class AlpacaPrompter:
         # returns the full prompt from instruction and optional input
         # if a label (=response, =output) is provided, it's also appended.
         if input:
-            res = self.system_prompt + self.turn_format.format(
-                instruction=instruction, input=input
-            )
+            res = self.prompt_input.format(instruction=instruction, input=input)
         else:
-            res = self.system_no_input_prompt + self.turn_no_input_format.format(
-                instruction=instruction
-            )
+            res = self.prompt_no_input.format(instruction=instruction)
         if output:
             res = f"{res}{output}"
         yield res
 
+    def get_response(self, output: str) -> str:
+        return output.split(self.response_split)[1].strip()
+
 
 class UnpromptedPrompter(AlpacaPrompter):
     """
@@ -86,10 +93,7 @@ class MultipleChoiceExplainPrompter(AlpacaPrompter):
     """
 
     system_prompt = (
-        "Choose the answer that best answers the question. Explain your reasoning.\n"
-    )
-    system_no_input_prompt = (
-        "Choose the answer that best answers the question. Explain your reasoning.\n"
+        "Choose the answer that best answers the question. Explain your reasoning."
     )
 
 
@@ -98,12 +102,7 @@ class MultipleChoiceConcisePrompter(AlpacaPrompter):
     Prompter for multiple choice concise
     """
 
-    system_prompt = "Choose the answer that best answers the question. Be concise in your response.\n\n"
-    system_no_input_prompt = "Choose the answer that best answers the question. Be concise in your response.\n\n"
-
-    def match_prompt_style(self):
-        self.turn_format = "USER: {instruction}\n{input}\nASSISTANT:"
-        self.turn_no_input_format = "USER: {instruction}\nASSISTANT:"
+    prompt_input = "Choose the answer that best answers the question. Be concise in your response.\n\nUSER: {instruction}\n{input}\nASSISTANT:\n"
 
 
 class SummarizeTLDRPrompter(AlpacaPrompter):
@@ -111,12 +110,9 @@ class SummarizeTLDRPrompter(AlpacaPrompter):
     Prompter for summarize TLDR
     """
 
-    system_prompt = ""
-    system_no_input_prompt = ""
-
-    def match_prompt_style(self):
-        self.turn_format = "USER: Summarize the following article as a TL;DR.\n{instruction}\n{input}\nASSISTANT:"
-        self.turn_no_input_format = "USER: Summarize the following article as a TL;DR.\n{instruction}\nASSISTANT:"
+    prompt_no_input = (
+        "USER: Summarize the following article as a TL;DR.\n{instruction}\nASSISTANT:"
+    )
 
 
 class CompletionPrompter:
@@ -132,6 +128,9 @@ class CompletionPrompter:
     ) -> Generator[str, None, None]:
         yield instruction
 
+    def get_response(self, output: str) -> str:
+        return output.strip()
+
 
 class GPTeacherPrompter(AlpacaPrompter):
     """
@@ -211,6 +210,9 @@ class ReflectAlpacaPrompter:
             res = f"{res}{label}"
         yield res
 
+    def get_response(self, output: str) -> str:
+        return output.split(self.response_split)[1].strip()
+
 
 class SeparatorStyle(Enum):
     """Different separator style."""
@@ -259,33 +261,34 @@ class Conversation:
         self.messages.append([role, message])
 
 
+conv_vicuna_v1_1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=["USER", "ASSISTANT"],
+    messages=[],
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2=" ",
+)
+
+
 class ShareGPTPrompter:  # pylint: disable=too-few-public-methods
     """
     A prompter that generates prompts for the ShareGPT
     """
 
-    def __init__(self, prompt_style=None, system_prompt: Optional[str] = None):
+    def __init__(self, prompt_style=None):
         if prompt_style != PromptStyle.CHAT.value:
             raise ValueError(
                 f"unsupported prompt_style for ShareGPTPrompter({prompt_style})"
             )
-        system: str = (
-            system_prompt
-            if system_prompt
-            else (
-                "A chat between a curious user and an artificial intelligence assistant. "
-                "The assistant gives helpful, detailed, and polite answers to the user's questions."
-            )
-        )
-        self._conversation = Conversation(
-            system=system,
-            roles=["USER", "ASSISTANT"],
-            messages=[],
-            offset=0,
-            sep_style=SeparatorStyle.TWO,
-            sep=" ",
-            sep2=" ",
-        )
+
+    # def match_prompt_style(self):
+    #     if self.prompt_style == PromptStyle.chat.value:
+    #         self.prompt_input = self.system_prompt + "USER: {instruction}\n{input}\nASSISTANT:"
+    #         self.prompt_no_input = self.system_no_input_prompt + "USER: {instruction}\nASSISTANT:"
+    #         self.response_split = "ASSISTANT:"
 
     def build_prompt(self, source) -> Generator[str, None, None]:
         # ignore the system prompt if provided
@@ -297,7 +300,7 @@ class ShareGPTPrompter:  # pylint: disable=too-few-public-methods
             # also happens on the data splitting leaving empty conversations
             raise IndexError
 
-        conv = self._conversation.copy()
+        conv = conv_vicuna_v1_1.copy()
         roles = {"human": conv.roles[0], "gpt": conv.roles[1]}
 
         try:

src/axolotl/utils/data.py

@@ -1,12 +1,12 @@
 """Module containing data utilities"""
-import functools
+
 import logging
 from hashlib import md5
 from pathlib import Path
 from typing import List, Tuple, Union
 
 import torch
-from datasets import Dataset, DatasetDict, load_dataset, load_from_disk
+from datasets import Dataset, DatasetDict, IterableDataset, load_dataset, load_from_disk
 from huggingface_hub import hf_hub_download
 from transformers import PreTrainedTokenizerBase
 
@@ -79,13 +79,6 @@ def load_tokenized_prepared_datasets(
     else:
         logging.info(f"Unable to find prepared dataset in {prepared_ds_path}")
         logging.info("Loading raw datasets...")
-
-        if cfg.seed:
-            seed = cfg.seed
-        else:
-            logging.info("No seed provided, using default seed of 42")
-            seed = 42
-
         datasets = []
         # pylint: disable=invalid-name
         for d in cfg.datasets:
@@ -135,11 +128,11 @@ def load_tokenized_prepared_datasets(
             # support for using a subset of the data
             if d.shards:
                 if "train" in ds:
-                    ds = ds.shuffle(seed=seed)["train"].shard(
+                    ds = ds.shuffle(seed=42)["train"].shard(
                         num_shards=d.shards, index=0
                     )
                 else:
-                    ds = ds.shuffle(seed=seed).shard(num_shards=d.shards, index=0)
+                    ds = ds.shuffle(seed=42).shard(num_shards=d.shards, index=0)
             d_type = d.type
             d_type_split = d_type.split(":")
             d_base_type = d_type_split[0]
@@ -240,21 +233,14 @@ def load_tokenized_prepared_datasets(
                 ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
                 datasets.append(ds_wrapper)
             else:
-                suffix = ""
-                if ":load_" in d.type:
-                    suffix = f" Did you mean {d.type.replace(':load_', '.load_')}?"
-                logging.error(
-                    f"unhandled prompt tokenization strategy: {d.type}. {suffix}"
-                )
-                raise ValueError(
-                    f"unhandled prompt tokenization strategy: {d.type} {suffix}"
-                )
+                logging.error(f"unhandled prompt tokenization strategy: {d.type}")
+                raise ValueError(f"unhandled prompt tokenization strategy: {d.type}")
         logging.info("tokenizing, merging, and shuffling master dataset")
 
         samples: List[int] = []
         for d in datasets:
             samples = samples + list(d)
-        dataset = Dataset.from_list(samples).shuffle(seed=seed)
+        dataset = Dataset.from_list(samples).shuffle(seed=42)
         if cfg.local_rank == 0:
             logging.info(
                 f"Saving merged prepared dataset to disk... {prepared_ds_path}"
@@ -406,116 +392,32 @@ def load_prepare_datasets(
     return train_dataset, eval_dataset
 
 
-def encode_pretraining(tokenizer, max_tokens, examples):
-    res = tokenizer(
-        examples["text"],
-        truncation=True,
-        max_length=max_tokens - 2,
-        add_special_tokens=True,
-    )
-    # Convert to PyTorch tensors
-    input_ids = [torch.tensor(seq) for seq in res["input_ids"]]
-    attention_mask = [torch.tensor(seq) for seq in res["attention_mask"]]
-    new_input_ids = []
-    new_attention_mask = []
-    # Append EOS and PAD tokens to input_ids, and correct attention_mask
-    for i, _ in enumerate(input_ids):
-        input_ids[i] = torch.cat(
-            (
-                input_ids[i],
-                torch.tensor([tokenizer.eos_token_id, tokenizer.pad_token_id]),
-            ),
-            dim=0,
-        )
-        attention_mask[i] = torch.cat((attention_mask[i], torch.tensor([1, 0])), dim=0)
+class PretrainingDatasetWrapper(IterableDataset):
+    """
+    Wrapper for pretraining dataset that avoids loading the dataset into memory
+    """
 
-    # Concatenate tokens so that their lengths are less than max_tokens
-    buffer_input_ids = torch.tensor([], dtype=torch.long)
-    buffer_attention_mask = torch.tensor([], dtype=torch.long)
+    def __init__(self, tokenizer, dataset_path, max_tokens=2048):
+        self.tokenizer = tokenizer
+        self.dataset_path = dataset_path
+        self.max_tokens = max_tokens
 
-    for ids, mask in zip(input_ids, attention_mask):
-        if buffer_input_ids.numel() == max_tokens:
-            new_input_ids.append(buffer_input_ids)
-            new_attention_mask.append(buffer_attention_mask)
-            buffer_input_ids = torch.tensor([], dtype=torch.long)
-            buffer_attention_mask = torch.tensor([], dtype=torch.long)
-            buffer_input_ids = torch.cat((buffer_input_ids, ids), dim=0)
-            buffer_attention_mask = torch.cat((buffer_attention_mask, mask), dim=0)
-        elif buffer_input_ids.numel() + ids.numel() <= max_tokens:
-            buffer_input_ids = torch.cat((buffer_input_ids, ids), dim=0)
-            buffer_attention_mask = torch.cat((buffer_attention_mask, mask), dim=0)
-        else:
-            buffer_input_ids = torch.cat(
-                (
-                    buffer_input_ids,
-                    torch.full(
-                        (max_tokens - buffer_input_ids.numel(),),
-                        tokenizer.pad_token_id,
-                        dtype=torch.long,
-                    ),
-                ),
-                dim=0,
-            )
-            buffer_attention_mask = torch.cat(
-                (
-                    buffer_attention_mask,
-                    torch.full(
-                        (max_tokens - buffer_attention_mask.numel(),),
-                        0,
-                        dtype=torch.long,
-                    ),
-                ),
-                dim=0,
-            )
-            new_input_ids.append(buffer_input_ids)
-            new_attention_mask.append(buffer_attention_mask)
-            buffer_input_ids = torch.tensor([], dtype=torch.long)
-            buffer_attention_mask = torch.tensor([], dtype=torch.long)
-
-            buffer_input_ids = torch.cat((buffer_input_ids, ids), dim=0)
-            buffer_attention_mask = torch.cat((buffer_attention_mask, mask), dim=0)
-
-    if buffer_input_ids.numel() > 0:  # for any leftover tokens
-        while buffer_input_ids.numel() < max_tokens:  # make all sequences equal in size
-            buffer_input_ids = torch.cat(
-                (
-                    buffer_input_ids,
-                    torch.full(
-                        (max_tokens - buffer_input_ids.numel(),),
-                        tokenizer.pad_token_id,
-                        dtype=torch.long,
-                    ),
-                ),
-                dim=0,
-            )
-            buffer_attention_mask = torch.cat(
-                (
-                    buffer_attention_mask,
-                    torch.full(
-                        (max_tokens - buffer_attention_mask.numel(),),
-                        0,
-                        dtype=torch.long,
-                    ),
-                ),
-                dim=0,
-            )
-        new_input_ids.append(buffer_input_ids)
-        new_attention_mask.append(buffer_attention_mask)
-
-    ret = {
-        "input_ids": [seq.tolist() for seq in new_input_ids],
-        "labels": [seq.tolist() for seq in new_input_ids],
-        "attention_mask": [seq.tolist() for seq in new_attention_mask],
-    }
-
-    logging.debug(len(ret["input_ids"]))
-    return ret
+    def __iter__(self):
+        buffer = []
+        for sample in load_dataset(
+            self.dataset_path,
+        )["train"].shuffle():
+            buffer += self.tokenizer(sample["text"])["input_ids"]
+            buffer += [self.tokenizer.eos_token_id]
+            while len(buffer) > self.max_tokens:
+                input_ids = torch.tensor(buffer[: self.max_tokens])
+                yield {
+                    "input_ids": input_ids,
+                    "attention_mask": torch.ones(input_ids.size()),
+                    "labels": input_ids,
+                }
+                buffer = buffer[self.max_tokens :]
 
 
-def load_pretraining_dataset(path, tokenizer, max_tokens=2048, seed=42):
-    encode = functools.partial(encode_pretraining, tokenizer, max_tokens)
-    dataset = load_dataset(path, streaming=True, split="train")
-    dataset = dataset.shuffle(seed=seed, buffer_size=10_000)
-    # TODO dynamically figure out which columns/features to remove
-    dataset = dataset.map(encode, batched=True, remove_columns=["text", "meta"])
-    return dataset
+def load_pretraining_dataset(path, tokenizer, max_tokens=2048):
+    return PretrainingDatasetWrapper(tokenizer, path, max_tokens=max_tokens)

src/axolotl/utils/models.py

@@ -11,16 +11,22 @@ import bitsandbytes as bnb
 import torch
 import transformers
 from optimum.bettertransformer import BetterTransformer
-from transformers import (  # noqa: F401
+from transformers import PreTrainedModel  # noqa: F401
+from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
     LlamaConfig,
-    PreTrainedModel,
-    PreTrainedTokenizerBase,
 )
 
+try:
+    from transformers import LlamaForCausalLM
+except ImportError:
+    logging.warning(
+        "This version of transformers does not support Llama. Consider upgrading."
+    )
+
 from axolotl.prompt_tokenizers import LLAMA_DEFAULT_PAD_TOKEN
 
 if TYPE_CHECKING:
@@ -34,20 +40,15 @@ def load_tokenizer(
     tokenizer_type,
     cfg,
 ):
-    use_fast = True  # this is the default
-    if cfg.tokenizer_use_fast is not None:
-        use_fast = cfg.tokenizer_use_fast
     if tokenizer_type:
         tokenizer = getattr(transformers, tokenizer_type).from_pretrained(
             tokenizer_config,
             trust_remote_code=cfg.trust_remote_code or False,
-            use_fast=use_fast,
         )
     else:
         tokenizer = AutoTokenizer.from_pretrained(
             tokenizer_config,
             trust_remote_code=cfg.trust_remote_code or False,
-            use_fast=use_fast,
         )
 
     logging.debug(f"EOS: {tokenizer.eos_token_id} / {tokenizer.eos_token}")
@@ -75,58 +76,45 @@ def load_tokenizer(
 
 
 def load_model(
-    base_model, base_model_config, model_type, tokenizer, cfg, adapter="lora"
+    base_model,
+    base_model_config,
+    model_type,
+    tokenizer,
+    cfg,
+    adapter="lora",
+    inference=False,
 ):
-    # type: (str, str, str, PreTrainedTokenizerBase, DictDefault, Optional[str]) -> Tuple[PreTrainedModel, Optional[PeftConfig]]
+    # type: (str, str, str, str, DictDefault, Optional[str], bool) -> Tuple[PreTrainedModel, Optional[PeftConfig]]
     """
     Load a model from a base model and a model type.
     """
 
     # TODO refactor as a kwarg
     load_in_8bit = cfg.load_in_8bit
-    cfg.is_llama_derived_model = "llama" in base_model or (
+    is_llama_derived_model = "llama" in base_model or (
         cfg.model_type and "llama" in cfg.model_type.lower()
     )
 
-    if cfg.is_llama_derived_model and cfg.flash_attention:
-        if cfg.device not in ["mps", "cpu"] and not cfg.inference:
+    if is_llama_derived_model and cfg.flash_attention:
+        if cfg.device not in ["mps", "cpu"] and inference is False:
             from axolotl.flash_attn import replace_llama_attn_with_flash_attn
 
             logging.info("patching with flash attention")
             replace_llama_attn_with_flash_attn()
-    elif cfg.is_llama_derived_model and cfg.xformers_attention:
+    elif is_llama_derived_model and cfg.xformers_attention:
         from axolotl.monkeypatch.llama_attn_hijack_xformers import (
             hijack_llama_attention,
         )
 
         logging.info("patching with xformers attention")
         hijack_llama_attention()
-    elif cfg.is_llama_derived_model and cfg.sdp_attention:
+    elif is_llama_derived_model and cfg.sdp_attention:
         from axolotl.monkeypatch.llama_attn_hijack_xformers import (
             hijack_llama_sdp_attention,
         )
 
         logging.info("patching with sdp attention")
         hijack_llama_sdp_attention()
-    elif cfg.is_llama_derived_model and cfg.landmark_attention:
-        from axolotl.monkeypatch.llama_landmark_attn import (
-            MEM_TOKEN,
-            patch_llama_with_landmark_attn,
-        )
-
-        logging.info("patching with landmark attention")
-        patch_llama_with_landmark_attn()
-
-        # Note: This might overwrite previous additional_special_tokens
-        tokenizer.add_special_tokens({"additional_special_tokens": [MEM_TOKEN]})
-
-    if cfg.is_llama_derived_model and cfg.xpos_rope:
-        from axolotl.monkeypatch.xpos_rope_llama_monkey_patch import (
-            replace_llama_rope_with_xpos_rope,
-        )
-
-        logging.info("patching with xpos rope")
-        replace_llama_rope_with_xpos_rope()
 
     if cfg.bf16 or cfg.bfloat16:
         torch_dtype = torch.bfloat16
@@ -141,18 +129,11 @@ def load_model(
             )
 
             replace_peft_model_with_int4_lora_model()
+        from peft import prepare_model_for_int8_training
     except Exception as err:
         logging.exception(err)
         raise err
 
-    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.adapter == "qlora" and cfg.load_in_4bit:
         model_kwargs["quantization_config"] = BitsAndBytesConfig(
@@ -164,7 +145,7 @@ def load_model(
             bnb_4bit_quant_type="nf4",
         )
     try:
-        if cfg.gptq and cfg.is_llama_derived_model:
+        if cfg.gptq and is_llama_derived_model:
             from alpaca_lora_4bit.autograd_4bit import load_llama_model_4bit_low_ram
             from huggingface_hub import snapshot_download
 
@@ -202,9 +183,7 @@ def load_model(
                 else True,
             )
             load_in_8bit = False
-        elif cfg.is_llama_derived_model:
-            from transformers import LlamaForCausalLM
-
+        elif is_llama_derived_model and "LlamaForCausalLM" in globals():
             config = LlamaConfig.from_pretrained(base_model_config)
             model = LlamaForCausalLM.from_pretrained(
                 base_model,
@@ -256,22 +235,6 @@ def load_model(
                 base_model,
                 trust_remote_code=cfg.trust_remote_code or False,
             )
-            # Shouldn't be a problem most of the time. will obviously error if the model doesn't support this
-            # when training starts
-            if (
-                hasattr(config, "max_seq_len")
-                and config.max_seq_len
-                and cfg.sequence_len > config.max_seq_len
-            ):
-                config.max_seq_len = cfg.sequence_len
-                logging.warning(f"increasing context length to {cfg.sequence_len}")
-            elif (
-                hasattr(config, "max_sequence_length")
-                and config.max_sequence_length
-                and cfg.sequence_len > config.max_sequence_length
-            ):
-                config.max_sequence_length = cfg.sequence_len
-                logging.warning(f"increasing context length to {cfg.sequence_len}")
             model = AutoModelForCausalLM.from_pretrained(
                 base_model,
                 config=config,
@@ -290,7 +253,6 @@ def load_model(
         model = AutoModelForCausalLM.from_pretrained(
             base_model,
             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=torch_dtype,
             device_map=cfg.device_map,
             trust_remote_code=cfg.trust_remote_code or False,
@@ -300,11 +262,7 @@ def load_model(
     embeddings_len = math.ceil(len(tokenizer) / 32) * 32
     model.resize_token_embeddings(embeddings_len)
 
-    if (
-        hasattr(model.config, "max_position_embeddings")
-        and model.config.max_position_embeddings
-        and cfg.sequence_len >= model.config.max_position_embeddings
-    ):
+    if cfg.sequence_len >= model.config.max_position_embeddings:
         logging.warning(
             f"increasing model.config.max_position_embeddings to {cfg.sequence_len}"
         )
@@ -314,10 +272,8 @@ def load_model(
         (cfg.adapter == "lora" and load_in_8bit)
         or (cfg.adapter == "qlora" and cfg.load_in_4bit)
     ):
-        logging.info("converting PEFT model w/ prepare_model_for_kbit_training")
-        model = prepare_model_for_kbit_training(
-            model, use_gradient_checkpointing=cfg.gradient_checkpointing
-        )
+        logging.info("converting PEFT model w/ prepare_model_for_int8_training")
+        model = prepare_model_for_int8_training(model)
 
     model, lora_config = load_adapter(model, cfg, adapter)
 
@@ -390,6 +346,7 @@ def load_llama_adapter(model, cfg):
         model = PeftModel.from_pretrained(
             model,
             cfg.lora_model_dir,
+            device_map=cfg.device_map,
             torch_dtype=torch.float16,
         )
     else:
@@ -451,7 +408,8 @@ def load_lora(model, cfg):
         model = PeftModel.from_pretrained(
             model,
             cfg.lora_model_dir,
-            is_trainable=not cfg.inference,
+            device_map=cfg.device_map,
+            # torch_dtype=torch.float16,
         )
     else:
         model = get_peft_model(model, lora_config)

src/axolotl/utils/tokenization.py

@@ -34,5 +34,3 @@ def check_example_labels(example, tokenizer):
 
     logging.info(" ".join(colored_tokens))
     logging.info("\n\n\n")
-
-    return " ".join(colored_tokens)

src/axolotl/utils/trainer.py

@@ -66,6 +66,8 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
         if cfg.logging_steps is not None
         else max(min(int(0.005 * total_num_steps), 10), 1)
     )
+    save_steps = cfg.save_steps
+    eval_steps = cfg.eval_steps
 
     training_arguments_kwargs = {}
     if cfg.bf16 == "full":
@@ -76,10 +78,6 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
     training_arguments_kwargs["tf32"] = cfg.tf32
     training_arguments_kwargs["warmup_steps"] = warmup_steps
     training_arguments_kwargs["logging_steps"] = logging_steps
-
-    if cfg.seed:
-        training_arguments_kwargs["seed"] = cfg.seed
-
     if cfg.gradient_checkpointing:
         if cfg.gptq:
             from alpaca_lora_4bit.gradient_checkpointing import (
@@ -115,19 +113,6 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
             # TODO search Path("./") for one
             training_arguments_kwargs["deepspeed"] = "./ds_config.json"
 
-    if cfg.adam_beta1:
-        training_arguments_kwargs["adam_beta1"] = cfg.adam_beta1
-    if cfg.adam_beta2:
-        training_arguments_kwargs["adam_beta2"] = cfg.adam_beta2
-    if cfg.adam_epsilon:
-        training_arguments_kwargs["adam_epsilon"] = cfg.adam_epsilon
-    if cfg.max_grad_norm:
-        training_arguments_kwargs["max_grad_norm"] = cfg.max_grad_norm
-
-    if cfg.push_to_hub_model_id:
-        training_arguments_kwargs["push_to_hub_model_id"] = cfg.push_to_hub_model_id
-        training_arguments_kwargs["push_to_hub"] = True
-
     training_args = transformers.TrainingArguments(
         per_device_train_batch_size=cfg.micro_batch_size,
         per_device_eval_batch_size=cfg.eval_batch_size
@@ -138,16 +123,16 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
         num_train_epochs=cfg.num_epochs,
         learning_rate=cfg.learning_rate,
         evaluation_strategy="steps" if cfg.val_set_size > 0 else "no",
-        save_strategy="steps" if cfg.save_steps else "epoch",
-        eval_steps=cfg.eval_steps if cfg.val_set_size > 0 else None,
-        save_steps=cfg.save_steps,
+        save_strategy="steps" if save_steps else "epoch",
+        eval_steps=eval_steps if cfg.val_set_size > 0 else None,
+        save_steps=save_steps,
         output_dir=cfg.output_dir,
         save_total_limit=3,
         load_best_model_at_end=(
             cfg.load_best_model_at_end is not False
             and cfg.val_set_size > 0
-            and cfg.save_steps
-            and cfg.save_steps % cfg.eval_steps == 0
+            and save_steps
+            and save_steps % eval_steps == 0
             and cfg.load_in_8bit is not True
         )
         or False,
@@ -245,6 +230,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
         callbacks.append(SavePeftModelCallback)
 
     if hasattr(model, "use_bettertransformer") and model.use_bettertransformer is True:
+        logging.info("Setting up SaveBetterTransformerModelCallback.")
         callbacks.append(SaveBetterTransformerModelCallback)
 
     data_collator_kwargs = {
@@ -255,26 +241,6 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
     else:
         data_collator_kwargs["pad_to_multiple_of"] = 8
 
-    if cfg.is_llama_derived_model and cfg.landmark_attention:
-        from functools import partial
-
-        from axolotl.monkeypatch.llama_landmark_attn import (
-            add_mem_tokens,
-            get_mem_id,
-            set_model_mem_id,
-        )
-
-        set_model_mem_id(model, tokenizer)
-
-        logging.info("Adding landmark attention tokens to dataset")
-
-        for dataset in [train_dataset, eval_dataset]:
-            dataset = dataset.map(
-                partial(add_mem_tokens, mem_freq=50, mem_id=get_mem_id(tokenizer)),
-                batched=False,
-                num_proc=32,
-            )
-
     trainer_cls = (
         OneCycleLRSchedulerTrainer
         if cfg.lr_scheduler == "one_cycle" and (cfg.fsdp or cfg.adapter == "qlora")

src/axolotl/utils/validation.py

@@ -10,12 +10,6 @@ def validate_config(cfg):
         raise ValueError(
             "please set only one of gradient_accumulation_steps or batch_size"
         )
-    if cfg.batch_size:
-        logging.warning(
-            "%s\n%s",
-            "batch_size is not recommended. Please use gradient_accumulation_steps instead.",
-            "To calculate the equivalent gradient_accumulation_steps, divide batch_size / micro_batch_size / number of gpus.",
-        )
     if cfg.load_4bit:
         raise ValueError(
             "cfg.load_4bit parameter has been deprecated and replaced by cfg.gptq"
@@ -56,14 +50,6 @@ def validate_config(cfg):
             "Require cfg.hf_use_auth_token to be True for push_dataset_to_hub"
         )
 
-    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.base_model and "mpt" in cfg.base_model.lower()
-    ) and cfg.gradient_checkpointing:
-        raise ValueError("gradient_checkpointing is not supported for MPT models")
-
     if cfg.flash_optimum is True:
         if cfg.adapter:
             logging.warning(
@@ -81,17 +67,6 @@ def validate_config(cfg):
             raise ValueError(
                 f"flash_optimum for BetterTransformers may not be used with {torch.__version__}"
             )
-
-    if cfg.pretraining_dataset and cfg.group_by_length:
-        logging.warning(
-            "You probably want to disable group_by_length as it will force a streamed dataset to download completely."
-        )
-
-    if any([cfg.adamw_beta1, cfg.adamw_beta2, cfg.adamw_epsilon]) and (
-        not cfg.optimizer or "adamw" not in cfg.optimizer
-    ):
-        logging.warning("adamw hyperparameters found, but no adamw optimizer set")
-
     # TODO
     # MPT 7b
     # https://github.com/facebookresearch/bitsandbytes/issues/25

src/axolotl/utils/wandb.py

@@ -15,5 +15,3 @@ def setup_wandb_env_vars(cfg):
             os.environ["WANDB_LOG_MODEL"] = cfg.wandb_log_model
         if cfg.wandb_run_id and len(cfg.wandb_run_id) > 0:
             os.environ["WANDB_RUN_ID"] = cfg.wandb_run_id
-    else:
-        os.environ["WANDB_DISABLED"] = "true"

tests/test_prompt_tokenizers.py

@@ -6,16 +6,8 @@ from pathlib import Path
 
 from transformers import AutoTokenizer
 
-from axolotl.prompt_strategies.alpaca_chat import NoSystemPrompter
-from axolotl.prompt_strategies.alpaca_w_system import (
-    InstructionWSystemPromptTokenizingStrategy,
-    SystemDataPrompter,
-)
-from axolotl.prompt_tokenizers import (
-    AlpacaPromptTokenizingStrategy,
-    ShareGPTPromptTokenizingStrategy,
-)
-from axolotl.prompters import AlpacaPrompter, PromptStyle, ShareGPTPrompter
+from axolotl.prompt_tokenizers import ShareGPTPromptTokenizingStrategy
+from axolotl.prompters import ShareGPTPrompter
 
 logging.basicConfig(level="INFO")
 
@@ -37,6 +29,7 @@ class TestPromptTokenizationStrategies(unittest.TestCase):
         )
 
     def test_sharegpt_integration(self):
+        print(Path(__file__).parent)
         with open(
             Path(__file__).parent / "fixtures/conversation.json", encoding="utf-8"
         ) as fin:
@@ -60,79 +53,6 @@ class TestPromptTokenizationStrategies(unittest.TestCase):
             self.assertEqual(len(example[fields]), len(tokenized_conversation[fields]))
             self.assertEqual(example[fields], tokenized_conversation[fields])
 
-    def test_no_sys_prompt(self):
-        """
-        tests the interface between the user and assistant parts
-        """
-        prompter = NoSystemPrompter()
-        # pylint: disable=duplicate-code
-        strat = AlpacaPromptTokenizingStrategy(
-            prompter,
-            self.tokenizer,
-            False,
-            2048,
-        )
-        sample = {
-            "instruction": "hello cruel. lorem ipsum dolor sit amet.",
-            "output": "world!",
-        }
-        example = strat.tokenize_prompt(sample)
-        world_idx = example["input_ids"].index(3186)
-        assert example["labels"][world_idx] == 3186
-        assert example["labels"][world_idx - 1] == -100
-
-    def test_alpaca(self):
-        """
-        tests the interface between the user and assistant parts
-        """
-        # pylint: disable=duplicate-code
-        prompter = AlpacaPrompter()
-        strat = AlpacaPromptTokenizingStrategy(
-            prompter,
-            self.tokenizer,
-            False,
-            2048,
-        )
-        sample = {"instruction": "hello!", "output": "Hi! How can I help?"}
-        example = strat.tokenize_prompt(sample)
-        world_idx = example["input_ids"].index(6324)
-        assert example["labels"][world_idx] == 6324
-        assert example["labels"][world_idx - 1] == -100
-
-
-class InstructionWSystemPromptTokenizingStrategyTest(unittest.TestCase):
-    """
-    Test class for prompt tokenization strategies with sys prompt from the dataset
-    """
-
-    def setUp(self) -> None:
-        # pylint: disable=duplicate-code
-        self.tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
-        self.tokenizer.add_special_tokens(
-            {
-                "bos_token": "<s>",
-                "eos_token": "</s>",
-                "unk_token": "<unk>",
-            }
-        )
-
-    def test_system_alpaca(self):
-        prompter = SystemDataPrompter(PromptStyle.CHAT.value)
-        strat = InstructionWSystemPromptTokenizingStrategy(
-            prompter,
-            self.tokenizer,
-            False,
-            2048,
-        )
-        sample = {
-            "system": "use cot",
-            "instruction": "hello!",
-            "output": "Hi! How can I help?",
-        }
-        example = strat.tokenize_prompt(sample)
-        assert example["input_ids"][0:3] == [1, 671, 20118]  # <s>use cot
-        assert example["input_ids"][3] == 11889  # USER
-
 
 if __name__ == "__main__":
     unittest.main()

tests/test_prompters.py

@@ -2,13 +2,7 @@
 
 import unittest
 
-from axolotl.prompt_strategies.alpaca_w_system import SystemDataPrompter
-from axolotl.prompters import (
-    AlpacaPrompter,
-    MultipleChoiceExplainPrompter,
-    PromptStyle,
-    UnpromptedPrompter,
-)
+from axolotl.prompters import AlpacaPrompter, PromptStyle
 
 
 class AlpacaPrompterTest(unittest.TestCase):
@@ -61,64 +55,3 @@ class AlpacaPrompterTest(unittest.TestCase):
         assert "### Response:" not in res
         assert "USER:" in res
         assert "ASSISTANT:" in res
-
-    def test_system_prompt(self):
-        prompter = SystemDataPrompter(prompt_style=PromptStyle.CHAT.value)
-        res = next(
-            prompter.build_prompt_w_system(
-                "use cot", "tell me a joke about the following", "alpacas"
-            )
-        )
-        assert "use cot" in res
-        assert res.startswith("use cot")
-        assert "### Instruction:" not in res
-        assert "### Input:" not in res
-        assert "alpacas" in res
-        assert "### Response:" not in res
-        assert "USER:" in res
-        assert "ASSISTANT:" in res
-
-
-class UnpromptedPrompterTest(unittest.TestCase):
-    """
-    Test class for UnpromptedPrompter with no system prompts
-    """
-
-    def test_prompt_style_w_none(self):
-        prompter = UnpromptedPrompter(prompt_style=None)
-        res = next(prompter.build_prompt("tell me a joke"))
-        assert "### Instruction:" in res
-        assert "tell me a joke" in res
-        assert res.startswith("###")
-
-    def test_prompt_style_w_instruct(self):
-        prompter = UnpromptedPrompter(prompt_style=PromptStyle.INSTRUCT.value)
-        res = next(
-            prompter.build_prompt("tell me a joke about the following", "alpacas")
-        )
-        assert "### Instruction:" in res
-        assert "tell me a joke" in res
-        assert res.startswith("###")
-
-    def test_prompt_style_w_chat(self):
-        prompter = UnpromptedPrompter(prompt_style=PromptStyle.CHAT.value)
-        res = next(
-            prompter.build_prompt("tell me a joke about the following", "alpacas")
-        )
-        assert "USER:" in res
-        assert "tell me a joke" in res
-        assert res.startswith("USER:")
-
-
-class MultipleChoiceExplainPrompterTest(unittest.TestCase):
-    """
-    Test class for MultipleChoiceExplainPrompter
-    """
-
-    def test_prompt_style_w_chat(self):
-        prompter = MultipleChoiceExplainPrompter(prompt_style=PromptStyle.CHAT.value)
-        res = next(prompter.build_prompt("choose one", "- A\n- B\n- C", "C"))
-        assert "USER:" in res
-        assert "choose one" in res
-        assert "Choose the answer that best answers the question." in res
-        assert "- A\n- B\n- C" in res

tests/test_tokenizers.py

@@ -1,31 +0,0 @@
-"""
-Test cases for the tokenizer loading
-"""
-import unittest
-
-from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_tokenizer
-
-
-class TestTokenizers(unittest.TestCase):
-    """
-    test class for the load_tokenizer fn
-    """
-
-    def test_default_use_fast(self):
-        cfg = DictDefault({})
-        tokenizer = load_tokenizer("huggyllama/llama-7b", None, cfg)
-        assert "Fast" in tokenizer.__class__.__name__
-
-    def test_dont_use_fast(self):
-        cfg = DictDefault(
-            {
-                "tokenizer_use_fast": False,
-            }
-        )
-        tokenizer = load_tokenizer("huggyllama/llama-7b", None, cfg)
-        assert "Fast" not in tokenizer.__class__.__name__
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/test_validation.py

@@ -1,8 +1,6 @@
 """Module for testing the validation module"""
 
-import logging
 import unittest
-from typing import Optional
 
 import pytest
 
@@ -15,12 +13,6 @@ class ValidationTest(unittest.TestCase):
     Test the validation module
     """
 
-    _caplog: Optional[pytest.LogCaptureFixture] = None
-
-    @pytest.fixture(autouse=True)
-    def inject_fixtures(self, caplog):
-        self._caplog = caplog
-
     def test_load_4bit_deprecate(self):
         cfg = DictDefault(
             {
@@ -31,17 +23,6 @@ class ValidationTest(unittest.TestCase):
         with pytest.raises(ValueError):
             validate_config(cfg)
 
-    def test_batch_size_unused_warning(self):
-        cfg = DictDefault(
-            {
-                "batch_size": 32,
-            }
-        )
-
-        with self._caplog.at_level(logging.WARNING):
-            validate_config(cfg)
-            assert "batch_size is not recommended" in self._caplog.records[0].message
-
     def test_qlora(self):
         base_cfg = DictDefault(
             {
@@ -165,151 +146,3 @@ class ValidationTest(unittest.TestCase):
         )
 
         validate_config(cfg)
-
-    def test_falcon_fsdp(self):
-        regex_exp = r".*FSDP is not supported for falcon models.*"
-
-        # Check for lower-case
-        cfg = DictDefault(
-            {
-                "base_model": "tiiuae/falcon-7b",
-                "fsdp": ["full_shard", "auto_wrap"],
-            }
-        )
-
-        with pytest.raises(ValueError, match=regex_exp):
-            validate_config(cfg)
-
-        # Check for upper-case
-        cfg = DictDefault(
-            {
-                "base_model": "Falcon-7b",
-                "fsdp": ["full_shard", "auto_wrap"],
-            }
-        )
-
-        with pytest.raises(ValueError, match=regex_exp):
-            validate_config(cfg)
-
-        cfg = DictDefault(
-            {
-                "base_model": "tiiuae/falcon-7b",
-            }
-        )
-
-        validate_config(cfg)
-
-    def test_mpt_gradient_checkpointing(self):
-        regex_exp = r".*gradient_checkpointing is not supported for MPT models*"
-
-        # Check for lower-case
-        cfg = DictDefault(
-            {
-                "base_model": "mosaicml/mpt-7b",
-                "gradient_checkpointing": True,
-            }
-        )
-
-        with pytest.raises(ValueError, match=regex_exp):
-            validate_config(cfg)
-
-    def test_flash_optimum(self):
-        cfg = DictDefault(
-            {
-                "flash_optimum": True,
-                "adapter": "lora",
-            }
-        )
-
-        with self._caplog.at_level(logging.WARNING):
-            validate_config(cfg)
-            assert any(
-                "BetterTransformers probably doesn't work with PEFT adapters"
-                in record.message
-                for record in self._caplog.records
-            )
-
-        cfg = DictDefault(
-            {
-                "flash_optimum": True,
-            }
-        )
-
-        with self._caplog.at_level(logging.WARNING):
-            validate_config(cfg)
-            assert any(
-                "probably set bfloat16 or float16" in record.message
-                for record in self._caplog.records
-            )
-
-        cfg = DictDefault(
-            {
-                "flash_optimum": True,
-                "fp16": True,
-            }
-        )
-        regex_exp = r".*AMP is not supported.*"
-
-        with pytest.raises(ValueError, match=regex_exp):
-            validate_config(cfg)
-
-        cfg = DictDefault(
-            {
-                "flash_optimum": True,
-                "bf16": True,
-            }
-        )
-        regex_exp = r".*AMP is not supported.*"
-
-        with pytest.raises(ValueError, match=regex_exp):
-            validate_config(cfg)
-
-    def test_adamw_hyperparams(self):
-        cfg = DictDefault(
-            {
-                "optimizer": None,
-                "adamw_epsilon": 0.0001,
-            }
-        )
-
-        with self._caplog.at_level(logging.WARNING):
-            validate_config(cfg)
-            assert any(
-                "adamw hyperparameters found, but no adamw optimizer set"
-                in record.message
-                for record in self._caplog.records
-            )
-
-        cfg = DictDefault(
-            {
-                "optimizer": "adafactor",
-                "adamw_beta1": 0.0001,
-            }
-        )
-
-        with self._caplog.at_level(logging.WARNING):
-            validate_config(cfg)
-            assert any(
-                "adamw hyperparameters found, but no adamw optimizer set"
-                in record.message
-                for record in self._caplog.records
-            )
-
-        cfg = DictDefault(
-            {
-                "optimizer": "adamw_bnb_8bit",
-                "adamw_beta1": 0.0001,
-                "adamw_beta2": 0.0001,
-                "adamw_epsilon": 0.0001,
-            }
-        )
-
-        validate_config(cfg)
-
-        cfg = DictDefault(
-            {
-                "optimizer": "adafactor",
-            }
-        )
-
-        validate_config(cfg)