From e0d149478344c1d2115a30bd901ebd1f4d198a6f Mon Sep 17 00:00:00 2001 From: Thomas Viehmann Date: Mon, 2 Dec 2024 14:20:37 +0100 Subject: [PATCH 1/6] import thunder.tests.litgpt_model.Config in tests (#1503) --- thunder/tests/test_jit_general.py | 4 +++- thunder/tests/test_torch_compile_executor.py | 3 ++- thunder/tests/test_transforms.py | 5 +++-- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/thunder/tests/test_jit_general.py b/thunder/tests/test_jit_general.py index 3c5f4c4d0f..c7e30a0c58 100644 --- a/thunder/tests/test_jit_general.py +++ b/thunder/tests/test_jit_general.py @@ -680,6 +680,8 @@ def test_litgpt_variants(name, device): if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") + if device == "cuda" and name == "falcon-40b-like": + pytest.skip("NVFuser reenable when https://github.com/NVIDIA/Fuser/issues/3505 is fixed, Thunder issue #1504") if device == "cuda" and name == "falcon-7b-like": pytest.skip("NVFuser reenable when https://github.com/NVIDIA/Fuser/issues/3292 is fixed") @@ -783,7 +785,7 @@ def sample(logits): ("cpu", "cuda"), ) def test_tom_overrides_proxy(device): - from litgpt.config import Config + from thunder.tests.litgpt_model import Config from litgpt.model import GPT if device == "cuda" and not torch.cuda.is_available(): diff --git a/thunder/tests/test_torch_compile_executor.py b/thunder/tests/test_torch_compile_executor.py index c0bd1b351d..ea02dc8b5d 100644 --- a/thunder/tests/test_torch_compile_executor.py +++ b/thunder/tests/test_torch_compile_executor.py @@ -20,6 +20,7 @@ def test_supported_ops_are_in_pytorch_executor(): # appropriate visual studio config. @pytest.mark.skipif(not is_inductor_supported() or platform.system() == "Windows", reason="inductor unsupported") def test_torch_compile_litgpt(): + from thunder.tests.litgpt_model import Config from litgpt.model import GPT model = GPT.from_name("llama1-like", n_layer=1) @@ -40,7 +41,7 @@ def test_torch_compile_litgpt(): @requiresCUDA @pytest.mark.skipif(not device_supports_bf16(torch.device("cuda")), reason="bf16 is not supported") def test_torch_compile_cat_nvfuser_phi2_tanh(): - from litgpt.config import Config + from thunder.tests.litgpt_model import Config from litgpt.model import GPT device = torch.device("cuda") diff --git a/thunder/tests/test_transforms.py b/thunder/tests/test_transforms.py index f6b56bc90d..42a0ba0e6f 100644 --- a/thunder/tests/test_transforms.py +++ b/thunder/tests/test_transforms.py @@ -74,7 +74,8 @@ def _test_equal_nvtx_push_and_pop(trc): @requiresCUDA def test_materialization(): from thunder.transforms import MaterializationTransform - from litgpt.config import Config + from thunder.tests.litgpt_model import Config + from litgpt.model import GPT config = Config.from_name("llama2-like") @@ -121,7 +122,7 @@ def test_materialization(): def test_quantization_on_meta(): from thunder.transforms import MaterializationTransform from thunder.transforms.quantization import BitsAndBytesLinearQuant4bit, get_bitsandbytes_executor - from litgpt.config import Config + from thunder.tests.litgpt_model import Config from litgpt.model import GPT bitsandbytes_executor = get_bitsandbytes_executor() From 15c48efc7fa2befd628ffa0e86a3e2d17816b9de Mon Sep 17 00:00:00 2001 From: Yan Wang Date: Mon, 2 Dec 2024 17:56:34 +0100 Subject: [PATCH 2/6] Add output node if it does not exist in the split module (#1480) Fixes #1476. Add a workaround to make ThunderFX work with an older version of PyTorch by going through all submodules of split_module and adding an output node if it's missing. --- thunder/dynamo/splitter.py | 7 +++++++ thunder/tests/test_dynamo.py | 29 +++++++++++++++++++++++++---- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/thunder/dynamo/splitter.py b/thunder/dynamo/splitter.py index 4b455f60b6..587729dde3 100644 --- a/thunder/dynamo/splitter.py +++ b/thunder/dynamo/splitter.py @@ -138,6 +138,13 @@ def callback(node) -> int: original_split_gm: torch.fx.GraphModule = split_module( gm, root_m=None, split_callback=callback, keep_original_order=True, keep_original_node_name=True ) + + # Workaround for the Torch bug https://github.com/pytorch/pytorch/pull/139275 + for submodule in original_split_gm.children(): + if not submodule.graph.find_nodes(op="output"): + submodule.graph.output(()) + if not original_split_gm.graph.find_nodes(op="output"): + original_split_gm.graph.output(()) split_gm = copy.deepcopy(original_split_gm) def is_thunder_supported_partition(node: torch.fx.Node) -> bool: diff --git a/thunder/tests/test_dynamo.py b/thunder/tests/test_dynamo.py index cc740ff408..dd03580991 100644 --- a/thunder/tests/test_dynamo.py +++ b/thunder/tests/test_dynamo.py @@ -449,10 +449,6 @@ def func(x): IS_WINDOWS, reason="torch.compile Windows support is still WIP - https://github.com/pytorch/pytorch/issues/122094", ), - pytest.mark.skipif( - LooseVersion(torch.__version__) < LooseVersion("2.6.0"), - reason="Skip until the Torch bug is fixed - https://github.com/pytorch/pytorch/pull/139275", - ), pytest.mark.skipif( version_between(torch.__version__, min_ver="2.6.0dev0", max_ver="2.6.0a99"), reason="https://github.com/Lightning-AI/lightning-thunder/issues/1471", @@ -864,3 +860,28 @@ def forward(self, x): cmd = "pytest" if use_pytest_benchmark else "python" result1 = run([cmd, s1], capture_output=True, text=True) assert result1.returncode == 0, f"Reproducer {s1} failed with return code {result1.returncode}" + + +def test_deepcopy_graph_module(): + class MyModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + y = x + 1 + + m = MyModule() + gm = torch.fx.symbolic_trace(m) + n = gm.graph.find_nodes(op="output") + gm.graph.erase_node(n[0]) + import thunder + + _, subgraph_info = thunder.dynamo.splitter._splitter(gm, thunder.jit, thunder.jit, []) + original_split_gm = subgraph_info.original_split_graph_module + assert original_split_gm.graph.find_nodes(op="output") + for subm in original_split_gm.children(): + assert subm.graph.find_nodes(op="output") + import copy + + # No assertion error + copy_gm = copy.deepcopy(original_split_gm) From 8eddb9f8b82cad139b6a3ad03f392945af026dc6 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 2 Dec 2024 20:04:11 +0100 Subject: [PATCH 3/6] Bump Lightning-AI/utilities from 0.11.8 to 0.11.9 (#1499) Co-authored-by: Thomas Viehmann --- .github/workflows/ci-checks.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci-checks.yml b/.github/workflows/ci-checks.yml index 3559ec9713..fbc48f8262 100644 --- a/.github/workflows/ci-checks.yml +++ b/.github/workflows/ci-checks.yml @@ -11,17 +11,17 @@ concurrency: jobs: precommit-run: - uses: Lightning-AI/utilities/.github/workflows/check-precommit.yml@v0.11.8 + uses: Lightning-AI/utilities/.github/workflows/check-precommit.yml@v0.11.9 with: python-version: "3.10" check-schema: - uses: Lightning-AI/utilities/.github/workflows/check-schema.yml@v0.11.8 + uses: Lightning-AI/utilities/.github/workflows/check-schema.yml@v0.11.9 with: azure-dir: ".azure" check-package: - uses: Lightning-AI/utilities/.github/workflows/check-package.yml@v0.11.8 + uses: Lightning-AI/utilities/.github/workflows/check-package.yml@v0.11.9 with: actions-ref: v0.11.8 import-name: "thunder" From 20216b8e57faf6bddc75b7fa7d92d61e64852e3f Mon Sep 17 00:00:00 2001 From: jjsjann123 Date: Mon, 2 Dec 2024 11:10:26 -0800 Subject: [PATCH 4/6] support dynamic getitem in dictionary for symbolic values (#1450) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- thunder/core/interpreter.py | 4 +++- thunder/tests/test_jit_general.py | 31 +++++++++++++++++++++++++++---- 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/thunder/core/interpreter.py b/thunder/core/interpreter.py index 992d1e27ec..6557df1802 100644 --- a/thunder/core/interpreter.py +++ b/thunder/core/interpreter.py @@ -2513,7 +2513,9 @@ def __getitem__(self, key): except Exception as e: return do_raise(e) - populate_single_dict_item_wrapper(uv, self, key.value) + from thunder.core.proxies import Proxy + + populate_single_dict_item_wrapper(uv, self, key if isinstance(key.value, Proxy) else key.value) v = self.item_wrappers[key.value] assert uv is v.value or uv is v.original_value, f"value for {key.value} out of sync {uv} {v.value}" return v diff --git a/thunder/tests/test_jit_general.py b/thunder/tests/test_jit_general.py index c7e30a0c58..b9e304e681 100644 --- a/thunder/tests/test_jit_general.py +++ b/thunder/tests/test_jit_general.py @@ -1029,7 +1029,6 @@ def forward(self, x): ids=("remove_duplicate=False", "remove_duplicate=True"), ) def test_named_params_and_named_buffers(prefix, recurse, remove_duplicate): - buffer_tensor = torch.tensor([1.0]) class SubMod(torch.nn.Module): @@ -1143,7 +1142,6 @@ def test_custom_autograd_function(): from torch.testing._internal.common_utils import gradcheck class MyFunction(torch.autograd.Function): - @staticmethod def forward(ctx, x: torch.Tensor) -> torch.Tensor: return x * 2.0 @@ -1206,7 +1204,6 @@ def forward(self, x): def test_autograd_function_apply(): - def forward(ctx, x): saved_for_backward = (x,) return x.sin(), saved_for_backward @@ -1275,7 +1272,6 @@ def my_sin_with_wrong_backward(x): def test_autograd_function_empty_forward(): - class Fn(torch.autograd.Function): @staticmethod def forward(self, x): @@ -1464,3 +1460,30 @@ def foo(a): expected = foo(a) assert_close(actual, expected) + + +def test_cache_symbolic_values_dict(): + def foo(a, v): + return a[v].relu() + + jfoo = thunder.jit(foo, cache="symbolic values") + + a = { + 2: torch.randn(2, 3, 8, requires_grad=True, device="cpu"), + 5: torch.randn(4, 8, requires_grad=True, device="cpu"), + } + + actual = jfoo(a, 2) + expected = foo(a, 2) + + assert_close(actual, expected) + + b = { + "a": torch.randn(2, 8, requires_grad=True, device="cpu"), + "b": torch.randn(7, requires_grad=True, device="cpu"), + } + + actual = jfoo(b, "b") + expected = foo(b, "b") + + assert_close(actual, expected) From 7da6bd3669581272047cf74b3feb170872343f06 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 2 Dec 2024 20:15:39 +0100 Subject: [PATCH 5/6] Bump pypa/gh-action-pypi-publish from 1.11.0 to 1.12.2 (#1497) Co-authored-by: Thomas Viehmann --- .github/workflows/release-nightly.yml | 4 ++-- .github/workflows/release-pypi.yml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/release-nightly.yml b/.github/workflows/release-nightly.yml index 8df4e8ad4d..241e047bb8 100644 --- a/.github/workflows/release-nightly.yml +++ b/.github/workflows/release-nightly.yml @@ -33,7 +33,7 @@ jobs: # We do this, since failures on test.pypi aren't that bad - name: Publish to Test PyPI if: github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' - uses: pypa/gh-action-pypi-publish@v1.11.0 + uses: pypa/gh-action-pypi-publish@v1.12.2 with: user: __token__ password: ${{ secrets.test_pypi_password }} @@ -41,7 +41,7 @@ jobs: - name: Publish distribution 📦 to PyPI if: github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' - uses: pypa/gh-action-pypi-publish@v1.11.0 + uses: pypa/gh-action-pypi-publish@v1.12.2 with: user: __token__ password: ${{ secrets.pypi_password }} diff --git a/.github/workflows/release-pypi.yml b/.github/workflows/release-pypi.yml index 68ad9c5455..b4f0353994 100644 --- a/.github/workflows/release-pypi.yml +++ b/.github/workflows/release-pypi.yml @@ -26,7 +26,7 @@ jobs: # We do this, since failures on test.pypi aren't that bad - name: Publish to Test PyPI if: startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release' - uses: pypa/gh-action-pypi-publish@v1.11.0 + uses: pypa/gh-action-pypi-publish@v1.12.2 with: user: __token__ password: ${{ secrets.test_pypi_password }} @@ -34,7 +34,7 @@ jobs: - name: Publish distribution 📦 to PyPI if: startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release' - uses: pypa/gh-action-pypi-publish@v1.11.0 + uses: pypa/gh-action-pypi-publish@v1.12.2 with: user: __token__ password: ${{ secrets.pypi_password }} From 995de343d63cfe1e326b14bdffd8374646a922fd Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 2 Dec 2024 20:17:20 +0100 Subject: [PATCH 6/6] Bump transformers from 4.46.2 to 4.46.3 (#1495) Co-authored-by: Thomas Viehmann --- requirements/test.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/test.txt b/requirements/test.txt index 9ea8bc8ab6..05efd615e3 100644 --- a/requirements/test.txt +++ b/requirements/test.txt @@ -18,7 +18,7 @@ pandas # thunder/benchmarks/test_benchmark_litgpt.py xlsxwriter # thunder/benchmarks/test_benchmark_litgpt.py jsonargparse # thunder/benchmarks/benchmark_litgpt.py bitsandbytes==0.42.0 # fixed version! -transformers==4.46.2 # for test_networks.py +transformers==4.46.3 # for test_networks.py # Installs JAX on Linux and MacOS jaxlib; sys_platform == 'linux' or sys_platform == 'darwin' # required for jax, see https://github.com/google/jax#installation