Merge branch 'master' into tohtana/fix_zero_init_patch

microsoft · Aug 21, 2024 · ac3eb90 · ac3eb90
2 parents d555266 + b65ea50
commit ac3eb90
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),