Merge pull request #1 from Pomilon/fix-ci-tests-16470124072715661906

Fix tests and model compatibility for CI/CD

Author: Pomilon

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
 ]
 dependencies = [
     "torch>=2.0.0",
-    "transformers>=4.30.0",
+    "transformers>=4.30.0,<5.0.0",
     "safetensors>=0.3.0",
     "accelerate>=0.20.0",
     "peft>=0.4.0",

requirements.txt

@@ -1,5 +1,5 @@
 torch>=2.0.0
-transformers>=4.30.0
+transformers>=4.30.0,<5.0.0
 safetensors>=0.3.0
 accelerate>=0.20.0
 peft>=0.4.0

src/lema/models/llama.py

@@ -11,7 +11,15 @@ def __init__(self, config: Dict[str, Any]):
         self.hf_config = LlamaConfig(**config)
         if getattr(self.hf_config, "_attn_implementation", None) is None:
             self.hf_config._attn_implementation = config.get("attn_implementation", "eager")
-        self.rotary_emb = LlamaRotaryEmbedding(self.hf_config)
+
+        try:
+            self.rotary_emb = LlamaRotaryEmbedding(self.hf_config)
+        except TypeError:
+            self.rotary_emb = LlamaRotaryEmbedding(
+                self.hf_config.hidden_size // self.hf_config.num_attention_heads,
+                max_position_embeddings=self.hf_config.max_position_embeddings
+            )
+
         self.layer_pool: List[nn.Module] = []
         self.param_mappings: Dict[int, List[tuple]] = {}
         self._max_pool_size = 8

tests/test_core_components.py

@@ -1,9 +1,9 @@
 import pytest
 import torch
 import torch.nn as nn
-from src.lema.core.gbi import GlobalBinaryIndex
-from src.lema.core.lora import LoRAManager, LoRAWrapper
-from src.lema.core.memory import TripleBufferManager
+from lema.core.gbi import GlobalBinaryIndex
+from lema.core.lora import LoRAManager, LoRAWrapper
+from lema.core.memory import TripleBufferManager
 
 # Mocking
 class MockAdapter:

tests/test_gradient_equivalence.py

@@ -2,10 +2,10 @@
 import torch.nn as nn
 from transformers import GPT2Config, GPT2LMHeadModel
 from safetensors.torch import save_file
-from src.lema.core.gbi import GlobalBinaryIndex
-from src.lema.models.gpt2 import GPT2Adapter
-from src.lema.engine.trainer import LemaTrainer
-from src.lema.config import LemaConfig, MemoryStrategy
+from lema.core.gbi import GlobalBinaryIndex
+from lema.models.gpt2 import GPT2Adapter
+from lema.engine.trainer import LemaTrainer
+from lema.config import LemaConfig, MemoryStrategy
 import os
 import pytest
 

tests/test_llama_forward.py

@@ -1,6 +1,6 @@
 import torch
 import torch.nn as nn
-from src.lema.models.llama import LlamaAdapter
+from lema.models.llama import LlamaAdapter
 from transformers import LlamaConfig
 
 def test_llama_adapter_forward():