use yet-another-deepspeed branch from transformers#37324

requirements.txt

@@ -12,9 +12,9 @@ liger-kernel==0.5.5
 packaging==23.2
 
 peft==0.15.0
-transformers==4.50.3
+transformers @ git+https://github.com/huggingface/transformers.git@yet-another-deepspeed
 tokenizers>=0.21.1
-accelerate==1.5.2
+accelerate==1.6.0
 datasets==3.5.0
 deepspeed==0.15.4
 trl==0.16.0

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

@@ -56,7 +56,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,
@@ -193,7 +193,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 +390,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 +403,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,
@@ -469,7 +469,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|>",
@@ -483,7 +483,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,
@@ -566,7 +566,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|>",
@@ -639,7 +639,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|>",
@@ -712,7 +712,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|>",