Merge branch 'main' into dist-ckp

Author: flybird11111

colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py

@@ -118,15 +118,11 @@ def _model_sharder(
         for name, param in model.named_parameters():
             if param is None:
                 continue
-
-            if gather_dtensor:
-                # 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)
-            else:
-                param_ = param
-
+                
+            # Gather tensor pieces when using tensor parallel.
+            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")