diff --git a/thunder/constants.py b/thunder/constants.py index 5a49f674a0..29d6e8c292 100644 --- a/thunder/constants.py +++ b/thunder/constants.py @@ -5,3 +5,4 @@ # PyTorch version constants # `use_base_version` so that dev builds like "2.13.0a0+gitabc123" compare as "2.13.0" _TORCH_GREATER_EQUAL_2_13 = compare_version("torch", operator.ge, "2.13.0", use_base_version=True) +_TORCH_GREATER_EQUAL_2_14 = compare_version("torch", operator.ge, "2.14.0", use_base_version=True) diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py index 95b38e3d94..18fb6792dc 100644 --- a/thunder/core/jit_ext.py +++ b/thunder/core/jit_ext.py @@ -1009,7 +1009,11 @@ def _generate_random_str_id() -> str: continue trace_of_forward.bound_symbols.append(bsym.from_bsym()) with tracectx(trace_of_forward): - prims.python_return(*(sequencify(output))) + # Keep the fwd graph's output structure. Since torch 2.12 it returns a tuple, which the caller indexes + if isinstance(output, (tuple, list)): + prims.python_return(tuple(output)) + else: + prims.python_return(output) # See NOTE: `autograd_function_apply` and `no_grad` interaction for details about # `thunder.torch.call_higher_order_function_and_consider_outer_autograd_setting` diff --git a/thunder/tests/framework.py b/thunder/tests/framework.py index 7fa3b1485b..5265ddfe1a 100644 --- a/thunder/tests/framework.py +++ b/thunder/tests/framework.py @@ -58,6 +58,12 @@ class NOTHING: ) IS_WINDOWS = platform.system() == "Windows" +# NOTE Running these models in bfloat16 on CPU intermittently kills the process on Windows with an +# illegal instruction, inside eager torch rather than anything of ours. The fp32 variants of the +# same models still run there, and these still run on Linux and macOS. +# TODO: Remove this skip once the crash is fixed upstream, or pinned to a oneDNN/MKL code path we can avoid. +WINDOWS_GPT_CRASH_REASON = "bfloat16 CPU model intermittently crashes the process on Windows (0xc000001d)" + def _bitsandbytes_available(): if not package_available("bitsandbytes"): diff --git a/thunder/tests/opinfos.py b/thunder/tests/opinfos.py index f6997926de..a39cc5fa2c 100644 --- a/thunder/tests/opinfos.py +++ b/thunder/tests/opinfos.py @@ -23,6 +23,7 @@ import thunder.core.prims as prims from thunder.core.pytree import tree_map from thunder.core.symbol import Symbol +from thunder.constants import _TORCH_GREATER_EQUAL_2_14 import thunder.executors as executors from thunder.tests.framework import _all_devicetypes, custom_comparator, IS_WINDOWS from thunder.tests.make_tensor import make_tensor, make_tensor_like @@ -8574,6 +8575,10 @@ def group_norm_sample_generator(op, device, dtype, requires_grad, **kwargs): if torch.device(device).type == "cuda" and ndim >= 3 and num_channels == 0: continue + # torch 2.14 rejects empty channel or spatial dims (an empty batch is still fine) + if _TORCH_GREATER_EQUAL_2_14 and (num_channels == 0 or math.prod(inner_dims) == 0): + continue + a = make(shape) for weight, bias in itertools.product((False, True), repeat=2): diff --git a/thunder/tests/test_interpreter.py b/thunder/tests/test_interpreter.py index c545954069..fe36f26e64 100644 --- a/thunder/tests/test_interpreter.py +++ b/thunder/tests/test_interpreter.py @@ -11,7 +11,7 @@ import pytest import torch from torch.testing import assert_close -from thunder.tests.framework import IS_WINDOWS +from thunder.tests.framework import IS_WINDOWS, WINDOWS_GPT_CRASH_REASON import thunder from thunder.core.interpreter import ( @@ -3413,6 +3413,7 @@ def test_nanogpt_block(jit): assert_close(result, fn(*args, **kwargs)) +@pytest.mark.skipif(IS_WINDOWS, reason=WINDOWS_GPT_CRASH_REASON) def test_nanogpt(jit): from thunder.benchmarks import NanoGPTBenchmark, NanoGPTConfig, _nanogpt_configs @@ -3428,6 +3429,7 @@ def test_nanogpt(jit): assert_close(result, fn(*args, **kwargs)) +@pytest.mark.skipif(IS_WINDOWS, reason=WINDOWS_GPT_CRASH_REASON) def test_litgpt(jit): from thunder.benchmarks import LitGPTBenchmark from thunder.tests.litgpt_model import Config diff --git a/thunder/tests/test_jit_general.py b/thunder/tests/test_jit_general.py index 15a62703ee..22b42e33d4 100644 --- a/thunder/tests/test_jit_general.py +++ b/thunder/tests/test_jit_general.py @@ -14,7 +14,7 @@ import thunder -from thunder.tests.framework import requiresCUDA, IS_WINDOWS +from thunder.tests.framework import requiresCUDA, IS_WINDOWS, WINDOWS_GPT_CRASH_REASON from thunder.core.options import CACHE_OPTIONS import thunder.core.prims as prims from thunder import pytorch_executor, nvfuser_executor @@ -47,6 +47,9 @@ def _autograd_function_apply_kwargs(args_tensor_mask, non_differentiable_idx=Non kwargs = {} if _HAS_ARGS_TENSOR_MASK: kwargs["args_tensor_mask"] = args_tensor_mask + else: + # required since torch 2.11; the tests here save only x + kwargs["saved_for_backward_idx"] = [0] if non_differentiable_idx is not None: kwargs["non_differentiable_idx"] = non_differentiable_idx return kwargs @@ -581,6 +584,7 @@ def h(d, c): assert args_names == ("a", "b", "c", "d") +@pytest.mark.skipif(IS_WINDOWS, reason=WINDOWS_GPT_CRASH_REASON) def test_litgpt(): from thunder.benchmarks import LitGPTBenchmark from thunder.tests.litgpt_model import Config @@ -661,6 +665,7 @@ def test_nanogpt_mlp(): assert_close(result, module(*args, **kwargs)) +@pytest.mark.skipif(IS_WINDOWS, reason=WINDOWS_GPT_CRASH_REASON) def test_nanogpt(): from thunder.benchmarks import NanoGPTBenchmark, NanoGPTConfig, _nanogpt_configs diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py index 74ae2ba269..417043b2b0 100644 --- a/thunder/torch/__init__.py +++ b/thunder/torch/__init__.py @@ -6715,6 +6715,8 @@ def autograd_function_apply( *args: Any, args_tensor_mask: Sequence[bool] | None = None, non_differentiable_idx: Sequence[int] | None = None, + saved_for_backward_idx: Sequence[int] | None = None, # unused, passed by torch >= 2.11 + dirty_idx: Sequence[int] | None = None, # unused, passed by torch >= 2.13 ) -> TensorProxy | tuple[TensorProxy, ...]: # TODO: Remove this once this autograd API becomes stable. # On stable PyTorch, fwd expects ctx as first argument @@ -6733,6 +6735,8 @@ def augmented_forward_autograd_function_apply( *args: Any, args_tensor_mask: Sequence[bool] | None = None, non_differentiable_idx: Sequence[int] | None = None, + saved_for_backward_idx: Sequence[int] | None = None, # unused, passed by torch >= 2.11 + dirty_idx: Sequence[int] | None = None, # unused, passed by torch >= 2.13 ) -> tuple[TensorProxy | tuple[TensorProxy, ...], tuple[Any, ...]]: # TODO: Remove this once this autograd API becomes stable. # On stable PyTorch, fwd expects ctx as first argument