Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Can鈥檛 save SparseTensor for custom device #121797

Open
huihoaan opened this issue Mar 13, 2024 · 2 comments 路 May be fixed by #126384
Open

Can鈥檛 save SparseTensor for custom device #121797

huihoaan opened this issue Mar 13, 2024 · 2 comments 路 May be fixed by #126384
Labels
module: sparse Related to torch.sparse triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module

Comments

@huihoaan
Copy link
Contributor

huihoaan commented Mar 13, 2024 •
edited by pytorch-bot bot

馃悰 Describe the bug

Our backend has support to create SparseTensor, but we find that SparseTensor in PrivateUse1 can not be saved directly, like:

import torch
a = torch.tensor([[0, 2.], [3, 0]]).to_sparse().to("privateuseone")
torch.save(a)

It will raise error TypeError: can't convert Sparse layout tensor to numpy. Use Tensor.dense() first.
I found that the following code was causing the problem:

pytorch/torch/_tensor.py

Lines 262 to 273 in e99fa00

if self.device.type in ["xla", "mtia", "ort"] or (
not torch._C._has_storage(self)
and self.device.type == torch._C._get_privateuse1_backend_name()
):
# Convert BFloat16 tesors to Float32 before conversion to numpy, as numpy doesn't
# support BFloat16. The rebuild tensor from numpy takes in the original self.dtype,
# this would reconstruct the BFloat16 tensor from numpy.
numpy_tensor = (
self.cpu().numpy()
if self.dtype != torch.bfloat16
else self.cpu().to(torch.float32).numpy()
)

SparseTensor doesn't have storage, so above code will be running. And self.cpu().numpy() will failed in

pytorch/torch/csrc/utils/tensor_numpy.cpp

Lines 133 to 138 in e99fa00

TORCH_CHECK_TYPE(
tensor.layout() == Layout::Strided,
"can't convert ",
c10::str(tensor.layout()).c_str(),
" layout tensor to numpy. ",
"Use Tensor.dense() first.");

I think here are two ways to solve this problem, maybe you have other advice?

pytorch/torch/_tensor.py

Lines 263 to 264 in e99fa00

not torch._C._has_storage(self)
and self.device.type == torch._C._get_privateuse1_backend_name()

  1. Delete the branch judgment in the _tensor.py and let PrivateUse1 goto the follow branchs like other devices
  2. Add is_sparse, is_neg, to the branch judgment for PrivateUse1, because numpy() cannot handle these tensors

Versions

torch/main

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer @jcaip
@drisspg drisspg added module: sparse Related to torch.sparse triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module labels Mar 13, 2024
@huihoaan
Copy link
Contributor Author

@albanD Do you know, why has this judgment for PrivateUse1. Shouldn't it be aligned with other devices, rather than as a special backend like ["xla", "mtia", "ort"].

pytorch/torch/_tensor.py

Lines 262 to 265 in e99fa00

if self.device.type in ["xla", "mtia", "ort"] or (
not torch._C._has_storage(self)
and self.device.type == torch._C._get_privateuse1_backend_name()
):

@albanD
Copy link
Collaborator

albanD commented Apr 2, 2024

I guess @heidongxianhua would know since they added it?

@huihoaan huihoaan linked a pull request May 16, 2024 that will close this issue
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
module: sparse Related to torch.sparse triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

make torch.save(PrivateUse1's Tensor) more generalizable huihoaan/pytorch
3 participants
@albanD @drisspg @huihoaan