akankshakshurai/cuda-error-resolution-analysis huggingface dataset:
Use saved_tensors_hooks seem can't move the tensor from gpu to cpu

YeZhengMao: <p>I want to implement checkpoint use offload instead of recomputing to save GPU memory.<br>
The code is like this(using pytorch stable version 2.0.1):</p>
<pre><code class="lang-python">def pack_hook(x: torch.Tensor):
    device = x.device
    x = x.detach().cpu()
    return (device, x)
        
def unpack_hook(t):
    device, data = t
    return data.to(device)
    
for epoch in range(0, 100):
	with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
		output = model.forward(input)
	output.backward()
	torch.cuda.empty_cache()
</code></pre>
<p>When the first epoch was done, the GPU memory didn’t reduce. So when running the second epoch, raise CUDA OOM.</p>

ptrblck: <p>How did you check the memory usage? Did you compare the allocated, reserved, or both memory stats?</p>

YeZhengMao: <pre><code class="lang-python">import torch


def pack_hook(x: torch.Tensor):
    device = x.device
    x = x.detach().cpu()
    return (device, x)


def unpack_hook(t):
    device, data = t
    return data.to(device)


torch.cuda.reset_peak_memory_stats()
# about 2G
a = torch.ones((64, 4096, 4096), dtype=torch.float16,
               device="cuda:1", requires_grad=True) * 2
b = torch.ones((64, 4096, 4096), dtype=torch.float16,
               device="cuda:1", requires_grad=True) * 3


def fn():
    y = a * b
    return y


print(f"max allocated {torch.cuda.max_memory_allocated(1)/1024/1024/1024}G")
print(f"max reserved {torch.cuda.max_memory_reserved(1)/1024/1024/1024}G")

torch.cuda.reset_peak_memory_stats()
with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
    output = fn()
    print(f"max allocated {torch.cuda.max_memory_allocated(1)/1024/1024/1024}G")
    print(f"max reserved {torch.cuda.max_memory_reserved(1)/1024/1024/1024}G")
    torch.cuda.reset_peak_memory_stats()
    output.mean().backward()
    print(f"max allocated {torch.cuda.max_memory_allocated(1)/1024/1024/1024}G")
    print(f"max reserved {torch.cuda.max_memory_reserved(1)/1024/1024/1024}G")

print("ALL DONE.")
</code></pre>
<p>I used this code to test it, and below is the output:</p>
<pre><code class="lang-auto">max allocated 8.0G
max reserved 8.0G

max allocated 10.0G
max reserved 10.0G

max allocated 20.000000953674316G
max reserved 20.001953125G

ALL DONE.
</code></pre>
<p>and if I delete the code <code>with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):</code><br>
The output is:</p>
<pre><code class="lang-auto">max allocated 8.0G
max reserved 8.0G
max allocated 10.0G
max reserved 10.0G
max allocated 16.000000953674316G
max reserved 16.001953125G
ALL DONE.
</code></pre>
<p>seem the pack and unpack hook function will increase the memory usage?</p>
Contrastive sentence embeddings loss implementation

gonzrubio: <p>I’m looking to implement the supervised <a href="https://arxiv.org/abs/2104.08821" rel="noopener nofollow ugc">SimCSE</a> loss using positive and negative pairs.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/a/4a23244125a3f9bab931db0dd29ea478cd3679d8.png" alt="image" data-base62-sha1="azQCEi3QQ6yHJAdbpiPqQgB5nXa" width="498" height="132"></p>
<p>Is there a way to vectorize the naive implementation below?</p>
<pre><code class="lang-auto">import torch
import torch.nn.functional as F


batch_size = 4
feature_dim = 1024

h = torch.randn(batch_size, 3, feature_dim)

temp = 1

num = torch.exp(
    F.cosine_similarity(h[:, 0, :], h[:, 1, :], dim=1) / temp
    )

denom = torch.empty_like(num, device=num.device)
for j in range(batch_size):
    denomjj = 0
    for jj in range(batch_size):
        denomjj += torch.exp(F.cosine_similarity(h[j, 0, :], h[jj, 1, :], dim=0) / temp)
        denomjj += torch.exp(F.cosine_similarity(h[j, 0, :], h[jj, 2, :], dim=0) / temp)
    denom[j] = denomjj

loss = -torch.log(num / denom)
</code></pre>