update peft to 0.15.1

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

@@ -24,6 +24,13 @@ jobs:
       fail-fast: false
       matrix:
         include:
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.6.0
+            axolotl_extras: vllm
+            num_gpus: 2
+            nightly_build: "true"
           - cuda: 124
             cuda_version: 12.4.1
             python_version: "3.11"
@@ -38,13 +45,6 @@ jobs:
             axolotl_extras: vllm
             num_gpus: 2
             nightly_build: "true"
-          - cuda: 124
-            cuda_version: 12.4.1
-            python_version: "3.11"
-            pytorch: 2.6.0
-            axolotl_extras: vllm
-            num_gpus: 2
-            nightly_build: "true"
     runs-on: [self-hosted, modal]
     timeout-minutes: 120
     steps:

.github/workflows/tests.yml

@@ -211,7 +211,7 @@ jobs:
           - cuda: 124
             cuda_version: 12.4.1
             python_version: "3.11"
-            pytorch: 2.5.1
+            pytorch: 2.6.0
             num_gpus: 1
             axolotl_extras: vllm
     steps:
@@ -258,7 +258,7 @@ jobs:
           - cuda: 124
             cuda_version: 12.4.1
             python_version: "3.11"
-            pytorch: 2.6.0
+            pytorch: 2.5.1
             num_gpus: 1
             axolotl_extras: vllm
     steps:

requirements.txt

@@ -11,13 +11,14 @@ liger-kernel==0.5.5
 
 packaging==23.2
 
-peft==0.15.0
+peft==0.15.1
 transformers==4.51.0
 tokenizers>=0.21.1
 accelerate==1.6.0
 datasets==3.5.0
 deepspeed>=0.15.4
 trl==0.16.1
+hf_xet==1.0.0
 
 optimum==1.16.2
 hf_transfer

src/axolotl/monkeypatch/multipack.py

@@ -13,6 +13,7 @@ from axolotl.monkeypatch.utils import get_unpad_data
 SUPPORTED_MULTIPACK_MODEL_TYPES = [
     "mllama_text_model",
     "llama",
+    "llama4",
     "mistral",
     "mixtral",
     "qwen2",

tests/e2e/multigpu/test_llama.py

@@ -7,9 +7,11 @@ import os
 from pathlib import Path
 
 import pytest
+import transformers
 import yaml
 from accelerate.test_utils import execute_subprocess_async
 from huggingface_hub import snapshot_download
+from packaging import version
 from transformers.testing_utils import get_torch_dist_unique_port
 
 from axolotl.utils.dict import DictDefault
@@ -28,6 +30,10 @@ def download_model():
     snapshot_download("HuggingFaceTB/SmolLM2-135M")
 
 
+def transformers_version_eq(required_version):
+    return version.parse(transformers.__version__) == version.parse(required_version)
+
+
 class TestMultiGPULlama:
     """
     Test case for Llama models using LoRA
@@ -56,7 +62,7 @@ class TestMultiGPULlama:
                 ],
                 "num_epochs": 1,
                 "max_steps": 2,
-                "micro_batch_size": 4,
+                "micro_batch_size": 1,
                 "gradient_accumulation_steps": 4,
                 # "gradient_checkpointing": True,
                 "output_dir": temp_dir,
@@ -108,7 +114,7 @@ class TestMultiGPULlama:
                 "lora_alpha": 16,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "val_set_size": 0.01,
+                "val_set_size": 0.05,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
                 },
@@ -116,6 +122,7 @@ class TestMultiGPULlama:
                     {
                         "path": "tatsu-lab/alpaca",
                         "type": "alpaca",
+                        "split": "train[:20%]",
                     },
                 ],
                 "num_epochs": 1,
@@ -193,7 +200,7 @@ class TestMultiGPULlama:
                 ],
                 "num_epochs": 1,
                 "max_steps": 2,
-                "micro_batch_size": 4,
+                "micro_batch_size": 2,
                 "gradient_accumulation_steps": 4,
                 # "gradient_checkpointing": True,
                 "output_dir": temp_dir,
@@ -390,7 +397,7 @@ class TestMultiGPULlama:
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "sample_packing": True,
                 "pad_to_sequence_len": True,
-                "sequence_len": 2048,
+                "sequence_len": 1024,
                 "val_set_size": 0.01,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
@@ -403,7 +410,7 @@ class TestMultiGPULlama:
                 ],
                 "num_epochs": 1,
                 "max_steps": 2,
-                "micro_batch_size": 4,
+                "micro_batch_size": 2,
                 "gradient_accumulation_steps": 2,
                 # "gradient_checkpointing": True,
                 "output_dir": temp_dir,
@@ -551,7 +558,7 @@ class TestMultiGPULlama:
                 "sample_packing": True,
                 "eval_sample_packing": False,
                 "pad_to_sequence_len": True,
-                "sequence_len": 2048,
+                "sequence_len": 1024,
                 "val_set_size": 0.01,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
@@ -565,7 +572,7 @@ class TestMultiGPULlama:
                 ],
                 "num_epochs": 1,
                 "max_steps": 2,
-                "micro_batch_size": 4,
+                "micro_batch_size": 2,
                 "gradient_accumulation_steps": 2,
                 # "gradient_checkpointing": True,
                 "output_dir": temp_dir,
@@ -612,8 +619,11 @@ class TestMultiGPULlama:
             temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss is too high"
         )
 
-    @pytest.mark.skip(
+    # TODO: remove skip once deepspeed regression is fixed
-        reason="ds-zero3 broken in main until transformers#37281 resolved"
+    # 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",
@@ -651,7 +661,7 @@ class TestMultiGPULlama:
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "sample_packing": True,
                 "pad_to_sequence_len": True,
-                "sequence_len": 2048,
+                "sequence_len": 1024,
                 "val_set_size": 0.01,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
@@ -724,7 +734,7 @@ class TestMultiGPULlama:
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "sample_packing": True,
                 "pad_to_sequence_len": True,
-                "sequence_len": 2048,
+                "sequence_len": 1024,
                 "val_set_size": 0.01,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
@@ -797,7 +807,7 @@ class TestMultiGPULlama:
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "sample_packing": True,
                 "pad_to_sequence_len": True,
-                "sequence_len": 2048,
+                "sequence_len": 1024,
                 "val_set_size": 0.01,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
@@ -885,7 +895,7 @@ class TestMultiGPULlama:
                 "sample_packing": True,
                 "bf16": True,
                 "save_safetensors": True,
-                "deepspeed": str(AXOLOTL_ROOT / "deepspeed_configs/zero1.json"),
+                # "deepspeed": str(AXOLOTL_ROOT / "deepspeed_configs/zero1.json"),
                 "use_tensorboard": True,
             }
         )

tests/e2e/multigpu/test_ray.py

@@ -31,7 +31,7 @@ class TestMultiGPURay:
         cfg = DictDefault(
             {
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "sequence_len": 2048,
+                "sequence_len": 1024,
                 "adapter": "lora",
                 "lora_r": 8,
                 "lora_alpha": 16,
@@ -94,8 +94,8 @@ class TestMultiGPURay:
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "sample_packing": True,
                 "pad_to_sequence_len": True,
-                "sequence_len": 2048,
+                "sequence_len": 1024,
-                "val_set_size": 0.05,
+                "val_set_size": 0.01,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
                 },