GDS Swapping Fix (#6386)

Fixing gds api call Co-authored-by: Ubuntu <deepspeed@H100-VM2.shlnn55tgwve1eacvp21ie45dg.jx.internal.cloudapp.net>
microsoft · Aug 20, 2024 · b65ea50 · b65ea50
1 parent 96393f5
commit b65ea50
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py b/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py
@@ -121,7 +121,7 @@ def _configure_aio(self, ds_config):
                                        dtype=self.dtype,
                                        device=get_accelerator().device_name(),
                                        requires_grad=False)
-            self.aio_read_handle.new_device_locked_tensor(self.buffers)
+            self.aio_read_handle.pin_device_tensor(self.buffers)
         else:
             self.buffers = get_accelerator().pin_memory(torch.empty(int(self.aligned_elements_per_buffer *
                                                                         self.param_buffer_count),