Save 16bit model weights This method saves the 16bit model weights at the desired destination. Arguments: save_dir: Required. Directory for saving the model save_filename: Optional. Filename to save to. Defaults to ``pytorch_model.bin``
(self, save_dir, save_filename="pytorch_model.bin", exclude_frozen_parameters=False)
| 5484 | return self.save_16bit_model(save_dir, save_filename) |
| 5485 | |
| 5486 | def save_16bit_model(self, save_dir, save_filename="pytorch_model.bin", exclude_frozen_parameters=False): |
| 5487 | """ |
| 5488 | Save 16bit model weights |
| 5489 | |
| 5490 | This method saves the 16bit model weights at the desired destination. |
| 5491 | |
| 5492 | Arguments: |
| 5493 | save_dir: Required. Directory for saving the model |
| 5494 | save_filename: Optional. Filename to save to. Defaults to ``pytorch_model.bin`` |
| 5495 | exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state. |
| 5496 | |
| 5497 | Returns: |
| 5498 | ``True`` when a model has been saved, ``False`` otherwise. It will not be saved if |
| 5499 | stage3_gather_16bit_weights_on_model_save is ``False``. |
| 5500 | |
| 5501 | Important: all processes must call this method and not just the process with rank 0. It is |
| 5502 | because the processes need to work in sync to gather the weights. This method will hang |
| 5503 | waiting to synchronize with other processes if it's called just for the process with rank 0. |
| 5504 | |
| 5505 | """ |
| 5506 | |
| 5507 | path = os.path.join(save_dir, save_filename) |
| 5508 | |
| 5509 | if self.zero_optimization_partition_weights(): |
| 5510 | self._raise_if_autoep_zero3_consolidated_export("save_16bit_model") |
| 5511 | if self.zero_gather_16bit_weights_on_model_save(): |
| 5512 | # consolidation is expensive in time and memory and therefore isn't a default |
| 5513 | state_dict = self._zero3_consolidated_16bit_state_dict( |
| 5514 | exclude_frozen_parameters=exclude_frozen_parameters) |
| 5515 | else: |
| 5516 | # the model will be bogus if not consolidated so don't confuse the user by saving it |
| 5517 | logger.info( |
| 5518 | f"Did not save the model {path} because stage3_gather_16bit_weights_on_model_save is False") |
| 5519 | return False |
| 5520 | else: |
| 5521 | state_dict = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters) |
| 5522 | |
| 5523 | tag = f"global_step{self.global_steps}" |
| 5524 | tag = str(tag) |
| 5525 | commit_info = CheckpointCommitInfo(tag=tag, save_dir=save_dir, save_latest=False) |
| 5526 | self.checkpoint_engine.create(commit_info) |
| 5527 | |
| 5528 | if dist.get_rank() == 0: |
| 5529 | self.checkpoint_engine.makedirs(save_dir, exist_ok=True) |
| 5530 | logger.info(f"Saving model weights to {path}, tag: {tag}") |
| 5531 | self.checkpoint_engine.save(state_dict, path) |
| 5532 | |
| 5533 | self.checkpoint_engine.commit(commit_info) |
| 5534 | |
| 5535 | return True |
| 5536 | |
| 5537 | def empty_partition_cache(self): |
| 5538 | """ |