make sure the model config loader respects the model_revision too (#1160

) [skip-ci]
axolotl-ai-cloud · Jan 22, 2024 · fccb542 · fccb542
1 parent 2ce5c0d
commit fccb542
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -76,10 +76,15 @@ def load_model_config(cfg):
     if not model_config_name and cfg.tokenizer_config:
         model_config_name = cfg.tokenizer_config
     trust_remote_code = cfg.trust_remote_code is True
+    config_kwargs = {}
+    if cfg.model_revision:
+        config_kwargs["revision"] = cfg.model_revision
 
     try:
         model_config = AutoConfig.from_pretrained(
-            model_config_name, trust_remote_code=trust_remote_code
+            model_config_name,
+            trust_remote_code=trust_remote_code,
+            **config_kwargs,
         )
     except ValueError as err:
         if "mamba" in model_config_name: