From bbedfdf761adc1d45ea3a982df819b0ad921b4c7 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Mon, 25 Apr 2022 02:55:11 +0900 Subject: [PATCH 01/10] Add test --- tests/utilities/test_meta.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/tests/utilities/test_meta.py b/tests/utilities/test_meta.py index c55f04591f139..8c81fbe328090 100644 --- a/tests/utilities/test_meta.py +++ b/tests/utilities/test_meta.py @@ -13,9 +13,11 @@ # limitations under the License. from torch import nn +from pytorch_lightning.trainer.trainer import Trainer from pytorch_lightning.core.lightning import LightningModule from pytorch_lightning.utilities.meta import init_meta_context, is_on_meta_device, materialize_module from tests.helpers.runif import RunIf +from tests.helpers.boring_model import BoringModel class MLP(nn.Module): @@ -24,7 +26,7 @@ def __init__(self, num_layers: int): self.layer = nn.Sequential(*[nn.Linear(1, 1) for _ in range(num_layers)] + [nn.Dropout(), nn.LayerNorm(1)]) -class BoringModel(LightningModule): +class SimpleBoringModel(LightningModule): def __init__(self, num_layers: int): super().__init__() self.save_hyperparameters() @@ -48,7 +50,7 @@ def test_init_meta_context(): assert not is_on_meta_device(mlp) assert not is_on_meta_device(nn.Module()) - model = BoringModel(4) + model = SimpleBoringModel(4) assert model.layer[0].weight.device.type == "meta" materialize_module(model) assert model.layer[0].weight.device.type == "cpu" @@ -68,3 +70,17 @@ def test_init_meta_context(): m = nn.Linear(in_features=1, out_features=1) assert m.weight.device.type == "cpu" + + +@RunIf(min_torch="1.10.0", standalone=True) +def test_fit_with_init_meta_context(): + """Test fitting a model initialized on meta device.""" + with init_meta_context(): + model = BoringModel() + trainer = Trainer( + fast_dev_run=True, + enable_progress_bar=False, + enable_model_summary=False, + enable_checkpointing=False, + ) + trainer.fit(model) From 4b8d64214cc6cea621b8f890d1c0ee3267a68768 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Mon, 25 Apr 2022 02:55:57 +0900 Subject: [PATCH 02/10] Update etst --- tests/utilities/test_meta.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/utilities/test_meta.py b/tests/utilities/test_meta.py index 8c81fbe328090..c1da1fb665d42 100644 --- a/tests/utilities/test_meta.py +++ b/tests/utilities/test_meta.py @@ -72,7 +72,7 @@ def test_init_meta_context(): assert m.weight.device.type == "cpu" -@RunIf(min_torch="1.10.0", standalone=True) +@RunIf(min_torch="1.10.0") def test_fit_with_init_meta_context(): """Test fitting a model initialized on meta device.""" with init_meta_context(): From 5b1f60192ff01a99d999e2c9d780fdce38064efd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 24 Apr 2022 17:58:23 +0000 Subject: [PATCH 03/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/utilities/test_meta.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/utilities/test_meta.py b/tests/utilities/test_meta.py index c1da1fb665d42..83b134b94c6ab 100644 --- a/tests/utilities/test_meta.py +++ b/tests/utilities/test_meta.py @@ -13,11 +13,11 @@ # limitations under the License. from torch import nn -from pytorch_lightning.trainer.trainer import Trainer from pytorch_lightning.core.lightning import LightningModule +from pytorch_lightning.trainer.trainer import Trainer from pytorch_lightning.utilities.meta import init_meta_context, is_on_meta_device, materialize_module -from tests.helpers.runif import RunIf from tests.helpers.boring_model import BoringModel +from tests.helpers.runif import RunIf class MLP(nn.Module): From ec6a02c18b9d7bf2c0853e458041faa1a47b1859 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Mon, 25 Apr 2022 03:12:05 +0900 Subject: [PATCH 04/10] Revert "Update etst" This reverts commit 4b8d64214cc6cea621b8f890d1c0ee3267a68768. --- tests/utilities/test_meta.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/utilities/test_meta.py b/tests/utilities/test_meta.py index c1da1fb665d42..8c81fbe328090 100644 --- a/tests/utilities/test_meta.py +++ b/tests/utilities/test_meta.py @@ -72,7 +72,7 @@ def test_init_meta_context(): assert m.weight.device.type == "cpu" -@RunIf(min_torch="1.10.0") +@RunIf(min_torch="1.10.0", standalone=True) def test_fit_with_init_meta_context(): """Test fitting a model initialized on meta device.""" with init_meta_context(): From f65d836de34cb8d4bf52901e012ecebd43e409a6 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Tue, 26 Apr 2022 11:18:10 +0900 Subject: [PATCH 05/10] Temporarily add init_meta_context testing in CI --- .github/workflows/ci_test-full.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/ci_test-full.yml b/.github/workflows/ci_test-full.yml index 9e2e1554ca21d..c4f9300d846eb 100644 --- a/.github/workflows/ci_test-full.yml +++ b/.github/workflows/ci_test-full.yml @@ -136,6 +136,11 @@ jobs: run: | python requirements/check-avail-extras.py + # FIXME + - name: Test init_meta_context + run: | + PL_RUN_STANDALONE_TESTS=1 python -m pytest tests/utilities/test_meta.py::test_fit_with_init_meta_context -vv + - name: Tests run: | # NOTE: do not include coverage report here, see: https://github.com/nedbat/coveragepy/issues/1003 From ca6fe0480d2fb445d4145ceb5c825dd4452b7927 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Wed, 11 May 2022 19:16:10 +0900 Subject: [PATCH 06/10] Don't set materialized child to child's child --- pytorch_lightning/utilities/meta.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/utilities/meta.py b/pytorch_lightning/utilities/meta.py index d14f111e8759a..a5edcfb300188 100644 --- a/pytorch_lightning/utilities/meta.py +++ b/pytorch_lightning/utilities/meta.py @@ -186,7 +186,7 @@ def materialize_module(root_module: nn.Module) -> nn.Module: if not materialize_fn or isinstance(child, (Sequential, ModuleList, ModuleDict)): materialize_module(child) else: - setattr(child, name, materialize_fn()) + setattr(root_module, name, materialize_fn()) return root_module From 238aa26eae1a5c649d9c1af3093eecdd9fca7749 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Thu, 12 May 2022 14:23:03 +0900 Subject: [PATCH 07/10] Revert "Temporarily add init_meta_context testing in CI" This reverts commit f65d836de34cb8d4bf52901e012ecebd43e409a6. --- .github/workflows/ci_test-full.yml | 5 ----- 1 file changed, 5 deletions(-) diff --git a/.github/workflows/ci_test-full.yml b/.github/workflows/ci_test-full.yml index 6c237bc2e02d7..84ca60fd30f2e 100644 --- a/.github/workflows/ci_test-full.yml +++ b/.github/workflows/ci_test-full.yml @@ -132,11 +132,6 @@ jobs: run: | python requirements/check-avail-extras.py - # FIXME - - name: Test init_meta_context - run: | - PL_RUN_STANDALONE_TESTS=1 python -m pytest tests/utilities/test_meta.py::test_fit_with_init_meta_context -vv - - name: Tests run: | # NOTE: do not include coverage report here, see: https://github.com/nedbat/coveragepy/issues/1003 From ceb0abeb70a42940517864c1b26b1bf8e7cbaa84 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Thu, 12 May 2022 14:34:28 +0900 Subject: [PATCH 08/10] Update CHANGELOG --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c803239d5fccb..70f3398ee5732 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -188,7 +188,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Fixed an issue with unsupported torch.inference_mode() on hpu backends by making it use no_grad ([#13014](https://github.com/PyTorchLightning/pytorch-lightning/pull/13014)) -- +- Fixed `materialize_module` setting a module's child recursively ([#12870](https://github.com/PyTorchLightning/pytorch-lightning/pull/12870)) - From 0771babaa028c65898b0bb21707bbfcda0172a49 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Fri, 13 May 2022 17:32:24 +0900 Subject: [PATCH 09/10] Update test --- tests/utilities/test_meta.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/utilities/test_meta.py b/tests/utilities/test_meta.py index 738ede05f31de..fb1b27b7bc731 100644 --- a/tests/utilities/test_meta.py +++ b/tests/utilities/test_meta.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import pytest from torch import nn from pytorch_lightning.core.lightning import LightningModule @@ -73,14 +74,12 @@ def test_init_meta_context(): @RunIf(min_torch="1.10.0", standalone=True) -def test_fit_with_init_meta_context(): - """Test fitting a model initialized on meta device.""" +def test_materialize_module_recursive_child(): + """Test materialize_module doesn't set a child recursively to a model instantiated within init_meta_context.""" with init_meta_context(): model = BoringModel() - trainer = Trainer( - fast_dev_run=True, - enable_progress_bar=False, - enable_model_summary=False, - enable_checkpointing=False, - ) - trainer.fit(model) + + materialize_module(model) + + with pytest.raises(AttributeError, match="'Linear' object has no attribute 'layer'"): + model.layer.layer From 3c80761b714fd50c323e1fd7a38cdd7195f057c3 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Fri, 13 May 2022 20:18:25 +0900 Subject: [PATCH 10/10] Remove unused import --- tests/utilities/test_meta.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/utilities/test_meta.py b/tests/utilities/test_meta.py index fb1b27b7bc731..34a2c3d04ebf6 100644 --- a/tests/utilities/test_meta.py +++ b/tests/utilities/test_meta.py @@ -15,7 +15,6 @@ from torch import nn from pytorch_lightning.core.lightning import LightningModule -from pytorch_lightning.trainer.trainer import Trainer from pytorch_lightning.utilities.meta import init_meta_context, is_on_meta_device, materialize_module from tests.helpers.boring_model import BoringModel from tests.helpers.runif import RunIf