From c30bede2248552b37ee9a5825c78210a3e2ff08a Mon Sep 17 00:00:00 2001 From: flybird11111 <1829166702@qq.com> Date: Wed, 27 Nov 2024 14:54:40 +0800 Subject: [PATCH] fix --- colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py b/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py index 4c1defe27556..9373379e5ee5 100644 --- a/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py +++ b/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py @@ -112,8 +112,6 @@ def _model_sharder( pinned_state_dicts[prefix + name] = torch.empty_like(param_, pin_memory=True, device="cpu") pinned_state_dicts[prefix + name].copy_(param_) param_ = pinned_state_dicts[prefix + name] - else: - param_ = param_.cpu() block, block_size = state_dict_sharder.append_param(prefix + name, param_) if block is not None: yield block, block_size @@ -128,8 +126,6 @@ def _model_sharder( pinned_state_dicts[prefix + name] = torch.empty_like(param_, pin_memory=True, device="cpu") pinned_state_dicts[prefix + name].copy_(buffer) buffer = pinned_state_dicts[prefix + name] - else: - buffer = buffer.cpu() block, block_size = state_dict_sharder.append_param(prefix + name, buffer) if block is not None: yield block, block_size @@ -146,8 +142,6 @@ def _model_sharder( pinned_state_dicts[extra_state_key] = torch.empty_like(param_, pin_memory=True, device="cpu") pinned_state_dicts[extra_state_key].copy_(extra_state) extra_state = pinned_state_dicts[extra_state_key] - else: - extra_state = extra_state.cpu() block, block_size = state_dict_sharder.append_param(extra_state_key, extra_state) if block is not None: yield block, block_size