chore: refactor normalize_attn to use mapping and loop

* fsdp embeddings should be float32 per comment

* patch peft to not upcast everything

* add tabs back to code check

* fix import

* add configurable option and fix check

* add check for dtypes

* move embeddings test to patch dir

* fix test

* fix comment and logic

.github/workflows/tests-nightly.yml

@@ -18,96 +18,9 @@ jobs:
         env:
           SKIP: no-commit-to-branch
 
-  preload-cache:
-    name: Preload HF cache
-    runs-on: ubuntu-latest
-    strategy:
-      fail-fast: false
-      matrix:
-        python_version: ["3.11"]
-        pytorch_version: ["2.6.0"]
-    timeout-minutes: 20
-
-    env:
-      AXOLOTL_IS_CI_CACHE_PRELOAD: "1"
-
-    steps:
-      - name: Check out repository code
-        uses: actions/checkout@v4
-
-      - name: Restore HF cache
-        id: hf-cache-restore
-        uses: actions/cache/restore@v4
-        with:
-          path: |
-            /home/runner/.cache/huggingface/hub/datasets--*
-            /home/runner/.cache/huggingface/hub/models--*
-          key: ${{ runner.os }}-hf-hub-cache-v2
-
-      - name: Setup Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python_version }}
-          cache: 'pip' # caching pip dependencies
-
-      - name: upgrade pip
-        run: |
-          pip3 install --upgrade pip
-          pip3 install --upgrade packaging==23.2 setuptools==75.8.0 wheel
-
-      - name: Install PyTorch
-        run: |
-          pip3 install torch==${{ matrix.pytorch_version }}
-
-      - name: Install dependencies
-        run: |
-          pip3 show torch
-          pip3 install --no-build-isolation -U -e .
-          python scripts/unsloth_install.py | sh
-          python scripts/cutcrossentropy_install.py | sh
-          pip3 install -r requirements-dev.txt -r requirements-tests.txt
-
-      - name: Make sure PyTorch version wasn't clobbered
-        run: |
-          python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__"
-
-      - name: Ensure axolotl CLI was installed
-        run: |
-          axolotl --help
-
-      - name: Pre-Download dataset fixture
-        run: |
-          huggingface-cli download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures
-
-      - name: Run tests
-        run: |
-          pytest -v tests/conftest.py
-
-      - name: Upload coverage to Codecov
-        uses: codecov/codecov-action@v5
-        with:
-          token: ${{ secrets.CODECOV_TOKEN }}
-          files: ./coverage.xml
-          flags: unittests,pytorch-${{ matrix.pytorch_version }}
-          fail_ci_if_error: false
-
-      - name: cleanup pip cache
-        run: |
-          find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
-
-      - name: Save HF cache
-        id: hf-cache
-        uses: actions/cache/save@v4
-        with:
-          path: |
-            /home/runner/.cache/huggingface/hub/datasets--*
-            /home/runner/.cache/huggingface/hub/models--*
-          key: ${{ steps.hf-cache-restore.outputs.cache-primary-key }}
-
   pytest:
     name: PyTest
     runs-on: ubuntu-latest
-    needs: [preload-cache]
     strategy:
       fail-fast: false
       max-parallel: 2

.github/workflows/tests.yml

@@ -329,6 +329,18 @@ jobs:
       fail-fast: false
       matrix:
         include:
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.6.0
+            num_gpus: 1
+            axolotl_extras: llmcompressor
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.4.1
+            num_gpus: 1
+            axolotl_extras:
           - cuda: 124
             cuda_version: 12.4.1
             python_version: "3.11"

docs/config.qmd

@@ -612,7 +612,6 @@ lr_div_factor: # Learning rate div factor
 # - optimi_adamw
 # - ao_adamw_8bit
 # - ao_adamw_fp8
-# - came_pytorch
 optimizer:
 # Dictionary of arguments to pass to the optimizer
 optim_args:

docs/custom_integrations.qmd

@@ -49,7 +49,8 @@ sections = [
     ("Knowledge Distillation (KD)", "kd"),
     ("Liger Kernels", "liger"),
     ("Language Model Evaluation Harness (LM Eval)", "lm_eval"),
-    ("Spectrum", "spectrum")
+    ("Spectrum", "spectrum"),
+    ("LLMCompressor", "llm_compressor")
 ]
 
 for section_name, folder_name in sections:

examples/cerebras/btlm-ft.yml

@@ -59,9 +59,7 @@ gradient_checkpointing: false
 resume_from_checkpoint:
 logging_steps: 1
 
-flash_attention: true
+attention: flash
-sdp_attention:
-flash_optimum:
 
 gptq_groupsize:
 gptq_model_v1:

examples/cerebras/qlora.yml

@@ -39,8 +39,7 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-xformers_attention: true
+attention: xformers
-flash_attention:
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 10

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

