[Core] Support pre-extracted nemo checkpoint for restoration

[titu1994]

Signed-off-by: smajumdar titu1994@gmail.com

What does this PR do ?

Adds support to SaveRestoreConnector to employ a pre-extracted nemo directory path as a restoration directory.
This is required for extremely large language model restoration, which cannot be extracted partially during multinode inference.

Collection: [Core, Megatron]

Changelog

• Add specific line by line info of high level changes in this PR.

Note

• During restore_from, the path of the nemo file is ignored in its entirety since the actual pre-extracted directory is passed.
• The filepath passed to restore_from in such a case is not preserved, it is overwritten by the value of connector.model_extracted_dir inside AppState.

Usage

# Extract a nemo file apriori to restoration
# Here we will use the SaveRestoreConnector's method to extract the checkpoint to a given directory 

connector = MySaveRestoreConnector()
MySaveRestoreConnector._unpack_nemo_file(nemo_filepath, extracted_directory)

# Then update SaveRestoreConnector with path of the extracted directory
connector.model_extracted_dir = extracted_directory

# Call restore as always
# note, we pass in the "old" nemo_filepath, stored somewhere other than the extracted directory
restored_model = ModelPT.restore_from(nemo_filepath, save_restore_connector=connector)

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

Who can review?

@okuchaiev @ericharper

[okuchaiev]

/blossom-ci

[okuchaiev]

/blossom-ci

[okuchaiev]

/blossom-ci

[okuchaiev]

what if the user provided absolute path outside of their homedir?

[titu1994]

expanduser() only expands ~ in a user provided path, so if absolute path is given then as long as theres no ~ in it, it will be correct expansion.

[okuchaiev]

thanks for adding unittests!