Don't replace unused variables with None (#2396)

Author: beverlylytle

notebooks/liger_kernel.ipynb

@@ -251,7 +251,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "id": "32cd98f0-a36f-4e01-8ae4-6e36adf2699b",
    "metadata": {},
    "outputs": [],
@@ -322,11 +322,12 @@
     "        while bound_symbols:\n",
     "            bsym = bound_symbols.pop(0)\n",
     "            if bsym.sym == litgpt_apply_rope:\n",
-    "                for i, bsym2 in enumerate(bound_symbols):\n",
+    "                while bound_symbols:\n",
+    "                    bsym2 = bound_symbols.pop(0)\n",
     "                    assert not any(o is bsym.output for o in bsym2.flat_outs)\n",
     "                    if bsym2.sym == litgpt_apply_rope:\n",
     "                        break\n",
-    "                bsym2 = bound_symbols.pop(i)\n",
+    "                    new_compute_trace.bound_symbols.append(bsym2.from_bsym())\n",
     "                assert bsym2.sym == litgpt_apply_rope\n",
     "\n",
     "                output = (bsym.output, bsym2.output)\n",

thunder/core/symbol.py

@@ -14,18 +14,16 @@
 from collections.abc import Sequence
 
 import thunder.core.baseutils as baseutils
+from thunder.core.baseutils import BoundSymbolInterface, TagBase
 import thunder.core.codeutils as codeutils
 from thunder.core.codeutils import Printable, Positions
-from thunder.core.baseutils import BoundSymbolInterface, TagBase
-from thunder.core.utils import FrozenDict, make_hashable
-from thunder.core.pytree import tree_flatten_with_dataclass, tree_unflatten, tree_map
-from thunder.core.proxies import Proxy, TensorProxy, variableify, CollectionProxy, ProxyTag
 from thunder.core.compile_data import get_compile_data
+import thunder.core.prims as prims
+from thunder.core.proxies import Proxy, TensorProxy, variableify, CollectionProxy, ProxyTag
+from thunder.core.pytree import tree_flatten, tree_flatten_with_dataclass, tree_unflatten, tree_map
+from thunder.core.trace import get_tracectx, VariableInterface
+from thunder.core.utils import FrozenDict, make_hashable
 
-from thunder.core.trace import (
-    get_tracectx,
-    VariableInterface,
-)
 
 #
 # Support for querying "traceable" functions
@@ -314,6 +312,17 @@ def __call__(self, *args, **kwargs):
         else:
             trace.push_scope(subsymbols)
             result = self.meta(*args, **kwargs)
+
+            # To avoid passing an arg directly to output, we make a shallow_copy
+            flat_results, spec = tree_flatten(result)
+            flat_args, _ = tree_flatten((args, kwargs))
+            for i, result_ in enumerate(flat_results):
+                for arg in flat_args:
+                    if arg is result_ and isinstance(arg, Proxy):
+                        flat_results[i] = prims.shallow_copy(arg)
+
+            result = tree_unflatten(flat_results, spec)
+
             trace.pop_scope()
 
         cd = get_compile_data()

thunder/core/transform_common.py

@@ -162,29 +162,14 @@ def dce(trace: Trace, needed_proxies: None | set[Variable] = None) -> Trace:
         else:
             needed = False
 
-        # NOTE This block is run even if we know we're preserving the operation, because it
-        #   may mark some of the operation's outputs as unused
-        some_unused = False
         for out in bsym.flat_proxy_outs:
             if variableify(out) in needed_proxies and producer_map[out] == bsym:
                 needed = True
-            else:
-                some_unused = True
+                break
 
         if needed:
             nbsym: BoundSymbol = bsym
 
-            # Replaces unused Proxy outputs with None
-            if some_unused:
-
-                def _helper(x):
-                    if isinstance(x, Proxy) and (variableify(x) not in needed_proxies or producer_map[x] != bsym):
-                        return None
-                    return x
-
-                nbsym_output = tree_map(_helper, bsym.output)
-                nbsym = bsym.from_bsym(output=nbsym_output)
-
             # Eliminates no-op subsymbols
             # NOTE In general editing subsymbols doesn't do anything, but no-op subsymbols are a pain
             #   for transforms to deal with. Transforms typically look for a "flattened" version of an

thunder/core/transforms.py

@@ -750,6 +750,16 @@ def _cat_prim_grad(tensors: list[TensorProxy], /, dim: int) -> TensorProxy:
 register_grad(pids.CAT, _cat_prim_grad)
 
 
+def _shallow_copy_prim_grad(a: TensorProxy) -> TensorProxy:
+    fwd = prims.shallow_copy(a)
+    g = get_grad(fwd)
+    put_grad(a, g)
+    return fwd
+
+
+register_grad(pids.SHALLOW_COPY, _shallow_copy_prim_grad)
+
+
 def _update_aliases_prim_grad(tensors: tuple[TensorProxy, ...]) -> tuple[TensorProxy, ...]:
     fwd_tensors = prims.update_aliases(tensors)
     for fwd_t, t in zip(fwd_tensors, tensors):

thunder/examine/memory_calculation.py

@@ -26,6 +26,7 @@
     "contiguous",
     "split",
     "torch_wait_prim_impl",
+    "shallow_copy",
 )
 
 # A registry of symbols that require special memory calculation;

thunder/executors/nvfuserex_impl.py

@@ -1764,6 +1764,16 @@ def clone(a: TensorProxy, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any:
 
 register_supported(PrimIDs.CLONE, clone, _elementwise_unary_check)
 
+
+def shallow_copy(a: TensorProxy, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any:
+    nva = getnv(a, fd, lc_to_nv_map)
+
+    return nva
+
+
+register_supported(PrimIDs.SHALLOW_COPY, shallow_copy, _elementwise_unary_check)
+
+
 # update_aliases is disabled.  nvfuser does not support it.
 # TODO: Enable this once nvfuser supports it.
 # def update_aliases(aliases: tuple[TensorProxy], *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any:

thunder/executors/torch_compile.py

@@ -227,6 +227,7 @@ def cuda_device_checker(*args, **kwargs):
     # parallel residual paths are used in the transformer block
     prims.div.id,
     prims.erf.id,
+    prims.shallow_copy.id,
 }
 torch_compile_cat_ex._implmap = {
     op: ImplInfo(checker=cuda_device_checker) for op in pytorch_ex.implmap if op in supported_ops

thunder/tests/opinfos.py

@@ -695,6 +695,12 @@ def _abs_torch(x: torch.Tensor | Number):
             dtypes=(datatypes.complex32,),
             devicetypes=(devices.DeviceType.CPU,),
         ),
+        DecorateInfo(
+            pytest.mark.skip(reason="PyTorch doesn't support abs for bool on cpu"),
+            "test_core_vs_torch_consistency",
+            dtypes=(datatypes.bool8,),
+            devicetypes=(devices.DeviceType.CPU,),
+        ),
         # Ref - https://github.com/Lightning-AI/lightning-thunder/issues/2363
         DecorateInfo(
             pytest.mark.skip,

thunder/tests/test_core.py

@@ -2111,28 +2111,6 @@ def func(qkv):
     assert "scaled_dot_product_attention" in tuple(bsym.sym.id for bsym in traces[-1].bound_symbols)
 
 
-@instantiate(dtypes=NOTHING)
-def test_no_passthrough_symbol(executor, device, _):
-    # A test case for the situation reported in
-    # "backward trace contains symbols not present in forward that cause
-    # NotImplementedError"
-    # When an operation simply passes through its input, we should not
-    # add it to the trace.
-
-    def func(x):
-        return x.type_as(x)
-
-    x = make_tensor((2, 2), device=device, dtype=torch.float32)
-    compiled = executor.make_callable(func)
-    out = compiled(x)
-    assert out is x
-    initial_trace_with_dce = thunder.last_traces(compiled)[3]
-    assert "Constructed by Dead Code Elimination" in str(initial_trace_with_dce)
-    assert len(initial_trace_with_dce.bound_symbols) == 2
-    assert initial_trace_with_dce.bound_symbols[0].sym.id == prims.PrimIDs.UNPACK_TRIVIAL
-    assert initial_trace_with_dce.bound_symbols[1].sym.id == prims.PrimIDs.RETURN
-
-
 @instantiate(
     dtypes=NOTHING,
     # https://github.com/Lightning-AI/lightning-thunder/issues/946

thunder/tests/test_examine_memory.py

@@ -113,5 +113,5 @@ def test_nanogpt_block():
     # We are checking the estimated memory against a fixed value for consistency.
     assert max_mem_fw[0] == 381754368
     assert sum(max_mem_fw[1].values()) == 375462912
-    assert max_mem_bw[0] == 641097728
-    assert sum(max_mem_bw[1].values()) == 440474624
+    assert max_mem_bw[0] == 741761024
+    assert sum(max_mem_bw[1].values()) == 541137920