@@ -45,7 +45,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

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

@@ -46,7 +46,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

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

@@ -45,7 +45,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

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

@@ -46,7 +46,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

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

@@ -45,7 +45,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

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

@@ -46,7 +46,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/cohere/command-r-7b-qlora.yml

@@ -49,7 +49,8 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_ratio: 0.1
 evals_per_epoch:

examples/colab-notebooks/colab-axolotl-example.ipynb

@@ -112,9 +112,7 @@
     "early_stopping_patience:\n",
     "resume_from_checkpoint:\n",
     "logging_steps: 1\n",
-    "xformers_attention:\n",
+    "attention: sdpa\n",
-    "flash_attention: false\n",
-    "sdp_attention: true\n",
     "\n",
     "warmup_steps: 1\n",
     "max_steps: 25\n",

examples/dbrx/16bit-lora.yaml

@@ -52,7 +52,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch:

examples/dbrx/8bit-lora.yaml

@@ -55,7 +55,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch:

examples/dbrx/fft-ds-zero3.yaml

@@ -39,7 +39,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch:

examples/deepseek-v2/fft-fsdp-16b.yaml

@@ -35,7 +35,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 100
 evals_per_epoch: 2

examples/deepseek-v2/qlora-fsdp-2_5.yaml

@@ -59,7 +59,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 100
 evals_per_epoch: 2

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

@@ -43,8 +43,7 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-xformers_attention: true
+attention: xformers
-flash_attention:
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 40

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

@@ -73,8 +73,7 @@ early_stopping_patience: 3
 resume_from_checkpoint:
 auto_resume_from_checkpoints: true
 logging_steps: 1
-xformers_attention: true
+attention: xformers
-flash_attention:
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 10

examples/falcon/config-7b.yml

@@ -40,8 +40,7 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-xformers_attention: true
+attention: xformers
-flash_attention:
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 40

examples/gemma/qlora.yml

@@ -47,7 +47,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_ratio: 0.1
 evals_per_epoch: 4

examples/gemma2/qlora.yml

@@ -53,7 +53,8 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_ratio: 0.1
 evals_per_epoch:

examples/gemma2/reward-model.yaml

@@ -43,7 +43,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_ratio: 0.1
 evals_per_epoch:

examples/gemma3/gemma-3-1b-qlora.yml

@@ -57,7 +57,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_ratio: 0.1
 evals_per_epoch:

examples/gemma3/gemma-3-4b-qlora.yml

@@ -51,8 +51,7 @@ gradient_checkpointing: true
 gradient_checkpointing_kwargs:
   use_reentrant: false
 logging_steps: 1
-flash_attention: true
+attention: flash
-eager_attention:
 
 warmup_ratio: 0.1
 evals_per_epoch: 1

examples/gemma3/gemma-3-4b-vision-qlora.yml

@@ -53,8 +53,7 @@ gradient_checkpointing: true
 gradient_checkpointing_kwargs:
   use_reentrant: false
 logging_steps: 1
-flash_attention: true
+attention: flash
-eager_attention:
 
 warmup_ratio: 0.1
 evals_per_epoch: 1

examples/gptj/qlora.yml

@@ -36,8 +36,7 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-xformers_attention: true
+attention: xformers
-flash_attention:
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 10

examples/jamba/qlora.yaml

@@ -47,7 +47,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch:

examples/jamba/qlora_deepspeed.yaml

@@ -46,7 +46,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch:

examples/jamba/qlora_fsdp_large.yaml

@@ -45,7 +45,8 @@ gradient_checkpointing: true
 gradient_checkpointing_kwargs:
   use_reentrant: true
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 1

examples/jeopardy-bot/config.yml

@@ -37,8 +37,7 @@ bf16: auto
 tf32: true
 resume_from_checkpoint:
 logging_steps: 5
-xformers_attention: true
+attention: xformers
-flash_attention:
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 20

examples/llama-2/fft_optimized.yml

@@ -42,7 +42,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 flash_attn_cross_entropy: false
 flash_attn_rms_norm: true
 flash_attn_fuse_qkv: false

examples/llama-2/gptq-lora.yml

@@ -53,9 +53,7 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention:
+attention: flash
-sdp_attention:
-flash_optimum:
 warmup_steps: 100
 evals_per_epoch: 4
 saves_per_epoch: 1

examples/llama-2/lisa.yml

@@ -46,7 +46,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 flash_attn_cross_entropy: false
 flash_attn_rms_norm: true
 flash_attn_fuse_qkv: false

examples/llama-2/loftq.yml

@@ -45,7 +45,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-2/lora.yml

@@ -45,7 +45,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-2/qlora-fsdp.yml

@@ -48,7 +48,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-2/qlora.yml

@@ -46,7 +46,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-2/relora.yml

