diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9f4bf369a..7f7960316 100644
--- a/CHANGELOG.md
+++ b/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
 
diff --git a/hf_olmo/configuration_olmo.py b/hf_olmo/configuration_olmo.py
index 4e076e9a6..405dacb2d 100644
--- a/hf_olmo/configuration_olmo.py
+++ b/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)
diff --git a/hf_olmo/modeling_olmo.py b/hf_olmo/modeling_olmo.py
index 1bf763917..15cdd000b 100644
--- a/hf_olmo/modeling_olmo.py
+++ b/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)
diff --git a/tests/hf_olmo/hf_olmo_test.py b/tests/hf_olmo/hf_olmo_test.py
index b0107ffec..bfe18e3ad 100644
--- a/tests/hf_olmo/hf_olmo_test.py
+++ b/tests/hf_olmo/hf_olmo_test.py
@@ -1,5 +1,7 @@
 import pytest
 import torch
+import transformers
+from packaging import version
 
 from olmo import BlockType, Tokenizer, TrainConfig
 from olmo.data import DataCollator
@@ -7,6 +9,10 @@
 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