use inference-runner corresponding to reference model (#346)

Author: RaymondLi0

fast_llm/engine/multi_stage/config.py

@@ -31,6 +31,7 @@
 
 if typing.TYPE_CHECKING:
     from fast_llm.engine.inference.huggingface import HuggingfaceBaseModelForCausalLM
+    from fast_llm.engine.inference.runner import InferenceRunner
     from fast_llm.engine.multi_stage.fast_llm_model import FastLLMModel
 
 logger = logging.getLogger(__name__)
@@ -241,6 +242,10 @@ def get_checkpoint_handler_class(cls, format: type[CheckpointFormat] | str) -> t
     def get_model_class(cls) -> type["FastLLMModel"]:
         raise NotImplementedError
 
+    @classmethod
+    def get_inference_runner_class(cls) -> type["InferenceRunner"]:
+        raise NotImplementedError
+
     @classmethod
     def get_huggingface_model_for_causal_lm_class(cls) -> type["HuggingfaceBaseModelForCausalLM"]:
         raise NotImplementedError

fast_llm/engine/training/config.py

@@ -32,7 +32,6 @@
 from fast_llm.utils import Assert
 
 if typing.TYPE_CHECKING:
-    from fast_llm.engine.inference.runner import InferenceRunner
     from fast_llm.engine.training.trainer import Trainer, TrainingEvaluator
 
 
@@ -403,10 +402,6 @@ def _setup(self):
     def get_trainer_class(cls) -> type["Trainer"]:
         raise NotImplementedError
 
-    @classmethod
-    def get_inference_runner_class(cls) -> type["InferenceRunner"]:
-        raise NotImplementedError
-
     def _get_runnable(self) -> typing.Callable[[], None]:
         from fast_llm.engine.distributed.distributed import Distributed
 

fast_llm/engine/training/trainer.py

@@ -142,7 +142,7 @@ def __init__(self, config: TrainerConfig):
         self._reference_models = {}
         for name, reference_config in self._config.reference_models.items():
             log_main_rank(f"Creating `{name} reference model...")
-            self._reference_models[name] = self._config.get_inference_runner_class()(
+            self._reference_models[name] = reference_config.model.get_inference_runner_class()(
                 reference_config.model.get_model_class()(reference_config.model)
             )
             self._multi_stage.base_model.add_reference_model(name, self._reference_models[name])

fast_llm/models/gpt/config.py

@@ -171,6 +171,12 @@ def get_model_class(cls) -> type["GPTModel"]:
 
         return GPTModel
 
+    @classmethod
+    def get_inference_runner_class(cls) -> type["GPTInferenceRunner"]:
+        from fast_llm.models.gpt.model import GPTInferenceRunner
+
+        return GPTInferenceRunner
+
     @classmethod
     def get_huggingface_model_for_causal_lm_class(cls) -> type["HuggingfaceGPTModelForCausalLM"]:
         from fast_llm.models.gpt.huggingface import HuggingfaceGPTModelForCausalLM
@@ -254,9 +260,3 @@ def get_trainer_class(cls) -> type["GPTTrainer"]:
         from fast_llm.models.gpt.trainer import GPTTrainer
 
         return GPTTrainer
-
-    @classmethod
-    def get_inference_runner_class(cls) -> type["GPTInferenceRunner"]:
-        from fast_llm.models.gpt.model import GPTInferenceRunner
-
-        return GPTInferenceRunner

fast_llm/models/ssm/config.py

@@ -173,6 +173,16 @@ def get_model_class(cls) -> type["HybridSSMModel"]:
 
         return HybridSSMModel
 
+    @classmethod
+    def get_inference_runner_class(cls) -> type["HybridSSMInferenceRunner"]:
+        from fast_llm.models.ssm.model import HybridSSMInferenceRunner
+
+        logger.warning(
+            "HybridSSMInferenceRunner only supports training-style forward pass. Use generate with cache disabled."
+        )
+
+        return HybridSSMInferenceRunner
+
     @classmethod
     def get_huggingface_model_for_causal_lm_class(cls) -> type["HuggingfaceHybridSSMModelForCausalLM"]:
         from fast_llm.models.ssm.huggingface import HuggingfaceHybridSSMModelForCausalLM
@@ -227,13 +237,3 @@ def _validate(self) -> None:
             Assert.none(reference_model.model.base_model.cross_entropy_splits)
             Assert.eq(reference_model.model.base_model.parallel_embeddings, self.model.base_model.parallel_embeddings)
             Assert.geq(reference_model.model.base_model.prediction_heads, self.model.base_model.prediction_heads)
-
-    @classmethod
-    def get_inference_runner_class(cls) -> type["HybridSSMInferenceRunner"]:
-        from fast_llm.models.ssm.model import HybridSSMInferenceRunner
-
-        logger.warning(
-            "HybridSSMInferenceRunner only supports training-style forward pass. Use generate with cache disabled."
-        )
-
-        return HybridSSMInferenceRunner