[Piecewise] Fix PCG issue for multimodal and embedding model that wraps language_model (#17290)

python/sglang/srt/model_executor/model_runner.py

@@ -2027,9 +2027,10 @@ class ModelRunner(ModelRunnerKVCacheMixin):
 
         # Collect attention layers and moe layers from the model
         self.model.model = resolve_language_model(self.model)
+        language_model = getattr(self.model, "language_model", self.model)
         self.attention_layers = []
         self.moe_layers = []
-        for layer in self.model.model.layers:
+        for layer in language_model.model.layers:
             if hasattr(layer, "self_attn"):
                 if hasattr(layer.self_attn, "attn"):
                     self.attention_layers.append(layer.self_attn.attn)

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -240,8 +240,11 @@ class PiecewiseCudaGraphRunner:
         set_graph_pool_id(get_global_graph_memory_pool())
 
         with enable_piecewise_cuda_graph():
+            language_model = getattr(
+                self.model_runner.model, "language_model", self.model_runner.model
+            )
             with patch_model(
-                self.model_runner.model.model, self.compile_config.compiler
+                language_model.model, self.compile_config.compiler
             ) as patched_model:
                 install_torch_compiled(
                     patched_model,