Skip to content

Commit 38fcd5f

Browse files
ericharperericharper
committed
Merge branch 'r2.0.0rc0' of github.com:NVIDIA/NeMo into r2.0.0rc0
2 parents fd36bcc + 2a2d985 commit 38fcd5f

File tree

119 files changed

+1008
-556
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

119 files changed

+1008
-556
lines changed

‎.github/workflows/cicd-main.yml‎

Lines changed: 83 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -132,8 +132,8 @@ jobs:
132132
apt-get update && apt-get install libsox-fmt-all -y && \
133133
popd
134134
135-
# AMMO installation
136-
pip install nvidia-ammo~=0.9.0 --extra-index-url https://pypi.nvidia.com --no-cache-dir
135+
# ModelOpt installation
136+
pip install nvidia-modelopt[torch]~=0.11.0 --extra-index-url https://pypi.nvidia.com --no-cache-dir
137137
138138
# PyTorch Lightning version
139139
python -c "import pytorch_lightning; print(pytorch_lightning.__version__)"
@@ -394,7 +394,7 @@ jobs:
394394
- name: Checkout repository
395395
uses: actions/checkout@v4
396396
- run: |
397-
python examples/nlp/language_modeling/megatron_llama_quantization.py \
397+
python examples/nlp/language_modeling/megatron_quantization.py \
398398
model_file=/home/TestData/nlp/megatron_llama/llama_ci.nemo \
399399
quantization.algorithm=null \
400400
model_save=/home/TestData/nlp/megatron_llama/ci_baseline
@@ -403,69 +403,70 @@ jobs:
403403
- uses: "NVIDIA/NeMo/.github/actions/cancel-workflow@main"
404404
if: "failure()"
405405

406-
# L2_PTQ_Llama2_FP8:
407-
# needs: [cicd-test-container-setup]
408-
# runs-on: self-hosted-azure
409-
# timeout-minutes: 10
410-
# container:
411-
# image: nemoci.azurecr.io/nemo_container_${{ github.run_id }}
412-
# options:
413-
# # --user 0:128
414-
# --device=/dev/nvidia0
415-
# --gpus all
416-
# --shm-size=8g
417-
# --env TRANSFORMERS_OFFLINE=0
418-
# --env HYDRA_FULL_ERROR=1
419-
# --volume /mnt/datadrive/TestData:/home/TestData
420-
# steps:
421-
# - name: Checkout repository
422-
# uses: actions/checkout@v4
423-
# - run: |
424-
# python examples/nlp/language_modeling/megatron_llama_quantization.py \
425-
# model_file=/home/TestData/nlp/megatron_llama/llama_ci.nemo \
426-
# tensor_model_parallel_size=2 \
427-
# trainer.devices=2 \
428-
# quantization.calib_dataset=/home/TestData/nlp/test_quantization/test.json \
429-
# quantization.algorithm=fp8 \
430-
# quantization.num_calib_size=8 \
431-
# inference.batch_size=2 \
432-
# export.inference_tensor_parallel=2 \
433-
# model_save=/home/TestData/nlp/megatron_llama/ci_fp8.qnemo
434-
435-
# rm -rf /home/TestData/nlp/megatron_llama/ci_fp8.qnemo
436-
# - uses: "NVIDIA/NeMo/.github/actions/cancel-workflow@main"
437-
# if: "failure()"
438-
439-
# L2_PTQ_Llama2_INT8_SQ:
440-
# needs: [cicd-test-container-setup]
441-
# runs-on: self-hosted-azure
442-
# timeout-minutes: 10
443-
# container:
444-
# image: nemoci.azurecr.io/nemo_container_${{ github.run_id }}
445-
# options:
446-
# # --user 0:128
447-
# --device=/dev/nvidia0
448-
# --gpus all
449-
# --shm-size=8g
450-
# --env TRANSFORMERS_OFFLINE=0
451-
# --env HYDRA_FULL_ERROR=1
452-
# --volume /mnt/datadrive/TestData:/home/TestData
453-
# steps:
454-
# - name: Checkout repository
455-
# uses: actions/checkout@v4
456-
# - run: |
457-
# python examples/nlp/language_modeling/megatron_llama_quantization.py \
458-
# model_file=/home/TestData/nlp/megatron_llama/llama_ci.nemo \
459-
# quantization.calib_dataset=/home/TestData/nlp/test_quantization/test.json \
460-
# quantization.algorithm=int8_sq \
461-
# quantization.num_calib_size=8 \
462-
# inference.batch_size=2 \
463-
# model_save=/home/TestData/nlp/megatron_llama/ci_int8_sq.qnemo
464-
465-
# rm -rf /home/TestData/nlp/megatron_llama/ci_int8_sq.qnemo
466-
# - uses: "NVIDIA/NeMo/.github/actions/cancel-workflow@main"
467-
# if: "failure()"
468-
406+
L2_PTQ_Llama2_FP8:
407+
needs: [cicd-test-container-setup]
408+
runs-on: self-hosted-azure
409+
timeout-minutes: 10
410+
container:
411+
image: nemoci.azurecr.io/nemo_container_${{ github.run_id }}
412+
options:
413+
# --user 0:128
414+
--device=/dev/nvidia0
415+
--gpus all
416+
--shm-size=8g
417+
--env TRANSFORMERS_OFFLINE=0
418+
--env HYDRA_FULL_ERROR=1
419+
--volume /mnt/datadrive/TestData:/home/TestData
420+
steps:
421+
- name: Checkout repository
422+
uses: actions/checkout@v4
423+
- run: |
424+
python examples/nlp/language_modeling/megatron_quantization.py \
425+
model_file=/home/TestData/nlp/megatron_llama/llama_ci.nemo \
426+
tensor_model_parallel_size=2 \
427+
trainer.devices=2 \
428+
quantization.calib_dataset=/home/TestData/nlp/test_quantization/test.json \
429+
quantization.algorithm=fp8 \
430+
quantization.num_calib_size=8 \
431+
inference.batch_size=2 \
432+
export.inference_tensor_parallel=2 \
433+
model_save=/home/TestData/nlp/megatron_llama/ci_fp8.qnemo
434+
435+
rm -rf /home/TestData/nlp/megatron_llama/ci_fp8.qnemo
436+
- uses: "NVIDIA/NeMo/.github/actions/cancel-workflow@main"
437+
if: "failure()"
438+
439+
L2_PTQ_Llama2_INT8_SQ:
440+
needs: [cicd-test-container-setup]
441+
runs-on: self-hosted-azure
442+
timeout-minutes: 10
443+
container:
444+
image: nemoci.azurecr.io/nemo_container_${{ github.run_id }}
445+
options:
446+
# --user 0:128
447+
--device=/dev/nvidia0
448+
--gpus all
449+
--shm-size=8g
450+
--env TRANSFORMERS_OFFLINE=0
451+
--env HYDRA_FULL_ERROR=1
452+
--volume /mnt/datadrive/TestData:/home/TestData
453+
steps:
454+
- name: Checkout repository
455+
uses: actions/checkout@v4
456+
- run: |
457+
python examples/nlp/language_modeling/megatron_quantization.py \
458+
model_file=/home/TestData/nlp/megatron_llama/llama_ci.nemo \
459+
quantization.calib_dataset=/home/TestData/nlp/test_quantization/test.json \
460+
quantization.algorithm=int8_sq \
461+
quantization.num_calib_size=8 \
462+
inference.batch_size=2 \
463+
model_save=/home/TestData/nlp/megatron_llama/ci_int8_sq.qnemo
464+
465+
rm -rf /home/TestData/nlp/megatron_llama/ci_int8_sq.qnemo
466+
- uses: "NVIDIA/NeMo/.github/actions/cancel-workflow@main"
467+
if: "failure()"
468+
469+
# TODO: investigate int4_awq stuck issues and restore the test
469470
#L2_PTQ_Llama2_INT4_AWQ:
470471
# needs: [cicd-test-container-setup]
471472
# runs-on: self-hosted-azure
@@ -484,7 +485,7 @@ jobs:
484485
# - name: Checkout repository
485486
# uses: actions/checkout@v4
486487
# - run: |
487-
# python examples/nlp/language_modeling/megatron_llama_quantization.py \
488+
# python examples/nlp/language_modeling/megatron_quantization.py \
488489
# model_file=/home/TestData/nlp/megatron_llama/llama_ci.nemo \
489490
# tensor_model_parallel_size=1 \
490491
# trainer.devices=1 \
@@ -2268,8 +2269,8 @@ jobs:
22682269
# stage('L2: Model Parallel Size 2 Megatron Text Classification') {
22692270
# when {
22702271
# anyOf{
2271-
# branch 'main'
2272-
# changeRequest target: 'main'
2272+
# branch 'r2.0.0rc0'
2273+
# changeRequest target: 'r2.0.0rc0'
22732274
# }
22742275
# }
22752276
# failFast true
@@ -2297,8 +2298,8 @@ jobs:
22972298
# stage('L2: Model Parallel Size 2 Megatron Autoresume') {
22982299
# when {
22992300
# anyOf{
2300-
# branch 'main'
2301-
# changeRequest target: 'main'
2301+
# branch 'r2.0.0rc0'
2302+
# changeRequest target: 'r2.0.0rc0'
23022303
# }
23032304
# }
23042305
# failFast true
@@ -2328,8 +2329,8 @@ jobs:
23282329
# stage('L2: Model Parallel Size 2 Megatron Evaluation from .nemo') {
23292330
# when {
23302331
# anyOf{
2331-
# branch 'main'
2332-
# changeRequest target: 'main'
2332+
# branch 'r2.0.0rc0'
2333+
# changeRequest target: 'r2.0.0rc0'
23332334
# }
23342335
# }
23352336
# failFast true
@@ -2349,8 +2350,8 @@ jobs:
23492350
# stage('L2: Model Parallel Size 2 Megatron Train from .nemo') {
23502351
# when {
23512352
# anyOf{
2352-
# branch 'main'
2353-
# changeRequest target: 'main'
2353+
# branch 'r2.0.0rc0'
2354+
# changeRequest target: 'r2.0.0rc0'
23542355
# }
23552356
# }
23562357
# failFast true
@@ -3464,8 +3465,8 @@ jobs:
34643465
# stage('L2: NMT Bottleneck Fallback') {
34653466
# when {
34663467
# anyOf {
3467-
# branch 'main'
3468-
# changeRequest target: 'main'
3468+
# branch 'r2.0.0rc0'
3469+
# changeRequest target: 'r2.0.0rc0'
34693470
# }
34703471
# }
34713472
# failFast true
@@ -3510,8 +3511,8 @@ jobs:
35103511
# stage('L2: NMT Bottleneck Architecture') {
35113512
# when {
35123513
# anyOf {
3513-
# branch 'main'
3514-
# changeRequest target: 'main'
3514+
# branch 'r2.0.0rc0'
3515+
# changeRequest target: 'r2.0.0rc0'
35153516
# }
35163517
# }
35173518
# failFast true
@@ -3591,8 +3592,8 @@ jobs:
35913592
# stage('L2: NMT Bottleneck LVM') {
35923593
# when {
35933594
# anyOf {
3594-
# branch 'main'
3595-
# changeRequest target: 'main'
3595+
# branch 'r2.0.0rc0'
3596+
# changeRequest target: 'r2.0.0rc0'
35963597
# }
35973598
# }
35983599
# failFast true
@@ -6484,7 +6485,7 @@ jobs:
64846485
Speech_Checkpoints_tests:
64856486
needs: [cicd-test-container-setup]
64866487
runs-on: self-hosted-azure
6487-
timeout-minutes: 10
6488+
timeout-minutes: 20
64886489
container:
64896490
image: nemoci.azurecr.io/nemo_container_${{ github.run_id }}
64906491
options:

‎Dockerfile‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -133,8 +133,6 @@ RUN for f in $(ls requirements*.txt); do pip3 install --disable-pip-version-chec
133133
RUN pip install flash-attn
134134
# install numba for latest containers
135135
RUN pip install numba>=0.57.1
136-
# install ammo
137-
RUN pip install nvidia-ammo~=0.9.0 --extra-index-url https://pypi.nvidia.com --no-cache-dir
138136

139137
# copy nemo source into a scratch image
140138
FROM scratch as nemo-src

‎docs/source/nlp/quantization.rst‎

Lines changed: 45 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,18 +10,60 @@ PTQ enables deploying a model in a low-precision format -- FP8, INT4, or INT8 --
1010

1111
Model quantization has two primary benefits: reduced model memory requirements and increased inference throughput.
1212

13-
In NeMo, quantization is enabled by the Nvidia AMMO library -- a unified algorithmic model optimization & deployment toolkit.
13+
In NeMo, quantization is enabled by the `NVIDIA TensorRT Model Optimizer (ModelOpt) <https://github.com/NVIDIA/TensorRT-Model-Optimizer>`_ library -- a library to quantize and compress deep learning models for optimized inference on GPUs.
1414

1515
The quantization process consists of the following steps:
1616

1717
1. Loading a model checkpoint using an appropriate parallelism strategy
1818
2. Calibrating the model to obtain appropriate algorithm-specific scaling factors
1919
3. Producing an output directory or .qnemo tarball with model config (json), quantized weights (safetensors) and tokenizer config (yaml).
2020

21-
Loading models requires using an AMMO spec defined in `megatron.core.inference.gpt.model_specs.py <https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/inference/gpt/model_specs.py>`_ module. Typically the calibration step is lightweight and uses a small dataset to obtain appropriate statistics for scaling tensors. The output directory produced (or a .qnemo tarball) is ready to be used to build a serving engine with the Nvidia TensorRT-LLM library. The engine build step is also available in NeMo project in ``nemo.deploy`` and ``nemo.export`` modules.
21+
Loading models requires using an ModelOpt spec defined in `nemo.collections.nlp.models.language_modeling.megatron.gpt_layer_modelopt_spec <https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/models/language_modeling/megatron/gpt_layer_modelopt_spec.py>`_ module. Typically the calibration step is lightweight and uses a small dataset to obtain appropriate statistics for scaling tensors. The output directory produced (or a .qnemo tarball) is ready to be used to build a serving engine with the Nvidia TensorRT-LLM library. The engine build step is also available in NeMo project in ``nemo.deploy`` and ``nemo.export`` modules.
2222

2323
Quantization algorithm can also be conveniently set to ``"null"`` to perform only the weights export step using default precision for TensorRT-LLM deployment. This is useful to obtain baseline performance and accuracy results for comparison.
2424

25+
Support Matrix
26+
^^^^^^^^^^^^^^
27+
28+
Table below presents verified model support matrix for popular LLM architectures. Each model entry also optionally provides a download link to a corresponding Nemo checkpoint for testing purposes. Support for other model families is experimental.
29+
30+
.. list-table:: Model Support Matrix
31+
:widths: 15 15 15 15
32+
:header-rows: 1
33+
34+
* - **Model Family**
35+
- **FP8**
36+
- **INT8_SQ**
37+
- **INT4_AWQ**
38+
* - Llama (1, 2, 3)
39+
- ✅
40+
- ✅
41+
- ✅
42+
* - Mistral
43+
- ✅
44+
- ✅
45+
- ✅
46+
* - `GPT-3 <https://huggingface.co/nvidia/GPT-2B-001>`_
47+
- ✅
48+
- ✅
49+
- ✅
50+
* - `Nemotron-3 8b <https://huggingface.co/nvidia/nemotron-3-8b-base-4k>`_
51+
- ✅
52+
- ✅
53+
- ✅
54+
* - Nemotron-4 15b
55+
- ✅
56+
- ✅
57+
- ✅
58+
* - StarCoder 2
59+
- ✅
60+
- ✅
61+
- ✅
62+
* - Gemma
63+
- ✅
64+
- ✅
65+
- ✅
66+
2567

2668
Example
2769
^^^^^^^
@@ -31,7 +73,7 @@ The script must be launched correctly with the number of processes equal to tens
3173

