pinning transformers version

tests/e2e/multigpu/test_llama.py

@@ -16,7 +16,7 @@ from transformers.testing_utils import get_torch_dist_unique_port
 
 from axolotl.utils.dict import DictDefault
 
-from tests.e2e.utils import check_tensorboard, require_torch_2_6_0
+from tests.e2e.utils import check_tensorboard, require_torch_2_6_0, require_torch_2_7_0
 
 LOG = logging.getLogger("axolotl.tests.e2e.multigpu")
 os.environ["WANDB_DISABLED"] = "true"
@@ -457,18 +457,12 @@ class TestMultiGPULlama:
             temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss is too high"
         )
 
-    @require_torch_2_6_0
-    @pytest.mark.parametrize(
-        "attention_backend",
-        ["flash", "flex"],
-    )
+    @require_torch_2_7_0
     @pytest.mark.parametrize(
         "fsdp_reshard_after_forward",
         [True, False],
     )
-    def test_fsdp2_packed(
-        self, temp_dir, attention_backend, fsdp_reshard_after_forward
-    ):
+    def test_fsdp2_packed_flash(self, temp_dir, fsdp_reshard_after_forward):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
@@ -509,13 +503,79 @@ class TestMultiGPULlama:
                     "fsdp_reshard_after_forward": fsdp_reshard_after_forward,
                 },
                 "use_tensorboard": True,
+                "flex_attention": True,
             }
         )
-        if attention_backend == "flash":
-            cfg.flash_attention = True
-        elif attention_backend == "flex":
-            cfg.flex_attention = True
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
 
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.1, "Train Loss is too high"
+        )
+
+    @require_torch_2_6_0
+    @pytest.mark.parametrize(
+        "fsdp_reshard_after_forward",
+        [True, False],
+    )
+    def test_fsdp2_packed_flex(self, temp_dir, fsdp_reshard_after_forward):
+        # pylint: disable=duplicate-code
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "sequence_len": 2048,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "gradient_checkpointing": True,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_8bit",
+                "lr_scheduler": "cosine",
+                "fsdp": [
+                    "auto_wrap",
+                ],
+                "fsdp_config": {
+                    "fsdp_version": 2,
+                    # "fsdp_forward_prefetch": True,  # not yet implemented in accelerate
+                    "fsdp_offload_params": False,
+                    "fsdp_cpu_ram_efficient_loading": False,
+                    "fsdp_transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
+                    "fsdp_state_dict_type": "SHARDED_STATE_DICT",
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "fsdp_reshard_after_forward": fsdp_reshard_after_forward,
+                },
+                "use_tensorboard": True,
+                "flash_attention": True,
+            }
+        )
         # write cfg to yaml file
         Path(temp_dir).mkdir(parents=True, exist_ok=True)
         with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
@@ -617,12 +677,6 @@ 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/solo/test_flex.py

@@ -14,7 +14,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_tensorboard, require_torch_2_6_0, with_temp_dir
+from ..utils import check_tensorboard, require_torch_2_7_0, with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -25,7 +25,7 @@ class TestPackedFlex(unittest.TestCase):
     Test case for Packed training of llama models
     """
 
-    @require_torch_2_6_0
+    @require_torch_2_7_0
     @with_temp_dir
     def test_loss_llama(self, temp_dir):
         # pylint: disable=duplicate-code

tests/e2e/utils.py

@@ -33,6 +33,18 @@ def with_temp_dir(test_func):
     return wrapper
 
 
+def require_torch_2_7_0(test_case):
+    """
+    Decorator marking a test that requires torch >= 2.7.0
+    """
+
+    def is_min_2_7_0():
+        torch_version = version.parse(torch.__version__)
+        return torch_version >= version.parse("2.7.0")
+
+    return unittest.skipUnless(is_min_2_7_0(), "test requires torch>=2.7.0")(test_case)
+
+
 def most_recent_subdir(path):
     base_path = Path(path)
     subdirectories = [d for d in base_path.iterdir() if d.is_dir()]