Lightning-AI / Lightning-AI/lightning-thunder
Revise memory clearing mechanism in the torch.autograd.Function integration
@nikitaved is already working on this.
Since Jun 17, 2024.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
🚀 Memory clearing mechanism with torch.autograd.Function integration
Today we pass "saved for backward" tensors to the generated backward function inside a Python list and within the generated function we clear the list after it's unpacked. This is required to remove references to tensors and let the CUDA memory be freed as soon as the tensor is not needed anymore later in the function with del calls. The unavoidable Python behavior is to hold a reference to all function arguments until the end of the function, any tensor passed as an argument or part of immutable containers passed as an argument like tuples cannot be freed, and a reference to that still exists even after a del call.
This memory problem was fixed in https://github.com/Lightning-AI/lightning-thunder/commit/afd69e9cf03ccc9e7290a8cb477828258c79103b.
Another way how people sometimes achieve the same effect is by swapping the .data attribute, here's one example from TransformerEngine. However, this variant depends on internal PyTorch attributes.
@nikitaved has an idea of implementing a special object wrapping PyTorch tensors that would delete the reference to the tensor when del is called with the goal of removing clear_mutable_collection from the trace. See the comment here https://github.com/Lightning-AI/lightning-thunder/pull/596#discussion_r1639853031.
- Why do we need to remove
clear_mutable_collectionwhich is also just a verbose way of callingdel list[:]? - What will the trace look like after the idea is implemented?
- How should the input to the backward function be preprocessed? What is the expected overhead if any over putting "saved for backward" into a list?
- How can this custom wrapper of PyTorch tensor work universally with all different extensions that are not registered for
__torch_function__or__torch_dispatch__?
Click here to see how the backward trace looks today.
When printing the backward trace you can notice `clear_mutable_collection` after a collection is unpacked:In [1]: import torch; import thunder;
In [2]: @thunder.jit
...: def func(a):
...: for _ in range(2):
...: a = a @ a
...: return a
...:
In [3]: a = torch.randn(3, 3, device="cuda", requires_grad=True)
In [4]: func(a);
In [5]: thunder.last_backward_traces(func)[-1]
def backward_fn(saved_for_backward, cotangents):
# saved_for_backward: "Collection"
# cotangents: "Collection"
C0, _, = saved_for_backward
clear_mutable_collection(saved_for_backward)
del saved_for_backward
t2, = cotangents
clear_mutable_collection(cotangents)
del cotangents
a, t0, = C0
clear_mutable_collection(C0)
del C0
t13 = torch.permute(a, (1, 0)) # t13: "cuda:0 f32[3, 3]"
# t13 = ltorch.permute(a, (1, 0)) # t13: "cuda:0 f32[3, 3]"
# t13 = prims.transpose(a, (1, 0)) # t13: "cuda:0 f32[3, 3]"
del a
t8 = torch.permute(t0, (1, 0)) # t8: "cuda:0 f32[3, 3]"
# t8 = ltorch.permute(t0, (1, 0)) # t8: "cuda:0 f32[3, 3]"
# t8 = prims.transpose(t0, (1, 0)) # t8: "cuda:0 f32[3, 3]"
del t0
t9 = torch.matmul(t2, t8) # t9: "cuda:0 f32[3, 3]"
# t9 = ltorch.matmul(t2, t8) # t9: "cuda:0 f32[3, 3]"
# t9 = prims.matmul(t2, t8) # t9: "cuda:0 f32[3, 3]"
t11 = torch.matmul(t8, t2) # t11: "cuda:0 f32[3, 3]"
# t11 = ltorch.matmul(t8, t2) # t11: "cuda:0 f32[3, 3]"
# t11 = prims.matmul(t8, t2) # t11: "cuda:0 f32[3, 3]"
del t8, t2
[t12] = nvFusion0(t11, t9)
# t12 = prims.add(t9, t11) # t12: "cuda:0 f32[3, 3]"
del t11, t9
t14 = torch.matmul(t12, t13) # t14: "cuda:0 f32[3, 3]"
# t14 = ltorch.matmul(t12, t13) # t14: "cuda:0 f32[3, 3]"
# t14 = prims.matmul(t12, t13) # t14: "cuda:0 f32[3, 3]"
t16 = torch.matmul(t13, t12) # t16: "cuda:0 f32[3, 3]"
# t16 = ltorch.matmul(t13, t12) # t16: "cuda:0 f32[3, 3]"
# t16 = prims.matmul(t13, t12) # t16: "cuda:0 f32[3, 3]"
del t13, t12
[t17] = nvFusion1(t14, t16)
# t17 = prims.add(t14, t16) # t17: "cuda:0 f32[3, 3]"
del t14, t16
return (t17,)
Here's documentation about reference count in Python:
the call mechanism guarantees to hold a reference to every argument for the duration of the call.
(from https://docs.python.org/3/c-api/intro.html#reference-counts) means that even if we delete a variable that was passed an argument inside our Python function a reference to the object still exists until the return statement. We can't free the memory of tensors passed as arguments until we exit the function.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.