Repository navigation
test(distributed): compare DTensor basic ops against eager, not torch.compile [14/n] - #2851
Merged
Merged
Conversation
The project metadata lives in pyproject.toml's [project] table, but `exclude *.toml` left it out of the sdist. A wheel built from that sdist comes out named `thunder`, without its summary, license or requires-python, and uv rejects it outright. pyproject.toml is the only toml at the root, so drop the exclude and let setuptools include it as it does by default.
torchaudio is not imported anywhere in the repo. CI pins it in lockstep with torch, but torchaudio stopped tracking torch after 2.11, so the pin no longer resolves and every job fails at dependency install.
pyproject.toml turns every FutureWarning into an error. Three come from code outside thunder and fail tests that have nothing wrong with them: - torch's test utils read the deprecated torch._dynamo.config.inline_inbuilt_nn_modules as a default argument, so importing them fails collection of every distributed test file. - torch.jit.script is deprecated, and torch.func.jvp, which the grads tests use as their reference, imports decompositions_for_jvp, which scripts its decompositions on import. The prefix covers both wordings torch emits. - bitsandbytes still calls torch._check_is_size, which torch deprecates. Ignore exactly those messages. Every other FutureWarning stays fatal.
torch 2.13 removed named tensors, so Tensor.align_as, align_to, has_names, refine_names and rename are gone. default_torch_ops.py lists them in a module-level dict, so `import thunder` fails with an AttributeError. Only register them on torch versions that still have them.
torch 2.13 reimplemented Function.apply in C, so Function.apply.__func__ no longer exists and `import thunder` fails at module scope. Swapping in the bound Function.apply would not work: it is a builtin_function_or_method, and MyFunc.apply is a fresh object that is neither equal nor hash-equal to it, so the identity-keyed lookaside map would never match a user's subclass and every autograd.Function would silently stop being traced. The classmethod_descriptor on Function is the one object all subclasses share, so key the lookaside on that, and have the interpreter resolve builtin classmethods to their descriptor.
bhimrazy
requested review from
k223kim,
lianakoleva and
tejapulagam
as code owners
September 24, 2026 13:31
bhimrazy
added this pull request to stack #2849
September 24, 2026 13:31
Since torch 2.12, the partitions the splitter hands to thunder.jit are _LazyGraphModules whose code has not been generated yet (on 2.11 they are plain GraphModules). Their first call, and with it GraphModule.recompile(), then runs inside thunder's interpreter, and tracing torch fx codegen fails on what the interpreter doesn't support, which broke every DynamoThunder test: TypeError: object of type 'MappingKeysView' has no len() TypeError: object of type 'generator' has no len() Call recompile_graph first, as the splitter already does for split_gm. It runs real_recompile() on lazy modules and is a plain recompile otherwise. Do the same for the partition's copy in original_split_gm, which ThunderCompilerGraphBenchmarking jits.
torch 2.14 rejects group_norm inputs with an empty channel or spatial
extent ("Expected number of channels to be greater than 0", "Expected HxW
to be greater than 0"); 2.13 and earlier accepted them. The sample
generator builds shapes from (0, 2) in every dim, so the torch reference
failed on inputs it now refuses. An empty batch is still accepted, so only
channels and the spatial dims are filtered.
The boundary is 2.14, not 2.13, hence the new constant.
torch 2.11 replaced args_tensor_mask on the autograd_function_apply higher-order op with a required saved_for_backward_idx (and 2.13 added an optional dirty_idx), so test_autograd_function_apply failed with KeyError: 'saved_for_backward_idx'. The test helper now passes it, and thunder's autograd_function_apply symbol accepts both new kwargs, which test_autograd_function_apply_with_no_grad passes through. thunder does not use either, same as non_differentiable_idx.
Since torch 2.12, the fwd subgraph dynamo builds returns its outputs as a tuple, and the caller indexes it: return ((y,), (l_x_,)) getitem = autograd_function_apply[0] python_return(*sequencify(output)) unpacked a one-element tuple to the bare tensor, so [0] indexed into the tensor instead, and thunderfx returned a scalar where eager gives f32[2]. test_higher_order_inplace_alias_update: AssertionError: The values for attribute 'shape' do not match: torch.Size([]) != torch.Size([2]). Return the structure the graph produced. Bare and multi-element outputs are unchanged.
Running these models in bfloat16 on CPU intermittently kills the Windows test process with an illegal instruction (0xc000001d). It happens in eager torch, before thunder runs, so pytest cannot catch it and the whole xdist worker dies. The fp32 variants of the same models still run on Windows, and these still run on Linux and macOS.
ipython[all] includes test_extra, which now pulls jupyter-ai -> jupyter-server-mcp -> fastmcp -> mcp, and with it newer PyJWT and cryptography. The GPU image ships both through apt, with no RECORD file, so pip cannot upgrade them and the notebooks job fails at install: error: uninstall-no-record-file Cannot uninstall cryptography 41.0.7 The notebooks need none of that. Keep every other extra of [all].
…ation_state TE ships no cu128 wheel for torch 2.8, so the job builds it from source. From 2.17, that build includes torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp, which torch 2.8 does not have, so the install fails. Cap at 2.16. TE 2.15 moved global_amax_buffer and the other amax buffers off FP8GlobalStateManager onto its quantization_state, so the test failed with AttributeError on 2.16. Read them from there, falling back to the class on older TE.
On torch nightly, test_core_vs_torch_consistency_cross_entropy with nvFuser takes 175-240s per dtype, and float16 keeps landing on the ops job's 240s limit: 240.18s call ...cross_entropy_nvfuser_cuda_thunder.dtypes.float16 Failed: Timeout (>240.0s) from pytest-timeout. It is the slowest test in test_ops by a wide margin (the next is ~90s), so raise the limit for it alone rather than for the whole job.
test_dtensor_basic_op built its reference with torch.compile, which put an inductor Triton launch in front of a test that is about thunder's lowering of the op over a DTensor. On the torch nightly image that launch fails: torch/_inductor/runtime/static_triton_launcher.py:291 RuntimeError: CUDA driver error: invalid argument and the other rank then waits until the 300s process timeout, so the same cause shows up as a crash in one variant and as timeouts in two others. Binding each rank to its device first does not change this; an eager reference fixes all three, as the opinfo tests already use.
bhimrazy
force-pushed
the
bhimrazy/dtensor-eager-reference
branch
from
September 24, 2026 13:44
4d8ab36 to
741efbf
Compare
lianakoleva
approved these changes
Sep 24, 2026
…aved-for-backward-idx
…/autograd-fwd-tuple-output
…ormer-engine-2.16
…ntropy-nvfuser-timeout
…tensor-eager-reference
Base automatically changed from
bhimrazy/cross-entropy-nvfuser-timeout
to
main
October 12, 2026 14:03
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Compares
test_dtensor_basic_opagainst eager instead oftorch.compile, as the opinfo tests in the same file already do.Problem: the test built its reference with
torch.compile, so an inductor Triton launch ran before thunder was reached. On torch nightly that launch fails withCUDA driver error: invalid argument, and the other rank then waits for the 300s process timeout. So one cause shows up as a crash in one variant and as timeouts in two others. The test is about thunder's lowering of the op over a DTensor, not inductor's.Checked on 2× L4 with the
pt_main-devimage: 3 of the 6 variants fail before (11 min, with the timeouts), and all 6 pass after (55s). Binding each rank to its device first doesn't change the result.Related CI:
all-tests.yamldistributedonpt_main-dev: the threetest_dtensor_basic_op_*failures.CI log (
distributed,pt_main-dev, on #2847)Part 14 of the breakdown of #2832.