-
Couldn't load subscription status.
- Fork 31k
Description
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
examplesfolder (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.