run 2.5.1 test without waiting for 1st e2e

* reset known modules that are patched on each test function end

* fix the llama model module name

* prevent unsloth patching multiple times

* pop classes out of the globals after reset

* fix tuple indexing

* manually workaround for llama fa2

.github/workflows/base.yml

@@ -22,36 +22,38 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - cuda: "121"
-            cuda_version: 12.1.1
-            cudnn_version: 8
-            python_version: "3.10"
-            pytorch: 2.3.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "121"
-            cuda_version: 12.1.1
-            cudnn_version: 8
-            python_version: "3.11"
-            pytorch: 2.3.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "124"
-            cuda_version: 12.4.1
-            cudnn_version: ""
-            python_version: "3.10"
-            pytorch: 2.4.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "124"
-            cuda_version: 12.4.1
-            cudnn_version: ""
-            python_version: "3.11"
-            pytorch: 2.4.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+#          - cuda: "121"
+#            cuda_version: 12.1.1
+#            cudnn_version: 8
+#            python_version: "3.10"
+#            pytorch: 2.3.1
+#            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+#            from_base_img: ""
+#            from_base_tag: ""
+#          - cuda: "121"
+#            cuda_version: 12.1.1
+#            cudnn_version: 8
+#            python_version: "3.11"
+#            pytorch: 2.3.1
+#            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+#            from_base_img: ""
+#            from_base_tag: ""
+#          - cuda: "124"
+#            cuda_version: 12.4.1
+#            cudnn_version: ""
+#            python_version: "3.11"
+#            pytorch: 2.4.1
+#            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+#            from_base_img: ""
+#            from_base_tag: ""
           - cuda: "124"
             cuda_version: 12.4.1
             cudnn_version: ""
             python_version: "3.11"
             pytorch: 2.5.1
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+            from_base_img: nvcr.io/nvidia/pytorch
+            from_base_tag: 24.10-py3
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -61,7 +63,7 @@ jobs:
         with:
           images: |
             winglian/axolotl-base
-            axolotlai/axolotl-base
+#            axolotlai/axolotl-base
       - name: Login to Docker Hub
         uses: docker/login-action@v2
         with:
@@ -74,7 +76,8 @@ jobs:
         with:
           context: .
           file: ./docker/Dockerfile-base
-          push: ${{ github.event_name != 'pull_request' }}
+          push: true
+#          push: ${{ github.event_name != 'pull_request' }}
           tags: ${{ steps.metadata.outputs.tags }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}
           build-args: |
