more gpt-neox long ctx fixes

fix bettertransformers save, force it to skip after saving correctly in callback
more tweaks to do pre-training with bettertransformers
2023-06-01 08:20:08 -04:00 · 2023-06-01 00:33:13 -04:00 · 2023-05-31 21:59:15 -04:00 · 2023-05-31 16:51:19 -04:00 · 2023-05-31 16:41:24 -04:00 · 2023-05-31 16:41:21 -04:00
47 changed files with 872 additions and 2327 deletions
--- a/FAQS.md
+++ b/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.
--- a/README.md
+++ b/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

@@ -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:
@@ -390,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
@@ -428,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:
@@ -500,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
@@ -527,7 +397,6 @@ Add below flag to train command above
 Please reduce any below
  - `micro_batch_size`
  - `eval_batch_size`
-  - `gradient_accumulation_steps`
  - `sequence_len`

 > RuntimeError: expected scalar type Float but found Half
@@ -538,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

@@ -552,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).
--- a/configs/accelerate/default_config.yaml
+++ b/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
--- a/configs/cerebras_1_3B_alpaca.yml
+++ b/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:
--- a/configs/galactica_1_3B.yml
+++ b/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>"
--- a/configs/llama_13B_alpaca.yml
+++ b/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:
--- a/configs/llama_65B_alpaca.yml
+++ b/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:
--- a/configs/llama_7B_4bit.yml
+++ b/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:
--- a/configs/llama_7B_alpaca.yml
+++ b/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:
--- a/examples/jeopardy-bot/config.yml
+++ b/examples/jeopardy-bot/config.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>"
--- a/configs/pythia_1_2B_alpaca.yml
+++ b/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
--- a/configs/quickstart.yml
+++ b/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:
--- a/configs/sample.yml
+++ b/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:
--- a/configs/stability_3b.yml
+++ b/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>"
--- a/configs/vicuna_13B_4bit_reflect.yml
+++ b/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
--- a/docker-compose.yaml
+++ b/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
--- a/docker/Dockerfile
+++ b/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 \
--- a/examples/cerebras/qlora.yml
+++ b/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|>"
--- a/examples/falcon/config-7b-lora.yml
+++ b/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:
--- a/examples/falcon/config-7b-qlora.yml
+++ b/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|>"
--- a/examples/falcon/config-7b.yml
+++ b/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:
--- a/examples/gptj/qlora.yml
+++ b/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|>"
--- a/examples/gptq-lora-7b/README.md
+++ b/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

 ```
--- a/examples/lora-openllama-3b/config.yml
+++ b/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:
--- a/examples/openllama-3b/README.md
+++ b/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
-```
--- a/examples/openllama-3b/config.yml
+++ b/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>"
--- a/examples/pythia-12b/README.md
+++ b/examples/pythia-12b/README.md
@@ -0,0 +1,10 @@
+# Python 12B
+
+- Single-GPU A100 only (?)
+
+```shell
+python scripts/finetune.py examples/pythia-12b/config.yml
+```
+
+⚠️ Multiple-GPU A100 - Doesn't seem to work with multi-gpu without causing OOM! ⚠️
+
--- a/examples/pythia-12b/config.yml
+++ b/examples/pythia-12b/config.yml
@@ -0,0 +1,49 @@
+base_model: EleutherAI/pythia-12b-deduped
+base_model_config: EleutherAI/pythia-12b-deduped
+base_model_ignore_patterns: pytorch*  # prefer safetensors
+model_type: GPTNeoXForCausalLM
+tokenizer_type: AutoTokenizer
+load_in_8bit: false
+load_in_4bit: false
+gptq: false
+device_map: auto
+datasets:
+  - path: vicgalle/alpaca-gpt4
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.05
+adapter:
+lora_model_dir:
+sequence_len: 2048
+max_packed_sequence_len: 2048
+lora_r: 64
+lora_alpha: 32
+lora_dropout: 0.0
+lora_target_modules:
+lora_target_linear: true
+lora_fan_in_fan_out: true  # pythia/GPTNeoX lora specific
+wandb_project: pythia-12b
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+output_dir: ./pythia-12b
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 5
+learning_rate: 0.00003
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+train_on_inputs: false
+group_by_length: false
+bf16: false
+fp16: false
+float16: true
+tf32: true
+flash_optimum: true
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:
+gradient_checkpointing: true
+fsdp:
+fsdp_transformer_layer_cls_to_wrap:
+collator_pad_to_longest: true
--- a/examples/qlora-openllama-3b/README.md
+++ b/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
+
+```
--- a/examples/qlora-openllama-3b/config.yml
+++ b/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
--- a/requirements.txt
+++ b/requirements.txt
@@ -11,6 +11,7 @@ sentencepiece
 wandb
 einops
 xformers
+optimum
 # qlora things
 bert-score==0.3.13
 evaluate==0.4.0
--- a/scripts/finetune.py
+++ b/scripts/finetune.py
@@ -12,13 +12,15 @@ from typing import Any, Dict, List, Optional, Union
 import fire
 import torch
 import yaml
-from transformers import GenerationConfig, TextStreamer
-
-from axolotl.utils.data import load_prepare_datasets
-from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_model, load_tokenizer

 # 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
+
+from axolotl.utils.data import load_prepare_datasets, load_pretraining_dataset
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.models import load_model, load_tokenizer
 from axolotl.utils.tokenization import check_dataset_labels
 from axolotl.utils.trainer import setup_trainer
 from axolotl.utils.validation import validate_config
@@ -48,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}
@@ -63,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(
@@ -118,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]))


@@ -174,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])
@@ -206,20 +185,28 @@ 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
-        train_dataset, eval_dataset = load_prepare_datasets(
-            tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH
-        )
+        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(
+                pretraining_dataset, tokenizer, max_tokens=cfg.sequence_len
+            )
+            train_dataset = Dataset.from_list(list(train_dataset))
+            eval_dataset = None

    if cfg.debug or "debug" in kwargs:
        logging.info("check_dataset_labels...")
@@ -243,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:
@@ -255,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
@@ -285,12 +282,15 @@ def train(

    # In case we want to stop early with ctrl+c, this is a nice to have to save the pretrained model
    if cfg.local_rank == 0:
+
+        def terminate_handler(_, __, model):
+            if cfg.flash_optimum:
+                model = BetterTransformer.reverse(model)
+            model.save_pretrained(cfg.output_dir)
+            sys.exit(0)
+
        signal.signal(
-            signal.SIGINT,
-            lambda signal, frame: (
-                model.save_pretrained(cfg.output_dir),
-                sys.exit(0),
-            ),
+            signal.SIGINT, lambda signum, frame: terminate_handler(signum, frame, model)
        )

    logging.info("Starting trainer...")
@@ -311,15 +311,21 @@ 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)
-    trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+    if cfg.flash_optimum:
+        with torch.backends.cuda.sdp_kernel(
+            enable_flash=True, enable_math=True, enable_mem_efficient=True
+        ):
+            trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+    else:
+        trainer.train(resume_from_checkpoint=resume_from_checkpoint)

    logging.info(f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}")

    # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
    # only save on rank 0, otherwise it corrupts output on multi-GPU when multiple processes attempt to write the same file
    if cfg.local_rank == 0:
+        if cfg.flash_optimum:
+            model = BetterTransformer.reverse(model)
        model.save_pretrained(cfg.output_dir)

    # trainer.save_model(cfg.output_dir)  # TODO this may be needed for deepspeed to work? need to review another time
--- a/src/axolotl/datasets.py
+++ b/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
--- a/src/axolotl/monkeypatch/llama_landmark_attn.py
+++ b/src/axolotl/monkeypatch/llama_landmark_attn.py
--- a/src/axolotl/monkeypatch/xpos_rope_llama_monkey_patch.py
+++ b/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
--- a/src/axolotl/prompt_strategies/alpaca_chat.py
+++ b/src/axolotl/prompt_strategies/alpaca_chat.py
@@ -18,15 +18,6 @@ def load(tokenizer, cfg):
    )


-class AlpacaConcisePrompter(AlpacaPrompter):
-    """
-    Alpaca Prompter extending the system prompt to ask for concise answers
-    """
-
-    system_prompt = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that concisely and appropriately completes the request.\n\n"
-    system_no_input_prompt = "Below is an instruction that describes a task. Write a response that appropriately and concisely completes the request.\n\n"
-
-
 class AlpacaQAPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
    """
    Tokenizing strategy for AlpacaQA
@@ -40,28 +31,6 @@ 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(
        AlpacaPrompter(PromptStyle.CHAT.value),
@@ -69,12 +38,3 @@ def load_qa(tokenizer, cfg):
        cfg.train_on_inputs,
        cfg.sequence_len,
    )
-
-
-def load_camel_ai(tokenizer, cfg):
-    return CamelAIPromptTokenizingStrategy(
-        AlpacaPrompter(PromptStyle.CHAT.value),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
--- a/src/axolotl/prompt_strategies/context_qa.py
+++ b/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.",
-        )
--- a/src/axolotl/prompt_strategies/sharegpt_jokes.py
+++ b/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"]},
-        ]
--- a/src/axolotl/prompt_strategies/sharegpt_simple.py
+++ b/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
--- a/src/axolotl/prompters.py
+++ b/src/axolotl/prompters.py
@@ -261,33 +261,28 @@ 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:
@@ -305,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:
--- a/src/axolotl/utils/callbacks.py
+++ b/src/axolotl/utils/callbacks.py
@@ -2,13 +2,14 @@

 import os

+from optimum.bettertransformer import BetterTransformer
 from transformers import (
    TrainerCallback,
    TrainerControl,
    TrainerState,
    TrainingArguments,
 )
-from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, IntervalStrategy


 class SavePeftModelCallback(TrainerCallback):  # pylint: disable=too-few-public-methods
@@ -30,3 +31,39 @@ class SavePeftModelCallback(TrainerCallback):  # pylint: disable=too-few-public-
        kwargs["model"].save_pretrained(peft_model_path)

        return control
+
+
+class SaveBetterTransformerModelCallback(
+    TrainerCallback
+):  # pylint: disable=too-few-public-methods
+    """Callback to save the BetterTransformer wrapped model"""
+
+    def on_step_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        # Save
+        if (
+            args.save_strategy == IntervalStrategy.STEPS
+            and args.save_steps > 0
+            and state.global_step % args.save_steps == 0
+        ):
+            control.should_save = True
+
+        if control.should_save:
+            checkpoint_folder = os.path.join(
+                args.output_dir,
+                f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}",
+            )
+
+            model = BetterTransformer.reverse(kwargs["model"])
+            model.save_pretrained(checkpoint_folder)
+            # FIXME - need to cleanup old checkpoints
+
+            # since we're saving here, we don't need the trainer loop to attempt to save too b/c
+            # the trainer will raise an exception since it can't save a BetterTransformer wrapped model
+            control.should_save = False
+        return control
--- a/src/axolotl/utils/data.py
+++ b/src/axolotl/utils/data.py
@@ -5,7 +5,8 @@ from hashlib import md5
 from pathlib import Path
 from typing import List, Tuple, Union

-from datasets import Dataset, DatasetDict, load_dataset, load_from_disk
+import torch
+from datasets import Dataset, DatasetDict, IterableDataset, load_dataset, load_from_disk
 from huggingface_hub import hf_hub_download
 from transformers import PreTrainedTokenizerBase

@@ -78,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:
@@ -134,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]
@@ -239,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}"
@@ -394,8 +381,43 @@ def load_prepare_datasets(
            index=cfg.dataset_shard_idx,
        )

-    dataset = dataset.train_test_split(test_size=cfg.val_set_size, shuffle=False)
-    train_dataset = dataset["train"]
-    eval_dataset = dataset["test"]
+    if cfg.val_set_size:
+        dataset = dataset.train_test_split(test_size=cfg.val_set_size, shuffle=False)
+        train_dataset = dataset["train"]
+        eval_dataset = dataset["test"]
+    else:
+        train_dataset = dataset
+        eval_dataset = None

    return train_dataset, eval_dataset
+
+
+class PretrainingDatasetWrapper(IterableDataset):
+    """
+    Wrapper for pretraining dataset that avoids loading the dataset into memory
+    """
+
+    def __init__(self, tokenizer, dataset_path, max_tokens=2048):
+        self.tokenizer = tokenizer
+        self.dataset_path = dataset_path
+        self.max_tokens = max_tokens
+
+    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):
+    return PretrainingDatasetWrapper(tokenizer, path, max_tokens=max_tokens)
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -10,8 +10,9 @@ from typing import TYPE_CHECKING, Optional, Tuple  # noqa: F401
 import bitsandbytes as bnb
 import torch
 import transformers
+from optimum.bettertransformer import BetterTransformer
 from transformers import PreTrainedModel  # noqa: F401
-from transformers import (  # noqa: F401
+from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
@@ -19,6 +20,13 @@ from transformers import (  # noqa: F401
    LlamaConfig,
 )

+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:
@@ -68,62 +76,49 @@ 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, AutoTokenizer, 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:
+    if cfg.bf16 or cfg.bfloat16:
        torch_dtype = torch.bfloat16
-    elif cfg.load_in_8bit or cfg.fp16:
+    elif cfg.load_in_8bit or cfg.fp16 or cfg.float16:
        torch_dtype = torch.float16
    else:
        torch_dtype = torch.float32
@@ -134,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(
@@ -157,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

@@ -195,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,
@@ -249,17 +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 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 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,
@@ -287,14 +262,18 @@ def load_model(
    embeddings_len = math.ceil(len(tokenizer) / 32) * 32
    model.resize_token_embeddings(embeddings_len)

+    if cfg.sequence_len >= model.config.max_position_embeddings:
+        logging.warning(
+            f"increasing model.config.max_position_embeddings to {cfg.sequence_len}"
+        )
+        model.config.max_position_embeddings = cfg.sequence_len
+
    if not cfg.gptq and (
        (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)

@@ -332,6 +311,9 @@ def load_model(
        logging.warning("there are no parameters that require gradient updates")
    model.config.use_cache = False

+    if cfg.flash_optimum:
+        model = BetterTransformer.transform(model)
+
    # TODO resume_from_checkpoint handling
    return model, lora_config

@@ -364,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:
@@ -425,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)
--- a/src/axolotl/utils/trainer.py
+++ b/src/axolotl/utils/trainer.py
@@ -16,7 +16,10 @@ from torch.optim.lr_scheduler import OneCycleLR
 from transformers import EarlyStoppingCallback, Trainer
 from transformers.trainer_pt_utils import get_parameter_names

-from axolotl.utils.callbacks import SavePeftModelCallback
+from axolotl.utils.callbacks import (
+    SaveBetterTransformerModelCallback,
+    SavePeftModelCallback,
+)
 from axolotl.utils.schedulers import InterpolatingLogScheduler


@@ -63,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":
@@ -73,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 (
@@ -122,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,
@@ -228,6 +229,10 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
    ]:  # only save in rank 0
        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 = {
        "padding": True,
    }
@@ -236,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")
--- a/src/axolotl/utils/validation.py
+++ b/src/axolotl/utils/validation.py
@@ -2,18 +2,14 @@

 import logging

+import torch
+

 def validate_config(cfg):
    if cfg.gradient_accumulation_steps and cfg.batch_size:
        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"
@@ -54,15 +50,31 @@ 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(
+                "BetterTransformers probably doesn't work with PEFT adapters"
+            )
+        if cfg.fp16 or cfg.bf16:
+            raise ValueError("AMP is not supported with BetterTransformer")
+        if cfg.float16 is not True and cfg.bloat16 is not True:
+            logging.warning(
+                "You should probably set bfloat16 or float16 to true to "
+                "load the model in float16 for BetterTransformers"
+            )
+        if int(torch.__version__.split(".")[0]) < 2:
+            logging.warning("torch>=2.0.0 required")
+            raise ValueError(
+                f"flash_optimum for BetterTransformers may not be used with {torch.__version__}"
+            )
    # TODO
    # MPT 7b
    # https://github.com/facebookresearch/bitsandbytes/issues/25
-    # no 8bit adamw w bf16
+    # no 8bit adaAmw w bf16
+
+    # GPT-NeoX
+    # evals broken when extending context len
+    # File "/root/miniconda3/envs/py3.9/lib/python3.9/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py", line 162, in forward                        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
+    # File "/root/miniconda3/envs/py3.9/lib/python3.9/site-packages/optimum/bettertransformer/models/attention.py", line 74, in gpt2_wrapped_scaled_dot_product
+    # attention_mask = causal_mask + attention_mask
+    # RuntimeError: The size of tensor a (2048) must match the size of tensor b (8132) at non-singleton dimension 3
--- a/src/axolotl/utils/wandb.py
+++ b/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"
--- a/tests/test_validation.py
+++ b/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,50 +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)
Author	SHA1	Message	Date
Wing Lian	6fcb73faaa	more gpt-neox long ctx fixes Some checks failed pre-commit / pre-commit (push) Has been cancelled Details PyTest / test (3.10) (push) Has been cancelled Details PyTest / test (3.9) (push) Has been cancelled Details	2023-06-01 08:20:08 -04:00
Wing Lian	a32cc1d021	fix bettertransformers save, force it to skip after saving correctly in callback	2023-06-01 00:33:13 -04:00
Wing Lian	86bd9fcff4	more tweaks to do pre-training with bettertransformers	2023-05-31 21:59:15 -04:00
Wing Lian	ed7531abb8	experimental expansion of ctx len	2023-05-31 16:51:19 -04:00
Wing Lian	bdb547b830	add validation/warning for bettertransformers and torch version	2023-05-31 16:41:24 -04:00
Wing Lian	8a37b43678	use pythia-12b, neox-20b is flaky	2023-05-31 16:41:21 -04:00
Wing Lian	28acebac36	add flash attn context for efficient training and attempt setting model to train mode:	2023-05-31 16:40:38 -04:00
Wing Lian	adea682316	add support for opimum bettertransformers	2023-05-31 16:39:35 -04:00