(self)
| 1517 | log_dist(f'DeepSpeed LR Scheduler = {self.lr_scheduler}', ranks=[0]) |
| 1518 | |
| 1519 | def _configure_checkpointing(self): |
| 1520 | # Enable optimization to parallelize checkpointing of DP state |
| 1521 | optimize_dp_state = not self.zero_optimization_partition_weights() |
| 1522 | self.checkpoint_engine = create_checkpoint_engine(config_params=self._config, |
| 1523 | groups=groups, |
| 1524 | zero_stage=self.zero_optimization_stage(), |
| 1525 | has_moe_layers=self.has_moe_layers, |
| 1526 | optimize_dp_state=optimize_dp_state) |
| 1527 | |
| 1528 | dp_rank = groups._get_sequence_data_parallel_rank() |
| 1529 | rank = self.local_rank if self.use_node_local_storage() else dp_rank |
| 1530 | |
| 1531 | # Determine if this data parallel process needs to store the model checkpoint |
| 1532 | if self.checkpoint_engine.is_data_parallel_writer(rank) \ |
| 1533 | or (self.zero_optimization_partition_weights() and self.is_first_weights_partition_group()): |
| 1534 | self.save_non_zero_checkpoint = True |
| 1535 | |
| 1536 | if hasattr(self.optimizer, 'dp_process_group'): |
| 1537 | param_rank = dist.get_rank(group=self.optimizer.dp_process_group) |
| 1538 | |
| 1539 | # Only the first parameter parallel process needs to store the |
| 1540 | # optimizer state checkpoints for zero |
| 1541 | self.save_zero_checkpoint = param_rank == dp_rank |
| 1542 | |
| 1543 | def _scheduler_from_config(self, optimizer): |
| 1544 | scheduler_name = self.scheduler_name() |
no test coverage detected