From 6f061899aa3f62d5cb24204c3b7d354c242f7fc4 Mon Sep 17 00:00:00 2001 From: amaurya Date: Wed, 10 Jul 2024 19:26:40 +0000 Subject: [PATCH] Fix memory leak for pipelined optimizer swapper --- deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py b/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py index 03dfe663fcb6..66a372877d38 100644 --- a/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py +++ b/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py @@ -8,6 +8,7 @@ from deepspeed.ops.op_builder import AsyncIOBuilder from deepspeed import comm as dist +import torch from deepspeed.runtime.swap_tensor.constants import * from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, print_object @@ -154,6 +155,8 @@ def swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors): def _complete_swap_out(self, swap_out_type): self.swap_ops[swap_out_type].wait() + for buffer in self.swap_ops[swap_out_type].state_buffers: + buffer = torch.Tensor() self.swap_buffer_manager.free(self.swap_ops[swap_out_type].allocated_buffers) self.swap_ops[swap_out_type] = None