Add section on ".qnemo" checkpoints

[janekl]

What does this PR do ?

Add section on ".qnemo" checkpoints to #9329.

Collection: [Note which collection this PR will affect]

Changelog

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

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

GitHub Actions CI

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

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first 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)

[yaoyu-33]

.qnemo would not support distributed checkpoint format? i.e. you saved with world_size 2 and have to load with world_size 2?

[janekl]

One thing to clarify is that these config.json + rank{i}.safetensors output is a TRT-LLM checkpoint. This should not be confused by distributed checkpoint in Nemo sense.

Anyway, the feature you asked for is not available in TRT-LLM currently. So to build a TRT-LLM engine with world_size=2 one needs to calibrate/quantize model to TRT-LLM checkpoint with world_size=2 and provide this as the input to trtllm-build command. In other words, world_size cannnot be changed at engine build.

[jgerh]

Edits for 1 - 21.

Checkpoints

This section presents the key functionalities of NVIDIA NeMo that pertain to checkpoint management.

Understand Checkpoint Formats

A .nemo checkpoint is essentially a tar file that combines various components of a trained model. These components include the model configurations (specified in a YAML file), the model weights, and other related artifacts such as tokenizer models or vocabulary files. This design simplifies tasks like sharing, loading, tuning, evaluating, and performing inference with the model.

On the other hand, the .ckpt file, generated during PyTorch Lightning training, contains both the model weights and the optimizer states. It is typically used to resume training from a paused state.

Sharded Model Weights

In both .nemo and .ckpt checkpoints, the model weights can be saved in either a regular format (as a single file named model_weights.ckpt within model parallelism folders) or a sharded format (where they are stored in a folder called model_weights).

Sharded model weights allow you to efficiently save and load the state of your training script across multiple GPUs or nodes. This approach avoids the necessity to modify model partitions when resuming tuning with a different model parallelism setup.

NeMo supports the distributed (sharded) checkpoint format from Megatron Core. In Megatron Core, there are two supported backends: Zarr-based and PyTorch-based.

[janekl]

@yaoyu-33 that would be sth for you to account for in the destination branch yuya/add_checkpoints_section

[jgerh]

Edits to 45-47

NVIDIA provides easy-to-use tools that enable users to convert community checkpoints into the NeMo format. These tools facilitate various operations, including resuming training, Supervised Fine-tuning (SFT), Parameter Efficient Fine-Tuning (PEFT), and deployment. Please consult our documentation for detailed instructions and guidelines. We provide comprehensive guides to assist both end users and developers.

[janekl]

@yaoyu-33 this is also sth for you to address here #9329, please have a look