feat: add cut_cross_entropy (#2091)

* feat: add cut_cross_entropy

* fix: add to input

* fix: remove from setup.py

* feat: refactor into an integration

* chore: ignore lint

* feat: add test for cce

* fix: set max_steps for liger test

* chore: Update base model following suggestion

Co-authored-by: Wing Lian <wing.lian@gmail.com>

* chore: update special_tokens following suggestion

Co-authored-by: Wing Lian <wing.lian@gmail.com>

* chore: remove with_temp_dir following comments

* fix: plugins aren't loaded

* chore: update quotes in error message

* chore: lint

* chore: lint

* feat: enable FA on test

* chore: refactor get_pytorch_version

* fix: lock cce commit version

* fix: remove subclassing UT

* fix: downcast even if not using FA and config check

* feat: add test to check different attentions

* feat: add install to CI

* chore: refactor to use parametrize for attention

* fix: pytest not detecting test

* feat: handle torch lower than 2.4

* fix args/kwargs to match docs

* use release version cut-cross-entropy==24.11.4

* fix quotes

* fix: use named params for clarity for modal builder

* fix: handle install from pip

* fix: test check only top level module install

* fix: re-add import check

* uninstall existing version if no transformers submodule in cce

* more dataset fixtures into the cache

---------

Co-authored-by: Wing Lian <wing.lian@gmail.com>
Co-authored-by: Wing Lian <wing@axolotl.ai>

tests/conftest.py

@@ -51,6 +51,22 @@ def download_mlabonne_finetome_100k_dataset():
     snapshot_download("mlabonne/FineTome-100k", repo_type="dataset")
 
 
+@pytest.fixture
+def download_argilla_distilabel_capybara_dpo_7k_binarized_dataset():
+    # download the dataset
+    snapshot_download(
+        "argilla/distilabel-capybara-dpo-7k-binarized", repo_type="dataset"
+    )
+
+
+@pytest.fixture
+def download_arcee_ai_distilabel_intel_orca_dpo_pairs_dataset():
+    # download the dataset
+    snapshot_download(
+        "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized", repo_type="dataset"
+    )
+
+
 @pytest.fixture
 def temp_dir():
     # Create a temporary directory

tests/e2e/integrations/liger.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, prepare_plugins
 from axolotl.utils.dict import DictDefault
 
 from ..utils import with_temp_dir
@@ -54,8 +54,10 @@ class LigerIntegrationTestCase(unittest.TestCase):
                 "lr_scheduler": "cosine",
                 "save_safetensors": True,
                 "bf16": "auto",
+                "max_steps": 10,
             }
         )
+        prepare_plugins(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
@@ -99,8 +101,10 @@ class LigerIntegrationTestCase(unittest.TestCase):
                 "lr_scheduler": "cosine",
                 "save_safetensors": True,
                 "bf16": "auto",
+                "max_steps": 10,
             }
         )
+        prepare_plugins(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

tests/e2e/integrations/test_cut_cross_entropy.py

@@ -0,0 +1,94 @@
+"""
+Simple end-to-end test for Cut Cross Entropy integration
+"""
+
+from pathlib import Path
+
+import pytest
+
+from axolotl.cli import load_datasets
+from axolotl.common.cli import TrainerCliArgs
+from axolotl.train import train
+from axolotl.utils import get_pytorch_version
+from axolotl.utils.config import normalize_config, prepare_plugins
+from axolotl.utils.dict import DictDefault
+
+# pylint: disable=duplicate-code
+
+
+@pytest.fixture()
+def min_cfg(temp_dir):
+    return {
+        "base_model": "HuggingFaceTB/SmolLM2-135M",
+        "plugins": [
+            "axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin",
+        ],
+        "cut_cross_entropy": True,
+        "sequence_len": 1024,
+        "val_set_size": 0.1,
+        "special_tokens": {
+            "pad_token": "<|endoftext|>",
+        },
+        "datasets": [
+            {
+                "path": "mhenrichsen/alpaca_2k_test",
+                "type": "alpaca",
+            },
+        ],
+        "num_epochs": 1,
+        "micro_batch_size": 8,
+        "gradient_accumulation_steps": 1,
+        "learning_rate": 0.00001,
+        "optimizer": "adamw_torch",
+        "output_dir": temp_dir,
+        "lr_scheduler": "cosine",
+        "save_safetensors": True,
+        "max_steps": 10,
+        "bf16": "auto",
+    }
+
+
+class TestCutCrossEntropyIntegration:
+    """
+    e2e tests for cut_cross_entropy integration with Axolotl
+    """
+
+    # pylint: disable=redefined-outer-name
+    def test_llama_w_cce(self, min_cfg, temp_dir):
+        cfg = DictDefault(min_cfg)
+        prepare_plugins(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+
+        major, minor, _ = get_pytorch_version()
+        if (major, minor) < (2, 4):
+            with pytest.raises(ImportError):
+                train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+        else:
+            train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+            assert (Path(temp_dir) / "model.safetensors").exists()
+
+    @pytest.mark.parametrize(
+        "attention_type",
+        ["flash_attention", "sdp_attention", "xformers_attention"],
+    )
+    def test_llama_w_cce_and_attention(self, min_cfg, temp_dir, attention_type):
+        cfg = DictDefault(
+            min_cfg
+            | {
+                attention_type: True,
+            }
+        )
+        prepare_plugins(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+
+        major, minor, _ = get_pytorch_version()
+        if (major, minor) < (2, 4):
+            with pytest.raises(ImportError):
+                train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+        else:
+            train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+            assert (Path(temp_dir) / "model.safetensors").exists()