Support nested NeMo models #5671

artbataev · 2022-12-19T14:13:59Z

What does this PR do ?

Adds support for nested NeMo models (with resources).

It's possible to instantiate and save models with NeMo submodules. Artifacts are correctly saved and restored.

Collection: [core]

Changelog

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

Usage

Example usage:

3 ways to instantiate child models:
- use subconfig directly
- use child_model_path with .nemo checkpoint path to load the model
- use pretrained model
child model can contain artifacts
child model config can be changed after parent model instantiation. Config will be saved when saving parent model

from nemo.core.classes import ModelPT

class ChildModel(ModelPT):
    ...  # implement necessary methods

class ParentModel(ModelPT):
    def __init__(self, cfg, trainer=None):
        super().__init__(cfg=cfg, trainer=trainer)

        # optionally annotate type for IDE autocompletion and type checking
        self.child_model: Optional[ChildModel]
        if cfg.get("child_model") is not None:
            # load directly from config
            # either if config provided initially, or automatically
            # after model restoration
            self.register_nemo_submodule(
                "child_model",
                config_field="child_model",
                model=ChildModel(self.cfg.child_model),
            )
        elif cfg.get('child_model_path') is not None:
            # load from .nemo model checkpoint
            # while saving, config will be automatically assigned/updated
            # in cfg.child_model
            self.register_nemo_submodule(
                "child_model",
                config_field="child_model",
                model=ChildModel.restore_from(self.cfg.child_model_path),
            )
        elif cfg.get('child_model_name') is not None:
            # load from pretrained model
            # while saving, config will be automatically assigned/updated
            # in cfg.child_model
            self.register_nemo_submodule(
                "child_model",
                config_field="child_model",
                model=ChildModel.from_pretrained(self.cfg.child_model_name),
            )
        else:
            self.child_model = None

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)

Signed-off-by: Vladimir Bataev <[email protected]>

github-actions · 2023-01-06T01:59:39Z

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

artbataev · 2023-01-06T09:13:35Z

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

The PR is still actual. Need a review / discussion (maybe after 9 Jan)

Signed-off-by: Vladimir Bataev <[email protected]>

artbataev · 2023-01-10T16:15:08Z

@titu1994, I reworked the approach to handle nested submodules artifacts automatically, register_submodule_artifacts is not required now.
This is because nested models can change their config + artifacts, e.g., change_vocabulary for BPE models changes model resources. Since we can't track changes in submodules from the parent module, I check submodules for changes when saving the model.

I added tests for:

serialization + deserialization
multi-nested model
multiple save-restore passes
cases when the child model changes its config

I'm not sure about model-parallel serialization. Do we have any examples to start from?

Also, I'm not sure about the is_model_being_restored flag.
I can add a simple context manager to preserve this flag when loading the child NeMo model. Is it a good solution? Do we have any tests for this flag?

Signed-off-by: Vladimir Bataev <[email protected]>

nemo/core/classes/modelPT.py

titu1994

Minor comments here and there, but overall incredible work !
Would not have thought of this solution, but it works very very well.
@ericharper after your review, we can merge this for Non-MP models. I think the refactor for MP models would be only minor changes which can be done in a subsequent PR (this PR is a blocker for two others)

titu1994 · 2023-01-13T22:14:52Z

nemo/core/connectors/save_restore_connector.py

@@ -37,7 +39,7 @@ def __init__(self) -> None:
        self._model_weights_ckpt = "model_weights.ckpt"
        self._model_extracted_dir = None

-    def save_to(self, model, save_path: str):
+    def save_to(self, model: nemo_classes.ModelPT, save_path: str):


Just use string here, do not use the actual class just for typing

I think we should avoid quotes, and use PEP563 as a more stable approach, since quotes are intended to use with forward declarations by original PEP484, and there is no forward declarations in most cases when quotes are used in NeMo.

But there is a problem with partial initialization here, so I used from __future__ import annotations. This is fully compatible with Python3.7+, widely recommended and adopted.

E.g. see in mypy, pydantic and so on.

You can't see it in old code (due to Python3.7+ limitations), but in new code it is also widely use it, e.g.:

Lightning (app)

TorchVision

To avoid long discussions here, I added quotes.

But I will appreciate if you look at alternative approach and discuss it (maybe separately).

Please use string name here. Pep is not an enforcement it is a suggestion. It does not make sense to mess around internals of Nemo just for an annotation.

nemo/core/classes/modelPT.py

nemo/core/connectors/save_restore_connector.py

titu1994 · 2023-01-13T22:32:51Z

nemo/core/connectors/save_restore_connector.py

-                    self._unpack_nemo_file(path2file=model_metadata.restoration_path, out_folder=archive_dir)
+                    # unpack all restorations paths (nemo checkpoints)
+                    # in nemo checkpoints all resources contain hash in name, so there should be no collisions
+                    for path in restoration_paths:


