Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Raise when in place operations occur on leafs requiring grad #1458

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
Open

Raise when in place operations occur on leafs requiring grad #1458

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
b56dd80
raise if is_leaf and require_grad in inplace operations
beverlylytle Nov 19, 2024
d79173c
Merge branch 'main' into check_inplace_leafs
beverlylytle Nov 19, 2024
e4ad972
restore wraps
beverlylytle Nov 20, 2024
551de30
add test and comment
beverlylytle Nov 21, 2024
a179d12
test thunder, not torch
beverlylytle Nov 22, 2024
a26ed67
add parameter to copy__meta
beverlylytle Nov 25, 2024
c279c14
Merge branch 'main' into check_inplace_leafs
beverlylytle Nov 26, 2024
590ef18
Merge branch 'main' into check_inplace_leafs
t-vi Nov 27, 2024
3fb1f53
apply function with compile data
beverlylytle Nov 27, 2024
cfd5143
pass grad_enabled bool instead of relying on compile data
beverlylytle Dec 12, 2024
ad518a5
Merge branch 'main' into check_inplace_leafs
beverlylytle Dec 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions thunder/core/prims.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4030,6 +4030,8 @@ def embedding_backward_meta(grad, indices, num_weights, padding_idx, scale_grad_
def copy__meta(
copy_from: TensorProxy,
copy_to: TensorProxy,
*,
grad_enabled: bool = False,
):
utils.check_type(copy_from, TensorProxy)
utils.check_type(copy_to, TensorProxy)
Expand Down
2 changes: 1 addition & 1 deletion thunder/core/transforms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1598,7 +1598,7 @@ def zeros_like(x):
prims.PrimIDs.LOG2: lambda x: (prims.log2(x), (x,)),
prims.PrimIDs.ZETA: lambda x, y: (prims.zeta(x, y), (x, y)),
prims.PrimIDs.FMOD: lambda x, y: (prims.fmod(x, y), (x, y)),
prims.PrimIDs.COPY_: lambda x, y: (prims.copy_(x, y), tuple()),
prims.PrimIDs.COPY_: lambda x, y, grad_enabled: (prims.copy_(x, y, grad_enabled=grad_enabled), tuple()),
prims.PrimIDs.CLONE: lambda x: (prims.clone(x), tuple()),
}

Expand Down
3 changes: 3 additions & 0 deletions thunder/executors/nvfuserex_impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2075,6 +2075,8 @@ def var_mean(
def _copy__check(
copy_from: TensorProxy,
copy_to: TensorProxy,
*,
grad_enabled: bool,
) -> bool:
return are_supported_tensors(copy_from, copy_to)

Expand All @@ -2085,6 +2087,7 @@ def copy_(
*,
fd: FusionDefinition,
lc_to_nv_map: dict,
grad_enabled: bool,
) -> Any:
nvcopy_from = getnv(copy_from, fd, lc_to_nv_map)
nvcopy_to = getnv(copy_to, fd, lc_to_nv_map)
Expand Down
20 changes: 7 additions & 13 deletions thunder/executors/torchex.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,32 +1,24 @@
from __future__ import annotations
import operator
import importlib
from dataclasses import replace
from contextlib import ContextDecorator
from functools import wraps, partial
from inspect import signature
from itertools import groupby
from functools import partial, wraps
from numbers import Number
from typing import TYPE_CHECKING
from collections.abc import Callable
from collections.abc import Hashable, Sequence
from collections.abc import Sequence
from types import ModuleType
from enum import Enum, auto

import torch
import math
from looseversion import LooseVersion

from thunder.core.compile_data import get_compile_data
import thunder.core.dtypes as dtypes
from thunder.core.dtypes import to_torch_dtype, to_dtype
import thunder.core.devices as devices
from thunder.core.devices import to_torch_device, to_device
import thunder.core.prims as prims
from thunder.core.trace import TraceCtx, set_tracectx, reset_tracectx, from_trace
from thunder.core.proxies import NumberProxy, TensorProxy, FutureTensorProxy, variableify, pytype
from thunder.core.pytree import tree_flatten, tree_unflatten
from thunder.core.symbol import Symbol, BoundSymbol
from thunder.core.proxies import NumberProxy, TensorProxy, FutureTensorProxy, pytype
from thunder.core.symbol import Symbol
from thunder.distributed.prims import DistributedReduceOps
import thunder.distributed.prims as dist_prims
import thunder.core.utils as utils
Expand Down Expand Up @@ -2192,7 +2184,9 @@ def is_float_type(self, input):
einops._backends._type2backend[TensorProxy] = EinopsThunderBackend()


def _copy__impl(copy_from, copy_to):
def _copy__impl(copy_from, copy_to, grad_enabled):
if grad_enabled and copy_to.is_leaf and copy_to.requires_grad:
raise RuntimeError("a leaf Variable that requires grad is being used in an in-place operation.")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am wondering if Symbol copy_ in thunder/torch/__init__.py is more appropriate location for the check.

lightning-thunder/thunder/torch/__init__.py

Lines 1961 to 1963 in 60f3ee1

@torchsymbol(torch.Tensor.copy_, is_method=True) # , tags=(prims.OpTags.IN_PLACE,))
def copy_(a, b, /):
return prims.copy_(b, a)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

a and b are proxies and it it not clear to me if a proxy knows that it is a leaf.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

They do not. It's only a PyTorch concept that's available at runtime inside _copy__impl.

Copy link
Collaborator

@kshitij12345 kshitij12345 Nov 22, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, previously I missed that the fix was in copy_impl. And since, it is happening at runtime, I am wondering if compile_data is actually available.

Quick test shows (see below) that it wouldn't be. So, we probably need a way to check if this copy was called under no_grad in users code (as PyTorch supports inplace of leaf tensors under no_grad, see comment).

Snippet to check if compile_data is available -

import torch
import thunder
from thunder.extend import OperatorExecutor
from thunder.core.compile_data import get_compile_data
from thunder.core.proxies import TensorProxy

ex = OperatorExecutor("ex")

def clone_impl(x):
    cd = get_compile_data()
    print(cd)  # None
    return x

clone = ex.register_operator("clone", meta=lambda x: TensorProxy(like=x), fn=clone_impl)

def fn(x):
    return clone(x)

x = torch.ones(3)

jfn = thunder.jit(fn)

jfn(x)
exec_trace = thunder.last_traces(jfn)[-1]
# print(exec_trace)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Indeed, compile_data was not available, but now it should be with the added context manager in thunder/init.py

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is still incorrect because as discussed in #1486, the value of compile_data.is_grad_enabled here would be that of last updated state which can lead to incorrectness when used outside of tracing context.

We can see the discrepancy here.

import torch
import thunder

x = torch.randn(3, 3, requires_grad=True)

@torch.no_grad
def fn(x):
  return x.add_(1)

fn(x)  # This works

thunder.jit(fn)(x)  # This raises error

So, whether the copy is in no_grad region needs to be captured during the tracing time.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, this is why I created the other issue. This PR fixes the leaf/grad issue when there is no annotation. When there is an annotation, another approach is required. This other approach may or may not involve using compile data in _copy__impl.

As far as I understand, compile data is the medium for passing around data such as whether grad is enabled. But as the other issue points out, compile data reflects the end state of a function call and not the "live" state, at least at the time it reaches _copy__impl. So I'm left with the questions "are there other mechanisms for passing around whether grad is enabled?" "where else in the execution is it simultaneously knowable that a (1) leaf tensor (2) requiring grad is being (3) copied when (4) grad is enabled?" "is it feasible/desirable to make the compile data more dynamic?" "is there a way to context-manage the tensors so that their requires_grad flags are set to False when the interpreter sees torch._C._set_grad_enabled(False), and then later restored, thereby obviating the need for the compile data for this check?" Do you have suggestions for a fix that addresses both issues? Or can we close out this issue and move the discussion to the more involved issue?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So to tackle - leaf tensor requiring grad being copied into when grad is enabled, I think similar to a previous commit,
we can update prims.copy to take a argument is_grad_enabled. With this, ltorch.copy will query cd.is_grad_enabled and call prims.copy by also passing this argument.

lightning-thunder/thunder/torch/__init__.py

Lines 1984 to 1986 in 9de5434

@torchsymbol(torch.Tensor.copy_, is_method=True) # , tags=(prims.OpTags.IN_PLACE,))
def copy_(a, b, /):
return prims.copy_(b, a)

With these changes, the copy_impl's signature will also change to accept is_grad_enabled and it will be called at runtime with a tensor which we can query if it is a leaf and also whether grad was enabled or not when calling that particular copy. Wdyt @beverlylytle?

Though, I am curious if there is another approach to this - cc: @IvanYashchuk

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's see what the CI thinks.

copy_to.copy_(copy_from)
return copy_to

Expand Down
15 changes: 14 additions & 1 deletion thunder/tests/test_inplace_copy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import thunder
import thunder.core.dtypes as datatypes
import thunder.torch as ttorch
from thunder.tests.framework import instantiate, nvFuserExecutor
from thunder.tests.framework import instantiate, nvFuserExecutor, TorchExecutor


@instantiate(dtypes=datatypes.all_dtypes - datatypes.float_8bit_dtypes)
Expand Down Expand Up @@ -178,3 +178,16 @@ def func(T0):
assert_close(a_ref, a)
for o, o_ref in zip(o_thunder, o_eager):
assert_close(o, o_ref)


@instantiate(executors=(TorchExecutor,), dtypes=datatypes.float_math_dtypes)
def test_inplace_copy_of_leaf_requiring_grad_fails(executor, device, dtype):
def fn(x):
x.copy_(x)

jitted_fn = executor.make_callable(fn)

tdtype = ttorch.to_torch_dtype(dtype)
a = make_tensor((4, 4), device=device, dtype=tdtype, requires_grad=True)
with pytest.raises(RuntimeError):
jitted_fn(a)
17 changes: 9 additions & 8 deletions thunder/tests/test_inplace_functionalization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -477,15 +477,15 @@ def f(xs, ys, z):
)
def test_inplace_to_tensors_with_grad(executor, device, _):
@torch.no_grad
beverlylytle marked this conversation as resolved.
Show resolved Hide resolved
def add_y(x, y):
x.add_(y, alpha=0.1)
def add_grad(x, y):
return x.add_(x.grad)

@torch.no_grad
def add_grad(x, y):
x.add_(x.grad, alpha=0.1)
def add_y(x, y):
return x.add_(y, alpha=0.1)

for f in (add_y, add_grad):
jitted_f = executor.make_callable(f)
for fn in (add_grad, add_y):
jitted_f = executor.make_callable(fn)
x = make_tensor((2, 2), device=device, dtype=torch.float32, requires_grad=True)
x.grad = make_tensor((2, 2), device=device, dtype=torch.float32)
y = make_tensor((2, 2), device=device, dtype=torch.float32)
Expand All @@ -495,7 +495,7 @@ def add_grad(x, y):
y_ref = y.clone().detach()

res = jitted_f(x, y)
res_ref = f(x_ref, y_ref)
res_ref = fn(x_ref, y_ref)

torch.testing.assert_close(x, x_ref)
torch.testing.assert_close(x.grad, x_ref.grad)
Expand Down Expand Up @@ -549,7 +549,8 @@ def single_tensor_adam(
ref_state_steps = [torch.tensor(1, device=device) for _ in range(2)]
single_tensor_adam(*ref_tensors, state_steps=ref_state_steps)

jitted = executor.make_callable(single_tensor_adam)
# torch.compile does not support accessing the ContextVariable compile data used in _copy__impl_
jitted = executor.make_callable(single_tensor_adam, torch_compile_fullgraph=False)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interesting that torch.compile creates a graph break when calling get on ContextVariable.

import torch
from contextvars import ContextVar

_compile_data = ContextVar("compile_data", default=(None, None))

def fn(x):
    _compile_data.get()
    return x + 1

torch.compile(fn, fullgraph=False)(torch.randn(3, 3))  # Works with GraphBreak at _compile_data.get()
torch.compile(fn, fullgraph=True)(torch.randn(3, 3))  # Fails

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does Thunder's Interpreter do? It probably fails

Copy link
Collaborator

@kshitij12345 kshitij12345 Nov 22, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thunder just burns the value in computation trace (if used) without having a corresponding check in prologue. (Will file an issue for the same).

Eg.

import torch
import thunder
from contextvars import ContextVar

_compile_data = ContextVar("compile_data", default=1)

def fn(x):
    v = _compile_data.get()
    return x + v

jfn = thunder.jit(fn)
o = jfn(torch.ones(3,))
print(o)  # tensor([2., 2., 2.])

_compile_data.set((2,))
o = jfn(torch.ones(3,))
print(o)  # tensor([2., 2., 2.])

print(thunder.last_prologue_traces(jfn)[-1])
# @torch.no_grad()
# @no_autocast
# def prologue(*args, **kwargs):
#   # args: "Any"
#   check_len(args, 1)
#     # prims.check_len(args, 1)
#   # kwargs: "Any"
#   check_len(kwargs, 0)
#     # prims.check_len(kwargs, 0)
#   x: "cpu f32[3]" = args[0]
#   check_tensor_metadata(x, (3,), 'cpu', torch.float32, False)
#     # prims.check_tensor_shape_and_metadata(x, (3,), 'cpu', torch.float32, False)
#   cache_info: "Any" = thunder._get_cache_info()
#   cache_info_default_dtype: "<class 'torch.dtype'>" = cache_info['default_dtype']
#   check_literal_like(cache_info_default_dtype, torch.float32)
#     # prims.check_literal_like(cache_info_default_dtype, torch.float32)
#   cache_info_default_device: "<class 'torch.device'>" = cache_info['default_device']
#   check_literal_like(cache_info_default_device, torch.device("cpu"))
#     # prims.check_literal_like(cache_info_default_device, torch.device("cpu"))
#   cache_info_is_autocast_enabled: "bool False" = cache_info['is_autocast_enabled']
#   check_number_type_and_value(cache_info_is_autocast_enabled, False)
#     # prims.check_number_type_and_value(cache_info_is_autocast_enabled, False)
#   cache_info_no_grad_sync: "bool False" = cache_info['no_grad_sync']
#   check_number_type_and_value(cache_info_no_grad_sync, False)
#     # prims.check_number_type_and_value(cache_info_no_grad_sync, False)
#   cache_info_alias_tensor_indices: "str" = cache_info['alias_tensor_indices']
#   check_string_value(cache_info_alias_tensor_indices, '')
#     # prims.check_string_value(cache_info_alias_tensor_indices, '')
#   cache_info_is_grad_enabled: "bool True" = cache_info['is_grad_enabled']
#   check_number_type_and_value(cache_info_is_grad_enabled, True)
#     # prims.check_number_type_and_value(cache_info_is_grad_enabled, True)
#   return ((x,), ())

print(thunder.last_traces(jfn)[-1])
# @torch.no_grad()
# @no_autocast
# def computation(x):
#   # x: "cpu f32[3]"
#   t0 = torch.add(x, 1, alpha=1)  # t0: "cpu f32[3]"
#     # t0 = ltorch.add(x, 1, alpha=1)  # t0: "cpu f32[3]"
#       # _ = prims.convert_element_type(1, float)
#       # t0 = prims.add(x, 1.0)  # t0: "cpu f32[3]"
#   return t0

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Issue filed at #1464

params, grads, exp_avgs, exp_avg_sqs = tensors

jitted(params, grads, exp_avgs, exp_avg_sqs, state_steps)
Expand Down
3 changes: 2 additions & 1 deletion thunder/torch/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1983,7 +1983,8 @@ def copysign_(a, b, /):

@torchsymbol(torch.Tensor.copy_, is_method=True) # , tags=(prims.OpTags.IN_PLACE,))
def copy_(a, b, /):
return prims.copy_(b, a)
cd = get_compile_data()
return prims.copy_(b, a, grad_enabled=cd.is_grad_enabled if cd is not None else False)


# TODO Implement div
Expand Down
Loading