From b5709cce66ac09c879054dba0576e8ab0d770f73 Mon Sep 17 00:00:00 2001 From: Max Kovalenko Date: Thu, 21 Nov 2024 18:01:13 +0200 Subject: [PATCH] Enable torch compile on _allgather_params (#6769) * Previosuly ZerO3 was crashing when trying to compile _allgather_params * Disabling grad solves the issue --- deepspeed/runtime/zero/partition_parameters.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py index 22a6746bb57c..cb0cd7c8017d 100755 --- a/deepspeed/runtime/zero/partition_parameters.py +++ b/deepspeed/runtime/zero/partition_parameters.py @@ -1882,6 +1882,7 @@ def _allgather_params_coalesced(self, param_list, hierarchy=0, quantize=False): return None + @torch.no_grad() def _allgather_params(self, param_list, hierarchy=0): if len(param_list) == 0: return