Restoration paths were tempdirs, do we recreate those tempdirs?

restoration_paths are only paths to .nemo checkpoints, since we don't use .nemo files inside parent .nemo file. Using .nemo in .nemo will still break the code and should be avoided.

I changed it to a set to unpack each checkpoint only once

We need to add a test where we attempt to register a Nemo file. Maybe check and raise an error ?

nemo/core/connectors/save_restore_connector.py

titu1994 · 2023-01-13T22:39:26Z

tests/core/test_save_restore.py

+                # path in config should be set to `None` to restore model from config after saving
+                self.cfg.child1_model_path = None
+                # config for child model should be stored in base model config
+                self.cfg.child1_model = self.child1_model.cfg


Is this for demonstration or do we need to ask model devs to do this? If so, can we automate it somehow.

If the model is loaded from checkpoint (see a comment about 2 ways to construct child models), the developer should assign a config.

I also thought about doing this automatically (since we can use __getattr__ on ModelPT to do this), but:

there are cases when the model will be directly constructed from config, so we don't need to assign it once more

this will break existing code: there is existing NLP code, which uses additional NeMo models as attributes, but doesn't save them (actually I had some errors when tried to do everything recursively). E.g. examples/nlp/token_classification/token_classification_train.py -> NLPModel.bert_model is assigned, but there is some tricky code to handle this submodel

if the model is not needed to be saved (assigned as attribute, but then deleted), the developer will have to delete cfg.child1_model directly after 'magic' construction of this attribute.

So, I think it's better to ask developers to do it manually if constructing child models from .nemo checkpoints.

Let's not modify getattr. Very annoying to do it correctly and brittle for future.

However it is also very easy to forget this step and completely break your implementation.

Add a test that checks this - a model which is nested but "forgets" to set the config. Raise an error.

What about doing this - inside of save_to, check if any module is ModelPT and if so, check the config is available other wise raise error

Or do this - inside save to check if the config is set, if yes don't do anything else otherwise use self.nested_module_i.cfg and set the config automatically inside of self.cfg

titu1994 · 2023-01-13T22:40:21Z

tests/core/test_save_restore.py

+    def __init__(self, cfg, trainer=None):
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        # child 1


Where do you register the model file? Or is that no longer necessary cause internally it will handle the module traversal ?

There are 2 ways to instantiate child model.

I added comments to MockModelWithChildren.

Variant 1 for creating nested NeMo model:

create a child model from .nemo checkpoint (from *_model_path)

assign config to indicate that this model should be saved/restored

when model is restored, parent and children will be restored from a signle parent .nemo checkpoint

