Merge pull request #6 from efeslab/fix-copy

Fix non-blocking transfer to CPU
efeslab · Apr 28, 2024 · 227715b · 227715b
2 parents 83c16a3 + f2c111c
commit 227715b
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/src/fiddler/mixtral.py b/src/fiddler/mixtral.py
@@ -543,9 +543,7 @@ def mixtral_forward(self, input_ids, position_ids, is_decode):
                     )
 
                     if not is_cuda:
-                        experts[i_expert] = experts[i_expert].to(
-                            "cpu", non_blocking=True
-                        )
+                        experts[i_expert] = experts[i_expert].to("cpu")
 
                     # end of one expert
 
@@ -629,10 +627,8 @@ def mixtral_forward(self, input_ids, position_ids, is_decode):
                     current_state = self.run_expert_at_cpu(
                         i_layer,
                         i_expert,
-                        current_state.to("cpu", non_blocking=True),
-                        routing_weights[top_2_list, idx_list, None].to(
-                            "cpu", non_blocking=True
-                        ),
+                        current_state.to("cpu"),
+                        routing_weights[top_2_list, idx_list, None].to("cpu"),
                     )
                     inps_after_experts.index_add_(
                         0,