Merge pull request #556 from allenai/shanea/make-hf-olmo-support-new-…

…transformers

Make hf_olmo support new transformers versions

CHANGELOG.md

@@ -35,6 +35,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Changed `tie_weights` method to a no-op as weight tying is handled in olmo/model.py
 - Fixed the size calculation for qk layer norm
 - Fixed pipeline test failure that occurs due to a bug in transformers version 4.39.1
+- Make `hf_olmo` compatible with transformers versions >=4.40.0
 
 ## [v0.2.5](https://github.com/allenai/OLMo/releases/tag/v0.2.5) - 2024-03-06
 

hf_olmo/configuration_olmo.py

@@ -2,6 +2,8 @@
 OLMo configuration
 """
 
+import transformers
+from packaging import version
 from transformers import AutoConfig, PretrainedConfig
 from transformers.utils import logging
 
@@ -37,5 +39,7 @@ def hidden_size(self):
         return self.d_model
 
 
-# Register the config class so that it is available for transformer pipelines, auto-loading etc.
-AutoConfig.register("olmo", OLMoConfig)
+if version.parse(transformers.__version__) < version.parse("4.40.0"):
+    # Register the config class so that it is available for transformer pipelines, auto-loading etc.
+    # OLMo is integrated directly in transformers from v4.40.0 onwards
+    AutoConfig.register("olmo", OLMoConfig)

hf_olmo/modeling_olmo.py

@@ -3,6 +3,8 @@
 from typing import List, Optional, Tuple, Union
 
 import torch
+import transformers
+from packaging import version
 from transformers import PreTrainedModel
 from transformers.cache_utils import Cache
 from transformers.modeling_outputs import CausalLMOutputWithPast
@@ -222,5 +224,7 @@ def resize_token_embeddings(
         return model_embeds
 
 
-# Register the model so that it is available for transformer pipelines, auto-loading, etc.
-AutoModelForCausalLM.register(OLMoConfig, OLMoForCausalLM)
+if version.parse(transformers.__version__) < version.parse("4.40.0"):
+    # Register the model so that it is available for transformer pipelines, auto-loading, etc.
+    # OLMo is integrated directly in transformers from v4.40.0 onwards
+    AutoModelForCausalLM.register(OLMoConfig, OLMoForCausalLM)

tests/hf_olmo/hf_olmo_test.py

@@ -1,12 +1,18 @@
 import pytest
 import torch
+import transformers
+from packaging import version
 
 from olmo import BlockType, Tokenizer, TrainConfig
 from olmo.data import DataCollator
 from olmo.model import OLMo
 from olmo.torch_util import seed_all
 
 
+@pytest.mark.skipif(
+    version.parse(transformers.__version__) >= version.parse("4.40.0"),
+    reason="hf_olmo auto classes are not compatible with transformers >=v4.40.0",
+)
 def test_auto_hf_classes(model_path: str):
     from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer