bump transformers to 4.51.3

.coveragerc

@@ -1,14 +0,0 @@
-[run]
-source = axolotl
-omit =
-    */tests/*
-    setup.py
-
-[report]
-exclude_lines =
-    pragma: no cover
-    def __repr__
-    raise NotImplementedError
-    if __name__ == .__main__.:
-    pass
-    raise ImportError

.github/workflows/base.yml

@@ -46,18 +46,6 @@ jobs:
             python_version: "3.11"
             pytorch: 2.6.0
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "126"
-            cuda_version: 12.6.3
-            cudnn_version: ""
-            python_version: "3.11"
-            pytorch: 2.7.0
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "128"
-            cuda_version: 12.6.3
-            cudnn_version: ""
-            python_version: "3.11"
-            pytorch: 2.7.0
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
           - cuda: "128"
             cuda_version: 12.8.1
             cudnn_version: ""

.github/workflows/main.yml

@@ -29,13 +29,8 @@ jobs:
             cuda_version: 12.4.1
             python_version: "3.11"
             pytorch: 2.6.0
-            axolotl_extras: vllm
+            axolotl_extras:
             is_latest: true
-          - cuda: 126
-            cuda_version: 12.6.3
-            python_version: "3.11"
-            pytorch: 2.7.0
-            axolotl_extras: vllm
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
@@ -98,11 +93,6 @@ jobs:
             pytorch: 2.6.0
             axolotl_extras:
             is_latest: true
-          - cuda: 126
-            cuda_version: 12.6.3
-            python_version: "3.11"
-            pytorch: 2.7.0
-            axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
@@ -148,7 +138,7 @@ jobs:
           - cuda: 124
             cuda_version: 12.4.1
             python_version: "3.11"
-            pytorch: 2.6.0
+            pytorch: 2.4.1
             axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:

.github/workflows/multi-gpu-e2e.yml

@@ -8,7 +8,6 @@ on:
       - 'setup.py'
       - 'pyproject.toml'
       - '.github/workflows/multi-gpu-e2e.yml'
-      - 'src/axolotl/core/trainers/mixins/sequence_parallel.py'
   workflow_dispatch:
   schedule:
     - cron: '0 0 * * 1,4'  # Runs at 00:00 UTC every monday & thursday
@@ -46,13 +45,6 @@ jobs:
             axolotl_extras: vllm
             num_gpus: 2
             nightly_build: "true"
-          - cuda: 126
-            cuda_version: 12.6.3
-            python_version: "3.11"
-            pytorch: 2.7.0
-            axolotl_extras:
-            num_gpus: 2
-            nightly_build: "true"
     runs-on: [self-hosted, modal]
     timeout-minutes: 120
     steps:
@@ -75,7 +67,6 @@ jobs:
           echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
           echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
           echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
-          echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
       - name: Run tests job on Modal
         run: |
           modal run cicd.multigpu

.github/workflows/tests-nightly.yml

@@ -147,7 +147,6 @@ jobs:
           echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
           echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
           echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
-          echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
       - name: Run tests job on Modal
         run: |
           modal run cicd.e2e_tests

.github/workflows/tests.yml

@@ -49,7 +49,7 @@ jobs:
       max-parallel: 2
       matrix:
         python_version: ["3.11"]
-        pytorch_version: ["2.4.1", "2.5.1", "2.6.0", "2.7.0"]
+        pytorch_version: ["2.4.1", "2.5.1", "2.6.0"]
     timeout-minutes: 20
 
     steps:
@@ -102,17 +102,9 @@ jobs:
 
       - name: Run tests
         run: |
-          pytest -v -n8 --dist loadfile --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli/ tests/ --cov=axolotl --cov-report=xml
-          pytest -v tests/patched/ --cov=axolotl --cov-append --cov-report=xml
-          pytest -v tests/cli/ --cov=axolotl --cov-append --cov-report=xml
-
-      - 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
+          pytest -v -n8 --dist loadfile --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli/ tests/
+          pytest -v tests/patched/
+          pytest -v tests/cli/
 
       - name: cleanup pip cache
         run: |
@@ -242,7 +234,6 @@ jobs:
           echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
           echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV
           echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
-          echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
       - name: Run tests job on Modal
         run: |
           modal run cicd.e2e_tests
@@ -270,12 +261,6 @@ jobs:
             pytorch: 2.5.1
             num_gpus: 1
             axolotl_extras: vllm
-          - cuda: 126
-            cuda_version: 12.6.3
-            python_version: "3.11"
-            pytorch: 2.7.0
-            num_gpus: 1
-            axolotl_extras:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -296,7 +281,6 @@ jobs:
           echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
           echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV
           echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
-          echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
       - name: Run tests job on Modal
         run: |
           modal run cicd.e2e_tests

.runpod/.gitignore

@@ -1,161 +0,0 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/#use-with-ide
-.pdm.toml
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
-pod/scripts/config.yaml

.runpod/Dockerfile

@@ -1,18 +0,0 @@
-FROM runpod/pytorch:3.10-2.0.0-117
-
-COPY .runpod/requirements.txt /requirements.txt
-RUN --mount=type=cache,target=/root/.cache/pip \
-    python3 -m pip install --upgrade pip && \
-    python3 -m pip install --upgrade -r /requirements.txt
-
-
-# Environment settings
-ARG BASE_VOLUME="/runpod-volume"
-ENV BASE_VOLUME=$BASE_VOLUME
-ENV HF_DATASETS_CACHE="${BASE_VOLUME}/huggingface-cache/datasets"
-ENV HUGGINGFACE_HUB_CACHE="${BASE_VOLUME}/huggingface-cache/hub"
-ENV TRANSFORMERS_CACHE="${BASE_VOLUME}/huggingface-cache/hub"
-
-COPY .runpod/src /src
-
-CMD ["python3", "/src/handler.py"]

.runpod/README.md

@@ -1,335 +0,0 @@
-<h1>LLM Post Training- Full fine-tune, LoRA, QLoRa etc. Llama/Mistral/Gemma and more</h1>
-
-# Configuration Options
-
-This document outlines all available configuration options for training models. The configuration can be provided as a JSON request.
-
-## Usage
-
-You can use these configuration Options:
-
-1. As a JSON request body:
-
-```json
-{
-  "input": {
-    "user_id": "user",
-    "model_id": "model-name",
-    "run_id": "run-id",
-    "credentials": {
-      "wandb_api_key": "", # add your Weights & biases key. TODO:  you will be able to set this in Enviornment variables.
-      "hf_token": "", # add your HF_token. TODO:  you will be able to set this in Enviornment variables.
-    },
-    "args": {
-      "base_model": "NousResearch/Llama-3.2-1B",
-      // ... other options
-    }
-  }
-}
-```
-
-## Configuration Options
-
-### Model Configuration
-
-| Option              | Description                                                                                   | Default              |
-| ------------------- | --------------------------------------------------------------------------------------------- | -------------------- |
-| `base_model`        | Path to the base model (local or HuggingFace)                                                 | Required             |
-| `base_model_config` | Configuration path for the base model                                                         | Same as base_model   |
-| `revision_of_model` | Specific model revision from HuggingFace hub                                                  | Latest               |
-| `tokenizer_config`  | Custom tokenizer configuration path                                                           | Optional             |
-| `model_type`        | Type of model to load                                                                         | AutoModelForCausalLM |
-| `tokenizer_type`    | Type of tokenizer to use                                                                      | AutoTokenizer        |
-| `hub_model_id`      | Repository ID where the model will be pushed on Hugging Face Hub (format: username/repo-name) | Optional             |
-
-## Model Family Identification
-
-| Option                     | Default | Description                    |
-| -------------------------- | ------- | ------------------------------ |
-| `is_falcon_derived_model`  | `false` | Whether model is Falcon-based  |
-| `is_llama_derived_model`   | `false` | Whether model is LLaMA-based   |
-| `is_qwen_derived_model`    | `false` | Whether model is Qwen-based    |
-| `is_mistral_derived_model` | `false` | Whether model is Mistral-based |
-
-## Model Configuration Overrides
-
-| Option                                          | Default    | Description                        |
-| ----------------------------------------------- | ---------- | ---------------------------------- |
-| `overrides_of_model_config.rope_scaling.type`   | `"linear"` | RoPE scaling type (linear/dynamic) |
-| `overrides_of_model_config.rope_scaling.factor` | `1.0`      | RoPE scaling factor                |
-
-### Model Loading Options
-
-| Option         | Description                   | Default |
-| -------------- | ----------------------------- | ------- |
-| `load_in_8bit` | Load model in 8-bit precision | false   |
-| `load_in_4bit` | Load model in 4-bit precision | false   |
-| `bf16`         | Use bfloat16 precision        | false   |
-| `fp16`         | Use float16 precision         | false   |
-| `tf32`         | Use tensor float 32 precision | false   |
-
-## Memory and Device Settings
-
-| Option             | Default   | Description             |
-| ------------------ | --------- | ----------------------- |
-| `gpu_memory_limit` | `"20GiB"` | GPU memory limit        |
-| `lora_on_cpu`      | `false`   | Load LoRA on CPU        |
-| `device_map`       | `"auto"`  | Device mapping strategy |
-| `max_memory`       | `null`    | Max memory per device   |
-
-## Training Hyperparameters
-
-| Option                        | Default   | Description                 |
-| ----------------------------- | --------- | --------------------------- |
-| `gradient_accumulation_steps` | `1`       | Gradient accumulation steps |
-| `micro_batch_size`            | `2`       | Batch size per GPU          |
-| `eval_batch_size`             | `null`    | Evaluation batch size       |
-| `num_epochs`                  | `4`       | Number of training epochs   |
-| `warmup_steps`                | `100`     | Warmup steps                |
-| `warmup_ratio`                | `0.05`    | Warmup ratio                |
-| `learning_rate`               | `0.00003` | Learning rate               |
-| `lr_quadratic_warmup`         | `false`   | Quadratic warmup            |
-| `logging_steps`               | `null`    | Logging frequency           |
-| `eval_steps`                  | `null`    | Evaluation frequency        |
-| `evals_per_epoch`             | `null`    | Evaluations per epoch       |
-| `save_strategy`               | `"epoch"` | Checkpoint saving strategy  |
-| `save_steps`                  | `null`    | Saving frequency            |
-| `saves_per_epoch`             | `null`    | Saves per epoch             |
-| `save_total_limit`            | `null`    | Maximum checkpoints to keep |
-| `max_steps`                   | `null`    | Maximum training steps      |
-
-### Dataset Configuration
-
-```yaml
-datasets:
-  - path: vicgalle/alpaca-gpt4 # HuggingFace dataset or TODO: You will be able to add the local path.
-    type: alpaca # Format type (alpaca, gpteacher, oasst, etc.)
-    ds_type: json # Dataset type
-    data_files: path/to/data # Source data files
-    train_on_split: train # Dataset split to use
-```
-
-## Chat Template Settings
-
-| Option                   | Default                          | Description            |
-| ------------------------ | -------------------------------- | ---------------------- |
-| `chat_template`          | `"tokenizer_default"`            | Chat template type     |
-| `chat_template_jinja`    | `null`                           | Custom Jinja template  |
-| `default_system_message` | `"You are a helpful assistant."` | Default system message |
-
-## Dataset Processing
-
-| Option                        | Default                    | Description                       |
-| ----------------------------- | -------------------------- | --------------------------------- |
-| `dataset_prepared_path`       | `"data/last_run_prepared"` | Path for prepared dataset         |
-| `push_dataset_to_hub`         | `""`                       | Push dataset to HF hub            |
-| `dataset_processes`           | `4`                        | Number of preprocessing processes |
-| `dataset_keep_in_memory`      | `false`                    | Keep dataset in memory            |
-| `shuffle_merged_datasets`     | `true`                     | Shuffle merged datasets           |
-| `dataset_exact_deduplication` | `true`                     | Deduplicate datasets              |
-
-## LoRA Configuration
-
-| Option                     | Default                | Description                    |
-| -------------------------- | ---------------------- | ------------------------------ |
-| `adapter`                  | `"lora"`               | Adapter type (lora/qlora)      |
-| `lora_model_dir`           | `""`                   | Directory with pretrained LoRA |
-| `lora_r`                   | `8`                    | LoRA attention dimension       |
-| `lora_alpha`               | `16`                   | LoRA alpha parameter           |
-| `lora_dropout`             | `0.05`                 | LoRA dropout                   |
-| `lora_target_modules`      | `["q_proj", "v_proj"]` | Modules to apply LoRA          |
-| `lora_target_linear`       | `false`                | Target all linear modules      |
-| `peft_layers_to_transform` | `[]`                   | Layers to transform            |
-| `lora_modules_to_save`     | `[]`                   | Modules to save                |
-| `lora_fan_in_fan_out`      | `false`                | Fan in/out structure           |
-
-## Optimization Settings
-
-| Option                    | Default | Description                |
-| ------------------------- | ------- | -------------------------- |
-| `train_on_inputs`         | `false` | Train on input prompts     |
-| `group_by_length`         | `false` | Group by sequence length   |
-| `gradient_checkpointing`  | `false` | Use gradient checkpointing |
-| `early_stopping_patience` | `3`     | Early stopping patience    |
-
-## Learning Rate Scheduling
-
-| Option                     | Default    | Description          |
-| -------------------------- | ---------- | -------------------- |
-| `lr_scheduler`             | `"cosine"` | Scheduler type       |
-| `lr_scheduler_kwargs`      | `{}`       | Scheduler parameters |
-| `cosine_min_lr_ratio`      | `null`     | Minimum LR ratio     |
-| `cosine_constant_lr_ratio` | `null`     | Constant LR ratio    |
-| `lr_div_factor`            | `null`     | LR division factor   |
-
-## Optimizer Settings
-
-| Option                 | Default      | Description         |
-| ---------------------- | ------------ | ------------------- |
-| `optimizer`            | `"adamw_hf"` | Optimizer choice    |
-| `optim_args`           | `{}`         | Optimizer arguments |
-| `optim_target_modules` | `[]`         | Target modules      |
-| `weight_decay`         | `null`       | Weight decay        |
-| `adam_beta1`           | `null`       | Adam beta1          |
-| `adam_beta2`           | `null`       | Adam beta2          |
-| `adam_epsilon`         | `null`       | Adam epsilon        |
-| `max_grad_norm`        | `null`       | Gradient clipping   |
-
-## Attention Implementations
-
-| Option                     | Default | Description                   |
-| -------------------------- | ------- | ----------------------------- |
-| `flash_optimum`            | `false` | Use better transformers       |
-| `xformers_attention`       | `false` | Use xformers                  |
-| `flash_attention`          | `false` | Use flash attention           |
-| `flash_attn_cross_entropy` | `false` | Flash attention cross entropy |
-| `flash_attn_rms_norm`      | `false` | Flash attention RMS norm      |
-| `flash_attn_fuse_qkv`      | `false` | Fuse QKV operations           |
-| `flash_attn_fuse_mlp`      | `false` | Fuse MLP operations           |
-| `sdp_attention`            | `false` | Use scaled dot product        |
-| `s2_attention`             | `false` | Use shifted sparse attention  |
-
-## Tokenizer Modifications
-
-| Option           | Default | Description                  |
-| ---------------- | ------- | ---------------------------- |
-| `special_tokens` | -       | Special tokens to add/modify |
-| `tokens`         | `[]`    | Additional tokens            |
-
-## Distributed Training
-
-| Option                  | Default | Description           |
-| ----------------------- | ------- | --------------------- |
-| `fsdp`                  | `null`  | FSDP configuration    |
-| `fsdp_config`           | `null`  | FSDP config options   |
-| `deepspeed`             | `null`  | Deepspeed config path |
-| `ddp_timeout`           | `null`  | DDP timeout           |
-| `ddp_bucket_cap_mb`     | `null`  | DDP bucket capacity   |
-| `ddp_broadcast_buffers` | `null`  | DDP broadcast buffers |
-
-<details>
-<summary><h3>Example Configuration Request:</h3></summary>
-
-Here's a complete example for fine-tuning a LLaMA model using LoRA:
-
-```json
-{
-  "input": {
-    "user_id": "user",
-    "model_id": "llama-test",
-    "run_id": "test-run",
-    "credentials": {
-      "wandb_api_key": "",
-      "hf_token": ""
-    },
-    "args": {
-      "base_model": "NousResearch/Llama-3.2-1B",
-      "load_in_8bit": false,
-      "load_in_4bit": false,
-      "strict": false,
-      "datasets": [
-        {
-          "path": "teknium/GPT4-LLM-Cleaned",
-          "type": "alpaca"
-        }
-      ],
-      "dataset_prepared_path": "last_run_prepared",
-      "val_set_size": 0.1,
-      "output_dir": "./outputs/lora-out",
-      "adapter": "lora",
-      "sequence_len": 2048,
-      "sample_packing": true,
-      "eval_sample_packing": true,
-      "pad_to_sequence_len": true,
-      "lora_r": 16,
-      "lora_alpha": 32,
-      "lora_dropout": 0.05,
-      "lora_target_modules": [
-        "gate_proj",
-        "down_proj",
-        "up_proj",
-        "q_proj",
-        "v_proj",
-        "k_proj",
-        "o_proj"
-      ],
-      "gradient_accumulation_steps": 2,
-      "micro_batch_size": 2,
-      "num_epochs": 1,
-      "optimizer": "adamw_8bit",
-      "lr_scheduler": "cosine",
-      "learning_rate": 0.0002,
-      "train_on_inputs": false,
-      "group_by_length": false,
-      "bf16": "auto",
-      "tf32": false,
-      "gradient_checkpointing": true,
-      "logging_steps": 1,
-      "flash_attention": true,
-      "loss_watchdog_threshold": 5,
-      "loss_watchdog_patience": 3,
-      "warmup_steps": 10,
-      "evals_per_epoch": 4,
-      "saves_per_epoch": 1,
-      "weight_decay": 0,
-      "hub_model_id": "runpod/llama-fr-lora",
-      "wandb_name": "test-run-1",
-      "wandb_project": "test-run-1",
-      "wandb_entity": "axo-test",
-      "special_tokens": {
-        "pad_token": "<|end_of_text|>"
-      }
-    }
-  }
-}
-```
-
-</details>
-
-### Advanced Features
-
-#### Wandb Integration
-
-- `wandb_project`: Project name for Weights & Biases
-- `wandb_entity`: Team name in W&B
-- `wandb_watch`: Monitor model with W&B
-- `wandb_name`: Name of the W&B run
-- `wandb_run_id`: ID for the W&B run
-
-#### Performance Optimization
-
-- `sample_packing`: Enable efficient sequence packing
-- `eval_sample_packing`: Use sequence packing during evaluation
-- `torch_compile`: Enable PyTorch 2.0 compilation
-- `flash_attention`: Use Flash Attention implementation
-- `xformers_attention`: Use xFormers attention implementation
-
-### Available Optimizers
-
-The following optimizers are supported:
-
-- `adamw_hf`: HuggingFace's AdamW implementation
-- `adamw_torch`: PyTorch's AdamW
-- `adamw_torch_fused`: Fused AdamW implementation
-- `adamw_torch_xla`: XLA-optimized AdamW
-- `adamw_apex_fused`: NVIDIA Apex fused AdamW
-- `adafactor`: Adafactor optimizer
-- `adamw_anyprecision`: Anyprecision AdamW
-- `adamw_bnb_8bit`: 8-bit AdamW from bitsandbytes
-- `lion_8bit`: 8-bit Lion optimizer
-- `lion_32bit`: 32-bit Lion optimizer
-- `sgd`: Stochastic Gradient Descent
-- `adagrad`: Adagrad optimizer
-
-## Notes
-
-- Set `load_in_8bit: true` or `load_in_4bit: true` for memory-efficient training
-- Enable `flash_attention: true` for faster training on modern GPUs
-- Use `gradient_checkpointing: true` to reduce memory usage
-- Adjust `micro_batch_size` and `gradient_accumulation_steps` based on your GPU memory
-
-For more detailed information, please refer to the [documentation](https://axolotl-ai-cloud.github.io/axolotl/docs/config.html).
-
-### Errors:
-
-- if you face any issues with the Flash Attention-2, Delete yoor worker and Re-start.

.runpod/hub.json

@@ -1,93 +0,0 @@
-{
-  "title": "Axolotl Fine-Tuning",
-  "description": "Serverless fine-tuning of open-source LLMs with Axolotl. Supports LoRA, QLoRA, DPO, and more using Hugging Face models and datasets.",
-  "type": "serverless",
-  "category": "language",
-  "iconUrl": "https://avatars.githubusercontent.com/u/167502477",
-  "config": {
-    "runsOn": "GPU",
-    "containerDiskInGb": 200,
-    "gpuCount": 1,
-    "allowedCudaVersions": [
-      "12.8",
-      "12.7",
-      "12.6",
-      "12.5",
-      "12.4"
-    ],
-    "presets": [],
-    "env": [
-      {
-        "key": "TOKENIZER",
-        "input": {
-          "name": "Tokenizer",
-          "type": "string",
-          "description": "Name or path of the Hugging Face tokenizer to use.",
-          "default": "",
-          "advanced": true
-        }
-      },
-      {
-        "key": "MAX_NUM_SEQS",
-        "input": {
-          "name": "Max Num Seqs",
-          "type": "number",
-          "description": "Maximum number of sequences per iteration.",
-          "default": 256,
-          "advanced": true
-        }
-      },
-      {
-        "key": "DISABLE_LOG_STATS",
-        "input": {
-          "name": "Disable Log Stats",
-          "type": "boolean",
-          "description": "Disable logging statistics.",
-          "default": false,
-          "trueValue": "true",
-          "falseValue": "false"
-        }
-      },
-      {
-        "key": "LOAD_FORMAT",
-        "input": {
-          "name": "Load Format",
-          "type": "string",
-          "description": "The format of the model weights to load.",
-          "default": "auto",
-          "options": [
-            {
-              "label": "auto",
-              "value": "auto"
-            },
-            {
-              "label": "pt",
-              "value": "pt"
-            },
-            {
-              "label": "safetensors",
-              "value": "safetensors"
-            },
-            {
-              "label": "npcache",
-              "value": "npcache"
-            },
-            {
-              "label": "dummy",
-              "value": "dummy"
-            },
-            {
-              "label": "tensorizer",
-              "value": "tensorizer"
-            },
-            {
-              "label": "bitsandbytes",
-              "value": "bitsandbytes"
-            }
-          ],
-          "advanced": true
-        }
-      }
-    ]
-  }
-}

.runpod/requirements.txt

@@ -1,15 +0,0 @@
-# Required Python packages get listed here, one per line.
-# Reccomended to lock the version number to avoid unexpected changes.
-
-# You can also install packages from a git repository, e.g.:
-# git+https://github.com/runpod/runpod-python.git
-# To learn more, see https://pip.pypa.io/en/stable/reference/requirements-file-format/
-runpod~=1.7.0
-huggingface_hub
-typing-extensions
-pydantic
-pydantic-settings
-hf-transfer
-setuptools
-numpy==2.0.0
-axolotl[flash-attn,deepspeed]

.runpod/src/config/config.yaml

@@ -1,577 +0,0 @@
-# # This is the huggingface model that contains *.pt, *.safetensors, or *.bin files
-# # This can also be a relative path to a model on disk
-# base_model: ./llama-7b-hf
-# # You can specify an ignore pattern if the model repo contains more than 1 model type (*.pt, etc)
-# base_model_ignore_patterns:
-# # If the base_model repo on hf hub doesn't include configuration .json files,
-# # You can set that here, or leave this empty to default to base_model
-# base_model_config: ./llama-7b-hf
-# # You can specify to choose a specific model revision from huggingface hub
-# model_revision:
-# # Optional tokenizer configuration override in case you want to use a different tokenizer
-# # than the one defined in the base model
-# tokenizer_config:
-# # If you want to specify the type of model to load, AutoModelForCausalLM is a good choice too
-# model_type: AutoModelForCausalLM
-# # Corresponding tokenizer for the model AutoTokenizer is a good choice
-# tokenizer_type: AutoTokenizer
-# # Trust remote code for untrusted source
-# trust_remote_code:
-# # use_fast option for tokenizer loading from_pretrained, default to True
-# tokenizer_use_fast:
-# # Whether to use the legacy tokenizer setting, defaults to True
-# tokenizer_legacy:
-# # Resize the model embeddings when new tokens are added to multiples of 32
-# # This is reported to improve training speed on some models
-# resize_token_embeddings_to_32x:
-
-# # Used to identify which the model is based on
-# 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
-# gptq_groupsize: 128 # group size
-# gptq_model_v1: false # v1 or v2
-
-# # This will attempt to quantize the model down to 8 bits and use adam 8 bit optimizer
-# load_in_8bit: true
-# # Use bitsandbytes 4 bit
-# load_in_4bit:
-
-# # Use CUDA bf16
-# bf16: true # bool or 'full' for `bf16_full_eval`. require >=ampere
-# # Use CUDA fp16
-# fp16: true
-# # Use CUDA tf32
-# tf32: true # require >=ampere
-
-# # No AMP (automatic mixed precision)
-# bfloat16: true # require >=ampere
-# 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
-#   - 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>
-#     ds_type: # Optional[str] (json|arrow|parquet|text|csv) defines the datatype when path is a file
-#     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.
-
-#   # Custom user prompt
-#   - path: repo
-#     type:
-#       # The below are defaults. only set what's needed.
-#       system_prompt: ""
-#       system_format: "{system}"
-#       field_system: system
-#       field_instruction: instruction
-#       field_input: input
-#       field_output: output
-
-#       # Customizable to be single line or multi-line
-#       # 'format' can include {input}
-#       format: |-
-#         User: {instruction} {input}
-#         Assistant:
-#       # 'no_input_format' cannot include {input}
-#       no_input_format: "{instruction} "
-
-#       # For `completion` datsets only, uses the provided field instead of `text` column
-#       field:
-
-# # Axolotl attempts to save the dataset as an arrow after packing the data together so
-# # subsequent training attempts load faster, relative path
-# dataset_prepared_path: data/last_run_prepared
-# # Push prepared dataset to hub
-# push_dataset_to_hub: # repo path
-# # The maximum number of processes to use while preprocessing your input dataset. This defaults to `os.cpu_count()`
-# # if not set.
-# dataset_processes: # defaults to os.cpu_count() if not set
-# # push checkpoints to hub
-# hub_model_id: # repo path to push finetuned model
-# # how to push checkpoints to hub
-# # https://huggingface.co/docs/transformers/v4.31.0/en/main_classes/trainer#transformers.TrainingArguments.hub_strategy
-# hub_strategy:
-# # Whether to use hf `use_auth_token` for loading datasets. Useful for fetching private datasets
-# # Required to be true when used in combination with `push_dataset_to_hub`
-# hf_use_auth_token: # boolean
-# # How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc. 0 for no eval.
-# val_set_size: 0.04
-# # Num shards for whole dataset
-# dataset_shard_num:
-# # Index of shard to use for whole dataset
-# dataset_shard_idx:
-
-# # The maximum length of an input to train with, this should typically be less than 2048
-# # as most models have a token/context limit of 2048
-# sequence_len: 2048
-# # Pad inputs so each step uses constant sized buffers
-# # This will reduce memory fragmentation and may prevent OOMs, by re-using memory more efficiently
-# pad_to_sequence_len:
-# # Max sequence length to concatenate training samples together up to
-# # Inspired by StackLLaMA. see https://huggingface.co/blog/stackllama#supervised-fine-tuning
-# # FutureWarning: This will soon be DEPRECATED
-# max_packed_sequence_len: 1024
-# # Use efficient multi-packing with block diagonal attention and per sequence position_ids. Recommend set to 'true'
-# sample_packing:
-# # Set to 'false' if getting errors during eval with sample_packing on.
-# eval_sample_packing:
-# # You can set these packing optimizations AFTER starting a training at least once.
-# # The trainer will provide recommended values for these values.
-# sample_packing_eff_est:
-# total_num_tokens:
-
-# # 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:
-
-# # 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:
-#   - q_proj
-#   - v_proj
-# #  - k_proj
-# #  - o_proj
-# #  - gate_proj
-# #  - down_proj
-# #  - up_proj
-# lora_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:
-# #  - embed_tokens
-# #  - lm_head
-
-# # Once you complete training, the model will be saved to the following directory.
-# # 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
-
-# # ReLoRA configuration
-# # Must use either 'lora' or 'qlora' adapter, and does not support fsdp or deepspeed
-# relora_steps: # Number of steps per ReLoRA restart
-# relora_warmup_steps: # Number of per-restart warmup steps
-# relora_cpu_offload: # True to perform lora weight merges on cpu during restarts, for modest gpu memory savings
-
-# # wandb configuration if you're using it
-# wandb_mode: # "offline" to save run metadata locally and not sync to the server, "disabled" to turn off wandb
-# wandb_project: # Your wandb project name
-# wandb_entity: # A wandb Team name if using a Team
-# wandb_watch:
-# 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
-# output_dir: ./completed-model
-
-# # Whether to use torch.compile and which backend to use
-# torch_compile:  # bool
-# torch_compile_backend:  # Optional[str]
-
-# # Training hyperparameters
-
-# # If greater than 1, backpropagation will be skipped and the gradients will be accumulated for the given number of steps.
-# 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
-# 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
-# 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.
-# # e.g., when 1 epoch is 1000 steps => `num_epochs: 2` and `max_steps: 100` will train for 100 steps
-# 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
-
-# # Save model as safetensors (require safetensors package)
-# save_safetensors:
-
-# # Whether to mask out or include the human's prompt from the training labels
-# train_on_inputs: false
-# # Group similarly sized data to minimize padding.
-# # May be slower to start, as it must download and sort the entire dataset.
-# # Note that training loss may have an oscillating pattern with this enabled.
-# group_by_length: false
-
-# # Whether to use gradient checkpointing https://huggingface.co/docs/transformers/v4.18.0/en/performance#gradient-checkpointing
-# gradient_checkpointing: false
-
-# # Stop training after this many evaluation losses have increased in a row
-# # https://huggingface.co/transformers/v4.2.2/_modules/transformers/trainer_callback.html#EarlyStoppingCallback
-# early_stopping_patience: 3
-
-# # Specify a scheduler and kwargs to use with the optimizer
-# lr_scheduler: # 'one_cycle' | 'log_sweep' | empty for cosine
-# lr_scheduler_kwargs:
-
-# # For one_cycle optim
-# lr_div_factor: # Learning rate div factor
-
-# # For log_sweep optim
-# log_sweep_min_lr:
-# log_sweep_max_lr:
-
-# # Specify optimizer
-# # Valid values are driven by the Transformers OptimizerNames class, see:
-# # https://github.com/huggingface/transformers/blob/95b374952dc27d8511541d6f5a4e22c9ec11fb24/src/transformers/training_args.py#L134
-# #
-# # Note that not all optimizers may be available in your environment, ex: 'adamw_anyprecision' is part of
-# # torchdistx, 'adamw_bnb_8bit' is part of bnb.optim.Adam8bit, etc. When in doubt, it is recommended to start with the optimizer used
-# # in the examples/ for your model and fine-tuning use case.
-# #
-# # Valid values for 'optimizer' include:
-# # - adamw_hf
-# # - adamw_torch
-# # - adamw_torch_fused
-# # - adamw_torch_xla
-# # - adamw_apex_fused
-# # - adafactor
-# # - adamw_anyprecision
-# # - sgd
-# # - adagrad
-# # - adamw_bnb_8bit
-# # - lion_8bit
-# # - lion_32bit
-# # - paged_adamw_32bit
-# # - paged_adamw_8bit
-# # - paged_lion_32bit
-# # - paged_lion_8bit
-# optimizer:
-# # Specify weight decay
-# weight_decay:
-# # adamw hyperparams
-# adam_beta1:
-# adam_beta2:
-# adam_epsilon:
-# # Gradient clipping max norm
-# max_grad_norm:
-
-# # Augmentation techniques
-# # NEFT https://arxiv.org/abs/2310.05914, set this to a number (paper default is 5) to add noise to embeddings
-# # currently only supported on Llama and Mistral
-# noisy_embedding_alpha:
-
-# # Whether to bettertransformers
-# flash_optimum:
-# # Whether to use xformers attention patch https://github.com/facebookresearch/xformers:
-# xformers_attention:
-# # Whether to use flash attention patch https://github.com/Dao-AILab/flash-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:
-# # Landmark attention (only llama)
-# landmark_attention:
-# # xpos RoPE see https://github.com/kaiokendev/cutoff-len-is-context-len/blob/main/util/xpos_rope_llama_monkey_patch.py
-# # LLaMA only
-# xpos_rope:
-
-# # Resume from a specific checkpoint dir
-# resume_from_checkpoint:
-# # If resume_from_checkpoint isn't set and you simply want it to start where it left off.
-# # Be careful with this being turned on between different models.
-# auto_resume_from_checkpoints: false
-
-# # Don't mess with this, it's here for accelerate and torchrun
-# local_rank:
-
-# # Add or change special tokens.
-# # If you add tokens here, you don't need to add them to the `tokens` list.
-# special_tokens:
-#   # bos_token: "<s>"
-#   # eos_token: "</s>"
-#   # unk_token: "<unk>"
-
-# # Add extra tokens.
-# tokens:
-
-# # FSDP
-# fsdp:
-# fsdp_config:
-
-# # Deepspeed config path. e.g., deepspeed/zero3.json
-# deepspeed:
-
-# # Advanced DDP Arguments
-# ddp_timeout:
-# ddp_bucket_cap_mb:
-# ddp_broadcast_buffers:
-
-# # Path to torch distx for optim 'adamw_anyprecision'
-# torchdistx_path:
-
-# # Set to HF dataset for type: 'completion' for streaming instead of pre-tokenize
-# pretraining_dataset:
-
-# # Debug mode
-# debug:
-
-# # Seed
-# seed:
-
-# # Allow overwrite yml config using from cli
-# strict:
-
-
-
-base_model: ${BASE_MODEL}
-base_model_ignore_patterns: ${BASE_MODEL_IGNORE_PATTERNS}
-base_model_config: ${BASE_MODEL_CONFIG}
-revision_of_model: ${REVISION_OF_MODEL}
-tokenizer_config: ${TOKENIZER_CONFIG}
-model_type: ${MODEL_TYPE}
-tokenizer_type: ${TOKENIZER_TYPE}
-trust_remote_code: ${TRUST_REMOTE_CODE}
-tokenizer_use_fast: ${TOKENIZER_USE_FAST}
-tokenizer_legacy: ${TOKENIZER_LEGACY}
-resize_token_embeddings_to_32x: ${RESIZE_TOKEN_EMBEDDINGS_TO_32X}
-
-is_falcon_derived_model: ${IS_FALCON_DERIVED_MODEL}
-is_llama_derived_model: ${IS_LLAMA_DERIVED_MODEL}
-is_qwen_derived_model: ${IS_QWEN_DERIVED_MODEL}
-is_mistral_derived_model: ${IS_MISTRAL_DERIVED_MODEL}
-
-overrides_of_model_config:
-  rope_scaling:
-    type: ${ROPE_SCALING_TYPE}
-    factor: ${ROPE_SCALING_FACTOR}
-
-bnb_config_kwargs:
-  llm_int8_has_fp16_weight: ${BNB_LLM_INT8_HAS_FP16_WEIGHT}
-  bnb_4bit_quant_type: ${BNB_4BIT_QUANT_TYPE}
-  bnb_4bit_use_double_quant: ${BNB_4BIT_USE_DOUBLE_QUANT}
-
-gptq: ${GPTQ}
-load_in_8bit: ${LOAD_IN_8BIT}
-load_in_4bit: ${LOAD_IN_4BIT}
-bf16: ${BF16}
-fp16: ${FP16}
-tf32: ${TF32}
-bfloat16: ${BFLOAT16}
-float16: ${FLOAT16}
-
-gpu_memory_limit: ${GPU_MEMORY_LIMIT}
-lora_on_cpu: ${LORA_ON_CPU}
-
-datasets:
-  - path: ${DATASET_PATH}
-    type: ${DATASET_TYPE}
-    ds_type: ${DATASET_DS_TYPE}
-    data_files: ${DATASET_DATA_FILES}
-    shards: ${DATASET_SHARDS}
-    name: ${DATASET_NAME}
-    train_on_split: ${DATASET_TRAIN_ON_SPLIT}
-    revision: ${DATASET_REVISION}
-    trust_remote_code: ${DATASET_TRUST_REMOTE_CODE}
-
-rl: ${RL}
-dpo_use_weighting: ${DPO_USE_WEIGHTING}
-
-chat_template: ${CHAT_TEMPLATE}
-chat_template_jinja: ${CHAT_TEMPLATE_JINJA}
-default_system_message: ${DEFAULT_SYSTEM_MESSAGE}
-dataset_prepared_path: ${DATASET_PREPARED_PATH}
-push_dataset_to_hub: ${PUSH_DATASET_TO_HUB}
-dataset_processes: ${DATASET_PROCESSES}
-dataset_keep_in_memory: ${DATASET_KEEP_IN_MEMORY}
-hub_model_id: ${HUB_MODEL_ID}
-hub_strategy: ${HUB_STRATEGY}
-hf_use_auth_token: ${HF_USE_AUTH_TOKEN}
-val_set_size: ${VAL_SET_SIZE}
-dataset_shard_num: ${DATASET_SHARD_NUM}
-dataset_shard_idx: ${DATASET_SHARD_IDX}
-
-sequence_len: ${SEQUENCE_LEN}
-pad_to_sequence_len: ${PAD_TO_SEQUENCE_LEN}
-sample_packing: ${SAMPLE_PACKING}
-eval_sample_packing: ${EVAL_SAMPLE_PACKING}
-sample_packing_eff_est: ${SAMPLE_PACKING_EFF_EST}
-total_num_tokens: ${TOTAL_NUM_TOKENS}
-sample_packing_group_size: ${SAMPLE_PACKING_GROUP_SIZE}
-sample_packing_bin_size: ${SAMPLE_PACKING_BIN_SIZE}
-
-batch_flattening: ${BATCH_FLATTENING}
-device_map: ${DEVICE_MAP}
-max_memory: ${MAX_MEMORY}
-
-adapter: ${ADAPTER}
-lora_model_dir: ${LORA_MODEL_DIR}
-
-lora_r: ${LORA_R}
-lora_alpha: ${LORA_ALPHA}
-lora_dropout: ${LORA_DROPOUT}
-lora_target_modules:
-  - ${LORA_TARGET_MODULES}
-lora_target_linear: ${LORA_TARGET_LINEAR}
-peft_layers_to_transform: ${PEFT_LAYERS_TO_TRANSFORM}
-lora_modules_to_save: ${LORA_MODULES_TO_SAVE}
-lora_fan_in_fan_out: ${LORA_FAN_IN_FAN_OUT}
-
-loraplus_lr_ratio: ${LORAPLUS_LR_RATIO}
-loraplus_lr_embedding: ${LORAPLUS_LR_EMBEDDING}
-
-peft:
-  loftq_config:
-    loftq_bits: ${LOFTQ_BITS}
-
-relora_steps: ${RELORA_STEPS}
-relora_warmup_steps: ${RELORA_WARMUP_STEPS}
-relora_anneal_steps: ${RELORA_ANNEAL_STEPS}
-relora_prune_ratio: ${RELORA_PRUNE_RATIO}
-relora_cpu_offload: ${RELORA_CPU_OFFLOAD}
-
-wandb_mode: ${WANDB_MODE}
-wandb_project: ${WANDB_PROJECT}
-wandb_entity: ${WANDB_ENTITY}
-wandb_watch: ${WANDB_WATCH}
-wandb_name: ${WANDB_NAME}
-wandb_run_id: ${WANDB_RUN_ID}
-wandb_log_model: ${WANDB_LOG_MODEL}
-
-mlflow_tracking_uri: ${MLFLOW_TRACKING_URI}
-mlflow_experiment_name: ${MLFLOW_EXPERIMENT_NAME}
-mlflow_run_name: ${MLFLOW_RUN_NAME}
-hf_mlflow_log_artifacts: ${HF_MLFLOW_LOG_ARTIFACTS}
-
-use_comet: ${USE_COMET}
-comet_api_key: ${COMET_API_KEY}
-comet_workspace: ${COMET_WORKSPACE}
-comet_project_name: ${COMET_PROJECT_NAME}
-comet_experiment_key: ${COMET_EXPERIMENT_KEY}
-comet_mode: ${COMET_MODE}
-comet_online: ${COMET_ONLINE}
-comet_experiment_config: ${COMET_EXPERIMENT_CONFIG}
-
-output_dir: ${OUTPUT_DIR}
-
-torch_compile: ${TORCH_COMPILE}
-torch_compile_backend: ${TORCH_COMPILE_BACKEND}
-
-gradient_accumulation_steps: ${GRADIENT_ACCUMULATION_STEPS}
-micro_batch_size: ${MICRO_BATCH_SIZE}
-eval_batch_size: ${EVAL_BATCH_SIZE}
-num_epochs: ${NUM_EPOCHS}
-warmup_steps: ${WARMUP_STEPS}
-warmup_ratio: ${WARMUP_RATIO}
-learning_rate: ${LEARNING_RATE}
-lr_quadratic_warmup: ${LR_QUADRATIC_WARMUP}
-logging_steps: ${LOGGING_STEPS}
-eval_steps: ${EVAL_STEPS}
-evals_per_epoch: ${EVALS_PER_EPOCH}
-save_strategy: ${SAVE_STRATEGY}
-save_steps: ${SAVE_STEPS}
-saves_per_epoch: ${SAVES_PER_EPOCH}
-save_total_limit: ${SAVE_TOTAL_LIMIT}
-max_steps: ${MAX_STEPS}
-
-eval_table_size: ${EVAL_TABLE_SIZE}
-eval_max_new_tokens: ${EVAL_MAX_NEW_TOKENS}
-eval_causal_lm_metrics: ${EVAL_CAUSAL_LM_METRICS}
-
-profiler_steps: ${PROFILER_STEPS}
-loss_watchdog_threshold: ${LOSS_WATCHDOG_THRESHOLD}
-loss_watchdog_patience: ${LOSS_WATCHDOG_PATIENCE}
-
-save_safetensors: ${SAVE_SAFETENSORS}
-train_on_inputs: ${TRAIN_ON_INPUTS}
-group_by_length: ${GROUP_BY_LENGTH}
-gradient_checkpointing: ${GRADIENT_CHECKPOINTING}
-early_stopping_patience: ${EARLY_STOPPING_PATIENCE}
-
-lr_scheduler: ${LR_SCHEDULER}
-lr_scheduler_kwargs: ${LR_SCHEDULER_KWARGS}
-cosine_min_lr_ratio: ${COSINE_MIN_LR_RATIO}
-cosine_constant_lr_ratio: ${COSINE_CONSTANT_LR_RATIO}
-lr_div_factor: ${LR_DIV_FACTOR}
-
-optimizer: ${OPTIMIZER}
-optim_args: ${OPTIM_ARGS}
-optim_target_modules: ${OPTIM_TARGET_MODULES}
-weight_decay: ${WEIGHT_DECAY}
-adam_beta1: ${ADAM_BETA1}
-adam_beta2: ${ADAM_BETA2}
-adam_epsilon: ${ADAM_EPSILON}
-max_grad_norm: ${MAX_GRAD_NORM}
-
-neftune_noise_alpha: ${NEFTUNE_NOISE_ALPHA}
-
-flash_optimum: ${FLASH_OPTIMUM}
-xformers_attention: ${XFORMERS_ATTENTION}
-flash_attention: ${FLASH_ATTENTION}
-flash_attn_cross_entropy: ${FLASH_ATTN_CROSS_ENTROPY}
-flash_attn_rms_norm: ${FLASH_ATTN_RMS_NORM}
-flash_attn_fuse_qkv: ${FLASH_ATTN_FUSE_QKV}
-flash_attn_fuse_mlp: ${FLASH_ATTN_FUSE_MLP}
-sdp_attention: ${SDP_ATTENTION}
-s2_attention: ${S2_ATTENTION}
-resume_from_checkpoint: ${RESUME_FROM_CHECKPOINT}
-auto_resume_from_checkpoints: ${AUTO_RESUME_FROM_CHECKPOINTS}
-
-local_rank: ${LOCAL_RANK}
-
-special_tokens:
-  bos_token: ${SPECIAL_TOKEN_BOS}
-  eos_token: ${SPECIAL_TOKEN_EOS}
-  unk_token: ${SPECIAL_TOKEN_UNK}
-  pad_token: ${SPECIAL_TOKEN_PAD}
-
-tokens: ${TOKENS}
-
-fsdp: ${FSDP}
-fsdp_config: ${FSDP_CONFIG}
-deepspeed: ${DEEPSPEED}
-
-ddp_timeout: ${DDP_TIMEOUT}
-ddp_bucket_cap_mb: ${DDP_BUCKET_CAP_MB}
-ddp_broadcast_buffers: ${DDP_BROADCAST_BUFFERS}
-
-torchdistx_path: ${TORCHDISTX_PATH}
-pretraining_dataset: ${PRETRAINING_DATASET}
-debug: ${DEBUG}
-seed: ${SEED}
-strict: ${STRICT}

.runpod/src/handler.py

@@ -1,64 +0,0 @@
-"""
-Runpod serverless entrypoint handler
-"""
-
-import os
-
-import runpod
-import yaml
-from huggingface_hub._login import login
-from train import train
-from utils import get_output_dir
-
-BASE_VOLUME = os.environ.get("BASE_VOLUME", "/runpod-volume")
-if not os.path.exists(BASE_VOLUME):
-    os.makedirs(BASE_VOLUME)
-
-logger = runpod.RunPodLogger()
-
-
-async def handler(job):
-    runpod_job_id = job["id"]
-    inputs = job["input"]
-    run_id = inputs.get("run_id", "default_run_id")
-    args = inputs.get("args", {})
-
-    # Set output directory
-    output_dir = os.path.join(BASE_VOLUME, get_output_dir(run_id))
-    args["output_dir"] = output_dir
-
-    # First save args to a temporary config file
-    config_path = "/workspace/test_config.yaml"
-
-    # Add run_name and job_id to args before saving
-    args["run_name"] = run_id
-    args["runpod_job_id"] = runpod_job_id
-
-    yaml_data = yaml.dump(args, default_flow_style=False)
-    with open(config_path, "w", encoding="utf-8") as file:
-        file.write(yaml_data)
-
-    # Handle credentials
-    credentials = inputs.get("credentials", {})
-
-    if "wandb_api_key" in credentials:
-        os.environ["WANDB_API_KEY"] = credentials["wandb_api_key"]
-    if "hf_token" in credentials:
-        os.environ["HF_TOKEN"] = credentials["hf_token"]
-
-    if os.environ.get("HF_TOKEN"):
-        login(token=os.environ["HF_TOKEN"])
-    else:
-        logger.info("No HF_TOKEN provided. Skipping login.")
-
-    logger.info("Starting Training.")
-    async for result in train(config_path):  # Pass the config path instead of args
-        logger.info(result)
-    logger.info("Training Complete.")
-
-    # Cleanup
-    del os.environ["WANDB_API_KEY"]
-    del os.environ["HF_TOKEN"]
-
-
-runpod.serverless.start({"handler": handler, "return_aggregate_stream": True})

.runpod/src/test_input.json

@@ -1,61 +0,0 @@
-{
-  "input": {
-    "user_id": "user",
-    "model_id": "llama-test",
-    "run_id": "llama-test",
-    "credentials": {
-      "wandb_api_key": "",
-      "hf_token": ""
-    },
-    "args": {
-      "base_model": "NousResearch/Meta-Llama-3-8B",
-      "model_type": "LlamaForCausalLM",
-      "tokenizer_type": "AutoTokenizer",
-      "load_in_8bit": true,
-      "load_in_4bit": false,
-      "strict": false,
-      "datasets": [
-        {
-          "path": "mhenrichsen/alpaca_2k_test",
-          "type": "alpaca"
-        }
-      ],
-      "val_set_size": 0.05,
-      "output_dir": "./outputs/lora-out",
-      "sequence_len": 4096,
-      "sample_packing": true,
-      "eval_sample_packing": false,
-      "pad_to_sequence_len": true,
-      "adapter": "lora",
-      "lora_r": 32,
-      "lora_alpha": 16,
-      "lora_dropout": 0.05,
-      "lora_target_linear": true,
-      "lora_modules_to_save": [
-        "embed_tokens",
-        "lm_head"
-      ],
-      "gradient_accumulation_steps": 4,
-      "micro_batch_size": 2,
-      "num_epochs": 1,
-      "optimizer": "adamw_bnb_8bit",
-      "lr_scheduler": "cosine",
-      "learning_rate": 0.0002,
-      "train_on_inputs": false,
-      "group_by_length": false,
-      "bf16": "auto",
-      "tf32": false,
-      "gradient_checkpointing": true,
-      "logging_steps": 1,
-      "flash_attention": true,
-      "warmup_steps": 1,
-      "evals_per_epoch": 1,
-      "eval_max_new_tokens": 128,
-      "saves_per_epoch": 1,
-      "weight_decay": 0.0,
-      "special_tokens": {
-        "pad_token": "<|end_of_text|>"
-      }
-    }
-  }
-}

.runpod/src/train.py

@@ -1,45 +0,0 @@
-"""
-Runpod train entrypoint
-"""
-
-import asyncio
-
-
-async def train(config_path: str, gpu_id: str = "0", preprocess: bool = True):
-    """
-    Run preprocessing (if enabled) and training with the given config file
-    :param config_path: Path to the YAML config file
-    :param gpu_id: GPU ID to use (default: "0")
-    :param preprocess: Whether to run preprocessing (default: True)
-
-    """
-    # First check if preprocessing is needed
-    if preprocess:
-        # Preprocess command
-        preprocess_cmd = (
-            f"CUDA_VISIBLE_DEVICES={gpu_id} axolotl preprocess {config_path}"
-        )
-        process = await asyncio.create_subprocess_shell(
-            preprocess_cmd,
-            stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.STDOUT,
-        )
-
-        if process.stdout is not None:
-            async for line in process.stdout:
-                yield f"Preprocessing: {line.decode().strip()}"
-        await process.wait()
-        yield "Preprocessing completed."
-    else:
-        yield "Skipping preprocessing step."
-
-    # Training command
-    train_cmd = f"axolotl train {config_path}"
-    process = await asyncio.create_subprocess_shell(
-        train_cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT
-    )
-
-    if process.stdout is not None:
-        async for line in process.stdout:
-            yield f"Training: {line.decode().strip()}"
-    await process.wait()

.runpod/src/utils.py

@@ -1,89 +0,0 @@
-"""
-Runpod launcher utils
-"""
-
-import os
-
-import yaml
-
-
-def get_output_dir(run_id):
-    path = f"fine-tuning/{run_id}"
-    return path
-
-
-def make_valid_config(input_args):
-    """
-    Creates and saves updated config file, returns the path to the new config
-    :param input_args: dict of input args
-    :return: str, path to the updated config file
-    """
-    # Load default config
-    with open("config/config.yaml", "r", encoding="utf-8") as fin:
-        all_args = yaml.safe_load(fin)
-
-    if not input_args:
-        print("No args provided, using defaults")
-    else:
-        all_args.update(input_args)
-
-    # Create updated config path
-    updated_config_path = "config/updated_config.yaml"
-
-    # Save updated config to new file
-    with open(updated_config_path, "w", encoding="utf-8") as f:
-        yaml.dump(all_args, f)
-
-    return updated_config_path
-
-
-def set_config_env_vars(args: dict):
-    """
-    Convert API arguments into environment variables.
-    Handles nested dictionaries, lists, and special values.
-
-    Args:
-        args (dict): The arguments dictionary from the API request
-    """
-
-    def process_value(value):
-        """Convert Python values to string format for environment variables"""
-        if value is None:
-            return ""
-        if isinstance(value, bool):
-            return str(value).lower()
-        if isinstance(value, (list, dict)):
-            return str(value)
-        return str(value)
-
-    def set_env_vars(data, prefix=""):
-        """Recursively set environment variables from nested dictionary"""
-        for key, value in data.items():
-            env_key = prefix + key.upper()
-
-            # Handle special cases
-            if isinstance(value, dict):
-                # For nested dictionaries (like special_tokens)
-                set_env_vars(value, f"{env_key}_")
-            elif isinstance(value, list):
-                # Handle list of dictionaries (like datasets)
-                if value and isinstance(value[0], dict):
-                    for i, item in enumerate(value):
-                        set_env_vars(item, f"{env_key}_{i}_")
-                else:
-                    # For simple lists (like lora_target_modules)
-                    os.environ[env_key] = process_value(value)
-            else:
-                # Handle all other cases
-                os.environ[env_key] = process_value(value)
-
-    # Clear any existing related environment variables
-    # This prevents old values from persisting
-    for key in list(os.environ.keys()):
-        if key.startswith(
-            ("BASE_MODEL", "MODEL_TYPE", "TOKENIZER_TYPE", "DATASET", "LORA_", "WANDB_")
-        ):
-            del os.environ[key]
-
-    # Set new environment variables
-    set_env_vars(args)

.runpod/tests.json

@@ -1,89 +0,0 @@
-{
-  "tests": [
-    {
-      "name": "quick_smoke_test_sft",
-      "input": {
-        "user_id": "user",
-        "model_id": "llama-test",
-        "run_id": "llama-test",
-        "credentials": {
-          "wandb_api_key": "",
-          "hf_token": ""
-        },
-        "args": {
-          "base_model": "NousResearch/Meta-Llama-3-8B",
-          "model_type": "LlamaForCausalLM",
-          "tokenizer_type": "AutoTokenizer",
-          "load_in_8bit": true,
-          "load_in_4bit": false,
-          "strict": false,
-          "datasets": [
-            {
-              "path": "mhenrichsen/alpaca_2k_test",
-              "type": "alpaca"
-            }
-          ],
-          "val_set_size": 0.05,
-          "output_dir": "./outputs/lora-out",
-          "sequence_len": 4096,
-          "sample_packing": true,
-          "eval_sample_packing": false,
-          "pad_to_sequence_len": true,
-          "adapter": "lora",
-          "lora_r": 32,
-          "lora_alpha": 16,
-          "lora_dropout": 0.05,
-          "lora_target_linear": true,
-          "lora_modules_to_save": [
-            "embed_tokens",
-            "lm_head"
-          ],
-          "gradient_accumulation_steps": 4,
-          "micro_batch_size": 2,
-          "num_epochs": 1,
-          "optimizer": "adamw_bnb_8bit",
-          "lr_scheduler": "cosine",
-          "learning_rate": 0.0002,
-          "train_on_inputs": false,
-          "group_by_length": false,
-          "bf16": "auto",
-          "tf32": false,
-          "gradient_checkpointing": true,
-          "logging_steps": 1,
-          "flash_attention": true,
-          "warmup_steps": 1,
-          "evals_per_epoch": 1,
-          "eval_max_new_tokens": 128,
-          "saves_per_epoch": 1,
-          "weight_decay": 0.0,
-          "special_tokens": {
-            "pad_token": "<|end_of_text|>"
-          }
-        }
-      },
-      "timeout": 100000
-    }
-  ],
-  "config": {
-    "gpuTypeId": "NVIDIA GeForce RTX 4090",
-    "gpuCount": 1,
-    "containerDiskInGb": 200,
-    "env": [
-      {
-        "key": "TOKENIZER",
-        "value": ""
-      },
-      {
-        "key": "DISABLE_LOG_STATS",
-        "value": "true"
-      }
-    ],
-    "allowedCudaVersions": [
-      "12.8",
-      "12.7",
-      "12.6",
-      "12.5",
-      "12.4"
-    ]
-  }
-}

README.md

@@ -9,7 +9,6 @@
 <p align="center">
     <img src="https://img.shields.io/github/license/axolotl-ai-cloud/axolotl.svg?color=blue" alt="GitHub License">
     <img src="https://github.com/axolotl-ai-cloud/axolotl/actions/workflows/tests.yml/badge.svg" alt="tests">
-    <a href="https://codecov.io/gh/axolotl-ai-cloud/axolotl"><img src="https://codecov.io/gh/axolotl-ai-cloud/axolotl/branch/main/graph/badge.svg" alt="codecov"></a>
     <a href="https://github.com/axolotl-ai-cloud/axolotl/releases"><img src="https://img.shields.io/github/release/axolotl-ai-cloud/axolotl.svg" alt="Releases"></a>
     <br/>
     <a href="https://github.com/axolotl-ai-cloud/axolotl/graphs/contributors"><img src="https://img.shields.io/github/contributors-anon/axolotl-ai-cloud/axolotl?color=yellow&style=flat-square" alt="contributors" style="height: 20px;"></a>

cicd/cicd.sh

@@ -3,53 +3,10 @@ set -e
 
 python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__"
 
-# Run unit tests with initial coverage report
-pytest -v --durations=10 -n8 \
-  --ignore=tests/e2e/ \
-  --ignore=tests/patched/ \
-  --ignore=tests/cli \
-  /workspace/axolotl/tests/ \
-  --cov=axolotl
-
-# Run lora kernels tests with coverage append
-pytest -v --durations=10 \
-  /workspace/axolotl/tests/e2e/patched/lora_kernels \
-  --cov=axolotl \
-  --cov-append
-
-# Run patched tests excluding lora kernels with coverage append
-pytest -v --durations=10 \
-  --ignore=tests/e2e/patched/lora_kernels \
-  /workspace/axolotl/tests/e2e/patched \
-  --cov=axolotl \
-  --cov-append
-
-# Run solo tests with coverage append
-pytest -v --durations=10 -n1 \
-  /workspace/axolotl/tests/e2e/solo/ \
-  --cov=axolotl \
-  --cov-append
-
-# Run integration tests with coverage append
-pytest -v --durations=10 \
-  /workspace/axolotl/tests/e2e/integrations/ \
-  --cov=axolotl \
-  --cov-append
-
-pytest -v --durations=10 /workspace/axolotl/tests/cli \
-  --cov=axolotl \
-  --cov-append
-
-# Run remaining e2e tests with coverage append and final report
-pytest -v --durations=10 \
-  --ignore=tests/e2e/solo/ \
-  --ignore=tests/e2e/patched/ \
-  --ignore=tests/e2e/multigpu/ \
-  --ignore=tests/e2e/integrations/ \
-  --ignore=tests/cli \
-  /workspace/axolotl/tests/e2e/ \
-  --cov=axolotl \
-  --cov-append \
-  --cov-report=xml:e2e-coverage.xml
-
-codecov upload-process -t $CODECOV_TOKEN -f e2e-coverage.xml -F e2e,pytorch-${PYTORCH_VERSION} || true
+pytest -v --durations=10 -n8 --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli /workspace/axolotl/tests/
+pytest -v --durations=10 /workspace/axolotl/tests/e2e/patched/lora_kernels  # running these with the other patches causes a failure
+pytest -v --durations=10 --ignore=tests/e2e/patched/lora_kernels /workspace/axolotl/tests/e2e/patched
+pytest -v --durations=10 -n1 /workspace/axolotl/tests/e2e/solo/
+pytest -v --durations=10 /workspace/axolotl/tests/e2e/integrations/
+pytest -v --durations=10 /workspace/axolotl/tests/cli
+pytest -v --durations=10 --ignore=tests/e2e/solo/ --ignore=tests/e2e/patched/ --ignore=tests/e2e/multigpu/ --ignore=tests/e2e/integrations/ --ignore=tests/cli /workspace/axolotl/tests/e2e/

cicd/e2e_tests.py

@@ -28,7 +28,6 @@ df_args = {
     "GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"),
     "GITHUB_SHA": os.environ.get("GITHUB_SHA", ""),
     "NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""),
-    "CODECOV_TOKEN": os.environ.get("CODECOV_TOKEN", ""),
     "HF_HOME": "/workspace/data/huggingface-cache/hub",
 }
 

cicd/multigpu.py

@@ -29,7 +29,6 @@ df_args = {
     "CUDA": os.environ.get("CUDA", "121"),
     "GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"),
     "GITHUB_SHA": os.environ.get("GITHUB_SHA", ""),
-    "CODECOV_TOKEN": os.environ.get("CODECOV_TOKEN", ""),
     "HF_HOME": "/workspace/data/huggingface-cache/hub",
 }
 

cicd/multigpu.sh

@@ -1,23 +1,6 @@
 #!/bin/bash
 set -e
 
-# Only run two tests at a time to avoid OOM on GPU (with coverage collection)
-pytest -v -n2 \
-  --ignore=/workspace/axolotl/tests/e2e/multigpu/solo/ \
-  --ignore=/workspace/axolotl/tests/e2e/multigpu/patched/ \
-  /workspace/axolotl/tests/e2e/multigpu/ \
-  --cov=axolotl
-
-# Run solo tests with coverage append
-pytest -v --durations=10 -n1 \
-  /workspace/axolotl/tests/e2e/multigpu/solo/ \
-  --cov=axolotl \
-  --cov-append
-
-pytest -v  --durations=10 -n1 /workspace/axolotl/tests/e2e/multigpu/patched/ \
-  --cov=axolotl \
-  --cov-append \
-  --cov-report=xml:multigpu-coverage.xml
-
-# Upload coverage to Codecov
-codecov upload-process -t $CODECOV_TOKEN -f multigpu-coverage.xml -F multigpu,docker-tests,pytorch-${PYTORCH_VERSION}
+# only run one test at a time so as not to OOM the GPU
+pytest -v  --durations=10 -n2 /workspace/axolotl/tests/e2e/multigpu/ --ignore=/workspace/axolotl/tests/e2e/multigpu/solo/
+pytest -v  --durations=10 -n1 /workspace/axolotl/tests/e2e/multigpu/solo/

codecov.yml

@@ -1,56 +0,0 @@
-codecov:
-  require_ci_to_pass: yes
-  notify:
-    wait_for_ci: true
-
-coverage:
-  precision: 2
-  round: down
-  range: "70...100"
-  status:
-    project:
-      default:
-        # basic
-        target: auto
-        threshold: 0%
-        base: auto
-        # advanced
-        branches: null
-        if_no_uploads: error
-        if_not_found: success
-        if_ci_failed: error
-        only_pulls: false
-        flags: null
-        paths: null
-    patch:
-      default:
-        # basic
-        target: auto
-        threshold: 0%
-        base: auto
-        # advanced
-        branches: null
-        if_no_uploads: error
-        if_not_found: success
-        if_ci_failed: error
-        only_pulls: false
-        flags: null
-        paths: null
-
-parsers:
-  gcov:
-    branch_detection:
-      conditional: yes
-      loop: yes
-      method: no
-      macro: no
-
-comment:
-  layout: "reach,diff,flags,files,footer"
-  behavior: default
-  require_changes: no
-  require_base: no
-  require_head: yes
-
-github_checks:
-  annotations: false

docker/Dockerfile-base

@@ -37,7 +37,3 @@ RUN git lfs install --skip-repo && \
     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
-
-RUN if [ "$PYTORCH_VERSION" = "2.7.0" ] ; then \
-        pip3 install flash-attn==2.7.4.post1; \
-    fi

docs/cli.qmd

@@ -199,17 +199,6 @@ output_dir: # Directory to save evaluation results
 
 See [LM Eval Harness](https://github.com/EleutherAI/lm-evaluation-harness) for more details.
 
-### delinearize-llama4
-
-Delinearizes a Llama 4 linearized model into a regular HuggingFace Llama 4 model. This only works with the non-quantized linearized model.
-
-```bash
-axolotl delinearize-llama4 --model path/to/model_dir --output path/to/output_dir
-```
-
-This would be necessary to use with other frameworks. If you have an adapter, merge it with the non-quantized linearized model before delinearizing.
-
-
 ## Legacy CLI Usage
 
 While the new Click-based CLI is preferred, Axolotl still supports the legacy module-based CLI:

docs/config.qmd

@@ -693,9 +693,6 @@ sequence_parallel_degree:
 # Optional; strides across the key dimension. Larger values use more memory but should make training faster.
 # Must evenly divide the number of KV heads in your model.
 heads_k_stride: 1
-# One of "varlen_llama3", "batch_ring", "batch_zigzag", "batch_stripe". Defaults to "varlen_llama3"
-# in the sample packing case, and "batch_ring" in the non-sample packing case.
-ring_attn_func:
 
 # Path to torch distx for optim 'adamw_anyprecision'
 torchdistx_path:

docs/docker.qmd

@@ -28,8 +28,6 @@ main-base-py{python_version}-cu{cuda_version}-{pytorch_version}
 
 Tags examples:
 
-- `main-base-py3.11-cu128-2.7.0`
-- `main-base-py3.11-cu126-2.7.0`
 - `main-base-py3.11-cu124-2.6.0`
 - `main-base-py3.11-cu124-2.5.1`
 - `main-base-py3.11-cu124-2.4.1`
@@ -52,7 +50,7 @@ Link: [Docker Hub](https://hub.docker.com/r/axolotlai/axolotl)
 # on push to main
 main-py{python_version}-cu{cuda_version}-{pytorch_version}
 
-# latest main (currently torch 2.6.0, python 3.11, cuda 12.4)
+# latest main (currently torch 2.5.1, python 3.11, cuda 12.4)
 main-latest
 
 # nightly build
@@ -70,7 +68,6 @@ There may be some extra tags appended to the image, like `-vllm` which installs
 
 Tags examples:
 
-- `main-py3.11-cu126-2.7.0`
 - `main-py3.11-cu124-2.6.0`
 - `main-py3.11-cu124-2.5.1`
 - `main-py3.11-cu124-2.4.1`

docs/installation.qmd

@@ -19,12 +19,6 @@ This guide covers all the ways you can install and set up Axolotl for your envir
 
 ## Installation Methods {#sec-installation-methods}
 
-::: {.callout-important}
-Please make sure to have Pytorch installed before installing Axolotl in your local environment.
-
-Follow the instructions at: [https://pytorch.org/get-started/locally/](https://pytorch.org/get-started/locally/)
-:::
-
 ### PyPI Installation (Recommended) {#sec-pypi}
 
 ```{.bash}

docs/sequence_parallelism.qmd

@@ -27,9 +27,6 @@ To enable sequence parallelism, add the following to your configuration file:
 sequence_parallel_degree: 4  # Split sequences across 4 GPUs
 # Optional; strides across the key dimension. Larger values use more memory but should make training faster.
 heads_k_stride: 1
-# Optional; one of "varlen_llama3", "batch_ring", "batch_zigzag", "batch_stripe". Defaults to
-# "varlen_llama3" when `sample_packing: true`, and "batch_ring" otherwise.
-ring_attn_func:
 ```
 
 The `sequence_parallel_degree` should be a divisor of the total number of GPUs. For example:

examples/glm4/qlora-32b.yaml

@@ -1,62 +0,0 @@
-base_model: THUDM/GLM-4-32B-0414
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-load_in_4bit: true
-
-datasets:
-  - path: teknium/GPT4-LLM-Cleaned
-    type: alpaca
-dataset_prepared_path: last_run_prepared
-val_set_size: 0
-output_dir: ./outputs/qlora-out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-eval_sample_packing: true
-pad_to_sequence_len: true
-
-lora_r: 16
-lora_alpha: 32
-lora_dropout: 0.05
-lora_target_modules:
-  - gate_proj
-  - down_proj
-  - up_proj
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 2
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-loss_watchdog_threshold: 5.0
-loss_watchdog_patience: 3
-
-warmup_steps: 10
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:

examples/llama-4/README.md

@@ -1,36 +1,16 @@
 # Llama 4 by Meta AI
 
-## Flash Attention vs Flex Attention
-
-While Flash Attention to support is "enabled" for Llama-4, the upstream implementation is not correct and usage of Flex Attention is recommended.
-
 ## Available Examples
 
 ### Llama 4 Scout 17Bx16Experts (109B)
+- [Multi-Modal/Vision QLoRA w/ FSDP1](./scout-vision-qlora-fsdp.yaml)
+- [Text Single GPU (H100) QLoRA](./scout-qlora-single-h100.yaml)
+- [Text Multi GPU QLoRA w/ FSDP1](./scout-qlora-fsdp1.yaml)
 
-Flex Attention
-- [Text Single GPU (H100) QLoRA](./scout-qlora-single-h100-flex.yaml)
-- [Text Multi GPU QLoRA w/ FSDP2](./scout-qlora-flexattn-fsdp2.yaml)
-
-[//]: # (Flash Attention (Do not use))
-
-[//]: # (- [Multi-Modal/Vision QLoRA w/ FSDP1](./scout-vision-qlora-fsdp.yaml))
-
-[//]: # (- [Text Single GPU (H100) QLoRA](./scout-qlora-single-h100.yaml))
-
-[//]: # (- [Text Multi GPU QLoRA w/ FSDP1](./scout-qlora-fsdp1.yaml))
-
-Our Single H100 implementation for Llama 4 Scout uses only 64.5GB VRAM for post-training with 4k context length @ 519 tokens/second. [WandB logs here](https://wandb.ai/axolotl-ai/llama4-flexattn-qlora/runs/wpie7dkj)
-Multi-GPU (4xH100) for Llama 4 Scout uses 62.8GB VRAM/GPU @ 4k contenxt length @ 280tps/gpu, [WandB logs here](https://wandb.ai/axolotl-ai/llama4-flexattn-qlora/runs/2lkezdj8)
+Our Single H100 implementation for Llama 4 Scout uses only 68.5GB VRAM for post-training with 4k context length @ 546 tokens/second. [WandB logs here](https://wandb.ai/axolotl-ai/llama4-sft/runs/zic56rhd)
 
 ### Llama 4 Maverick 17Bx128Experts (400B)
 
-Coming Soon
+- [Text Multi GPU QLoRA w/FSDP1](./maverick-qlora-fsdp1.yaml)
 
-## Delinearized Llama 4 Models
-
-We provide a script to delinearize Llama 4 linearized models into regular HuggingFace Llama 4 models.
-
-```bash
-axolotl delinearize-llama4 --model path/to/model_dir --output path/to/output_dir
-```
+Our 4xH100 implementation for Llama 4 Maverick uses 79.5GB VRAM/GPU for post-training with 4k context length @ 206 tokens/second. [WandB logs here.](https://wandb.ai/axolotl-ai/llama-sft/runs/siyvwuxc?nw=nwuserwinglian)

examples/llama-4/scout-qlora-flexattn-fsdp2.yaml

@@ -1,86 +0,0 @@
-base_model: axolotl-quants/Llama-4-Scout-17B-16E-Linearized-bnb-nf4-bf16
-model_type: Llama4ForConditionalGeneration
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-plugins:
-  - axolotl.integrations.liger.LigerPlugin
-
-liger_glu_activation: true
-liger_rms_norm: true
-liger_layer_norm: true
-
-llama4_linearized_experts: true
-load_in_4bit: true
-adapter: qlora
-lora_r: 32
-lora_alpha: 64
-lora_target_modules:
-  - self_attn.q_proj
-  - self_attn.k_proj
-  - self_attn.v_proj
-  - self_attn.o_proj
-  - shared_expert.gate_proj
-  - shared_expert.up_proj
-  - shared_expert.down_proj
-  # - experts.gate_projs.[0-9]+$
-  # - experts.up_projs.[0-9]+$
-  # - experts.down_projs.[0-9]+$
-lora_modules_to_save:
-  # - lm_head
-  # - embed_tokens
-
-chat_template: llama4
-datasets:
-  - path: mlabonne/FineTome-100k
-    type: chat_template
-    split: train[:20%]
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-sequence_len: 4096
-sample_packing: true
-pad_to_sequence_len: true
-
-gradient_accumulation_steps: 1
-micro_batch_size: 2
-num_epochs: 3
-optimizer: adamw_torch_4bit
-lr_scheduler: cosine
-learning_rate: 1e-4
-
-bf16: true
-tf32: true
-
-logging_steps: 1
-flex_attention: true
-flex_attn_compile_kwargs:
-  dynamic: false
-  mode: max-autotune-no-cudagraphs
-
-warmup_steps: 10
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
-fsdp:
-  - auto_wrap
-  - full_shard
-fsdp_config:
-  fsdp_version: 2
-  fsdp_offload_params: false
-  fsdp_cpu_ram_efficient_loading: true
-  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  fsdp_transformer_layer_cls_to_wrap: Llama4TextDecoderLayer
-  fsdp_state_dict_type: SHARDED_STATE_DICT
-  fsdp_sharding_strategy: FULL_SHARD
-  fsdp_reshard_after_forward: true
-  fsdp_activation_checkpointing: true
-special_tokens:
-  pad_token: <|finetune_right_pad_id|>
-  eos_token: <|eot|>

examples/llama-4/scout-qlora-single-h100-flex.yaml

@@ -1,84 +0,0 @@
-base_model: axolotl-quants/Llama-4-Scout-17B-16E-Linearized-bnb-nf4-bf16
-model_type: Llama4ForConditionalGeneration
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-plugins:
-  - axolotl.integrations.liger.LigerPlugin
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-liger_glu_activation: true
-liger_rms_norm: true
-liger_layer_norm: true
-
-llama4_linearized_experts: true  # needed with custom linearized experts model
-load_in_4bit: true
-adapter: qlora
-lora_r: 32
-lora_alpha: 64
-lora_target_modules:
-  - self_attn.q_proj
-  - self_attn.k_proj
-  - self_attn.v_proj
-  - self_attn.o_proj
-  - shared_expert.gate_proj
-  - shared_expert.up_proj
-  - shared_expert.down_proj
-  # - experts.gate_projs.[0-9]+$  # optionally train the moe experts
-  # - experts.up_projs.[0-9]+$
-  # - experts.down_projs.[0-9]+$
-lora_modules_to_save:
-  # - lm_head  # needed if modifying vocabulary
-  # - embed_tokens
-
-lora_mlp_kernel: true
-lora_qkv_kernel: true
-lora_o_kernel: true
-
-chat_template: llama4
-datasets:
-  - path: mlabonne/FineTome-100k
-    type: chat_template
-    split: train[:20%]
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-sequence_len: 4096  # up to 8k will work on a single H100
-sample_packing: true
-pad_to_sequence_len: true
-
-gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_torch_4bit
-lr_scheduler: cosine
-learning_rate: 1e-4
-
-bf16: true
-tf32: true
-
-torch_compile: true
-flex_attention: true
-flex_attn_compile_kwargs:
-  dynamic: false
-  mode: max-autotune-no-cudagraphs
-
-gradient_checkpointing: offload
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-
-logging_steps: 1
-warmup_steps: 20
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-weight_decay: 0.0
-special_tokens:
-  pad_token: <|finetune_right_pad_id|>
-  eos_token: <|eot|>

examples/llama-4/scout-vision-qlora-fsdp2-flex.yaml

@@ -1,89 +0,0 @@
-base_model: axolotl-quants/Llama-4-Scout-17B-16E-Linearized-bnb-nf4-bf16
-model_type: Llama4ForConditionalGeneration
-processor_type: Llama4Processor
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-# these 3 lines are needed for now to handle vision chat templates w images
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-
-sequence_len: 4096
-
-plugins:
-  - axolotl.integrations.liger.LigerPlugin
-
-liger_glu_activation: true
-liger_rms_norm: true
-liger_layer_norm: true
-
-llama4_linearized_experts: true  # use Axolotl's customized model
-load_in_4bit: true
-adapter: qlora
-lora_r: 32
-lora_alpha: 64
-lora_target_modules:
-  - self_attn.q_proj
-  - self_attn.k_proj
-  - self_attn.v_proj
-  - self_attn.o_proj
-  - shared_expert.gate_proj
-  - shared_expert.up_proj
-  - shared_expert.down_proj
-  - vision_adapter.mlp.fc1
-  - vision_adapter.mlp.fc2
-  # - experts.gate_projs.[0-9]+$
-  # - experts.up_projs.[0-9]+$
-  # - experts.down_projs.[0-9]+$
-lora_modules_to_save:
-  - lm_head
-  - embed_tokens
-
-chat_template: llama4
-datasets:
-  - path: HuggingFaceH4/llava-instruct-mix-vsft
-    type: chat_template
-    split: train[:1%]
-    field_messages: messages
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_torch_4bit
-lr_scheduler: cosine
-learning_rate: 1e-4
-
-bf16: true
-tf32: true
-
-logging_steps: 1
-flex_attention: true
-flex_attn_compile_kwargs:
-  dynamic: false
-  mode: max-autotune-no-cudagraphs
-
-warmup_steps: 10
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
-fsdp:
-  - auto_wrap
-  - full_shard
-fsdp_config:
-  fsdp_version: 2
-  fsdp_offload_params: false
-  fsdp_cpu_ram_efficient_loading: true
-  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  fsdp_transformer_layer_cls_to_wrap: Llama4TextDecoderLayer
-  fsdp_state_dict_type: SHARDED_STATE_DICT
-  fsdp_sharding_strategy: FULL_SHARD
-  fsdp_reshard_after_forward: true
-  fsdp_activation_checkpointing: true
-special_tokens:
-  pad_token: <|finetune_right_pad_id|>
-  eos_token: <|eot|>

requirements-dev.txt

@@ -1,6 +1,6 @@
+pre-commit
 black
 mypy
-pre-commit
 types-requests
 quartodoc
 jupyter

requirements-tests.txt

@@ -1,8 +1,5 @@
-codecov
-codecov-cli
 pytest
-pytest-cov
+pytest-xdist
 pytest-retry
 pytest-sugar
-pytest-xdist
 tbparse

requirements.txt

@@ -6,7 +6,7 @@ triton>=3.0.0
 mamba-ssm==1.2.0.post1
 xformers>=0.0.23.post1
 autoawq==0.2.7.post3
-liger-kernel==0.5.8
+liger-kernel==0.5.6
 # END section
 
 packaging==23.2
@@ -19,7 +19,6 @@ datasets==3.5.0
 deepspeed>=0.15.4
 trl==0.16.1
 hf_xet==1.0.0
-hqq==0.2.5
 
 optimum==1.16.2
 hf_transfer

scripts/cutcrossentropy_install.py

@@ -25,5 +25,5 @@ if cce_spec:
 
 print(
     UNINSTALL_PREFIX
-    + 'pip install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@bad6f7b49c75fdec69471abb71b4cddd0f0c6438"'
+    + 'pip install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@24fbe4b5dab9a6c250a014573613c1890190536c"'
 )

setup.py

@@ -51,7 +51,7 @@ def parse_requirements(extras_require_map):
             try:
                 torch_version = version("torch")
             except PackageNotFoundError:
-                torch_version = "2.6.0"  # default to torch 2.6
+                torch_version = "2.5.1"
             _install_requires.append(f"torch=={torch_version}")
 
             version_match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", torch_version)
@@ -64,16 +64,10 @@ def parse_requirements(extras_require_map):
             else:
                 raise ValueError("Invalid version format")
 
-            if (major, minor) >= (2, 7):
+            if (major, minor) >= (2, 6):
                 _install_requires.pop(_install_requires.index(xformers_version))
-                # _install_requires.append("xformers==0.0.29.post3")  # xformers seems to be hard pinned to 2.6.0
-                extras_require_map["vllm"] = ["vllm==0.8.3"]
-            elif (major, minor) >= (2, 6):
-                _install_requires.pop(_install_requires.index(xformers_version))
-                _install_requires.append(
-                    "xformers==0.0.29.post2"
-                )  # vllm needs post2 w torch 2.6
-                extras_require_map["vllm"] = ["vllm==0.8.3"]
+                _install_requires.append("xformers==0.0.29.post2")
+                extras_require_map["vllm"] = ["vllm==0.8.1"]
             elif (major, minor) >= (2, 5):
                 _install_requires.pop(_install_requires.index(xformers_version))
                 if patch == 0:

src/axolotl/cli/args.py

@@ -39,16 +39,16 @@ class TrainerCliArgs:
 class VllmServeCliArgs:
     """Dataclass with CLI arguments for `axolotl vllm-serve` command."""
 
-    tensor_parallel_size: Optional[int] = field(
-        default=None,
+    tensor_parallel_size: int = field(
+        default=1,
         metadata={"help": "Number of tensor parallel workers to use."},
     )
-    host: Optional[str] = field(
-        default=None,  # nosec B104
+    host: str = field(
+        default="0.0.0.0",  # nosec B104
         metadata={"help": "Host address to run the server on."},
     )
-    port: Optional[int] = field(
-        default=None,
+    port: int = field(
+        default=8000,
         metadata={"help": "Port to run the server on."},
     )
     gpu_memory_utilization: Optional[float] = field(

src/axolotl/cli/delinearize_llama4.py

@@ -1,156 +0,0 @@
-"""
-CLI tool to delinearize quantized/Linearized Llama-4 models.
-"""
-
-import os
-from pathlib import Path
-from typing import Generator, Union
-
-import fire
-import torch
-from accelerate import init_empty_weights
-from dotenv import load_dotenv
-from transformers import AutoProcessor
-
-
-def iter_convert_patched_to_hf(model_state_dict, num_experts) -> Generator:
-    keys = list(model_state_dict.keys())
-    for key in keys:
-        if ".feed_forward.experts." not in key:
-            yield key, model_state_dict[key]
-        if ".feed_forward.experts.gate_projs" in key:
-            # gate gets fused with up so skip the yield on this and we'll fuse it when asking for the up
-            continue
-        if ".feed_forward.experts.up_projs" in key:
-            if ".feed_forward.experts.up_projs.0." in key:
-                # handle the re-shape and fusing of gate and up, and conversion from linear to parameter
-                prefix = key.split(".up_projs.0.")[0]
-                key = f"{prefix}.gate_up_proj"
-                # grab all the up_projs and gate_projs across all experts
-                gate_stacked = torch.stack(
-                    [
-                        model_state_dict[
-                            f"{prefix}.gate_projs.{expert_idx}.weight"
-                        ].transpose(0, 1)
-                        for expert_idx in range(num_experts)
-                    ]
-                )
-                up_stacked = torch.stack(
-                    [
-                        model_state_dict[
-                            f"{prefix}.up_projs.{expert_idx}.weight"
-                        ].transpose(0, 1)
-                        for expert_idx in range(num_experts)
-                    ]
-                )
-                gate_up_proj = torch.cat((gate_stacked, up_stacked), dim=-1)
-                del gate_stacked, up_stacked
-                yield key, gate_up_proj
-            else:
-                del model_state_dict[key]
-                continue
-        if ".feed_forward.experts.down_projs" in key:
-            if ".feed_forward.experts.down_projs.0." in key:
-                # handle the re-shape and fusing of gate and up, and conversion from linear to parameter
-                prefix = key.split(".down_projs.0.")[0]
-                key = f"{prefix}.down_proj"
-                # grab all the down_projs across all experts
-                down_stacked = torch.stack(
-                    [
-                        model_state_dict[
-                            f"{prefix}.down_projs.{expert_idx}.weight"
-                        ].transpose(0, 1)
-                        for expert_idx in range(num_experts)
-                    ]
-                )
-                yield key, down_stacked
-            else:
-                del model_state_dict[key]
-                continue
-
-
-def do_cli(model: Union[Path, str], output: Union[Path, str]) -> None:
-    """
-    Convert a patched HF format Llama4 model (with separated projections)
-    back to the original HF format (with fused projections).
-
-    Args:
-        model: Path to the patched HF model
-        output: Path to save the converted model
-    """
-    print(f"Loading model from {model}")
-    from axolotl.monkeypatch.models.llama4.modeling import (
-        patch_llama4_linearized_modeling,
-    )
-
-    unpatch_llama4 = patch_llama4_linearized_modeling()
-    from transformers import Llama4ForConditionalGeneration
-
-    model_ = Llama4ForConditionalGeneration.from_pretrained(
-        model, torch_dtype=torch.bfloat16
-    )
-    processor = AutoProcessor.from_pretrained(model)
-    processor.save_pretrained(output)
-
-    device = model_.device.type
-    if device == "cuda":
-        print(
-            f"peak memory allocated: {torch.cuda.max_memory_allocated() / 1024**2} MB"
-        )
-        print(f"peak memory reserved: {torch.cuda.max_memory_reserved() / 1024**2} MB")
-    model_config = model_.config
-    config = model_.config.get_text_config()
-
-    # Get key dimensions from the config
-    hidden_size = config.hidden_size
-    intermediate_size = config.intermediate_size
-    num_experts = config.num_local_experts
-
-    print(
-        f"Model dimensions: hidden_size={hidden_size}, intermediate_size={intermediate_size}, num_experts={num_experts}"
-    )
-
-    # Create output directory if it doesn't exist
-    os.makedirs(output, exist_ok=True)
-
-    # Get state dict
-    state_dict = model_.state_dict()
-    del model_
-
-    # Create a new state dict for the converted model
-    converted_state_dict = {}
-
-    # First, copy all keys that don't need modification
-    for key, value in iter_convert_patched_to_hf(state_dict, num_experts):
-        converted_state_dict[key] = value
-
-    del state_dict
-    if device == "cuda":
-        torch.cuda.empty_cache()
-        print("State dict converted.")
-        print(
-            f"peak memory allocated: {torch.cuda.max_memory_allocated() / 1024**2} MB"
-        )
-        print(f"peak memory reserved: {torch.cuda.max_memory_reserved() / 1024**2} MB")
-    # Ideally re-load the model import to load the converted state dict
-    # Save the converted model
-    with init_empty_weights():
-        unpatch_llama4()
-        model_ = Llama4ForConditionalGeneration(model_config)
-
-    if device == "cuda":
-        print("State dict loaded into model.")
-        print(
-            f"peak memory allocated: {torch.cuda.max_memory_allocated() / 1024**2} MB"
-        )
-        print(f"peak memory reserved: {torch.cuda.max_memory_reserved() / 1024**2} MB")
-    model_.load_state_dict(converted_state_dict, strict=False, assign=True)
-    print(f"Saving converted model to {output}...")
-    model_.save_pretrained(output)
-
-    print(f"Model successfully converted and saved to {output}")
-
-
-if __name__ == "__main__":
-    load_dotenv()
-    fire.Fire(do_cli)

src/axolotl/cli/main.py

@@ -330,15 +330,6 @@ def vllm_serve(config: str, **cli_args: VllmServeCliArgs):
     do_vllm_serve(config, cli_args)
 
 
-@cli.command()
-@click.argument("model", type=click.Path(exists=True, path_type=str))
-@click.argument("output", type=click.Path(exists=False, path_type=str))
-def delinearize_llama4(model: str, output: str) -> None:
-    from axolotl.cli.delinearize_llama4 import do_cli as do_delinearize_llama4
-
-    do_delinearize_llama4(model, output)
-
-
 cli.add_command(lm_eval)
 
 

src/axolotl/cli/merge_lora.py

@@ -40,7 +40,6 @@ def do_merge_lora(*, cfg: DictDefault) -> None:
         LOG.warning("Error raised: %s", e)
 
     model.generation_config.do_sample = True
-    model.config.use_cache = True
 
     if cfg.local_rank == 0:
         LOG.info(f"Saving merged model to: {str(Path(cfg.output_dir) / 'merged')}...")

src/axolotl/core/trainer_builder.py

@@ -776,7 +776,6 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
         training_arguments_kwargs["sequence_parallel_degree"] = (
             self.cfg.sequence_parallel_degree
         )
-        training_arguments_kwargs["ring_attn_func"] = self.cfg.ring_attn_func
 
         if self.cfg.reward_model:
             training_args_cls = AxolotlRewardConfig
@@ -932,6 +931,8 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
                 collator = DataCollatorForSeq2Seq
 
         kwargs["return_tensors"] = "pt"
+        if issubclass(collator, DataCollatorForSeq2Seq):
+            kwargs["sequence_parallel_degree"] = training_args.sequence_parallel_degree
 
         return collator(
             *collator_args,
@@ -1037,20 +1038,15 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
         if self.cfg.dataset_processes:
             training_args_kwargs["dataset_num_proc"] = self.cfg.dataset_processes
 
-        if self.cfg.trl and self.cfg.trl.beta is not None:
-            training_args_kwargs["beta"] = self.cfg.trl.beta
-        elif self.cfg.rl_beta is not None:
-            training_args_kwargs["beta"] = self.cfg.rl_beta
-        elif self.cfg.orpo_alpha is not None:
+        if (self.cfg.trl and self.cfg.trl.beta) or self.cfg.rl_beta:
+            training_args_kwargs["beta"] = self.cfg.trl.beta or self.cfg.rl_beta
+        if self.cfg.orpo_alpha:
             # trl does some odd mapping of alpha to beta to reuse the beta parameter ???
             training_args_kwargs["beta"] = self.cfg.orpo_alpha
 
         if self.cfg.rpo_alpha is not None:
             training_args_kwargs["rpo_alpha"] = self.cfg.rpo_alpha
 
-        if self.cfg.use_wandb:
-            training_args_kwargs["run_name"] = self.cfg.wandb_name
-
         training_args_cls = None
         blocklist_args_kwargs = []
         if self.cfg.rl == "simpo":
@@ -1121,12 +1117,6 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
             **training_args_kwargs,
         )
 
-        # unset run_name so wandb sets up experiment names
-        if self.cfg.use_wandb and training_args.run_name == training_args.output_dir:
-            training_args.run_name = (  # pylint: disable=attribute-defined-outside-init
-                None
-            )
-
         return training_args
 
     def build(self, total_num_steps):

src/axolotl/core/trainers/base.py

@@ -371,15 +371,13 @@ class AxolotlTrainer(
                 num_items_in_batch=num_items_in_batch,
             )
 
-        loss = super().compute_loss(
+        return super().compute_loss(
             model,
             inputs,
             return_outputs=return_outputs,
             num_items_in_batch=num_items_in_batch,
         )
 
-        return loss
-
     @staticmethod
     def orpo_concatenate_inputs(inputs, label_pad_token=-100, pad_token=0, device=None):
         concatenated_batch = {}

src/axolotl/core/trainers/grpo/__init__.py

@@ -40,8 +40,8 @@ class GRPOStrategy:
 
         if trl.use_vllm:
             grpo_args_kwargs["use_vllm"] = trl.use_vllm
-            grpo_args_kwargs["vllm_server_host"] = trl.vllm_server_host or trl.vllm.host
-            grpo_args_kwargs["vllm_server_port"] = trl.vllm_server_port or trl.vllm.port
+            grpo_args_kwargs["vllm_server_host"] = trl.vllm_server_host
+            grpo_args_kwargs["vllm_server_port"] = trl.vllm_server_port
             if trl.vllm_server_timeout:
                 grpo_args_kwargs["vllm_server_timeout"] = trl.vllm_server_timeout
             if trl.vllm_guided_decoding_regex:

src/axolotl/core/trainers/mixins/__init__.py

@@ -6,4 +6,4 @@
 from .optimizer import OptimizerMixin
 from .rng_state_loader import RngLoaderMixin
 from .scheduler import SchedulerMixin
-from .sequence_parallel import SequenceParallelContextManager, SequenceParallelMixin
+from .sequence_parallel import SequenceParallelMixin

src/axolotl/core/trainers/mixins/sequence_parallel.py

@@ -1,86 +1,16 @@
-"""
-Module for Axolotl trainer sequence parallelism mixin and training context manager
-"""
+"""Module for Axolotl trainer sequence parallelism mixin"""
 
-import functools
 import logging
 
-import torch
 import torch.distributed as dist
 from datasets import Dataset
-from torch import nn
 from torch.utils.data import DistributedSampler, Sampler
-from torch.utils.hooks import RemovableHandle
 
-from axolotl.monkeypatch.attention.ring_attn import (
-    RingAttnFunc,
-    get_ring_attn_group,
-    update_ring_attn_params,
-)
+from axolotl.monkeypatch.attention.ring_attn import get_ring_attn_group
 
 LOG = logging.getLogger(__name__)
 
 
-def apply_sequence_parallelism(
-    batch: dict[str, torch.Tensor],
-    local_rank: int,
-    local_world_size: int,
-    ring_attn_func: RingAttnFunc,
-) -> dict[str, torch.Tensor]:
-    """
-    Apply sequence parallelism slicing to a batch.
-
-    Args:
-        batch: Batch dictionary (e.g., input_ids, attention_mask, etc.)
-        local_rank: Local rank in the sequence parallel group
-        local_world_size: World size of the sequence parallel group
-        ring_attn_func: The ring attention function to use
-
-    Returns:
-        Sliced batch dictionary.
-    """
-    # Update ring attention params if needed
-    if batch.get("position_ids") is not None:
-        update_ring_attn_params(position_ids=batch["position_ids"])
-
-    # Slice batch for sequence parallel processing
-    total_seq_len = batch["input_ids"].size(1)
-    for key in batch:
-        if (
-            key in batch
-            and isinstance(batch[key], torch.Tensor)
-            and batch[key].dim() > 1
-            and batch[key].size(1) == total_seq_len
-        ):
-
-            if ring_attn_func in [
-                RingAttnFunc.VARLEN_LLAMA3,
-                RingAttnFunc.BATCH_RING,
-            ]:
-                # Split in sequential fashion and grab this rank's chunk
-                batch[key] = (
-                    batch[key].chunk(local_world_size, dim=1)[local_rank].contiguous()
-                )
-            elif ring_attn_func is RingAttnFunc.BATCH_ZIGZAG:
-                chunks = batch[key].chunk(2 * local_world_size, dim=1)
-
-                # Take rank's chunk and opposing chunk for zigzag pattern
-                selected_chunks = [
-                    chunks[local_rank],
-                    chunks[2 * local_world_size - local_rank - 1],
-                ]
-                batch[key] = torch.cat(selected_chunks, dim=1).contiguous()
-            elif ring_attn_func is RingAttnFunc.BATCH_STRIPE:
-                # Split into striped data and stack
-                tensor = torch.stack(
-                    batch[key].split(local_world_size, dim=1),
-                    dim=1,
-                ).transpose(1, 2)
-                batch[key] = tensor[:, local_rank].contiguous()
-
-    return batch
-
-
 class SequenceParallelMixin:
     """
     Mixin class for sequence parallelism support in trainers.
@@ -157,157 +87,3 @@ class SequenceParallelMixin:
         return self._create_sequence_parallel_sampler(
             eval_dataset, shuffle=False, is_eval=True
         )
-
-
-class SequenceParallelContextManager:
-    """
-    Context manager for sequence parallelism operations.
-
-    This class provides a context that will automatically apply sequence parallelism
-    during model forward passes using a pre-forward hook, and gather outputs from
-    across the sequence parallelism group using a post-forward hook.
-    """
-
-    def __init__(
-        self,
-        model: nn.Module,
-        sequence_parallel_degree: int,
-        ring_attn_func: RingAttnFunc,
-    ):
-        self.model = model
-        self.sequence_parallel_degree = sequence_parallel_degree
-        self.ring_attn_func = ring_attn_func
-        self.process_group = get_ring_attn_group()
-
-        # Initialize sequence parallel group details
-        self.local_rank = dist.get_rank(self.process_group)
-        self.local_world_size = dist.get_world_size(self.process_group)
-
-        # Will store hook handles for removal
-        self.hook_handles: list[RemovableHandle] = []
-
-        # Create a partially applied version of the apply_sequence_parallelism function
-        # with pre-configured params
-        self.apply_sequence_parallelism = functools.partial(
-            apply_sequence_parallelism,
-            local_rank=self.local_rank,
-            local_world_size=self.local_world_size,
-            ring_attn_func=self.ring_attn_func,
-        )
-
-    def __enter__(self):
-        # Forward pre-hook to apply sequence parallelism
-        def sequence_parallel_pre_hook(_, args, kwargs):
-            # Apply sequence parallelism to kwargs
-            kwargs = self.apply_sequence_parallelism(batch=kwargs)
-            return args, kwargs
-
-        # Forward post-hook to gather outputs
-        def sequence_parallel_post_hook(_, __, output):
-            # Gather the sharded outputs
-            return self.gather_outputs(output)
-
-        # Register both hooks
-        self.hook_handles.append(
-            self.model.register_forward_pre_hook(
-                sequence_parallel_pre_hook, with_kwargs=True
-            )
-        )
-        self.hook_handles.append(
-            self.model.register_forward_hook(sequence_parallel_post_hook)
-        )
-
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        # Remove all hooks
-        for handle in self.hook_handles:
-            handle.remove()
-        self.hook_handles = []
-
-    def gather_outputs(self, output):
-        """Gather sharded outputs from all ranks and reconstruct the full tensor."""
-        # Handle different output formats (dict, tensor, etc.)
-        if isinstance(output, dict):
-            gathered_output = {}
-            for key, value in output.items():
-                if isinstance(value, torch.Tensor) and value.dim() > 1:
-                    # Gather logits or other sequence-sharded tensors
-                    gathered_value = self.gather_tensor(value)
-                    gathered_output[key] = gathered_value
-                else:
-                    gathered_value = value.clone()
-                    dist.all_reduce(
-                        gathered_value, op=dist.ReduceOp.SUM, group=self.process_group
-                    )
-                    gathered_output[key] = gathered_value
-            return gathered_output
-        if isinstance(output, torch.Tensor):
-            return self.gather_tensor(output)
-
-        return output
-
-    def gather_tensor(self, tensor):
-        """Gather a sharded tensor from all ranks."""
-        # Prepare tensors for all_gather
-        world_size = self.local_world_size
-
-        # Create list to store tensors from all ranks
-        gathered_tensors = [torch.zeros_like(tensor) for _ in range(world_size)]
-
-        # All-gather operation
-        dist.all_gather(gathered_tensors, tensor, group=self.process_group)
-
-        # Concatenate along sequence dimension (typically dim=1)
-        if self.ring_attn_func in [RingAttnFunc.VARLEN_LLAMA3, RingAttnFunc.BATCH_RING]:
-            # Simple concatenation for standard sharding
-            return torch.cat(gathered_tensors, dim=1)
-
-        if self.ring_attn_func is RingAttnFunc.BATCH_ZIGZAG:
-            # Each rank has a pattern of (rank, world_size*2-rank-1)
-            reconstituted_tensors = [None] * (world_size * 2)
-
-            # First, split each gathered tensor into its two chunks
-            for rank, gathered_tensor in enumerate(gathered_tensors):
-                # Each tensor contains two chunks in the sequence dimension
-                chunk_size = gathered_tensor.size(1) // 2
-                chunk1, chunk2 = gathered_tensor.split(chunk_size, dim=1)
-
-                # Place chunks in their original positions
-                reconstituted_tensors[rank] = chunk1
-                reconstituted_tensors[world_size * 2 - rank - 1] = chunk2
-
-            # Concatenate the reconstituted tensors in the correct order
-            return torch.cat(reconstituted_tensors, dim=1)
-
-        # Otherwise, RingAttnFunc.BATCH_STRIPE
-        # In striping, each rank has every world_size-th slice
-        batch_size = tensor.size(0)
-        hidden_dim = tensor.size(-1)
-
-        # First, determine the full sequence length
-        total_seq_len = 0
-        for t in gathered_tensors:
-            total_seq_len += t.size(1)
-
-        # Create a tensor to hold the unstriped result
-        result = torch.zeros(
-            batch_size,
-            total_seq_len,
-            hidden_dim,
-            dtype=tensor.dtype,
-            device=tensor.device,
-        )
-
-        # For each rank's tensor, distribute its slices to the correct positions
-        for rank, gathered_tensor in enumerate(gathered_tensors):
-            # The rank's tensor contains every world_size-th slice
-            # starting from its rank position
-            seq_len = gathered_tensor.size(1)
-            for i in range(seq_len):
-                # Calculate the position in the full tensor
-                pos = i * world_size + rank
-                if pos < total_seq_len:
-                    result[:, pos] = gathered_tensor[:, i]
-
-        return result

src/axolotl/core/training_args.py

@@ -9,8 +9,6 @@ from PIL.Image import Resampling
 from transformers import TrainingArguments
 from trl import CPOConfig, KTOConfig, ORPOConfig, PRMConfig, RewardConfig
 
-from axolotl.monkeypatch.attention.ring_attn.patch import RingAttnFunc
-
 
 @dataclass
 class AxolotlTrainingMixins:
@@ -220,12 +218,6 @@ class AxolotlTrainingMixins:
         default=1,
         metadata={"help": "The number of workers to use in sequence parallelism"},
     )
-    ring_attn_func: Optional[RingAttnFunc] = field(
-        default=None,
-        metadata={
-            "help": "The ring-flash-attn function to use in sequence parallelism"
-        },
-    )
 
     # multi-modal section
 

src/axolotl/integrations/cut_cross_entropy/README.md

@@ -12,14 +12,12 @@ See https://github.com/apple/ml-cross-entropy
 
 Run the following command to install `cut_cross_entropy[transformers]` if you don't have it already.
 
-- If you are in dev environment
 ```bash
+# if you are in dev environment
 python scripts/cutcrossentropy_install.py | sh
-```
 
-- If you are installing from pip
-```bash
-pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@bad6f7b49c75fdec69471abb71b4cddd0f0c6438"
+# if you are not in dev environment
+pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@24fbe4b5dab9a6c250a014573613c1890190536c"
 ```
 
 ## Usage
@@ -27,6 +25,8 @@ pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transform
 ```yaml
 plugins:
   - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+cut_cross_entropy: true
 ```
 
 ## Supported Models
@@ -45,8 +45,6 @@ plugins:
 - qwen2
 - cohere
 - cohere2
-- glm
-- glm4
 
 ## Citation
 

src/axolotl/integrations/cut_cross_entropy/__init__.py

@@ -33,7 +33,7 @@ LOG = logging.getLogger("axolotl.integrations.cut_cross_entropy")
 
 _CCE_INSTALL_MESSAGE = (
     "Please install cut_cross_entropy with transformers support using "
-    '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@bad6f7b49c75fdec69471abb71b4cddd0f0c6438"`'
+    '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@24fbe4b5dab9a6c250a014573613c1890190536c"`'
 )
 
 

src/axolotl/integrations/cut_cross_entropy/args.py

@@ -28,7 +28,7 @@ class CutCrossEntropyArgs(BaseModel):
     Input args for Cut Cross Entropy.
     """
 
-    cut_cross_entropy: Optional[bool] = True
+    cut_cross_entropy: Optional[bool] = None
 
     @model_validator(mode="before")
     @classmethod

src/axolotl/integrations/cut_cross_entropy/monkeypatch/glm4.py

@@ -1,57 +0,0 @@
-"""GLM 4 patch. GLM family inherits from Llama."""
-
-from types import MethodType
-
-import transformers
-from cut_cross_entropy.transformers.utils import (
-    PatchOptions,
-    TransformersModelT,
-)
-
-
-def patch_glm(
-    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
-    patch_options: PatchOptions,
-) -> TransformersModelT | None:
-
-    # Set the _PATCH_OPTS in the llama patch file
-    import cut_cross_entropy.transformers.llama as llama_patch
-
-    llama_patch._PATCH_OPTS = patch_options  # pylint: disable=protected-access
-
-    from cut_cross_entropy.transformers.llama import cce_forward
-    from transformers.models.glm import modeling_glm
-
-    if isinstance(maybe_model, transformers.PreTrainedModel):
-        assert isinstance(
-            maybe_model, modeling_glm.GlmForCausalLM
-        ), f"Expected a GlmForCausalLM model. Got {type(maybe_model)}."
-        maybe_model.forward = MethodType(cce_forward, maybe_model)
-        return maybe_model
-
-    modeling_glm.GlmForCausalLM.forward = cce_forward
-    return None
-
-
-def patch_glm4(
-    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
-    patch_options: PatchOptions,
-) -> TransformersModelT | None:
-
-    # Set the _PATCH_OPTS in the llama patch file
-    import cut_cross_entropy.transformers.llama as llama_patch
-
-    llama_patch._PATCH_OPTS = patch_options  # pylint: disable=protected-access
-
-    from cut_cross_entropy.transformers.llama import cce_forward
-    from transformers.models.glm4 import modeling_glm4
-
-    if isinstance(maybe_model, transformers.PreTrainedModel):
-        assert isinstance(
-            maybe_model, modeling_glm4.Glm4ForCausalLM
-        ), f"Expected a Glm4ForCausalLM model. Got {type(maybe_model)}."
-        maybe_model.forward = MethodType(cce_forward, maybe_model)
-        return maybe_model
-
-    modeling_glm4.Glm4ForCausalLM.forward = cce_forward
-    return None

src/axolotl/integrations/cut_cross_entropy/monkeypatch/llama4.py

@@ -165,7 +165,7 @@ def cce_forward(
 )
 def cce_forward_multimodal(
     self,
-    input_ids: torch.LongTensor | None = None,  # type: ignore
+    input_ids: torch.LongTensor | None = None,
     pixel_values: torch.FloatTensor | None = None,
     attention_mask: Optional[torch.Tensor] = None,
     position_ids: Optional[torch.LongTensor] = None,
@@ -254,7 +254,7 @@ def cce_forward_multimodal(
         )
 
     if inputs_embeds is None:
-        inputs_embeds = self.get_input_embeddings()(input_ids)  # type: ignore
+        inputs_embeds = self.get_input_embeddings()(input_ids)
 
     if pixel_values is not None:
         image_features = self.get_image_features(
@@ -263,13 +263,13 @@ def cce_forward_multimodal(
             vision_feature_select_strategy=vision_feature_select_strategy,
             image_sizes=image_sizes,
         )
-        original_inputs_embeds_shape = inputs_embeds.shape  # type: ignore
+        original_inputs_embeds_shape = inputs_embeds.shape
 
         vision_flat = image_features.view(-1, image_features.size(-1))
         projected_vision_flat = self.multi_modal_projector(vision_flat)
 
         special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1)
-        final_mask = special_image_mask.to(inputs_embeds.device)  # type: ignore
+        final_mask = special_image_mask.to(inputs_embeds.device)
         inputs_embeds = inputs_embeds.view(-1, inputs_embeds.size(-1))  # type: ignore
 
         final_mask_1d = final_mask[..., 0].reshape(-1)

src/axolotl/integrations/cut_cross_entropy/monkeypatch/patch.py

@@ -20,10 +20,6 @@ from axolotl.integrations.cut_cross_entropy.monkeypatch.gemma3 import (
     patch_gemma3,
     patch_gemma3_text,
 )
-from axolotl.integrations.cut_cross_entropy.monkeypatch.glm4 import (
-    patch_glm,
-    patch_glm4,
-)
 from axolotl.integrations.cut_cross_entropy.monkeypatch.llama4 import (
     patch_llama4,
     patch_llama4_text,
@@ -49,8 +45,6 @@ CUT_CROSS_ENTROPY_MODEL_MAPPING = {
     "qwen2": patch_qwen2,
     "cohere": patch_cohere,
     "cohere2": patch_cohere2,
-    "glm": patch_glm,
-    "glm4": patch_glm4,
 }
 
 

src/axolotl/integrations/liger/README.md

@@ -25,7 +25,7 @@ liger_fused_linear_cross_entropy: true
 - deepseek_v2
 - gemma
 - gemma2
-- gemma3
+- gemma3 (partial support, no support for FLCE yet)
 - granite
 - jamba
 - llama

src/axolotl/integrations/liger/__init__.py

@@ -21,6 +21,7 @@ It is designed to be performant, correct, and light-weight.
 import inspect
 import logging
 import sys
+from functools import partial
 
 from axolotl.integrations.base import BasePlugin
 
@@ -54,6 +55,7 @@ class LigerPlugin(BasePlugin):
             )
         from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss
         from liger_kernel.transformers.functional import liger_cross_entropy
+        from liger_kernel.transformers.geglu import LigerGEGLUMLP
         from liger_kernel.transformers.layer_norm import LigerLayerNorm
         from liger_kernel.transformers.monkey_patch import MODEL_TYPE_TO_APPLY_LIGER_FN
         from liger_kernel.transformers.rms_norm import LigerRMSNorm
@@ -139,6 +141,38 @@ class LigerPlugin(BasePlugin):
                 modeling_mod.CrossEntropyLoss = LigerCrossEntropyLoss
             if cfg.liger_fused_linear_cross_entropy:
                 modeling_mod.DeepseekV2ForCausalLM.forward = deepseekv2_lce_forward
+        elif cfg.model_config_type in ["gemma3", "gemma3_text"]:
+            from transformers.models.gemma3 import modeling_gemma3
+
+            if cfg.liger_rope:
+                modeling_gemma3.apply_rotary_pos_emb = liger_rotary_pos_emb
+            if cfg.liger_rms_norm:
+
+                def _liger_rms_norm_wrapper(dim, **kwargs):
+                    "Convert 'dim' keyword to 'hidden_size' to pass to LigerRMSNorm"
+                    return LigerRMSNorm(hidden_size=dim, **kwargs)
+
+                modeling_gemma3.Gemma3RMSNorm = partial(
+                    _liger_rms_norm_wrapper,
+                    offset=1.0,
+                    casting_mode="gemma",
+                    init_fn="zeros",
+                    in_place=False,
+                )
+            if cfg.liger_glu_activation:
+                modeling_gemma3.Gemma3MLP = LigerGEGLUMLP
+            if cfg.liger_layer_norm:
+                modeling_gemma3.nn.LayerNorm = LigerLayerNorm
+
+            if cfg.liger_cross_entropy:
+                from transformers.loss.loss_utils import nn
+
+                nn.functional.cross_entropy = liger_cross_entropy
+
+            if cfg.liger_fused_linear_cross_entropy:
+                raise NotImplementedError(
+                    "Fused linear cross entropy is not yet supported for Gemma3."
+                )
         elif cfg.model_config_type == "llama4":
             from axolotl.integrations.liger.models.llama4 import (
                 apply_liger_kernel_to_llama4,

src/axolotl/monkeypatch/accelerate/fsdp2.py

@@ -49,7 +49,7 @@ def fsdp2_load_full_state_dict(accelerator, model: torch.nn.Module, full_sd: dic
             )
             sharded_sd[param_name] = sharded_tensor
 
-    model.load_state_dict(sharded_sd, assign=True)
+    model.load_state_dict(sharded_sd)
 
 
 def patch_accelerate_fsdp_utils():

src/axolotl/monkeypatch/attention/flex_attn.py

@@ -7,11 +7,12 @@ import torch
 import transformers
 
 
-def patch_flex_wrapper(**flex_attn_compile_kwargs):
+def patch_flex_wrapper():
     # TODO remove this patch when transformers#37285 is merged and in a release
     is_torch_2_6 = torch.__version__.startswith("2.6")
+    is_transformers_below_4_51 = transformers.__version__ < "4.51.0"
 
-    if not is_torch_2_6:
+    if not (is_torch_2_6 and is_transformers_below_4_51):
         return
 
     from torch.nn.attention.flex_attention import flex_attention
@@ -31,24 +32,17 @@ def patch_flex_wrapper(**flex_attn_compile_kwargs):
                 cls._instance = super().__new__(cls)
             return cls._instance
 
-        @classmethod
-        def del_singleton(cls):
-            cls._instance = None
-
         @torch.compiler.disable(recursive=False)
-        def __init__(self, training):
+        def __init__(self):
             """
             Initialize or update the singleton instance.
             """
-            self.training = None
-            if not self._is_flex_compiled or training != self.training:
-                # In PyTorch 2.6.0, there's a known issue with flex attention compilation which may
-                # cause errors. The suggested fix is to compile with "max-autotune-no-cudagraphs"
-                # see https://github.com/pytorch/pytorch/issues/146260 for training
-                self.training = training
+            if not self._is_flex_compiled:
                 self._compiled_flex_attention = torch.compile(
                     flex_attention,
-                    **flex_attn_compile_kwargs,
+                    dynamic=False,
+                    mode="max-autotune-no-cudagraphs",
+                    fullgraph=True,
                 )
                 self._is_flex_compiled = True
 
@@ -56,22 +50,15 @@ def patch_flex_wrapper(**flex_attn_compile_kwargs):
             return self._compiled_flex_attention
 
     transformers.integrations.flex_attention.WrappedFlexAttention = WrappedFlexAttention
-    setattr(
-        sys.modules["transformers.integrations.flex_attention"],
-        "WrappedFlexAttention",
-        WrappedFlexAttention,
-    )
 
 
 def patch_flex_make_mask():
     is_torch_2_6 = torch.__version__.startswith("2.6")
+    is_transformers_eq_4_51 = transformers.__version__ == "4.51.0"
 
-    if not is_torch_2_6:
+    if not (is_torch_2_6 and is_transformers_eq_4_51):
         return
 
-    from torch.nn.attention.flex_attention import (
-        _DEFAULT_SPARSE_BLOCK_SIZE as flex_default_block_size,
-    )
     from torch.nn.attention.flex_attention import (
         BlockMask,
     )
@@ -117,16 +104,14 @@ def patch_flex_make_mask():
         if not query_length:
             query_length = total_seq_len
         attention_mask_2d = torch.nn.functional.pad(
-            attention_mask_2d,
-            value=0,
-            pad=(0, abs(total_seq_len - max(key_length, flex_default_block_size))),
+            attention_mask_2d, value=0, pad=(0, key_length)
         )
         device = attention_mask_2d.device
         document_ids = attention_mask_2d.clone()
 
         if attention_chunk_size is not None:
             # we create an arange, then we just // by chunk size to get [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]
-            chunk_idxs = (document_ids.clone().fill_(1).cumsum(-1) - 1) // (
+            document_ids = (document_ids.fill_(1).cumsum(-1) - 1) // (
                 attention_chunk_size
             )
 
@@ -153,18 +138,6 @@ def patch_flex_make_mask():
             final_mask = causal_mask & padding_mask & document_mask
             return final_mask
 
-        def chunk_causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx):
-            """
-            Combines the chunk mask with the causal mask for chunked attention.
-            """
-            chunk_mask = chunk_idxs[batch_idx, q_idx] == chunk_idxs[batch_idx, kv_idx]
-            causal_doc_mask = causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx)
-            return chunk_mask & causal_doc_mask
-
-        mask_mod_maybe_combined = (
-            causal_mask_mod if attention_chunk_size is None else chunk_causal_mask_mod
-        )
-
         if offsets is not None:
             q_offset = offsets[0]
             kv_offset = offsets[1]
@@ -172,10 +145,10 @@ def patch_flex_make_mask():
             def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
                 offset_q = q_idx + q_offset
                 offset_kv = kv_idx + kv_offset
-                return mask_mod_maybe_combined(batch_idx, head_idx, offset_q, offset_kv)
+                return causal_mask_mod(batch_idx, head_idx, offset_q, offset_kv)
 
         else:
-            mask_mod = mask_mod_maybe_combined
+            mask_mod = causal_mask_mod
         return create_block_causal_mask_flex(
             mask_mod=mask_mod,
             B=batch_size,
@@ -187,16 +160,11 @@ def patch_flex_make_mask():
         )
 
     for n in tuple(sys.modules):
-        if ".modeling_" in n:
+        if ".modeling_" in n and "llama4" not in n:
             if hasattr(sys.modules[n], "make_flex_block_causal_mask"):
                 sys.modules[n].make_flex_block_causal_mask = (
                     patched_make_flex_block_causal_mask
                 )
-                setattr(
-                    sys.modules[n],
-                    "make_flex_block_causal_mask",
-                    patched_make_flex_block_causal_mask,
-                )
 
     transformers.integrations.flex_attention.make_flex_block_causal_mask = (
         patched_make_flex_block_causal_mask

src/axolotl/monkeypatch/attention/ring_attn.py

@@ -6,8 +6,6 @@ package, specifically the `hf_adapter.substitute_hf_flash_attn` function to patc
 their sequence parallel version of Flash Attention 2.
 """
 
-from enum import Enum
-
 import torch
 import torch.distributed as dist
 from accelerate.logging import get_logger
@@ -18,7 +16,6 @@ from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
 configure_logging()
 LOG = get_logger(__name__)
 
-
 RING_ATTN_GROUP = None
 
 
@@ -43,22 +40,7 @@ def set_ring_attn_group(ring_attn_group: dist.ProcessGroup | None):
     RING_ATTN_GROUP = ring_attn_group
 
 
-class RingAttnFunc(str, Enum):
-    """Enum class for supported `ring-flash-attn` implementations"""
-
-    # VARLEN_RING = "varlen_ring"
-    # VARLEN_ZIGZAG = "varlen_zigzag"
-    VARLEN_LLAMA3 = "varlen_llama3"
-    BATCH_RING = "batch_ring"
-    BATCH_ZIGZAG = "batch_zigzag"
-    BATCH_STRIPE = "batch_stripe"
-
-
-def register_ring_attn(
-    sequence_parallel_degree: int,
-    heads_k_stride: int | None,
-    ring_attn_func: RingAttnFunc | None,
-):
+def register_ring_attn(sequence_parallel_degree: int, heads_k_stride: int | None):
     """
     Create ring attention group and substitute flash attn with ring flash attn.
 
@@ -66,9 +48,6 @@ def register_ring_attn(
         sequence_parallel_degree: Sequence parallelism factor.
         heads_k_stride: Sequence parallelism K head stride size. Passed
             through to `ring_flash_attn.substitute_hf_flash_attn`.
-        ring_attn_func: `ring_flash_attn` ring attention implemention. If sample
-            packing is enabled, it must be a `varlen` function; otherwise, it must be a
-            `batch` function.
     """
     if get_ring_attn_group() is not None:
         LOG.info("Ring attention already registered, exiting early...")
@@ -79,9 +58,7 @@ def register_ring_attn(
         f"each sequence will be processed across {sequence_parallel_degree} GPUs"
     )
 
-    rank = dist.get_rank()
     world_size = dist.get_world_size()
-
     assert sequence_parallel_degree <= world_size, (
         f"sequence_parallel_degree ({sequence_parallel_degree}) "
         f"must be less than or equal to world_size ({world_size})"
@@ -91,8 +68,10 @@ def register_ring_attn(
         f"must evenly divide world_size ({world_size})"
     )
 
-    # Assign ranks to sequence parallel groups
+    # Detailed logging of group formation
+    rank = dist.get_rank()
     group_assignments = {}
+
     for i in range(world_size // sequence_parallel_degree):
         ring_attn_ranks = list(
             range(
@@ -113,37 +92,35 @@ def register_ring_attn(
     if rank == 0:
         LOG.info(f"Sequence parallel group assignments: {group_assignments}")
 
-    if ring_attn_func is RingAttnFunc.VARLEN_LLAMA3:
-        from ring_flash_attn import substitute_hf_flash_attn
+    if heads_k_stride is None:
+        heads_k_stride = 1
 
-        substitute_hf_flash_attn(
-            process_group=get_ring_attn_group(), heads_k_stride=heads_k_stride or 1
-        )
-    elif ring_attn_func in [
-        RingAttnFunc.BATCH_RING,
-        RingAttnFunc.BATCH_ZIGZAG,
-        RingAttnFunc.BATCH_STRIPE,
-    ]:
-        from axolotl.monkeypatch.attention.ring_attn.adapters.batch import (
-            substitute_hf_flash_attn,
-        )
+    from ring_flash_attn import substitute_hf_flash_attn
 
-        substitute_hf_flash_attn(
-            process_group=get_ring_attn_group(),
-            ring_attn_func=ring_attn_func,
-        )
+    substitute_hf_flash_attn(
+        process_group=get_ring_attn_group(), heads_k_stride=heads_k_stride
+    )
 
 
-def update_ring_attn_params(position_ids: torch.Tensor | None):
+def update_ring_attn_params(batch: dict[str, torch.Tensor]):
     """
     Calculate the cumulative sequence lengths for the current forward pass and pass the
     value to the substituted `ring_flash_attn`.
 
     Args:
-        position_ids: Optional tensor of position IDs (for sample packed data).
+        batch: A dictionary with a batch of data. May or may not contain `position_ids`
+            data; if not, we compute it.
     """
     from ring_flash_attn import update_ring_flash_attn_params
 
+    input_ids = batch["input_ids"]
+    position_ids = batch.get("position_ids")
+    if position_ids is None:
+        seq_len = input_ids.shape[1]
+        position_ids = torch.arange(
+            0, seq_len, dtype=torch.long, device=input_ids.device
+        ).unsqueeze(0)
+
     cu_seqlens, _ = get_cu_seqlens_from_pos_ids(position_ids)
     cu_seqlens = cu_seqlens.squeeze().to(device=torch.cuda.current_device())
     update_ring_flash_attn_params(cu_seqlens, get_ring_attn_group())

src/axolotl/monkeypatch/attention/ring_attn/__init__.py

@@ -1,12 +0,0 @@
-"""Init for ring attention monkeypatch module"""
-
-# pylint: disable=unused-import
-# flake8: noqa
-
-from .patch import (
-    RingAttnFunc,
-    get_ring_attn_group,
-    register_ring_attn,
-    set_ring_attn_group,
-    update_ring_attn_params,
-)

src/axolotl/monkeypatch/attention/ring_attn/adapters/batch.py

@@ -1,192 +0,0 @@
-"""
-HuggingFace flash attention adapter for basic ring attention (batch API).
-
-Inspired by
-https://github.com/zhuzilin/ring-flash-attention/blob/ce9fd3935ca0e5f0592bb0826cbed18ec69da729/ring_flash_attn/adapters/hf_adapter.py.
-Our implementation closely follows the structure of that module, but we've minified it
-somewhat to support only the latest versions of transformers.
-"""
-
-# pylint: disable=protected-access,cyclic-import
-
-import os
-from typing import Callable
-
-import torch
-import torch.distributed as dist
-import transformers
-import transformers.modeling_flash_attention_utils
-from ring_flash_attn import (
-    ring_flash_attn_func,
-    stripe_flash_attn_func,
-    zigzag_ring_flash_attn_func,
-)
-from ring_flash_attn.adapters.hf_adapter import check_params
-from transformers.modeling_flash_attention_utils import (
-    _flash_supports_window_size,
-    is_flash_attn_greater_or_equal,
-)
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
-
-from axolotl.monkeypatch.attention.ring_attn.patch import RingAttnFunc
-
-RING_ATTN_FUNC_MAPPING = {
-    RingAttnFunc.BATCH_RING: ring_flash_attn_func,
-    RingAttnFunc.BATCH_ZIGZAG: zigzag_ring_flash_attn_func,
-    RingAttnFunc.BATCH_STRIPE: stripe_flash_attn_func,
-}
-
-
-def create_flash_attn_forward(
-    process_group: dist.ProcessGroup, ring_attn_func: RingAttnFunc
-) -> Callable:
-    """
-    Create a ring flash attention forward function compatible with HuggingFace's
-    interface.
-
-    Args:
-        process_group: A PyTorch distributed process group.
-        ring_attn_func: Function from `ring_flash_attention` to replace HF flash
-            attention with.
-
-    Returns:
-        A function that implements the ring flash attention forward pass with the
-            signature expected by HuggingFace Transformers.
-    """
-
-    # transformers 4.48+
-    # pylint: disable=unused-argument
-    def _flash_attention_forward(
-        query_states: torch.Tensor,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-        query_length: int,
-        is_causal: bool,
-        dropout: float = 0.0,
-        position_ids: torch.Tensor | None = None,
-        softmax_scale: float | None = None,
-        sliding_window: int | None = None,
-        use_top_left_mask: bool = False,
-        softcap: float | None = None,
-        deterministic: bool = None,
-        cu_seq_lens_q: torch.LongTensor | None = None,
-        cu_seq_lens_k: torch.LongTensor | None = None,
-        max_length_q: int | None = None,
-        max_length_k: int | None = None,
-        target_dtype: torch.dtype | None = None,
-        **kwargs,
-    ):
-        """
-        Calls the forward method of Ring Flash Attention.
-
-        Args:
-            query_states: Tensor containing the query vectors.
-            key_states: Tensor containing the key vectors.
-            value_states: Tensor containing the value vectors.
-            attention_mask: Not used in this implementation.
-            query_length: Integer representing the length of the query sequence.
-            is_causal: Boolean indicating whether to apply a causal mask to the attention.
-            dropout: Float representing the dropout probability. Default is 0.0.
-            position_ids: Not used in this implementation.
-            softmax_scale: Optional float value for the softmax scaling factor. Default is None.
-            sliding_window: Optional integer defining the size of the sliding attention window.
-                Default is None.
-            use_top_left_mask: Boolean indicating whether to use a top-left mask for the attention.
-                Default is False.
-            softcap: Not used in this implementation.
-            deterministic: Optional boolean to enforce deterministic computation. Default is None.
-            cu_seq_lens_q: Not used in this implementation.
-            cu_seq_lens_k: Not used in this implementation.
-            max_length_q: Not used in this implementation.
-            max_length_k: Not used in this implementation.
-            target_dtype: Not used in this implementation.
-            **kwargs: Additional keyword arguments. Not used in this implementation.
-
-        Returns:
-            torch.Tensor: The output of the attention mechanism, with shape
-                `[batch_size, query_length, num_heads, head_dim]`.
-        """
-        if not use_top_left_mask:
-            causal = is_causal
-        else:
-            causal = is_causal and query_length != 1
-
-        # Handle sliding window
-        use_sliding_windows = (
-            _flash_supports_window_size
-            and sliding_window is not None
-            and key_states.shape[1] > sliding_window
-        )
-        window_size = (
-            (sliding_window, sliding_window) if use_sliding_windows else (-1, -1)
-        )
-
-        # Handle deterministic mode
-        if is_flash_attn_greater_or_equal("2.4.1"):
-            if deterministic is None:
-                deterministic = (
-                    os.environ.get("FLASH_ATTENTION_DETERMINISTIC", "0") == "1"
-                )
-
-        # Call ring flash attention function
-        attn_output = RING_ATTN_FUNC_MAPPING[ring_attn_func](
-            query_states,
-            key_states,
-            value_states,
-            dropout_p=dropout,
-            softmax_scale=softmax_scale,
-            causal=causal,
-            window_size=window_size,
-            alibi_slopes=None,
-            deterministic=deterministic,
-            return_attn_probs=False,
-            group=process_group,
-        )
-
-        return attn_output
-
-    return _flash_attention_forward
-
-
-def substitute_hf_flash_attn(
-    process_group: dist.ProcessGroup, ring_attn_func: RingAttnFunc
-):
-    """
-    Substitute HuggingFace's flash attention implementation with ring-based implementation.
-
-    Args:
-        process_group: PyTorch distributed process group for communication.
-        ring_attn_func: Function from `ring_flash_attention` to replace HF flash
-            attention with.
-    """
-    try:
-        # Substitute flash attention
-        old_flash_attention_forward = (
-            transformers.modeling_flash_attention_utils._flash_attention_forward
-        )
-        new_flash_attention_forward = create_flash_attn_forward(
-            process_group=process_group, ring_attn_func=ring_attn_func
-        )
-
-        if check_params(old_flash_attention_forward, new_flash_attention_forward):
-            transformers.modeling_flash_attention_utils._flash_attention_forward = (
-                new_flash_attention_forward
-            )
-        else:
-            raise ValueError(
-                "The signature of the new flash attention forward function does not match the old one."
-            )
-    except Exception as exception:
-        raise ValueError(
-            f"The current transformer version {transformers.__version__} is not supported. "
-            "Please use pip install -U transformers to upgrade to the latest version. "
-            "If the code failed with the latest version, "
-            f"please file an issue."
-        ) from exception
-
-    # Register with ALL_ATTENTION_FUNCTIONS if available
-    if ALL_ATTENTION_FUNCTIONS is not None:
-        from ring_flash_attn.adapters.hf_adapter import flash_attention_forward
-
-        ALL_ATTENTION_FUNCTIONS["flash_attention_2"] = flash_attention_forward

src/axolotl/monkeypatch/models/llama4/modeling.py

@@ -93,20 +93,9 @@ def patch_llama4_linearized_modeling():
     """
     from transformers.models.llama4 import modeling_llama4
 
-    old_lamma_4_text_experts = modeling_llama4.Llama4TextExperts
     modeling_llama4.Llama4TextExperts = Llama4TextExperts
     setattr(
         sys.modules["transformers.models.llama4"],
         "Llama4TextExperts",
         Llama4TextExperts,
     )
-
-    def unpatch():
-        modeling_llama4.Llama4TextExperts = old_lamma_4_text_experts
-        setattr(
-            sys.modules["transformers.models.llama4"],
-            "Llama4TextExperts",
-            old_lamma_4_text_experts,
-        )
-
-    return unpatch

src/axolotl/monkeypatch/multipack.py

@@ -31,8 +31,6 @@ SUPPORTED_MULTIPACK_MODEL_TYPES = [
     "starcoder2",
     "deepseek_v2",
     "deepseek_v3",
-    "glm",
-    "glm4",
 ]
 
 

src/axolotl/monkeypatch/relora.py

@@ -272,7 +272,7 @@ class ReLoRAScheduler(LRScheduler):
         self.warmup_steps = warmup_steps
         self.anneal_steps = anneal_steps
         self.min_lr_scale = min_lr_scale
-        super().__init__(optimizer, inner_schedule.last_epoch)
+        super().__init__(optimizer, inner_schedule.last_epoch, inner_schedule.verbose)
 
     def get_lr(self) -> float:
         self.inner_schedule.last_epoch = self.last_epoch

src/axolotl/monkeypatch/trainer_eval_guard.py

@@ -1,78 +0,0 @@
-"""
-fix for FSDP2 evals when using torch.compile
-"""
-
-import inspect
-import logging
-
-from transformers import Trainer
-
-from axolotl.monkeypatch.utils import detab_code
-
-LOG = logging.getLogger(__name__)
-
-ORIGINAL_TRAINER_CODE = """
-    model.eval()
-"""
-
-PATCHED_TRAINER_CODE = """
-    if hasattr(model, "eval") and callable(model.eval):
-        self.model.eval()
-"""
-
-
-def get_evaluation_loop_code() -> str:
-    training_loop = inspect.getsource(Trainer.evaluation_loop)
-    return training_loop
-
-
-def check_evaluation_loop_is_patchable() -> bool:
-    eval_loop = get_evaluation_loop_code()
-    eval_loop, _ = detab_code(eval_loop)
-    return ORIGINAL_TRAINER_CODE in eval_loop
-
-
-def patch_evaluation_loop_for_fsdp2():
-    """
-    monkeypatch for fixing the eval loop for fsdp2 with torch.compile
-    """
-
-    try:
-        evaluation_loop = get_evaluation_loop_code()
-    except OSError:
-        return
-    Trainer._original_evaluation_loop = (  # pylint: disable=protected-access
-        evaluation_loop
-    )
-    evaluation_loop, _ = detab_code(evaluation_loop)
-    if ORIGINAL_TRAINER_CODE not in evaluation_loop:
-        return
-
-    evaluation_loop = evaluation_loop.replace(
-        ORIGINAL_TRAINER_CODE, PATCHED_TRAINER_CODE
-    )
-    evaluation_loop = evaluation_loop.replace(
-        "def evaluation_loop(",
-        "def _fixed_evaluation_loop(",
-        1,
-    )
-
-    # load imports necessary
-    import transformers.trainer
-
-    items_to_import = []
-    for item in dir(transformers.trainer):
-        if item in evaluation_loop:
-            items_to_import.append(item)
-
-    exec(  # pylint: disable=exec-used  # nosec B102
-        "from transformers.trainer import ("
-        + ", ".join(x for x in items_to_import)
-        + ")",
-        globals(),
-    )
-    exec(evaluation_loop, globals())  # pylint: disable=exec-used  # nosec B102
-    LOG.info("patching _inner_training_loop for fsdp optimizer save")
-    Trainer.evaluation_loop = (  # pylint: disable=protected-access
-        _fixed_evaluation_loop  # pylint: disable=undefined-variable  # noqa: F821
-    )

src/axolotl/train.py

@@ -6,7 +6,6 @@ import os
 import signal
 import sys
 import weakref
-from contextlib import nullcontext
 from pathlib import Path
 from typing import Any, Dict
 
@@ -26,9 +25,6 @@ from axolotl.contribs.lgpl import (  # pylint: disable = no-name-in-module
     fix_untrained_tokens,
 )
 from axolotl.core.trainer_builder import HFCausalTrainerBuilder, HFRLTrainerBuilder
-from axolotl.core.trainers.mixins.sequence_parallel import (
-    SequenceParallelContextManager,
-)
 from axolotl.logging_config import configure_logging
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.distributed import cleanup_distributed
@@ -85,11 +81,6 @@ def setup_model_and_tokenizer(
     # Apply freezing if specified
     if cfg.unfrozen_parameters:
         freeze_layers_except(model, cfg.unfrozen_parameters)
-        if any(
-            any(embed in param for embed in ["lm_head", "embed_tokens"])
-            for param in cfg.unfrozen_parameters
-        ):
-            model.enable_input_require_grads()
 
     return model, tokenizer, peft_config, processor
 
@@ -189,28 +180,16 @@ def execute_training(
         trainer: The configured trainer object.
         resume_from_checkpoint: Path to checkpoint to resume from, if applicable.
     """
-    # Define the context managers to use
-    flash_context = (
-        torch.backends.cuda.sdp_kernel(
+    LOG.info("Starting trainer...")
+    if cfg.flash_optimum:
+        with torch.backends.cuda.sdp_kernel(
+            # TODO configure these from the YAML w/ sdp_kernel_kwargs: ...
             enable_flash=True,
             enable_math=True,
             enable_mem_efficient=True,
-        )
-        if cfg.flash_optimum
-        else nullcontext()
-    )
-    sequence_parallel_context = (
-        SequenceParallelContextManager(
-            model=trainer.model,
-            sequence_parallel_degree=cfg.sequence_parallel_degree,
-            ring_attn_func=cfg.ring_attn_func,
-        )
-        if cfg.sequence_parallel_degree > 1
-        else nullcontext()
-    )
-
-    LOG.info("Starting trainer...")
-    with flash_context, sequence_parallel_context:
+        ):
+            trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+    else:
         trainer.train(resume_from_checkpoint=resume_from_checkpoint)
 
 

src/axolotl/utils/collators/batching.py

@@ -1,12 +1,19 @@
-"""Data collators for axolotl to pad labels and position_ids for packed sequences"""
+"""
+Data collators for axolotl to pad labels and position_ids for packed sequences. Also
+includes logic for handling sequence parallelism collation.
+"""
 
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Optional, Union
 
 import numpy as np
+import torch
+import torch.distributed as dist
 from transformers import PreTrainedTokenizerBase
 from transformers.utils import PaddingStrategy
 
+from axolotl.monkeypatch.attention.ring_attn import update_ring_attn_params
+
 
 @dataclass
 class DataCollatorForSeq2Seq:
@@ -41,16 +48,28 @@ class DataCollatorForSeq2Seq:
             The id to use when padding the labels (-100 will be automatically ignored by PyTorch loss functions).
         return_tensors (`str`):
             The type of Tensor to return. Allowable values are "np", "pt" and "tf".
+        sequence_parallel_degree (`int`):
+            The degree of sequence parallelism. Default to 1 for no sequence parallelism.
     """
 
     tokenizer: PreTrainedTokenizerBase
-    model: Any | None = None
-    padding: bool | str | PaddingStrategy = True
-    max_length: int | None = None
-    pad_to_multiple_of: int | None = None
+    model: Optional[Any] = None
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
     label_pad_token_id: int = -100
     position_pad_token_id: int = 0
     return_tensors: str = "pt"
+    sequence_parallel_degree: int = 1
+
+    def __post_init__(self):
+        if self.sequence_parallel_degree > 1:
+            from axolotl.monkeypatch.attention.ring_attn import get_ring_attn_group
+
+            # Get information about our position in the SP group
+            sp_group = get_ring_attn_group()
+            self.local_rank = dist.get_rank(group=sp_group)
+            self.local_world_size = dist.get_world_size(group=sp_group)
 
     def __call__(self, features, return_tensors=None):
         has_attn_mask = "attention_mask" in features[0].keys()
@@ -120,8 +139,40 @@ class DataCollatorForSeq2Seq:
             )
             features["decoder_input_ids"] = decoder_input_ids
 
+        if self.sequence_parallel_degree > 1:
+            features = self.apply_sequence_parallelism(features)
+
         return features
 
+    def apply_sequence_parallelism(
+        self, batch: dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        Apply sequence parallelism slicing to a batch.
+
+        Args:
+            batch: Batch dictionary from parent collator.
+
+        Returns:
+            Sliced batch dictionary.
+        """
+        # Get local (start, end) for sequence parallelism slicing
+        total_seq_len = batch["input_ids"].shape[1]
+        slice_size = total_seq_len // self.local_world_size
+        start = self.local_rank * slice_size
+        end = start + slice_size
+
+        # Update params for ring attention calculation
+        update_ring_attn_params(batch=batch)
+
+        # Slice batch for sequence parallel processing
+        keys_to_slice = ["input_ids", "attention_mask", "labels", "position_ids"]
+        for key in keys_to_slice:
+            if key in batch:
+                batch[key] = batch[key][:, start:end]
+
+        return batch
+
 
 @dataclass
 class BatchSamplerDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):

src/axolotl/utils/config/__init__.py

@@ -126,6 +126,9 @@ def normalize_config(cfg):
             with open(ds_config_path, encoding="utf-8") as f:
                 cfg.deepspeed = json.load(f)
 
+    if cfg.sequence_parallel_degree is None:
+        cfg.sequence_parallel_degree = 1
+
     if cfg.saves_per_epoch:
         save_steps = 1.0 / (cfg.saves_per_epoch * cfg.num_epochs)
         if save_steps < 1.0:  # prevent saves on every step

src/axolotl/utils/data/sft.py

@@ -3,7 +3,6 @@
 import functools
 import logging
 import os
-import tempfile
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
 
@@ -118,26 +117,9 @@ def prepare_dataset(cfg, tokenizer, processor=None, preprocess_iterable=None):
             cfg.pretraining_dataset[0]["type"] or "pretrain",
         )
 
-        # when letting accelerator dispatch batches from the main process, we don't need to load the dataset from
-        # other ranks, we just need to present a fake dataset
-        if (
-            cfg.accelerator_config
-            and cfg.accelerator_config.dispatch_batches
-            and not is_local_main_process()
-        ):
-            with tempfile.NamedTemporaryFile(mode="w+", delete=False) as f:
-                f.write("text\n")
-                f.write("lorem ipsum dolor sit amet\n")
-                # rewind the file pointer to the beginning so we can read it again
-                f.seek(0)
-                iter_ds = load_dataset(
-                    "csv", data_files=f.name, split="train", streaming=True
-                )
-        else:
-            iter_ds = load_dataset(
-                path, streaming=True, split=split, name=name, data_files=data_files
-            )
-
+        iter_ds = load_dataset(
+            path, streaming=True, split=split, name=name, data_files=data_files
+        )
         if skip:
             LOG.info(f"Skipping {skip} samples from the dataset")
             iter_ds = iter_ds.skip(skip)
@@ -350,23 +332,16 @@ def load_tokenized_prepared_datasets(
         if cfg.local_rank == 0 and not cfg.skip_prepare_dataset:
             LOG.info(f"Saving merged prepared dataset to disk... {prepared_ds_path}")
             if isinstance(dataset, IterableDataset):
-                num_workers = cfg.dataset_processes
 
-                def gen_from_iter_ds(_ds, worker_id: List[int], num_workers: List[int]):
-                    """Generator function to correctly splice the dataset for each worker"""
-                    for i, item in enumerate(_ds):
-                        if i % num_workers[0] == worker_id[0]:
-                            yield item
+                def gen_from_iter_ds(_ds, _=None):
+                    yield from _ds
 
                 ds_from_iter = Dataset.from_generator(
                     functools.partial(gen_from_iter_ds, dataset),
                     features=dataset.features,
-                    num_proc=num_workers,
+                    num_proc=cfg.dataset_processes,
                     split=split,
-                    gen_kwargs={
-                        "worker_id": list(range(num_workers)),
-                        "num_workers": [num_workers] * num_workers,
-                    },
+                    gen_kwargs={"_": list(range(cfg.dataset_processes))},
                 )
                 ds_from_iter.save_to_disk(str(prepared_ds_path))
             else:

src/axolotl/utils/freeze.py

@@ -2,14 +2,13 @@
 module to freeze/unfreeze parameters by name
 """
 
+import logging
 import re
 from typing import Callable, List, Tuple, Union
 
-from accelerate.logging import get_logger
-
 from axolotl.utils.distributed import is_main_process
 
-LOG = get_logger(__name__)
+LOG = logging.getLogger("axolotl.utils.freeze")
 
 
 def freeze_layers_except(model, regex_patterns):
@@ -185,7 +184,7 @@ class LayerNamePattern:
         """
         self.raw_pattern = pattern
         name_pattern, self.range = self._parse_pattern(pattern)
-        self.name_regex = re.compile(re.sub(r"\.(?!\+)", "\\.", name_pattern))
+        self.name_regex = re.compile(name_pattern.replace(".", "\\."))
 
     def match(self, name: str) -> bool:
         """

src/axolotl/utils/gradient_checkpointing/__init__.py

@@ -1,7 +1,5 @@
 """custom checkpointing utils"""
 
-from functools import partial
-
 from axolotl.utils.gradient_checkpointing.unsloth import (
     Unsloth_Offloaded_Gradient_Checkpointer,
 )
@@ -11,10 +9,6 @@ def hf_grad_checkpoint_offload_wrapper(
     decoder_layer, *args, use_reentrant=None
 ):  # pylint: disable=unused-argument
     return Unsloth_Offloaded_Gradient_Checkpointer.apply(
-        (
-            decoder_layer.func.__self__
-            if isinstance(decoder_layer, partial)
-            else decoder_layer.__self__
-        ),
+        decoder_layer.__self__,
         *args,
     )

src/axolotl/utils/models.py

@@ -542,17 +542,6 @@ class ModelLoader:
             from axolotl.monkeypatch.accelerate.fsdp2 import patch_accelerate_fsdp_utils
 
             patch_accelerate_fsdp_utils()
-
-        if self.cfg.flex_attention:
-            from axolotl.monkeypatch.attention.flex_attn import (
-                patch_flex_make_mask,
-                patch_flex_wrapper,
-            )
-
-            flex_attn_compile_kwargs = self.cfg.flex_attn_compile_kwargs or {}
-            patch_flex_wrapper(**flex_attn_compile_kwargs)
-            patch_flex_make_mask()
-
         # patch gemma3 conditional generation forward before loading plugins
         # as it could be overridden by plugins
         if self.cfg.model_config_type == "llama4":
@@ -655,7 +644,6 @@ class ModelLoader:
             register_ring_attn(
                 sequence_parallel_degree=self.cfg.sequence_parallel_degree,
                 heads_k_stride=self.cfg.heads_k_stride,
-                ring_attn_func=self.cfg.ring_attn_func,
             )
 
     def patch_attention(self) -> None:
@@ -917,6 +905,13 @@ class ModelLoader:
             self.model_config._attn_implementation = (  # pylint: disable=protected-access
                 "flex_attention"
             )
+            from axolotl.monkeypatch.attention.flex_attn import (
+                patch_flex_make_mask,
+                patch_flex_wrapper,
+            )
+
+            patch_flex_wrapper()
+            patch_flex_make_mask()
 
         elif self.cfg.flash_attention:
             if not self.cfg.sample_packing and self.cfg.s2_attention:
@@ -1120,7 +1115,7 @@ class ModelLoader:
 
         return skip_move_to_device
 
-    def adjust_model_config(self) -> None:
+    def ajust_model_config(self) -> None:
         if (
             hasattr(self.model, "config")
             and hasattr(self.model.config, "max_position_embeddings")
@@ -1280,7 +1275,7 @@ class ModelLoader:
         else:
             self.model.tie_weights()
 
-        self.adjust_model_config()
+        self.ajust_model_config()
 
         # log device memory usage
         if hasattr(self.model, "device") and self.model.device.type in (

src/axolotl/utils/schedulers.py

@@ -40,7 +40,7 @@ class RexLR(LRScheduler):
         self.max_lr = max_lr
         self.total_steps = total_steps
         self.num_warmup_steps = num_warmup_steps
-        self.last_step = max(last_step - 1, 0)
+        self.last_step = last_step - 1
 
         # Ensure each parameter group has an "initial_lr" key to avoid issues when resuming.
         for group in optimizer.param_groups:

src/axolotl/utils/schemas/config.py

@@ -225,7 +225,6 @@ class AxolotlInputConfig(
     sdp_attention: bool | None = None
     s2_attention: bool | None = None
     flex_attention: bool | None = None
-    flex_attn_compile_kwargs: dict[str, Any] | None = None
     flash_attention: bool | None = None
     flash_attn_cross_entropy: bool | None = None
     flash_attn_rms_norm: bool | None = None
@@ -259,7 +258,6 @@ class AxolotlInputConfig(
 
     sequence_parallel_degree: int | None = None
     heads_k_stride: int | None = None
-    ring_attn_func: str | None = None
 
     special_tokens: SpecialTokensConfig | None = None
     tokens: list[str] | None = None
@@ -660,7 +658,6 @@ class AxolotlInputConfig(
             data.get("val_set_size") == 0
             and (data.get("eval_steps") or data.get("eval_strategy"))
             and not data.get("test_datasets")
-            and data.get("eval_strategy") != "no"
         ):
             raise ValueError(
                 "eval_steps and eval_strategy are not supported with val_set_size == 0"
@@ -1149,19 +1146,21 @@ class AxolotlInputConfig(
 
         return data
 
-    @model_validator(mode="after")
-    def check_sequence_parallel_degree(self):
-        if not self.sequence_parallel_degree:
-            self.sequence_parallel_degree = 1
-        elif self.sequence_parallel_degree > 1:
-            if not self.flash_attention:
+    @field_validator("sequence_parallel_degree", mode="before")
+    @classmethod
+    def check_sequence_parallel_degree(cls, value, info):
+        if not value:
+            value = 1
+
+        if value > 1:
+            if not info.data.get("flash_attention"):
                 raise ValueError(
                     "flash_attention: true must be set with sequence_parallel_degree > 1"
                 )
 
-            if self.sample_packing and self.micro_batch_size > 1:
+            if not info.data["micro_batch_size"] == 1:
                 raise ValueError(
-                    "micro_batch_size must be set to 1 when sample_packing is enabled"
+                    "micro_batch_size must be set to 1 "
                     "due to a `ring-flash-attn` requirement"
                 )
 
@@ -1179,40 +1178,14 @@ class AxolotlInputConfig(
             # according to the proportion of non-padding tokens per rank.
             LOG.warning(
                 "Sequence parallelism (SP) is enabled with "
-                f"sequence_parallel_degree={self.sequence_parallel_degree}. "
-                "Please note that logged losses may differ slightly to the non-SP "
-                "losses due to transformers Trainer implementation details. "
-                "Please see https://github.com/axolotl-ai-cloud/axolotl/pull/2495#issuecomment-2784022042 "
+                f"sequence_parallel_degree={value}. Please note that logged losses may "
+                "differ slightly to the non-SP losses due to transformers Trainer "
+                "implementation details. Please see "
+                "https://github.com/axolotl-ai-cloud/axolotl/pull/2495#issuecomment-2784022042 "
                 "for more details."
             )
 
-        return self
-
-    @model_validator(mode="after")
-    def validate_ring_attn_func(self):
-        if getattr(self, "sequence_parallel_degree", 1) == 1:
-            return self
-
-        from axolotl.monkeypatch.attention.ring_attn.patch import RingAttnFunc
-
-        if self.ring_attn_func is not None:
-            valid_funcs = list(RingAttnFunc)
-            if self.ring_attn_func in valid_funcs:
-                self.ring_attn_func = RingAttnFunc(self.ring_attn_func)
-            else:
-                raise ValueError(
-                    f"ring_attn_func: {self.ring_attn_func} must be in {valid_funcs}"
-                )
-        else:
-            # Default ring attention function selection
-            sample_packing = getattr(self, "sample_packing", False)
-            self.ring_attn_func = (
-                RingAttnFunc.VARLEN_LLAMA3
-                if sample_packing
-                else RingAttnFunc.BATCH_RING
-            )
-
-        return self
+        return value
 
     @model_validator(mode="before")
     @classmethod
@@ -1303,14 +1276,11 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
         ):
             capabilities = data.get("capabilities")
             is_fsdp = data.get("fsdp") is not None
-            is_fsdp2 = (
-                data.get("fsdp_config") is not None
-                and str(data.get("fsdp_config").get("fsdp_version")) == "2"
-            )
-            if capabilities and capabilities.get("n_gpu", 0) > 1 and not is_fsdp2:
+
+            if capabilities and capabilities.get("n_gpu", 0) > 1:
                 if is_fsdp:
                     raise ValueError(
-                        "lora_mlp_kernel, lora_qkv_kernel, and lora_o_kernel are not compatible with FSDP1."
+                        "lora_mlp_kernel, lora_qkv_kernel, and lora_o_kernel are not compatible with FSDP."
                     )
         return data
 

src/axolotl/utils/schemas/vllm.py

@@ -36,11 +36,3 @@ class VllmConfig(BaseModel):
         default=None,
         json_schema_extra={"description": "Enable prefix caching for VLLM"},
     )
-    host: str | None = Field(
-        default="0.0.0.0",  # nosec B104
-        json_schema_extra={"description": "Host for the vLLM server to start on"},
-    )
-    port: int | None = Field(
-        default=8000,
-        json_schema_extra={"description": "Port of the vLLM server to start on"},
-    )

src/axolotl/utils/trainer.py

@@ -17,7 +17,6 @@ from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
 from transformers.utils import is_torch_bf16_gpu_available
 
 from axolotl.core.trainer_builder import HFCausalTrainerBuilder, HFRLTrainerBuilder
-from axolotl.monkeypatch.trainer_eval_guard import patch_evaluation_loop_for_fsdp2
 from axolotl.utils.distributed import reduce_and_broadcast
 from axolotl.utils.environment import check_cuda_p2p_ib_support
 from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
@@ -236,8 +235,7 @@ def drop_long_seq(sample, sequence_len=2048, min_sequence_len=2):
 
 
 def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
-    drop_attn_mask = cfg.model_config_type in ["mamba", "gemma3"]
-    if drop_attn_mask:
+    if cfg.model_config_type in ["mamba", "gemma3"]:
         LOG.info("dropping attention_mask column")
         train_dataset = train_dataset.remove_columns("attention_mask")
         if eval_dataset:
@@ -348,7 +346,7 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
                     load_from_cache_file=not cfg.is_preprocess,
                     desc="Add position_id column (PoSE)",
                 )
-    elif cfg.sample_packing:
+    elif cfg.sample_packing or cfg.sequence_parallel_degree > 1:
         drop_long_kwargs = {}
         if filter_map_kwargs:
             drop_long_kwargs["desc"] = "Add position_id column (Sample Packing)"
@@ -358,7 +356,7 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
             **filter_map_kwargs,
             **drop_long_kwargs,
         )
-        if cfg.eval_sample_packing:
+        if cfg.eval_sample_packing or cfg.sequence_parallel_degree > 1:
             if eval_dataset:
                 eval_dataset = eval_dataset.map(
                     add_position_ids,
@@ -528,13 +526,6 @@ def setup_torch_compile_env(cfg):
 def setup_deepspeed_env(cfg, stage=None):
     from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig
 
-    from axolotl.utils.distributed import distributed_state
-
-    if distributed_state and distributed_state.initialized:
-        raise RuntimeError(
-            "Distributed State already initialized before Deepspeed setup"
-        )
-
     os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
     os.environ["ACCELERATE_DEEPSPEED_CONFIG_FILE"] = cfg.deepspeed
     if stage:
@@ -634,12 +625,6 @@ def setup_trainer(
         A trainer instance (either `HFRLTrainer` or `HFCausalTrainer`) configured based
             on the provided parameters.
     """
-    if (
-        cfg.torch_compile
-        and cfg.fsdp_config
-        and str(cfg.fsdp_config.fsdp_version) == "2"
-    ):
-        patch_evaluation_loop_for_fsdp2()
     if cfg.rl:
         trainer_builder = HFRLTrainerBuilder(cfg, model, tokenizer, processor)
         trainer_builder.model_ref = model_ref

tests/conftest.py

@@ -193,14 +193,6 @@ def download_tiny_shakespeare_dataset():
     snapshot_download_w_retry("winglian/tiny-shakespeare", repo_type="dataset")
 
 
-@pytest.fixture(scope="session", autouse=True)
-def download_evolkit_kd_sample_dataset():
-    # download the dataset
-    snapshot_download_w_retry(
-        "axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample", repo_type="dataset"
-    )
-
-
 @pytest.fixture(scope="session", autouse=True)
 def download_deepseek_model_fixture():
     snapshot_download_w_retry("axolotl-ai-co/DeepSeek-V3-11M", repo_type="model")
@@ -216,16 +208,6 @@ def download_huggyllama_model_fixture():
     )
 
 
-@pytest.fixture(scope="session", autouse=True)
-def download_llama33_70b_model_fixture():
-    # download the tokenizer only
-    snapshot_download_w_retry(
-        "axolotl-ai-co/Llama-3.3-70B-Instruct-tokenizer",
-        repo_type="model",
-        allow_patterns=["*token*", "config.json"],
-    )
-
-
 @pytest.fixture(scope="session", autouse=True)
 def download_llama_1b_model_fixture():
     # download the tokenizer only
@@ -333,14 +315,6 @@ def download_llama2_model_fixture():
     )
 
 
-@pytest.fixture(scope="session", autouse=True)
-def download_llama32_1b_model_fixture():
-    snapshot_download_w_retry(
-        "osllmai-community/Llama-3.2-1B",
-        repo_type="model",
-    )
-
-
 @pytest.fixture
 @enable_hf_offline
 def tokenizer_huggyllama(

tests/e2e/integrations/test_cut_cross_entropy.py

@@ -8,7 +8,7 @@ from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils import get_pytorch_version
-from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
+from axolotl.utils.config import normalize_config, prepare_plugins
 from axolotl.utils.dict import DictDefault
 
 from ..utils import check_model_output_exists
@@ -56,7 +56,6 @@ class TestCutCrossEntropyIntegration:
     # pylint: disable=redefined-outer-name
     def test_llama_w_cce(self, min_cfg, temp_dir):
         cfg = DictDefault(min_cfg)
-        cfg = validate_config(cfg)
         prepare_plugins(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
@@ -102,7 +101,6 @@ class TestCutCrossEntropyIntegration:
                 "bf16": "auto",
             }
         )
-        cfg = validate_config(cfg)
         prepare_plugins(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
@@ -131,7 +129,6 @@ class TestCutCrossEntropyIntegration:
                 attention_type: True,
             }
         )
-        cfg = validate_config(cfg)
         prepare_plugins(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()

tests/e2e/integrations/test_liger.py

@@ -5,7 +5,7 @@ Simple end-to-end test for Liger integration
 from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
+from axolotl.utils.config import normalize_config, prepare_plugins
 from axolotl.utils.dict import DictDefault
 
 from tests.e2e.utils import check_model_output_exists, require_torch_2_4_1
@@ -54,7 +54,6 @@ class LigerIntegrationTestCase:
             }
         )
         # pylint: disable=duplicate-code
-        cfg = validate_config(cfg)
         prepare_plugins(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
@@ -101,7 +100,6 @@ class LigerIntegrationTestCase:
             }
         )
         # pylint: disable=duplicate-code
-        cfg = validate_config(cfg)
         prepare_plugins(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()

tests/e2e/multigpu/solo/__init__.py

@@ -1,2 +0,0 @@
-# Tests under this directory should get run "solo" on their own as they
-# seem to cause issues when run in the same batch as other tests.

tests/e2e/multigpu/solo/test_flex.py

@@ -49,20 +49,18 @@ class TestPackedFlex:
                 },
                 "datasets": [
                     {
-                        "path": "tatsu-lab/alpaca",
+                        "path": "vicgalle/alpaca-gpt4",
                         "type": "alpaca",
-                        "split": "train[:10%]",
                     },
                 ],
                 "num_epochs": 1,
                 "micro_batch_size": 2,
                 "gradient_accumulation_steps": 2,
-                "gradient_checkpointing": True,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 2,
+                "max_steps": 5,
                 "use_tensorboard": True,
                 "save_strategy": "no",
             }

tests/e2e/multigpu/solo/test_grpo.py

@@ -177,7 +177,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
             "NCCL_P2P_LEVEL": "LOC",
             **current_env,
             "CUDA_VISIBLE_DEVICES": "1",
-            "VLLM_USE_V1": "0",
         }
         vllm_process_id = start_vllm(
             cfg.base_model,
@@ -265,7 +264,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
             "NCCL_P2P_LEVEL": "LOC",  # nccl can be brittle, assume P2P isn't reliable
             **current_env,
             "CUDA_VISIBLE_DEVICES": "1",
-            "VLLM_USE_V1": "0",
         }
         vllm_process_id = start_vllm(
             cfg.base_model,

tests/e2e/multigpu/test_llama.py

@@ -621,6 +621,12 @@ class TestMultiGPULlama:
             temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss is too high"
         )
 
+    # TODO: remove skip once deepspeed regression is fixed
+    # see https://github.com/huggingface/transformers/pull/37324
+    @pytest.mark.skipif(
+        transformers_version_eq("4.51.0"),
+        reason="zero3 is not supported with transformers==4.51.0",
+    )
     @pytest.mark.parametrize(
         "gradient_accumulation_steps",
         [1, 2],

tests/e2e/multigpu/test_sp.py

@@ -3,14 +3,13 @@
 import os
 from pathlib import Path
 
-import pytest
 import yaml
 from accelerate.test_utils import execute_subprocess_async
 from transformers.testing_utils import get_torch_dist_unique_port
 
 from axolotl.utils.dict import DictDefault
 
-from ...utils import check_tensorboard
+from ..utils import check_tensorboard
 
 os.environ["WANDB_DISABLED"] = "true"
 
@@ -18,15 +17,8 @@ os.environ["WANDB_DISABLED"] = "true"
 class TestSequenceParallelism:
     """Test case for training with sequence parallelism enabled"""
 
-    def _run_sequence_parallel_test(
-        self,
-        temp_dir,
-        sample_packing=True,
-        micro_batch_size=1,
-        pad_to_sequence_len=True,
-        ring_attn_func=None,
-    ):
-        """Helper method to run sequence parallel tests with different configurations"""
+    def test_sequence_parallel_training(self, temp_dir):
+        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
@@ -35,9 +27,9 @@ class TestSequenceParallelism:
                 "strict": False,
                 "sequence_len": 2048,
                 "adapter": "qlora",
-                "sample_packing": sample_packing,
-                "eval_sample_packing": sample_packing,
-                "pad_to_sequence_len": pad_to_sequence_len,
+                "sample_packing": True,
+                "eval_sample_packing": True,
+                "pad_to_sequence_len": True,
                 "lora_r": 8,
                 "lora_alpha": 16,
                 "lora_dropout": 0.05,
@@ -53,7 +45,7 @@ class TestSequenceParallelism:
                 ],
                 "num_epochs": 1,
                 "max_steps": 8,
-                "micro_batch_size": micro_batch_size,
+                "micro_batch_size": 1,
                 "gradient_accumulation_steps": 2,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
@@ -69,7 +61,6 @@ class TestSequenceParallelism:
                 "weight_decay": 0.0,
                 "use_tensorboard": True,
                 "sequence_parallel_degree": 2,
-                "ring_attn_func": ring_attn_func,
             }
         )
 
@@ -95,35 +86,3 @@ class TestSequenceParallelism:
         check_tensorboard(
             temp_dir + "/runs", "train/train_loss", 2.6, "Train Loss is too high"
         )
-
-    @pytest.mark.parametrize(
-        "sample_packing, micro_batch_size, pad_to_sequence_len, ring_attn_func",
-        [
-            (True, 1, True, None),  # defaults to varlen_llama3 ring_attn_func
-            (False, 2, True, None),  # defaults to batch_ring ring_attn_func
-            (False, 2, True, "batch_zigzag"),
-            # (False, 2, False),  # not yet working
-        ],
-        ids=[
-            "sample_packing, varlen_llama3 ring_attn_func",
-            "no sample_packing, no pad_to_sequence_len, batch_ring ring_attn_func",
-            "no sample_packing, no pad_to_sequence_len, batch_zigzag ring_attn_func",
-            # "no sample_packing, pad_to_sequence_len",  # not yet working
-        ],
-    )
-    def test_sequence_parallel_training(
-        self,
-        temp_dir,
-        sample_packing,
-        micro_batch_size,
-        pad_to_sequence_len,
-        ring_attn_func,
-    ):
-        """Test sequence parallel training with different configurations"""
-        self._run_sequence_parallel_test(
-            temp_dir,
-            sample_packing=sample_packing,
-            micro_batch_size=micro_batch_size,
-            pad_to_sequence_len=pad_to_sequence_len,
-            ring_attn_func=ring_attn_func,
-        )

tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py

@@ -144,7 +144,7 @@ def test_swiglu_mlp_integration(small_llama_model):
 def test_geglu_model_integration():
     """Test GeGLU activation with Gemma model."""
     model = AutoModelForCausalLM.from_pretrained(
-        "mhenrichsen/gemma-2b", torch_dtype=torch.float16, device_map="cuda:0"
+        "mhenrichsen/gemma-2b", torch_dtype=torch.float16, device_map="auto"
     )
     peft_config = get_peft_config(
         {
@@ -347,7 +347,7 @@ def test_model_architecture(model_config):
     """Test LoRA kernel patches across different model architectures."""
     # Load model with appropriate dtype
     model = AutoModelForCausalLM.from_pretrained(
-        model_config["name"], torch_dtype=model_config["dtype"], device_map="cuda:0"
+        model_config["name"], torch_dtype=model_config["dtype"], device_map="auto"
     )
 
     # Apply LoRA configuration

tests/e2e/patched/test_4d_multipack_llama.py

@@ -9,7 +9,7 @@ import unittest
 from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
 from ..utils import check_model_output_exists, with_temp_dir
@@ -60,7 +60,6 @@ class Test4dMultipackLlama(unittest.TestCase):
                 "fp16": True,
             }
         )
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
@@ -105,7 +104,6 @@ class Test4dMultipackLlama(unittest.TestCase):
                 "fp16": True,
             }
         )
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

tests/e2e/patched/test_activation_checkpointing.py

@@ -1,77 +0,0 @@
-"""
-E2E tests for activation checkpointing
-"""
-
-import pytest
-import transformers
-from torch.utils.checkpoint import checkpoint
-
-from axolotl.cli.args import TrainerCliArgs
-from axolotl.common.datasets import load_datasets
-from axolotl.train import train
-from axolotl.utils.config import normalize_config, validate_config
-from axolotl.utils.dict import DictDefault
-
-from ..utils import check_model_output_exists
-
-
-@pytest.fixture()
-def fix_checkpoint_after_test():
-    yield
-    transformers.modeling_utils.checkpoint = checkpoint
-
-
-class TestActivationCheckpointing:
-    """
-    E2E tests for activation checkpointing
-    """
-
-    def test_activation_checkpointing_offload(
-        self,
-        temp_dir,
-        fix_checkpoint_after_test,  # pylint: disable=unused-argument,redefined-outer-name
-    ):
-        # pylint: disable=duplicate-code
-        cfg = DictDefault(
-            {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "sequence_len": 1024,
-                "val_set_size": 0.0,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                    "eos_token": "<|im_end|>",
-                },
-                "datasets": [
-                    {
-                        "chat_template": "chatml",
-                        "path": "mlabonne/FineTome-100k",
-                        "type": "chat_template",
-                        "split": "train[:10%]",
-                        "field_messages": "conversations",
-                        "message_field_role": "from",
-                        "message_field_content": "value",
-                    },
-                ],
-                "num_epochs": 1,
-                "max_steps": 5,
-                "micro_batch_size": 1,
-                "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_8bit",
-                "lr_scheduler": "cosine",
-                "flash_attention": True,
-                "sample_packing": True,
-                "bf16": True,
-                "save_safetensors": True,
-                "gradient_checkpointing": "offload",
-            }
-        )
-
-        cfg = validate_config(cfg)
-        normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
-
-        train(cfg=cfg, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)

tests/e2e/patched/test_falcon_samplepack.py

@@ -9,7 +9,7 @@ import unittest
 from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
 from ..utils import check_model_output_exists, with_temp_dir
@@ -63,7 +63,6 @@ class TestFalconPatched(unittest.TestCase):
                 "bf16": "auto",
             }
         )
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
@@ -104,7 +103,6 @@ class TestFalconPatched(unittest.TestCase):
                 "bf16": "auto",
             }
         )
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

tests/e2e/patched/test_fused_llama.py

@@ -12,7 +12,7 @@ from transformers.utils import is_torch_bf16_gpu_available
 from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
 from ..utils import check_model_output_exists, with_temp_dir
@@ -67,7 +67,6 @@ class TestFusedLlama(unittest.TestCase):
             cfg.bf16 = True
         else:
             cfg.fp16 = True
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

tests/e2e/patched/test_llama_s2_attention.py

@@ -11,7 +11,7 @@ import pytest
 from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
 from ..utils import check_model_output_exists, with_temp_dir
@@ -65,7 +65,6 @@ class TestLlamaShiftedSparseAttention(unittest.TestCase):
             }
         )
 
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
@@ -106,7 +105,6 @@ class TestLlamaShiftedSparseAttention(unittest.TestCase):
             }
         )
 
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

tests/e2e/patched/test_lora_llama_multipack.py

@@ -12,7 +12,7 @@ from transformers.utils import is_auto_gptq_available, is_torch_bf16_gpu_availab
 from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
 from ..utils import check_model_output_exists, with_temp_dir
@@ -70,7 +70,6 @@ class TestLoraLlama(unittest.TestCase):
         else:
             cfg.fp16 = True
 
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
@@ -121,7 +120,6 @@ class TestLoraLlama(unittest.TestCase):
                 "lr_scheduler": "cosine",
             }
         )
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

tests/e2e/patched/test_mistral_samplepack.py

@@ -9,7 +9,7 @@ import unittest
 from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
 from ..utils import check_model_output_exists, with_temp_dir
@@ -63,7 +63,6 @@ class TestMistral(unittest.TestCase):
                 "bf16": "auto",
             }
         )
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
@@ -105,7 +104,6 @@ class TestMistral(unittest.TestCase):
                 "bf16": "auto",
             }
         )
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

tests/e2e/patched/test_mixtral_samplepack.py

@@ -9,7 +9,7 @@ import unittest
 from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
 from ..utils import check_model_output_exists, with_temp_dir
@@ -60,7 +60,6 @@ class TestMixtral(unittest.TestCase):
                 "bf16": "auto",
             }
         )
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
@@ -99,7 +98,6 @@ class TestMixtral(unittest.TestCase):
                 "bf16": "auto",
             }
         )
-        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)