The true size of DTensor (Distributed Tensor in Tensor Parallel)

kaiyuyue: <aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/131342">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/131342" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/131342" target="_blank" rel="noopener nofollow ugc">size of DTensor</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2024-07-22" data-time="16:55:30" data-timezone="UTC">04:55PM - 22 Jul 24 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2024-07-22" data-time="17:17:03" data-timezone="UTC">05:17PM - 22 Jul 24 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/kaiyuyue" target="_blank" rel="noopener nofollow ugc">
          <img alt="kaiyuyue" src="https://avatars.githubusercontent.com/u/19852297?v=4" class="onebox-avatar-inline" width="20" height="20">
          kaiyuyue
        </a>
      </div>
    </div>

    <div class="labels">
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">```python
# torchrun --standalone --nnodes=1 --nproc-per-node=4 {THIS SCRIPT}.p<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">y

import os
import torch
import torch.distributed as dist
from torch.distributed._tensor import init_device_mesh, Shard, distribute_tensor

mesh = init_device_mesh("cuda", (int(os.environ["WORLD_SIZE"]),))
xx = torch.randn(100000, 88)
xd = distribute_tensor(xx, mesh, [Shard(dim=0)])

if dist.get_rank() == 0:
    print(torch.__version__)
print("rank:", dist.get_rank(), xx.shape, xd.shape, xd._local_tensor.shape)
```

The output is
```bash
2.3.1  # torch.__version__
rank: 0 torch.Size([100000, 88]) torch.Size([100000, 88]) torch.Size([25000, 88])
rank: 1 torch.Size([100000, 88]) torch.Size([100000, 88]) torch.Size([25000, 88])
rank: 2 torch.Size([100000, 88]) torch.Size([100000, 88]) torch.Size([25000, 88])
rank: 3 torch.Size([100000, 88]) torch.Size([100000, 88]) torch.Size([25000, 88])
```

I am not sure this is a bug or not. But after sharding the tensor `xx` to the DTensor `xd`, the naive torch.Size of `xd` is still the same as before. 

Now if I want to get the real size of `xd`, I need to get the `_local_tensor` out of DTensor Wrapper. 

If I miss sth. important, please correct me because I think that the `DTensor.shape` should be intuitive and directly output the size of `__local_tensor`. 

### Versions

```python
torch.__version__ = 2.3.1
```</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<pre><code class="lang-auto">print(x, x.shape, x.abs().sum())

DTensor(
    local_tensor=AsyncCollectiveTensor(tensor([[[ 0.1075, -2.3521,  0.0279, -1.4321,  1.8674,  0.0181, -1.0666,
           0.2343, -0.8745,  0.2822,  0.0908,  0.0743,  1.0759,  0.7753,
           1.4046, -1.2099, -0.4284,  0.3137, -0.4697, -0.8437,  1.3829,
          -0.5523,  1.3563,  0.7076, -0.4797,  0.8596,  0.0976, -0.7871,
           1.0103,  1.4723,  0.1845,  0.2991]]], device='cuda:1')), 
    device_mesh=DeviceMesh([0, 1], 
    mesh_dim_names=('tp',)), 
    placements=(Shard(dim=1),)
) 
torch.Size([1, 2, 32])  # shape still is the original one before sharding
DTensor(
    local_tensor=24.138418197631836, 
    device_mesh=DeviceMesh([0, 1], 
    mesh_dim_names=('tp',)), 
    placements=(_Partial(reduce_op=RedOpType.SUM),)
)
</code></pre>
<p>I have a question about the torch size of DTensor, <code>x</code> in the above example, which is sharded along with the row axis, from <code>[1, 2, 32]</code> to <code>[1, 1, 32]</code>. So the actual size should be <code>[1, 1, 32]</code>, can be verified by print it. But the output from <code>x.shape</code> is <code>[1, 2, 32]</code>. Can I know why? Thanks in advance.</p>

LilyAvo: <p>The true size of a DTensor can depend on several factors including its data type, the number of elements, and any additional metadata PyTorch stores. If you need precise memory usage, you might want to consider using the <code>element_size()</code> method and multiplying it by the number of elements in the tensor.</p>
<pre><code class="lang-auto">import torch

dtensor = torch.tensor([1.0, 2.0, 3.0])
size_in_bytes = dtensor.element_size() * dtensor.nelement()
print(f"Size of DTensor: {size_in_bytes} bytes")

</code></pre>
<p>Plus,  tools like torch.cuda.memory_allocated() can help measure memory usage if you’re working with CUDA tensors.</p>

kaiyuyue: <p>according to the source code of <code>DTensor</code>, we can get its size by getting its local_tensor out of wrapper.</p>
<pre data-code-wrap="python"><code class="lang-python">print(x)

DTensor(
    local_tensor=AsyncCollectiveTensor(tensor([[[ 0.1075, -2.3521,  0.0279, -1.4321,  1.8674,  0.0181, -1.0666,
           0.2343, -0.8745,  0.2822,  0.0908,  0.0743,  1.0759,  0.7753,
           1.4046, -1.2099, -0.4284,  0.3137, -0.4697, -0.8437,  1.3829,
          -0.5523,  1.3563,  0.7076, -0.4797,  0.8596,  0.0976, -0.7871,
           1.0103,  1.4723,  0.1845,  0.2991]]], device='cuda:1')), 
    device_mesh=DeviceMesh([0, 1], 
    mesh_dim_names=('tp',)), 
    placements=(Shard(dim=1),)
) 

print(x.shape)
torch.Size([1, 2, 32])  # this will output the original shape before sharding it

print(x._local_tensor.shape)
torch.Size([1, 1, 32])  # this gives the expected sharding shape
</code></pre>