[Codec] Finite scalar quantizer #7886

anteju · 2023-11-14T02:00:56Z

What does this PR do ?

This PR adds finite scalar quantizer as used in
Mentzer et al., Finite Scalar Quantization: VQ-VAE Made Simple (https://arxiv.org/abs/2309.15505v1)

Collection: TTS

Changelog

Added FiniteScalarQuantizer
Added unit tests

Usage

This can be used as a drop-in replacement for ResidualVectorQuantizer, for example using the following configuration

model:
  ...
  commit_loss_scale: 0.0 # need to be set to zero, since FSQ does not include commit loss
  ...
  vector_quantizer:
    _target_: nemo.collections.tts.modules.audio_codec_modules.FiniteScalarQuantizer
    num_levels: [8, 5, 5, 5]

For using larger vector, GroupFiniteScalarQuantizer can be used to repeat quantization pattern several times, for example using the following configuration

model:
  ...
  commit_loss_scale: 0.0 # need to be set to zero, since FSQ does not include commit loss
  ...
  vector_quantizer:
    _target_: nemo.collections.tts.modules.audio_codec_modules.GroupFiniteScalarQuantizer
    num_groups: 3
    num_levels_per_group: [3, 4]

This corresponds to FSQ with num_levels: [3, 4, 3, 4, 3, 4].

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)

anteju · 2023-11-14T02:01:14Z

jenkins

tests/collections/tts/modules/test_audio_codec_modules.py

rlangman · 2023-11-14T23:39:22Z

nemo/collections/tts/models/audio_codec.py

+            if len(vq_output_types) == 3 and vq_output_types[-1] == 'commit_loss':
+                self.vector_quantizer_has_commit_loss = True
+                logging.info('Vector quantizer supports commit loss.')
+            else:
+                self.vector_quantizer_has_commit_loss = False
+                logging.info('Vector quantizer does not support commit loss.')


Would it be simpler if we require the VQ to return commit loss and modify FSQ to return 0 commit loss?

We could do that as well, we could skip ~7 lines above.
A disadvantage may be that we could still set commit_loss_scale, even if the quantizer is always returning 0.0, which may be confusing.

I implemented something like that earlier:
an abstract base class VectorQuantizer with methods forward (including commit_loss), encode, decode and the same typecheck types.

If it is not much extra work, I think having the abstract base class would be better. It does not matter much now with only 2 implementations, but should help long term as we inevitably add more types of quantizers.

Let's merge this PR and we can factorize out to an abstract class in a follow-up.
Sounds good?

Sure, sounds good.

tests/collections/tts/modules/test_audio_codec_modules.py

Signed-off-by: Ante Jukić <ajukic@nvidia.com>

anteju · 2023-11-16T19:35:05Z

jenkins

rlangman

LGTM

rlangman · 2023-11-16T21:33:22Z

nemo/collections/tts/models/audio_codec.py

+            if len(vq_output_types) == 3 and vq_output_types[-1] == 'commit_loss':
+                self.vector_quantizer_has_commit_loss = True
+                logging.info('Vector quantizer supports commit loss.')
+            else:
+                self.vector_quantizer_has_commit_loss = False
+                logging.info('Vector quantizer does not support commit loss.')


If it is not much extra work, I think having the abstract base class would be better. It does not matter much now with only 2 implementations, but should help long term as we inevitably add more types of quantizers.

nithinraok

LGTM

