diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 3151f6264a12..14991c83df57 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/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], diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index a4c927f34c12..83391e506436 100755 --- a/src/transformers/trainer.py +++ b/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: