add emtpy_cache() after each padding (opendatahub-io#120)

prarit · Aug 6, 2024 · 98f31cd · 98f31cd
1 parent c034d5d
commit 98f31cd
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py
@@ -187,9 +187,11 @@ def process_weights_after_loading(self):
                 self.w13_weight = nn.Parameter(F.pad(self.w13_weight.data,
                                                      (0, 128), "constant", 0),
                                                requires_grad=False)
+                torch.cuda.empty_cache()
                 self.w2_weight = nn.Parameter(F.pad(self.w2_weight.data,
                                                     (0, 128), "constant", 0),
                                               requires_grad=False)
+                torch.cuda.empty_cache()
             return
 
         # If checkpoint is fp16, quantize here.