Skip to content

use clone() method before save tensor when converting Hf model to Megatron model #25911

@flower-with-safe

Description

@flower-with-safe

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions