[checkpointio] gather tensor before unpad it if the tensor is both pa…

…dded and distributed (#6168)
hpcaitech · Jan 21, 2025 · 97e60cb · 97e60cb
1 parent 5b094a8
commit 97e60cb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py b/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py
@@ -107,9 +107,9 @@ def _model_sharder(
             if param is None:
                 continue
             # Gather tensor pieces when using tensor parallel.
-            if is_padded_tensor(param):
-                param = to_unpadded_tensor(param)
             param_ = gather_distributed_param(param, keep_vars=False)
+            if is_padded_tensor(param_):
+                param_ = to_unpadded_tensor(param_)
             if pinned_state_dicts is not None:
                 if (prefix + name) not in pinned_state_dicts:
                     pinned_state_dicts[prefix + name] = torch.empty_like(param_, pin_memory=True, device="cpu")