fix: CUDA error when inferencing with Falcon-40B base model (#992)

2026-07-26 13:37:27 +08:00 · 2023-09-10 01:39:02 -07:00
parent a62de9ecfd
commit 898285c9bf
1 changed files with 2 additions and 1 deletions
@@ -114,8 +114,9 @@ class ModelConfig:
        # Note: for falcon, when new_decoder_architecture is True, the
        # multi_query flag is ignored and we use n_head_kv for the number of
        # KV heads.
+        falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"]
        new_decoder_arch_falcon = (
-            self.hf_config.model_type == "falcon"
+            self.hf_config.model_type in falcon_model_types
            and getattr(self.hf_config, "new_decoder_architecture", False))
        if not new_decoder_arch_falcon and getattr(self.hf_config,
                                                   "multi_query", False):