System Info

when converting a huggingface model to megatron model, we need to split the model into TP * PP parts.
when TP is 1 and PP is X, each part of splited model occupies 1 / X storage space of the original model.
However, when TP is x and pp is 1, each part of the splited model occupies exactly the sample storage space as the original model.
This is due to when using torch.chunk() method to split tensor and then use torch.save method to save tensor, each saved tenor is a view of the original tensor and occupies the same storage space.

using the clone() method to clone tensor can solve this problem and save storage space.

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

import torch

a = torch.rand(1024, 512)
torch.save(a, "./full.pt")
print(a.size())
out = torch.chunk(a, 4, dim=0)

for i in range(4):
    print(out[i].size())
    torch.save(out[i], "./" + str(i) + "sub.pt")
    torch.save(out[i].clone(), "./" + str(i) + "sub_clone.pt")`

Expected behavior

the size of the pt file explains.