Save policy, optimizer, and scheduler state to disk, gathering from all processes and saving only on the rank 0 process.
(self, output_dir=None, metrics=None)
| 495 | return self.policy.clip_grad_norm_(self.config.max_grad_norm).item() |
| 496 | |
| 497 | def save(self, output_dir=None, metrics=None): |
| 498 | """Save policy, optimizer, and scheduler state to disk, gathering from all processes and saving only on the rank 0 process.""" |
| 499 | save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) |
| 500 | with FSDP.state_dict_type(self.policy, StateDictType.FULL_STATE_DICT, state_dict_config=save_policy): |
| 501 | policy_state_dict = self.policy.state_dict() |
| 502 | |
| 503 | if self.rank == 0: |
| 504 | self.write_state_dict(self.example_counter, policy_state_dict, metrics, 'policy.pt', output_dir) |
| 505 | del policy_state_dict |
| 506 | dist.barrier() |
| 507 | |
| 508 | save_policy = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True) |
| 509 | with FSDP.state_dict_type(self.policy, StateDictType.FULL_STATE_DICT, optim_state_dict_config=save_policy): |
| 510 | optimizer_state_dict = FSDP.optim_state_dict(self.policy, self.optimizer) |
| 511 | |
| 512 | if self.rank == 0: |
| 513 | self.write_state_dict(self.example_counter, optimizer_state_dict, metrics, 'optimizer.pt', output_dir) |
| 514 | del optimizer_state_dict |
| 515 | dist.barrier() |
| 516 | |
| 517 | if self.rank == 0: |
| 518 | scheduler_state_dict = self.scheduler.state_dict() |
| 519 | self.write_state_dict(self.example_counter, scheduler_state_dict, metrics, 'scheduler.pt', output_dir) |
| 520 | dist.barrier() |
| 521 | |
| 522 | |
| 523 | class TensorParallelTrainer(BasicTrainer): |
no test coverage detected