Add e2e test for ia3 ft

Co-authored-by: NanoCode012 <kevinvong@rocketmail.com>

.github/workflows/main.yml

@@ -23,7 +23,6 @@ jobs:
             python_version: "3.10"
             pytorch: 2.0.1
             axolotl_extras:
-            is_latest: true
           - cuda: 118
             cuda_version: 11.8.0
             python_version: "3.10"
@@ -55,9 +54,7 @@ jobs:
             PYTORCH_VERSION=${{ matrix.pytorch }}
           file: ./docker/Dockerfile
           push: ${{ github.event_name != 'pull_request' }}
-          tags: |
-            ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
-            ${{ (matrix.is_latest) && format('{0}-latest', steps.metadata.outputs.tags) || '' }}
+          tags: ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}
   build-axolotl-runpod:
     needs: build-axolotl

.github/workflows/tests.yml

@@ -71,9 +71,8 @@ jobs:
 
       - name: Install dependencies
         run: |
-          pip3 install --extra-index-url https://download.pytorch.org/whl/cu118 -U torch==2.0.1
           pip3 uninstall -y transformers accelerate
-          pip3 install -U -e .[flash-attn,mamba-ssm]
+          pip3 install -U -e .[flash-attn]
           pip3 install -r requirements-tests.txt
 
       - name: Run e2e tests

.mypy.ini

@@ -8,9 +8,6 @@ ignore_missing_imports = True
 [mypy-axolotl.monkeypatch.*]
 ignore_errors = True
 
-[mypy-axolotl.models.mixtral.*]
-ignore_errors = True
-
 [mypy-axolotl.models.phi.*]
 ignore_errors = True
 

.pylintrc

@@ -12,3 +12,4 @@ generated-members=numpy.*, torch.*
 disable=missing-function-docstring, line-too-long, import-error,
     too-many-arguments, too-many-locals, too-many-statements, too-many-branches, too-few-public-methods,
     too-many-instance-attributes, fixme, import-outside-toplevel, logging-fstring-interpolation,
+    too-many-boolean-expressions,

README.md

@@ -25,15 +25,14 @@ Features:
 - [Installation](#installation)
   - [Docker](#docker)
   - [Conda/Pip venv](#condapip-venv)
-  - [Runpod](#runpod)
   - [LambdaLabs](#lambdalabs)
   - [Windows](#windows)
-  - [Launching on public clouds via SkyPilot](#launching-on-public-clouds-via-skypilot)
 - [Dataset](#dataset)
   - [How to Add Custom Prompts](#how-to-add-custom-prompts)
   - [How to Use Custom Pretokenized Dataset](#how-to-use-your-custom-pretokenized-dataset)
 - [Config](#config)
   - [Train](#train)
+  - [Training w/ Deepspeed](#training-with-deepspeed)
   - [Inference](#inference)
   - [Merge LORA to Base](#merge-lora-to-base)
 - [Common Errors](#common-errors-)
@@ -65,21 +64,17 @@ Features:
 
 ## Axolotl supports
 
-|             | fp16/fp32 | lora | qlora | gptq | gptq w/flash attn | flash attn | xformers attn |
-|-------------|:----------|:-----|-------|------|-------------------|------------|--------------|
-| llama       | ✅         | ✅    | ✅     | ✅             | ✅                 | ✅          | ✅            |
-| Mistral     | ✅         | ✅    | ✅     | ✅             | ✅                 | ✅          | ✅            |
-| Mixtral-MoE | ✅         | ✅    | ✅     | ❓             | ❓                 | ❓          | ❓            |
-| Pythia      | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
-| cerebras    | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
-| btlm        | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
-| mpt         | ✅         | ❌    | ❓     | ❌             | ❌                 | ❌          | ❓            |
-| falcon      | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
-| gpt-j       | ✅         | ✅    | ✅     | ❌             | ❌                 | ❓          | ❓            |
-| XGen        | ✅         | ❓    | ✅     | ❓             | ❓                 | ❓          | ✅            |
-| phi         | ✅         | ✅    | ✅     | ❓             | ❓                 | ❓          | ❓            |
-| RWKV        | ✅         | ❓    | ❓     | ❓             | ❓                 | ❓          | ❓            |
-| Qwen        | ✅         | ✅    | ✅     | ❓             | ❓                 | ❓          | ❓            |
+|          | fp16/fp32 | lora | qlora | gptq | gptq w/flash attn | flash attn | xformers attn |
+|----------|:----------|:-----|-------|------|-------------------|------------|--------------|
+| llama    | ✅         | ✅    | ✅     | ✅             | ✅                 | ✅          | ✅            |
+| Pythia   | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
+| cerebras | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
+| btlm     | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
+| mpt      | ✅         | ❌    | ❓     | ❌             | ❌                 | ❌          | ❓            |
+| falcon   | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
+| gpt-j    | ✅         | ✅    | ✅     | ❌             | ❌                 | ❓          | ❓            |
+| XGen     | ✅         | ❓    | ✅     | ❓             | ❓                 | ❓          | ✅            |
+| phi      | ✅         | ✅    | ✅     | ❓             | ❓                 | ❓          | ❓            |
 
 
 ## Quickstart ⚡
@@ -88,29 +83,20 @@ Get started with Axolotl in just a few steps! This quickstart guide will walk yo
 
 **Requirements**: Python >=3.9 and Pytorch >=2.0.
 
-`pip3 install "axolotl[flash-attn,deepspeed] @ git+https://github.com/OpenAccess-AI-Collective/axolotl"`
-
-### For developers
 ```bash
 git clone https://github.com/OpenAccess-AI-Collective/axolotl
 cd axolotl
 
 pip3 install packaging
 pip3 install -e '.[flash-attn,deepspeed]'
-```
+pip3 install -U git+https://github.com/huggingface/peft.git
 
-### Usage
-```bash
 # finetune lora
 accelerate launch -m axolotl.cli.train examples/openllama-3b/lora.yml
 
 # inference
 accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
-    --lora_model_dir="./lora-out"
-
-# gradio
-accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
-    --lora_model_dir="./lora-out" --gradio
+    --peft_model_dir="./lora-out"
 ```
 
 ## Installation
@@ -121,6 +107,7 @@ accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
   ```bash
   docker run --gpus '"all"' --rm -it winglian/axolotl:main-py3.10-cu118-2.0.1
   ```
+  - `winglian/axolotl-runpod:main-latest`: for runpod or use this [direct link](https://runpod.io/gsc?template=v2ickqhz9s&ref=6i7fkpdz)
 
   Or run on the current files for development:
 
@@ -128,27 +115,6 @@ accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
   docker compose up -d
   ```
 
-  <details>
-
-  <summary>Docker advanced</summary>
-
-  A more powerful Docker command to run would be this:
-
-  ```bash
-  docker run --privileged --gpus '"all"' --shm-size 10g --rm -it --name axolotl --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --mount type=volume,src=axolotl,target=/workspace/axolotl -v ${HOME}/.cache/huggingface:/root/.cache/huggingface winglian/axolotl:main-py3.10-cu118-2.0.1
-  ```
-
-  It additionally:
-  * Prevents memory issues when running e.g. deepspeed (e.g. you could hit SIGBUS/signal 7 error) through `--ipc` and `--ulimit` args.
-  * Persists the downloaded HF data (models etc.) and your modifications to axolotl code through `--mount`/`-v` args.
-  * The `--name` argument simply makes it easier to refer to the container in vscode (`Dev Containers: Attach to Running Container...`) or in your terminal.
-  * The `--privileged` flag gives all capabilities to the container.
-  * The `--shm-size 10g` argument increases the shared memory size. Use this if you see `exitcode: -7` errors using deepspeed.
-
-  [More information on nvidia website](https://docs.nvidia.com/deeplearning/frameworks/user-guide/index.html#setincshmem)
-
-  </details>
-
 #### Conda/Pip venv
   1. Install python >=**3.9**
 
@@ -165,10 +131,6 @@ accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
         ```
         Get the token at huggingface.co/settings/tokens
 
-#### Runpod
-
-Use `winglian/axolotl-runpod:main-latest` or use this [direct link](https://runpod.io/gsc?template=v2ickqhz9s&ref=6i7fkpdz)
-
 #### LambdaLabs
   <details>
 
@@ -216,28 +178,6 @@ Use `winglian/axolotl-runpod:main-latest` or use this [direct link](https://runp
 #### Windows
 Please use WSL or Docker!
 
-
-#### Launching on public clouds via SkyPilot
-To launch on GPU instances (both on-demand and spot instances) on 7+ clouds (GCP, AWS, Azure, OCI, and more), you can use [SkyPilot](https://skypilot.readthedocs.io/en/latest/index.html):
-```bash
-pip install "skypilot-nightly[gcp,aws,azure,oci,lambda,kubernetes,ibm,scp]"  # choose your clouds
-sky check
-```
-Get the [example YAMLs](https://github.com/skypilot-org/skypilot/tree/master/llm/axolotl) of using Axolotl to finetune `mistralai/Mistral-7B-v0.1`:
-```
-git clone https://github.com/skypilot-org/skypilot.git
-cd skypilot/llm/axolotl
-```
-Use one command to launch:
-```bash
-# On-demand
-HF_TOKEN=xx sky launch axolotl.yaml --env HF_TOKEN
-
-# Managed spot (auto-recovery on preemption)
-HF_TOKEN=xx BUCKET=<unique-name> sky spot launch axolotl-spot.yaml --env HF_TOKEN --env BUCKET
-```
-
-
 ### Dataset
 
 Axolotl supports a variety of dataset formats. Below are some of the formats you can use.
@@ -247,7 +187,7 @@ Have dataset(s) in one of the following format (JSONL recommended):
   ```json
   {"instruction": "...", "input": "...", "output": "..."}
   ```
-- `sharegpt`: conversations where `from` is `human`/`gpt`. (optional: `system` to override default system prompt)
+- `sharegpt`: conversations where `from` is `human`/`gpt`
   ```json
   {"conversations": [{"from": "...", "value": "..."}]}
   ```
@@ -416,13 +356,6 @@ See [examples](examples) for quick start. It is recommended to duplicate and mod
         - typescript
       type: ... # unimplemented custom format
 
-  # fastchat conversation
-  # See 'conversation' options: https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
-  datasets:
-    - path: ...
-      type: sharegpt
-      conversation: chatml
-
   # local
   datasets:
     - path: data.jsonl # or json
@@ -434,12 +367,6 @@ See [examples](examples) for quick start. It is recommended to duplicate and mod
     - path: knowrohit07/know_sql
       type: context_qa.load_v2
       train_on_split: validation
-
-  # loading from s3 or gcs
-  # s3 creds will be loaded from the system default and gcs only supports public access
-  dataset:
-    - path: s3://path_to_ds # Accepts folder with arrow/parquet or file path like above. Supports s3, gcs.
-      ...
   ```
 
 - loading
@@ -457,17 +384,17 @@ See [examples](examples) for quick start. It is recommended to duplicate and mod
 - lora
   ```yaml
   adapter: lora # qlora or leave blank for full finetune
-  lora_r: 8
-  lora_alpha: 16
-  lora_dropout: 0.05
-  lora_target_modules:
+  peft_r: 8
+  peft_alpha: 16
+  peft_dropout: 0.05
+  peft_target_modules:
     - q_proj
     - v_proj
   ```
 
 <details>
 
-<summary>All yaml options (click me)</summary>
+<summary>All yaml options</summary>
 
 ```yaml
 # This is the huggingface model that contains *.pt, *.safetensors, or *.bin files
@@ -502,15 +429,6 @@ is_falcon_derived_model:
 is_llama_derived_model:
 # Please note that if you set this to true, `padding_side` will be set to "left" by default
 is_mistral_derived_model:
-is_qwen_derived_model:
-
-# optional overrides to the base model configuration
-model_config:
-  # RoPE Scaling https://github.com/huggingface/transformers/pull/24653
-  rope_scaling:
-    type: # linear | dynamic
-    factor: # float
-
 
 # Whether you are training a 4-bit GPTQ quantized model
 gptq: true
@@ -535,7 +453,7 @@ float16: true
 
 # A list of one or more datasets to finetune the model with
 datasets:
-  # HuggingFace dataset repo | s3://,gs:// path | "json" for local dataset, make sure to fill data_files
+  # HuggingFace dataset repo | "json" for local dataset, make sure to fill data_files
   - path: vicgalle/alpaca-gpt4
   # The type of prompt to use for training. [alpaca, sharegpt, gpteacher, oasst, reflection]
     type: alpaca # format | format:<prompt_style> (chat/instruct) | <prompt_strategies>.load_<load_fn>
@@ -543,12 +461,7 @@ datasets:
     data_files: # Optional[str] path to source data files
     shards: # Optional[int] number of shards to split data into
     name: # Optional[str] name of dataset configuration to load
-    train_on_split: train # Optional[str] name of dataset split to load from
-
-    # Optional[str] fastchat conversation type, only used with type: sharegpt
-    conversation:  # Options (see Conversation 'name'): https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
-    field_human: # Optional[str]. Human key to use for conversation.
-    field_model: # Optional[str]. Assistant key to use for conversation.
+    conversation:  # Optional[str] fastchat conversation type, only used with type: sharegpt
 
   # Custom user prompt
   - path: repo
@@ -614,25 +527,19 @@ eval_sample_packing:
 sample_packing_eff_est:
 total_num_tokens:
 
-# Passed through to transformers when loading the model when launched without accelerate
-# Use `sequential` when training w/ model parallelism to limit memory
-device_map:
-# Defines the max memory usage per gpu on the system. Passed through to transformers when loading the model.
-max_memory:
-
 # If you want to use 'lora' or 'qlora' or 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.
 # This means after training, if you want to test the model, you should set this to the value of `lora_out_dir`.
-lora_model_dir:
+peft_model_dir:
 
 # LoRA hyperparameters
 # For more details about the following options, see:
 # https://www.anyscale.com/blog/fine-tuning-llms-lora-or-full-parameter-an-in-depth-analysis-with-llama-2
-lora_r: 8
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_modules:
+peft_r: 8
+peft_alpha: 16
+peft_dropout: 0.05
+peft_target_modules:
   - q_proj
   - v_proj
 #  - k_proj
@@ -640,13 +547,13 @@ lora_target_modules:
 #  - gate_proj
 #  - down_proj
 #  - up_proj
-lora_target_linear: # If true, will target all linear layers
+peft_target_linear: # if true, will target all linear layers
 
 # If you added new tokens to the tokenizer, you may need to save some LoRA modules because they need to know the new tokens.
 # For LLaMA and Mistral, you need to save `embed_tokens` and `lm_head`. It may vary for other models.
 # `embed_tokens` converts tokens to embeddings, and `lm_head` converts embeddings to token probabilities.
 # https://github.com/huggingface/peft/issues/334#issuecomment-1561727994
-lora_modules_to_save:
+peft_modules_to_save:
 #  - embed_tokens
 #  - lm_head
 
@@ -654,7 +561,8 @@ lora_modules_to_save:
 # If you merge the adapter to the base model, a subdirectory `merged` will be created under this directory.
 # Make sure `lora_model_dir` points to this directory if you want to use the trained model.
 lora_out_dir:
-lora_fan_in_fan_out: false
+peft_fan_in_fan_out: false
+peft_feedforward_modules:  # ffn modules for IA3, for llama down projection
 
 # ReLoRA configuration
 # Must use either 'lora' or 'qlora' adapter, and does not support fsdp or deepspeed
@@ -667,8 +575,7 @@ wandb_mode: # "offline" to save run metadata locally and not sync to the server,
 wandb_project: # Your wandb project name
 wandb_entity: # A wandb Team name if using a Team
 wandb_watch:
-wandb_name: # Set the name of your wandb run
-wandb_run_id: # Set the ID of your wandb run
+wandb_run_id: # Set the name of your wandb run
 wandb_log_model: # "checkpoint" to log model to wandb Artifacts every `save_steps` or "end" to log only at the end of training
 
 # Where to save the full-finetuned model to
@@ -685,15 +592,14 @@ gradient_accumulation_steps: 1
 # The number of samples to include in each batch. This is the number of samples sent to each GPU.
 micro_batch_size: 2
 eval_batch_size:
-num_epochs: 4
-warmup_steps: 100  # cannot use with warmup_ratio
-warmup_ratio: 0.05  # cannot use with warmup_steps
+num_epochs: 3
+warmup_steps: 100
 learning_rate: 0.00003
 lr_quadratic_warmup:
 logging_steps:
 save_strategy: # Set to `no` to skip checkpoint saves
 save_steps: # Leave empty to save at each epoch
-eval_steps: # Leave empty to eval at each epoch, integers for every N steps. decimal for fraction of total steps
+eval_steps: # Leave empty to eval at each epoch
 save_total_limit: # Checkpoints saved at a time
 # Maximum number of iterations to train for. It precedes num_epochs which means that
 # if both are set, num_epochs will not be guaranteed.
@@ -703,9 +609,6 @@ max_steps:
 eval_table_size: # Approximate number of predictions sent to wandb depending on batch size. Enabled above 0. Default is 0
 eval_table_max_new_tokens: # Total number of tokens generated for predictions sent to wandb. Default is 128
 
-loss_watchdog_threshold: # High loss value, indicating the learning has broken down (a good estimate is ~2 times the loss at the start of training)
-loss_watchdog_patience: # Number of high-loss steps in a row before the trainer aborts (default: 3)
-
 # Save model as safetensors (require safetensors package)
 save_safetensors:
 
@@ -782,8 +685,6 @@ xformers_attention:
 flash_attention:
 flash_attn_cross_entropy:  # Whether to use flash-attention cross entropy implementation - advanced use only
 flash_attn_rms_norm:  # Whether to use flash-attention rms norm implementation - advanced use only
-flash_attn_fuse_qkv: # Whether to fuse QKV into a single operation
-flash_attn_fuse_mlp: # Whether to fuse part of the MLP into a single operation
 # Whether to use scaled-dot-product attention
 # https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
 sdp_attention:
@@ -792,6 +693,10 @@ 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:
+# RoPE Scaling https://github.com/huggingface/transformers/pull/24653
+rope_scaling:
+  type: # linear | dynamic
+  factor: # float
 
 # Resume from a specific checkpoint dir
 resume_from_checkpoint:
@@ -909,41 +814,14 @@ Run
 accelerate launch -m axolotl.cli.train your_config.yml
 ```
 
-#### Preprocess dataset
-
-You can optionally pre-tokenize dataset with the following before finetuning.
-This is recommended for large datasets.
-
-- Set `push_dataset_to_hub: hf_user/repo` to push it to Huggingface.
-- Use `--debug` to see preprocessed examples.
-
-```bash
-python -m axolotl.cli.preprocess your_config.yml
-```
-
 #### Multi-GPU
 
-Below are the options available in axolotl for training with multiple GPUs. Note that DeepSpeed
-is the recommended multi-GPU option currently because FSDP may experience
-[loss instability](https://github.com/huggingface/transformers/issues/26498).
-
-##### DeepSpeed
-
-Deepspeed is an optimization suite for multi-gpu systems allowing you to train much larger models than you
-might typically be able to fit into your GPU's VRAM. More information about the various optimization types
-for deepspeed is available at https://huggingface.co/docs/accelerate/main/en/usage_guides/deepspeed#what-is-integrated
-
-We provide several default deepspeed JSON configurations for ZeRO stage 1, 2, and 3.
-
-```yaml
-deepspeed: deepspeed/zero1.json
+You can optionally pre-tokenize dataset with the following before finetuning:
+```bash
+CUDA_VISIBLE_DEVICES="" accelerate launch -m axolotl.cli.train your_config.yml --prepare_ds_only
 ```
 
-```shell
-accelerate launch -m axolotl.cli.train examples/llama-2/config.py --deepspeed deepspeed/zero1.json
-```
-
-##### FSDP
+##### Config
 
 - llama FSDP
 ```yaml
@@ -964,17 +842,35 @@ wandb_mode:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 ```
 
+### Training with Deepspeed
+
+Deepspeed is an optimization suite for multi-gpu systems allowing you to train much larger models than you
+might typically be able to fit into your GPU's VRAM. More information about the various optimization types
+for deepspeed is available at https://huggingface.co/docs/accelerate/main/en/usage_guides/deepspeed#what-is-integrated
+
+We provide several default deepspeed JSON configurations for ZeRO stage 1, 2, and 3.
+
+```shell
+accelerate launch -m axolotl.cli.train examples/llama-2/config.py --deepspeed deepspeed/zero1.json
+```
+
+or
+
+```yaml
+deepspeed: deepspeed/zero1.json
+```
+
 ### Inference
 
 Pass the appropriate flag to the train command:
 
 - Pretrained LORA:
   ```bash
-  python -m axolotl.cli.inference examples/your_config.yml --lora_model_dir="./lora-output-dir"
+  python -m axolotl.cli.inference examples/your_config.yml --peft_model_dir="./lora-output-dir"
   ```
 - Full weights finetune:
   ```bash
@@ -985,10 +881,6 @@ Pass the appropriate flag to the train command:
   cat /tmp/prompt.txt | python -m axolotl.cli.inference examples/your_config.yml \
     --base_model="./completed-model" --prompter=None --load_in_8bit=True
   ```
--- With gradio hosting
-  ```bash
-  python -m axolotl.cli.inference examples/your_config.yml --gradio
-  ```
 
 Please use `--sample_packing False` if you have it on and receive the error similar to below:
 
@@ -999,7 +891,7 @@ Please use `--sample_packing False` if you have it on and receive the error simi
 Add below flag to train command above
 
 ```bash
-python3 -m axolotl.cli.merge_lora examples/your_config.yml --lora_model_dir="./completed-model" --load_in_8bit=False --load_in_4bit=False
+python3 -m axolotl.cli.merge_lora examples/your_config.yml --peft_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
@@ -1010,8 +902,6 @@ CUDA_VISIBLE_DEVICES="" python3 -m axolotl.cli.merge_lora ...
 
 ## Common Errors 🧰
 
-See also the [FAQ's](./docs/faq.md).
-
 > If you encounter a 'Cuda out of memory' error, it means your GPU ran out of memory during the training process. Here's how to resolve it:
 
 Please reduce any below

deepspeed/zero1.json

@@ -24,6 +24,16 @@
       "weight_decay": "auto"
     }
   },
+  "scheduler": {
+    "type": "WarmupDecayLR",
+    "params": {
+      "warmup_min_lr": "auto",
+      "warmup_max_lr": "auto",
+      "warmup_num_steps": "auto",
+      "warmup_type": "linear",
+      "total_num_steps": "auto"
+    }
+  },
   "gradient_accumulation_steps": "auto",
   "train_batch_size": "auto",
   "train_micro_batch_size_per_gpu": "auto",

deepspeed/zero2.json

@@ -28,6 +28,16 @@
       "weight_decay": "auto"
     }
   },
+  "scheduler": {
+    "type": "WarmupDecayLR",
+    "params": {
+      "warmup_min_lr": "auto",
+      "warmup_max_lr": "auto",
+      "warmup_num_steps": "auto",
+      "warmup_type": "linear",
+      "total_num_steps": "auto"
+    }
+  },
   "gradient_accumulation_steps": "auto",
   "train_batch_size": "auto",
   "train_micro_batch_size_per_gpu": "auto",

deepspeed/zero3.json

@@ -1,6 +1,14 @@
 {
   "zero_optimization": {
     "stage": 3,
+    "offload_optimizer": {
+      "device": "cpu",
+      "pin_memory": true
+    },
+    "offload_param": {
+      "device": "cpu",
+      "pin_memory": true
+    },
     "overlap_comm": true,
     "contiguous_gradients": true,
     "sub_group_size": 0,
@@ -32,6 +40,15 @@
       "weight_decay": "auto"
     }
   },
+  "scheduler": {
+    "type": "WarmupLR",
+    "params": {
+      "warmup_min_lr": "auto",
+      "warmup_max_lr": "auto",
+      "warmup_num_steps": "auto",
+      "warmup_type": "linear"
+    }
+  },
   "gradient_accumulation_steps": "auto",
   "train_batch_size": "auto",
   "train_micro_batch_size_per_gpu": "auto",

docker/Dockerfile

@@ -21,9 +21,9 @@ WORKDIR /workspace/axolotl
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN sed -i "s/torch==.*/torch==$PYTORCH_VERSION/" requirements.txt
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[deepspeed,flash-attn,$AXOLOTL_EXTRAS]; \
+        pip install -e .[flash-attn,$AXOLOTL_EXTRAS]; \
     else \
-        pip install -e .[deepspeed,flash-attn]; \
+        pip install -e .[flash-attn]; \
     fi
 
 # fix so that git fetch/pull from remote works

docker/Dockerfile-base

@@ -10,10 +10,8 @@ ENV PATH="/root/miniconda3/bin:${PATH}"
 ARG PYTHON_VERSION="3.9"
 ARG PYTORCH_VERSION="2.0.1"
 ARG CUDA="118"
-ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"
 
 ENV PYTHON_VERSION=$PYTHON_VERSION
-ENV TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST
 
 RUN apt-get update \
     && apt-get install -y wget git build-essential ninja-build git-lfs libaio-dev && rm -rf /var/lib/apt/lists/* \
@@ -29,9 +27,47 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
 WORKDIR /workspace
 
 RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
-    python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} deepspeed-kernels --extra-index-url https://download.pytorch.org/whl/cu$CUDA
+    python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} --extra-index-url https://download.pytorch.org/whl/cu$CUDA
 
-RUN git lfs install --skip-repo && \
-    pip3 install awscli && \
+FROM base-builder AS deepspeed-builder
+
+ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"
+
+WORKDIR /workspace
+
+RUN git clone https://github.com/microsoft/DeepSpeed.git && \
+    cd DeepSpeed && \
+    MAX_CONCURRENCY=8 DS_BUILD_SPARSE_ATTN=0 DS_BUILD_OPS=1 DS_BUILD_EVOFORMER_ATTN=0 python3 setup.py bdist_wheel
+
+FROM base-builder AS bnb-builder
+
+WORKDIR /workspace
+ARG CUDA="118"
+ENV CUDA=$CUDA
+ARG MAX_JOBS="-1"
+ENV MAX_JOBS=$MAX_JOBS
+
+RUN git clone https://github.com/TimDettmers/bitsandbytes.git && \
+    cd bitsandbytes && \
+    CUDA_VERSION=$CUDA make cuda11x && \
+    python setup.py bdist_wheel
+
+FROM base-builder
+
+ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"
+ENV TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST
+
+RUN mkdir -p /workspace/builds
+COPY --from=bnb-builder /workspace/bitsandbytes /workspace/builds/bitsandbytes
+
+RUN mkdir -p /workspace/wheels/bitsandbytes
+COPY --from=deepspeed-builder /workspace/DeepSpeed/dist/deepspeed-*.whl wheels
+COPY --from=bnb-builder /workspace/bitsandbytes/dist/bitsandbytes-*.whl wheels
+COPY --from=bnb-builder /workspace/bitsandbytes/bitsandbytes/libbitsandbytes*.so wheels/bitsandbytes
+
+RUN pip3 install wheels/deepspeed-*.whl
+RUN cd /workspace/builds/bitsandbytes && python3 setup.py install
+RUN git lfs install --skip-repo
+RUN pip3 install awscli && \
     # The base image ships with `pydantic==1.8.2` which is not working
     pip3 install -U --no-cache-dir pydantic==1.10.10

docker/Dockerfile-runpod

@@ -4,7 +4,6 @@ FROM winglian/axolotl:$BASE_TAG
 ENV HF_DATASETS_CACHE="/workspace/data/huggingface-cache/datasets"
 ENV HUGGINGFACE_HUB_CACHE="/workspace/data/huggingface-cache/hub"
 ENV TRANSFORMERS_CACHE="/workspace/data/huggingface-cache/hub"
-ENV HF_HOME="/workspace/data/huggingface-cache/hub"
 
 COPY scripts/runpod-entrypoint.sh /root/runpod-entrypoint.sh
 

docs/faq.md

@@ -1,18 +0,0 @@
-# Axolotl FAQ's
-
-
-> The trainer stopped and hasn't progressed in several minutes.
-
-Usually an issue with the GPU's communicating with each other. See the [NCCL doc](../docs/nccl.md)
-
-> Exitcode -9
-
-This usually happens when you run out of system RAM.
-
-> Exitcode -7 while using deepspeed
-
-Try upgrading deepspeed w: `pip install -U deepspeed`
-
-> AttributeError: 'DummyOptim' object has no attribute 'step'
-
-You may be using deepspeed with single gpu. Please don't set `deepspeed:` in yaml or cli.

examples/cerebras/btlm-ft.yml

@@ -1,4 +1,5 @@
 base_model: cerebras/btlm-3b-8k-base
+base_model_config: cerebras/btlm-3b-8k-base
 model_type: AutoModelForCausalLM
 tokenizer_type: GPT2Tokenizer
 trust_remote_code: true
@@ -14,10 +15,10 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path: last_prepared_run
-val_set_size: 0.05
+val_set_size: 0.01
 
 adapter:
-lora_model_dir:
+peft_model_dir:
 sequence_len: 2048
 max_packed_sequence_len:
 sample_packing: false
@@ -35,7 +36,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 output_dir: btlm-out

examples/cerebras/qlora.yml

@@ -1,4 +1,5 @@
 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
@@ -7,9 +8,9 @@ datasets:
   - path: teknium/GPT4-LLM-Cleaned
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 sequence_len: 2048
 max_packed_sequence_len: 2048
 lora_r: 16
@@ -24,7 +25,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./qlora-out
 batch_size: 4
@@ -49,7 +50,7 @@ flash_attention:
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 save_steps:
 debug:
 deepspeed:

examples/code-llama/13b/lora.yml

@@ -1,4 +1,5 @@
 base_model: codellama/CodeLlama-13b-hf
+base_model_config: codellama/CodeLlama-13b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
 is_llama_derived_model: true
@@ -11,7 +12,7 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./lora-out
 
 sequence_len: 4096
@@ -19,7 +20,7 @@ sample_packing: true
 pad_to_sequence_len: true
 
 adapter: lora
-lora_model_dir:
+peft_model_dir:
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
@@ -29,12 +30,12 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -54,7 +55,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 save_steps:
 debug:
 deepspeed:

examples/code-llama/13b/qlora.yml

@@ -1,4 +1,5 @@
 base_model: codellama/CodeLlama-13b-hf
+base_model_config: codellama/CodeLlama-13b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
 is_llama_derived_model: true
@@ -11,11 +12,11 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./qlora-out
 
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 
 sequence_len: 4096
 sample_packing: true
@@ -31,12 +32,12 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
+num_epochs: 3
 optimizer: paged_adamw_32bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -56,7 +57,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 save_steps:
 debug:
 deepspeed:

examples/code-llama/34b/lora.yml

@@ -1,4 +1,5 @@
 base_model: codellama/CodeLlama-34b-hf
+base_model_config: codellama/CodeLlama-34b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
 is_llama_derived_model: true
@@ -11,7 +12,7 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./lora-out
 
 sequence_len: 4096
@@ -19,7 +20,7 @@ sample_packing: true
 pad_to_sequence_len: true
 
 adapter: lora
-lora_model_dir:
+peft_model_dir:
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
@@ -29,12 +30,12 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -54,7 +55,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 save_steps:
 debug:
 deepspeed:

examples/code-llama/34b/qlora.yml

@@ -1,4 +1,5 @@
 base_model: codellama/CodeLlama-34b-hf
+base_model_config: codellama/CodeLlama-34b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
 is_llama_derived_model: true
@@ -11,11 +12,11 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./qlora-out
 
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 
 sequence_len: 4096
 sample_packing: true
@@ -31,12 +32,12 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
+num_epochs: 3
 optimizer: paged_adamw_32bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -56,7 +57,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 save_steps:
 debug:
 deepspeed:

examples/code-llama/7b/lora.yml

@@ -1,4 +1,5 @@
 base_model: codellama/CodeLlama-7b-hf
+base_model_config: codellama/CodeLlama-7b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
 is_llama_derived_model: true
@@ -11,7 +12,7 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./lora-out
 
 sequence_len: 4096
@@ -19,7 +20,7 @@ sample_packing: true
 pad_to_sequence_len: true
 
 adapter: lora
-lora_model_dir:
+peft_model_dir:
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
@@ -29,12 +30,12 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -54,7 +55,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 save_steps:
 debug:
 deepspeed:

examples/code-llama/7b/qlora.yml

@@ -1,4 +1,5 @@
 base_model: codellama/CodeLlama-7b-hf
+base_model_config: codellama/CodeLlama-7b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
 is_llama_derived_model: true
@@ -11,11 +12,11 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./qlora-out
 
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 
 sequence_len: 4096
 sample_packing: true
@@ -31,12 +32,12 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
+num_epochs: 3
 optimizer: paged_adamw_32bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -56,7 +57,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 save_steps:
 debug:
 deepspeed:

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

@@ -1,4 +1,5 @@
 base_model: tiiuae/falcon-7b
+base_model_config: tiiuae/falcon-7b
 trust_remote_code: true
 model_type: AutoModelForCausalLM
 tokenizer_type: AutoTokenizer
@@ -12,9 +13,9 @@ datasets:
   - path: teknium/GPT4-LLM-Cleaned
     type: alpaca:chat
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 adapter: lora
-lora_model_dir:
+peft_model_dir:
 sequence_len: 2048
 max_packed_sequence_len:
 lora_r: 16
@@ -26,7 +27,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./falcon-7b
 batch_size: 2

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

@@ -1,6 +1,7 @@
 # 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
@@ -18,10 +19,10 @@ datasets:
       - Chain-of-Thought/formatted_cot_data/gsm8k_train.json
     type: "alpaca:chat"
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 # enable QLoRA
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 sequence_len: 2048
 max_packed_sequence_len:
 
@@ -40,7 +41,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./qlora-out
 
@@ -53,7 +54,7 @@ output_dir: ./qlora-out
 # decrease if OOM, increase for max VRAM utilization
 micro_batch_size: 1
 gradient_accumulation_steps: 2
-num_epochs: 4
+num_epochs: 3
 # Optimizer for QLoRA
 optimizer: paged_adamw_32bit
 torchdistx_path:

examples/falcon/config-7b.yml

@@ -1,4 +1,5 @@
 base_model: tiiuae/falcon-7b
+base_model_config: tiiuae/falcon-7b
 trust_remote_code: true
 model_type: AutoModelForCausalLM
 tokenizer_type: AutoTokenizer
@@ -12,9 +13,9 @@ datasets:
   - path: teknium/GPT4-LLM-Cleaned
     type: alpaca:chat
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 adapter:
-lora_model_dir:
+peft_model_dir:
 sequence_len: 2048
 max_packed_sequence_len:
 lora_r: 64
@@ -26,7 +27,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./falcon-7b
 batch_size: 2

examples/gptj/qlora.yml

@@ -1,4 +1,5 @@
 base_model: EleutherAI/gpt-j-6b
+base_model_config: EleutherAI/gpt-j-6b
 load_in_8bit: false
 load_in_4bit: true
 strict: false
@@ -7,9 +8,9 @@ datasets:
   - path: teknium/GPT4-LLM-Cleaned
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 sequence_len: 2048
 max_packed_sequence_len:
 lora_r: 8
@@ -21,7 +22,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./qlora-out
 gradient_accumulation_steps: 2
@@ -46,7 +47,7 @@ flash_attention:
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 save_steps:
 debug:
 deepspeed:

examples/jeopardy-bot/config.yml

@@ -1,4 +1,5 @@
 base_model: huggyllama/llama-7b
+base_model_config: huggyllama/llama-7b
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 load_in_8bit: false
@@ -8,7 +9,7 @@ datasets:
 dataset_prepared_path:
 val_set_size: 0.02
 adapter:
-lora_model_dir:
+peft_model_dir:
 sequence_len: 512
 max_packed_sequence_len:
 lora_r:
@@ -19,12 +20,12 @@ lora_fan_in_fan_out: false
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./jeopardy-bot-7b
 gradient_accumulation_steps: 1
 micro_batch_size: 1
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine

examples/llama-2/README.md

@@ -9,16 +9,12 @@ gradient_accumulation_steps: 2
 micro_batch_size: 1
 
 ```shell
-accelerate launch -m axolotl.cli.train examples/llama-2/qlora.yml
+accelerate launch scripts/finetune.py examples/llama-2/qlora.yml
+
 ```
 or
 
 ```shell
-accelerate launch -m axolotl.cli.train examples/llama-2/lora.yml
-```
+accelerate launch scripts/finetune.py examples/llama-2/lora.yml
 
-To launch a full finetuning with 16-bit precision:
-
-```shell
-accelerate launch -m axolotl.cli.train examples/llama-2/fft_optimized.yml
 ```

examples/llama-2/gptq-lora.yml

@@ -1,4 +1,5 @@
 base_model: TheBloke/Llama-2-7B-GPTQ
+base_model_config: TheBloke/Llama-2-7B-GPTQ
 is_llama_derived_model: false
 gptq: true
 gptq_disable_exllama: true
@@ -15,9 +16,9 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 adapter: lora
-lora_model_dir:
+peft_model_dir:
 sequence_len: 4096
 sample_packing:
 lora_r: 8
@@ -32,12 +33,12 @@ lora_target_linear:
 lora_fan_in_fan_out:
 wandb_project:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./model-out
 gradient_accumulation_steps: 1
 micro_batch_size: 1
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_torch
 adam_beta2: 0.95
 adam_eps: 0.00001

examples/llama-2/ia3.yml

@@ -1,9 +1,10 @@
-base_model: NousResearch/Llama-2-7b-hf
+base_model: meta-llama/Llama-2-7b-hf
+base_model_config: meta-llama/Llama-2-7b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 is_llama_derived_model: true
 
-load_in_8bit: false
+load_in_8bit: true
 load_in_4bit: false
 strict: false
 
@@ -11,30 +12,32 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path: last_run_prepared
-val_set_size: 0.05
-output_dir: ./out
+val_set_size: 0.01
+output_dir: ./ia3-out
 
 sequence_len: 4096
 sample_packing: true
 pad_to_sequence_len: true
 
-adapter:
-lora_model_dir:
-lora_r:
-lora_alpha:
-lora_dropout:
-lora_target_linear:
-lora_fan_in_fan_out:
+adapter: ia3
+peft_model_dir:
+peft_target_modules:
+  - k_proj
+  - v_proj
+  - down_proj
+peft_feedforward_modules:
+  - down_proj
+peft_fan_in_fan_out: false
 
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 1
+micro_batch_size: 2
+num_epochs: 5
 optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -52,17 +55,14 @@ local_rank:
 logging_steps: 1
 xformers_attention:
 flash_attention: true
-flash_attn_cross_entropy: false
-flash_attn_rms_norm: true
-flash_attn_fuse_qkv: false
-flash_attn_fuse_mlp: true
 
-warmup_steps: 100
+warmup_steps: 10
 eval_steps: 0.05
 eval_table_size:
+eval_table_max_new_tokens:
 save_steps:
 debug:
-deepspeed: #deepspeed/zero2.json # multi-gpu only
+deepspeed:
 weight_decay: 0.1
 fsdp:
 fsdp_config:

examples/llama-2/lora.yml

@@ -1,4 +1,5 @@
 base_model: NousResearch/Llama-2-7b-hf
+base_model_config: NousResearch/Llama-2-7b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 is_llama_derived_model: true
@@ -11,7 +12,7 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./lora-out
 
 sequence_len: 4096
@@ -19,7 +20,7 @@ sample_packing: true
 pad_to_sequence_len: true
 
 adapter: lora
-lora_model_dir:
+peft_model_dir:
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
@@ -29,12 +30,12 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -54,7 +55,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 eval_table_size:
 eval_table_max_new_tokens: 128
 save_steps:

examples/llama-2/qlora.yml

@@ -1,4 +1,5 @@
 base_model: NousResearch/Llama-2-7b-hf
+base_model_config: NousResearch/Llama-2-7b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 is_llama_derived_model: true
@@ -11,11 +12,11 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./qlora-out
 
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 
 sequence_len: 4096
 sample_packing: true
@@ -31,12 +32,12 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
+num_epochs: 3
 optimizer: paged_adamw_32bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -56,7 +57,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 eval_table_size:
 save_steps:
 debug:

examples/llama-2/relora.yml

@@ -1,4 +1,5 @@
 base_model: NousResearch/Llama-2-7b-hf
+base_model_config: NousResearch/Llama-2-7b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 is_llama_derived_model: true
@@ -11,11 +12,11 @@ datasets:
   - path: teknium/GPT4-LLM-Cleaned
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./relora-out
 
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 
 sequence_len: 4096
 sample_packing: true
@@ -35,12 +36,12 @@ relora_cpu_offload: false
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 4
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -60,7 +61,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 save_steps: 50
 debug:
 deepspeed:

examples/llama-2/tiny-llama.yml

@@ -1,4 +1,5 @@
-base_model: PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T
+base_model: PY007/TinyLlama-1.1B-step-50K-105b
+base_model_config: PY007/TinyLlama-1.1B-step-50K-105b
 
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
@@ -12,14 +13,14 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./lora-out
 
 sequence_len: 4096
 sample_packing: true
 
 adapter: lora
-lora_model_dir:
+peft_model_dir:
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
@@ -29,12 +30,12 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002
@@ -54,7 +55,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps: 20
 eval_table_size:
 save_steps:
 debug:

examples/mamba/config.yml

@@ -1,61 +0,0 @@
-base_model: state-spaces/mamba-2.8b
-model_type: MambaLMHeadModel
-tokenizer_type: AutoTokenizer
-tokenizer_config: EleutherAI/gpt-neox-20b
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
-
-datasets:
-  - path: mhenrichsen/alpaca_2k_test
-    type: alpaca
-dataset_prepared_path:
-val_set_size: 0.0
-output_dir: ./out
-
-sequence_len: 2048
-sample_packing: false
-pad_to_sequence_len: false
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 2
-optimizer: paged_adamw_8bit
-lr_scheduler: cosine
-learning_rate: 5e-5
-
-train_on_inputs: false
-group_by_length: true
-
-bf16: true
-fp16: false
-tf32: true
-
-gradient_checkpointing: false
-early_stopping_patience:
-resume_from_checkpoint:
-local_rank:
-logging_steps: 1
-xformers_attention:
-flash_attention:
-
-warmup_steps: 10
-eval_steps:
-eval_table_size:
-eval_table_max_new_tokens: 128
-save_steps:  0.25
-debug:
-deepspeed:
-weight_decay: 0.0
-fsdp:
-fsdp_config:
-special_tokens:
-tokens:
-save_safetensors: False

examples/mistral/config.yml

@@ -1,4 +1,5 @@
 base_model: mistralai/Mistral-7B-v0.1
+base_model_config: mistralai/Mistral-7B-v0.1
 model_type: MistralForCausalLM
 tokenizer_type: LlamaTokenizer
 is_mistral_derived_model: true
@@ -11,7 +12,7 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./out
 
 sequence_len: 8192
@@ -21,12 +22,12 @@ pad_to_sequence_len: true
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.000005
@@ -46,8 +47,8 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
-eval_table_size:
+eval_steps: 20
+eval_table_size: 5
 eval_table_max_new_tokens: 128
 save_steps:
 debug:

examples/mistral/mixtral.yml

@@ -1,79 +0,0 @@
-base_model: DiscoResearch/mixtral-7b-8expert
-model_type: MixtralForCausalLM
-tokenizer_type: LlamaTokenizer
-trust_remote_code: true
-
-load_in_8bit: false
-load_in_4bit: true
-strict: false
-
-datasets:
-  - path: tatsu-lab/alpaca
-    type: alpaca
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.0
-output_dir: ./qlora-out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 4096
-sample_packing: true
-pad_to_sequence_len: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-lora_fan_in_fan_out:
-#lora_target_modules:
-#  - gate
-#  - q_proj
-#  - k_proj
-#  - v_proj
-#  - o_proj
-#  - w1
-#  - w2
-#  - w3
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 2
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-train_on_inputs: false
-group_by_length: false
-bf16: true
-fp16: false
-tf32: false
-
-gradient_checkpointing: true
-early_stopping_patience:
-resume_from_checkpoint:
-local_rank:
-logging_steps: 1
-xformers_attention:
-flash_attention: true
-
-loss_watchdog_threshold: 5.0
-loss_watchdog_patience: 3
-
-warmup_steps: 10
-eval_steps:
-eval_table_size:
-eval_table_max_new_tokens: 128
-save_steps:
-debug:
-deepspeed: deepspeed/zero2.json
-weight_decay: 0.0
-fsdp:
-fsdp_config:
-special_tokens:

examples/mistral/qlora.yml

@@ -1,4 +1,5 @@
 base_model: mistralai/Mistral-7B-v0.1
+base_model_config: mistralai/Mistral-7B-v0.1
 model_type: MistralForCausalLM
 tokenizer_type: LlamaTokenizer
 is_mistral_derived_model: true
@@ -11,7 +12,7 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path: last_run_prepared
-val_set_size: 0.05
+val_set_size: 0.01
 output_dir: ./qlora-out
 
 adapter: qlora
@@ -38,7 +39,7 @@ lora_target_modules:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 4
@@ -62,12 +63,9 @@ logging_steps: 1
 xformers_attention:
 flash_attention: true
 
-loss_watchdog_threshold: 5.0
-loss_watchdog_patience: 3
-
 warmup_steps: 10
-eval_steps: 0.05
-eval_table_size:
+eval_steps: 20
+eval_table_size: 5
 eval_table_max_new_tokens: 128
 save_steps:
 debug:

examples/mpt-7b/config.yml

@@ -1,4 +1,5 @@
 base_model: mosaicml/mpt-7b
+base_model_config: mosaicml/mpt-7b
 tokenizer_type: AutoTokenizer
 trust_remote_code: true  # required for mpt as their model class is not merged into transformers yet
 load_in_8bit: false
@@ -8,7 +9,7 @@ datasets:
 dataset_prepared_path:
 val_set_size: 0.02
 adapter:
-lora_model_dir:
+peft_model_dir:
 sequence_len: 2048
 max_packed_sequence_len:
 lora_r: 8
@@ -21,12 +22,12 @@ lora_fan_in_fan_out: false
 wandb_project: mpt-alpaca-7b
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./mpt-alpaca-7b
 gradient_accumulation_steps: 1
 micro_batch_size: 1
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine

examples/openllama-3b/config.yml

@@ -1,4 +1,5 @@
 base_model: openlm-research/open_llama_3b_v2
+base_model_config: openlm-research/open_llama_3b_v2
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 load_in_8bit: false
@@ -11,7 +12,7 @@ datasets:
 dataset_prepared_path:
 val_set_size: 0.02
 adapter:
-lora_model_dir:
+peft_model_dir:
 sequence_len: 1024
 sample_packing: true
 lora_r:
@@ -23,7 +24,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./openllama-out
 gradient_accumulation_steps: 1

examples/openllama-3b/lora.yml

@@ -1,4 +1,5 @@
 base_model: openlm-research/open_llama_3b_v2
+base_model_config: openlm-research/open_llama_3b_v2
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 load_in_8bit: true
@@ -11,7 +12,7 @@ datasets:
 dataset_prepared_path:
 val_set_size: 0.02
 adapter: lora
-lora_model_dir:
+peft_model_dir:
 sequence_len: 1024
 sample_packing: true
 lora_r: 8
@@ -29,7 +30,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./lora-out
 gradient_accumulation_steps: 1

examples/openllama-3b/qlora.yml

@@ -1,4 +1,5 @@
 base_model: openlm-research/open_llama_3b_v2
+base_model_config: openlm-research/open_llama_3b_v2
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 load_in_8bit: false
@@ -9,9 +10,9 @@ datasets:
   - path: teknium/GPT4-LLM-Cleaned
     type: alpaca
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 sequence_len: 1024
 sample_packing: true
 lora_r: 8
@@ -23,7 +24,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./qlora-out
 gradient_accumulation_steps: 1

examples/phi/phi-ft.yml

@@ -1,5 +1,6 @@
 base_model: microsoft/phi-1_5
-model_type: PhiForCausalLM
+base_model_config: microsoft/phi-1_5
+model_type: MixFormerSequentialForCausalLM
 tokenizer_type: AutoTokenizer
 is_llama_derived_model: false
 trust_remote_code: true
@@ -21,7 +22,7 @@ sample_packing: true
 pad_to_sequence_len:
 
 adapter:
-lora_model_dir:
+peft_model_dir:
 lora_r:
 lora_alpha:
 lora_dropout:
@@ -31,7 +32,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 1

examples/phi/phi-qlora.yml

@@ -1,4 +1,5 @@
 base_model: microsoft/phi-1_5
+base_model_config: microsoft/phi-1_5
 model_type: AutoModelForCausalLM
 tokenizer_type: AutoTokenizer
 is_llama_derived_model: false
@@ -21,7 +22,7 @@ sample_packing: false  # not CURRENTLY compatible with LoRAs
 pad_to_sequence_len:
 
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 lora_r: 64
 lora_alpha: 32
 lora_dropout: 0.05
@@ -31,7 +32,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 
 gradient_accumulation_steps: 1

examples/pythia-12b/config.yml

@@ -1,4 +1,5 @@
 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
@@ -12,7 +13,7 @@ datasets:
 dataset_prepared_path:
 val_set_size: 0.05
 adapter:
-lora_model_dir:
+peft_model_dir:
 sequence_len: 2048
 max_packed_sequence_len: 2048
 lora_r: 64
@@ -24,7 +25,7 @@ lora_fan_in_fan_out: true  # pythia/GPTNeoX lora specific
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./pythia-12b
 gradient_accumulation_steps: 1

examples/pythia/lora.yml

@@ -1,4 +1,5 @@
 base_model: EleutherAI/pythia-1.4b-deduped
+base_model_config: EleutherAI/pythia-1.4b-deduped
 load_in_8bit: true
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
@@ -6,7 +7,7 @@ datasets:
 dataset_prepared_path:
 val_set_size: 0.05
 adapter: lora
-lora_model_dir:
+peft_model_dir:
 sequence_len: 512
 lora_r: 16
 lora_alpha: 32
@@ -18,12 +19,12 @@ lora_fan_in_fan_out: true  # pythia/GPTNeoX lora specific
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./lora-alpaca-pythia
 gradient_accumulation_steps: 1
 micro_batch_size: 4
-num_epochs: 4
+num_epochs: 3
 learning_rate: 0.00001
 train_on_inputs: false
 group_by_length: false
@@ -33,5 +34,5 @@ early_stopping_patience:
 resume_from_checkpoint:
 local_rank:
 weight_decay: 0.1
-eval_steps: 0.05
+eval_steps: 20
 logging_steps: 1

examples/qwen/lora.yml

@@ -1,68 +0,0 @@
-base_model: Qwen/Qwen-7B
-model_type: AutoModelForCausalLM
-tokenizer_type: AutoTokenizer
-
-is_qwen_derived_model: true
-trust_remote_code: true
-
-load_in_8bit: true
-load_in_4bit: false
-strict: false
-
-datasets:
-  - path: mhenrichsen/alpaca_2k_test
-    type: alpaca
-dataset_prepared_path:
-val_set_size: 0.05
-output_dir: ./lora-out
-
-sequence_len: 2048  # supports up to 8192
-sample_packing: false
-pad_to_sequence_len:
-
-adapter: lora
-lora_model_dir:
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-lora_fan_in_fan_out:
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 4
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-train_on_inputs: false
-group_by_length: false
-bf16: true
-fp16: false
-tf32: false
-
-gradient_checkpointing: false
-early_stopping_patience:
-resume_from_checkpoint:
-local_rank:
-logging_steps: 1
-xformers_attention:
-flash_attention:
-
-warmup_steps: 10
-eval_steps: 0.05
-eval_table_size:
-eval_table_max_new_tokens: 128
-save_steps:
-debug:
-deepspeed:
-weight_decay: 0.0
-fsdp:
-fsdp_config:
-special_tokens:

examples/qwen/qlora.yml

@@ -1,68 +0,0 @@
-base_model: Qwen/Qwen-7B
-model_type: AutoModelForCausalLM
-tokenizer_type: AutoTokenizer
-
-is_qwen_derived_model: true
-trust_remote_code: true
-
-load_in_8bit: false
-load_in_4bit: true
-strict: false
-
-datasets:
-  - path: mhenrichsen/alpaca_2k_test
-    type: alpaca
-dataset_prepared_path:
-val_set_size: 0.05
-output_dir: ./lora-out
-
-sequence_len: 2048  # supports up to 8192
-sample_packing: false
-pad_to_sequence_len:
-
-adapter: qlora
-lora_model_dir:
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-lora_fan_in_fan_out:
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 4
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-train_on_inputs: false
-group_by_length: false
-bf16: true
-fp16: false
-tf32: false
-
-gradient_checkpointing: false
-early_stopping_patience:
-resume_from_checkpoint:
-local_rank:
-logging_steps: 1
-xformers_attention:
-flash_attention:
-
-warmup_steps: 10
-eval_steps: 0.05
-eval_table_size:
-eval_table_max_new_tokens: 128
-save_steps:
-debug:
-deepspeed:
-weight_decay: 0.0
-fsdp:
-fsdp_config:
-special_tokens:

examples/redpajama/config-3b.yml

@@ -1,4 +1,5 @@
 base_model: togethercomputer/RedPajama-INCITE-Chat-3B-v1
+base_model_config: togethercomputer/RedPajama-INCITE-Chat-3B-v1
 model_type: GPTNeoXForCausalLM
 tokenizer_type: AutoTokenizer
 trust_remote_code:
@@ -9,7 +10,7 @@ datasets:
 dataset_prepared_path:
 val_set_size: 0.02
 adapter:
-lora_model_dir:
+peft_model_dir:
 sequence_len: 2048
 max_packed_sequence_len:
 lora_r: 8
@@ -22,12 +23,12 @@ lora_fan_in_fan_out: false
 wandb_project: redpajama-alpaca-3b
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./redpajama-alpaca-3b
 batch_size: 4
 micro_batch_size: 1
-num_epochs: 4
+num_epochs: 3
 optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine

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

@@ -1,4 +1,5 @@
 base_model: replit/replit-code-v1-3b
+base_model_config: replit/replit-code-v1-3b
 trust_remote_code: true
 load_in_8bit: false
 datasets:
@@ -7,7 +8,7 @@ datasets:
 dataset_prepared_path:
 val_set_size: 0.05
 adapter: lora
-lora_model_dir:
+peft_model_dir:
 sequence_len: 2048
 max_packed_sequence_len:
 lora_r: 8
@@ -21,12 +22,12 @@ lora_fan_in_fan_out:
 wandb_project: lora-replit
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./lora-replit
 batch_size: 8
 micro_batch_size: 1
-num_epochs: 4
+num_epochs: 3
 optimizer:
 torchdistx_path:
 lr_scheduler:

examples/xgen-7b/xgen-7b-8k-qlora.yml

@@ -1,6 +1,7 @@
 # An example finetuning Saleforce's XGen-7b model with 8k context using qlora
 # on Tim Dettmer's Guanaco dataset.
 base_model: Salesforce/xgen-7b-8k-base
+base_model_config: Salesforce/xgen-7b-8k-base
 trust_remote_code: true
 model_type: AutoModelForCausalLM
 tokenizer_type: AutoTokenizer
@@ -16,10 +17,10 @@ datasets:
       - openassistant_best_replies_train.jsonl
     type: "completion"
 dataset_prepared_path:
-val_set_size: 0.05
+val_set_size: 0.01
 # enable QLoRA
 adapter: qlora
-lora_model_dir:
+peft_model_dir:
 sequence_len: 8192
 max_packed_sequence_len:
 
@@ -38,7 +39,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_name:
+wandb_run_id:
 wandb_log_model:
 output_dir: ./qlora-out
 
@@ -51,7 +52,7 @@ output_dir: ./qlora-out
 # decrease if OOM, increase for max VRAM utilization
 micro_batch_size: 1
 gradient_accumulation_steps: 1
-num_epochs: 4
+num_epochs: 3
 # Optimizer for QLoRA
 optimizer: paged_adamw_32bit
 torchdistx_path:

requirements.txt

@@ -1,22 +1,23 @@
+--extra-index-url https://download.pytorch.org/whl/cu118
 --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
-auto-gptq==0.5.1
+torch==2.0.1
+auto-gptq
 packaging
-peft==0.6.0
-transformers @ git+https://github.com/huggingface/transformers.git@df5c5c62ae253055336f5bb0828ca8e3e15ab6bd
-tokenizers==0.15.0
+peft @ git+https://github.com/huggingface/peft.git
+transformers @ git+https://github.com/huggingface/transformers.git@bd6205919aad4d3a2300a39a98a642f1cc3a5348
 bitsandbytes>=0.41.1
-accelerate==0.24.1
+accelerate @ git+https://github.com/huggingface/accelerate@80da9cfb09bb3cc9f1b385cb55d6b90d025a5fd9
 deepspeed
 addict
 fire
 PyYAML>=6.0
-datasets>=2.15.0
-flash-attn==2.3.3
+datasets
+flash-attn>=2.3.0
 sentencepiece
 wandb
 einops
-xformers==0.0.22
-optimum==1.13.2
+xformers>=0.0.22
+optimum
 hf_transfer
 colorama
 numba
@@ -30,10 +31,3 @@ scikit-learn==1.2.2
 pynvml
 art
 fschat==0.2.29
-gradio==3.50.2
-tensorboard
-
-# remote filesystems
-s3fs
-gcsfs
-# adlfs

scripts/finetune.py

@@ -45,6 +45,8 @@ def do_cli(config: Path = Path("examples/"), **kwargs):
         shard(cfg=parsed_cfg, cli_args=parsed_cli_args)
     else:
         dataset_meta = load_datasets(cfg=parsed_cfg, cli_args=parsed_cli_args)
+        if parsed_cli_args.prepare_ds_only:
+            return
         train(cfg=parsed_cfg, cli_args=parsed_cli_args, dataset_meta=dataset_meta)
 
 

setup.py

@@ -46,13 +46,10 @@ setup(
     dependency_links=dependency_links,
     extras_require={
         "flash-attn": [
-            "flash-attn==2.3.3",
+            "flash-attn>=2.3.0",
         ],
         "deepspeed": [
             "deepspeed",
         ],
-        "mamba-ssm": [
-            "mamba-ssm==1.0.1",
-        ],
     },
 )

src/axolotl/cli/__init__.py

@@ -6,10 +6,8 @@ import os
 import random
 import sys
 from pathlib import Path
-from threading import Thread
 from typing import Any, Dict, List, Optional, Union
 
-import gradio as gr
 import torch
 import yaml
 
@@ -18,7 +16,7 @@ from accelerate.commands.config import config_args
 from art import text2art
 from huggingface_hub import HfApi
 from huggingface_hub.utils import LocalTokenNotFoundError
-from transformers import GenerationConfig, TextIteratorStreamer, TextStreamer
+from transformers import GenerationConfig, TextStreamer
 
 from axolotl.common.cli import TrainerCliArgs, load_model_and_tokenizer
 from axolotl.logging_config import configure_logging
@@ -29,7 +27,6 @@ from axolotl.utils.dict import DictDefault
 from axolotl.utils.distributed import is_main_process
 from axolotl.utils.models import load_tokenizer
 from axolotl.utils.tokenization import check_dataset_labels
-from axolotl.utils.trainer import prepare_optim_env
 from axolotl.utils.wandb_ import setup_wandb_env_vars
 
 project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
@@ -47,7 +44,7 @@ def print_axolotl_text_art(suffix=None):
     ascii_text = "  axolotl"
     if suffix:
         ascii_text += f"  x  {suffix}"
-    ascii_art = text2art(ascii_text, font=font)
+    ascii_art = text2art(" axolotl", font=font)
 
     if is_main_process():
         print(ascii_art)
@@ -72,7 +69,7 @@ def do_merge_lora(
 
     LOG.info("running merge of LoRA with base model")
     model = model.merge_and_unload()
-    model.to(dtype=cfg.torch_dtype)
+    model.to(dtype=torch.float16)
 
     if cfg.local_rank == 0:
         LOG.info(f"saving merged model to: {str(Path(cfg.output_dir) / 'merged')}")
@@ -156,91 +153,6 @@ def do_inference(
         print(tokenizer.decode(generated["sequences"].cpu().tolist()[0]))
 
 
-def do_inference_gradio(
-    *,
-    cfg: DictDefault,
-    cli_args: TrainerCliArgs,
-):
-    model, tokenizer = load_model_and_tokenizer(cfg=cfg, cli_args=cli_args)
-    prompter = cli_args.prompter
-    default_tokens = {"unk_token": "<unk>", "bos_token": "<s>", "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
-        )
-
-    model = model.to(cfg.device)
-
-    def generate(instruction):
-        if not instruction:
-            return
-        if prompter_module:
-            # pylint: disable=stop-iteration-return
-            prompt: str = next(
-                prompter_module().build_prompt(instruction=instruction.strip("\n"))
-            )
-        else:
-            prompt = instruction.strip()
-        batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
-
-        model.eval()
-        with torch.no_grad():
-            generation_config = GenerationConfig(
-                repetition_penalty=1.1,
-                max_new_tokens=1024,
-                temperature=0.9,
-                top_p=0.95,
-                top_k=40,
-                bos_token_id=tokenizer.bos_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-                pad_token_id=tokenizer.pad_token_id,
-                do_sample=True,
-                use_cache=True,
-                return_dict_in_generate=True,
-                output_attentions=False,
-                output_hidden_states=False,
-                output_scores=False,
-            )
-            streamer = TextIteratorStreamer(tokenizer)
-            generation_kwargs = {
-                "inputs": batch["input_ids"].to(cfg.device),
-                "generation_config": generation_config,
-                "streamer": streamer,
-            }
-
-            thread = Thread(target=model.generate, kwargs=generation_kwargs)
-            thread.start()
-
-            all_text = ""
-
-            for new_text in streamer:
-                all_text += new_text
-                yield all_text
-
-    demo = gr.Interface(
-        fn=generate,
-        inputs="textbox",
-        outputs="text",
-        title=cfg.get("gradio_title", "Axolotl Gradio Interface"),
-    )
-    demo.queue().launch(show_api=False, share=True)
-
-
 def choose_config(path: Path):
     yaml_files = list(path.glob("*.yml"))
 
@@ -297,8 +209,6 @@ def load_cfg(config: Path = Path("examples/"), **kwargs):
 
     validate_config(cfg)
 
-    prepare_optim_env(cfg)
-
     normalize_config(cfg)
 
     setup_wandb_env_vars(cfg)
@@ -312,9 +222,7 @@ def load_datasets(
 ) -> TrainDatasetMeta:
     tokenizer = load_tokenizer(cfg)
 
-    train_dataset, eval_dataset, total_num_steps, prompters = prepare_dataset(
-        cfg, tokenizer
-    )
+    train_dataset, eval_dataset, total_num_steps = prepare_dataset(cfg, tokenizer)
 
     if cli_args.debug or cfg.debug:
         LOG.info("check_dataset_labels...")
@@ -330,10 +238,6 @@ def load_datasets(
             text_only=cli_args.debug_text_only,
         )
 
-        LOG.info("printing prompters...")
-        for prompter in prompters:
-            LOG.info(prompter)
-
     return TrainDatasetMeta(
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,

src/axolotl/cli/inference.py

@@ -6,16 +6,11 @@ from pathlib import Path
 import fire
 import transformers
 
-from axolotl.cli import (
-    do_inference,
-    do_inference_gradio,
-    load_cfg,
-    print_axolotl_text_art,
-)
+from axolotl.cli import do_inference, load_cfg, print_axolotl_text_art
 from axolotl.common.cli import TrainerCliArgs
 
 
-def do_cli(config: Path = Path("examples/"), gradio=False, **kwargs):
+def do_cli(config: Path = Path("examples/"), **kwargs):
     # pylint: disable=duplicate-code
     print_axolotl_text_art()
     parsed_cfg = load_cfg(config, **kwargs)
@@ -26,10 +21,7 @@ def do_cli(config: Path = Path("examples/"), gradio=False, **kwargs):
     )
     parsed_cli_args.inference = True
 
-    if gradio:
-        do_inference_gradio(cfg=parsed_cfg, cli_args=parsed_cli_args)
-    else:
-        do_inference(cfg=parsed_cfg, cli_args=parsed_cli_args)
+    do_inference(cfg=parsed_cfg, cli_args=parsed_cli_args)
 
 
 if __name__ == "__main__":

src/axolotl/cli/preprocess.py

@@ -1,53 +0,0 @@
-"""
-CLI to run training on a model
-"""
-import logging
-from pathlib import Path
-
-import fire
-import transformers
-from colorama import Fore
-
-from axolotl.cli import (
-    check_accelerate_default_config,
-    check_user_token,
-    load_cfg,
-    load_datasets,
-    print_axolotl_text_art,
-)
-from axolotl.common.cli import PreprocessCliArgs
-from axolotl.common.const import DEFAULT_DATASET_PREPARED_PATH
-
-LOG = logging.getLogger("axolotl.cli.preprocess")
-
-
-def do_cli(config: Path = Path("examples/"), **kwargs):
-    # pylint: disable=duplicate-code
-    print_axolotl_text_art()
-    parsed_cfg = load_cfg(config, **kwargs)
-    check_accelerate_default_config()
-    check_user_token()
-    parser = transformers.HfArgumentParser((PreprocessCliArgs))
-    parsed_cli_args, _ = parser.parse_args_into_dataclasses(
-        return_remaining_strings=True
-    )
-    if not parsed_cfg.dataset_prepared_path:
-        msg = (
-            Fore.RED
-            + "preprocess CLI called without dataset_prepared_path set, "
-            + f"using default path: {DEFAULT_DATASET_PREPARED_PATH}"
-            + Fore.RESET
-        )
-        LOG.warning(msg)
-        parsed_cfg.dataset_prepared_path = DEFAULT_DATASET_PREPARED_PATH
-
-    _ = load_datasets(cfg=parsed_cfg, cli_args=parsed_cli_args)
-    LOG.info(
-        Fore.GREEN
-        + f"Success! Preprocessed data path: `dataset_prepared_path: {parsed_cfg.dataset_prepared_path}`"
-        + Fore.RESET
-    )
-
-
-if __name__ == "__main__":
-    fire.Fire(do_cli)

src/axolotl/cli/train.py

@@ -6,6 +6,7 @@ from pathlib import Path
 
 import fire
 import transformers
+from colorama import Fore
 
 from axolotl.cli import (
     check_accelerate_default_config,
@@ -15,6 +16,7 @@ from axolotl.cli import (
     print_axolotl_text_art,
 )
 from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.const import DEFAULT_DATASET_PREPARED_PATH
 from axolotl.train import train
 
 LOG = logging.getLogger("axolotl.cli.train")
@@ -30,7 +32,18 @@ def do_cli(config: Path = Path("examples/"), **kwargs):
     parsed_cli_args, _ = parser.parse_args_into_dataclasses(
         return_remaining_strings=True
     )
+    if parsed_cli_args.prepare_ds_only and not parsed_cfg.dataset_prepared_path:
+        msg = (
+            Fore.RED
+            + "--prepare_ds_only called without dataset_prepared_path set."
+            + Fore.RESET
+        )
+        LOG.warning(msg)
+        parsed_cfg.dataset_prepared_path = DEFAULT_DATASET_PREPARED_PATH
+
     dataset_meta = load_datasets(cfg=parsed_cfg, cli_args=parsed_cli_args)
+    if parsed_cli_args.prepare_ds_only:
+        return
     train(cfg=parsed_cfg, cli_args=parsed_cli_args, dataset_meta=dataset_meta)
 
 

src/axolotl/common/cli.py

@@ -25,22 +25,11 @@ class TrainerCliArgs:
     debug_num_examples: int = field(default=5)
     inference: bool = field(default=False)
     merge_lora: bool = field(default=False)
+    prepare_ds_only: bool = field(default=False)
     prompter: Optional[str] = field(default=None)
     shard: bool = field(default=False)
 
 
-@dataclass
-class PreprocessCliArgs:
-    """
-    dataclass representing arguments for preprocessing only
-    """
-
-    debug: bool = field(default=False)
-    debug_text_only: bool = field(default=False)
-    debug_num_examples: int = field(default=1)
-    prompter: Optional[str] = field(default=None)
-
-
 def load_model_and_tokenizer(
     *,
     cfg: DictDefault,

src/axolotl/core/trainer_builder.py

@@ -1,798 +0,0 @@
-"""
-Builder for the training args and trainer
-"""
-
-import abc
-import importlib
-import logging
-import math
-import sys
-from abc import abstractmethod
-from dataclasses import dataclass, field
-from functools import partial
-from pathlib import Path
-from typing import Optional
-
-import torch
-import transformers
-from datasets import Dataset
-from torch.optim.lr_scheduler import OneCycleLR
-from torch.utils.data import BatchSampler, DataLoader, RandomSampler, SequentialSampler
-from transformers import EarlyStoppingCallback, Trainer, TrainingArguments
-from transformers.trainer_utils import seed_worker
-
-from axolotl.monkeypatch.relora import ReLoRACallback, ReLoRAScheduler
-from axolotl.utils.callbacks import (
-    EvalFirstStepCallback,
-    GPUStatsCallback,
-    LossWatchDogCallback,
-    SaveAxolotlConfigtoWandBCallback,
-    SaveBetterTransformerModelCallback,
-    bench_eval_callback_factory,
-    log_prediction_callback_factory,
-)
-from axolotl.utils.collators import (
-    BatchSamplerDataCollatorForSeq2Seq,
-    MambaDataCollator,
-)
-from axolotl.utils.samplers import MultipackBatchSampler
-from axolotl.utils.schedulers import get_cosine_schedule_with_quadratic_warmup
-
-try:
-    import torch._dynamo  # pylint: disable=ungrouped-imports
-except ImportError:
-    pass
-
-LOG = logging.getLogger("axolotl.core.trainer_builder")
-
-
-@dataclass
-class AxolotlTrainingArguments(TrainingArguments):
-    """
-    Extend the base TrainingArguments for axolotl helpers
-    """
-
-    model_type: Optional[str] = field(
-        default=None, metadata={"help": "HF model configuration model_type."}
-    )
-    lr_quadratic_warmup: bool = field(
-        default=False,
-        metadata={"help": "Use quadratic warmup for cosine scheduling."},
-    )
-    sample_packing: bool = field(
-        default=False,
-        metadata={"help": "Use sample packing for efficient training."},
-    )
-    eval_sample_packing: Optional[bool] = field(
-        default=None,
-        metadata={"help": "Use sample packing for efficient evals."},
-    )
-    sample_packing_efficiency: float = field(
-        default=1.0,
-        metadata={"help": "Sample packing efficiency for calculating batch length."},
-    )
-    max_seq_length: int = field(
-        default=2048,
-        metadata={"help": "The maximum sequence length the model can handle"},
-    )
-    sample_packing_seq_len_multiplier: int = field(
-        default=1,
-        metadata={"help": "the multiplier for the max len for packed sequences"},
-    )
-    relora_steps: Optional[int] = field(
-        default=None,
-        metadata={"help": "how often to reset for ReLoRA"},
-    )
-    relora_warmup_steps: Optional[int] = field(
-        default=None,
-        metadata={"help": "how many warmup steps to take after reset for ReLoRA"},
-    )
-    bench_split: Optional[str] = field(
-        default="eval", metadata={"help": "The benchmark split to run on"}
-    )
-    bench_dataset: Optional[str] = field(
-        default="pharaouk/dharma-1/dharma_1_mini.json",
-        metadata={
-            "help": "Benchmark dataset to use: options are `mmlu-zs`, `mmlu-fs`, or the full path to the dataset file"
-        },
-    )
-    do_bench_eval: Optional[bool] = field(
-        default=False, metadata={"help": "Whether to run the Benchmark evaluation."}
-    )
-    max_bench_samples: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": "If set, only evaluates on `max_bench_samples` of the benchmark dataset."
-        },
-    )
-    bench_source_max_len: int = field(
-        default=2048, metadata={"help": "Maximum source sequence length for bench."}
-    )
-    dataloader_prefetch_factor: Optional[int] = field(
-        default=None,
-        metadata={"help": "prefetch_factor argument to the dataloader"},
-    )
-
-
-class AxolotlTrainer(Trainer):
-    """
-    Extend the base Trainer for axolotl helpers
-    """
-
-    args = None  # type: AxolotlTrainingArguments
-
-    def __init__(self, *args, num_epochs=1, bench_data_collator=None, **kwargs):
-        self.num_epochs = num_epochs
-        self.bench_data_collator = bench_data_collator
-        super().__init__(*args, **kwargs)
-
-    def create_scheduler(
-        self, num_training_steps: int, optimizer: torch.optim.Optimizer = None
-    ):
-        """
-        Setup the scheduler. The optimizer of the trainer must have been set up either before this method is called or
-        passed as an argument.
-
-        Args:
-            num_training_steps (int): The number of training steps to do.
-            optimizer (torch.optim.Optimizer): The training optimizer
-        """
-
-        # fmt: off
-        if self.lr_scheduler is None:  # type: ignore  # pylint: disable=access-member-before-definition
-            # fmt: on
-            if (
-                self.args.lr_scheduler_type == "cosine"
-                and self.args.lr_quadratic_warmup is True
-            ):
-                self.lr_scheduler = get_cosine_schedule_with_quadratic_warmup(  # pylint: disable=attribute-defined-outside-init
-                    optimizer,
-                    num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
-                    num_training_steps=num_training_steps,
-                )
-            else:
-                return super().create_scheduler(num_training_steps, optimizer)
-        return self.lr_scheduler
-
-    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
-        if self.args.sample_packing:
-            return MultipackBatchSampler(
-                RandomSampler(self.train_dataset),
-                self.args.train_batch_size,
-                drop_last=True,
-                batch_max_len=self._train_batch_size * self.args.max_seq_length,
-                lengths=(
-                    self.train_dataset.data.column("position_ids")
-                    .to_pandas()
-                    .apply(lambda x: x[-1] + 1)
-                    .values
-                ),
-                packing_efficiency_estimate=self.args.sample_packing_efficiency,
-            )
-        return super()._get_train_sampler()
-
-    def _get_eval_sampler(
-        self, eval_dataset: Dataset
-    ) -> Optional[torch.utils.data.Sampler]:
-        if self.args.sample_packing and self.args.eval_sample_packing is not False:
-            return MultipackBatchSampler(
-                SequentialSampler(eval_dataset),
-                self.args.per_device_eval_batch_size,
-                drop_last=True,
-                batch_max_len=self.args.eval_batch_size * self.args.max_seq_length,
-                lengths=(
-                    eval_dataset.data.column("position_ids")
-                    .to_pandas()
-                    .apply(lambda x: x[-1] + 1)
-                    .values
-                ),
-                packing_efficiency_estimate=self.args.sample_packing_efficiency,
-            )
-        return super()._get_eval_sampler(eval_dataset)
-
-    def get_train_dataloader(self) -> DataLoader:
-        if self.args.sample_packing:
-            train_dataset = self.train_dataset
-            train_dataset = train_dataset.remove_columns(["length"])
-            data_collator = self.data_collator
-            dataloader_params = {
-                "batch_size": self._train_batch_size,
-                "collate_fn": data_collator,
-                "num_workers": self.args.dataloader_num_workers,
-                "pin_memory": self.args.dataloader_pin_memory,
-            }
-            if self.args.dataloader_prefetch_factor:
-                dataloader_params[
-                    "prefetch_factor"
-                ] = self.args.dataloader_prefetch_factor
-
-            sampler = self._get_train_sampler()
-            if isinstance(sampler, BatchSampler):
-                dataloader_params["batch_sampler"] = sampler
-                del dataloader_params["batch_size"]
-            else:
-                dataloader_params["sampler"] = sampler
-                dataloader_params["drop_last"] = self.args.dataloader_drop_last
-            dataloader_params["worker_init_fn"] = seed_worker
-
-            self.accelerator.even_batches = False
-            return self.accelerator.prepare_data_loader(
-                DataLoader(train_dataset, **dataloader_params)
-            )
-        return super().get_train_dataloader()
-
-    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
-        if self.args.sample_packing and self.args.eval_sample_packing is not False:
-            eval_dataset = (
-                eval_dataset if eval_dataset is not None else self.eval_dataset
-            )
-
-            eval_sampler = self._get_eval_sampler(eval_dataset)
-            eval_dataset = eval_dataset.remove_columns(["length"])
-            data_collator = self.data_collator
-            dataloader_params = {
-                "batch_size": self.args.eval_batch_size,
-                "collate_fn": data_collator,
-                "num_workers": self.args.dataloader_num_workers,
-                "pin_memory": self.args.dataloader_pin_memory,
-            }
-            if self.args.dataloader_prefetch_factor:
-                dataloader_params[
-                    "prefetch_factor"
-                ] = self.args.dataloader_prefetch_factor
-
-            if isinstance(eval_sampler, BatchSampler):
-                dataloader_params["batch_sampler"] = eval_sampler
-                del dataloader_params["batch_size"]
-            else:
-                dataloader_params["sampler"] = eval_sampler
-                dataloader_params["drop_last"] = self.args.dataloader_drop_last
-
-            self.accelerator.even_batches = False
-            return self.accelerator.prepare_data_loader(
-                DataLoader(eval_dataset, **dataloader_params)
-            )
-        return super().get_eval_dataloader(eval_dataset)
-
-    def _get_bench_sampler(
-        self, bench_dataset: Dataset
-    ) -> Optional[torch.utils.data.Sampler]:
-        if self.args.world_size <= 1:
-            return SequentialSampler(bench_dataset)
-        return None
-
-    def get_bench_dataloader(
-        self,
-        bench_dataset: Dataset,
-    ) -> DataLoader:
-        dataloader_params = {
-            "batch_size": self.args.eval_batch_size,
-            "collate_fn": self.bench_data_collator,
-            "num_workers": self.args.dataloader_num_workers,
-            "pin_memory": self.args.dataloader_pin_memory,
-        }
-        if self.args.dataloader_prefetch_factor:
-            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
-
-        if not isinstance(bench_dataset, torch.utils.data.IterableDataset):
-            dataloader_params["sampler"] = self._get_bench_sampler(bench_dataset)
-            dataloader_params["drop_last"] = self.args.dataloader_drop_last
-
-        return DataLoader(bench_dataset, **dataloader_params)
-        # return self.accelerator.prepare(DataLoader(bench_dataset, **dataloader_params))
-
-    def compute_loss(self, model, inputs, return_outputs=False):
-        # use one's weighted cross entropy loss calc
-        # if self.args.sample_packing:
-        #     labels = inputs.pop("labels")
-        #     outputs = model(**inputs)
-        #     loss = trainer_weighted_loss(outputs, labels, shift_labels=True)
-        #     return (loss, outputs) if return_outputs else loss
-        return super().compute_loss(model, inputs, return_outputs=return_outputs)
-
-
-class AxolotlMambaTrainer(AxolotlTrainer):
-    """
-    Mamba specific trainer to handle loss calculation
-    """
-
-    def compute_loss(
-        self,
-        model,
-        inputs,
-        return_outputs=False,  # pylint: disable=unused-argument
-    ):
-        input_ids = inputs.pop("input_ids")
-        lm_logits = model(input_ids).logits
-
-        labels = input_ids.to(lm_logits.device)
-        shift_logits = lm_logits[:, :-1, :].contiguous()
-        labels = labels[:, 1:].contiguous()
-
-        loss_fct = torch.nn.CrossEntropyLoss()
-        lm_loss = loss_fct(
-            shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1)
-        )
-
-        return lm_loss
-
-
-class OneCycleLRSchedulerTrainer(AxolotlTrainer):
-    """
-    Trainer subclass that uses the OneCycleLR scheduler
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.lr_scheduler = None
-
-    def create_scheduler(
-        self,
-        num_training_steps: int,
-        optimizer: Optional[torch.optim.Optimizer] = None,
-    ):
-        optimizer = self.optimizer if optimizer is None else optimizer
-        num_warmup_steps = self.args.get_warmup_steps(num_training_steps)
-        pct_start = num_warmup_steps / num_training_steps
-
-        self.lr_scheduler = OneCycleLR(
-            optimizer,
-            max_lr=self.args.learning_rate,
-            total_steps=num_training_steps,
-            pct_start=pct_start,
-            div_factor=6,
-        )
-
-        return self.lr_scheduler
-
-
-class ReLoRATrainer(AxolotlTrainer):
-    """
-    Trainer subclass that uses the OneCycleLR scheduler
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.lr_scheduler = None
-
-    def create_scheduler(
-        self,
-        num_training_steps: int,
-        optimizer: Optional[torch.optim.Optimizer] = None,
-    ):
-        optimizer = self.optimizer if optimizer is None else optimizer
-        lr_scheduler = super().create_scheduler(num_training_steps, optimizer)
-
-        if self.args.relora_steps:
-            warmup_steps = (
-                self.args.relora_warmup_steps if self.args.relora_warmup_steps else 10
-            )
-            self.lr_scheduler = ReLoRAScheduler(
-                optimizer,
-                lr_scheduler,
-                self.args.relora_steps,
-                warmup_steps,
-            )
-        else:
-            self.lr_scheduler = lr_scheduler
-
-        return self.lr_scheduler
-
-
-class TrainerBuilderBase(abc.ABC):
-    """
-    Base class for trainer builder
-    """
-
-    _train_dataset = None
-    _eval_dataset = None
-
-    def __init__(self, cfg, model, tokenizer):
-        self.cfg = cfg
-        self.model = model
-        self.tokenizer = tokenizer
-
-    @property
-    def train_dataset(self):
-        return self._train_dataset
-
-    @train_dataset.setter
-    def train_dataset(self, dataset):
-        self._train_dataset = dataset
-
-    @property
-    def eval_dataset(self):
-        return self._eval_dataset
-
-    @eval_dataset.setter
-    def eval_dataset(self, dataset):
-        self._eval_dataset = dataset
-
-    @abstractmethod
-    def build(self, total_num_steps):
-        pass
-
-    @abstractmethod
-    def get_callbacks(self):
-        pass
-
-    @abstractmethod
-    def get_post_trainer_create_callbacks(self, trainer):
-        """
-        Callbacks added after the trainer is created, usually b/c these need access to the trainer
-        """
-
-
-class HFCausalTrainerBuilder(TrainerBuilderBase):
-    """
-    Build the HuggingFace training args/trainer for Causal models
-    """
-
-    def hook_pre_create_training_args(self, training_arguments_kwargs):
-        # TODO
-        return training_arguments_kwargs
-
-    def hook_post_create_training_args(self, training_arguments):
-        # TODO
-        return training_arguments
-
-    def hook_pre_create_trainer(self, trainer_kwargs, trainer_cls):
-        # TODO
-        return trainer_kwargs, trainer_cls
-
-    def hook_post_create_trainer(self, trainer):
-        # TODO
-        return trainer
-
-    def get_callbacks(self):
-        callbacks = []
-        callbacks.append(GPUStatsCallback(self.cfg))
-        callbacks.append(EvalFirstStepCallback)
-
-        if self.cfg.relora_steps:
-            callbacks.append(ReLoRACallback(self.cfg))
-
-        if (
-            hasattr(self.model, "use_bettertransformer")
-            and self.model.use_bettertransformer is True
-        ):
-            callbacks.append(SaveBetterTransformerModelCallback)
-
-        if self.cfg.use_wandb:
-            callbacks.append(
-                SaveAxolotlConfigtoWandBCallback(self.cfg.axolotl_config_path)
-            )
-
-        if self.cfg.loss_watchdog_threshold is not None:
-            callbacks.append(LossWatchDogCallback(self.cfg))
-
-        return callbacks
-
-    def get_post_trainer_create_callbacks(self, trainer):
-        callbacks = []
-        if self.cfg.use_wandb and self.cfg.eval_table_size > 0:
-            LogPredictionCallback = log_prediction_callback_factory(
-                trainer, self.tokenizer
-            )
-            callbacks.append(LogPredictionCallback(self.cfg))
-
-        if self.cfg.do_bench_eval:
-            callbacks.append(bench_eval_callback_factory(trainer, self.tokenizer))
-
-        if self.cfg.early_stopping_patience:
-            early_stop_cb = EarlyStoppingCallback(
-                self.cfg.early_stopping_patience,
-            )
-            callbacks.append(early_stop_cb)
-
-        return callbacks
-
-    def _get_trainer_cls(self):
-        if self.cfg.lr_scheduler == "one_cycle" and (
-            self.cfg.fsdp or self.cfg.adapter == "qlora"
-        ):
-            return OneCycleLRSchedulerTrainer
-        if self.cfg.relora_steps:
-            return ReLoRATrainer
-        if self.cfg.model_config_type == "mamba":
-            return AxolotlMambaTrainer
-        return AxolotlTrainer
-
-    def build(self, total_num_steps):
-        warmup_steps = None
-        if self.cfg.warmup_steps is not None:
-            warmup_steps = self.cfg.warmup_steps
-        elif self.cfg.warmup_ratio is not None:
-            warmup_steps = max(int(self.cfg.warmup_ratio * total_num_steps), 0)
-        else:
-            warmup_steps = min(int(0.03 * total_num_steps), 100)
-
-        logging_steps = (
-            self.cfg.logging_steps
-            if self.cfg.logging_steps is not None
-            else max(min(int(0.005 * total_num_steps), 10), 1)
-        )
-
-        training_arguments_kwargs = {}
-        if self.cfg.bf16 == "full":
-            training_arguments_kwargs["bf16_full_eval"] = True
-        else:
-            training_arguments_kwargs["bf16"] = self.cfg.bf16
-        training_arguments_kwargs["fp16"] = (
-            self.cfg.fp16 and not self.cfg.bf16
-        ) or False
-        training_arguments_kwargs["tf32"] = self.cfg.tf32
-        training_arguments_kwargs["warmup_steps"] = warmup_steps
-        training_arguments_kwargs["logging_steps"] = logging_steps
-
-        if self.cfg.seed:
-            training_arguments_kwargs["seed"] = self.cfg.seed
-
-        if self.cfg.gradient_checkpointing:
-            training_arguments_kwargs[
-                "gradient_checkpointing"
-            ] = self.cfg.gradient_checkpointing
-        if self.cfg.fsdp:
-            training_arguments_kwargs["fsdp"] = self.cfg.fsdp
-            if self.cfg.fsdp_config:
-                training_arguments_kwargs["fsdp_config"] = dict(self.cfg.fsdp_config)
-
-        # deepspeed
-        if self.cfg.deepspeed:
-            training_arguments_kwargs["deepspeed"] = self.cfg.deepspeed
-
-        if self.cfg.lr_quadratic_warmup is not None:
-            training_arguments_kwargs[
-                "lr_quadratic_warmup"
-            ] = self.cfg.lr_quadratic_warmup
-
-        if self.cfg.adam_beta1:
-            training_arguments_kwargs["adam_beta1"] = self.cfg.adam_beta1
-        if self.cfg.adam_beta2:
-            training_arguments_kwargs["adam_beta2"] = self.cfg.adam_beta2
-        if self.cfg.adam_epsilon:
-            training_arguments_kwargs["adam_epsilon"] = self.cfg.adam_epsilon
-        if self.cfg.max_grad_norm:
-            training_arguments_kwargs["max_grad_norm"] = self.cfg.max_grad_norm
-
-        if self.cfg.hub_model_id:
-            training_arguments_kwargs["hub_model_id"] = self.cfg.hub_model_id
-            training_arguments_kwargs["push_to_hub"] = True
-            training_arguments_kwargs["hub_private_repo"] = True
-
-            if self.cfg.hub_strategy:
-                training_arguments_kwargs["hub_strategy"] = self.cfg.hub_strategy
-
-        if self.cfg.save_safetensors is not None:
-            training_arguments_kwargs["save_safetensors"] = self.cfg.save_safetensors
-
-        if self.cfg.sample_packing_eff_est:
-            training_arguments_kwargs[
-                "sample_packing_efficiency"
-            ] = self.cfg.sample_packing_eff_est
-
-        if self.cfg.dataloader_pin_memory is not None:
-            training_arguments_kwargs[
-                "dataloader_pin_memory"
-            ] = self.cfg.dataloader_pin_memory
-        if self.cfg.dataloader_num_workers is not None:
-            training_arguments_kwargs[
-                "dataloader_num_workers"
-            ] = self.cfg.dataloader_num_workers
-        if self.cfg.dataloader_prefetch_factor is not None:
-            training_arguments_kwargs[
-                "dataloader_prefetch_factor"
-            ] = self.cfg.dataloader_prefetch_factor
-
-        if self.cfg.val_set_size == 0:
-            # no eval set, so don't eval
-            training_arguments_kwargs["evaluation_strategy"] = "no"
-        elif self.cfg.eval_steps:
-            training_arguments_kwargs["evaluation_strategy"] = "steps"
-            training_arguments_kwargs["eval_steps"] = self.cfg.eval_steps
-        elif self.cfg.evaluation_strategy:
-            training_arguments_kwargs[
-                "evaluation_strategy"
-            ] = self.cfg.evaluation_strategy
-        else:
-            # we have an eval set, but no steps defined, default to use epoch
-            training_arguments_kwargs["evaluation_strategy"] = "epoch"
-
-        if self.cfg.save_steps:
-            training_arguments_kwargs["save_strategy"] = "steps"
-            training_arguments_kwargs["save_steps"] = self.cfg.save_steps
-        elif self.cfg.save_strategy:
-            training_arguments_kwargs["save_strategy"] = self.cfg.save_strategy
-        else:
-            # default to saving each epoch if not defined
-            training_arguments_kwargs["save_strategy"] = "epoch"
-
-        if self.cfg.do_bench_eval:
-            training_arguments_kwargs["do_bench_eval"] = self.cfg.do_bench_eval
-            if self.cfg.bench_dataset:
-                training_arguments_kwargs["bench_dataset"] = self.cfg.bench_dataset
-        if self.cfg.metric_for_best_model:
-            training_arguments_kwargs[
-                "metric_for_best_model"
-            ] = self.cfg.metric_for_best_model
-        if self.cfg.greater_is_better:
-            training_arguments_kwargs["greater_is_better"] = self.cfg.greater_is_better
-
-        if self.cfg.torch_compile:
-            if torch.__version__ < "2.1.0":  # pylint: disable=protected-access
-                LOG.warning("torch>=2.1.0 required for torch_compile to work properly")
-            elif torch._dynamo:  # pylint: disable=protected-access
-                torch._dynamo.config.suppress_errors = (  # pylint: disable=protected-access
-                    True
-                )
-                training_arguments_kwargs["torch_compile"] = self.cfg.torch_compile
-                if self.cfg.torch_compile_backend:
-                    training_arguments_kwargs[
-                        "torch_compile_backend"
-                    ] = self.cfg.torch_compile_backend
-
-        # DDP Config
-        if self.cfg.ddp_timeout:
-            training_arguments_kwargs["ddp_timeout"] = self.cfg.ddp_timeout
-        # see https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html
-        if self.cfg.ddp_bucket_cap_mb:
-            training_arguments_kwargs["ddp_bucket_cap_mb"] = self.cfg.ddp_bucket_cap_mb
-        if self.cfg.ddp_broadcast_buffers is not None:
-            training_arguments_kwargs[
-                "ddp_broadcast_buffers"
-            ] = self.cfg.ddp_broadcast_buffers
-
-        # these are all the "standard" kwargs that are def used
-        training_arguments_kwargs["max_steps"] = (
-            total_num_steps if self.cfg.max_steps else -1
-        )
-        training_arguments_kwargs["max_seq_length"] = self.cfg.sequence_len
-        training_arguments_kwargs[
-            "per_device_train_batch_size"
-        ] = self.cfg.micro_batch_size
-        training_arguments_kwargs[
-            "per_device_eval_batch_size"
-        ] = self.cfg.eval_batch_size
-        training_arguments_kwargs[
-            "gradient_accumulation_steps"
-        ] = self.cfg.gradient_accumulation_steps
-        training_arguments_kwargs[
-            "eval_accumulation_steps"
-        ] = self.cfg.gradient_accumulation_steps
-        training_arguments_kwargs["num_train_epochs"] = self.cfg.num_epochs
-        training_arguments_kwargs["learning_rate"] = self.cfg.learning_rate
-        training_arguments_kwargs["output_dir"] = self.cfg.output_dir
-        training_arguments_kwargs["save_total_limit"] = (
-            self.cfg.save_total_limit if self.cfg.save_total_limit else 4
-        )
-        training_arguments_kwargs["load_best_model_at_end"] = (
-            (
-                self.cfg.load_best_model_at_end is not False
-                or self.cfg.early_stopping_patience
-            )
-            and self.cfg.val_set_size > 0
-            and self.cfg.save_steps
-            and self.cfg.eval_steps
-            and self.cfg.save_steps % self.cfg.eval_steps == 0
-        ) or False
-        training_arguments_kwargs["ddp_find_unused_parameters"] = (
-            False if self.cfg.ddp else None
-        )
-        training_arguments_kwargs["group_by_length"] = self.cfg.group_by_length
-        training_arguments_kwargs["report_to"] = "wandb" if self.cfg.use_wandb else None
-        training_arguments_kwargs["run_name"] = (
-            self.cfg.wandb_name if self.cfg.use_wandb else None
-        )
-        training_arguments_kwargs["optim"] = (
-            self.cfg.optimizer if self.cfg.optimizer else "adamw_hf"
-        )
-        training_arguments_kwargs["lr_scheduler_type"] = (
-            self.cfg.lr_scheduler
-            if self.cfg.lr_scheduler
-            and self.cfg.lr_scheduler not in ("one_cycle", "log_sweep")
-            else "cosine"
-        )
-        training_arguments_kwargs["weight_decay"] = (
-            self.cfg.weight_decay if self.cfg.weight_decay is not None else 0.0
-        )
-        training_arguments_kwargs["sample_packing"] = (
-            self.cfg.sample_packing if self.cfg.sample_packing else False
-        )
-        training_arguments_kwargs["eval_sample_packing"] = (
-            self.cfg.sample_packing
-            if self.cfg.eval_sample_packing is not False
-            else False
-        )
-        training_arguments_kwargs[
-            "sample_packing_seq_len_multiplier"
-        ] = self.cfg.micro_batch_size
-        training_arguments_kwargs["relora_steps"] = self.cfg.relora_steps
-        training_arguments_kwargs["relora_warmup_steps"] = self.cfg.relora_warmup_steps
-        training_arguments_kwargs = self.hook_pre_create_training_args(
-            training_arguments_kwargs
-        )
-        training_arguments_kwargs["model_type"] = self.cfg.model_config_type
-        training_args = (
-            AxolotlTrainingArguments(  # pylint: disable=unexpected-keyword-arg
-                **training_arguments_kwargs,
-            )
-        )
-        training_args = self.hook_post_create_training_args(training_args)
-        trainer_kwargs = {}
-
-        if self.cfg.optimizer == "adamw_anyprecision":
-            if Path(self.cfg.torchdistx_path).exists():
-                sys.path.append(self.cfg.torchdistx_path)
-                importlib.import_module("torchdistx")
-
-        data_collator_kwargs = {
-            "padding": True,  # True/"longest" is the default
-        }
-        if self.cfg.pad_to_sequence_len:
-            data_collator_kwargs["pad_to_multiple_of"] = 64 * math.ceil(
-                self.cfg.sequence_len / 64
-            )
-        else:
-            # A100 is best at 64, while others at 8. Let's use the larger so we don't have to check
-            # https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html
-            data_collator_kwargs["pad_to_multiple_of"] = 64
-
-        if self.cfg.is_llama_derived_model and self.cfg.landmark_attention:
-            from axolotl.monkeypatch.llama_landmark_attn import (
-                add_mem_tokens,
-                get_mem_id,
-                set_model_mem_id,
-            )
-
-            set_model_mem_id(self.model, self.tokenizer)
-
-            LOG.info("Adding landmark attention tokens to dataset")
-
-            for dataset in [self.train_dataset, self.eval_dataset]:
-                dataset = dataset.map(
-                    partial(
-                        add_mem_tokens, mem_freq=50, mem_id=get_mem_id(self.tokenizer)
-                    ),
-                    batched=False,
-                    num_proc=32,
-                )
-
-        trainer_cls = self._get_trainer_cls()
-        trainer_kwargs, trainer_cls = self.hook_pre_create_trainer(
-            trainer_kwargs, trainer_cls
-        )
-        trainer = trainer_cls(
-            model=self.model,
-            train_dataset=self.train_dataset,
-            eval_dataset=self.eval_dataset,
-            args=training_args,
-            data_collator=self.build_collator(**data_collator_kwargs),
-            bench_data_collator=transformers.DataCollatorForSeq2Seq(
-                self.tokenizer,
-                return_tensors="pt",
-                **data_collator_kwargs,
-            ),
-            callbacks=self.get_callbacks(),
-            num_epochs=self.cfg.num_epochs,
-            **trainer_kwargs,
-        )
-        trainer = self.hook_post_create_trainer(trainer)
-        for callback in self.get_post_trainer_create_callbacks(trainer):
-            trainer.add_callback(callback)
-
-        if self.cfg.deepspeed and self.cfg.sample_packing:
-            trainer.accelerator.state.deepspeed_plugin.deepspeed_config[
-                "train_micro_batch_size_per_gpu"
-            ] = self.cfg.micro_batch_size
-
-        return trainer
-
-    def build_collator(self, **kwargs):
-        if self.cfg.model_config_type == "mamba":
-            return MambaDataCollator(tokenizer=self.tokenizer)
-
-        return BatchSamplerDataCollatorForSeq2Seq(
-            self.tokenizer,
-            return_tensors="pt",
-            **kwargs,
-        )

src/axolotl/datasets.py

@@ -2,7 +2,7 @@
 
 import logging
 import os
-from typing import List, Optional
+from typing import List
 
 import torch
 from datasets import Dataset, IterableDataset
@@ -30,20 +30,14 @@ class TokenizedPromptDataset(Dataset):
         self,
         prompt_tokenizer: PromptTokenizingStrategy,
         dataset: IterableDataset,
-        process_count: Optional[int] = None,
         **kwargs,
     ):
         self.prompt_tokenizer = prompt_tokenizer
-        self.process_count = process_count
         super().__init__(self.process(dataset).data, **kwargs)
 
     def process(self, dataset):
         features = dataset.features.keys()
-        num_proc = (
-            min(64, self.process_count)
-            if self.process_count
-            else min(64, os.cpu_count())
-        )
+        num_proc = min(64, os.cpu_count())
         map_kwargs = {}
         if self.prompt_tokenizer.supports_batched:
             map_kwargs["batched"] = True

src/axolotl/models/mamba/__init__.py

@@ -1,12 +0,0 @@
-"""
-Modeling module for Mamba models
-"""
-
-
-def fix_mamba_attn_for_loss():
-    from mamba_ssm.models import mixer_seq_simple
-
-    from .modeling_mamba import MambaLMHeadModel as MambaLMHeadModelFixed
-
-    mixer_seq_simple.MambaLMHeadModel = MambaLMHeadModelFixed
-    return mixer_seq_simple.MambaLMHeadModel  # pylint: disable=invalid-name

src/axolotl/models/mamba/configuration_mamba.py

@@ -1,42 +0,0 @@
-"""
-HF Transformers MambaConfig
-"""
-from transformers import PretrainedConfig
-
-
-class MambaConfig(PretrainedConfig):
-    """
-    modeling configuration for state space model/mamba
-    """
-
-    model_type = "mamba"
-
-    def __init__(
-        self,
-        vocab_size=50280,
-        d_model=2560,
-        n_layer=64,
-        rms_norm=True,
-        residual_in_fp32=True,
-        fused_add_norm=True,
-        pad_vocab_size_multiple=8,
-        pad_token_id=50277,
-        bos_token_id=0,
-        eos_token_id=0,
-        tie_word_embeddings=False,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.n_layer = n_layer
-        self.rms_norm = rms_norm
-        self.residual_in_fp32 = residual_in_fp32
-        self.fused_add_norm = fused_add_norm
-        self.pad_vocab_size_multiple = pad_vocab_size_multiple
-        super().__init__(
-            pad_token_id=pad_token_id,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
-        )

src/axolotl/models/mamba/modeling_mamba.py

@@ -1,128 +0,0 @@
-# pylint: skip-file
-import os
-from collections import namedtuple
-from functools import partial
-from typing import Optional, Union
-
-import torch
-from mamba_ssm.models.mixer_seq_simple import MixerModel, _init_weights
-from mamba_ssm.utils.generation import GenerationMixin
-from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf
-from torch import nn
-from torch.nn import CrossEntropyLoss
-
-from axolotl.models.mamba.configuration_mamba import MambaConfig
-
-
-class MambaLMHeadModel(nn.Module, GenerationMixin):
-    def __init__(
-        self,
-        d_model: int,
-        n_layer: int,
-        vocab_size: int,
-        initializer_cfg=None,
-        pad_vocab_size_multiple: int = 1,
-        device=None,
-        dtype=None,
-        **backbone_kwargs,
-    ) -> None:
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        if vocab_size % pad_vocab_size_multiple != 0:
-            vocab_size += pad_vocab_size_multiple - (
-                vocab_size % pad_vocab_size_multiple
-            )
-        self.config = MambaConfig(
-            vocab_size=vocab_size,
-            d_model=d_model,
-            n_layer=n_layer,
-            pad_vocab_size_multiple=pad_vocab_size_multiple,
-        )
-        self.backbone = MixerModel(
-            d_model=d_model,
-            n_layer=n_layer,
-            vocab_size=vocab_size,
-            initializer_cfg=initializer_cfg,
-            **backbone_kwargs,
-            **factory_kwargs,
-        )
-        self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)
-
-        # Initialize weights and apply final processing
-        self.apply(
-            partial(
-                _init_weights,
-                n_layer=n_layer,
-                **(initializer_cfg if initializer_cfg is not None else {}),
-            )
-        )
-        self.tie_weights()
-
-    def tie_weights(self):
-        self.lm_head.weight = self.backbone.embedding.weight
-
-    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
-        return self.backbone.allocate_inference_cache(
-            batch_size, max_seqlen, dtype=dtype, **kwargs
-        )
-
-    def forward(
-        self,
-        input_ids,
-        position_ids=None,
-        inference_params=None,
-        num_last_tokens=0,
-        labels=None,
-        **kwargs,
-    ):
-        """
-        "position_ids" is just to be compatible with Transformer generation. We don't use it.
-        num_last_tokens: if > 0, only return the logits for the last n tokens
-        """
-        hidden_states = self.backbone(input_ids, inference_params=inference_params)
-        if num_last_tokens > 0:
-            hidden_states = hidden_states[:, -num_last_tokens:]
-        lm_logits = self.lm_head(hidden_states)
-
-        CausalLMOutput = namedtuple("CausalLMOutput", ["logits"])
-        return CausalLMOutput(logits=lm_logits)
-
-        loss = None
-        if labels is not None:
-            logits = lm_logits
-            # Shift so that tokens < n predict n
-            shift_logits = logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            # Flatten the tokens
-            loss_fct = CrossEntropyLoss()
-            shift_logits = shift_logits.view(-1, self.config.vocab_size)
-            shift_labels = shift_labels.view(-1)
-            # Enable model parallelism
-            shift_labels = shift_labels.to(shift_logits.device)
-            loss = loss_fct(shift_logits, shift_labels)
-            CausalLMOutput = namedtuple("CausalLMOutput", ["logits", "loss"])
-            print(loss)
-            return CausalLMOutput(logits=lm_logits, loss=loss)
-
-        else:
-            CausalLMOutput = namedtuple("CausalLMOutput", ["logits"])
-            return CausalLMOutput(logits=lm_logits)
-
-    def save_pretrained(
-        self,
-        save_directory: Union[str, os.PathLike],
-        state_dict: Optional[dict] = None,
-        safe_serialization: Optional[bool] = None,  # pylint: disable=unused-argument
-    ):
-        if state_dict is None:
-            state_dict = self.state_dict()
-        torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin"))
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name, device=None, dtype=None, **kwargs):
-        config = load_config_hf(pretrained_model_name)
-        model = cls(**config, device=device, dtype=dtype, **kwargs)
-        model.load_state_dict(
-            load_state_dict_hf(pretrained_model_name, device={"": device}, dtype=dtype)
-        )
-        return model

src/axolotl/models/mixtral/__init__.py

@@ -1,9 +0,0 @@
-"""
-Custom modeling code for mixtral
-"""
-
-from .configuration_moe_mistral import MixtralConfig  # noqa
-from .modeling_moe_mistral import (  # noqa
-    MixtralForCausalLM,
-    replace_mixtral_mlp_with_swiglu,
-)

src/axolotl/models/mixtral/configuration_moe_mistral.py

@@ -1,154 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Mistral AI and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Mistral model configuration"""
-
-from transformers.configuration_utils import PretrainedConfig
-from transformers.utils import logging
-
-logger = logging.get_logger(__name__)
-
-MISTRAL_PRETRAINED_CONFIG_ARCHIVE_MAP = {
-    "mistralai/Mistral-7B-v0.1": "https://huggingface.co/mistralai/Mistral-7B-v0.1/resolve/main/config.json",
-    "mistralai/Mistral-7B-Instruct-v0.1": "https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1/resolve/main/config.json",
-}
-
-
-class MixtralConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`MistralModel`]. It is used to instantiate an
-    Mistral model according to the specified arguments, defining the model architecture. Instantiating a configuration
-    with the defaults will yield a similar configuration to that of the Mistral-7B-v0.1 or Mistral-7B-Instruct-v0.1.
-
-    [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1)
-    [mistralai/Mistral-7B-Instruct-v0.1](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1)
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 32000):
-            Vocabulary size of the Mistral model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed when calling [`MistralModel`]
-        hidden_size (`int`, *optional*, defaults to 4096):
-            Dimension of the hidden representations.
-        intermediate_size (`int`, *optional*, defaults to 14336):
-            Dimension of the MLP representations.
-        num_hidden_layers (`int`, *optional*, defaults to 32):
-            Number of hidden layers in the Transformer encoder.
-        num_attention_heads (`int`, *optional*, defaults to 32):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        num_key_value_heads (`int`, *optional*, defaults to 8):
-            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
-            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
-            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
-            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
-            by meanpooling all the original heads within that group. For more details checkout [this
-            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `8`.
-        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
-            The non-linear activation function (function or string) in the decoder.
-        max_position_embeddings (`int`, *optional*, defaults to `4096*32`):
-            The maximum sequence length that this model might ever be used with. Mistral's sliding window attention
-            allows sequence of up to 4096*32 tokens.
-        initializer_range (`float`, *optional*, defaults to 0.02):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
-            The epsilon used by the rms normalization layers.
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether or not the model should return the last key/values attentions (not used by all models). Only
-            relevant if `config.is_decoder=True`.
-        pad_token_id (`int`, *optional*):
-            The id of the padding token.
-        bos_token_id (`int`, *optional*, defaults to 1):
-            The id of the "beginning-of-sequence" token.
-        eos_token_id (`int`, *optional*, defaults to 2):
-            The id of the "end-of-sequence" token.
-        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
-            Whether the model's input and output word embeddings should be tied.
-        rope_theta (`float`, *optional*, defaults to 10000.0):
-            The base period of the RoPE embeddings.
-        sliding_window (`int`, *optional*, defaults to 4096):
-            Sliding window attention window size. If not specified, will default to `4096`.
-        attention_dropout (`float`, *optional*, defaults to 0.0):
-            The dropout ratio for the attention probabilities.
-
-    ```python
-    >>> from transformers import MistralModel, MistralConfig
-
-    >>> # Initializing a Mistral 7B style configuration
-    >>> configuration = MixtralConfig()
-
-    >>> # Initializing a model from the Mistral 7B style configuration
-    >>> model = MixtralModel(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```"""
-
-    model_type = "mistral"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size=32000,
-        hidden_size=4096,
-        intermediate_size=14336,
-        num_hidden_layers=32,
-        num_attention_heads=32,
-        num_key_value_heads=8,
-        hidden_act="silu",
-        max_position_embeddings=4096 * 32,
-        initializer_range=0.02,
-        rms_norm_eps=1e-6,
-        use_cache=True,
-        pad_token_id=None,
-        bos_token_id=1,
-        eos_token_id=2,
-        tie_word_embeddings=False,
-        rope_theta=10000.0,
-        attention_dropout=0.0,
-        num_experts_per_token=2,
-        num_experts=8,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-
-        # for backward compatibility
-        if num_key_value_heads is None:
-            num_key_value_heads = num_attention_heads
-
-        self.num_key_value_heads = num_key_value_heads
-        self.hidden_act = hidden_act
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.attention_dropout = attention_dropout
-        self.num_experts = num_experts
-        self.num_experts_per_token = num_experts_per_token
-
-        # pylint: disable=duplicate-code
-        super().__init__(
-            pad_token_id=pad_token_id,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
-        )

src/axolotl/models/phi/__init__.py

@@ -3,6 +3,4 @@ MixFormers model architecture used for phi models
 """
 
 from .configuration_mixformer_sequential import MixFormerSequentialConfig  # noqa
-from .configuration_phi import PhiConfig  # noqa
 from .modeling_mixformer_sequential import MixFormerSequentialForCausalLM  # noqa
-from .modeling_phi import PhiForCausalLM  # noqa

src/axolotl/models/phi/configuration_phi.py

@@ -1,65 +0,0 @@
-# pylint: skip-file
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT license.
-
-import math
-from typing import Optional
-
-from transformers import PretrainedConfig
-
-
-class PhiConfig(PretrainedConfig):
-    """Phi configuration."""
-
-    model_type = "phi"
-    attribute_map = {
-        "max_position_embeddings": "n_positions",
-        "hidden_size": "n_embd",
-        "num_attention_heads": "n_head",
-        "num_hidden_layers": "n_layer",
-    }
-
-    def __init__(
-        self,
-        vocab_size: int = 50304,
-        n_positions: int = 2048,
-        n_embd: int = 1024,
-        n_layer: int = 20,
-        n_inner: Optional[int] = None,
-        n_head: int = 16,
-        n_head_kv: Optional[int] = None,
-        rotary_dim: Optional[int] = 32,
-        activation_function: Optional[str] = "gelu_new",
-        flash_attn: bool = False,
-        flash_rotary: bool = False,
-        fused_dense: bool = False,
-        attn_pdrop: float = 0.0,
-        embd_pdrop: float = 0.0,
-        resid_pdrop: float = 0.0,
-        layer_norm_epsilon: float = 1e-5,
-        initializer_range: float = 0.02,
-        tie_word_embeddings: bool = False,
-        pad_vocab_size_multiple: int = 64,
-        **kwargs
-    ) -> None:
-        self.vocab_size = int(
-            math.ceil(vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple
-        )
-        self.n_positions = n_positions
-        self.n_embd = n_embd
-        self.n_layer = n_layer
-        self.n_inner = n_inner
-        self.n_head = n_head
-        self.n_head_kv = n_head_kv
-        self.rotary_dim = min(rotary_dim, n_embd // n_head)
-        self.activation_function = activation_function
-        self.flash_attn = flash_attn
-        self.flash_rotary = flash_rotary
-        self.fused_dense = fused_dense
-        self.attn_pdrop = attn_pdrop
-        self.embd_pdrop = embd_pdrop
-        self.resid_pdrop = resid_pdrop
-        self.layer_norm_epsilon = layer_norm_epsilon
-        self.initializer_range = initializer_range
-
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)

src/axolotl/monkeypatch/llama_attn_hijack_flash.py

@@ -13,18 +13,12 @@ import transformers
 from einops import rearrange
 from flash_attn.bert_padding import pad_input, unpad_input
 from transformers.modeling_outputs import BaseModelOutputWithPast
-from transformers.models.llama.modeling_llama import LlamaAttention
 from transformers.models.llama.modeling_llama import (
     LlamaDecoderLayer as OriginalLlamaDecoderLayer,
 )
-from transformers.models.llama.modeling_llama import (
-    LlamaMLP,
-    apply_rotary_pos_emb,
-    repeat_kv,
-)
-from xformers.ops import SwiGLU
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv
 
-from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids, set_module_name
+from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
 
 try:
     from flash_attn.flash_attn_interface import (  # pylint: disable=ungrouped-imports
@@ -44,28 +38,6 @@ except ImportError:
 LOG = logging.getLogger("axolotl")
 
 
-def replace_llama_mlp_with_swiglu(model):
-    for name, module in model.named_modules():
-        if isinstance(module, LlamaMLP):
-            mlp = FusedMLP(
-                module.config, module.gate_proj, module.up_proj, module.down_proj
-            )
-            set_module_name(model, name, mlp)
-
-
-def replace_llama_qkv_with_fused(model):
-    for name, module in model.named_modules():
-        if isinstance(module, LlamaAttention):
-            qkv = FusedAttention(
-                module.config,
-                module.q_proj,
-                module.k_proj,
-                module.v_proj,
-                module.o_proj,
-            )
-            set_module_name(model, name, qkv)
-
-
 def replace_llama_attn_with_flash_attn(
     packed: Optional[bool] = False,
     cross_entropy: Optional[bool] = False,
@@ -114,92 +86,6 @@ def replace_llama_attn_with_flash_attn(
             )
 
 
-class FusedAttention(LlamaAttention):
-    """
-    Fused QKV Attention layer for incrementally improved training efficiency
-    """
-
-    def __init__(
-        self,
-        config,
-        q: torch.nn.Linear,  # pylint: disable=invalid-name
-        k: torch.nn.Linear,  # pylint: disable=invalid-name
-        v: torch.nn.Linear,  # pylint: disable=invalid-name
-        o: torch.nn.Linear,  # pylint: disable=invalid-name
-    ):
-        super().__init__(config)
-        self.config = config
-        self.init_device = next(iter(q.state_dict().values())).device
-
-        # define equivalent fused qkv projection
-        self.out_features: List[int] = [q.out_features, k.out_features, v.out_features]
-        self.qkv_proj = torch.nn.Linear(
-            q.in_features, sum(self.out_features), device=self.init_device, bias=False
-        )
-        self.o_proj = o
-
-        # overwrite initialized weights with pretrained weights
-        self.qkv_proj.weight.data = torch.cat(
-            (q.weight.data, k.weight.data, v.weight.data), dim=0
-        )
-
-    def _post_training(self, model, name):
-        q_proj, k_proj, v_proj = torch.split(
-            self.qkv_proj.weight.data, self.out_features, dim=0
-        )
-
-        new_attn = LlamaAttention(self.config)
-        new_attn.q_proj.weight.data = q_proj
-        new_attn.k_proj.weight.data = k_proj
-        new_attn.v_proj.weight.data = v_proj
-        new_attn.o_proj.weight.data = self.o_proj.weight.data
-
-        set_module_name(model, name, new_attn)
-
-
-class FusedMLP(torch.nn.Module):
-    """
-    Fused MLP layer for incrementally improved training efficiency
-    """
-
-    def __init__(
-        self,
-        config,
-        gate_proj: torch.nn.Linear,
-        up_proj: torch.nn.Linear,
-        down_proj: torch.nn.Linear,
-    ):
-        super().__init__()
-        self.config = config
-        self.swiglu = SwiGLU(
-            in_features=config.hidden_size,
-            hidden_features=config.intermediate_size,
-            bias=False,
-            _pack_weights=True,
-        )
-        # overwrite initialized weights with pretrained weights
-        self.swiglu.w12.weight.data = torch.cat(
-            (gate_proj.weight.data, up_proj.weight.data), dim=0
-        )
-        self.swiglu.w3.weight.data = down_proj.weight.data
-
-    def _post_training(self, model, name):
-        w1, w2 = torch.split(  # pylint: disable=invalid-name
-            self.swiglu.w12.weight.data, self.config.intermediate_size, dim=0
-        )
-
-        # Assign the split weights back to the original layers
-        new_mlp = LlamaMLP(self.config)
-        new_mlp.gate_proj.weight.data = w1
-        new_mlp.up_proj.weight.data = w2
-        new_mlp.down_proj.weight.data = self.swiglu.w3.weight.data
-
-        set_module_name(model, name, new_mlp)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:  # pylint: disable=invalid-name
-        return self.swiglu(x)
-
-
 # Disable the transformation of the attention mask in LlamaModel as the flash attention
 # requires the attention mask to be the same as the key_padding_mask
 def _prepare_decoder_attention_mask(
@@ -230,6 +116,8 @@ def flashattn_forward(
     attention_mask: [bsz, q_len]
     """
     # pylint: disable=duplicate-code
+    original_dtype = hidden_states.dtype
+
     bsz, q_len, _ = hidden_states.size()
 
     if not hasattr(self, "pretraining_tp"):
@@ -261,14 +149,16 @@ def flashattn_forward(
         value_states = torch.cat(value_states, dim=-1)
 
     else:
-        if isinstance(self, FusedAttention):
-            query_states, key_states, value_states = self.qkv_proj(hidden_states).split(
-                self.out_features, dim=-1
-            )
-        else:
-            query_states = self.q_proj(hidden_states)
-            key_states = self.k_proj(hidden_states)
-            value_states = self.v_proj(hidden_states)
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+    if query_states.dtype == torch.float32:
+        query_states = query_states.to(dtype=original_dtype)
+    if key_states.dtype == torch.float32:
+        key_states = key_states.to(dtype=original_dtype)
+    if value_states.dtype == torch.float32:
+        value_states = value_states.to(dtype=original_dtype)
 
     query_states = query_states.view(
         bsz, q_len, self.num_heads, self.head_dim
@@ -321,8 +211,6 @@ def flashattn_forward(
         # only on first autoregressive step q,k,v have same seqlen
         is_causal = key_states.shape == query_states.shape
 
-    dropout_rate = 0.0 if not self.training else getattr(self, "attention_dropout", 0.0)
-
     if cu_seqlens is not None and max_seqlen is not None and cu_seqlens.dim() == 1:
         # special handling using sample packing
         qkv = torch.stack(
@@ -332,12 +220,7 @@ def flashattn_forward(
         qkv = rearrange(qkv, "b s ... -> (b s) ...")
 
         output = flash_attn_varlen_qkvpacked_func(
-            qkv,
-            cu_seqlens,
-            max_seqlen,
-            dropout_p=dropout_rate,
-            softmax_scale=None,
-            causal=True,
+            qkv, cu_seqlens, max_seqlen, 0.0, softmax_scale=None, causal=True
         )
         output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
     elif query_states.shape == key_states.shape:
@@ -360,7 +243,7 @@ def flashattn_forward(
             qkv_unpad,
             cu_seqlens_q,
             max_seqlen_q,
-            dropout_p=dropout_rate,
+            0.0,
             softmax_scale=None,
             causal=is_causal,
         )
@@ -373,7 +256,6 @@ def flashattn_forward(
             output = flash_attn_kvpacked_func(
                 query_states,
                 torch.stack([key_states, value_states], 2),
-                dropout_p=dropout_rate,
                 causal=is_causal,
             )
         else:
@@ -406,7 +288,7 @@ def flashattn_forward(
                 cu_seqlens_k,
                 max_seqlen_q,
                 max_seqlen_k,
-                dropout_p=dropout_rate,
+                0.0,
                 softmax_scale=None,
                 causal=is_causal,
             )
@@ -436,6 +318,10 @@ def flashattn_forward(
     else:
         attn_output = self.o_proj(attn_output)
 
+    # handle conversion back for IA3
+    if attn_output.dtype == torch.float32:
+        attn_output = attn_output.to(dtype=original_dtype)
+
     return attn_output, None, past_key_value
 
 
@@ -629,6 +515,7 @@ def llama_model_forward(
     )
 
     hidden_states = inputs_embeds
+    original_dtype = hidden_states.dtype
 
     if self.gradient_checkpointing and self.training:
         if use_cache:
@@ -686,6 +573,10 @@ def llama_model_forward(
 
         hidden_states = layer_outputs[0]
 
+        # handle conversion back for IA3
+        if hidden_states.dtype == torch.float32:
+            hidden_states = hidden_states.to(dtype=original_dtype)
+
         if use_cache:
             next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
 

src/axolotl/monkeypatch/llama_attn_hijack_sdp.py

@@ -25,8 +25,6 @@ def sdp_attention_forward(
     past_key_value: Optional[Tuple[torch.Tensor]] = None,
     output_attentions: bool = False,
     use_cache: bool = False,
-    padding_mask: Optional[torch.LongTensor] = None,  # pylint: disable=unused-argument
-    **kwargs,  # pylint: disable=unused-argument
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
     # pylint: disable=duplicate-code
     bsz, q_len, _ = hidden_states.size()

src/axolotl/monkeypatch/llama_attn_hijack_xformers.py

@@ -29,8 +29,6 @@ def xformers_forward(
     past_key_value: Optional[Tuple[torch.Tensor]] = None,
     output_attentions: bool = False,
     use_cache: bool = False,
-    padding_mask: Optional[torch.LongTensor] = None,  # pylint: disable=unused-argument
-    **kwargs,  # pylint: disable=unused-argument
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
     # pylint: disable=duplicate-code
     bsz, q_len, _ = hidden_states.size()

src/axolotl/monkeypatch/llama_embeddings_hijack.py

@@ -0,0 +1,40 @@
+"""
+patch to add noisy embeddings per https://arxiv.org/abs/2310.05914
+"""
+
+import torch
+import transformers.models.llama.modeling_llama
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+def replace_llama_embeddings_with_uniform_distribution(noise_alpha=5):
+    # pylint: disable=duplicate-code
+    def noised_embed(orig_embed, noise_alpha, model):
+        def new_func(input_ids):
+            # during training, we add noise to the embedding
+            # during generation, we don't add noise to the embedding
+            if model.training:
+                embed_init = orig_embed(input_ids)
+                dims = torch.tensor(embed_init.size(1) * embed_init.size(2))
+                mag_norm = noise_alpha / torch.sqrt(dims)
+                return embed_init + torch.zeros_like(embed_init).uniform_(
+                    -mag_norm, mag_norm
+                )
+            return orig_embed(input_ids)
+
+        return new_func
+
+    def post_init(orig_post_init):
+        def new_func(self):
+            orig_post_init(self)
+            self.embed_tokens.forward = noised_embed(
+                self.embed_tokens.forward, noise_alpha, self
+            )
+
+        return new_func
+
+    transformers.models.llama.modeling_llama.LlamaModel.post_init = post_init(
+        transformers.models.llama.modeling_llama.LlamaModel.post_init
+    )

src/axolotl/monkeypatch/mistral_attn_hijack_flash.py

@@ -201,8 +201,6 @@ def flashattn_forward(
         # only on first autoregressive step q,k,v have same seqlen
         is_causal = key_states.shape == query_states.shape
 
-    dropout_rate = 0.0 if not self.training else getattr(self, "attention_dropout", 0.0)
-
     if cu_seqlens is not None and max_seqlen is not None and cu_seqlens.dim() == 1:
         # special handling using sample packing
         qkv = torch.stack(
@@ -215,7 +213,7 @@ def flashattn_forward(
             qkv,
             cu_seqlens,
             max_seqlen,
-            dropout_p=dropout_rate,
+            0.0,
             softmax_scale=None,
             causal=True,
             window_size=window_size,
@@ -241,7 +239,7 @@ def flashattn_forward(
             qkv_unpad,
             cu_seqlens_q,
             max_seqlen_q,
-            dropout_p=dropout_rate,
+            0.0,
             softmax_scale=None,
             causal=is_causal,
             window_size=window_size,
@@ -255,7 +253,6 @@ def flashattn_forward(
             output = flash_attn_kvpacked_func(
                 query_states,
                 torch.stack([key_states, value_states], 2),
-                dropout_p=dropout_rate,
                 causal=is_causal,
                 window_size=window_size,
             )
@@ -289,7 +286,7 @@ def flashattn_forward(
                 cu_seqlens_k,
                 max_seqlen_q,
                 max_seqlen_k,
-                dropout_p=dropout_rate,
+                0.0,
                 softmax_scale=None,
                 causal=is_causal,
                 window_size=window_size,

src/axolotl/monkeypatch/mistral_embeddings_hijack.py

@@ -0,0 +1,40 @@
+"""
+patch to add noisy embeddings per https://arxiv.org/abs/2310.05914
+"""
+
+import torch
+import transformers.models.mistral.modeling_mistral
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+def replace_mistral_embeddings_with_uniform_distribution(noise_alpha=5):
+    # pylint: disable=duplicate-code
+    def noised_embed(orig_embed, noise_alpha, model):
+        def new_func(input_ids):
+            # during training, we add noise to the embedding
+            # during generation, we don't add noise to the embedding
+            if model.training:
+                embed_init = orig_embed(input_ids)
+                dims = torch.tensor(embed_init.size(1) * embed_init.size(2))
+                mag_norm = noise_alpha / torch.sqrt(dims)
+                return embed_init + torch.zeros_like(embed_init).uniform_(
+                    -mag_norm, mag_norm
+                )
+            return orig_embed(input_ids)
+
+        return new_func
+
+    def post_init(orig_post_init):
+        def new_func(self):
+            orig_post_init(self)
+            self.embed_tokens.forward = noised_embed(
+                self.embed_tokens.forward, noise_alpha, self
+            )
+
+        return new_func
+
+    transformers.models.mistral.modeling_mistral.MistralModel.post_init = post_init(
+        transformers.models.mistral.modeling_mistral.MistralModel.post_init
+    )

src/axolotl/monkeypatch/neft_embeddings.py

@@ -1,65 +0,0 @@
-"""
-patches implemented through the trainer hooks to enable NEFT/noisy embeddings per https://arxiv.org/abs/2310.05914
-"""
-import torch
-from peft import PeftModel
-from transformers import PreTrainedModel
-
-
-def patch_neft(alpha, model):
-    embeddings = None
-    if isinstance(model, PreTrainedModel):
-        embeddings = model.get_input_embeddings()
-    if isinstance(model, PeftModel):
-        embeddings = model.base_model.get_input_embeddings()
-    if not embeddings:
-        raise ValueError(f"unhandled model class for neft: {model.__class__.__name__}")
-    embeddings.noisy_embedding_alpha = alpha
-    old_forward = embeddings.forward
-
-    # This hack seems to be needed to properly use a custom forward pass
-    # all credits to: https://discuss.pytorch.org/t/how-can-i-replace-the-forward-method-of-a-predefined-torchvision-model-with-my-customized-forward-function/54224/11
-    bound_method = neft_forward.__get__(  # pylint: disable=no-value-for-parameter
-        embeddings, embeddings.__class__
-    )
-    setattr(embeddings, "forward", bound_method)
-
-    embeddings._old_forward = old_forward  # pylint: disable=protected-access
-    return model
-
-
-def unpatch_neft(model):
-    embeddings = None
-    if isinstance(model, PreTrainedModel):
-        embeddings = model.get_input_embeddings()
-    if isinstance(model, PeftModel):
-        embeddings = model.base_model.get_input_embeddings()
-    if not embeddings:
-        raise ValueError(f"unhandled model class for neft: {model.__class__.__name__}")
-    if hasattr(embeddings, "_old_forward"):
-        embeddings.forward = embeddings._old_forward  # pylint: disable=protected-access
-        del embeddings._old_forward  # pylint: disable=protected-access
-        del embeddings.noisy_embedding_alpha
-
-
-def neft_forward(self, inputs: torch.Tensor):
-    embeddings = self._old_forward(inputs)  # pylint: disable=protected-access
-
-    if self.training:
-        dims = torch.tensor(embeddings.size(1) * embeddings.size(2))
-        mag_norm = self.noisy_embedding_alpha / torch.sqrt(dims)
-        embeddings = embeddings + torch.zeros_like(embeddings).uniform_(
-            -mag_norm, mag_norm
-        )
-
-    return embeddings
-
-
-def pretrain_hook(cfg, trainer):
-    if cfg.noisy_embedding_alpha:
-        trainer.model = patch_neft(cfg.noisy_embedding_alpha, trainer.model)
-
-
-def post_train_hook(cfg, trainer):
-    if cfg.noisy_embedding_alpha:
-        unpatch_neft(trainer.model)

src/axolotl/monkeypatch/utils.py

@@ -101,16 +101,3 @@ def get_cu_seqlens_from_pos_ids(position_ids):
         max_seq_lens.append(max_seq_len)
 
     return torch.stack(results).to(dtype=torch.int32), torch.stack(max_seq_lens)
-
-
-def set_module_name(model, name, value):
-    if "." in name:
-        parent_name = name.rsplit(".", 1)[0]
-        child_name = name[len(parent_name) + 1 :]
-        parent = model.get_submodule(parent_name)
-    else:
-        parent_name = ""
-        parent = model
-        child_name = name
-
-    setattr(parent, child_name, value)

src/axolotl/prompt_strategies/sharegpt.py

@@ -13,7 +13,7 @@ register_conv_template(
         system_message="You are a helpful assistant.",
         roles=["<|im_start|>user", "<|im_start|>assistant"],
         sep_style=SeparatorStyle.CHATML,
-        sep="<|im_end|>",
+        sep="<|im_end|>\n",
     )
 )
 
@@ -24,7 +24,7 @@ def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
     )
     field_human = ds_cfg["field_human"] if ds_cfg and "field_human" in ds_cfg else None
     field_model = ds_cfg["field_model"] if ds_cfg and "field_model" in ds_cfg else None
-    strategy = SimpleShareGPTPromptTokenizingStrategy(
+    return SimpleShareGPTPromptTokenizingStrategy(
         ShareGPTPrompterV2(
             conversation=conversation,
             role_key_model=field_model,
@@ -34,9 +34,6 @@ def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
         cfg.train_on_inputs,
         cfg.sequence_len,
     )
-    if ds_cfg and "strict" in ds_cfg:
-        strategy.strict = ds_cfg["strict"]
-    return strategy
 
 
 def load_role(tokenizer, cfg):
@@ -62,26 +59,8 @@ class SimpleShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
     basic sharegpt strategy to grab conversations from the sample row
     """
 
-    _strict = True
-
-    @property
-    def strict(self):
-        return self._strict
-
-    @strict.setter
-    def strict(self, strict):
-        self._strict = strict
-
     def get_conversation_thread(self, prompt):
-        conversations = prompt["conversations"]
-        if self.strict:
-            return conversations
-        # remap roles - allow for assistant turn
-        role_map = {"human": "human", "assistant": "gpt", "gpt": "gpt"}
-        turns = [
-            {"from": role_map[t["from"]], "value": t["value"]} for t in conversations
-        ]
-        return turns
+        return prompt["conversations"]
 
 
 class SimpleRoleShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):

src/axolotl/prompt_tokenizers.py

@@ -45,8 +45,6 @@ class PromptTokenizingStrategy(abc.ABC):
         self.prompter = prompter
         self.tokenizer: PreTrainedTokenizer = tokenizer
         self.train_on_inputs = train_on_inputs
-        # sequence_len and max_length can be different for CompletionPromptTokenizingStrategy.
-        # TODO: Document how they are different.
         self.sequence_len = sequence_len
         self.max_length = sequence_len
 
@@ -61,31 +59,34 @@ class PromptTokenizingStrategy(abc.ABC):
     def _tokenize(
         self, prompt: str, add_eos_token: bool = True, strip_bos_token: bool = False
     ) -> BatchEncoding:
-        empty = BatchEncoding(data={"input_ids": [], "attention_mask": []})
+        result: BatchEncoding
         if not prompt:
             LOG.warning("Empty text requested for tokenization.")
-            return empty
-
-        result = self.tokenizer(
-            prompt,
-            truncation=True,
-            max_length=self.max_length,
-            padding=False,
-            return_tensors=None,
-        )
+            result = BatchEncoding(data={"input_ids": [], "attention_mask": []})
+        else:
+            result = self.tokenizer(
+                prompt,
+                truncation=True,
+                max_length=self.max_length,
+                padding=False,
+                return_tensors=None,
+            )
         if len(result["input_ids"]) == 0:
             LOG.warning("Tokenizer result is empty. You may want to audit your dataset")
-            return empty
-
         if (
-            result["input_ids"][-1] != self.tokenizer.eos_token_id
+            len(result["input_ids"]) > 0
+            and result["input_ids"][-1] != self.tokenizer.eos_token_id
             and len(result["input_ids"]) < self.max_length
             and add_eos_token
         ):
             result["input_ids"].append(self.tokenizer.eos_token_id)
             result["attention_mask"].append(1)
 
-        if result["input_ids"][0] == self.tokenizer.bos_token_id and strip_bos_token:
+        if (
+            len(result["input_ids"]) > 0
+            and result["input_ids"][0] == self.tokenizer.bos_token_id
+            and strip_bos_token
+        ):
             result["input_ids"] = result["input_ids"][1:]
             result["attention_mask"] = result["attention_mask"][1:]
 
@@ -121,7 +122,7 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy):
         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"] = [IGNORE_INDEX] * user_prompt_len
+            tokenized_prompt["labels"] = [-100] * user_prompt_len
         tokenized_res_prompt = self._tokenize(
             response, strip_bos_token=True, add_eos_token=True
         )
@@ -245,7 +246,6 @@ class ReflectionPromptTokenizingStrategy(PromptTokenizingStrategy):
         raise NotImplementedError
 
     def tokenize_prompt(self, prompt):
-        # pylint: disable=duplicate-code
         (
             instruction,
             input,  # pylint: disable=redefined-builtin
@@ -270,7 +270,7 @@ class ReflectionPromptTokenizingStrategy(PromptTokenizingStrategy):
             user_prompt_len = len(tokenized_user_prompt["input_ids"])
             # TODO this could be sped up using numpy array slicing
             tokenized_full_prompt["labels"] = [
-                IGNORE_INDEX
+                -100
             ] * user_prompt_len + tokenized_full_prompt["labels"][user_prompt_len:]
 
         return tokenized_full_prompt
@@ -334,7 +334,6 @@ class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy):
         return prompt["conversations"]
 
     def tokenize_prompt(self, prompt):
-        # Initial values. We will append to these as we go through the conversation.
         result, current_len = tokenize_prompt_default()
         conversation: Conversation = (
             self.prompter._conversation.copy()  # pylint: disable=protected-access
@@ -356,67 +355,62 @@ class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy):
             for _, part in enumerate(
                 self.prompter.build_prompt(self.get_conversation_thread(prompt))
             ):
-                if not isinstance(part, tuple):
-                    LOG.warning(f"expected tuple, got {part}")
-                    continue
-
-                user, assistant = conversation.roles
-                role, content = part
-
-                # Uses "in" because role contains extra characters
-                if user in role:
-                    role = (
-                        role.replace(role_remap[0]["from"], role_remap[0]["to"])
-                        if role_remap
-                        else role
-                    )
-                    turn = role + content
-                    # this is still the user query, we should
-                    if not content.strip():
-                        LOG.warning(f"user turn has empty text: {prompt}")
-                    res = self._tokenize(
-                        turn,
-                        add_eos_token=False,
-                        strip_bos_token=True,
-                    )
-                    # everything from this is masked out from the labels
-                    labels = [IGNORE_TOKEN_ID] * len(res["input_ids"])
-                elif assistant in role:
-                    # TODO label assistant token/tokens w/ IGNORE_TOKEN_ID
-                    role = (
-                        role.replace(role_remap[1]["from"], role_remap[1]["to"])
-                        if role_remap
-                        else role
-                    )
-                    turn = role + content
-                    # this should be the assistant response, should end with an eos token
-                    if not content.strip():
-                        LOG.warning(f"assistant turn has empty text: {prompt}")
-                    res = self._tokenize(
-                        turn,
-                        add_eos_token=True,
-                        strip_bos_token=True,
-                    )
-                    role_res = self._tokenize(
-                        role.rstrip(),
-                        add_eos_token=False,
-                        strip_bos_token=True,
-                    )
-                    # not masked out from labels
-                    labels = copy.deepcopy(res["input_ids"])
-                    len_role = len(role_res["input_ids"])
-                    labels[:len_role] = [IGNORE_TOKEN_ID] * min(len_role, len(labels))
-                elif role == "":
-                    turn = content
-                    # this is only ever the first part, should include the bos token and the user query
-                    res = self._tokenize(
-                        turn, add_eos_token=False, strip_bos_token=False
-                    )
-                    # everything from this is masked out from the labels
-                    labels = [IGNORE_TOKEN_ID] * len(res["input_ids"])
-                else:
-                    LOG.warning(f"unhandled role: {role}")
-                    continue
+                if isinstance(part, tuple):
+                    if conversation.roles[0] in part[0]:
+                        role = (
+                            part[0].replace(role_remap[0]["from"], role_remap[0]["to"])
+                            if role_remap
+                            else part[0]
+                        )
+                        turn = role + part[1]
+                        # this is still the user query, we should
+                        if not part[1].strip():
+                            LOG.warning(f"user turn has empty text: {prompt}")
+                        res = self._tokenize(
+                            turn,
+                            add_eos_token=False,
+                            strip_bos_token=True,
+                        )
+                        # everything from this is masked out from the labels
+                        labels = [IGNORE_TOKEN_ID] * len(res["input_ids"])
+                    elif conversation.roles[1] in part[0]:
+                        # TODO label assistant token/tokens w/ IGNORE_TOKEN_ID
+                        role = (
+                            part[0].replace(role_remap[1]["from"], role_remap[1]["to"])
+                            if role_remap
+                            else part[0]
+                        )
+                        turn = role + part[1]
+                        # this should be the assistant response, should end with an eos token
+                        if not part[1].strip():
+                            LOG.warning(f"assistant turn has empty text: {prompt}")
+                        res = self._tokenize(
+                            turn,
+                            add_eos_token=True,
+                            strip_bos_token=True,
+                        )
+                        role_res = self._tokenize(
+                            role.rstrip(),
+                            add_eos_token=False,
+                            strip_bos_token=True,
+                        )
+                        # not masked out from labels
+                        labels = copy.deepcopy(res["input_ids"])
+                        len_role = len(role_res["input_ids"])
+                        labels[:len_role] = [IGNORE_TOKEN_ID] * min(
+                            len_role, len(labels)
+                        )
+                    elif part[0] == "":
+                        turn = part[1]
+                        # this is only ever the first part, should include the bos token and the user query
+                        res = self._tokenize(
+                            turn, add_eos_token=False, strip_bos_token=False
+                        )
+                        # everything from this is masked out from the labels
+                        labels = [IGNORE_TOKEN_ID] * len(res["input_ids"])
+                    else:
+                        LOG.warning(f"unhandled role: {part[0]}")
+                        continue
 
                 # pylint: disable=duplicate-code
                 result, current_len = parse_tokenized_to_result(
@@ -430,6 +424,38 @@ class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy):
         except (KeyError, AssertionError, IndexError) as err:
             raise InvalidDataException(str(err)) from err
 
+    def _tokenize(self, prompt, add_eos_token=True, strip_bos_token=False):
+        if not prompt.strip():
+            LOG.warning("Empty text requested for tokenization.")
+            result = BatchEncoding(data={"input_ids": [], "attention_mask": []})
+        else:
+            result = self.tokenizer(
+                prompt,
+                truncation=True,
+                max_length=self.sequence_len,
+                padding=False,
+                return_tensors=None,
+            )
+        if (
+            len(result["input_ids"]) > 0
+            and result["input_ids"][-1] != self.tokenizer.eos_token_id
+            and len(result["input_ids"]) < self.sequence_len
+            and add_eos_token
+        ):
+            result["input_ids"].append(self.tokenizer.eos_token_id)
+            result["attention_mask"].append(1)
+
+        if (
+            len(result["input_ids"]) > 0
+            and result["input_ids"][0] == self.tokenizer.bos_token_id
+            and strip_bos_token
+        ):
+            result["input_ids"] = result["input_ids"][1:]
+            result["attention_mask"] = result["attention_mask"][1:]
+
+        result["labels"] = result["input_ids"].copy()
+        return result
+
 
 def tokenize_prompt_default() -> Tuple[Dict[str, List[int]], int]:
     """

src/axolotl/prompters.py

@@ -4,12 +4,10 @@ import logging
 from enum import Enum
 from typing import Generator, Optional, Union
 
-from colorama import Fore
 from fastchat.conversation import Conversation, get_conv_template
 
 LOG = logging.getLogger("axolotl")
 IGNORE_TOKEN_ID = -100
-REPR_TEMPLATE = "\n<start>\n" + Fore.CYAN + "{full_prompt}" + Fore.RESET + "\n<end>\n"
 
 
 class PromptStyle(Enum):
@@ -22,13 +20,7 @@ class PromptStyle(Enum):
     CHATML = "chatml"
 
 
-class Prompter:
-    """
-    Base prompter class for all prompters
-    """
-
-
-class AlpacaPrompter(Prompter):
+class AlpacaPrompter:
     """
     Base class for alpaca prompters
     """
@@ -63,38 +55,29 @@ class AlpacaPrompter(Prompter):
             )
             self.system_format = "<|im_start|>system\n{system}<|im_end|>\n"
 
-    def _build_result(self, instruction, input_text, output):
-        # returns the full prompt from instruction and optional input
-        # if a label (=response, =output) is provided, it's also appended.
-        if input_text:
-            res = (
-                self.system_format.format(system=self.system_prompt)
-                if self.system_prompt
-                else ""
-            ) + self.turn_format.format(instruction=instruction, input=input_text)
-        else:
-            res = (
-                self.system_format.format(system=self.system_no_input_prompt)
-                if self.system_no_input_prompt
-                else ""
-            ) + self.turn_no_input_format.format(instruction=instruction)
-        if output:
-            res = f"{res}{output}"
-
-        return res
-
     def build_prompt(
         self,
         instruction: str,
         input: Union[None, str] = None,  # pylint: disable=redefined-builtin
         output: Union[None, str] = None,
     ) -> Generator[str, None, None]:
-        yield self._build_result(instruction, input, output)
-
-    def __repr__(self) -> str:
-        return REPR_TEMPLATE.format(
-            full_prompt=self._build_result("{instruction}", "{input}", "{output}")
-        )
+        # 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_format.format(system=self.system_prompt)
+                if self.system_prompt
+                else ""
+            ) + self.turn_format.format(instruction=instruction, input=input)
+        else:
+            res = (
+                self.system_format.format(system=self.system_no_input_prompt)
+                if self.system_prompt
+                else ""
+            ) + self.turn_no_input_format.format(instruction=instruction)
+        if output:
+            res = f"{res}{output}"
+        yield res
 
 
 class UnpromptedPrompter(AlpacaPrompter):
@@ -165,7 +148,7 @@ class NomicGPT4AllPrompter(AlpacaPrompter):
     """
 
 
-class ReflectAlpacaPrompter(Prompter):
+class ReflectAlpacaPrompter:
     """
     Prompter for ReflectAlpaca
     """
@@ -208,14 +191,14 @@ class ReflectAlpacaPrompter(Prompter):
             )
             self.response_split = "ASSISTANT:"
 
-    def _build_result(
+    def build_prompt(
         self,
         instruction: str,
         input: Union[None, str] = None,  # pylint: disable=redefined-builtin
         output: Union[None, str] = None,
         reflection: Union[None, str] = None,
         corrected: 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:
@@ -229,30 +212,7 @@ class ReflectAlpacaPrompter(Prompter):
                 corrected=corrected,
             )
             res = f"{res}{label}"
-
-        return res
-
-    def build_prompt(
-        self,
-        instruction: str,
-        input: Union[None, str] = None,  # pylint: disable=redefined-builtin
-        output: Union[None, str] = None,
-        reflection: Union[None, str] = None,
-        corrected: Union[None, str] = None,
-    ) -> Generator[str, None, None]:
-        # pylint: disable=duplicate-code
-        yield self._build_result(
-            instruction,
-            input,
-            output,
-            reflection,
-            corrected,
-        )
-
-    def __repr__(self) -> str:
-        return REPR_TEMPLATE.format(
-            full_prompt=self._build_result("{instruction}", "{input}", "{output}")
-        )
+        yield res
 
 
 SHAREGPT_ASSERTION_FAILED_ROLE = (
@@ -260,7 +220,7 @@ SHAREGPT_ASSERTION_FAILED_ROLE = (
 )
 
 
-class ShareGPTPrompter(Prompter):  # pylint: disable=too-few-public-methods
+class ShareGPTPrompter:  # pylint: disable=too-few-public-methods
     """
     A prompter that generates prompts for the ShareGPT
     """
@@ -287,7 +247,7 @@ class ShareGPTPrompter(Prompter):  # pylint: disable=too-few-public-methods
         if role_key_model:
             self.role_key_model = role_key_model
 
-    def _build_result(self, source):
+    def build_prompt(self, source) -> Generator[str, None, None]:
         if len(source) < 2:
             # If there isn't a back and forth conversation, ignore it
             # also happens on the data splitting leaving empty conversations
@@ -322,20 +282,11 @@ class ShareGPTPrompter(Prompter):  # pylint: disable=too-few-public-methods
                 LOG.warning(f"{SHAREGPT_ASSERTION_FAILED_ROLE}: {sentence}")
             conv.append_message(role, sentence["value"])
 
-        return conv.get_turns()
-
-    def build_prompt(self, source) -> Generator[str, None, None]:
-        turns = self._build_result(source)
-
-        for part in turns:
+        for part in conv.get_turns():
             if part[0] and not part[1]:
                 LOG.warning(f"role with empty message: {part[0]}")
             yield part
 
-    def __repr__(self) -> str:
-        turns = self._build_result([{"from": "{from}", "value": "{value}"}])
-        return "\n".join([REPR_TEMPLATE.format(full_prompt=part) for part in turns])
-
 
 class ShareGPTPrompterV2(ShareGPTPrompter):
     """
@@ -353,15 +304,3 @@ class ShareGPTPrompterV2(ShareGPTPrompter):
             role_key_human=role_key_human,
             role_key_model=role_key_model,
         )
-
-
-class UnsupportedPrompter(Prompter):
-    """
-    A dummy class for custom prompters
-    """
-
-    def __init__(self) -> None:
-        pass
-
-    def __repr__(self):
-        return "Pre-tokenized or custom dataset types are unsupported for logging"

src/axolotl/train.py

@@ -1,5 +1,6 @@
 """Prepare and train a model on a dataset. Can also infer from a model or merge lora"""
 
+import logging
 import os
 import signal
 import sys
@@ -9,14 +10,11 @@ from typing import Optional
 
 import torch
 import transformers.modelcard
-from accelerate.logging import get_logger
 from datasets import Dataset
 from optimum.bettertransformer import BetterTransformer
-from transformers.deepspeed import is_deepspeed_zero3_enabled
 
 from axolotl.common.cli import TrainerCliArgs
 from axolotl.logging_config import configure_logging
-from axolotl.monkeypatch import neft_embeddings
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.models import load_model, load_tokenizer
 from axolotl.utils.trainer import setup_trainer
@@ -26,7 +24,7 @@ src_dir = os.path.join(project_root, "src")
 sys.path.insert(0, src_dir)
 
 configure_logging()
-LOG = get_logger("axolotl.train")
+LOG = logging.getLogger("axolotl.train")
 
 
 @dataclass
@@ -41,13 +39,13 @@ class TrainDatasetMeta:
 
 
 def train(
-    *, cfg: DictDefault, cli_args: TrainerCliArgs, dataset_meta: TrainDatasetMeta
+    *,
+    cfg: DictDefault,
+    cli_args: TrainerCliArgs,
+    dataset_meta: TrainDatasetMeta,
 ):
     # load the tokenizer first
-    LOG.debug(
-        f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}",
-        main_process_only=True,
-    )
+    LOG.info(f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}")
     tokenizer = load_tokenizer(cfg)
 
     train_dataset = dataset_meta.train_dataset
@@ -55,10 +53,7 @@ def train(
     total_num_steps = dataset_meta.total_num_steps
 
     # Load the model and tokenizer
-    msg = "loading model"
-    if cfg.adapter:
-        msg += " and peft_config..."
-    LOG.debug(msg)
+    LOG.info("loading model and (optionally) peft_config...")
     model, peft_config = load_model(cfg, tokenizer, inference=cli_args.inference)
 
     safe_serialization = cfg.save_safetensors is True
@@ -82,8 +77,7 @@ def train(
         cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps
     )
 
-    if hasattr(model, "config"):
-        model.config.use_cache = False
+    model.config.use_cache = False
 
     # go ahead and presave, so we have the adapter config available to inspect
     if peft_config:
@@ -93,8 +87,7 @@ def train(
     if not Path(cfg.output_dir).is_dir():
         os.makedirs(cfg.output_dir, exist_ok=True)
     tokenizer.save_pretrained(str(Path(cfg.output_dir)))
-    if hasattr(model, "config"):
-        model.config.save_pretrained(str(Path(cfg.output_dir)))
+    model.config.save_pretrained(str(Path(cfg.output_dir)))
 
     # 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:
@@ -116,7 +109,6 @@ def train(
     if cfg.group_by_length:
         LOG.info("hang tight... sorting dataset for group_by_length")
 
-    pretrain_hooks(cfg, trainer)
     if cfg.flash_optimum:
         with torch.backends.cuda.sdp_kernel(
             enable_flash=True, enable_math=True, enable_mem_efficient=True
@@ -124,15 +116,9 @@ def train(
             trainer.train(resume_from_checkpoint=resume_from_checkpoint)
     else:
         trainer.train(resume_from_checkpoint=resume_from_checkpoint)
-    post_train_hooks(cfg, trainer)
 
     LOG.info(f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}")
 
-    # post training
-    for name, module in model.named_modules():
-        if hasattr(module, "_post_training"):
-            module._post_training(model, name)  # pylint: disable=protected-access
-
     if trainer.is_fsdp_enabled:
         trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
         LOG.info("Set FSDP state dict type to FULL_STATE_DICT for saving.")
@@ -148,22 +134,6 @@ def train(
     # only save on rank 0, otherwise it corrupts output on multi-GPU when multiple processes attempt to write the same file
     if cfg.fsdp:
         trainer.save_model(cfg.output_dir)
-    elif cfg.deepspeed and is_deepspeed_zero3_enabled():
-        # Copied over from: https://github.com/huggingface/accelerate/blob/5ae611118057232f441055f7ef9ba0b0f2b8d533/docs/source/usage_guides/deepspeed.md#saving-and-loading
-        trainer.accelerator.wait_for_everyone()
-        unwrapped_model = trainer.accelerator.unwrap_model(trainer.model_wrapped)
-
-        # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if
-        # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or
-        # `zero3_save_16bit_model` is True in DeepSpeed Plugin.
-        # For Zero Stages 1 and 2, models are saved as usual in the output directory.
-        # The model name saved is `pytorch_model.bin`
-        unwrapped_model.save_pretrained(
-            cfg.output_dir,
-            is_main_process=trainer.accelerator.is_main_process,
-            save_function=trainer.accelerator.save,
-            state_dict=trainer.accelerator.get_state_dict(trainer.model_wrapped),
-        )
     elif cfg.local_rank == 0:
         if cfg.flash_optimum:
             model = BetterTransformer.reverse(model)
@@ -174,23 +144,3 @@ def train(
         trainer.create_model_card(model_name=cfg.output_dir.lstrip("./"))
 
     return model, tokenizer
-
-
-def pretrain_hooks(cfg, trainer):
-    """
-    Run hooks right before kicking off the training
-    :param cfg:
-    :param trainer:
-    :return:
-    """
-    neft_embeddings.pretrain_hook(cfg, trainer)
-
-
-def post_train_hooks(cfg, trainer):
-    """
-    Run hooks right after training completes
-    :param cfg:
-    :param trainer:
-    :return:
-    """
-    neft_embeddings.post_train_hook(cfg, trainer)

src/axolotl/utils/callbacks.py

@@ -37,7 +37,7 @@ from axolotl.utils.distributed import (
 )
 
 if TYPE_CHECKING:
-    from axolotl.core.trainer_builder import AxolotlTrainingArguments
+    from axolotl.utils.trainer import AxolotlTrainingArguments
 
 LOG = logging.getLogger("axolotl.callbacks")
 IGNORE_INDEX = -100
@@ -124,36 +124,6 @@ class GPUStatsCallback(
         return control
 
 
-class LossWatchDogCallback(TrainerCallback):
-    """Callback to track loss and stop training if loss is too high"""
-
-    def __init__(self, cfg):
-        self.cfg = cfg
-        self.logged = False
-        self.violations = 0
-        self.threshold = cfg.loss_watchdog_threshold
-        self.patience = cfg.loss_watchdog_patience or 3
-
-    def on_step_end(
-        self,
-        _args: TrainingArguments,
-        state: TrainerState,
-        control: TrainerControl,
-        **_kwargs,
-    ):
-        if len(state.log_history) > 0 and "loss" in state.log_history[-1]:
-            if state.log_history[-1]["loss"] > self.threshold:
-                self.violations += 1
-                if self.violations >= self.patience:
-                    LOG.warning(
-                        "Loss is too high, stopping training (loss_watchdog_threshold)"
-                    )
-                    control.should_training_stop = True
-            else:
-                self.violations = 0
-        return control
-
-
 def bench_eval_callback_factory(trainer, tokenizer):
     accuracy = evaluate.load("accuracy")
     abcd_idx = [

src/axolotl/utils/collators.py

@@ -2,16 +2,12 @@
 DataCollator for axolotl to pad labels and position_ids for packed sequences
 """
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Sequence, Union
+from typing import Any, Optional, Union
 
 import numpy as np
-import torch
-import transformers
 from transformers import PreTrainedTokenizerBase
 from transformers.utils import PaddingStrategy
 
-IGNORE_INDEX = -100
-
 
 @dataclass
 class DataCollatorForSeq2Seq:
@@ -123,58 +119,3 @@ class DataCollatorForSeq2Seq:
             features["decoder_input_ids"] = decoder_input_ids
 
         return features
-
-
-@dataclass
-class BatchSamplerDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
-    """
-    Collator for multipack specific to the using the BatchSampler
-    """
-
-    def __call__(self, features, return_tensors=None):
-        chunked_data = {}
-        for feature in features[0].keys():
-            if feature == "length":
-                continue
-            if feature == "attention_mask":
-                arrays = [
-                    (1) * np.array(item[feature])
-                    for item in features
-                    if feature in item
-                ]
-                chunked_data[feature] = np.concatenate(arrays)
-            else:
-                arrays = [
-                    np.array(item[feature]) for item in features if feature in item
-                ]
-                chunked_data[feature] = np.concatenate(arrays)
-        features = [chunked_data]
-        return super().__call__(features, return_tensors=return_tensors)
-
-
-@dataclass
-class MambaDataCollator:
-    """
-    Collator for State Space Models (Mamba)
-    """
-
-    tokenizer: transformers.PreTrainedTokenizer
-
-    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
-        input_ids, labels = tuple(
-            [torch.LongTensor(instance[key]) for instance in instances]
-            for key in ("input_ids", "labels")
-        )
-        input_ids = torch.nn.utils.rnn.pad_sequence(
-            input_ids,
-            batch_first=True,
-            padding_value=self.tokenizer.pad_token_id,
-        )
-        labels = torch.nn.utils.rnn.pad_sequence(
-            labels, batch_first=True, padding_value=IGNORE_INDEX
-        )
-
-        return {
-            "input_ids": input_ids,
-            "labels": labels,
-        }

src/axolotl/utils/config.py

@@ -27,7 +27,7 @@ def choose_device(cfg):
 
     cfg.device = get_device()
     if cfg.world_size == 1:
-        cfg.device_map = cfg.device_map or "auto"
+        cfg.device_map = "auto"
     else:
         if cfg.device.startswith("cuda"):
             cfg.device_map = {"": torch.cuda.current_device()}
@@ -79,9 +79,6 @@ def normalize_config(cfg):
 
     cfg.dataset_processes = cfg.dataset_processes or os.cpu_count()
 
-    if not cfg.base_model_config:
-        cfg.base_model_config = cfg.base_model
-
     model_config = load_model_config(cfg)
     cfg.model_config_type = model_config.model_type
 
@@ -122,24 +119,20 @@ def normalize_config(cfg):
         or (cfg.model_type and "mistral" in cfg.model_type.lower())
     )
 
-    cfg.is_qwen_derived_model = (
-        (
-            hasattr(model_config, "model_type")
-            and model_config.model_type
-            in [
-                "qwen",
-            ]
-        )
-        or cfg.is_qwen_derived_model
-        or "qwen" in cfg.base_model.lower()
-        or (cfg.model_type and "qwen" in cfg.model_type.lower())
-    )
-
-    if isinstance(cfg.learning_rate, str):
-        cfg.learning_rate = float(cfg.learning_rate)
-
     log_gpu_memory_usage(LOG, "baseline", cfg.device)
 
+    if cfg.adapter is not None:
+        for key in list(cfg.keys()):
+            if key.startswith("lora_"):
+                new_key = key.replace("lora_", "peft_")
+                LOG.warning(
+                    PendingDeprecationWarning(
+                        f"{key} soon to be deprecated. please use {new_key}"
+                    )
+                )
+                cfg[new_key] = cfg[key]
+                del cfg[key]
+
 
 def validate_config(cfg):
     if is_torch_bf16_gpu_available():
@@ -178,11 +171,7 @@ def validate_config(cfg):
             "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.eval_batch_size
-        and cfg.micro_batch_size
-        and cfg.eval_batch_size != cfg.micro_batch_size
-    ):
+    if cfg.eval_batch_size != cfg.micro_batch_size:
         LOG.warning(
             "eval_batch_size != micro_batch_size. This can lead to VRAM instability."
         )
@@ -212,14 +201,11 @@ def validate_config(cfg):
             if not cfg.load_in_4bit:
                 raise ValueError("Require cfg.load_in_4bit to be True for qlora")
 
-        if cfg.flash_attn_fuse_qkv or cfg.flash_attn_fuse_mlp:
-            raise ValueError("Fused modules are not supported with QLoRA")
-
     if not cfg.load_in_8bit and cfg.adapter == "lora":
-        LOG.warning("We recommend setting `load_in_8bit: true` for LORA finetuning")
+        LOG.warning("We recommend setting `load_in_8bit: true` for LoRA finetuning")
 
-    if cfg.adapter == "lora" and (cfg.flash_attn_fuse_qkv or cfg.flash_attn_fuse_mlp):
-        raise ValueError("Fused modules are not supported with LoRA")
+    if not cfg.load_in_8bit and cfg.adapter == "ia3":
+        LOG.warning("We recommend setting `load_in_8bit: true` for IA3 finetuning")
 
     if cfg.relora_steps:
         if cfg.adapter not in ("lora", "qlora"):
@@ -234,9 +220,6 @@ def validate_config(cfg):
         if cfg.lr_scheduler == "one_cycle":
             raise ValueError("ReLoRA is not compatible with the one_cycle scheduler")
 
-        if cfg.flash_attn_fuse_qkv or cfg.flash_attn_fuse_mlp:
-            raise ValueError("Fused modules are not supported with ReLoRA")
-
     if cfg.trust_remote_code:
         LOG.warning(
             "`trust_remote_code` is set to true. Please make sure that you reviewed the remote code/model."
@@ -371,39 +354,6 @@ def validate_config(cfg):
             "eval_steps and evaluation_strategy are not supported with val_set_size == 0"
         )
 
-    if (
-        cfg.sample_packing
-        and cfg.eval_table_size
-        and cfg.eval_sample_packing is not False
-    ):
-        raise ValueError(
-            "eval_table_size and eval_sample_packing are not supported together with sample_packing. Please set 'eval_sample_packing' to false."
-        )
-
-    if not cfg.adapter and (cfg.load_in_8bit or cfg.load_in_4bit):
-        raise ValueError(
-            "load_in_8bit and load_in_4bit are not supported without setting an adapter."
-            "If you want to full finetune, please turn off load_in_8bit and load_in_4bit."
-        )
-
-    if cfg.rope_scaling:
-        LOG.warning("`rope_scaling` should now be be a key under `model_config`")
-
-    if cfg.warmup_steps and cfg.warmup_ratio:
-        raise ValueError("warmup_steps and warmup_ratio are mutually exclusive")
-
-    if cfg.is_qwen_derived_model and cfg.gradient_checkpointing:
-        LOG.warning(
-            "Gradient checkpointing is broken for Qwen models for transformers>=4.35.0, except main branch."
-        )
-
-    if cfg.wandb_run_id and not cfg.wandb_name:
-        cfg.wandb_name = cfg.wandb_run_id
-
-        LOG.warning(
-            "wandb_run_id sets the ID of the run. If you would like to set the name, please use wandb_name instead."
-        )
-
     # TODO
     # MPT 7b
     # https://github.com/facebookresearch/bitsandbytes/issues/25

src/axolotl/utils/data.py

@@ -34,10 +34,8 @@ from axolotl.prompters import (
     JeopardyPrompter,
     MultipleChoiceConcisePrompter,
     MultipleChoiceExplainPrompter,
-    Prompter,
     ReflectAlpacaPrompter,
     SummarizeTLDRPrompter,
-    UnsupportedPrompter,
 )
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.distributed import is_main_process, zero_first
@@ -57,10 +55,9 @@ def md5(to_hash: str, encoding: str = "utf-8") -> str:
 
 
 def prepare_dataset(cfg, tokenizer):
-    prompters = []
     if not cfg.pretraining_dataset:
         with zero_first(is_main_process()):
-            train_dataset, eval_dataset, prompters = load_prepare_datasets(
+            train_dataset, eval_dataset = load_prepare_datasets(
                 tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH
             )
     else:
@@ -73,33 +70,25 @@ def prepare_dataset(cfg, tokenizer):
         # 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")
         eval_dataset = None
-        return train_dataset, eval_dataset, cfg.max_steps, prompters
+        return train_dataset, eval_dataset, cfg.max_steps
 
     with zero_first(is_main_process()):
         train_dataset, eval_dataset = process_datasets_for_packing(
             cfg, train_dataset, eval_dataset, tokenizer
         )
-
-    if eval_dataset and cfg.sample_packing and cfg.eval_sample_packing is not False:
-        total_eval_steps = calculate_total_num_steps(cfg, eval_dataset, update=False)
-        if total_eval_steps == 0:
-            raise ValueError(
-                "eval dataset split is too small for sample_packing. You should set `eval_sample_packing: False`. "
-            )
-
     if cfg.max_steps:
         total_num_steps = min(
-            calculate_total_num_steps(cfg, train_dataset), cfg.max_steps
+            calculate_total_num_steps(cfg, train_dataset, tokenizer), cfg.max_steps
         )
         LOG.info(f"Maximum number of steps set at {total_num_steps}")
     else:
-        total_num_steps = calculate_total_num_steps(cfg, train_dataset)
-    return train_dataset, eval_dataset, total_num_steps, prompters
+        total_num_steps = calculate_total_num_steps(cfg, train_dataset, tokenizer)
+    return train_dataset, eval_dataset, total_num_steps
 
 
 def load_tokenized_prepared_datasets(
     tokenizer, cfg, default_dataset_prepared_path
-) -> Tuple[DatasetDict, List[Prompter]]:
+) -> DatasetDict:
     tokenizer_name = tokenizer.__class__.__name__
     ds_hash = str(
         md5(
@@ -107,12 +96,7 @@ def load_tokenized_prepared_datasets(
                 str(cfg.sequence_len)
                 + "@"
                 + "|".join(
-                    sorted(
-                        [
-                            f"{d.path}:{d.type}:{d.shards}:{d.conversation}"
-                            for d in cfg.datasets
-                        ]
-                    )
+                    sorted([f"{d.path}:{d.type}:{d.shards}" for d in cfg.datasets])
                 )
                 + "|"
                 + tokenizer_name
@@ -125,7 +109,6 @@ def load_tokenized_prepared_datasets(
         else Path(default_dataset_prepared_path) / ds_hash
     )
     dataset = None
-    prompters = []
     use_auth_token = cfg.hf_use_auth_token
     try:
         if cfg.push_dataset_to_hub:
@@ -164,13 +147,13 @@ def load_tokenized_prepared_datasets(
                     yield dataset
 
         # pylint: disable=invalid-name
-        for config_dataset in for_d_in_datasets(cfg.datasets):
+        for d in for_d_in_datasets(cfg.datasets):
             ds: Union[Dataset, DatasetDict] = None
             ds_from_hub = False
             try:
                 load_dataset(
-                    config_dataset.path,
-                    name=config_dataset.name,
+                    d.path,
+                    name=d.name,
                     streaming=True,
                     token=use_auth_token,
                 )
@@ -178,85 +161,34 @@ def load_tokenized_prepared_datasets(
             except (FileNotFoundError, ConnectionError):
                 pass
 
-            ds_from_cloud = False
-            storage_options = {}
-            remote_file_system = None
-            if config_dataset.path.startswith("s3://"):
-                try:
-                    import aiobotocore.session  # type: ignore
-                    import s3fs  # type: ignore
-                except ImportError as exc:
-                    raise ImportError(
-                        "s3:// paths require aiobotocore and s3fs to be installed"
-                    ) from exc
-
-                # Takes credentials from ~/.aws/credentials for default profile
-                s3_session = aiobotocore.session.AioSession(profile="default")
-                storage_options = {"session": s3_session}
-                remote_file_system = s3fs.S3FileSystem(**storage_options)
-            elif config_dataset.path.startswith(
-                "gs://"
-            ) or config_dataset.path.startswith("gcs://"):
-                try:
-                    import gcsfs  # type: ignore
-                except ImportError as exc:
-                    raise ImportError(
-                        "gs:// or gcs:// paths require gcsfs to be installed"
-                    ) from exc
-
-                # gcsfs will use default credentials from the environment else anon
-                # https://gcsfs.readthedocs.io/en/latest/#credentials
-                storage_options = {"token": None}
-                remote_file_system = gcsfs.GCSFileSystem(**storage_options)
-            # TODO: Figure out how to get auth creds passed
-            # elif config_dataset.path.startswith("adl://") or config_dataset.path.startswith("abfs://"):
-            #     try:
-            #         import adlfs
-            #     except ImportError as exc:
-            #        raise ImportError(
-            #            "adl:// or abfs:// paths require adlfs to be installed"
-            #        ) from exc
-
-            #     # Gen 1
-            #     storage_options = {
-            #         "tenant_id": TENANT_ID,
-            #         "client_id": CLIENT_ID,
-            #         "client_secret": CLIENT_SECRET,
-            #     }
-            #     # Gen 2
-            #     storage_options = {
-            #         "account_name": ACCOUNT_NAME,
-            #         "account_key": ACCOUNT_KEY,
-            #     }
-
-            #     remote_file_system = adlfs.AzureBlobFileSystem(**storage_options)
-            try:
-                if remote_file_system and remote_file_system.exists(
-                    config_dataset.path
-                ):
-                    ds_from_cloud = True
-            except (FileNotFoundError, ConnectionError):
-                pass
-
             # prefer local dataset, even if hub exists
-            local_path = Path(config_dataset.path)
+            local_path = Path(d.path)
             if local_path.exists():
                 if local_path.is_dir():
                     # TODO dirs with arrow or parquet files could be loaded with `load_from_disk`
                     ds = load_dataset(
-                        config_dataset.path,
-                        name=config_dataset.name,
-                        data_files=config_dataset.data_files,
+                        d.path,
+                        name=d.name,
+                        data_files=d.data_files,
                         streaming=False,
                         split=None,
                     )
                 elif local_path.is_file():
-                    ds_type = get_ds_type(config_dataset)
-
+                    ds_type = "json"
+                    if d.ds_type:
+                        ds_type = d.ds_type
+                    elif ".parquet" in d.path:
+                        ds_type = "parquet"
+                    elif ".arrow" in d.path:
+                        ds_type = "arrow"
+                    elif ".csv" in d.path:
+                        ds_type = "csv"
+                    elif ".txt" in d.path:
+                        ds_type = "text"
                     ds = load_dataset(
                         ds_type,
-                        name=config_dataset.name,
-                        data_files=config_dataset.path,
+                        name=d.name,
+                        data_files=d.path,
                         streaming=False,
                         split=None,
                     )
@@ -266,41 +198,25 @@ def load_tokenized_prepared_datasets(
                     )
             elif ds_from_hub:
                 ds = load_dataset(
-                    config_dataset.path,
-                    name=config_dataset.name,
+                    d.path,
+                    name=d.name,
                     streaming=False,
-                    data_files=config_dataset.data_files,
+                    data_files=d.data_files,
                     token=use_auth_token,
                 )
-            elif ds_from_cloud and remote_file_system:
-                if remote_file_system.isdir(config_dataset.path):
-                    ds = load_from_disk(
-                        config_dataset.path,
-                        storage_options=storage_options,
-                    )
-                elif remote_file_system.isfile(config_dataset.path):
-                    ds_type = get_ds_type(config_dataset)
-                    ds = load_dataset(
-                        ds_type,
-                        name=config_dataset.name,
-                        data_files=config_dataset.path,
-                        streaming=False,
-                        split=None,
-                        storage_options=storage_options,
-                    )
             else:
-                if isinstance(config_dataset.data_files, str):
+                if isinstance(d.data_files, str):
                     fp = hf_hub_download(
-                        repo_id=config_dataset.path,
+                        repo_id=d.path,
                         repo_type="dataset",
-                        filename=config_dataset.data_files,
+                        filename=d.data_files,
                     )
-                elif isinstance(config_dataset.data_files, list):
+                elif isinstance(d.data_files, list):
                     fp = []
-                    for file in config_dataset.data_files:
+                    for file in d.data_files:
                         fp.append(
                             hf_hub_download(
-                                repo_id=config_dataset.path,
+                                repo_id=d.path,
                                 repo_type="dataset",
                                 filename=file,
                             )
@@ -310,27 +226,21 @@ def load_tokenized_prepared_datasets(
                         "data_files must be either a string or list of strings"
                     )
                 ds = load_dataset(
-                    "json",
-                    name=config_dataset.name,
-                    data_files=fp,
-                    streaming=False,
-                    split=None,
+                    "json", name=d.name, data_files=fp, streaming=False, split=None
                 )
             if not ds:
                 raise ValueError("unhandled dataset load")
             # support for using a subset of the data
-            if config_dataset.shards:
+            if d.shards:
                 if "train" in ds:
                     ds = ds.shuffle(seed=seed)["train"].shard(
-                        num_shards=config_dataset.shards, index=0
+                        num_shards=d.shards, index=0
                     )
                 else:
-                    ds = ds.shuffle(seed=seed).shard(
-                        num_shards=config_dataset.shards, index=0
-                    )
+                    ds = ds.shuffle(seed=seed).shard(num_shards=d.shards, index=0)
 
             d_base_type = d_prompt_style = None
-            d_type = config_dataset.type
+            d_type = d.type
             if isinstance(d_type, str):
                 d_type_split = d_type.split(":")
                 d_base_type = d_type_split[0]
@@ -339,26 +249,108 @@ def load_tokenized_prepared_datasets(
                 ds = ds["train"]
             elif (
                 isinstance(ds, DatasetDict)
-                and config_dataset.train_on_split
-                and config_dataset.train_on_split in ds
+                and d.train_on_split
+                and d.train_on_split in ds
             ):
-                ds = ds[config_dataset.train_on_split]
+                ds = ds[d.train_on_split]
             elif isinstance(ds, DatasetDict):
                 raise ValueError(
-                    f"no train split found for dataset {config_dataset.path}, you may specify a split with 'train_on_split: `"
+                    f"no train split found for dataset {d.path}, you may specify a split with 'train_on_split: `"
+                )
+            if (
+                "input_ids" in ds.features
+                and "attention_mask" in ds.features
+                and "labels" in ds.features
+            ):
+                # dataset is already tokenized, just drop it straight in
+                datasets.append(ds)
+            elif isinstance(d.type, DictDefault):
+                ds_strategy = load("user_defined", tokenizer, cfg, d.type.to_dict())
+                ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
+                datasets.append(ds_wrapper)
+            elif ds_strategy := load(d.type, tokenizer, cfg, d):
+                ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
+                datasets.append(ds_wrapper)
+            elif d_base_type == "alpaca":
+                ds_strategy = AlpacaPromptTokenizingStrategy(
+                    AlpacaPrompter(d_prompt_style),
+                    tokenizer,
+                    cfg.train_on_inputs,
+                    cfg.sequence_len,
+                )
+                ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
+                datasets.append(ds_wrapper)
+            elif d_base_type == "explainchoice":
+                ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy(
+                    MultipleChoiceExplainPrompter(d_prompt_style),
+                    tokenizer,
+                    cfg.train_on_inputs,
+                    cfg.sequence_len,
+                )
+                ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
+                datasets.append(ds_wrapper)
+            elif d_base_type == "concisechoice":
+                ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy(
+                    MultipleChoiceConcisePrompter(d_prompt_style),
+                    tokenizer,
+                    cfg.train_on_inputs,
+                    cfg.sequence_len,
+                )
+                ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
+                datasets.append(ds_wrapper)
+            elif d_base_type == "summarizetldr":
+                ds_strategy = SummarizeTLDRPromptTokenizingStrategy(
+                    SummarizeTLDRPrompter(d_prompt_style),
+                    tokenizer,
+                    cfg.train_on_inputs,
+                    cfg.sequence_len,
+                )
+                ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
+                datasets.append(ds_wrapper)
+            elif d_base_type == "jeopardy":
+                ds_strategy = JeopardyPromptTokenizingStrategy(
+                    JeopardyPrompter(d_prompt_style),
+                    tokenizer,
+                    cfg.train_on_inputs,
+                    cfg.sequence_len,
+                )
+                ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
+                datasets.append(ds_wrapper)
+            elif d_base_type == "oasst":
+                ds_strategy = OpenAssistantPromptTokenizingStrategy(
+                    AlpacaPrompter(d_prompt_style),
+                    tokenizer,
+                    cfg.train_on_inputs,
+                    cfg.sequence_len,
+                )
+                ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
+                datasets.append(ds_wrapper)
+            elif d_base_type == "gpteacher":
+                ds_strategy = GPTeacherPromptTokenizingStrategy(
+                    GPTeacherPrompter(d_prompt_style),
+                    tokenizer,
+                    cfg.train_on_inputs,
+                    cfg.sequence_len,
+                )
+                ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
+                datasets.append(ds_wrapper)
+            elif d_base_type == "reflection":
+                ds_strategy = AlpacaReflectionPTStrategy(
+                    ReflectAlpacaPrompter(d_prompt_style),
+                    tokenizer,
+                    cfg.train_on_inputs,
+                    cfg.sequence_len,
+                )
+                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_')}?"
+                LOG.error(f"unhandled prompt tokenization strategy: {d.type}. {suffix}")
+                raise ValueError(
+                    f"unhandled prompt tokenization strategy: {d.type} {suffix}"
                 )
-
-            dataset_wrapper, dataset_prompter = get_dataset_wrapper(
-                config_dataset=config_dataset,
-                dataset=ds,
-                tokenizer=tokenizer,
-                cfg=cfg,
-                d_base_type=d_base_type,
-                d_prompt_style=d_prompt_style,
-            )
-            datasets.append(dataset_wrapper)
-            prompters.append(dataset_prompter)
-
         LOG.info("merging datasets")
         dataset = concatenate_datasets(datasets)
 
@@ -376,32 +368,14 @@ def load_tokenized_prepared_datasets(
                     f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True
                 )
 
-    return dataset, prompters
-
-
-def get_ds_type(config_dataset: DictDefault):
-    """
-    Get the dataset type from the path if it's not specified
-    """
-    ds_type = "json"
-    if config_dataset.ds_type:
-        ds_type = config_dataset.ds_type
-    elif ".parquet" in config_dataset.path:
-        ds_type = "parquet"
-    elif ".arrow" in config_dataset.path:
-        ds_type = "arrow"
-    elif ".csv" in config_dataset.path:
-        ds_type = "csv"
-    elif ".txt" in config_dataset.path:
-        ds_type = "text"
-    return ds_type
+    return dataset
 
 
 def load_prepare_datasets(
     tokenizer: PreTrainedTokenizerBase,
     cfg,
     default_dataset_prepared_path,
-) -> Tuple[Dataset, Dataset, List[Prompter]]:
+) -> Tuple[Dataset, Dataset]:
     max_packed_sequence_len = (
         cfg.max_packed_sequence_len if cfg.max_packed_sequence_len else cfg.sequence_len
     )
@@ -410,7 +384,6 @@ def load_prepare_datasets(
     )  # make sure we don't accidentally set it larger than sequence_len
 
     tokenizer_name = tokenizer.__class__.__name__
-    prompters: List[Prompter] = []
     if cfg.max_packed_sequence_len is not None:
         # see if we can go ahead and load the stacked dataset
         seed = f"@{str(cfg.seed)}" if cfg.seed else ""
@@ -466,7 +439,7 @@ def load_prepare_datasets(
                     f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True
                 )
         else:
-            dataset, prompters = load_tokenized_prepared_datasets(
+            dataset = load_tokenized_prepared_datasets(
                 tokenizer, cfg, default_dataset_prepared_path
             )
 
@@ -508,7 +481,7 @@ def load_prepare_datasets(
                         private=True,
                     )
     else:
-        dataset, prompters = load_tokenized_prepared_datasets(
+        dataset = load_tokenized_prepared_datasets(
             tokenizer, cfg, default_dataset_prepared_path
         )
 
@@ -544,13 +517,14 @@ def load_prepare_datasets(
         train_fingerprint = md5(to_hash_train)
         test_fingerprint = md5(to_hash_test)
 
-        dataset = dataset.train_test_split(
-            test_size=cfg.val_set_size,
-            shuffle=False,
-            seed=cfg.seed or 42,
-            train_new_fingerprint=train_fingerprint,
-            test_new_fingerprint=test_fingerprint,
-        )
+        with zero_first(is_main_process()):
+            dataset = dataset.train_test_split(
+                test_size=cfg.val_set_size,
+                shuffle=False,
+                seed=cfg.seed or 42,
+                train_new_fingerprint=train_fingerprint,
+                test_new_fingerprint=test_fingerprint,
+            )
 
         train_dataset = dataset["train"]
         eval_dataset = dataset["test"]
@@ -558,144 +532,7 @@ def load_prepare_datasets(
         train_dataset = dataset
         eval_dataset = None
 
-    return train_dataset, eval_dataset, prompters
-
-
-def get_dataset_wrapper(
-    config_dataset, dataset, tokenizer, cfg, d_base_type, d_prompt_style
-):
-    dataset_wrapper = None
-    dataset_prompter = None
-
-    if (
-        "input_ids" in dataset.features
-        and "attention_mask" in dataset.features
-        and "labels" in dataset.features
-    ):
-        # dataset is already tokenized, just drop it straight in
-        dataset_prompter = UnsupportedPrompter()
-        dataset_wrapper = dataset
-    elif isinstance(config_dataset.type, DictDefault):
-        ds_strategy = load(
-            "user_defined", tokenizer, cfg, config_dataset.type.to_dict()
-        )
-        dataset_prompter = UnsupportedPrompter()
-        dataset_wrapper = TokenizedPromptDataset(
-            ds_strategy, dataset, process_count=cfg.dataset_processes
-        )
-    elif ds_strategy := load(config_dataset.type, tokenizer, cfg, config_dataset):
-        dataset_prompter = UnsupportedPrompter()
-        dataset_wrapper = TokenizedPromptDataset(
-            ds_strategy, dataset, process_count=cfg.dataset_processes
-        )
-    elif d_base_type == "alpaca":
-        dataset_prompter = AlpacaPrompter(d_prompt_style)
-        ds_strategy = AlpacaPromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy, dataset, process_count=cfg.dataset_processes
-        )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "explainchoice":
-        dataset_prompter = MultipleChoiceExplainPrompter(d_prompt_style)
-        ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy, dataset, process_count=cfg.dataset_processes
-        )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "concisechoice":
-        dataset_prompter = MultipleChoiceConcisePrompter(d_prompt_style)
-        ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy, dataset, process_count=cfg.dataset_processes
-        )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "summarizetldr":
-        dataset_prompter = SummarizeTLDRPrompter(d_prompt_style)
-        ds_strategy = SummarizeTLDRPromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy, dataset, process_count=cfg.dataset_processes
-        )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "jeopardy":
-        dataset_prompter = JeopardyPrompter(d_prompt_style)
-        ds_strategy = JeopardyPromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy, dataset, process_count=cfg.dataset_processes
-        )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "oasst":
-        dataset_prompter = AlpacaPrompter(d_prompt_style)
-        ds_strategy = OpenAssistantPromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy, dataset, process_count=cfg.dataset_processes
-        )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "gpteacher":
-        dataset_prompter = GPTeacherPrompter(d_prompt_style)
-        ds_strategy = GPTeacherPromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy, dataset, process_count=cfg.dataset_processes
-        )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "reflection":
-        dataset_prompter = ReflectAlpacaPrompter(d_prompt_style)
-        ds_strategy = AlpacaReflectionPTStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy, dataset, process_count=cfg.dataset_processes
-        )
-        dataset_wrapper = ds_wrapper
-    else:
-        suffix = ""
-        if ":load_" in config_dataset.type:
-            suffix = f" Did you mean {config_dataset.type.replace(':load_', '.load_')}?"
-        LOG.error(
-            f"unhandled prompt tokenization strategy: {config_dataset.type}. {suffix}"
-        )
-        raise ValueError(
-            f"unhandled prompt tokenization strategy: {config_dataset.type} {suffix}"
-        )
-
-    return dataset_wrapper, dataset_prompter
+    return train_dataset, eval_dataset
 
 
 def encode_pretraining(

src/axolotl/utils/dataloader.py

@@ -0,0 +1,302 @@
+# pylint: skip-file
+import hashlib
+import itertools
+import logging
+import math
+from typing import Any, Callable, List, Union
+
+import numba
+import numpy as np
+from torch.utils.data import DistributedSampler, Sampler
+
+LOG = logging.getLogger("axolotl.utils.dataloader")
+
+
+@numba.njit
+def ffd_check(a: np.ndarray, c: int, n: int):
+    # First-fit-decreasing bin packing
+    # Check if a[] could fit in n bins with capacity c
+    # https://en.wikipedia.org/wiki/First-fit-decreasing_bin_packing
+
+    a = np.sort(a)[::-1]
+    bins = np.full((n,), c, dtype=a.dtype)
+    for size in a:
+        not_found = True
+        for idx in range(n):
+            if bins[idx] >= size:
+                bins[idx] -= size
+                not_found = False
+                break
+
+        if not_found:
+            return False
+
+    return True
+
+
+@numba.njit
+def ffd_with_result(a: np.ndarray, c: int, start_index: int):
+    # First-fit-decreasing bin packing (with result return)
+
+    indices = np.argsort(a)[::-1]
+    a = a[indices]
+
+    bins: List[Any] = []
+    bins_result: List[Any] = []
+    for a_id, size in enumerate(a):
+        add_new = True
+        for idx in range(len(bins)):
+            if bins[idx] >= size:
+                bins[idx] -= size
+                bins_result[idx].append(indices[a_id] + start_index)
+                add_new = False
+                break
+
+        if add_new:
+            bins.append(c - size)
+            bins_result.append([indices[a_id] + start_index])
+
+    return bins_result, len(a)
+
+
+@numba.njit
+def allocate(
+    lengths: np.ndarray, lengths_cumsum: np.ndarray, rank: int, c: int, n: int
+):
+    """
+    :param lengths: array of lengths of each sample
+    :param lengths_cumsum: cumulative sum of consecutive lengths
+    :param rank: rank for this process
+    :param c: length of tokens per batch
+    :param n: number of ranks
+    :return:
+    """
+    # Dynamic batch allocator, similar to Multifit
+    # https://en.wikipedia.org/wiki/Multifit_algorithm
+    # ~99.5% efficiency on OpenChat training set (12 * 2048 ctx len)
+
+    s = 0
+    start_index = 0
+    result = []
+    result_totseqs = []
+
+    while True:
+        # binary search [left, right)
+        left = 1
+        right = 1 + np.searchsorted(lengths_cumsum[start_index:], s + c * n, "right")
+
+        while right - left > 1:
+            mid = (left + right) // 2
+            if ffd_check(lengths[start_index : start_index + mid], c, n):
+                left = mid
+            else:
+                right = mid
+
+        # use length left
+        batch, tot_seqs = ffd_with_result(
+            lengths[start_index : start_index + left], c, start_index
+        )
+        if len(batch) < n:
+            break
+
+        start_index += left
+        s = lengths_cumsum[start_index - 1]
+
+        # add local rank
+        result.append(batch[rank])
+        # add total seqs for all ranks
+        result_totseqs.append(tot_seqs)
+        # yield batch[rank], tot_seqs, s, len(result) * c * n
+    return result, result_totseqs, s, len(result) * c * n
+
+
+def chunk(iterable, n):
+    """
+    Chunk data into tuples of length n
+    """
+    # batched('ABCDEFG', 3) --> ABC DEF G
+    if n < 1:
+        raise ValueError("n must be at least one")
+    it = iter(iterable)
+    while batch := tuple(itertools.islice(it, n)):
+        yield batch
+
+
+def hash_indices(lst: List[int]) -> str:
+    # Convert the list of integers to a string representation
+    concatenated = ",".join(map(str, lst))
+
+    # Generate the hash
+    sha256 = hashlib.sha256()
+    sha256.update(concatenated.encode())
+
+    return sha256.hexdigest()
+
+
+class MultipackDistributedDataloader:
+    """Unpadded data loading using Multipack.
+    Adapted from https://github.com/imoneoi/openchat/blob/v3_fix_mle_loss/ochat/training_deepspeed/multipack_dataloader.py
+    Approximate (at most ~1.22x) the optimal solution of the identical-machines scheduling problem, which is NP-hard.
+    """
+
+    def __init__(
+        self,
+        dataset: Any,
+        collate_fn: Callable,
+        seq_max_length: int = 2048,
+        batch_size: int = 1,
+        sampler: Union[Sampler, DistributedSampler] = None,
+        packing_efficiency_estimate: float = 1.0,
+        sample_packing_seq_len_multiplier: int = 1,
+        device_count: int = 1,
+    ):
+        # Dataset
+        self.dataset = dataset
+        self.lengths = (
+            dataset.data.column("position_ids")
+            .to_pandas()
+            .apply(lambda x: x[-1] + 1)
+            .values
+        )
+        assert isinstance(self.lengths, np.ndarray)
+        assert batch_size % sample_packing_seq_len_multiplier == 0
+        assert batch_size >= sample_packing_seq_len_multiplier
+        self.sampler = sampler
+        self.batch_size = batch_size
+        self.sample_packing_seq_len_multiplier = sample_packing_seq_len_multiplier
+        self.seq_max_length = seq_max_length
+        self.batch_max_length = batch_size * seq_max_length
+        self.collate_fn = collate_fn
+
+        self.num_replicas = 1
+        self.rank = 0
+
+        # statistics
+        self.eff_total_used = 0
+        self.eff_total_slots = 0
+        self.packing_efficiency_estimate = packing_efficiency_estimate or 1.0
+        self.device_count = device_count
+
+    def generate_batches(self, set_stats=False):
+        LOG.info("generating packed batches")
+        if self.sampler:
+            indices = [idx for idx in self.sampler]
+        else:
+            indices = range(0, len(self.dataset))
+
+        LOG.info(hash_indices(indices))
+        lengths = self.lengths[indices]
+        lengths_cumsum = np.cumsum(lengths)
+
+        batches, totseqs, total_used, total_slots = allocate(
+            lengths=lengths,
+            lengths_cumsum=lengths_cumsum,
+            rank=self.rank,
+            # c=self.batch_max_length,
+            c=self.seq_max_length * self.sample_packing_seq_len_multiplier,
+            n=self.num_replicas,
+        )
+
+        batches = [[indices[b_idx] for b_idx in batch] for batch in batches]
+
+        # statistics
+        if set_stats:
+            self.eff_total_used += total_used
+            self.eff_total_slots += total_slots
+
+        return batches, totseqs
+
+    def __iter__(self):
+        if hasattr(self.sampler, "set_epoch"):
+            new_epoch = self.sampler.epoch + 1
+            self.sampler.set_epoch(new_epoch)
+            LOG.info(f"calling sampler.set_epoch({new_epoch})")
+        all_batches, _ = self.generate_batches(set_stats=True)
+        features = self.dataset.features.keys()
+        len_remaining = self._len_est()
+        for batches in chunk(
+            all_batches, self.batch_size // self.sample_packing_seq_len_multiplier
+        ):
+            chunked_data = []
+            attn_mask_cum_idx = 0
+            for batch in batches:
+                concatenated = {}
+                batched_data = [self.dataset[batch_idx] for batch_idx in batch]
+                for feature in features:
+                    if feature == "length":
+                        continue
+                    if feature == "attention_mask":
+                        arrays = [
+                            (attn_mask_cum_idx + idx + 1) * np.array(item[feature])
+                            for idx, item in enumerate(batched_data)
+                            if feature in item
+                        ]
+                        attn_mask_cum_idx += len(batched_data)
+                        concatenated[feature] = np.concatenate(arrays)
+                    else:
+                        arrays = [
+                            np.array(item[feature])
+                            for item in batched_data
+                            if feature in item
+                        ]
+                        concatenated[feature] = np.concatenate(arrays)
+                chunked_data.append(concatenated)
+            yield self.collate_fn(chunked_data)
+            len_remaining -= 1
+            if not len_remaining:
+                return
+        # yield a no-op for cases where we don't have any data left to pack
+        for i in range(0, len_remaining):
+            yield self.collate_fn(
+                [
+                    {
+                        "input_ids": [0],
+                        "labels": [-100],
+                        "attention_mask": [True],
+                        "position_ids": [0],
+                    }
+                ]
+            )
+
+    def _len_est(self):
+        lengths_sum = np.sum(self.lengths)
+        lengths_sum_per_device = lengths_sum // self.device_count
+        LOG.info(
+            f"packing_efficiency_estimate: {self.packing_efficiency_estimate} "
+            f"total_num_tokens per device: {lengths_sum_per_device}"
+        )
+
+        # shave off 1% + 1 for dealing with variance in packing from random sampler to sampler
+        return (
+            math.floor(
+                0.99
+                * lengths_sum_per_device
+                / self.packing_efficiency_estimate
+                // self.seq_max_length
+                // self.batch_size
+            )
+            - 1
+        )
+
+    def __len__(self):
+        # this doesn't return the actual length b/c with distributed samplers, not all dataloaders get
+        # the same share of total tokens
+        # if not self.eff_total_used:
+        #     batches, _ = self.generate_batches(set_stats=True)
+        # LOG.info(
+        #     f"packing_efficiency_estimate: {self.packing_efficiency_estimate} "
+        #     f"actual packing efficiency: {self.efficiency()}"
+        # )
+        return max(1, self._len_est())
+
+    def len_w_stats(self):
+        if not self.eff_total_used:
+            batches, _ = self.generate_batches(set_stats=True)
+        LOG.info(
+            f"packing_efficiency_estimate: {self.packing_efficiency_estimate} "
+            f"actual packing efficiency: {self.efficiency()}"
+        )
+        return max(1, self._len_est())
+
+    def efficiency(self):
+        return self.eff_total_used / self.eff_total_slots

src/axolotl/utils/distributed.py

@@ -50,17 +50,6 @@ def get_world_size():
     return int(os.getenv("WORLD_SIZE", "1"))
 
 
-@contextmanager
-def zero_only():
-    """
-    Context manager that only runs the enclosed block on the main rank.
-    """
-    if is_main_process():
-        yield
-    else:
-        yield None
-
-
 @contextmanager
 def zero_first(is_main):
     """

src/axolotl/utils/models.py

@@ -4,7 +4,6 @@ import math
 import os
 from typing import Optional, Tuple  # noqa: F401
 
-import addict
 import bitsandbytes as bnb
 import torch
 import transformers
@@ -18,11 +17,11 @@ from transformers import (  # noqa: F401
     AutoTokenizer,
     BitsAndBytesConfig,
     GPTQConfig,
+    LlamaConfig,
     PreTrainedModel,
     PreTrainedTokenizerBase,
 )
 
-from axolotl.models.mamba import fix_mamba_attn_for_loss
 from axolotl.prompt_tokenizers import LLAMA_DEFAULT_EOS_TOKEN
 from axolotl.utils.bench import log_gpu_memory_usage
 from axolotl.utils.dict import DictDefault
@@ -30,57 +29,12 @@ from axolotl.utils.dict import DictDefault
 LOG = logging.getLogger("axolotl")
 
 
-def check_model_config(cfg: DictDefault, model_config: AutoConfig):
-    quant_config_exists = hasattr(model_config, "quantization_config")
-    quant_config_method_is_gptq = (
-        quant_config_exists
-        and "quant_method" in model_config.quantization_config
-        and model_config.quantization_config["quant_method"] == "gptq"
-    )
-
-    if cfg.gptq and not quant_config_method_is_gptq:
-        raise ValueError(
-            "model_config.quantization_config is not set or quant_method is not set to gptq. "
-            "Please make sure to point to a GPTQ model."
-        )
-
-    if not cfg.gptq and quant_config_exists:
-        raise ValueError(
-            "model_config.quantization_config is set but `gptq` flag is not. "
-            "Please use the `gptq` flag to train quantized model or point to a non-quantized model."
-        )
-
-
 def load_model_config(cfg):
     model_config_name = cfg.base_model_config or cfg.base_model
-    trust_remote_code = cfg.trust_remote_code is True
-    model_type = cfg.model_type
-
-    if model_type == "MixtralForCausalLM":
-        from axolotl.models.mixtral.configuration_moe_mistral import MixtralConfig
-
-        model_config = MixtralConfig.from_pretrained(model_config_name)
-    else:
-        try:
-            model_config = AutoConfig.from_pretrained(
-                model_config_name, trust_remote_code=trust_remote_code
-            )
-        except ValueError as err:
-            if "mamba" in model_config_name:
-                return addict.Dict(
-                    {
-                        "model_type": "mamba",
-                    }
-                )
-            raise err
-
-    if cfg.model_config:
-        for key, val in cfg.model_config.items():
-            setattr(model_config, key, val)
-
-    check_model_config(cfg, model_config)
-
-    return model_config
+    trust_remote_code: bool = False or cfg.trust_remote_code
+    return AutoConfig.from_pretrained(
+        model_config_name, trust_remote_code=trust_remote_code
+    )
 
 
 def load_tokenizer(cfg):
@@ -97,7 +51,7 @@ def load_tokenizer(cfg):
     if cfg.tokenizer_type:
         tokenizer_cls = getattr(transformers, cfg.tokenizer_type)
 
-    tokenizer_config = cfg.tokenizer_config or cfg.base_model_config or cfg.base_model
+    tokenizer_config = cfg.tokenizer_config or cfg.base_model_config
     tokenizer = tokenizer_cls.from_pretrained(
         tokenizer_config,
         trust_remote_code=cfg.trust_remote_code or False,
@@ -111,7 +65,6 @@ def load_tokenizer(cfg):
             "LlamaTokenizer",
             "LlamaTokenizerFast",
             "CodeLlamaTokenizer",
-            "CodeLlamaTokenizerFast",
         ]
         and hasattr(tokenizer, "pad_token")
         and not tokenizer.pad_token
@@ -119,6 +72,11 @@ def load_tokenizer(cfg):
         # set a pad_token, but use eos_token so we don't add a new token
         tokenizer.pad_token = LLAMA_DEFAULT_EOS_TOKEN
 
+    LOG.debug(f"EOS: {tokenizer.eos_token_id} / {tokenizer.eos_token}")
+    LOG.debug(f"BOS: {tokenizer.bos_token_id} / {tokenizer.bos_token}")
+    LOG.debug(f"PAD: {tokenizer.pad_token_id} / {tokenizer.pad_token}")
+    LOG.debug(f"UNK: {tokenizer.unk_token_id} / {tokenizer.unk_token}")
+
     if tokenizer.__class__.__name__ == "GPTNeoXTokenizerFast":
         tokenizer.add_special_tokens({"pad_token": "[PAD]"})
         os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -127,40 +85,11 @@ def load_tokenizer(cfg):
     if cfg.is_mistral_derived_model and cfg.flash_attention and not cfg.sample_packing:
         tokenizer.padding_side = "left"
 
-    # Qwen base only has single token, so we need to set the special tokens
-    if cfg.is_qwen_derived_model:
-        token_ids = ["bos_token_id", "eos_token_id", "pad_token_id", "unk_token_id"]
-        for attr_name in token_ids:
-            if getattr(tokenizer, attr_name) is None:
-                setattr(tokenizer, attr_name, tokenizer.eod_id)
-
-        token_names = ["bos_token", "eos_token", "pad_token", "unk_token"]
-        for attr_name in token_names:
-            if getattr(tokenizer, attr_name) is None:
-                setattr(tokenizer, attr_name, "<|endoftext|>")
-
     if cfg.special_tokens:
         for k, val in cfg.special_tokens.items():
             tokenizer.add_special_tokens(
                 {k: AddedToken(val, rstrip=False, lstrip=False, normalized=False)}
             )
-
-        # If we add bos_token and eos_token, we need to update the post processor to
-        # handle them correctly.
-        # https://github.com/huggingface/transformers/pull/24132
-        bos_or_eos_in_special_tokens = (
-            "bos_token" in cfg.special_tokens and "eos_token" in cfg.special_tokens
-        )
-        if (
-            tokenizer.__class__.__name__
-            in (
-                "LlamaTokenizerFast",
-                "CodeLlamaTokenizerFast",
-            )
-            and bos_or_eos_in_special_tokens
-        ):
-            tokenizer.update_post_processor()
-
     if cfg.tokens:
         tokenizer.add_tokens(
             [
@@ -169,11 +98,6 @@ def load_tokenizer(cfg):
             ]
         )
 
-    LOG.debug(f"EOS: {tokenizer.eos_token_id} / {tokenizer.eos_token}")
-    LOG.debug(f"BOS: {tokenizer.bos_token_id} / {tokenizer.bos_token}")
-    LOG.debug(f"PAD: {tokenizer.pad_token_id} / {tokenizer.pad_token}")
-    LOG.debug(f"UNK: {tokenizer.unk_token_id} / {tokenizer.unk_token}")
-
     return tokenizer
 
 
@@ -186,6 +110,7 @@ def load_model(
     Load a model for a given configuration and tokenizer.
     """
     base_model = cfg.base_model
+    base_model_config = cfg.base_model_config
     model_type = cfg.model_type
     model_config = load_model_config(cfg)
 
@@ -255,6 +180,26 @@ def load_model(
         LOG.info("patching with flash attention")
         replace_mistral_attn_with_flash_attn(packed=cfg.sample_packing)
 
+    if cfg.is_llama_derived_model and cfg.noisy_embedding_alpha:
+        from axolotl.monkeypatch.llama_embeddings_hijack import (
+            replace_llama_embeddings_with_uniform_distribution,
+        )
+
+        LOG.info("patching with noisy embeddings")
+        replace_llama_embeddings_with_uniform_distribution(
+            noise_alpha=cfg.noisy_embedding_alpha
+        )
+
+    if cfg.is_mistral_derived_model and cfg.noisy_embedding_alpha:
+        from axolotl.monkeypatch.mistral_embeddings_hijack import (
+            replace_mistral_embeddings_with_uniform_distribution,
+        )
+
+        LOG.info("patching with noisy embeddings")
+        replace_mistral_embeddings_with_uniform_distribution(
+            noise_alpha=cfg.noisy_embedding_alpha
+        )
+
     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,
@@ -276,7 +221,6 @@ def load_model(
     model_kwargs = {}
 
     model_kwargs["device_map"] = cfg.device_map
-    model_kwargs["max_memory"] = cfg.max_memory
     model_kwargs["torch_dtype"] = cfg.torch_dtype
 
     if cfg.model_revision:
@@ -308,35 +252,26 @@ def load_model(
             or cfg.is_falcon_derived_model
             or cfg.is_mistral_derived_model
         ):
-            # TODO enable once properly supported in transformers
-            # model_kwargs["attn_implementation"] = "flash_attention_2"
-            model_kwargs["use_flash_attention_2"] = True  # legacy, to be deprecated
+            model_kwargs["use_flash_attention_2"] = True
 
     try:
         if cfg.is_llama_derived_model and not cfg.trust_remote_code and not cfg.gptq:
             from transformers import LlamaForCausalLM
 
+            config_kwargs = {}
+            if cfg.rope_scaling:
+                config_kwargs["rope_scaling"] = cfg.rope_scaling
+            config = LlamaConfig.from_pretrained(
+                base_model_config,
+                **config_kwargs,
+            )
             model = LlamaForCausalLM.from_pretrained(
                 base_model,
-                config=model_config,
+                config=config,
                 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,
                 **model_kwargs,
             )
-
-            if cfg.flash_attention and not inference:
-                from axolotl.monkeypatch.llama_attn_hijack_flash import (
-                    replace_llama_mlp_with_swiglu,
-                    replace_llama_qkv_with_fused,
-                )
-
-                if cfg.flash_attn_fuse_mlp:
-                    LOG.info("patching with SwiGLU")
-                    replace_llama_mlp_with_swiglu(model)
-
-                if cfg.flash_attn_fuse_qkv:
-                    LOG.info("patching with fused QKV")
-                    replace_llama_qkv_with_fused(model)
         # elif model_type == "GPTNeoXForCausalLM" and cfg.flash_attention:
         #     This is a WIP, still an issue with the backward pass
         #     RuntimeError: grad can be implicitly created only for scalar outputs
@@ -363,116 +298,92 @@ def load_model(
         #         device=cfg.device,
         #     )
         #     model.train() # sets to train instead of eval mode
-        elif model_type == "PhiForCausalLM":
-            from axolotl.models.phi import PhiForCausalLM
+        elif model_type == "MixFormerSequentialForCausalLM":
+            from axolotl.models.phi import MixFormerSequentialForCausalLM
 
-            model = PhiForCausalLM.from_pretrained(
+            model = MixFormerSequentialForCausalLM.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,
                 **model_kwargs,
             )
-        elif model_type == "MixtralForCausalLM":
-            from axolotl.models.mixtral import (
-                MixtralForCausalLM,
-                replace_mixtral_mlp_with_swiglu,
-            )
-
-            model = MixtralForCausalLM.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,
-                **model_kwargs,
-            )
-
-            if cfg.flash_attn_fuse_mlp:
-                LOG.info("Mixtral MoE: Replacing experts with SwiGLU")
-                replace_mixtral_mlp_with_swiglu(model)
-
-        elif model_type == "MambaLMHeadModel":
-            # FIXME this is janky at best and hacked together to make it work
-            MambaLMHeadModel = fix_mamba_attn_for_loss()  # pylint: disable=invalid-name
-
-            model_kwargs["dtype"] = model_kwargs["torch_dtype"]
-            model_kwargs["device"] = torch.cuda.current_device()
-            del model_kwargs["torch_dtype"]
-            del model_kwargs["device_map"]
-            del model_kwargs["max_memory"]
-
-            model = MambaLMHeadModel.from_pretrained(
-                base_model,
-                **model_kwargs,
-            )
         elif model_type and not cfg.trust_remote_code:
             if cfg.gptq:
                 model = AutoModelForCausalLM.from_pretrained(
                     base_model,
-                    config=model_config,
                     trust_remote_code=cfg.trust_remote_code or False,
                     **model_kwargs,
                 )
             else:
                 model = getattr(transformers, model_type).from_pretrained(
                     base_model,
-                    config=model_config,
                     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,
                     trust_remote_code=cfg.trust_remote_code or False,
                     **model_kwargs,
                 )
         else:
+            config = AutoConfig.from_pretrained(
+                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(model_config, "max_seq_len")
-                and model_config.max_seq_len
-                and cfg.sequence_len > model_config.max_seq_len
+                hasattr(config, "max_seq_len")
+                and config.max_seq_len
+                and cfg.sequence_len > config.max_seq_len
             ):
-                model_config.max_seq_len = cfg.sequence_len
+                config.max_seq_len = cfg.sequence_len
                 LOG.warning(f"increasing context length to {cfg.sequence_len}")
             elif (
-                hasattr(model_config, "max_sequence_length")
-                and model_config.max_sequence_length
-                and cfg.sequence_len > model_config.max_sequence_length
+                hasattr(config, "max_sequence_length")
+                and config.max_sequence_length
+                and cfg.sequence_len > config.max_sequence_length
             ):
-                model_config.max_sequence_length = cfg.sequence_len
+                config.max_sequence_length = cfg.sequence_len
                 LOG.warning(f"increasing context length to {cfg.sequence_len}")
             if cfg.gptq:
                 model = AutoModelForCausalLM.from_pretrained(
                     base_model,
-                    config=model_config,
+                    config=config,
                     trust_remote_code=cfg.trust_remote_code or False,
                     **model_kwargs,
                 )
             else:
                 model = AutoModelForCausalLM.from_pretrained(
                     base_model,
-                    config=model_config,
+                    config=config,
                     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,
                     trust_remote_code=cfg.trust_remote_code or False,
                     **model_kwargs,
                 )
     except Exception as err:  # pylint: disable=broad-exception-caught
+        LOG.error(
+            "Exception raised attempting to load model, retrying with AutoModelForCausalLM"
+        )
         LOG.exception(err)
-        raise err
+        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,
+            trust_remote_code=cfg.trust_remote_code or False,
+            **model_kwargs,
+        )
 
     embeddings_len = (
         math.ceil(len(tokenizer) / 32) * 32
         if cfg.resize_token_embeddings_to_32x
         else len(tokenizer)
     )
-    if (
-        hasattr(model, "get_input_embeddings")
-        and model.get_input_embeddings().num_embeddings < embeddings_len
-    ):
+    if model.get_input_embeddings().num_embeddings < embeddings_len:
         model.resize_token_embeddings(embeddings_len)
     else:
         model.tie_weights()
 
     if (
-        hasattr(model, "config")
-        and hasattr(model.config, "max_position_embeddings")
+        hasattr(model.config, "max_position_embeddings")
         and model.config.max_position_embeddings
         and cfg.sequence_len > model.config.max_position_embeddings
     ):
@@ -481,23 +392,7 @@ def load_model(
         )
         model.config.max_position_embeddings = cfg.sequence_len
 
-    if (
-        hasattr(model, "config")
-        and hasattr(model.config, "bos_token_id")
-        and model.config.bos_token_id
-        and model.config.bos_token_id != tokenizer.bos_token_id
-    ):
-        model.config.bos_token_id = tokenizer.bos_token_id
-
-    if (
-        hasattr(model, "config")
-        and hasattr(model.config, "eos_token_id")
-        and model.config.eos_token_id
-        and model.config.eos_token_id != tokenizer.eos_token_id
-    ):
-        model.config.eos_token_id = tokenizer.eos_token_id
-
-    if hasattr(model, "device") and model.device.type == "cuda":
+    if model.device.type == "cuda":
         log_gpu_memory_usage(LOG, "after model load", model.device)
 
     # make sure these are fp32 per Ramesh et al. (2021)
@@ -511,28 +406,21 @@ def load_model(
             if hasattr(module, "weight"):
                 module.to(torch.float32)
 
-    needs_fa2_dtype = cfg.adapter or cfg.fsdp
-    skip_prepare_model_for_kbit_training = False
+    require_peft: bool = False
+    if cfg.adapter in ["lora", "qlora", "ia3"]:
+        require_peft = True
 
-    if cfg.model_config_type == "qwen" and cfg.adapter == "lora":
-        # Qwen doesn't play nicely with LoRA if this is enabled
-        skip_prepare_model_for_kbit_training = True
-
-    if (cfg.adapter == "lora" and load_in_8bit) or (
-        cfg.adapter == "qlora" and cfg.load_in_4bit
-    ):
+    if require_peft:
         LOG.info("converting PEFT model w/ prepare_model_for_kbit_training")
         if cfg.gradient_checkpointing:
             model.gradient_checkpointing_enable()
-        if not skip_prepare_model_for_kbit_training:
-            model = prepare_model_for_kbit_training(
-                model, use_gradient_checkpointing=cfg.gradient_checkpointing
-            )
-        needs_fa2_dtype = True
+        model = prepare_model_for_kbit_training(
+            model, use_gradient_checkpointing=cfg.gradient_checkpointing
+        )
 
     # LlamaRMSNorm layers are in fp32 after kbit_training or full finetune, so we need to
     # convert them back to fp16/bf16 for flash-attn compatibility.
-    if needs_fa2_dtype or (cfg.flash_attention and cfg.is_llama_derived_model):
+    if require_peft or cfg.fsdp or (cfg.flash_attention and cfg.is_llama_derived_model):
         LOG.info("converting modules to %s for flash attention", cfg.torch_dtype)
         for name, module in model.named_modules():
             if "norm" in name:
@@ -541,12 +429,19 @@ def load_model(
                 if hasattr(module, "weight"):
                     module.to(cfg.torch_dtype)
 
-    model, lora_config = load_adapter(model, cfg, cfg.adapter)
+    model, peft_config = load_adapter(model, cfg, cfg.adapter)
 
     if cfg.ddp and not load_in_8bit:
         model.to(f"cuda:{cfg.local_rank}")
 
-    if torch.cuda.device_count() > 1 and int(os.getenv("WORLD_SIZE", "1")) == 1:
+    if (
+        torch.cuda.device_count() > 1
+        and int(os.getenv("WORLD_SIZE", "1")) > 1
+        and (cfg.load_in_4bit)
+    ):
+        # llama is PROBABLY model parallelizable, but the default isn't that it is
+        # so let's only set it for the 4bit, see
+        # https://github.com/johnsmith0031/alpaca_lora_4bit/blob/08b3fca4a4a9e0d3945be1bab4529f100a428636/finetune.py#L130-L133
         setattr(model, "is_parallelizable", True)
         setattr(model, "model_parallel", True)
 
@@ -556,8 +451,7 @@ def load_model(
             requires_grad.append(f"{name}: {param.requires_grad}")
     if len(requires_grad) == 0:
         LOG.warning("there are no parameters that require gradient updates")
-    if hasattr(model, "config"):
-        model.config.use_cache = False
+    model.config.use_cache = False
 
     if cfg.flash_optimum:
         model = BetterTransformer.transform(model)
@@ -566,7 +460,7 @@ def load_model(
         log_gpu_memory_usage(LOG, "after adapters", model.device)
 
     # TODO resume_from_checkpoint handling
-    return model, lora_config
+    return model, peft_config
 
 
 def load_adapter(model, cfg, adapter, inference=False):
@@ -576,6 +470,8 @@ def load_adapter(model, cfg, adapter, inference=False):
         return model, None
     if hasattr(model, "enable_input_require_grads"):
         model.enable_input_require_grads()
+    if adapter == "ia3":
+        return load_ia3(model, cfg, inference=inference)
     if adapter in ["lora", "qlora"]:
         return load_lora(model, cfg, inference=inference)
     if adapter == "llama-adapter":
@@ -594,11 +490,11 @@ def load_llama_adapter(model, cfg):
         task_type="CAUSAL_LM",
     )
 
-    if cfg.lora_model_dir:
+    if cfg.peft_model_dir:
         LOG.debug("Loading pretained PEFT - llama_adapter")
         model = PeftModel.from_pretrained(
             model,
-            cfg.lora_model_dir,
+            cfg.peft_model_dir,
             torch_dtype=torch.float16,
         )
     else:
@@ -611,7 +507,7 @@ def load_llama_adapter(model, cfg):
 
 def find_all_linear_names(model):
     cls = (bnb.nn.Linear4bit, bnb.nn.Linear8bitLt, torch.nn.Linear, QuantLinear)
-    lora_module_names = set()
+    peft_module_names = set()
     for name, module in model.named_modules():
         if (
             isinstance(module, cls)
@@ -619,12 +515,12 @@ def find_all_linear_names(model):
             and module.__class__.__name__ not in ("LlamaLinearScalingRotaryEmbedding",)
         ):
             names = name.split(".")
-            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+            peft_module_names.add(names[0] if len(names) == 1 else names[-1])
 
-    if "lm_head" in lora_module_names:  # needed for 16-bit
-        lora_module_names.remove("lm_head")
+    if "lm_head" in peft_module_names:  # needed for 16-bit
+        peft_module_names.remove("lm_head")
 
-    return list(lora_module_names)
+    return list(peft_module_names)
 
 
 def load_lora(model, cfg, inference=False):
@@ -632,34 +528,68 @@ def load_lora(model, cfg, inference=False):
 
     from peft import LoraConfig, PeftModel, get_peft_model
 
-    lora_target_modules = list(cfg.lora_target_modules or [])
+    peft_target_modules = list(cfg.peft_target_modules or [])
 
-    if cfg.lora_target_linear:
+    if cfg.peft_target_linear:
         linear_names = find_all_linear_names(model)
         LOG.info(f"found linear modules: {repr(linear_names)}")
-        lora_target_modules = list(set(lora_target_modules + linear_names))
+        peft_target_modules = list(set(peft_target_modules + linear_names))
 
-    lora_config = LoraConfig(
-        r=cfg.lora_r,
-        lora_alpha=cfg.lora_alpha,
-        target_modules=lora_target_modules,
-        lora_dropout=cfg.lora_dropout,
-        fan_in_fan_out=cfg.lora_fan_in_fan_out,
-        modules_to_save=cfg.lora_modules_to_save if cfg.lora_modules_to_save else None,
+    peft_config = LoraConfig(
+        r=cfg.peft_r,
+        lora_alpha=cfg.peft_alpha,
+        target_modules=peft_target_modules,
+        lora_dropout=cfg.peft_dropout,
+        fan_in_fan_out=cfg.peft_fan_in_fan_out,
+        modules_to_save=cfg.peft_modules_to_save if cfg.peft_modules_to_save else None,
         bias="none",
         task_type="CAUSAL_LM",
     )
 
-    if cfg.lora_model_dir:
+    if cfg.peft_model_dir:
         LOG.debug("Loading pretained PEFT - LoRA")
         model = PeftModel.from_pretrained(
             model,
-            cfg.lora_model_dir,
+            cfg.peft_model_dir,
             is_trainable=(not inference),
         )
     else:
-        model = get_peft_model(model, lora_config)
+        model = get_peft_model(model, peft_config)
 
     model.print_trainable_parameters()
 
-    return model, lora_config
+    return model, peft_config
+
+
+def load_ia3(model, cfg, inference=False):
+    # type: (PreTrainedModel, DictDefault, bool) -> Tuple[PreTrainedModel, Optional[PeftConfig]]
+
+    from peft import IA3Config, PeftModel, get_peft_model
+
+    peft_config_kwargs = {}
+    if cfg.peft_init_ia3_weights is not None:
+        peft_config_kwargs["init_ia3_weights"] = cfg.peft_init_ia3_weights
+    if cfg.peft_fan_in_fan_out is not None:
+        peft_config_kwargs["fan_in_fan_out"] = cfg.peft_fan_in_fan_out
+
+    peft_config = IA3Config(
+        target_modules=cfg.peft_target_modules,
+        feedforward_modules=cfg.peft_feedforward_modules,
+        modules_to_save=cfg.peft_modules_to_save,
+        task_type="CAUSAL_LM",
+        **peft_config_kwargs,
+    )
+
+    if cfg.peft_model_dir:
+        LOG.debug("Loading pretained PEFT - IA3")
+        model = PeftModel.from_pretrained(
+            model,
+            cfg.peft_model_dir,
+            is_trainable=(not inference),
+        )
+    else:
+        model = get_peft_model(model, peft_config)
+
+    model.print_trainable_parameters()
+
+    return model, peft_config

src/axolotl/utils/samplers/__init__.py

@@ -1,4 +0,0 @@
-"""
-axolotl samplers module
-"""
-from .multipack import MultipackBatchSampler  # noqa: F401

src/axolotl/utils/samplers/multipack.py

@@ -1,196 +0,0 @@
-# pylint: skip-file
-"""
-Multipack Batch Sampler
-"""
-import logging
-import math
-import os
-from typing import Any, Iterable, List, Union
-
-import numba
-import numpy as np
-from torch.utils.data import BatchSampler, Sampler
-
-LOG = logging.getLogger("axolotl.utils.samplers.multipack")
-
-
-@numba.njit
-def ffd_check(a: np.ndarray, c: int, n: int):
-    # First-fit-decreasing bin packing
-    # Check if a[] could fit in n bins with capacity c
-    # https://en.wikipedia.org/wiki/First-fit-decreasing_bin_packing
-
-    a = np.sort(a)[::-1]
-    bins = np.full((n,), c, dtype=a.dtype)
-    for size in a:
-        not_found = True
-        for idx in range(n):
-            if bins[idx] >= size:
-                bins[idx] -= size
-                not_found = False
-                break
-
-        if not_found:
-            return False
-
-    return True
-
-
-@numba.njit
-def ffd_with_result(a: np.ndarray, c: int, start_index: int):
-    # First-fit-decreasing bin packing (with result return)
-
-    indices = np.argsort(a)[::-1]
-    a = a[indices]
-
-    bins: List[Any] = []
-    bins_result: List[Any] = []
-    for a_id, size in enumerate(a):
-        add_new = True
-        for idx in range(len(bins)):
-            if bins[idx] >= size:
-                bins[idx] -= size
-                bins_result[idx].append(indices[a_id] + start_index)
-                add_new = False
-                break
-
-        if add_new:
-            bins.append(c - size)
-            bins_result.append([indices[a_id] + start_index])
-
-    return bins_result
-
-
-@numba.njit
-def allocate(
-    lengths: np.ndarray, lengths_cumsum: np.ndarray, rank: int, c: int, n: int
-):
-    # Dynamic batch allocator, similar to Multifit
-    # https://en.wikipedia.org/wiki/Multifit_algorithm
-    # ~99.5% efficiency on OpenChat training set (12 * 2048 ctx len)
-
-    s = 0
-    start_index = 0
-    result = []
-
-    while True:
-        # binary search [l, r)
-        left = 1
-        right = 1 + np.searchsorted(lengths_cumsum[start_index:], s + c * n, "right")
-
-        while right - left > 1:
-            mid = (left + right) // 2
-            if ffd_check(lengths[start_index : start_index + mid], c, n):
-                left = mid
-            else:
-                right = mid
-
-        # use length l
-        batch = ffd_with_result(
-            lengths[start_index : start_index + left], c, start_index
-        )
-        assert len(batch) <= n
-        if len(batch) < n:
-            break
-
-        start_index += left
-        s = lengths_cumsum[start_index - 1]
-
-        # add local rank
-        result.append(batch[rank])
-
-    return result, s, len(result) * c * n
-
-
-class MultipackBatchSampler(BatchSampler):
-    """
-    Batch Sampler class for multipack
-    """
-
-    def __init__(
-        self,
-        sampler: Union[Sampler[int], Iterable[int]],
-        batch_size: int,
-        drop_last: bool,
-        batch_max_len: int,
-        lengths: np.ndarray,
-        packing_efficiency_estimate: float = 1.0,
-    ):
-        super().__init__(sampler, batch_size, drop_last)
-        self.batch_size = None
-        self.batch_max_len = batch_max_len
-        self.lengths: np.ndarray = lengths
-        self.packing_efficiency_estimate = packing_efficiency_estimate or 1.0
-
-        assert isinstance(self.lengths, np.ndarray)
-
-        self.epoch = 0
-
-        # statistics
-        self.eff_total_used = 0
-        self.eff_total_slots = 0
-
-    def set_epoch(self, epoch: int):
-        self.epoch = epoch
-
-    def generate_batches(self, set_stats=False):
-        indices = [idx for idx in self.sampler]
-
-        lengths = self.lengths[indices]
-        lengths_cumsum = np.cumsum(lengths)
-
-        batches, total_used, total_slots = allocate(
-            lengths=lengths,
-            lengths_cumsum=lengths_cumsum,
-            rank=0,
-            c=self.batch_max_len,
-            n=1,
-        )
-
-        batches = [[indices[b_idx] for b_idx in batch] for batch in batches]
-
-        # statistics
-        if set_stats:
-            self.eff_total_used += total_used
-            self.eff_total_slots += total_slots
-
-        return batches
-
-    def __iter__(self):
-        batches = self.generate_batches(set_stats=True)
-        return iter(batches)
-
-    def num_batches(self):
-        batches = self.generate_batches(set_stats=True)
-        return len(batches)
-
-    def efficiency(self):
-        return self.eff_total_used / self.eff_total_slots
-
-    def __len__(self):
-        self.num_batches()
-        return self._len_est()
-
-    def _len_est(self):
-        world_size = int(os.getenv("WORLD_SIZE", "1"))
-        lengths_sum = np.sum(self.lengths)
-        lengths_sum_per_device = lengths_sum // world_size
-        LOG.info(
-            f"packing_efficiency_estimate: {self.packing_efficiency_estimate} "
-            f"total_num_tokens per device: {lengths_sum_per_device}"
-        )
-
-        # shave off 1% + 1 for dealing with variance in packing from random sampler to sampler
-        return max(
-            0,
-            (
-                world_size
-                * math.floor(
-                    0.99
-                    * lengths_sum_per_device
-                    / self.packing_efficiency_estimate
-                    // self.batch_max_len
-                )
-                - 1
-            ),
-        )

src/axolotl/utils/tokenization.py

@@ -34,5 +34,6 @@ def check_example_labels(example, tokenizer, text_only=False):
     delimiter = "" if text_only else " "
     LOG.info(delimiter.join(colored_tokens))
     LOG.info("\n\n\n")
+    print(" ".join(colored_tokens))
 
     return " ".join(colored_tokens)

src/axolotl/utils/trainer.py

@@ -1,22 +1,51 @@
 """Module containing the Trainer class and related functions"""
+import importlib
+import logging
 import math
 import os
+import sys
 from contextlib import contextmanager
+from dataclasses import dataclass, field
 from functools import partial
-from typing import List
+from pathlib import Path
+from typing import List, Optional, Union
 
 import numpy as np
 import torch
 import torch.cuda
-from accelerate.logging import get_logger
-from datasets import set_caching_enabled
-from torch.utils.data import DataLoader, RandomSampler
+import torch.distributed as dist
+import transformers
+from datasets import Dataset, set_caching_enabled
+from torch.optim.lr_scheduler import OneCycleLR
+from torch.utils.data import (
+    DataLoader,
+    DistributedSampler,
+    RandomSampler,
+    SequentialSampler,
+)
+from transformers import EarlyStoppingCallback, Trainer, TrainingArguments
+from transformers.trainer_pt_utils import SequentialDistributedSampler
 
-from axolotl.core.trainer_builder import HFCausalTrainerBuilder
-from axolotl.utils.distributed import is_main_process, reduce_and_broadcast, zero_first
-from axolotl.utils.samplers import MultipackBatchSampler
+from axolotl.monkeypatch.relora import ReLoRACallback, ReLoRAScheduler
+from axolotl.utils.callbacks import (
+    EvalFirstStepCallback,
+    GPUStatsCallback,
+    SaveAxolotlConfigtoWandBCallback,
+    SaveBetterTransformerModelCallback,
+    bench_eval_callback_factory,
+    log_prediction_callback_factory,
+)
+from axolotl.utils.collators import DataCollatorForSeq2Seq
+from axolotl.utils.dataloader import MultipackDistributedDataloader
+from axolotl.utils.distributed import (
+    is_distributed,
+    is_main_process,
+    reduce_and_broadcast,
+    zero_first,
+)
+from axolotl.utils.schedulers import get_cosine_schedule_with_quadratic_warmup
 
-LOG = get_logger("axolotl")
+LOG = logging.getLogger("axolotl")
 
 
 @torch.jit.script
@@ -81,6 +110,269 @@ def trainer_weighted_loss(model_output, labels, shift_labels=True):
     return weighted_cross_entropy(logits, labels, weights)
 
 
+@dataclass
+class AxolotlTrainingArguments(TrainingArguments):
+    """
+    Extend the base TrainingArguments for axolotl helpers
+    """
+
+    lr_quadratic_warmup: bool = field(
+        default=False,
+        metadata={"help": "Use quadratic warmup for cosine scheduling."},
+    )
+    sample_packing: bool = field(
+        default=False,
+        metadata={"help": "Use sample packing for efficient training."},
+    )
+    eval_sample_packing: Optional[bool] = field(
+        default=None,
+        metadata={"help": "Use sample packing for efficient evals."},
+    )
+    sample_packing_efficiency: float = field(
+        default=1.0,
+        metadata={"help": "Sample packing efficiency for calculating batch length."},
+    )
+    max_seq_length: int = field(
+        default=2048,
+        metadata={"help": "The maximum sequence length the model can handle"},
+    )
+    sample_packing_seq_len_multiplier: int = field(
+        default=1,
+        metadata={"help": "the multiplier for the max len for packed sequences"},
+    )
+    relora_steps: Optional[int] = field(
+        default=None,
+        metadata={"help": "how often to reset for ReLoRA"},
+    )
+    relora_warmup_steps: Optional[int] = field(
+        default=None,
+        metadata={"help": "how many warmup steps to take after reset for ReLoRA"},
+    )
+    bench_split: Optional[str] = field(
+        default="eval", metadata={"help": "The benchmark split to run on"}
+    )
+    bench_dataset: Optional[str] = field(
+        default="pharaouk/dharma-1/dharma_1_mini.json",
+        metadata={
+            "help": "Benchmark dataset to use: options are `mmlu-zs`, `mmlu-fs`, or the full path to the dataset file"
+        },
+    )
+    do_bench_eval: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to run the Benchmark evaluation."}
+    )
+    max_bench_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "If set, only evaluates on `max_bench_samples` of the benchmark dataset."
+        },
+    )
+    bench_source_max_len: int = field(
+        default=2048, metadata={"help": "Maximum source sequence length for bench."}
+    )
+
+
+class AxolotlTrainer(Trainer):
+    """
+    Extend the base Trainer for axolotl helpers
+    """
+
+    args = None  # type: AxolotlTrainingArguments
+
+    def __init__(self, *args, bench_data_collator=None, **kwargs):
+        self.bench_data_collator = bench_data_collator
+        super().__init__(*args, **kwargs)
+
+    def create_scheduler(
+        self, num_training_steps: int, optimizer: torch.optim.Optimizer = None
+    ):
+        """
+        Setup the scheduler. The optimizer of the trainer must have been set up either before this method is called or
+        passed as an argument.
+
+        Args:
+            num_training_steps (int): The number of training steps to do.
+            optimizer (torch.optim.Optimizer): The training optimizer
+        """
+
+        # fmt: off
+        if self.lr_scheduler is None:  # type: ignore  # pylint: disable=access-member-before-definition
+            # fmt: on
+            if (
+                self.args.lr_scheduler_type == "cosine"
+                and self.args.lr_quadratic_warmup is True
+            ):
+                self.lr_scheduler = get_cosine_schedule_with_quadratic_warmup(  # pylint: disable=attribute-defined-outside-init
+                    optimizer,
+                    num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
+                    num_training_steps=num_training_steps,
+                )
+            else:
+                return super().create_scheduler(num_training_steps, optimizer)
+        return self.lr_scheduler
+
+    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
+        if self.args.world_size > 1 and self.args.sample_packing:
+            return DistributedSampler(
+                self.train_dataset,
+                num_replicas=self.args.world_size,
+                rank=self.args.process_index,
+                seed=self.args.seed,
+            )
+        return super()._get_train_sampler()
+
+    def _get_eval_sampler(
+        self, eval_dataset: Dataset
+    ) -> Optional[torch.utils.data.Sampler]:
+        if (
+            self.args.world_size > 1
+            and self.args.sample_packing
+            and self.args.eval_sample_packing is not False
+        ):
+            return SequentialDistributedSampler(
+                eval_dataset,
+                num_replicas=self.args.world_size,
+                rank=self.args.process_index,
+                batch_size=self.args.per_device_eval_batch_size,
+            )
+        return super()._get_eval_sampler(eval_dataset)
+
+    def get_train_dataloader(self) -> Union[DataLoader, MultipackDistributedDataloader]:
+        if self.args.sample_packing:
+            train_sampler = self._get_train_sampler()
+            return self.accelerator.prepare(
+                MultipackDistributedDataloader(
+                    self.train_dataset,
+                    batch_size=self._train_batch_size,
+                    seq_max_length=self.args.max_seq_length,
+                    collate_fn=self.data_collator,
+                    sampler=train_sampler,
+                    packing_efficiency_estimate=self.args.sample_packing_efficiency,
+                    sample_packing_seq_len_multiplier=self.args.sample_packing_seq_len_multiplier,
+                    device_count=int(os.environ.get("WORLD_SIZE", 1)),
+                )
+            )
+        return super().get_train_dataloader()
+
+    def get_eval_dataloader(
+        self, eval_dataset: Optional[Dataset] = None
+    ) -> Union[DataLoader, MultipackDistributedDataloader]:
+        if self.args.sample_packing and self.args.eval_sample_packing is not False:
+            eval_dataset = (
+                eval_dataset if eval_dataset is not None else self.eval_dataset
+            )
+
+            eval_sampler = self._get_eval_sampler(eval_dataset)
+            return self.accelerator.prepare(
+                MultipackDistributedDataloader(
+                    eval_dataset,
+                    batch_size=self.args.eval_batch_size,
+                    seq_max_length=self.args.max_seq_length,
+                    collate_fn=self.data_collator,
+                    sampler=eval_sampler,
+                    packing_efficiency_estimate=self.args.sample_packing_efficiency,
+                    sample_packing_seq_len_multiplier=self.args.eval_batch_size,
+                    device_count=int(os.environ.get("WORLD_SIZE", 1)),
+                )
+            )
+        return super().get_eval_dataloader(eval_dataset)
+
+    def _get_bench_sampler(
+        self, bench_dataset: Dataset
+    ) -> Optional[torch.utils.data.Sampler]:
+        if self.args.world_size <= 1:
+            return SequentialSampler(bench_dataset)
+        return None
+
+    def get_bench_dataloader(
+        self,
+        bench_dataset: Dataset,
+    ) -> Union[DataLoader, MultipackDistributedDataloader]:
+        dataloader_params = {
+            "batch_size": self.args.eval_batch_size,
+            "collate_fn": self.bench_data_collator,
+            "num_workers": self.args.dataloader_num_workers,
+            "pin_memory": self.args.dataloader_pin_memory,
+        }
+
+        if not isinstance(bench_dataset, torch.utils.data.IterableDataset):
+            dataloader_params["sampler"] = self._get_bench_sampler(bench_dataset)
+            dataloader_params["drop_last"] = self.args.dataloader_drop_last
+
+        return DataLoader(bench_dataset, **dataloader_params)
+        # return self.accelerator.prepare(DataLoader(bench_dataset, **dataloader_params))
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        # use one's weighted cross entropy loss calc
+        # if self.args.sample_packing:
+        #     labels = inputs.pop("labels")
+        #     outputs = model(**inputs)
+        #     loss = trainer_weighted_loss(outputs, labels, shift_labels=True)
+        #     return (loss, outputs) if return_outputs else loss
+        return super().compute_loss(model, inputs, return_outputs=return_outputs)
+
+
+class OneCycleLRSchedulerTrainer(AxolotlTrainer):
+    """
+    Trainer subclass that uses the OneCycleLR scheduler
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.lr_scheduler = None
+
+    def create_scheduler(
+        self,
+        num_training_steps: int,
+        optimizer: Optional[torch.optim.Optimizer] = None,
+    ):
+        optimizer = self.optimizer if optimizer is None else optimizer
+        num_warmup_steps = self.args.get_warmup_steps(num_training_steps)
+        pct_start = num_warmup_steps / num_training_steps
+
+        self.lr_scheduler = OneCycleLR(
+            optimizer,
+            max_lr=self.args.learning_rate,
+            total_steps=num_training_steps,
+            pct_start=pct_start,
+            div_factor=6,
+        )
+
+        return self.lr_scheduler
+
+
+class ReLoRATrainer(AxolotlTrainer):
+    """
+    Trainer subclass that uses the OneCycleLR scheduler
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.lr_scheduler = None
+
+    def create_scheduler(
+        self,
+        num_training_steps: int,
+        optimizer: Optional[torch.optim.Optimizer] = None,
+    ):
+        optimizer = self.optimizer if optimizer is None else optimizer
+        lr_scheduler = super().create_scheduler(num_training_steps, optimizer)
+
+        if self.args.relora_steps:
+            warmup_steps = (
+                self.args.relora_warmup_steps if self.args.relora_warmup_steps else 10
+            )
+            self.lr_scheduler = ReLoRAScheduler(
+                optimizer,
+                lr_scheduler,
+                self.args.relora_steps,
+                warmup_steps,
+            )
+        else:
+            self.lr_scheduler = lr_scheduler
+
+        return self.lr_scheduler
+
+
 def add_position_ids(sample):
     sample_len = len(sample["input_ids"])
     sample["position_ids"] = torch.arange(len(sample["input_ids"]))
@@ -131,10 +423,8 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset, tokenizer):
                     )
 
         # Phi doesn't want the attention_mask feature when training
-        if (
-            "CodeGenTokenizer" in tokenizer.__class__.__name__
-            or (cfg.is_mistral_derived_model and cfg.flash_attention)
-            or cfg.model_config_type == "mamba"
+        if "CodeGenTokenizer" in tokenizer.__class__.__name__ or (
+            cfg.is_mistral_derived_model and cfg.flash_attention
         ):
             train_dataset = train_dataset.remove_columns("attention_mask")
             if eval_dataset:
@@ -143,37 +433,30 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset, tokenizer):
     return train_dataset, eval_dataset
 
 
-def calculate_total_num_steps(cfg, train_dataset, update=True):
-    if not cfg.total_num_tokens:
-        total_num_tokens = np.sum(
-            train_dataset.data.column("input_ids")
-            .to_pandas()
-            .apply(lambda x: len(x))  # pylint: disable=unnecessary-lambda
-            .values
-        )
-        LOG.debug(f"total_num_tokens: {total_num_tokens}", main_process_only=True)
-        if update:
-            cfg.total_num_tokens = total_num_tokens
-
-    skip_estimates = cfg.model_config_type == "mamba"
-
-    if not skip_estimates and not cfg.total_supervised_tokens:
-        total_supervised_tokens = (
-            train_dataset.data.column("labels")
-            .to_pandas()
-            .apply(lambda x: np.sum(np.array(x) != -100))
-            .sum()
-        )
-        LOG.debug(
-            f"`total_supervised_tokens: {total_supervised_tokens}`",
-            main_process_only=True,
-        )
-        if update:
-            cfg.total_supervised_tokens = total_supervised_tokens
-
-    if not skip_estimates and cfg.sample_packing:
+def calculate_total_num_steps(cfg, train_dataset, tokenizer):
+    if cfg.sample_packing:
         # we have to drop anything longer then sequence len otherwise
         # flash attention with position ids fails
+        if not cfg.total_num_tokens:
+            LOG.info("calculating total_num_tokens")
+            total_num_tokens = np.sum(
+                train_dataset.data.column("input_ids")
+                .to_pandas()
+                .apply(lambda x: len(x))  # pylint: disable=unnecessary-lambda
+                .values
+            )
+            LOG.info(f"total_num_tokens: {total_num_tokens}")
+            cfg.total_num_tokens = total_num_tokens
+
+        if not cfg.total_supervised_tokens:
+            total_supervised_tokens = (
+                train_dataset.data.column("labels")
+                .to_pandas()
+                .apply(lambda x: np.sum(np.array(x) != -100))
+                .sum()
+            )
+            LOG.info(f"`total_supervised_tokens: {total_supervised_tokens}`")
+            cfg.total_supervised_tokens = total_supervised_tokens
 
         if cfg.sample_packing_eff_est:
             total_num_steps = (
@@ -191,41 +474,40 @@ def calculate_total_num_steps(cfg, train_dataset, update=True):
                 )
                 * cfg.num_epochs
             )
-            LOG.debug(
-                f"total_num_tokens: {cfg.total_num_tokens}, total_num_steps: {total_num_steps}",
-                main_process_only=True,
+            LOG.info(
+                f"total_num_tokens: {cfg.total_num_tokens}, total_num_steps: {total_num_steps}"
             )
         else:
-            sampler = MultipackBatchSampler(
-                sampler=RandomSampler(train_dataset),
-                batch_size=cfg.micro_batch_size,
-                drop_last=True,
-                batch_max_len=cfg.micro_batch_size
-                * (cfg.max_packed_sequence_len or cfg.sequence_len),
-                lengths=(
-                    train_dataset.data.column("position_ids")
-                    .to_pandas()
-                    .apply(lambda x: x[-1] + 1)
-                    .values
-                ),
-            )
+            if cfg.world_size > 1 and is_distributed():
+                sampler = DistributedSampler(
+                    train_dataset,
+                    num_replicas=cfg.world_size,
+                    rank=dist.get_rank(),
+                    seed=cfg.seed or 42,
+                )
+            else:
+                sampler = RandomSampler(train_dataset)
 
-            data_loader = DataLoader(
-                train_dataset.remove_columns(["length"]),
-                batch_sampler=sampler,
+            data_loader = MultipackDistributedDataloader(
+                train_dataset,
+                batch_size=cfg.micro_batch_size,
+                seq_max_length=cfg.max_packed_sequence_len or cfg.sequence_len,
+                collate_fn=DataCollatorForSeq2Seq(
+                    tokenizer,
+                    return_tensors="pt",
+                    padding="longest",
+                ),
+                sampler=sampler,
+                packing_efficiency_estimate=cfg.sample_packing_eff_est,
+                sample_packing_seq_len_multiplier=cfg.micro_batch_size,
+                device_count=int(os.environ.get("WORLD_SIZE", 1)),
             )
-            data_loader_len = len(data_loader)
-            actual_eff = sampler.efficiency()
-            LOG.debug(f"data_loader_len: {data_loader_len}", main_process_only=True)
+            data_loader_len = data_loader.len_w_stats()
+            actual_eff = data_loader.efficiency()
+            LOG.info(f"data_loader_len: {data_loader_len}")
             # FIXME: is there a bug here somewhere? the total num steps depends
             # on the agreed on value for sample_packing_eff_est
-            total_num_steps = int(
-                math.floor(
-                    data_loader_len
-                    * cfg.num_epochs
-                    / int(os.environ.get("WORLD_SIZE", 1))
-                )
-            )
+            total_num_steps = int(math.floor(data_loader_len * cfg.num_epochs))
 
             def calc_sample_packing_eff_est(estimates: List[float]):
                 LOG.info(f"sample_packing_eff_est across ranks: {repr(estimates)}")
@@ -238,22 +520,13 @@ def calculate_total_num_steps(cfg, train_dataset, update=True):
             sample_packing_eff_est = (
                 math.ceil(sample_packing_actual_eff_all * 100.0) / 100.0
             )
-            if update:
-                cfg.sample_packing_eff_est = sample_packing_eff_est
-            LOG.debug(
-                f"sample_packing_eff_est: {cfg.sample_packing_eff_est}",
-                main_process_only=True,
-            )
+            cfg.sample_packing_eff_est = sample_packing_eff_est
+            LOG.info(f"sample_packing_eff_est: {cfg.sample_packing_eff_est}")
     else:
         total_num_steps = int(
-            math.ceil(
-                len(train_dataset)
-                * cfg.num_epochs
-                / int(os.environ.get("WORLD_SIZE", 1))
-                / cfg.batch_size
-            )
+            math.ceil(len(train_dataset) * cfg.num_epochs / cfg.batch_size)
         )
-    LOG.debug(f"total_num_steps: {total_num_steps}", main_process_only=True)
+    LOG.info(f"total_num_steps: {total_num_steps}")
     return total_num_steps
 
 
@@ -271,16 +544,251 @@ def setup_fsdp_envs(cfg):
         ] = cfg.fsdp_config.fsdp_transformer_layer_cls_to_wrap
 
 
-def prepare_optim_env(cfg):
+def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps):
     if cfg.fsdp:
         setup_fsdp_envs(cfg)
     elif cfg.deepspeed:
         os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
 
+    warmup_steps = (
+        cfg.warmup_steps
+        if cfg.warmup_steps is not None
+        else min(int(0.03 * total_num_steps), 100)
+    )
+    logging_steps = (
+        cfg.logging_steps
+        if cfg.logging_steps is not None
+        else max(min(int(0.005 * total_num_steps), 10), 1)
+    )
 
-def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps):
-    trainer_builder = HFCausalTrainerBuilder(cfg, model, tokenizer)
-    trainer_builder.train_dataset = train_dataset
-    trainer_builder.eval_dataset = eval_dataset
+    training_arguments_kwargs = {}
+    if cfg.bf16 == "full":
+        training_arguments_kwargs["bf16_full_eval"] = True
+    else:
+        training_arguments_kwargs["bf16"] = cfg.bf16
+    training_arguments_kwargs["fp16"] = (cfg.fp16 and not cfg.bf16) or False
+    training_arguments_kwargs["tf32"] = cfg.tf32
+    training_arguments_kwargs["warmup_steps"] = warmup_steps
+    training_arguments_kwargs["logging_steps"] = logging_steps
 
-    return trainer_builder.build(total_num_steps)
+    if cfg.seed:
+        training_arguments_kwargs["seed"] = cfg.seed
+
+    if cfg.gradient_checkpointing:
+        training_arguments_kwargs["gradient_checkpointing"] = cfg.gradient_checkpointing
+    if cfg.fsdp:
+        training_arguments_kwargs["fsdp"] = cfg.fsdp
+        if cfg.fsdp_config:
+            training_arguments_kwargs["fsdp_config"] = dict(cfg.fsdp_config)
+
+    # deepspeed
+    if cfg.deepspeed:
+        training_arguments_kwargs["deepspeed"] = cfg.deepspeed
+
+    if cfg.lr_quadratic_warmup is not None:
+        training_arguments_kwargs["lr_quadratic_warmup"] = cfg.lr_quadratic_warmup
+
+    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.hub_model_id:
+        training_arguments_kwargs["hub_model_id"] = cfg.hub_model_id
+        training_arguments_kwargs["push_to_hub"] = True
+        training_arguments_kwargs["hub_private_repo"] = True
+
+        if cfg.hub_strategy:
+            training_arguments_kwargs["hub_strategy"] = cfg.hub_strategy
+
+    if cfg.save_safetensors:
+        training_arguments_kwargs["save_safetensors"] = cfg.save_safetensors
+
+    if cfg.sample_packing_eff_est:
+        training_arguments_kwargs[
+            "sample_packing_efficiency"
+        ] = cfg.sample_packing_eff_est
+
+    if cfg.eval_steps:
+        training_arguments_kwargs["evaluation_strategy"] = "steps"
+        training_arguments_kwargs["eval_steps"] = cfg.eval_steps
+    elif cfg.evaluation_strategy:
+        training_arguments_kwargs["evaluation_strategy"] = cfg.evaluation_strategy
+    elif cfg.val_set_size == 0:
+        # no eval set, so don't eval
+        training_arguments_kwargs["evaluation_strategy"] = "no"
+    else:
+        # we have an eval set, but no steps defined, default to use epoch
+        training_arguments_kwargs["evaluation_strategy"] = "epoch"
+
+    if cfg.save_steps:
+        training_arguments_kwargs["save_strategy"] = "steps"
+        training_arguments_kwargs["save_steps"] = cfg.save_steps
+    elif cfg.save_strategy:
+        training_arguments_kwargs["save_strategy"] = cfg.save_strategy
+    else:
+        # default to saving each epoch if not defined
+        training_arguments_kwargs["save_strategy"] = "epoch"
+
+    if cfg.do_bench_eval:
+        training_arguments_kwargs["do_bench_eval"] = cfg.do_bench_eval
+        if cfg.bench_dataset:
+            training_arguments_kwargs["bench_dataset"] = cfg.bench_dataset
+    if cfg.metric_for_best_model:
+        training_arguments_kwargs["metric_for_best_model"] = cfg.metric_for_best_model
+    if cfg.greater_is_better:
+        training_arguments_kwargs["greater_is_better"] = cfg.greater_is_better
+
+    if cfg.torch_compile:
+        if torch.__version__ < "2.1.0":  # pylint: disable=protected-access
+            LOG.warning("torch>=2.1.0 required for torch_compile to work properly")
+        else:
+            import torch._dynamo  # pylint: disable=redefined-outer-name
+
+            torch._dynamo.config.suppress_errors = (  # pylint: disable=protected-access
+                True
+            )
+            training_arguments_kwargs["torch_compile"] = cfg.torch_compile
+            if cfg.torch_compile_backend:
+                training_arguments_kwargs[
+                    "torch_compile_backend"
+                ] = cfg.torch_compile_backend
+
+    # DDP Config
+    if cfg.ddp_timeout:
+        training_arguments_kwargs["ddp_timeout"] = cfg.ddp_timeout
+    # see https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html
+    if cfg.ddp_bucket_cap_mb:
+        training_arguments_kwargs["ddp_bucket_cap_mb"] = cfg.ddp_bucket_cap_mb
+    if cfg.ddp_broadcast_buffers is not None:
+        training_arguments_kwargs["ddp_broadcast_buffers"] = cfg.ddp_broadcast_buffers
+
+    training_args = AxolotlTrainingArguments(  # pylint: disable=unexpected-keyword-arg
+        max_steps=total_num_steps if cfg.max_steps else -1,
+        max_seq_length=cfg.sequence_len,
+        per_device_train_batch_size=cfg.micro_batch_size,
+        per_device_eval_batch_size=cfg.eval_batch_size,
+        gradient_accumulation_steps=cfg.gradient_accumulation_steps,
+        eval_accumulation_steps=cfg.gradient_accumulation_steps,
+        num_train_epochs=cfg.num_epochs,
+        learning_rate=cfg.learning_rate,
+        output_dir=cfg.output_dir,
+        save_total_limit=cfg.save_total_limit if cfg.save_total_limit else 4,
+        load_best_model_at_end=(
+            (cfg.load_best_model_at_end is not False or cfg.early_stopping_patience)
+            and cfg.val_set_size > 0
+            and cfg.save_steps
+            and cfg.eval_steps
+            and cfg.save_steps % cfg.eval_steps == 0
+        )
+        or False,
+        ddp_find_unused_parameters=False if cfg.ddp else None,
+        group_by_length=cfg.group_by_length,
+        report_to="wandb" if cfg.use_wandb else None,
+        run_name=cfg.wandb_run_id if cfg.use_wandb else None,
+        optim=cfg.optimizer if cfg.optimizer else "adamw_hf",
+        lr_scheduler_type=cfg.lr_scheduler
+        if cfg.lr_scheduler and cfg.lr_scheduler not in ("one_cycle", "log_sweep")
+        else "cosine",
+        weight_decay=cfg.weight_decay if cfg.weight_decay is not None else 0.0,
+        sample_packing=cfg.sample_packing if cfg.sample_packing else False,
+        eval_sample_packing=cfg.eval_sample_packing,
+        sample_packing_seq_len_multiplier=cfg.micro_batch_size,
+        relora_steps=cfg.relora_steps,
+        relora_warmup_steps=cfg.relora_warmup_steps,
+        **training_arguments_kwargs,
+    )
+
+    trainer_kwargs = {}
+
+    if cfg.optimizer == "adamw_anyprecision":
+        if Path(cfg.torchdistx_path).exists():
+            sys.path.append(cfg.torchdistx_path)
+            importlib.import_module("torchdistx")
+
+    callbacks = []
+    callbacks.append(GPUStatsCallback(cfg))
+    callbacks.append(EvalFirstStepCallback)
+
+    if cfg.relora_steps:
+        callbacks.append(ReLoRACallback(cfg))
+
+    if hasattr(model, "use_bettertransformer") and model.use_bettertransformer is True:
+        callbacks.append(SaveBetterTransformerModelCallback)
+
+    data_collator_kwargs = {
+        "padding": True,  # True/"longest" is the default
+    }
+    if cfg.pad_to_sequence_len:
+        data_collator_kwargs["pad_to_multiple_of"] = 64 * math.ceil(
+            cfg.sequence_len / 64
+        )
+    else:
+        # A100 is best at 64, while others at 8. Let's use the larger so we don't have to check
+        # https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html
+        data_collator_kwargs["pad_to_multiple_of"] = 64
+
+    if cfg.is_llama_derived_model and cfg.landmark_attention:
+        from axolotl.monkeypatch.llama_landmark_attn import (
+            add_mem_tokens,
+            get_mem_id,
+            set_model_mem_id,
+        )
+
+        set_model_mem_id(model, tokenizer)
+
+        LOG.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 = AxolotlTrainer
+    if cfg.lr_scheduler == "one_cycle" and (cfg.fsdp or cfg.adapter == "qlora"):
+        trainer_cls = OneCycleLRSchedulerTrainer
+    elif cfg.relora_steps:
+        trainer_cls = ReLoRATrainer
+    trainer = trainer_cls(
+        model=model,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        args=training_args,
+        data_collator=DataCollatorForSeq2Seq(
+            tokenizer,
+            return_tensors="pt",
+            **data_collator_kwargs,
+        ),
+        bench_data_collator=transformers.DataCollatorForSeq2Seq(
+            tokenizer,
+            return_tensors="pt",
+            **data_collator_kwargs,
+        ),
+        callbacks=callbacks,
+        **trainer_kwargs,
+    )
+
+    if cfg.use_wandb and cfg.eval_table_size > 0:
+        LogPredictionCallback = log_prediction_callback_factory(trainer, tokenizer)
+        trainer.add_callback(LogPredictionCallback(cfg))
+
+    if cfg.use_wandb:
+        trainer.add_callback(SaveAxolotlConfigtoWandBCallback(cfg.axolotl_config_path))
+
+    if cfg.do_bench_eval:
+        trainer.add_callback(bench_eval_callback_factory(trainer, tokenizer))
+
+    # TODO on_save callback to sync checkpoints to GCP/AWS in background
+    if cfg.early_stopping_patience:
+        early_stop_cb = EarlyStoppingCallback(
+            cfg.early_stopping_patience,
+        )
+        trainer.add_callback(early_stop_cb)
+
+    return trainer

src/axolotl/utils/wandb_.py

@@ -2,20 +2,20 @@
 
 import os
 
-from axolotl.utils.dict import DictDefault
 
-
-def setup_wandb_env_vars(cfg: DictDefault):
-    for key in cfg.keys():
-        if key.startswith("wandb_"):
-            value = cfg.get(key, "")
-
-            if value and isinstance(value, str) and len(value) > 0:
-                os.environ[key.upper()] = value
-
-    # Enable wandb if project name is present
-    if cfg.wandb_project and len(cfg.wandb_project) > 0:
+def setup_wandb_env_vars(cfg):
+    if cfg.wandb_mode and cfg.wandb_mode == "offline":
+        os.environ["WANDB_MODE"] = cfg.wandb_mode
+    elif cfg.wandb_project and len(cfg.wandb_project) > 0:
+        os.environ["WANDB_PROJECT"] = cfg.wandb_project
         cfg.use_wandb = True
-        os.environ.pop("WANDB_DISABLED", None)  # Remove if present
+        if cfg.wandb_entity and len(cfg.wandb_entity) > 0:
+            os.environ["WANDB_ENTITY"] = cfg.wandb_entity
+        if cfg.wandb_watch and len(cfg.wandb_watch) > 0:
+            os.environ["WANDB_WATCH"] = cfg.wandb_watch
+        if cfg.wandb_log_model and len(cfg.wandb_log_model) > 0:
+            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/e2e/test_fused_llama.py

@@ -1,73 +0,0 @@
-"""
-E2E tests for lora llama
-"""
-
-import logging
-import os
-import unittest
-from pathlib import Path
-
-from transformers.utils import is_torch_bf16_gpu_available
-
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
-from axolotl.train import train
-from axolotl.utils.config import normalize_config
-from axolotl.utils.dict import DictDefault
-
-from .utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
-
-
-class TestFusedLlama(unittest.TestCase):
-    """
-    Test case for Llama models using Fused layers
-    """
-
-    @with_temp_dir
-    def test_fft_packing(self, temp_dir):
-        # pylint: disable=duplicate-code
-        cfg = DictDefault(
-            {
-                "base_model": "JackFram/llama-68m",
-                "flash_attention": True,
-                "flash_attn_fuse_qkv": True,
-                "flash_attn_fuse_mlp": True,
-                "sample_packing": True,
-                "sequence_len": 1024,
-                "val_set_size": 0.1,
-                "special_tokens": {
-                    "unk_token": "<unk>",
-                    "bos_token": "<s>",
-                    "eos_token": "</s>",
-                },
-                "datasets": [
-                    {
-                        "path": "mhenrichsen/alpaca_2k_test",
-                        "type": "alpaca",
-                    },
-                ],
-                "num_epochs": 2,
-                "micro_batch_size": 2,
-                "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
-                "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
-            }
-        )
-        if is_torch_bf16_gpu_available():
-            cfg.bf16 = True
-        else:
-            cfg.fp16 = True
-        normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
-
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/test_lora_llama.py

@@ -4,6 +4,7 @@ E2E tests for lora llama
 
 import logging
 import os
+import tempfile
 import unittest
 from pathlib import Path
 
@@ -13,8 +14,6 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
-
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
@@ -24,12 +23,17 @@ class TestLoraLlama(unittest.TestCase):
     Test case for Llama models using LoRA
     """
 
-    @with_temp_dir
-    def test_lora(self, temp_dir):
+    def test_lora(self):
+        """
+        support for legacy lora_ configs
+        :return:
+        """
         # pylint: disable=duplicate-code
+        output_dir = tempfile.mkdtemp()
         cfg = DictDefault(
             {
                 "base_model": "JackFram/llama-68m",
+                "base_model_config": "JackFram/llama-68m",
                 "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
@@ -53,7 +57,7 @@ class TestLoraLlama(unittest.TestCase):
                 "num_epochs": 2,
                 "micro_batch_size": 8,
                 "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
+                "output_dir": output_dir,
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_torch",
                 "lr_scheduler": "cosine",
@@ -64,14 +68,110 @@ class TestLoraLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        assert (Path(output_dir) / "adapter_model.bin").exists()
 
-    @with_temp_dir
-    def test_lora_packing(self, temp_dir):
+    def test_lora_peft(self):
+        """
+        support for legacy lora_ configs
+        :return:
+        """
         # pylint: disable=duplicate-code
+        output_dir = tempfile.mkdtemp()
         cfg = DictDefault(
             {
                 "base_model": "JackFram/llama-68m",
+                "base_model_config": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "peft_r": 32,
+                "peft_alpha": 64,
+                "peft_dropout": 0.05,
+                "peft_target_linear": True,
+                "val_set_size": 0.1,
+                "special_tokens": {
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 2,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": output_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+            }
+        )
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+
+        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+        assert (Path(output_dir) / "adapter_model.bin").exists()
+
+    def test_ia3_peft(self):
+        """
+        support for IA3 peft
+        :return:
+        """
+        # pylint: disable=duplicate-code
+        output_dir = tempfile.mkdtemp()
+        cfg = DictDefault(
+            {
+                "base_model": "JackFram/llama-68m",
+                "base_model_config": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "ia3",
+                "peft_r": 32,
+                "peft_alpha": 64,
+                "peft_dropout": 0.05,
+                "peft_target_modules": ["k_proj", "v_proj", "down_proj"],
+                "peft_feedforward_modules": ["down_proj"],
+                "val_set_size": 0.1,
+                "special_tokens": {
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 2,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": output_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+            }
+        )
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+
+        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+        assert (Path(output_dir) / "adapter_model.bin").exists()
+
+    def test_lora_packing(self):
+        # pylint: disable=duplicate-code
+        output_dir = tempfile.mkdtemp()
+        cfg = DictDefault(
+            {
+                "base_model": "JackFram/llama-68m",
+                "base_model_config": "JackFram/llama-68m",
                 "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "sample_packing": True,
@@ -97,11 +197,10 @@ class TestLoraLlama(unittest.TestCase):
                 "num_epochs": 2,
                 "micro_batch_size": 8,
                 "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
+                "output_dir": output_dir,
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_torch",
                 "lr_scheduler": "cosine",
-                "bf16": True,
             }
         )
         normalize_config(cfg)
@@ -109,14 +208,15 @@ class TestLoraLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        assert (Path(output_dir) / "adapter_model.bin").exists()
 
-    @with_temp_dir
-    def test_lora_gptq(self, temp_dir):
+    def test_lora_gptq(self):
         # pylint: disable=duplicate-code
+        output_dir = tempfile.mkdtemp()
         cfg = DictDefault(
             {
                 "base_model": "TheBlokeAI/jackfram_llama-68m-GPTQ",
+                "base_model_config": "TheBlokeAI/jackfram_llama-68m-GPTQ",
                 "model_type": "AutoModelForCausalLM",
                 "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
@@ -146,7 +246,7 @@ class TestLoraLlama(unittest.TestCase):
                 "save_steps": 0.5,
                 "micro_batch_size": 8,
                 "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
+                "output_dir": output_dir,
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_torch",
                 "lr_scheduler": "cosine",
@@ -157,4 +257,4 @@ class TestLoraLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        assert (Path(output_dir) / "adapter_model.bin").exists()

tests/e2e/test_mamba.py

@@ -1,65 +0,0 @@
-"""
-E2E tests for lora llama
-"""
-
-import logging
-import os
-import unittest
-from pathlib import Path
-
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
-from axolotl.train import train
-from axolotl.utils.config import normalize_config
-from axolotl.utils.dict import DictDefault
-
-from .utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
-
-
-class TestMistral(unittest.TestCase):
-    """
-    Test case for Llama models using LoRA
-    """
-
-    @with_temp_dir
-    def test_fft(self, temp_dir):
-        # pylint: disable=duplicate-code
-        cfg = DictDefault(
-            {
-                "base_model": "state-spaces/mamba-130m",
-                "model_type": "MambaLMHeadModel",
-                "tokenizer_type": "AutoTokenizer",
-                "tokenizer_config": "EleutherAI/gpt-neox-20b",
-                "flash_attention": False,
-                "sequence_len": 1024,
-                "load_in_8bit": False,
-                "val_set_size": 0.0,
-                "datasets": [
-                    {
-                        "path": "mhenrichsen/alpaca_2k_test",
-                        "type": "alpaca",
-                    },
-                ],
-                "gradient_checkpointing": False,
-                "num_epochs": 2,
-                "micro_batch_size": 2,
-                "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
-                "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": None,
-                "save_safetensors": False,
-            }
-        )
-        normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
-
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/test_mistral.py

@@ -4,6 +4,7 @@ E2E tests for lora llama
 
 import logging
 import os
+import tempfile
 import unittest
 from pathlib import Path
 
@@ -15,8 +16,6 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
-
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
@@ -26,12 +25,13 @@ class TestMistral(unittest.TestCase):
     Test case for Llama models using LoRA
     """
 
-    @with_temp_dir
-    def test_lora(self, temp_dir):
+    def test_lora(self):
         # pylint: disable=duplicate-code
+        output_dir = tempfile.mkdtemp()
         cfg = DictDefault(
             {
                 "base_model": "openaccess-ai-collective/tiny-mistral",
+                "base_model_config": "openaccess-ai-collective/tiny-mistral",
                 "flash_attention": True,
                 "sequence_len": 1024,
                 "load_in_8bit": True,
@@ -55,7 +55,7 @@ class TestMistral(unittest.TestCase):
                 "num_epochs": 2,
                 "micro_batch_size": 2,
                 "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
+                "output_dir": output_dir,
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_torch",
                 "lr_scheduler": "cosine",
@@ -69,14 +69,15 @@ class TestMistral(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        assert (Path(output_dir) / "adapter_model.bin").exists()
 
-    @with_temp_dir
-    def test_ft(self, temp_dir):
+    def test_ft(self):
         # pylint: disable=duplicate-code
+        output_dir = tempfile.mkdtemp()
         cfg = DictDefault(
             {
                 "base_model": "openaccess-ai-collective/tiny-mistral",
+                "base_model_config": "openaccess-ai-collective/tiny-mistral",
                 "flash_attention": True,
                 "sequence_len": 1024,
                 "val_set_size": 0.1,
@@ -94,7 +95,7 @@ class TestMistral(unittest.TestCase):
                 "num_epochs": 2,
                 "micro_batch_size": 2,
                 "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
+                "output_dir": output_dir,
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_torch",
                 "lr_scheduler": "cosine",
@@ -112,4 +113,4 @@ class TestMistral(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        assert (Path(output_dir) / "pytorch_model.bin").exists()

tests/e2e/test_mistral_samplepack.py

@@ -4,6 +4,7 @@ E2E tests for lora llama
 
 import logging
 import os
+import tempfile
 import unittest
 from pathlib import Path
 
@@ -15,8 +16,6 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
-
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
@@ -26,12 +25,13 @@ class TestMistral(unittest.TestCase):
     Test case for Llama models using LoRA
     """
 
-    @with_temp_dir
-    def test_lora_packing(self, temp_dir):
+    def test_lora_packing(self):
         # pylint: disable=duplicate-code
+        output_dir = tempfile.mkdtemp()
         cfg = DictDefault(
             {
                 "base_model": "openaccess-ai-collective/tiny-mistral",
+                "base_model_config": "openaccess-ai-collective/tiny-mistral",
                 "flash_attention": True,
                 "sample_packing": True,
                 "sequence_len": 1024,
@@ -56,7 +56,7 @@ class TestMistral(unittest.TestCase):
                 "num_epochs": 2,
                 "micro_batch_size": 2,
                 "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
+                "output_dir": output_dir,
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_torch",
                 "lr_scheduler": "cosine",
@@ -70,14 +70,15 @@ class TestMistral(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        assert (Path(output_dir) / "adapter_model.bin").exists()
 
-    @with_temp_dir
-    def test_ft_packing(self, temp_dir):
+    def test_ft_packing(self):
         # pylint: disable=duplicate-code
+        output_dir = tempfile.mkdtemp()
         cfg = DictDefault(
             {
                 "base_model": "openaccess-ai-collective/tiny-mistral",
+                "base_model_config": "openaccess-ai-collective/tiny-mistral",
                 "flash_attention": True,
                 "sample_packing": True,
                 "sequence_len": 1024,
@@ -96,7 +97,7 @@ class TestMistral(unittest.TestCase):
                 "num_epochs": 2,
                 "micro_batch_size": 2,
                 "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
+                "output_dir": output_dir,
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_torch",
                 "lr_scheduler": "cosine",
@@ -114,4 +115,4 @@ class TestMistral(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        assert (Path(output_dir) / "pytorch_model.bin").exists()

tests/e2e/test_phi.py

@@ -4,8 +4,8 @@ E2E tests for lora llama
 
 import logging
 import os
+import tempfile
 import unittest
-from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -13,8 +13,6 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
-
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
@@ -24,14 +22,14 @@ class TestPhi(unittest.TestCase):
     Test case for Llama models using LoRA
     """
 
-    @with_temp_dir
-    def test_ft(self, temp_dir):
+    def test_ft(self):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
                 "base_model": "microsoft/phi-1_5",
+                "base_model_config": "microsoft/phi-1_5",
                 "trust_remote_code": True,
-                "model_type": "PhiForCausalLM",
+                "model_type": "MixFormerSequentialForCausalLM",
                 "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 512,
                 "sample_packing": False,
@@ -55,7 +53,7 @@ class TestPhi(unittest.TestCase):
                 "num_epochs": 1,
                 "micro_batch_size": 1,
                 "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
+                "output_dir": tempfile.mkdtemp(),
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
@@ -67,16 +65,15 @@ class TestPhi(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
 
-    @with_temp_dir
-    def test_ft_packed(self, temp_dir):
+    def test_ft_packed(self):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
                 "base_model": "microsoft/phi-1_5",
+                "base_model_config": "microsoft/phi-1_5",
                 "trust_remote_code": True,
-                "model_type": "PhiForCausalLM",
+                "model_type": "MixFormerSequentialForCausalLM",
                 "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 512,
                 "sample_packing": True,
@@ -100,7 +97,7 @@ class TestPhi(unittest.TestCase):
                 "num_epochs": 1,
                 "micro_batch_size": 1,
                 "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
+                "output_dir": tempfile.mkdtemp(),
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
@@ -112,4 +109,3 @@ class TestPhi(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()