Commit 61f72a8c authored by Anthony Chen's avatar Anthony Chen Committed by Facebook GitHub Bot
Browse files

Enable preemption checkpointing for d2go FSDPCheckpointer

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/567

As title.

Reviewed By: tglik

Differential Revision: D46383823

fbshipit-source-id: b5f80f55eb37ddc4e0918a349840b451f2b4b094
parent 3ecf8806
...@@ -139,6 +139,8 @@ class FSDPCheckpointer(QATCheckpointer): ...@@ -139,6 +139,8 @@ class FSDPCheckpointer(QATCheckpointer):
Add support for saving sharding models and optimizers. Add support for saving sharding models and optimizers.
The rest of the code is copied from implementation in the superclass The rest of the code is copied from implementation in the superclass
""" """
# checkpoint_type is used to annotate preemption checkpoints for internal checkpointer. Ignore it here
kwargs.pop("checkpoint_type", None)
# If no sharding, only the main process enters the saving codepath; # If no sharding, only the main process enters the saving codepath;
# otherwise, all processes need to call state_dict() to enable state broadcasting among ranks # otherwise, all processes need to call state_dict() to enable state broadcasting among ranks
if not isinstance(self.model, FSDPWrapper): if not isinstance(self.model, FSDPWrapper):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment