(self)
| 1460 | log_dist(f'DeepSpeed LR Scheduler = {self.lr_scheduler}', ranks=[0]) |
| 1461 | |
| 1462 | def _configure_checkpointing(self): |
| 1463 | # Enable optimization to parallelize checkpointing of DP state |
| 1464 | optimize_dp_state = not self.zero_optimization_partition_weights() |
| 1465 | self.checkpoint_engine = create_checkpoint_engine(config_params=self._config, |
| 1466 | groups=groups, |
| 1467 | zero_stage=self.zero_optimization_stage(), |
| 1468 | has_moe_layers=self.has_moe_layers, |
| 1469 | optimize_dp_state=optimize_dp_state) |
| 1470 | |
| 1471 | dp_rank = groups._get_sequence_data_parallel_rank() |
| 1472 | rank = self.local_rank if self.use_node_local_storage() else dp_rank |
| 1473 | |
| 1474 | # Determine if this data parallel process needs to store the model checkpoint |
| 1475 | if self.checkpoint_engine.is_data_parallel_writer(rank) \ |
| 1476 | or (self.zero_optimization_partition_weights() and self.is_first_weights_partition_group()): |
| 1477 | self.save_non_zero_checkpoint = True |
| 1478 | |
| 1479 | if hasattr(self.optimizer, 'dp_process_group'): |
| 1480 | param_rank = dist.get_rank(group=self.optimizer.dp_process_group) |
| 1481 | |
| 1482 | # Only the first parameter parallel process needs to store the |
| 1483 | # optimizer state checkpoints for zero |
| 1484 | self.save_zero_checkpoint = param_rank == dp_rank |
| 1485 | |
| 1486 | def _scheduler_from_config(self, optimizer): |
| 1487 | scheduler_name = self.scheduler_name() |
no test coverage detected