S3 Dirpath + Async Uploading Support for Default Checkpoints

[alxzhang-amazon]

What does this PR do ?

This PR creates a new CheckpointIO which allows users to upload checkpoints directly to S3.

Add a one line overview of what this PR aims to accomplish.

Collection: [NLP]

Changelog

• Adds a new S3CheckpointIO to enable strategy to upload default checkpoints to S3.
• Adds a checkpoint_file_utils file which includes identifies existing paths
• • used in S3CheckpointIO to cleanup existing checkpoints
• Adds a S3Utils class that contains helper methods for interacting with S3.
• • Writes files using CRT built in due to installing boto3[crt]: reference
• Update exp_manager check_resume to only run on rank 0 when using an S3 dirpath (prevents throttling S3 due to check_resume operations).
• Add NeMoCheckpointConnector to exp_manager to call resume_start using the broadcasted checkpoint path.
• Add setup() override in NeMoModelCheckpoint to broadcast the trainer.ckpt_path since only Rank 0 has it after check_resume
• Updated requirements_nlp.txt to include support for S3 filesystem as well as CRT support for faster uploading.

Usage

• To upload checkpoints to S3, add the checkpoint connector to the trainer, and apply the S3CheckpointIO to your existing strategy.
• Update the dirpath field in your config to an s3 address.

Example from megatron_gpt_pretraining.py example

def main(cfg) -> None:
    logging.info("\n\n************** Experiment configuration ***********")
    logging.info(f'\n{OmegaConf.to_yaml(cfg)}')
 
    trainer = MegatronTrainerBuilder(cfg).create_trainer()
    trainer._checkpoint_connector = NeMoCheckpointConnector(trainer)
    exp_manager(trainer, cfg.exp_manager)
 
    model = MegatronGPTModel(cfg.model, trainer)
    trainer.fit(model)
 
if __name__ == '__main__':
    main()

async_checkpointing = self.cfg.checkpointing.get('enable_async_checkpointing', False)
dirpath = self.cfg.exp_manager.checkpoint_callback_params.get('dirpath')
s3_checkpoint_io = S3CheckpointIO(dirpath=dirpath, async_checkpointing=async_checkpointing)
return NLPDDPStrategy(
    no_ddp_communication_hook=True,
    checkpoint_io=s3_checkpoint_io,
    gradient_as_bucket_view=self.cfg.model.gradient_as_bucket_view,
    find_unused_parameters=False,
    nccl_communicator_config_path=self.cfg.model.get('nccl_communicator_config_path', None),
    sharp=self.cfg.model.get('sharp', False),
)

Config Updates:

checkpoint_callback_params:
    dirpath: s3://mstar-eks-dev-us-east-2/alxzhang/nemo123/1n/checkpoints
 
checkpointing:
  # write_concurrency * tp * pp * 1.15 (buffer) should be within 3500 S3 TPS limit per partition
  max_write_concurrency: 10
  # read_concurrency * tp * pp * 1.15 (buffer) should be within 5500 S3 TPS limit per partition
  max_read_concurrency: 15
  chunk_size_MB: 64
  # enables asynchronous checkpoint writing to S3
  enable_async_checkpointing: False

Jenkins CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

There's no need to comment jenkins on the PR to trigger Jenkins CI.
The GitHub Actions CI will run automatically when the PR is opened.
To run CI on an untrusted fork, a NeMo user with write access must click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[mikolajblaz]

I added some high-level comments.

How does this PR relate to this one: NVIDIA/Megatron-LM#748?

[alxzhang-amazon]

I added some high-level comments.

How does this PR relate to this one: NVIDIA/Megatron-LM#748?

I added some high-level comments.

How does this PR relate to this one: NVIDIA/Megatron-LM#748?

The PR linked is for supporting S3 checkpointing for the distributed checkpoint format.
This PR is for the legacy format, but also can help the linked PR by enabling handling of S3 dirpaths in the exp_manager

[ericharper]

LGTM. Thanks!

[ericharper]

LGTM. Thanks!