test(distributed): compare DTensor basic ops against eager, not torch.compile [14/n]

.lightning/workflows/transformer-engine.yaml

@@ -25,6 +25,8 @@ run: |
   pip install . -U -q -r requirements/test.txt
   # Need to explicitly point to cudnn.h as it is installed at a non-standard location
   # Ref: https://github.com/NVIDIA/TransformerEngine/issues/918#issuecomment-2187703769
-  CPLUS_INCLUDE_PATH="/usr/local/lib/python3.12/dist-packages/nvidia/cudnn/include/" pip install --no-build-isolation 'transformer_engine[pytorch]'
+  # TE builds from source here, and from 2.17 on that needs a torch header this image's torch 2.8 lacks
+  # TODO: drop the <=2.16 pin once this image is on torch >= 2.11, which ships that header
+  CPLUS_INCLUDE_PATH="/usr/local/lib/python3.12/dist-packages/nvidia/cudnn/include/" pip install --no-build-isolation 'transformer_engine[pytorch]<=2.16'
   pip list # for debugging purposes
   pytest thunder/tests/${test_file} -v -rs

MANIFEST.in

@@ -8,7 +8,6 @@ recursive-include pl_sandbox *.md
 include LICENSE
 
 # Exclude build configs
-exclude *.toml
 exclude *.svg
 exclude *.yml
 exclude *.yaml

pyproject.toml

@@ -74,6 +74,13 @@ markers = [
 ]
 filterwarnings = [
     "error::FutureWarning",
+    # todo: drop these once upstream stops raising them
+    # torch/testing/_internal/common_utils.py reads this config at import time
+    "ignore:torch._dynamo.config.inline_inbuilt_nn_modules is deprecated:FutureWarning",
+    # torch.func.jvp scripts its decompositions on import; the prefix also covers the 3.14+ wording
+    "ignore:`torch.jit.script` is:FutureWarning",
+    # called by bitsandbytes
+    "ignore:_check_is_size will be removed:FutureWarning",
 ]
 timeout = 900
 # xfail_strict = true  # todo

requirements/devel.txt

@@ -6,4 +6,3 @@ datasets>=3.5.0
 peft>=0.15.2
 
 torchvision
-torchaudio

requirements/notebooks.txt

@@ -1,7 +1,8 @@
 # just reuse the base requirements
 -r base.txt
 
-ipython[all] ~=8.37.0
+# all but the test extras, which pull in jupyter-ai and upgrade apt-owned packages in the GPU image
+ipython[black,doc,kernel,matplotlib,nbconvert,nbformat,notebook,parallel,qtconsole] ~=8.37.0
 numpy
 liger-kernel == 0.4.0
 cuda-python >=12.6.1, <14.0.0

thunder/constants.py

@@ -0,0 +1,8 @@
+import operator
+
+from lightning_utilities.core.imports import compare_version
+
+# 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)

thunder/core/interpreter.py

@@ -33,6 +33,7 @@
     NoneType,
     BuiltinFunctionType,
     BuiltinMethodType,
+    ClassMethodDescriptorType,
     MethodDescriptorType,
     MethodWrapperType,
     WrapperDescriptorType,
@@ -7027,6 +7028,15 @@ def _impl(fn, *args, **kwargs):
         assert is_opaque(fn)
         slf = fn.__self__
 
+        # NOTE Builtin classmethods (e.g. torch.autograd.Function.apply since torch 2.13) dispatch
+        #   through their descriptor, the one object all subclasses share, so lookasides can match it.
+        if isinstance(slf, type) and not is_pycapsule(slf):
+            for klass in slf.__mro__:
+                descriptor = klass.__dict__.get(fn.__name__)
+                if isinstance(descriptor, ClassMethodDescriptorType):
+                    wrapped_slf = _interpret_call(getattr, wrapped_fn, wrap_const("__self__"))
+                    return _interpret_call(wrap_const(descriptor), wrapped_slf, *args, **kwargs)
+
         if slf is not None and not isinstance(slf, (type, ModuleType)) and not is_pycapsule(slf):
             # NOTE: we need to walk the mro because we need to deal with super().foo
             #       using the qualname is not good here because Python qualnames come with no

thunder/core/jit_ext.py

@@ -794,7 +794,16 @@ def _convert_pytorchfunc_to_thundertrace(
     return trace, sequencify(wrapped_func_result)[0].provenance
 
 
-@register_general_jit_lookaside(torch.autograd.function.Function.apply.__func__)
+# The key must be the object every subclass shares: __func__ for a Python classmethod,
+# the descriptor since torch 2.13 reimplemented Function.apply in C.
+_torch_autograd_function_apply = getattr(
+    torch.autograd.function.Function.apply,
+    "__func__",
+    torch.autograd.function.Function.__dict__.get("apply"),
+)
+
+
+@register_general_jit_lookaside(_torch_autograd_function_apply)
 def _general_jit_torch_autograd_function_apply_lookaside(obj: Any, *args, **kwargs):
     """Encapsulate forward into a bsym, define and register augmented fwd and bwd.
 
@@ -1000,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`

thunder/dynamo/splitter.py

@@ -230,14 +230,17 @@ def is_thunder_supported_partition(node: torch.fx.Node) -> bool:
             example_input_metadatas.append(list(example_input_metadata))
             # Replace PyTorch operators within the checkpointed function with the corresponding Thunder operators
             checkpoint_converter(split_gm, graph_module)
+            # Since torch 2.12 this is a lazy GraphModule: generate its code now, so thunder doesn't trace fx codegen
+            recompile_graph(graph_module)
 
             jit_fn = thunder_jit(graph_module, is_differentiable_outputs=is_differentiable_outputs)
             # Update the node name from "submod_*" to "thunder_*" for more user-friendly names
             update_node_and_submodule(split_gm, node, node.name.replace("submod", "thunder"), jit_fn)
             thunder_compiled_fns.append(jit_fn)
-            submodule_to_compiled_fns[getattr(original_split_gm, node_name)] = CompiledFunction(
-                jit_fn, CompilerType.THUNDER
-            )
+            original_graph_module = getattr(original_split_gm, node_name)
+            # ThunderCompilerGraphBenchmarking jits this copy too
+            recompile_graph(original_graph_module)
+            submodule_to_compiled_fns[original_graph_module] = CompiledFunction(jit_fn, CompilerType.THUNDER)
         elif node.name.startswith("submod"):  # For inductor
             graph_module = getattr(split_gm, node.name)
 

thunder/tests/distributed/test_dtensor.py

@@ -157,7 +157,8 @@ def test_dtensor_basic_op(self, executor, fn_key):
         dim_size = 16
 
         def _helper(fn, in_dtensor, w_dtensor):
-            expected = torch.compile(fn)(in_dtensor, w_dtensor)
+            # NOTE Eager, like the opinfo tests: compiling it fails in inductor on torch nightly (pytorch#152639)
+            expected = fn(in_dtensor, w_dtensor)
             tmodel = thunder.jit(fn, executors=executors_map[executor].executors_list())
             actual = tmodel(in_dtensor, w_dtensor)
 

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"):

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):
@@ -9632,6 +9637,12 @@ def cross_entropy_error_generator(op, device, dtype=torch.float32, **kwargs):
             pytest.mark.skip,
             "test_vjp_correctness",
         ),
+        DecorateInfo(
+            pytest.mark.timeout(360),
+            "test_core_vs_torch_consistency",
+            executors=("nvfuser",),
+            devicetypes=(devices.DeviceType.CUDA,),
+        ),
     ),
 )
 nn_ops.append(cross_entropy_opinfo)

thunder/tests/test_auto_register_torchops.py

@@ -5,6 +5,7 @@
 import pytest
 import thunder
 import thunder.torch.default_torch_ops as ops
+from thunder.constants import _TORCH_GREATER_EQUAL_2_13
 from thunder.torch import _get_torch_function_name
 import torch
 
@@ -248,3 +249,13 @@ def fn_none(a):
         cfn(a)
         ops = thunder.get_auto_registered_torch_op_names(cfn)
         assert expect == ops
+
+
+def test_named_tensor_ops_follow_torch():
+    # torch 2.13 removed named tensors, and this table is built at import time.
+    named_tensor_methods = {"align_as", "align_to", "has_names", "refine_names", "rename"}
+    registered = {fn.__name__ for fn in ops.torch_auto_registered_ops[torch.Tensor]}
+    if _TORCH_GREATER_EQUAL_2_13:
+        assert not (named_tensor_methods & registered)
+    else:
+        assert named_tensor_methods <= registered

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

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
 

thunder/tests/test_transformer_engine_executor.py

@@ -635,12 +635,12 @@ def te_model(x):
         te_scales = []
         te_amax_hist = []
         with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
-            for buffer_key, _ in FP8GlobalStateManager.global_amax_buffer.items():
+            # TE >= 2.15 keeps these buffers on quantization_state
+            qstate = getattr(FP8GlobalStateManager, "quantization_state", FP8GlobalStateManager)
+            for buffer_key, _ in qstate.global_amax_buffer.items():
                 # needs to clone the tensors because TE will use in-place copy to modify it any time autocast is called
-                te_scales += [t.detach().clone() for t in FP8GlobalStateManager.global_scale_buffer[buffer_key]]
-                te_amax_hist += [
-                    t.detach().clone() for t in FP8GlobalStateManager.global_amax_history_buffer[buffer_key]
-                ]
+                te_scales += [t.detach().clone() for t in qstate.global_scale_buffer[buffer_key]]
+                te_amax_hist += [t.detach().clone() for t in qstate.global_amax_history_buffer[buffer_key]]
 
         # Make sure that the global state manager has been reset and
         # that there are only the buffers we need and not more

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

thunder/torch/default_torch_ops.py

@@ -1,5 +1,21 @@
 import torch
 
+from thunder.constants import _TORCH_GREATER_EQUAL_2_13
+
+# torch 2.13 removed named tensors (pytorch/pytorch#173895), and this table is built at import
+# time, so referencing the removed methods would break `import thunder` there.
+_named_tensor_methods = (
+    []
+    if _TORCH_GREATER_EQUAL_2_13
+    else [
+        torch.Tensor.align_as,
+        torch.Tensor.align_to,
+        torch.Tensor.has_names,
+        torch.Tensor.refine_names,
+        torch.Tensor.rename,
+    ]
+)
+
 torch_auto_registered_ops = {
     torch: [
         torch._native_multi_head_attention,
@@ -358,8 +374,8 @@
         torch.Tensor.addmv,
         torch.Tensor.addr,
         torch.Tensor.adjoint,
-        torch.Tensor.align_as,
-        torch.Tensor.align_to,
+        # align_as, align_to, has_names, refine_names and rename, on torch versions that still have them
+        *_named_tensor_methods,
         torch.Tensor.aminmax,
         torch.Tensor.angle,
         torch.Tensor.arccos,
@@ -432,7 +448,6 @@
         torch.Tensor.greater,
         torch.Tensor.greater_equal,
         torch.Tensor.half,
-        torch.Tensor.has_names,
         torch.Tensor.heaviside,
         torch.Tensor.histc,
         torch.Tensor.histogram,
@@ -520,8 +535,6 @@
         torch.Tensor.quantile,
         torch.Tensor.rad2deg,
         torch.Tensor.ravel,
-        torch.Tensor.refine_names,
-        torch.Tensor.rename,
         torch.Tensor.renorm,
         torch.Tensor.reshape_as,
         torch.Tensor.resize,