[Question] Concerns about gradient propagation through .cuda() operations in CPU→GPU transfers

[lejelly]

Hello!
Thank you for sharing your excellent research and implementation.
Your paper and code have been incredibly informative, and I'm learning a great deal about this novel approach to model merging using task vectors.

---

Context

Looking at the repository code, I notice that in the AdaMerging class, the weight coefficients (like alpha) are being composed with task difference parameters on CPU, then moved to GPU using params = tuple(p.cuda(0) for p in params) before the forward pass.

AdaMerging/src/main_task_wise_adamerging.py

Lines 120 to 126 in cea1165

	def forward(self, inp, dataset_name):
	alph = self.lambdas()
	params = tuple(sum(tuple(pi * lambdasi for pi, lambdasi in zip(p, alph[0].cpu()))) for j, p in enumerate(zip(*self.paramslist)))
	params = tuple(p.cuda(0) for p in params)
	load_weights(self.model, self.names, params)
	feature = self.model(inp)

Typically, in PyTorch's autograd mechanism, operations like .cpu() or .cuda() break the computation graph at the device transfer point, preventing gradient propagation.
The current flow appears to be:

1. Compose pretrained_model_parameter + Σ alpha * Taskvector_k on CPU
2. Move parameters back to GPU using .cuda()
3. Perform forward/backward passes
4. repeat 1~3

This sequence suggests that during backpropagation (.backward()), gradients might notflow back to alpha (or lambdas_raw).

---

Questions

In this implementation, it appears that gradients to the learnable alpha parameter might be blocked by the .cpu() → .cuda() operations. Am I missing something in my understanding?

If there's a specific mechanism or technique you've implemented to ensure gradient flow, I would greatly appreciate learning about it.

---

Thank you again for sharing your excellent research.
I would be very grateful for any insights you can provide when you have the time.

[EnnengYang]

Hi,

Thank you for your thoughtful question and for carefully analyzing the implementation!

I noticed that your core concern is whether device switching of variables in Pytorch will block the computation graph/gradient flow?

---

Will switching between .cpu() and .cuda() block Pytorch's computational graph?

From my knowledge, operations such as .cpu() or .gpu() in pytorch do not block the gradient of tensors in the computation graph.

An example is as follows:

Code:

import torch

# Define a tensor
lambdas_raw = torch.nn.Parameter(torch.tensor([[1.0, 2.0, 3.0]], requires_grad=True).cuda())
print("lambdas_raw.grad_fn:", lambdas_raw.grad_fn) # Leaf tensor, so its grad_fn attribute is None

# 1. Transfer lambdas_raw to CPU
lambdas_raw_cpu = lambdas_raw.cpu()
print("lambdas_raw_cpu.grad_fn:", lambdas_raw_cpu.grad_fn)
# 2. Perform simple operations on the CPU, such as squaring
result_cpu = lambdas_raw_cpu ** 2
print("result_cpu.grad_fn:", result_cpu.grad_fn)
# 3. Transfer the result back to the GPU
result_gpu = result_cpu.cuda()
print("result_gpu.grad_fn:", result_gpu.grad_fn)

# Define a simple loss: e.g., sum over all elements
def loss_func(x):
    return x.sum()
loss = loss_func(result_gpu)
print("loss.grad_fn:", loss.grad_fn)

# Back propagation, calculate the gradient
loss.backward()

# Output the gradient of lambdas_raw
print("lambdas_raw.grad:", lambdas_raw.grad)

Output:

lambdas_raw.grad_fn: None
lambdas_raw_cpu.grad_fn: <ToCopyBackward0 object at 0x7e9fe65078e0>
result_cpu.grad_fn: <PowBackward0 object at 0x7e9fe65078e0>
result_gpu.grad_fn: <ToCopyBackward0 object at 0x7e9fe65078e0>
loss.grad_fn: <SumBackward0 object at 0x7e9fe65078e0>
lambdas_raw.grad: tensor([[2., 4., 6.]], device='cuda:0')

We observed that:

• Although multiple device conversion operations such as .cpu() and .cuda() are performed on the lambdas_raw variable, we can still obtain the gradient of lambdas_raw (that is, the output of lambdas_raw.grad is not None).
• By directly accessing the grad_fn attributes of all intermediate tensors (e.g., lambdas_raw_cpu, result_cpu, result_gpu, loss), we observed that their grad_fn attributes are not None, but a backpropagation function object, such as <ToCopyBackward0 object at 0x7e9fe65078e0>.

This shows that executing .cpu() and .cuda() does not affect gradient calculation and update.

---

What operations will block Pytorch's computational graph?

In addition, from my knowledge, only the following three operations in Pytorch will block gradients/destroy computational graphs:

• Explicit .detach() operation: lambdas_raw=lambdas_raw.detach()
• Calculation with torch.no_grad(): with torch.no_grad(): lambdas_raw_ = lambdas_raw
• Extracting values ​​with .item(): lambdas_raw=lambdas_raw.item() will be directly converted to Python scalars, causing lambdas_raw to be unable to obtain gradients.

---

I hope these replies can solve your concerns.
Let me know if you have any further questions—I’d be happy to discuss!

[lejelly]

@EnnengYang

Thank you so much for your detailed and informative response! The sample code you provided was extremely helpful - I was able to reproduce the gradient flow behavior you described, which really cleared up my misunderstanding about PyTorch's device operations.

I sincerely appreciate you taking the time to explain both the specific implementation details and the broader context about which operations do and don't block PyTorch's computational graph. This has been very educational for me.

I'm currently working on extending your fascinating approach to LLM applications. I want to express my gratitude for making your research public and being so open to discussion - it's incredibly valuable for the research community.

Thank you again!