Fix model forward grad (#628)

python/sglang/srt/models/chatglm.py

@@ -360,6 +360,7 @@ class ChatGLMForCausalLM(nn.Module):
         self.logits_processor = LogitsProcessor(config)
         self.sampler = Sampler()
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/dbrx.py

@@ -368,6 +368,7 @@ class DbrxForCausalLM(nn.Module):
         )
         self.logits_processor = LogitsProcessor(config)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/grok.py

@@ -601,6 +601,7 @@ class Grok1ModelForCausalLM(nn.Module):
         # Monkey patch _prepare_weights to load pre-sharded weights
         setattr(DefaultModelLoader, "_prepare_weights", _prepare_presharded_weights)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/llama2.py

@@ -275,6 +275,7 @@ class LlamaForCausalLM(nn.Module):
         self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
         self.logits_processor = LogitsProcessor(config)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/llama_classification.py

@@ -31,6 +31,7 @@ class LlamaForClassification(nn.Module):
         )
         self.eos_token_id = config.eos_token_id
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/llava.py

@@ -95,6 +95,7 @@ class LlavaLlamaForCausalLM(nn.Module):
 
         return image_features
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.LongTensor,

python/sglang/srt/models/llavavid.py

@@ -106,6 +106,7 @@ class LlavaVidForCausalLM(nn.Module):
 
         return image_features
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.LongTensor,

python/sglang/srt/models/minicpm.py

@@ -283,6 +283,7 @@ class MiniCPMForCausalLM(nn.Module):
 
         self.logits_processor = LogitsProcessor(config)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/mixtral.py

@@ -460,6 +460,7 @@ class MixtralForCausalLM(nn.Module):
         self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
         self.logits_processor = LogitsProcessor(config)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/mixtral_quant.py

@@ -322,6 +322,7 @@ class QuantMixtralForCausalLM(nn.Module):
         self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
         self.logits_processor = LogitsProcessor(config)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/qwen.py

@@ -237,6 +237,7 @@ class QWenLMHeadModel(nn.Module):
         self.lm_head = ParallelLMHead(vocab_size, config.hidden_size)
         self.logits_processor = LogitsProcessor(config)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/qwen2.py

@@ -261,6 +261,7 @@ class Qwen2ForCausalLM(nn.Module):
         self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
         self.logits_processor = LogitsProcessor(config)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/qwen2_moe.py

@@ -355,6 +355,7 @@ class Qwen2MoeForCausalLM(nn.Module):
         self.logits_processor = LogitsProcessor(config)
         self.sampler = Sampler()
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,

python/sglang/srt/models/stablelm.py

@@ -235,6 +235,7 @@ class StableLmForCausalLM(nn.Module):
         self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
         self.logits_processor = LogitsProcessor(config)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,