megatron.core.tensor_parallel.random.checkpoint

megatron.core.tensor_parallel.random.checkpoint(function, distribute_saved_activations, *args)

Checkpoint a model or part of the model. This has been directly copied from torch.utils.checkpoint.