What's the recommended approach to enable the asynchronous checkpoint using NeMo?

[felixwqp]

Context: there are several ways to enable the async checkpoint in NeMo:

1. this document Logging and Checkpointing uses nemo.lightning.pytorch.callbacks.ModelCheckpoint by injecting to trainer callback or NeMoLogger.
2. YAML config exp_manager.create_checkpoint_callback uses the nemo.utils.callbacks.NeMoModelCheckpoint

Question: what’s the preferred way to enable checkpoint in NeMo?

My understanding is

1. nemo.lightning.pytorch.callbacks.ModelCheckpoint is the recommended for NeMo 2.0 using PyTorch Lightning’s modular abstractions
2. exp_manager.create_checkpoint_callback(nemo.utils.callbacks.NeMoModelCheckpoint) is the NeMo 1.0 fashion to enable the checkpoint using the yaml file with exp_manager configured.

can some NeMo checkpoint experts provide some guidance?

[ashors1]

Hi, thanks for the question! Your understanding is correct. NeMoModelCheckpoint is only used in NeMo 1.0. The NeMo 2.0 version is ModelCheckpoint. To enable async checkpointing in 2.0, set ckpt_async_save=True from MegatronStrategy.