Signed-off-by: Chen Cui <chcui@nvidia.com> support packed dataset Signed-off-by: Chen Cui <chcui@nvidia.com> [Codec] Finite scalar quantizer (NVIDIA#7886) * Finite scalar quantizer Signed-off-by: Ante Jukić <ajukic@nvidia.com> * Updated test Signed-off-by: Ante Jukić <ajukic@nvidia.com> --------- Signed-off-by: Ante Jukić <ajukic@nvidia.com> upgrade to latest mcore and TE (NVIDIA#7908) * reimport module Signed-off-by: dimapihtar <dpihtar@gmail.com> * update mcore and TE commits Signed-off-by: dimapihtar <dpihtar@gmail.com> --------- Signed-off-by: dimapihtar <dpihtar@gmail.com> Tar codec (NVIDIA#7867) added missing torch import (NVIDIA#7913) Signed-off-by: David Mosallanezhad <dmosallanezh@nvidia.com> add cpu init check (NVIDIA#7889) Signed-off-by: Chen Cui <chcui@nvidia.com> Fix pinned triton version (NVIDIA#7925) * Fix pinned triton version Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> * Remove comment Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Change README Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> * Remove flash-attn in Dockerfile Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> * Revert Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> --------- Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> fix tp_overlap config var name (NVIDIA#7928) Signed-off-by: Xiaowei Ren <xren@nvidia.com> add Dutch P&C FC model info (NVIDIA#7892) * add Dutch P&C FC model info Signed-off-by: zhehuaichen <dian.chenzhehuai@gmail.com> * update order of the results Signed-off-by: zhehuaichen <dian.chenzhehuai@gmail.com> --------- Signed-off-by: zhehuaichen <dian.chenzhehuai@gmail.com> fix issues with convert_nemo_llama_to_hf.py (NVIDIA#7922) [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix collate_fn bug for TP > 1 Signed-off-by: Chen Cui <chcui@nvidia.com> [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci make packed dataset work Signed-off-by: Chen Cui <chcui@nvidia.com> fix nan bug Signed-off-by: Chen Cui <chcui@nvidia.com> [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci support answer only loss Signed-off-by: Chen Cui <chcui@nvidia.com> [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci account for padding in cu_seqlens during dataloading for attn kernel Signed-off-by: Chen Cui <chcui@nvidia.com> fix path for answer_only_loss = false Signed-off-by: Chen Cui <chcui@nvidia.com> Modify GPTSFTPackedDataset to respond to pad_to_max_length setting Signed-off-by: Valerie Sarge <vsarge@nvidia.com>

Signed-off-by: Chen Cui <chcui@nvidia.com> support packed dataset Signed-off-by: Chen Cui <chcui@nvidia.com> [Codec] Finite scalar quantizer (NVIDIA#7886) * Finite scalar quantizer Signed-off-by: Ante Jukić <ajukic@nvidia.com> * Updated test Signed-off-by: Ante Jukić <ajukic@nvidia.com> --------- Signed-off-by: Ante Jukić <ajukic@nvidia.com> upgrade to latest mcore and TE (NVIDIA#7908) * reimport module Signed-off-by: dimapihtar <dpihtar@gmail.com> * update mcore and TE commits Signed-off-by: dimapihtar <dpihtar@gmail.com> --------- Signed-off-by: dimapihtar <dpihtar@gmail.com> Tar codec (NVIDIA#7867) added missing torch import (NVIDIA#7913) Signed-off-by: David Mosallanezhad <dmosallanezh@nvidia.com> add cpu init check (NVIDIA#7889) Signed-off-by: Chen Cui <chcui@nvidia.com> Fix pinned triton version (NVIDIA#7925) * Fix pinned triton version Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> * Remove comment Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Change README Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> * Remove flash-attn in Dockerfile Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> * Revert Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> --------- Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> fix tp_overlap config var name (NVIDIA#7928) Signed-off-by: Xiaowei Ren <xren@nvidia.com> add Dutch P&C FC model info (NVIDIA#7892) * add Dutch P&C FC model info Signed-off-by: zhehuaichen <dian.chenzhehuai@gmail.com> * update order of the results Signed-off-by: zhehuaichen <dian.chenzhehuai@gmail.com> --------- Signed-off-by: zhehuaichen <dian.chenzhehuai@gmail.com> fix issues with convert_nemo_llama_to_hf.py (NVIDIA#7922) [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci fix collate_fn bug for TP > 1 Signed-off-by: Chen Cui <chcui@nvidia.com> [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci make packed dataset work Signed-off-by: Chen Cui <chcui@nvidia.com> fix nan bug Signed-off-by: Chen Cui <chcui@nvidia.com> [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci support answer only loss Signed-off-by: Chen Cui <chcui@nvidia.com> [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci account for padding in cu_seqlens during dataloading for attn kernel Signed-off-by: Chen Cui <chcui@nvidia.com> fix path for answer_only_loss = false Signed-off-by: Chen Cui <chcui@nvidia.com>

* Finite scalar quantizer Signed-off-by: Ante Jukić <ajukic@nvidia.com> * Updated test Signed-off-by: Ante Jukić <ajukic@nvidia.com> --------- Signed-off-by: Ante Jukić <ajukic@nvidia.com> Signed-off-by: Piotr Żelasko <petezor@gmail.com>

* Finite scalar quantizer Signed-off-by: Ante Jukić <ajukic@nvidia.com> * Updated test Signed-off-by: Ante Jukić <ajukic@nvidia.com> --------- Signed-off-by: Ante Jukić <ajukic@nvidia.com>

anteju marked this pull request as ready for review November 14, 2023 02:01

github-actions bot added the TTS label Nov 14, 2023

anteju requested review from rlangman, nithinraok and KunalDhawan November 14, 2023 02:01

github-advanced-security bot found potential problems Nov 14, 2023

View reviewed changes

tests/collections/tts/modules/test_audio_codec_modules.py Fixed Show fixed Hide fixed

anteju force-pushed the pr/fsq branch 2 times, most recently from 003735a to 78df06d Compare November 15, 2023 00:12

rlangman reviewed Nov 15, 2023

View reviewed changes

anteju force-pushed the pr/fsq branch from 78df06d to b483fb4 Compare November 15, 2023 21:59

anteju added 2 commits November 16, 2023 10:25

Finite scalar quantizer

955e58a

Signed-off-by: Ante Jukić <ajukic@nvidia.com>

Updated test

15404f1

Signed-off-by: Ante Jukić <ajukic@nvidia.com>

anteju force-pushed the pr/fsq branch from 6d56bb9 to 15404f1 Compare November 16, 2023 18:25

anteju requested a review from rlangman November 16, 2023 18:25

rlangman approved these changes Nov 16, 2023

View reviewed changes

nithinraok approved these changes Nov 17, 2023

View reviewed changes

anteju merged commit 76e5bdf into NVIDIA:main Nov 17, 2023
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Codec] Finite scalar quantizer #7886

[Codec] Finite scalar quantizer #7886

anteju commented Nov 14, 2023 •

edited

Loading

anteju commented Nov 14, 2023

rlangman Nov 14, 2023 •

edited

Loading

anteju Nov 16, 2023

rlangman Nov 16, 2023

anteju Nov 17, 2023

rlangman Nov 17, 2023

anteju commented Nov 16, 2023

rlangman left a comment

rlangman Nov 16, 2023

nithinraok left a comment

[Codec] Finite scalar quantizer #7886

[Codec] Finite scalar quantizer #7886

Conversation

anteju commented Nov 14, 2023 • edited Loading

What does this PR do ?

Changelog

Usage

Before your PR is "Ready for review"

Who can review?

Additional Information

anteju commented Nov 14, 2023

rlangman Nov 14, 2023 • edited Loading

Choose a reason for hiding this comment

anteju Nov 16, 2023

Choose a reason for hiding this comment

rlangman Nov 16, 2023

Choose a reason for hiding this comment

anteju Nov 17, 2023

Choose a reason for hiding this comment

rlangman Nov 17, 2023

Choose a reason for hiding this comment

anteju commented Nov 16, 2023

rlangman left a comment

Choose a reason for hiding this comment

rlangman Nov 16, 2023

Choose a reason for hiding this comment

nithinraok left a comment

Choose a reason for hiding this comment

anteju commented Nov 14, 2023 •

edited

Loading

rlangman Nov 14, 2023 •

edited

Loading