Implementation:Mlfoundations Open flamingo Save checkpoint

Template:Metadata

Overview

Concrete tool for saving FSDP-aware training checkpoints with trainable-only parameter filtering provided by the OpenFlamingo training module.

Description

The save_checkpoint() function handles both FSDP and non-FSDP checkpoint saving. For FSDP: sets state dict type to FULL_STATE_DICT, gathers model and optimizer state to rank 0. For both modes: filters state dict to only trainable parameters (via filter_state_dict_to_trainable), saves to {run_name}/checkpoint_{epoch}.pt, optionally uploads to wandb, and optionally deletes the previous checkpoint.

Usage

Called at the end of each training epoch in the main training loop.

Code Reference

Source

Repository https://github.com/mlfoundations/open_flamingo, File: open_flamingo/train/train_utils.py Lines L336-375

Signature

def save_checkpoint(model, optimizer, lr_scheduler, epoch, args):
    """
    Save training checkpoint.

    Args:
        model: Flamingo model (may be FSDP wrapped)
        optimizer: AdamW optimizer
        lr_scheduler: learning rate scheduler
        epoch: current epoch number
        args: training arguments with run_name, fsdp, rank, delete_previous_checkpoint,
              report_to_wandb, save_checkpoints_to_wandb flags
    """

Import

from open_flamingo.train.train_utils import save_checkpoint

I/O Contract

Inputs

Name	Type	Required	Description
model	Flamingo/FSDP	Yes	Trained model
optimizer	AdamW	Yes	Optimizer state
lr_scheduler	LRScheduler	Yes	Scheduler state
epoch	int	Yes	Current epoch
args	Namespace	Yes	Training config with run_name, fsdp, rank, and checkpoint flags

Outputs

Artifact	Description
Checkpoint file at {run_name}/checkpoint_{epoch}.pt	Contains model_state_dict, optimizer_state_dict, lr_scheduler_state_dict, epoch

Usage Examples

Saving checkpoint after each epoch:

from open_flamingo.train.train_utils import save_checkpoint

for epoch in range(num_epochs):
    train_one_epoch(model, optimizer, lr_scheduler, data_loader, epoch, args)

    # Save checkpoint at end of each epoch
    save_checkpoint(model, optimizer, lr_scheduler, epoch, args)

Related Pages

Principle:Mlfoundations_Open_flamingo_Distributed_Checkpointing

• Environment:Mlfoundations_Open_flamingo_PyTorch_CUDA_Distributed