@@ -48,7 +48,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-3-vision/lora-11b.yaml

@@ -50,8 +50,7 @@ tf32: true
 
 gradient_checkpointing: true
 logging_steps: 1
-flash_attention: true
+attention: flash
-eager_attention:
 
 warmup_ratio: 0.1
 evals_per_epoch: 1

examples/llama-3/fft-8b-liger-fsdp.yaml

@@ -49,7 +49,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 100
 evals_per_epoch: 2

examples/llama-3/fft-8b.yaml

@@ -34,7 +34,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 100
 evals_per_epoch: 2

examples/llama-3/instruct-dpo-lora-8b.yml

@@ -61,7 +61,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-3/instruct-lora-8b.yml

@@ -56,7 +56,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-3/lora-1b-deduplicate-dpo.yml

@@ -77,7 +77,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-3/lora-1b-deduplicate-sft.yml

@@ -53,7 +53,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-3/lora-1b-kernels.yml

@@ -54,7 +54,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/llama-3/lora-1b-ray.yml

@@ -48,7 +48,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/llama-3/lora-1b-sample-packing-sequentially.yml

@@ -55,7 +55,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-3/lora-1b.yml

@@ -48,7 +48,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/llama-3/lora-8b.yml

@@ -49,7 +49,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-3/qlora-1b-kto.yaml

@@ -53,7 +53,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 20
 evals_per_epoch: 4

examples/llama-3/qlora-1b.yml

@@ -51,7 +51,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/llama-3/qlora-fsdp-405b.yaml

@@ -39,7 +39,8 @@ gradient_checkpointing: true
 gradient_checkpointing_kwargs:
   use_reentrant: true
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-3/qlora-fsdp-70b.yaml

@@ -48,7 +48,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-3/qlora.yml

@@ -46,7 +46,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/llama-3/sparse-finetuning.yaml

@@ -0,0 +1,77 @@
+base_model: neuralmagic/Sparse-Llama-3.1-8B-2of4
+
+plugins:
+  - axolotl.integrations.llm_compressor.LLMCompressorPlugin
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.05
+output_dir: ./outputs/out
+
+sequence_len: 4096
+sample_packing: true
+pad_to_sequence_len: true
+eval_sample_packing: false
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 8
+micro_batch_size: 1
+num_epochs: 1
+optimizer: paged_adamw_8bit
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+train_on_inputs: false
+group_by_length: false
+bf16: auto
+fp16:
+tf32: false
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+early_stopping_patience:
+resume_from_checkpoint:
+logging_steps: 1
+xformers_attention:
+flash_attention: true
+
+warmup_steps: 100
+evals_per_epoch: 2
+eval_table_size:
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:
+special_tokens:
+  pad_token: <|end_of_text|>
+
+llmcompressor:
+  recipe:
+    finetuning_stage:
+      finetuning_modifiers:
+        ConstantPruningModifier:
+          targets: [
+            're:.*q_proj.weight',
+            're:.*k_proj.weight',
+            're:.*v_proj.weight',
+            're:.*o_proj.weight',
+            're:.*gate_proj.weight',
+            're:.*up_proj.weight',
+            're:.*down_proj.weight',
+          ]
+          start: 0
+  save_compressed: true

examples/llama-4/README.md

@@ -34,5 +34,3 @@ We provide a script to delinearize Llama 4 linearized models into regular Huggin
 ```bash
 axolotl delinearize-llama4 --model path/to/model_dir --output path/to/output_dir
 ```
-
-Note: This only works with the non-quantized linearized model. If you have an adapter, merge it with the *non-quantized linearized* model before delinearizing.

examples/llava/lora-7b.yaml

@@ -46,8 +46,7 @@ tf32: true
 
 gradient_checkpointing: true
 logging_steps: 1
-flash_attention: true
+attention: flash
-eager_attention:
 
 warmup_ratio: 0.1
 evals_per_epoch: 1

examples/mamba/config.yml

@@ -39,7 +39,7 @@ tf32: true
 gradient_checkpointing: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention:
+attention: eager
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/mistral/bigstral-ds-zero3.yaml

@@ -42,7 +42,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 save_total_limit: 1
 save_steps:

examples/mistral/config.yml

@@ -36,7 +36,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/mistral/lora-mps.yml

@@ -53,8 +53,7 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: false
+attention: sdpa
-sdp_attention: true
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/mistral/lora.yml

@@ -54,7 +54,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/mistral/mistral-dpo-qlora.yml

@@ -71,7 +71,7 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: false
+attention: eager
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/mistral/mistral-qlora-fsdp.yml

@@ -51,7 +51,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/mistral/mistral-qlora-orpo.yml

@@ -59,7 +59,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/mistral/mistral-small-3.1-24B-lora.yml

@@ -48,9 +48,7 @@ tf32: true
 
 gradient_checkpointing: true
 logging_steps: 1
-flash_attention: false # PixtralVisionModel does not support Flash Attention 2.0 yet.
+attention: eager  # PixtralVisionModel does not support Flash Attention 2.0 yet.
-eager_attention:
-
 warmup_ratio: 0.1
 evals_per_epoch: 1
 saves_per_epoch: 1

examples/mistral/mixtral-8x22b-qlora-fsdp.yml

@@ -49,7 +49,8 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/mistral/mixtral-qlora-fsdp.yml

@@ -51,7 +51,8 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/mistral/mixtral.yml

@@ -69,7 +69,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/mistral/mixtral_22.yml

@@ -40,7 +40,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 save_total_limit: 1
 save_steps:

examples/mistral/qlora.yml

@@ -54,7 +54,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3

examples/mpt-7b/config.yml

@@ -39,7 +39,7 @@ bf16: auto
 tf32: true
 resume_from_checkpoint:
 logging_steps: 5
-flash_attention:
+attention: eager
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 20

examples/openllama-3b/config.yml

@@ -39,7 +39,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 20

examples/openllama-3b/lora.yml

@@ -47,7 +47,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 20

examples/openllama-3b/qlora.yml

@@ -40,7 +40,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 20

examples/phi/phi-ft.yml

@@ -48,7 +48,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: True
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 100
 evals_per_epoch: 4

examples/phi/phi-qlora.yml

@@ -51,7 +51,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: True
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 100
 evals_per_epoch: 4

examples/phi/phi2-ft.yml

@@ -48,7 +48,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: True
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 100
 evals_per_epoch: 4

examples/phi/phi3-ft-fsdp.yml

@@ -49,7 +49,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 100
 evals_per_epoch: 4

examples/phi/phi3-ft.yml

@@ -44,7 +44,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: True
 early_stopping_patience: 3
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 eval_steps: 1000
 save_steps: 5000

examples/pixtral/lora-12b.yml

@@ -46,8 +46,7 @@ tf32: true
 
 gradient_checkpointing: true
 logging_steps: 1
-flash_attention: false # PixtralVisionModel does not support Flash Attention 2.0 yet
+attention: eager  # PixtralVisionModel does not support Flash Attention 2.0 yet
-eager_attention:
 
 warmup_ratio: 0.1
 evals_per_epoch: 1

examples/qwen/lora.yml

@@ -47,7 +47,7 @@ tf32: false
 gradient_checkpointing: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention:
+attention: eager
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/qwen/qlora.yml

@@ -47,7 +47,7 @@ tf32: false
 gradient_checkpointing: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention:
+attention: flash
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/qwen/qwen2-moe-lora.yaml

@@ -43,7 +43,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/qwen/qwen2-moe-qlora.yaml

@@ -46,7 +46,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/qwen2-vl/lora-7b.yaml

@@ -46,8 +46,7 @@ tf32: true
 
 gradient_checkpointing: true
 logging_steps: 1
-flash_attention: true
+attention: flash
-eager_attention:
 
 warmup_ratio: 0.1
 evals_per_epoch: 1

examples/qwen2/dpo.yaml

@@ -49,7 +49,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/qwen2/prm.yaml

@@ -47,7 +47,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_ratio: 0.1
 evals_per_epoch:

examples/qwen2/qlora-fsdp.yaml

@@ -47,7 +47,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 10
 evals_per_epoch: 4

examples/qwen2/reward-model.yaml

@@ -43,7 +43,8 @@ gradient_checkpointing_kwargs:
   use_reentrant: false
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_ratio: 0.1
 evals_per_epoch:

examples/redpajama/config-3b.yml

@@ -40,7 +40,7 @@ bf16: auto
 tf32: true
 resume_from_checkpoint:
 logging_steps: 5
-flash_attention:
+attention: flash
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 20

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

@@ -38,7 +38,7 @@ tf32: true
 gradient_checkpointing:
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention:
+attention: eager
 gptq_groupsize:
 gptq_model_v1:
 warmup_steps: 20

examples/stablelm-2/1.6b/fft.yml

@@ -44,7 +44,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 flash_attn_cross_entropy: false
 flash_attn_rms_norm: true
 flash_attn_fuse_qkv: false

examples/stablelm-2/1.6b/lora.yml

@@ -47,7 +47,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 flash_attn_cross_entropy: false
 flash_attn_rms_norm: true
 

examples/starcoder2/qlora.yml

@@ -46,7 +46,8 @@ tf32: false
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: true
+attention: flash
+
 
 warmup_steps: 20
 evals_per_epoch: 4

examples/tiny-llama/lora-mps.yml

@@ -47,7 +47,7 @@ tf32: true
 gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
-flash_attention: false
+attention: eager
 
 warmup_steps: 10
 evals_per_epoch: 0