Data_ptr when moving across devices

username: zphang
cooked: <p>This question wraps together several question I have about how <code>data_ptr</code> is expected to work.</p>
<p>I discovered today that moving the same tensor to the same device twice leads to different <code>.data_ptr</code>s. This makes sense, since could be copying the tensors over separately.</p>
<pre><code class="lang-python">x = torch.tensor([1]).cuda()
x.cpu().data_ptr() == x.cpu().data_ptr()
# False
</code></pre>
<p>My issue is that this complicates matters when we expect the <code>data_ptr</code> to be the same. For instance, for models that share weights in different modules, such as word embeddings in input and output layers. In these cases, the input and output embeddings would have two entries in the model <code>state_dict</code>, but both would have the same <code>data_ptr</code>. However, if we use the typical recipe for moving <code>state_dict</code>s across devices:</p>
<pre><code class="lang-python">new_state_dict = {k: v.to(device) for k, v in state_dict.items()}
</code></pre>
<p>The input and output embeddings will now have different <code>data_ptr</code>s (in addition to taking up twice the memory).</p>
<p>My question is: is there a built-in method that handles this correctly, such that the moved dictionary of tensors correctly retains tensors with the same <code>data_ptr</code>s?</p>

username: ptrblck
cooked: <p>How did you apply the weight sharing? Could you share the code for it?<br>
If you just stored the parameter once and apply the functional API call using this parameter, there should be no problem I think.</p>

username: zphang
cooked: <p>Here’s an artificial example</p>
<pre><code class="lang-python">import torch
import torch.nn as nn

class ExampleModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(2, 2, bias=False)
        self.layer2 = nn.Linear(2, 2, bias=False)
        self.layer3 = nn.Linear(2, 2, bias=False)
        self.layer3 = self.layer1
        # Same result if you do
        #     self.layer3.weight = self.layer1.weight
    
# Moving model to device: works!
example_module = ExampleModule().cuda()
state_dict = example_module.cpu().state_dict()
assert state_dict["layer1.weight"].data_ptr() == state_dict["layer3.weight"].data_ptr()  
# passes

# Moving individual tensors: doesn't work
example_module = ExampleModule().cuda()
state_dict = {k: v.cpu() for k, v in example_module.state_dict().items()}
assert state_dict["layer1.weight"].data_ptr() == state_dict["layer3.weight"].data_ptr()  
# fails
</code></pre>
<p>In the first case, moving the whole model to a device correctly handles this. But there are other use-cases where you don’t want to move the model (e.g. if you’re just storing a checkpoint of the <code>state_dict</code>). In cases where you need to manually move the weights to a different device, the <code>data_ptr</code>s end up being different.</p>

username: ptrblck
cooked: <p>Thanks for the code example.<br>
I think the observed failure is due to the creation of a new <code>dict</code>, which will trigger a cope, wouldn’t it?</p>
<pre><code class="lang-python">a = {'a': 1, 'b': 2}
a_ref = a
a_copy = {k: v for k, v in a.items()}

a_copy['a'] = 3
print(a, a_ref, a_copy)
&gt; {'a': 1, 'b': 2} {'a': 1, 'b': 2} {'a': 3, 'b': 2}

a_ref['a'] = 4
print(a, a_ref, a_copy)
&gt; {'a': 4, 'b': 2} {'a': 4, 'b': 2} {'a': 3, 'b': 2}
</code></pre>
<p>In the first example, you are comparing the same internal <code>state_dict</code>, while the second code snippet created copies of all values.</p>

username: zphang
cooked: <p>I think this issue is somewhat unrelated to copying/references. In my examples, I am always comparing within a dictionary.</p>
<p>The question I’m asking is: what is a way to make a copy while still having resultant tensors that share the same data_ptrs (within the copy)?</p>
<p>For what it’s worth, I know that <code>copy.deepcopy</code> handles this correctly. However, I do not know of a solution that works when I need to move the <code>state_dict</code> across devices (other than moving the whole model to a different device first).</p>

username: zphang
cooked: <p>Bumping this. (Another way to frame this question is: What is the right and recommended way to move state_dicts across devices?)</p>