|
Checkpoints |
|
=========== |
|
|
|
In this section, we present four key functionalities of NVIDIA NeMo related to checkpoint management: |
|
|
|
1. **Checkpoint Loading**: Use the :code:`restore_from()` method to load local ``.nemo`` checkpoint files. |
|
2. **Partial Checkpoint Conversion**: Convert partially-trained ``.ckpt`` checkpoints to the ``.nemo`` format. |
|
3. **Community Checkpoint Conversion**: Convert checkpoints from community sources, like HuggingFace, into the ``.nemo`` format. |
|
4. **Model Parallelism Adjustment**: Adjusting model parallelism is crucial when training models that surpass the memory capacity of a single GPU, such as the NVGPT 5B version, LLaMA2 7B version, or larger models. NeMo incorporates both tensor (intra-layer) and pipeline (inter-layer) model parallelisms. For a deeper understanding, refer to "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM" (`link <https://arxiv.org/pdf/2104.04473.pdf>`_). This tool assists in modifying model parallelism. After downloading and converting a community checkpoint to the ``.nemo`` format, if a user wishes to fine-tune it further, this adjustment might become essential. |
|
|
|
Understanding Checkpoint Formats |
|
-------------------------------- |
|
|
|
A ``.nemo`` checkpoint is fundamentally a tar file that bundles the model configurations (given as a YAML file), model weights, and other pertinent artifacts like tokenizer models or vocabulary files. This consolidated design streamlines sharing, loading, tuning, evaluating, and inference. |
|
|
|
On the other hand, the ``.ckpt`` file, created during PyTorch Lightning training, contains only the model weights and the optimizer states, which is used to pick up training from a pause. |
|
|
|
The subsequent sections elucidate instructions for the functionalities above, specifically tailored for deploying fully trained checkpoints for assessment or additional fine-tuning. |
|
|
|
Loading Local Checkpoints |
|
------------------------- |
|
|
|
By default, NeMo saves checkpoints of trained models in the ``.nemo`` format. To save a model manually during training, use: |
|
|
|
.. code-block:: python |
|
|
|
model.save_to(<checkpoint_path>.nemo) |
|
|
|
To load a local ``.nemo`` checkpoint: |
|
|
|
.. code-block:: python |
|
|
|
import nemo.collections.multimodal as nemo_multimodal |
|
model = nemo_multimodal.models.<MODEL_BASE_CLASS>.restore_from(restore_path="<path/to/checkpoint/file.nemo>") |
|
|
|
Replace `<MODEL_BASE_CLASS>` with the appropriate MM model class. |
|
|
|
Converting Local Checkpoints |
|
---------------------------- |
|
|
|
Only the last checkpoint is automatically saved in the ``.nemo`` format. If intermediate training checkpoints evaluation is required, a ``.nemo`` conversion might be necessary. For this, refer to the script at `<ADD convert_ckpt_to_nemo.py PATH>`: |
|
|
|
.. code-block:: python |
|
|
|
python -m torch.distributed.launch --nproc_per_node=<tensor_model_parallel_size> * <pipeline_model_parallel_size> \ |
|
examples/multimodal/convert_ckpt_to_nemo.py \ |
|
--checkpoint_folder <path_to_PTL_checkpoints_folder> \ |
|
--checkpoint_name <checkpoint_name> \ |
|
--nemo_file_path <path_to_output_nemo_file> \ |
|
--tensor_model_parallel_size <tensor_model_parallel_size> \ |
|
--pipeline_model_parallel_size <pipeline_model_parallel_size> |
|
|
|
Converting Community Checkpoints |
|
-------------------------------- |
|
|
|
There is no support for converting community checkpoints to NeMo ViT. |
|
|
|
Model Parallelism Adjustment |
|
---------------------------- |
|
|
|
ViT Checkpoints |
|
^^^^^^^^^^^^^^^^ |
|
|
|
To adjust model parallelism from original model parallelism size to a new model parallelism size (Note: NeMo ViT currently only supports `pipeline_model_parallel_size=1`): |
|
|
|
.. code-block:: bash |
|
|
|
python examples/nlp/language_modeling/megatron_change_num_partitions.py \ |
|
--model_file=/path/to/source.nemo \ |
|
--target_file=/path/to/target.nemo \ |
|
--tensor_model_parallel_size=??? \ |
|
--target_tensor_model_parallel_size=??? \ |
|
--pipeline_model_parallel_size=-1 \ |
|
--target_pipeline_model_parallel_size=1 \ |
|
--precision=32 \ |
|
--model_class="nemo.collections.vision.models.megatron_vit_classification_models.MegatronVitClassificationModel" \ |
|
--tp_conversion_only |
|
|