@@ -84,3 +87,5 @@ jobs:
             PYTHON_VERSION=${{ matrix.python_version }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
             TORCH_CUDA_ARCH_LIST=${{ matrix.torch_cuda_arch_list }}
+            BASE_IMAGE=${{ matrix.from_base_img || '' }}
+            BASE_TAG=${{ matrix.from_base_tag || '' }}

.github/workflows/tests.yml

@@ -148,63 +148,64 @@ jobs:
         run: |
           find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
 
-  docker-e2e-tests-1st:
-    if: ${{ ! contains(github.event.commits[0].message, '[skip e2e]') && github.repository_owner == 'axolotl-ai-cloud' }}
-    # this job needs to be run on self-hosted GPU runners...
-    runs-on: [self-hosted, modal]
-    timeout-minutes: 90
-    needs: [pre-commit, pytest, pytest-sdist]
-
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - cuda: 124
-            cuda_version: 12.4.1
-            python_version: "3.11"
-            pytorch: 2.4.1
-            num_gpus: 1
-            axolotl_extras:
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-      - name: Install Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-      - name: Install Modal
-        run: |
-          python -m pip install --upgrade pip
-          pip install modal==0.63.64 jinja2
-      - name: Update env vars
-        run: |
-          echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV
-          echo "PYTORCH_VERSION=${{ matrix.pytorch}}" >> $GITHUB_ENV
-          echo "AXOLOTL_ARGS=${{ matrix.axolotl_args}}" >> $GITHUB_ENV
-          echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV
-          echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
-          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
-      - name: Run tests job on Modal
-        run: |
-          modal run cicd.tests
+#  docker-e2e-tests-1st:
+#    if: ${{ ! contains(github.event.commits[0].message, '[skip e2e]') && github.repository_owner == 'axolotl-ai-cloud' }}
+#    # this job needs to be run on self-hosted GPU runners...
+#    runs-on: [self-hosted, modal]
+#    timeout-minutes: 90
+#    needs: [pre-commit, pytest, pytest-sdist]
+#
+#    strategy:
+#      fail-fast: false
+#      matrix:
+#        include:
+#          - cuda: 124
+#            cuda_version: 12.4.1
+#            python_version: "3.11"
+#            pytorch: 2.4.1
+#            num_gpus: 1
+#            axolotl_extras:
+#    steps:
+#      - name: Checkout
+#        uses: actions/checkout@v4
+#      - name: Install Python
+#        uses: actions/setup-python@v5
+#        with:
+#          python-version: "3.10"
+#      - name: Install Modal
+#        run: |
+#          python -m pip install --upgrade pip
+#          pip install modal==0.63.64 jinja2
+#      - name: Update env vars
+#        run: |
+#          echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV
+#          echo "PYTORCH_VERSION=${{ matrix.pytorch}}" >> $GITHUB_ENV
+#          echo "AXOLOTL_ARGS=${{ matrix.axolotl_args}}" >> $GITHUB_ENV
+#          echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV
+#          echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
+#          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
+#      - name: Run tests job on Modal
+#        run: |
+#          modal run cicd.tests
 
   docker-e2e-tests:
     if: github.repository_owner == 'axolotl-ai-cloud'
     # this job needs to be run on self-hosted GPU runners...
     runs-on: [self-hosted, modal]
     timeout-minutes: 90
-    needs: [pre-commit, pytest, docker-e2e-tests-1st]
+#    needs: [pre-commit, pytest, docker-e2e-tests-1st]
+    needs: [pre-commit, pytest]
 
     strategy:
       fail-fast: false
       matrix:
         include:
-          - cuda: 121
-            cuda_version: 12.1.1
-            python_version: "3.10"
-            pytorch: 2.3.1
-            num_gpus: 1
-            axolotl_extras: mamba-ssm
+#          - cuda: 121
+#            cuda_version: 12.1.1
+#            python_version: "3.10"
+#            pytorch: 2.3.1
+#            num_gpus: 1
+#            axolotl_extras: mamba-ssm
           - cuda: 124
             cuda_version: 12.4.1
             python_version: "3.11"
@@ -224,7 +225,7 @@ jobs:
           pip install modal==0.63.64 jinja2
       - name: Update env vars
         run: |
-          echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV
+          echo "BASE_TAG=pr-2139-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV
           echo "PYTORCH_VERSION=${{ matrix.pytorch}}" >> $GITHUB_ENV
           echo "AXOLOTL_ARGS=${{ matrix.axolotl_args}}" >> $GITHUB_ENV
           echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV

cicd/Dockerfile.jinja

@@ -1,4 +1,4 @@
-FROM axolotlai/axolotl-base:{{ BASE_TAG }}
+FROM winglian/axolotl-base:{{ BASE_TAG }}
 
 ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
 ENV AXOLOTL_EXTRAS="{{ AXOLOTL_EXTRAS }}"

docker/Dockerfile

@@ -1,5 +1,6 @@
+ARG BASE_IMAGE=axolotlai/axolotl-base
 ARG BASE_TAG=main-base
-FROM axolotlai/axolotl-base:$BASE_TAG
+FROM $BASE_IMAGE:$BASE_TAG
 
 ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
 ARG AXOLOTL_EXTRAS=""

docker/Dockerfile-base

@@ -3,7 +3,10 @@ ARG CUDNN_VERSION="8"
 ARG UBUNTU_VERSION="22.04"
 ARG MAX_JOBS=4
 
-FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder
+ARG BASE_IMAGE=nvidia/cuda
+ARG DEFAULT_TAG=${CUDA_VERSION}-cudnn${CUDNN_VERSION}-devel-ubuntu${UBUNTU_VERSION}
+ARG BASE_TAG=""
+FROM ${BASE_IMAGE:-nvidia/cuda}:${BASE_TAG:-${DEFAULT_TAG}} AS base-builder
 
 ENV PATH="/root/miniconda3/bin:${PATH}"
 

src/axolotl/monkeypatch/trainer_grad_accum.py

@@ -3,14 +3,14 @@ fix for FSDP gradient accumulation
 see https://github.com/huggingface/transformers/pull/35128
 """
 import inspect
+import logging
 
-from accelerate.logging import get_logger
 from transformers import LlamaForCausalLM
 from transformers.trainer import Trainer
 
 from axolotl.monkeypatch.unsloth_ import detab_code
 
-LOG = get_logger("axolotl.monkeypatch.trainer_grad_accum")
+LOG = logging.getLogger("axolotl.monkeypatch.trainer_grad_accum")
 
 ORIGINAL_CONTEXT_CODE = """
     with self.compute_loss_context_manager():
@@ -67,7 +67,7 @@ PATCHED_LLAMA_FCLM_CODE = """
     return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
     # remove num_items_in_batch otherwise self.model attempts to pass it to flash_attention
-    num_items_in_batch = kwargs.pop("num_items_in_batch")
+    num_items_in_batch = kwargs.pop("num_items_in_batch", None)
 
     # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
     outputs = self.model(
@@ -111,12 +111,17 @@ def patch_training_step_for_ga():
     monkeypatch for fixing the training loop for gradient accumulation
     """
 
-    training_step = get_training_step_code()
+    try:
+        training_step = get_training_step_code()
+    except OSError:
+        return
     Trainer._original_training_step = training_step  # pylint: disable=protected-access
     training_step, _ = detab_code(training_step)
-    assert (
-        ORIGINAL_CONTEXT_CODE in training_step
-    ), "Original training_step code not found"
+    if ORIGINAL_CONTEXT_CODE not in training_step:
+        return
+    # assert (
+    #     ORIGINAL_CONTEXT_CODE in training_step
+    # ), "Original training_step code not found"
 
     training_step = training_step.replace(ORIGINAL_CONTEXT_CODE, PATCHED_CONTEXT_CODE)
     training_step = training_step.replace(
@@ -140,7 +145,7 @@ def patch_training_step_for_ga():
         globals(),
     )
     exec(training_step, globals())  # pylint: disable=exec-used  # nosec B102
-    LOG.info("patching training_step", main_process_only=True)
+    LOG.info("patching training_step")
     Trainer.training_step = (  # pylint: disable=protected-access
         _fixed_training_step  # pylint: disable=undefined-variable  # noqa: F821
     )
@@ -164,10 +169,15 @@ def patch_forward_for_ga():
     monkeypatch for fixing the training loop for gradient accumulation
     """
 
-    forward = get_model_forward_code()
+    try:
+        forward = get_model_forward_code()
+    except OSError:
+        return
     LlamaForCausalLM._original_forward = forward  # pylint: disable=protected-access
     forward, _ = detab_code(forward)
-    assert ORIGINAL_LLAMA_FCLM_CODE in forward, "Original forward code not found"
+    if ORIGINAL_LLAMA_FCLM_CODE not in forward:
+        return
+    # assert ORIGINAL_LLAMA_FCLM_CODE in forward, "Original forward code not found"
 
     forward = forward.replace(ORIGINAL_LLAMA_FCLM_CODE, PATCHED_LLAMA_FCLM_CODE)
     forward = forward.replace(
@@ -191,7 +201,7 @@ def patch_forward_for_ga():
         globals(),
     )
     exec(forward, globals())  # pylint: disable=exec-used  # nosec B102
-    LOG.info("patching forward", main_process_only=True)
+    LOG.info("patching forward")
     LlamaForCausalLM.forward = (  # pylint: disable=protected-access
         _fixed_forward  # pylint: disable=undefined-variable  # noqa: F821
     )

src/axolotl/monkeypatch/unsloth_.py

@@ -9,10 +9,7 @@ import torch
 from accelerate.logging import get_logger
 from peft import PeftModelForCausalLM
 from torch import nn
-from transformers.models.llama.modeling_llama import (
-    LlamaFlashAttention2,
-    LlamaForCausalLM,
-)
+from transformers.models.llama.modeling_llama import LlamaFlashAttention2
 
 LOG = get_logger("axolotl.monkeypatch.unsloth")
 
@@ -55,11 +52,6 @@ def original_apply_o(self, hidden_states):
     return attn_output
 
 
-def get_forward_code() -> str:
-    forward = inspect.getsource(LlamaForCausalLM.forward)
-    return forward
-
-
 def get_self_attn_code() -> str:
     forward = inspect.getsource(LlamaFlashAttention2.forward)
     return forward
@@ -102,12 +94,22 @@ def integrate_cross_entropy_loss_patch(model_type: str = "llama") -> None:
 
 
 def detab_code(code: str) -> Tuple[str, str]:
-    spaces = re.match(r"([\s\t]{1,})", code).group(0)
-    code = re.sub(r"^" + spaces, "", code, flags=re.MULTILINE)
+    try:
+        spaces = re.match(r"([\s\t]{1,})", code).group(0)
+        code = re.sub(r"^" + spaces, "", code, flags=re.MULTILINE)
+    except AttributeError:
+        return code, ""
     return code, spaces
 
 
+self_attn_lora_patched = False  # pylint: disable=invalid-name
+
+
 def patch_self_attn_lora():
+    global self_attn_lora_patched  # pylint: disable=global-statement
+    if self_attn_lora_patched:
+        # prevent patching multiple times
+        return
     self_attn_forward = get_self_attn_code()
     LlamaFlashAttention2._original_forward = (  # pylint: disable=protected-access
         self_attn_forward
@@ -139,6 +141,7 @@ def patch_self_attn_lora():
         globals(),
     )
     exec(self_attn_forward, globals())  # pylint: disable=exec-used  # nosec B102
+    self_attn_lora_patched = True
     LOG.info("patching unsloth attn lora", main_process_only=True)
     LlamaFlashAttention2.forward = (
         unsloth_attn_forward  # pylint: disable=undefined-variable  # noqa: F821

src/axolotl/utils/config/__init__.py

@@ -153,7 +153,7 @@ def normalize_config(cfg):
     cfg.is_llama_derived_model = (
         (
             hasattr(model_config, "model_type")
-            and model_config.model_type == ["llama", "mllama_text_model"]
+            and model_config.model_type in ["llama", "mllama_text_model"]
         )
         or cfg.is_llama_derived_model
         or "llama" in cfg.base_model.lower()

src/axolotl/utils/config/models/input/v0_4_1/__init__.py

@@ -1432,20 +1432,6 @@ class AxolotlInputConfig(
                 )
         return data
 
-    @model_validator(mode="before")
-    @classmethod
-    def notify_qlora_unsloth(cls, data):
-        if (
-            data.get("unsloth_lora_mlp")
-            or data.get("unsloth_lora_qkv")
-            or data.get("unsloth_lora_o")
-        ):
-            LOG.info(
-                "Unsloth may not be well supported with the latest version of Transformers, "
-                "resulting in loss that is incorrect."
-            )
-        return data
-
     @model_validator(mode="before")
     @classmethod
     def check_torch_compile_deepspeed(cls, data):

tests/conftest.py

@@ -2,7 +2,9 @@
 shared pytest fixtures
 """
 import functools
+import importlib
 import shutil
+import sys
 import tempfile
 import time
 
@@ -113,3 +115,30 @@ def temp_dir():
     yield _temp_dir
     # Clean up the directory after the test
     shutil.rmtree(_temp_dir)
+
+
+@pytest.fixture(scope="function", autouse=True)
+def cleanup_monkeypatches():
+    from transformers.models.llama.modeling_llama import LlamaFlashAttention2
+
+    original_fa2_forward = LlamaFlashAttention2.forward
+    # monkey patches can happen inside the tests
+    yield
+    # Reset LlamaFlashAttention2 forward
+    LlamaFlashAttention2.forward = original_fa2_forward
+
+    # Reset other known monkeypatches
+    modules_to_reset: list[tuple[str, list[str]]] = [
+        ("transformers.models.llama.modeling_llama", ["LlamaFlashAttention2"]),
+        ("transformers.trainer",),
+        ("transformers.loss.loss_utils",),
+    ]
+    for module_name_tuple in modules_to_reset:
+        module_name = module_name_tuple[0]
+        module = importlib.import_module(module_name)
+        sys.modules[module_name] = module
+        importlib.reload(sys.modules[module_name])
+        if len(module_name_tuple) > 1:
+            module_globals = module_name_tuple[1]
+            for module_global in module_globals:
+                globals().pop(module_global, None)

tests/e2e/patched/test_unsloth_qlora.py

@@ -20,7 +20,6 @@ os.environ["WANDB_DISABLED"] = "true"
 
 
 # pylint: disable=duplicate-code
-@pytest.mark.skip(reason="latest unsloth doesn't work with latest transformers")
 class TestUnslothQLoRA:
     """
     Test class for Unsloth QLoRA Llama models
@@ -37,6 +36,9 @@ class TestUnslothQLoRA:
                 "sequence_len": 1024,
                 "sample_packing": sample_packing,
                 "flash_attention": True,
+                "unsloth_lora_mlp": True,
+                "unsloth_lora_qkv": True,
+                "unsloth_lora_o": True,
                 "load_in_4bit": True,
                 "adapter": "qlora",
                 "lora_r": 16,
@@ -83,6 +85,9 @@ class TestUnslothQLoRA:
             {
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "sequence_len": 1024,
+                "unsloth_lora_mlp": True,
+                "unsloth_lora_qkv": True,
+                "unsloth_lora_o": True,
                 "sample_packing": False,
                 "load_in_4bit": True,
                 "adapter": "qlora",
@@ -134,6 +139,9 @@ class TestUnslothQLoRA:
             {
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "sequence_len": 1024,
+                "unsloth_lora_mlp": True,
+                "unsloth_lora_qkv": True,
+                "unsloth_lora_o": True,
                 "sample_packing": False,
                 "load_in_4bit": True,
                 "adapter": "qlora",