def _checkpoint_model(self,
task,
train_phase_idx,
mode_frequency,
mode_num,
mode="phase"):
"""
Checkpoint model. Can be called in 3 possible scenarios:
1. If training becomes NaN, then we checkpoint the model to facilitate debugging
2. After every N epochs (CHECKPOINT_FREQ), model state is checkpointed.
3. If user wants to checkpoint during the epoch (ie. after every few training
iterations, the model state is checkpointed.)
Args:
task: Self-supervision task that hold information about training iteration,
epoch number etc.
train_phase_idx (int): current training phase number. Starts from 0
mode_frequency (int): mode can be "phase" or "iteration". Frequency
of checkpointing for the given mode
mode_num (int): for the checkpointing mode (phase or iteration), the number
of phase or iteration at which checkpointing is being done
"""
phase_idx = task.phase_idx
# num_train_phases = num_epochs * num_phases_per_epoch
# For OSS use, num_train_phases will be equal to num_epochs
num_train_phases = task.num_train_phases
# check if we need to checkpoint this phase
is_checkpointing_phase = is_checkpoint_phase(mode_num, mode_frequency,
train_phase_idx,
num_train_phases, mode)
is_final_train_phase = ((train_phase_idx == (num_train_phases - 1))
and task.train and mode == "phase")
# handle checkpoint:
if task.train and (is_final_train_phase or is_checkpointing_phase):
# - if sharded state consolidate the state
# /!\ All the ranks have to participate
if hasattr(task.optimizer,
"consolidate_state_dict") and mode != "phase":
logging.info(
f"[{mode}: {mode_num}] Consolidating sharded state on all replicas"
)
task.optimizer.consolidate_state_dict()
# Depending on whether we are in FSDP mode or not
# - save the checkpoint on the primary rank
# - save the sharded checkpoint on all ranks
if is_primary() or isinstance(task.base_model, FSDP):
checkpoint_folder = task.checkpoint_folder
logging.info(
f"[{mode}: {mode_num}] Saving checkpoint to {checkpoint_folder}"
)
model_state_dict = task.get_classy_state()
# phase_idx is already incremented at the beginning of phase but if we
# are checkpointing at an iteration in the middle of phase, we should not
# save the incremented phase_idx as it will incorrectly assume that model
# trained for that phase already.
if mode == "iteration":
model_state_dict[
"phase_idx"] = model_state_dict["phase_idx"] - 1
if task.train:
train_phase_idx = train_phase_idx - 1
model_state_dict["train_phase_idx"] = train_phase_idx
restart_phase = phase_idx - 1
restart_iteration = task.iteration
# When loading from a phase checkpoint:
else:
restart_phase = phase_idx
restart_iteration = task.iteration
checkpoint_content = {
"phase_idx": restart_phase,
"iteration": restart_iteration,
"loss": task.loss.state_dict(),
"iteration_num": task.local_iteration_num,
"train_phase_idx": train_phase_idx,
"classy_state_dict": model_state_dict,
}
checkpoint_writer = CheckpointWriter(
checkpoint_folder=checkpoint_folder,
is_final_train_phase=is_final_train_phase,
mode=mode,
mode_num=mode_num,
backend=task.config["CHECKPOINT"]["BACKEND"],
)
if isinstance(task.base_model, FSDP):
_, rank = get_machine_local_and_dist_rank()
checkpoint_writer.save_sharded_checkpoint(
content=checkpoint_content,
shard_rank=rank,
world_size=self.world_size,
)
else:
checkpoint_writer.save_consolidated_checkpoint(
checkpoint_content)
def _checkpoint_model(self,
task,
train_phase_idx,
mode_frequency,
mode_num,
mode="phase"):
"""
Checkpoint model. Can be called in 3 possible scenarios:
1. If training becomes NaN, then we checkpoint the model to facilitate debugging
2. After every N epochs (CHECKPOINT_FREQ), model state is checkpointed.
3. If user wants to checkpoint during the epoch (ie. after every few training
iterations, the model state is checkpointed.)
Args:
task: Self-supervision task that hold information about training iteration,
epoch number etc.
train_phase_idx (int): current training phase number. Starts from 0
mode_frequency (int): mode can be "phase" or "iteration". Frequency
of checkpointing for the given mode
mode_num (int): for the checkpointing mode (phase or iteration), the number
of phase or iteration at which checkpointing is being done
"""
phase_idx = task.phase_idx
num_epochs = task.num_epochs
# check if we need to checkpoint this phase
is_checkpointing_phase = is_checkpoint_phase(mode_num, mode_frequency,
train_phase_idx,
num_epochs, mode)
is_final_train_phase = ((train_phase_idx == (num_epochs - 1))
and task.train and mode == "phase")
# handle checkpoint:
if task.train and (is_final_train_phase or is_checkpointing_phase):
# - if sharded state consolidate the state
# /!\ All the ranks have to participate
if hasattr(task.optimizer,
"consolidate_state_dict") and mode != "phase":
logging.info(
f"[{mode}: {mode_num}] Consolidating sharded state on all replicas"
)
task.optimizer.consolidate_state_dict()
# Model's state dict may need to be obtained on all ranks if we are running
# with FSDP since all_gather needs to happen here.
model_state_dict = None
if isinstance(task.base_model, FSDP):
model_state_dict = task.get_classy_state()
# - save the checkpoint on the primary rank
if is_primary():
checkpoint_folder = task.checkpoint_folder
logging.info(
f"[{mode}: {mode_num}] Saving checkpoint to {checkpoint_folder}"
)
if model_state_dict is None:
model_state_dict = task.get_classy_state()
# phase_idx is already incremented at the beginning of phase but if we
# are checkpointing at an iteration in the middle of phase, we should not
# save the incremented phase_idx as it will incorrectly assume that model
# trained for that phase already.
if mode == "iteration":
phase_idx = phase_idx - 1
model_state_dict[
"phase_idx"] = model_state_dict["phase_idx"] - 1
if task.train:
train_phase_idx = train_phase_idx - 1
model_state_dict["train_phase_idx"] = train_phase_idx
checkpoint_task = {
"phase_idx": phase_idx,
"iteration": task.iteration,
"loss": task.loss.state_dict(),
"iteration_num": task.local_iteration_num,
"train_phase_idx": train_phase_idx,
# TODO (Min): change the key to model_state_dict but we need to be careful
# about backward compatibilities.
"classy_state_dict": model_state_dict,
}
ckpt_name = f"model_{mode}{mode_num}.torch"
if is_final_train_phase:
ckpt_name = f"model_final_checkpoint_{mode}{mode_num}.torch"
backend = task.config["CHECKPOINT"]["BACKEND"]
assert backend == "disk", "Only disk BACKEND supported"
save_checkpoint(checkpoint_folder,
checkpoint_task,
checkpoint_file=ckpt_name)
logging.info(
f"Saved checkpoint: {checkpoint_folder}/{ckpt_name}")
# we create the checkpoint symlink and use this symlink to load
# checkpoints. This helps ensure that the checkpoint we load from
# are valid. It's a particularly useful feature for resuming trainings.
logging.info("Creating symlink...")
symlink_dest_file = f"{checkpoint_folder}/checkpoint.torch"
source_file = f"{checkpoint_folder}/{ckpt_name}"
create_file_symlink(source_file, symlink_dest_file)
logging.info(f"Created symlink: {symlink_dest_file}")