refactor: remove MTP-specific function requirement from non-MTP models.

Author: yingxudeng

xllm/core/framework/model/causal_lm.h

@@ -36,6 +36,43 @@ limitations under the License.
 
 namespace xllm {
 
+namespace detail {
+template <typename T, typename = void>
+struct has_get_lm_head : std::false_type {};
+
+template <typename T>
+struct has_get_lm_head<T,
+                       std::void_t<decltype(std::declval<T>().get_lm_head())>>
+    : std::true_type {};
+
+template <typename T, typename = void>
+struct has_set_lm_head : std::false_type {};
+
+template <typename T>
+struct has_set_lm_head<T,
+                       std::void_t<decltype(std::declval<T>().set_lm_head(
+                           std::declval<layer::LmHead&>()))>> : std::true_type {
+};
+
+template <typename T, typename = void>
+struct has_get_word_embedding : std::false_type {};
+
+template <typename T>
+struct has_get_word_embedding<
+    T,
+    std::void_t<decltype(std::declval<T>().get_word_embedding())>>
+    : std::true_type {};
+
+template <typename T, typename = void>
+struct has_set_word_embedding : std::false_type {};
+
+template <typename T>
+struct has_set_word_embedding<
+    T,
+    std::void_t<decltype(std::declval<T>().set_word_embedding(
+        std::declval<layer::WordEmbedding&>()))>> : std::true_type {};
+}  // namespace detail
+
 class CausalLM : public torch::nn::Module {
  public:
   ~CausalLM() override = default;
@@ -65,10 +102,23 @@ class CausalLM : public torch::nn::Module {
 
   virtual const torch::TensorOptions& options() const = 0;
 
-  virtual layer::LmHead get_lm_head() = 0;
-  virtual void set_lm_head(layer::LmHead& head) = 0;
-  virtual layer::WordEmbedding get_word_embedding() = 0;
-  virtual void set_word_embedding(layer::WordEmbedding& embedding) = 0;
+  // MTP-specific interface.
+  virtual layer::LmHead get_lm_head() {
+    LOG(FATAL)
+        << "Method 'get_lm_head' is not implemented/supported by this model.";
+  }
+  virtual void set_lm_head(layer::LmHead& head) {
+    LOG(FATAL)
+        << "Method 'set_lm_head' is not implemented/supported by this model.";
+  }
+  virtual layer::WordEmbedding get_word_embedding() {
+    LOG(FATAL) << "Method 'get_word_embedding' is not implemented/supported by "
+                  "this model.";
+  }
+  virtual void set_word_embedding(layer::WordEmbedding& embedding) {
+    LOG(FATAL) << "Method 'set_word_embedding' is not implemented/supported by "
+                  "this model.";
+  }
 };
 
 template <typename Model>
@@ -102,16 +152,36 @@ class CausalLMImpl : public CausalLM {
     return model_->update_expert_weight(layer_id);
   }
 
-  layer::LmHead get_lm_head() override { return model_->get_lm_head(); };
+  layer::LmHead get_lm_head() override {
+    if constexpr (detail::has_get_lm_head<Model>::value) {
+      return model_->get_lm_head();
+    } else {
+      return CausalLM::get_lm_head();
+    }
+  };
 
-  void set_lm_head(layer::LmHead& head) override { model_->set_lm_head(head); };
+  void set_lm_head(layer::LmHead& head) override {
+    if constexpr (detail::has_set_lm_head<Model>::value) {
+      model_->set_lm_head(head);
+    } else {
+      CausalLM::set_lm_head(head);
+    }
+  };
 
   layer::WordEmbedding get_word_embedding() override {
-    return model_->get_word_embedding();
+    if constexpr (detail::has_get_word_embedding<Model>::value) {
+      return model_->get_word_embedding();
+    } else {
+      return CausalLM::get_word_embedding();
+    }
   };
 
   void set_word_embedding(layer::WordEmbedding& embedding) override {
-    model_->set_word_embedding(embedding);
+    if constexpr (detail::has_set_word_embedding<Model>::value) {
+      model_->set_word_embedding(embedding);
+    } else {
+      CausalLM::set_word_embedding(embedding);
+    }
   };
 
   torch::Device device() const override { return options_.device(); }

xllm/core/framework/model/causal_vlm.h

@@ -63,16 +63,36 @@ class CausalVLMImpl : public CausalVLM {
 
   virtual void update_expert_weight(int32_t layer_id) { return; }
 
-  layer::LmHead get_lm_head() override { return model_->get_lm_head(); };
+  layer::LmHead get_lm_head() override {
+    if constexpr (detail::has_get_lm_head<Model>::value) {
+      return model_->get_lm_head();
+    } else {
+      return CausalLM::get_lm_head();
+    }
+  };
 
-  void set_lm_head(layer::LmHead& head) override { model_->set_lm_head(head); };
+  void set_lm_head(layer::LmHead& head) override {
+    if constexpr (detail::has_set_lm_head<Model>::value) {
+      model_->set_lm_head(head);
+    } else {
+      CausalLM::set_lm_head(head);
+    }
+  };
 
   layer::WordEmbedding get_word_embedding() override {
-    return model_->get_word_embedding();
+    if constexpr (detail::has_get_word_embedding<Model>::value) {
+      return model_->get_word_embedding();
+    } else {
+      return CausalLM::get_word_embedding();
+    }
   };
 
   void set_word_embedding(layer::WordEmbedding& embedding) override {
-    model_->set_word_embedding(embedding);
+    if constexpr (detail::has_set_word_embedding<Model>::value) {
+      model_->set_word_embedding(embedding);
+    } else {
+      CausalLM::set_word_embedding(embedding);
+    }
   };
 
   torch::Device device() const override { return options_.device(); }