Trainer: make checkpoints work under tensor parallelism

src/transformers/distributed/tensor_parallel.py

@@ -829,6 +829,25 @@ def apply_tensor_parallelism(model, tp_mesh):
     return model
 
 
+def shard_state_dict_for_load(state_dict: dict[str, torch.Tensor], model) -> dict[str, torch.Tensor]:
+    """Split full checkpoint tensors back over the mesh the model's parameters live on.
+
+    The inverse of `gather_state_dict_for_save`: a checkpoint holds whole tensors, while a
+    tensor-parallel model holds a shard of each one, so loading the file as it stands fails on every
+    sharded parameter's shape. Every rank must call this, since `distribute_tensor` is collective.
+    """
+    sharded = {}
+    # `remove_duplicate=False`, or a tied output embedding is missing from the map and loads unsharded
+    parameters = dict(model.named_parameters(remove_duplicate=False))
+    parameters.update(model.named_buffers(remove_duplicate=False))
+    for key, tensor in state_dict.items():
+        target = parameters.get(key)
+        if isinstance(target, torch.Tensor) and isinstance(target.data, DTensor) and not isinstance(tensor, DTensor):
+            tensor = distribute_tensor(tensor.to(target.data.device), target.data.device_mesh, target.data.placements)
+        sharded[key] = tensor
+    return sharded
+
+
 def gather_state_dict_for_save(
     state_dict: dict[str, torch.Tensor],
     _tp_plan: dict[str, str],

src/transformers/trainer.py

@@ -3433,9 +3433,12 @@ def _save_optimizer_and_scheduler(self, output_dir: str) -> None:
             save_fsdp_optimizer(
                 self.accelerator.state.fsdp_plugin, self.accelerator, self.optimizer, self.model, output_dir
             )
+        elif self.get_tp_size() > 1:
+            os.makedirs(output_dir, exist_ok=True)
+            torch.save(self.optimizer.state_dict(), self._get_optimizer_file(output_dir))
         elif self.args.should_save:
             # deepspeed.save_checkpoint above saves model/optim/sched
-            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
+            torch.save(self.optimizer.state_dict(), self._get_optimizer_file(output_dir))
 
         # Save SCHEDULER & SCALER
         is_deepspeed_custom_scheduler = self.is_deepspeed_enabled and not isinstance(
@@ -3564,6 +3567,12 @@ def _load_from_checkpoint(self, resume_from_checkpoint: str, model: nn.Module |
                     check_torch_load_is_safe()
                     state_dict = torch.load(weights_file, map_location="cpu", weights_only=True)
 
+                # A checkpoint holds whole tensors while a tensor-parallel model holds a shard of each
+                if self.get_tp_size() > 1:
+                    from .distributed.tensor_parallel import shard_state_dict_for_load
+
+                    state_dict = shard_state_dict_for_load(state_dict, model)
+
                 # workaround for FSDP bug https://github.com/pytorch/pytorch/issues/82963
                 # which takes *args instead of **kwargs
                 load_result = model.load_state_dict(state_dict, False)
@@ -3753,6 +3762,19 @@ def _load_rng_state(self, checkpoint: str | None) -> None:
         if is_torch_musa_available():
             set_rng_state_for_device("MUSA", torch.musa, checkpoint_rng_state, is_distributed)
 
+    def _get_optimizer_file(self, checkpoint: str) -> str:
+        """Path of the optimizer state to write or read.
+
+        Under tensor parallelism each rank holds a different shard of every parameter, so its optimizer
+        state is its own and cannot be read from another rank's file. The shapes match across ranks, so
+        a single shared file loads without error and silently gives every rank rank 0's moments.
+        """
+        if self.get_tp_size() > 1:
+            return os.path.join(
+                checkpoint, f"rank{self.args.process_index}-of-{self.args.world_size}-{OPTIMIZER_NAME}"
+            )
+        return os.path.join(checkpoint, OPTIMIZER_NAME)
+
     def _load_optimizer_and_scheduler(self, checkpoint: str | None) -> None:
         """If optimizer and scheduler states exist, load them."""
         if checkpoint is None:
@@ -3787,7 +3809,7 @@ def _load_optimizer_and_scheduler(self, checkpoint: str | None) -> None:
         )
         checkpoint_file_exists = (
             glob.glob(os.path.join(checkpoint, f"rank*-of-{self.args.world_size}-{OPTIMIZER_NAME}"))
-            if self.is_fsdp_xla_v1_enabled
+            if self.is_fsdp_xla_v1_enabled or self.get_tp_size() > 1
             else checkpoint_file_exists
         )
         if checkpoint_file_exists and os.path.isfile(os.path.join(checkpoint, SCHEDULER_NAME)):
@@ -3847,7 +3869,7 @@ def opt_load_hook(mod, opt):
                         check_torch_load_is_safe()
                         self.optimizer.load_state_dict(
                             torch.load(
-                                os.path.join(checkpoint, OPTIMIZER_NAME), map_location=map_location, weights_only=True
+                                self._get_optimizer_file(checkpoint), map_location=map_location, weights_only=True
                             )
                         )
                 with warnings.catch_warnings(record=True) as caught_warnings: