It should never be the case that you have owning Python references that don't increase the C++ refcount.

That was also what I thought, but after adding a print before this line, I got:

Single GPU Functional
Loss: -0.1287534236907959
Grad: [tensor([[0.2500, 0.2500, 0.2500, 0.2500]], device='cuda:0')]

Multi-GPU Functional
BroadcastBackward@0x56135dc84c50 shared_ptr use_count is 1, python refcnt is 3
BroadcastBackward@0x56135dc84c50 shared_ptr use_count is 2, python refcnt is 2
Loss: -0.1287534236907959
Grad: [tensor([[0.2500, 0.2500, 0.2500, 0.2500]], device='cuda:0'), tensor([0.], device='cuda:0')]

Multi-GPU Inline
BroadcastBackward@0x56135dc84c50 shared_ptr use_count is 1, python refcnt is 2
BroadcastBackward@0x56135dc84c50 shared_ptr use_count is 1, python refcnt is 1
Loss: -0.1287534236907959
Grad: [tensor([[0.2500, 0.2500, 0.2500, 0.2500]], device='cuda:0'), tensor([0.], device='cuda:0')]
BroadcastBackward@0x56139536c600 shared_ptr use_count is 1, python refcnt is 3
BroadcastBackward@0x56139536c600 shared_ptr use_count is 1, python refcnt is 2
BroadcastBackward@0x56139536c600 shared_ptr use_count is 2, python refcnt is 1
BroadcastBackward@0x56139536c600 shared_ptr use_count is 1, python refcnt is 1

Edit: added BroadcastBackward function object address

It still looks weird to me that the inline case is playing with two different BroadcastBackward objects @apaszke any thought?

I think it is because the loss object created in "Multi-GPU Functional" was deleted when the var is assigned again in the "Multi-GPU Inline" case.

Ok, but that means we're not managing the memory correctly, because we have multiple shared_ptrs that point to the same object, but use different control blocks! This is a much more serious bug, and while this patch might be a sufficient workaround for the issue you linked, there's a much bigger issue causing those problems.

@apaszke any suggestion for fixes? Greg and I discussed about using enable_shared_from_this, but it seems the implementation explicitly avoids shared_from_this to prevent holding a shared_ptr in PyFunction.

Functions already inherit from enable_shared_from_this. Can't we just use that?

@apaszke I was worried about the comments below. Could you please elaborate some more? Correct me if I get it wrong: If we use shared_from_this, PyFunction will hold a weak_ptr as a member, which requires that a shared_ptr of PyFunction is created and held somewhere. Do we have that shared_ptr created elsewhere? If no, will the behavior be undefined?

// Similar to shared_from_this. There's a problem that the Python object
// and its cdata depend on each other being alive, so we can't keep
// shared_ptrs as members, but we'd like to be able to manage the lifetime of
// the objects using shared_ptrs in the C++ graph. This returns a new
// shared_ptr, which will decrement the Python reference count when it's
// destructed. WARNING: it's generally not safe to create weak_ptrs from
// these shared_ptrs since multiple shared_ptrs may control the same underlying
// object.

Using shared_from_this in THPFunction_asFunction does not work, because there is no shared_ptr pointing to it:

Multi-GPU Functional
Traceback (most recent call last):
  File "bug_test.py", line 32, in <module>
    loss = gradient_penalty(multigpu_net, x)
  File "bug_test.py", line 10, in gradient_penalty
    output = netD(x)
  File "/home/shenli/project/pytorch/torch/nn/modules/module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/shenli/project/pytorch/torch/nn/parallel/data_parallel.py", line 139, in forward
    inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
  File "/home/shenli/project/pytorch/torch/nn/parallel/data_parallel.py", line 150, in scatter
    return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)
  File "/home/shenli/project/pytorch/torch/nn/parallel/scatter_gather.py", line 35, in scatter_kwargs
    inputs = scatter(inputs, target_gpus, dim) if inputs else []
  File "/home/shenli/project/pytorch/torch/nn/parallel/scatter_gather.py", line 28, in scatter
    return scatter_map(inputs)
  File "/home/shenli/project/pytorch/torch/nn/parallel/scatter_gather.py", line 15, in scatter_map
    return list(zip(*map(scatter_map, obj)))
  File "/home/shenli/project/pytorch/torch/nn/parallel/scatter_gather.py", line 13, in scatter_map
    return Scatter.apply(target_gpus, None, dim, obj)
RuntimeError: bad_weak_ptr

As multiple people reported hitting this bug, shall we have a consensus on the fixes? @apaszke @gchanan @resistor @ssnl

