Adding the ability to specify a checkpointing for models **before** aggregation. #128
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…g functionality to allow for pre- and post-aggregation (server-side) checkpointing. Also allowing for more generic checkpointing functionality. That is, given a loss value and a dictionary of metrics, users can define an arbitrary scoring function on those objects to produce a checkpoint. The best loss checkpointer is a specific instantiation of this type of checkpointer.
emersodb
requested review from
lotif,
fatemetkl,
jewelltaylor,
sanaAyrml and
yc7z
emersodb
commented
Apr 23, 2024
| checkpoint_dir, checkpoint_name, checkpoint_score_function=loss_score_function, maximize=False | ||
| ) | ||
|
|
||
| def maybe_checkpoint(self, model: nn.Module, loss: float, metrics: Dict[str, Scalar]) -> None: |
There was a problem hiding this comment.
Note, I'm overriding this method to replace the logging with something more specific. If anyone has any better ideas on how to do this, let me know.
emersodb
commented
Apr 23, 2024
| @@ -631,8 +657,6 @@ def validate(self) -> Tuple[float, Dict[str, Scalar]]: | |||
| metrics = self.val_metric_manager.compute() | |||
| self._handle_logging(loss_dict, metrics, is_validation=True) | |||
|
|
|||
| # Checkpoint based on loss which is output of user defined compute_loss method | |||
| self._maybe_checkpoint(loss_dict["checkpoint"]) | |||
There was a problem hiding this comment.
Note that I moved this into the evaluate function rather than the validate function, as we may not want to checkpoint every time we validate etc.
fatemetkl
approved these changes
May 6, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Note: There are a scary number of files changed because this affects the functionality of BasicClient in a way that requires slight modifications of many of the inheriting classes. Most of the changes should be quite minor in those downstream classes.
PR Type
Feature
Clickup Ticket(s):
Refactoring the client-side checkpointing functionality to allow for pre- and post-aggregation (server-side) checkpointing. That is, the user can specify a checkpointer for models after the weights have been aggregated by the server (still supported, but was the only default option), or prior to this happening (i.e. right after training). If either of the checkpointers are not specified, then checkpointing at that time is skipped. Note, this allows us to mimic "fine-tuning" of models. That is, the models final training is exclusively local.
For example, with FedProx, the model checkpointed would actually be a "personal" model for each client (potentially after several preceding rounds of aggregation).
Also allowing for more generic checkpointing functionality. That is, given a loss value and a dictionary of metrics, users can define an arbitrary scoring function on those objects to produce a checkpoint. The best loss checkpointer is a specific instantiation of this type of checkpointer.
Tests Added
Added some new tests to cover the new ClientSideCheckpointModule functionality and also the new functionality associated with the TorchCheckpointer child classes.