Variant 2 for creating nested NeMo model: load directly from config (can be instantiated via MyModel(child_cfg) or universally ModelPT.from_config_dict(child_cfg) (if _target_ is provided).

In both cases no .nemo file inside parent .nemo model will be saved (this is still incompatible with NeMo and global AppState!). So, after loading submodels from .nemo checkpoints, a solid model is constructed.

Keep in mind, that synchronization is only done in save_to:

aggregate and update subconfigs

aggregate and save artifacts

Then add a check inside of register_artifact that checks if it's a .Nemo file and raise an error saying that Nemo file itsel cannot be registered. Add a test case for this

Then add a check inside of register_artifact that checks if it's a .Nemo file and raise an error saying that Nemo file itsel cannot be registered. Add a test case for this

Firstly added, but then reverted, since such a check breaks existing code. See
tests/collections/asr/test_asr_rnnt_encoder_model_bpe.py::TestEncDecRNNTBPEModel::test_save_restore_nested_model

Lets add this check, and either remove or update the faulty RNNT test. That one was my initial attempt to hack together nested model support and is no longer in the correct format of this PR

titu1994 · 2023-01-13T22:53:55Z

tests/core/test_save_restore.py

+            parent = ModelPT.restore_from(parent_path)
+            # check model is transparent, child models can be accessed and can be saved/restored separately
+            _ = self.__test_restore_elsewhere(parent.child1_model, map_location='cpu')
+            child2 = self.__test_restore_elsewhere(parent.child2_model, map_location='cpu')


Amazing test, incredible work here !

ericharper · 2023-01-14T00:43:57Z

tests/core/test_save_restore.py

+        if cfg.get('child1_model_path') is None and cfg.get('child1_model') is None:
+            self.child1_model = None
+        else:
+            # if `child1_model_path` is set, model should be restored from nemo checkpoint


here "nemo checkpoint" -> "nemo model"?

This is a checkpoint, so I changed it to ".nemo model checkpoint"

I also added a comment above to clarify options how child model is constructed.

ericharper · 2023-01-14T00:45:25Z

tests/core/test_save_restore.py

+        if cfg.get('child2_model_path') is None and cfg.get('child2_model') is None:
+            self.child2_model = None
+        else:
+            # if `child2_model_path` is set, model should be restored from nemo checkpoint


"nemo checkpoint" -> "nemo model" ?

This is a checkpoint, so I changed it to ".nemo model checkpoint"

ericharper · 2023-01-14T01:58:29Z

tests/core/test_save_restore.py

+                child2_path = os.path.join(tmpdir_child, 'child2.nemo')
+                child2.save_to(child2_path)
+
+                # create model with children using saved "nemo" checkpoints


"nemo" checkpoints -> nemo models

".nemo model checkpoints"

tests/core/test_save_restore.py

Signed-off-by: Vladimir Bataev <[email protected]>

titu1994 · 2023-01-20T19:54:50Z

Disallowing registering .nemo file as an artifact breaks existing code, so I reverted changes related to this check. See tests/collections/asr/test_asr_rnnt_encoder_model_bpe.py::TestEncDecRNNTBPEModel::test_save_restore_nested_model

Enforce this check, remove/update that test. That test is a proxy of the old way I was trying to hack together multi model support.

titu1994

This is ready to merge (after adding back the nemo file checks).
@ericharper for final review

Amazing work !

docs/source/core/core.rst

titu1994 · 2023-01-20T19:58:05Z

docs/source/core/core.rst

+                self.register_nemo_submodule(
+                    "child_model",
+                    config_field="child_model",
+                    model=ChildModel(self.cfg.child_model),


add the trainer=trainer part

titu1994 · 2023-01-20T19:58:49Z

docs/source/core/core.rst

+                # either if config provided initially, or automatically
+                # after model restoration
+                self.register_nemo_submodule(
+                    "child_model",


Since we are using kwargs, add the key for the first arg

docs/source/core/core.rst

nemo/core/classes/modelPT.py

titu1994 · 2023-01-20T20:02:31Z

tests/core/test_save_restore.py

+    def __init__(self, cfg, trainer=None):
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        # child 1


Lets add this check, and either remove or update the faulty RNNT test. That one was my initial attempt to hack together nested model support and is no longer in the correct format of this PR

Signed-off-by: Vladimir Bataev <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Vladimir Bataev <[email protected]>

…o support_nested_models

artbataev · 2023-01-20T20:29:50Z

Disallowing registering .nemo file as an artifact breaks existing code, so I reverted changes related to this check. See tests/collections/asr/test_asr_rnnt_encoder_model_bpe.py::TestEncDecRNNTBPEModel::test_save_restore_nested_model

Enforce this check, remove/update that test. That test is a proxy of the old way I was trying to hack together multi model support.

I added the check and removed the test

ericharper

LGTM. Thanks for many changes, iterations, and documentation!

Signed-off-by: Vladimir Bataev <[email protected]>

artbataev · 2023-01-23T09:05:48Z

@titu1994, @ericharper
Thank you so much for your comments and discussion!

I finally updated the documentation + fixed the faulty test for NestedRNNTModel in tests/collections/asr/test_asr_rnnt_encoder_model_bpe.py instead of removing it.

I think now it's ready to be merged)

nemo/core/classes/modelPT.py

SeanNaren · 2023-01-23T15:17:38Z

LGTM as well, absolutely bonkers you made this work! I hope this promotes new awesome clean multi-model pipelines in NeMo 🚀

Co-authored-by: Sean Naren <[email protected]> Signed-off-by: Vladimir Bataev <[email protected]>

Signed-off-by: Vladimir Bataev <[email protected]>

ericharper

LGTM. Thanks!

Nested NeMo models support Signed-off-by: Vladimir Bataev <[email protected]> Co-authored-by: Oleksii Kuchaiev <[email protected]> Co-authored-by: Eric Harper <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Sean Naren <[email protected]>

artbataev added 2 commits December 19, 2022 18:04

Nested NeMo models support: save-restore with artifacts

21f9e79

Signed-off-by: Vladimir Bataev <[email protected]>

Fix import

a830aee

Signed-off-by: Vladimir Bataev <[email protected]>

github-actions bot added the core Changes to NeMo Core label Dec 19, 2022

artbataev requested review from titu1994, ericharper and okuchaiev December 19, 2022 14:14

artbataev mentioned this pull request Dec 19, 2022

Hybrid ASR-TTS models #5659

Merged

14 tasks

artbataev and others added 7 commits December 19, 2022 20:31

Improve tests

cd494d1

Signed-off-by: Vladimir Bataev <[email protected]>

Add test for double nested model

a19197c

Signed-off-by: Vladimir Bataev <[email protected]>

Fix target for mock model with children

f8519a4

Signed-off-by: Vladimir Bataev <[email protected]>

MockModel config: revert to original

3295d75

Signed-off-by: Vladimir Bataev <[email protected]>

Add comments. Clean up

b114920

Signed-off-by: Vladimir Bataev <[email protected]>

Merge branch 'main' into support_nested_models

7b661d3

Add docstring. Test number of artifacts

b722181

Signed-off-by: Vladimir Bataev <[email protected]>

github-actions bot added the stale label Jan 6, 2023

artbataev removed the stale label Jan 6, 2023

artbataev added 4 commits January 9, 2023 20:22

Merge branch 'main' into support_nested_models

35ffb35

Handle cases when child model can change config

f7d42f4

Signed-off-by: Vladimir Bataev <[email protected]>

Add test for multiple test-restore passes

b8e0f82

Signed-off-by: Vladimir Bataev <[email protected]>

Merge branch 'main' into support_nested_models

47e5e89

artbataev added 2 commits January 10, 2023 23:05

Do not use artifacts from models not in config

3d5e13d

Signed-off-by: Vladimir Bataev <[email protected]>

Fix for the case when artifacts attribute is not assigned

49d2895

Signed-off-by: Vladimir Bataev <[email protected]>

artbataev commented Jan 11, 2023

View reviewed changes

nemo/core/classes/modelPT.py Show resolved Hide resolved

titu1994 reviewed Jan 13, 2023

View reviewed changes

ericharper reviewed Jan 14, 2023

View reviewed changes

github-advanced-security bot found potential problems Jan 19, 2023

View reviewed changes

tests/core/test_save_restore.py Fixed Show resolved Hide resolved

Fix unused import

96ffd43

Signed-off-by: Vladimir Bataev <[email protected]>

titu1994 reviewed Jan 20, 2023

View reviewed changes

artbataev and others added 4 commits January 21, 2023 00:18

Disallow registering .nemo file as an artifact

75c9622

Signed-off-by: Vladimir Bataev <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

d498a26

for more information, see https://pre-commit.ci

Fix docs

07dc349

Signed-off-by: Vladimir Bataev <[email protected]>

Merge branch 'support_nested_models' of github.com:artbataev/NeMo int…

050a57c

…o support_nested_models

Merge branch 'main' into support_nested_models

0cae024

ericharper previously approved these changes Jan 20, 2023

View reviewed changes

artbataev added 2 commits January 23, 2023 11:54

Merge branch 'main' into support_nested_models

0589c33

Revert test for nested RNNT model, fix according to the new approach

764f76f

Signed-off-by: Vladimir Bataev <[email protected]>

artbataev dismissed ericharper’s stale review via 764f76f January 23, 2023 08:48

artbataev added 2 commits January 23, 2023 12:55

Fix documentation

5bf9b0e

Signed-off-by: Vladimir Bataev <[email protected]>

Fix documentation

d929d3f

Signed-off-by: Vladimir Bataev <[email protected]>

Merge branch 'main' into support_nested_models

0e68a46

SeanNaren reviewed Jan 23, 2023

View reviewed changes

nemo/core/classes/modelPT.py Outdated Show resolved Hide resolved

SeanNaren reviewed Jan 23, 2023

View reviewed changes

nemo/core/classes/modelPT.py Outdated Show resolved Hide resolved

artbataev and others added 2 commits January 23, 2023 20:18

Fix documentation for register_nemo_submodule

80a6706

Co-authored-by: Sean Naren <[email protected]> Signed-off-by: Vladimir Bataev <[email protected]>

Protect inner nemo submodules mapping

eb6ec43

Signed-off-by: Vladimir Bataev <[email protected]>

ericharper approved these changes Jan 23, 2023

View reviewed changes

artbataev merged commit 97973c5 into NVIDIA:main Jan 23, 2023

Support nested NeMo models #5671

Support nested NeMo models #5671

Conversation

artbataev commented Dec 19, 2022 • edited Loading

What does this PR do ?

Changelog

Usage

Before your PR is "Ready for review"

Who can review?

Additional Information

github-actions bot commented Jan 6, 2023

artbataev commented Jan 6, 2023

artbataev commented Jan 10, 2023

titu1994 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

artbataev Jan 16, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

artbataev Jan 16, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

artbataev Jan 16, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

artbataev Jan 16, 2023 • edited Loading

Choose a reason for hiding this comment

ericharper Jan 14, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

titu1994 commented Jan 20, 2023

titu1994 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

artbataev commented Jan 20, 2023

ericharper left a comment

Choose a reason for hiding this comment

artbataev commented Jan 23, 2023

SeanNaren commented Jan 23, 2023

ericharper left a comment

Choose a reason for hiding this comment

artbataev commented Dec 19, 2022 •

edited

Loading

artbataev Jan 16, 2023 •

edited

Loading

artbataev Jan 16, 2023 •

edited

Loading

artbataev Jan 16, 2023 •

edited

Loading

artbataev Jan 16, 2023 •

edited

Loading

ericharper Jan 14, 2023 •

edited

Loading