S3 Dirpath + Async Uploading Support for Default Checkpoints
State: closed | Pull request: yes
Labels: name: NLP

name: Run CICD

# 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](https://aws.amazon.com/blogs/storage/improving-amazon-s3-throughput-for-the-aws-cli-and-boto3-with-the-aws-common-runtime/)
- 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
```python
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**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] 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**:
- [x] 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](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.
 
# Additional Information
* Related to # (issue)


> I added some high-level comments.
> 
> How does this PR relate to this one: [NVIDIA/Megatron-LM#748](https://github.com/NVIDIA/Megatron-LM/pull/748)?



> I added some high-level comments.
> 
> How does this PR relate to this one: [NVIDIA/Megatron-LM#748](https://github.com/NVIDIA/Megatron-LM/pull/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

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

This PR was closed because it has been inactive for 7 days since being marked as stale.