3274
.. code-block:: bash
3375
34-
torchrun --nproc-per-node 8 examples/nlp/language_modeling/megatron_llama_quantization.py \
76+
torchrun --nproc-per-node 8 examples/nlp/language_modeling/megatron_quantization.py \
3577
model_file=llama2-70b-base-bf16.nemo \
3678
tensor_model_parallel_size=8 \
3779
pipeline_model_parallel_size=1 \

‎docs/source/starthere/intro.rst‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -96,13 +96,13 @@ This section details the steps to clone and install the Megatron Core.
9696
git checkout a5415fcfacef2a37416259bd38b7c4b673583675 && \
9797
pip install .
9898
99-
AMMO Installation
99+
Model Optimizer Installation
100100

101-
This final step involves installing the AMMO package.
101+
This final step involves installing the Model Optimizer package.
102102

103103
.. code-block:: bash
104104
105-
pip install nvidia-ammo~=0.7.0 --extra-index-url https://pypi.nvidia.com --no-cache-dir
105+
pip install nvidia-modelopt[torch]~=0.11.0 --extra-index-url https://pypi.nvidia.com
106106
107107
108108
.. code-block:: bash

‎examples/asr/transcribe_speech.py‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
3030
from nemo.collections.asr.parts.submodules.multitask_decoding import MultiTaskDecoding, MultiTaskDecodingConfig
3131
from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig
32+
from nemo.collections.asr.parts.submodules.rnnt_greedy_decoding import GreedyBatchedRNNTInferConfig
3233
from nemo.collections.asr.parts.utils.eval_utils import cal_write_wer
3334
from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
3435
from nemo.collections.asr.parts.utils.transcribe_utils import (
@@ -121,9 +122,9 @@ class TranscriptionConfig:
121122
pretrained_name: Optional[str] = None # Name of a pretrained model
122123
audio_dir: Optional[str] = None # Path to a directory which contains audio files
123124
dataset_manifest: Optional[str] = None # Path to dataset's JSON manifest
124-
channel_selector: Optional[
125-
Union[int, str]
126-
] = None # Used to select a single channel from multichannel audio, or use average across channels
125+
channel_selector: Optional[Union[int, str]] = (
126+
None # Used to select a single channel from multichannel audio, or use average across channels
127+
)
127128
audio_key: str = 'audio_filepath' # Used to override the default audio key in dataset_manifest
128129
eval_config_yaml: Optional[str] = None # Path to a yaml file of config of evaluation
129130
presort_manifest: bool = True # Significant inference speedup on short-form data due to padding reduction
@@ -161,7 +162,10 @@ class TranscriptionConfig:
161162
ctc_decoding: CTCDecodingConfig = CTCDecodingConfig()
162163

163164
# Decoding strategy for RNNT models
164-
rnnt_decoding: RNNTDecodingConfig = RNNTDecodingConfig(fused_batch_size=-1)
165+
# enable CUDA graphs for transcription
166+
rnnt_decoding: RNNTDecodingConfig = RNNTDecodingConfig(
167+
fused_batch_size=-1, greedy=GreedyBatchedRNNTInferConfig(use_cuda_graph_decoder=True)
168+
)
165169

166170
# Decoding strategy for AED models
167171
multitask_decoding: MultiTaskDecodingConfig = MultiTaskDecodingConfig()
@@ -407,7 +411,10 @@ def autocast(dtype=None):
407411
override_cfg.augmentor = augmentor
408412
override_cfg.text_field = cfg.gt_text_attr_name
409413
override_cfg.lang_field = cfg.gt_lang_attr_name
410-
transcriptions = asr_model.transcribe(audio=filepaths, override_config=override_cfg,)
414+
transcriptions = asr_model.transcribe(
415+
audio=filepaths,
416+
override_config=override_cfg,
417+
)
411418

412419
if cfg.dataset_manifest is not None:
413420
logging.info(f"Finished transcribing from manifest file: {cfg.dataset_manifest}")

0 commit comments

Comments
 (0)