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