#14960 is another bug report for this reference count issue.

@mrshenli where are those problematic autograd::Function instances created? One thing we could try would be to try to use shared_from_this() and create a new shared_ptr if that fails. That way we'd always have a single one pointing to that object.

@apaszke

One thing we could try would be to try to use shared_from_this() and create a new shared_ptr if that fails. That way we'd always have a single one pointing to that object.

Let me try it.

The stack trace is as follows.

replicate.py: replicate
python_function.cpp: THPFunction_apply -> process_outputs -> _wrap_outputs
python_variable.cpp: THPVariable_NewWithVar -> THPFunction_asFunction

UPDATE: this comment was written before your change to use shared_from_this(). As you will se below, I think it will break the fix from #1705.

Looking back at the history it feels like since #1705 , THPFunction are now completely managed by the refcount of their python object. A shared pointer use count is meaningless.
Also the "Similar to shared_from_this." comment is outdated and should be removed. It was the implementation before #1705 that was similar to it (storing a weakptr).

@apaszke we actually want a different shared pointer here. In particular to make sure that each THPVariable -> THPFunction contribute one refcount to the THPFunction object (the link between the two is via a shared pointer).

I'm not sure we need to add python related logic to Functions though.
We can have a use_count() method on Function that counts how many objects smart pointers point to it.

• For regular Function, it returns 1 as a single shared pointer ever use it, and we don't do pointer logic with it.
• For PyFunction, it returns Py_REFCNT(obj). In this case there is always a THPFunction associated with this PyFunction that is stored in obj and it's refcount represent how many use there are (from shared pointers used in cpp + actual python references).
  Based on this, you can delete a Function if and only if the Function.use_count() == 1 and you are the only one using your own smart pointer. Your condition would become, fn->use_count() == 1 && fn.use_count() == 1: if the Function is only referenced by a single smart pointer and I only have 1 ref left.

@albanD I agree that the current fix using shared_from_this() breaks what #1705 tries to achieve. Your proposed solution SGTM. I will wait for 1-2 days for others to comment. If no objection, I will revise this PR. Thanks!

Ok, while I agree that the current solution might regress #1705, having us check Python refcounts in pure C++ code (that should have no dependency on Python) is an example of a tail wagging the dog... For me it's clear that we absolutely need to use a single shared_ptr, or we will have similar bugs in the future too. What we should do instead is have a more sophisticated deleter. In particular, it should no longer be a simple invocation of Py_DECREF, but should be an object that stores how many decrefs should it perform when it is destroyed. We can bump this by doing something like py_function_ptr.get_deleter()->bump_decref();.

@malvika2147 and I were looking at the refcount situation in PyFunction, and I think the whole ownership situation here is a bit screwy.

Let us introduce our dramatis personae:

• THPFunction - a PyObject that represents a ctx object that is passed to our custom Python autograd functions. It contains a PyFunction
• PyFunction - a subclass of Function which says how to compute the actual backward function. Note that this is a closure in the sense that it is associated with a ctx object which actually contains the saved tensors. It contains a non-owning pointer to THPFunction

OK, do you see something wrong with this situation? The ownership here is totally backwards! It is PyFunction that should own the THPFunction, because a PyFunction is a closure that should own its context. To put it in other words, it is OK if a PyFunction gets freed before THPFunction, but not vice versa.

Once we fix this, we can fix the stupidity, e.g., here:

// torch/csrc/autograd/python_variable.cpp [unix, utf-8] [ln:57 col:22] 
    if (auto fn = dynamic_cast<PyFunction*>(v->cdata.grad_fn_unsafe())) {
      // Create a new reference to the THPFunction. This ensures that ref count
      // of the THPFunction is at least the number of referring THPVariables.
      const auto output_nr = v->cdata.output_nr();
      auto grad_fn = THPFunction_asFunction((THPFunction*)fn->obj);
      v->cdata.set_gradient_edge({std::move(grad_fn), output_nr});
    }

I also suspect that there is no need to have torch.autograd.Function be backed by an actual C++ class (unless we desperately want ctx to be a subclass of the Function you defined, but I don't see any good reason why we should work so hard to actually make this true.)

It might be easier to fix this once we apply @yf225's patch at #22762, since I suspect the reason we got into this situation in the first place was because of legacy autograd function API.

@ezyang @malvika2147

Thanks for the investigation. If you plan to work on this, feel free to abandon this PR or hijack it. It will take me a while to get back to this.

This should be addressed by #22983 per @ezyang.