v6d-io / v6d

vineyard (v6d): an in-memory immutable data manager. (Project under CNCF, TAG-Storage)
https://v6d.io
Apache License 2.0
829 stars 122 forks source link

Support bfloat16 type for pytorch module #1871

Closed dashanji closed 5 months ago

dashanji commented 5 months ago

Describe your problem

When loading the Qwen Model, get the following error.

  1. Download the safetensor file.
curl -LO https://huggingface.co/Qwen/Qwen-14B-Chat/resolve/main/model-00003-of-00015.safetensors
  1. Load the safetensor to vineyard.
import safetensors
import safetensors.torch

import vineyard
import vineyard.contrib.ml.torch as vineyard_torch

with open('model-00003-of-00015.safetensors', 'rb') as f:
  state_dict = safetensors.torch.load(f.read())

client = vineyard.connect("vineyard.sock")
with vineyard_torch.torch_context(client):
    object_id = client.put(state_dict)

with vineyard_torch.torch_context(client):
    client.get(object_id)

Then error happened.

Traceback (most recent call last):
  File "/opt/caoye/test1.py", line 12, in <module>
    object_id = client.put(state_dict)
  File "/opt/caoye/v6d/python/vineyard/core/client.py", line 696, in put
    return put(self, value, builder, persist, name, **kwargs)
  File "/opt/caoye/v6d/python/vineyard/core/builder.py", line 197, in put
    meta = get_current_builders().run(client, value, **kwargs)
  File "/opt/caoye/v6d/python/vineyard/core/builder.py", line 100, in run
    return self._factory[ty](client, value, **kw)
  File "/opt/caoye/v6d/python/vineyard/contrib/ml/torch.py", line 280, in torch_module_builder
    tensor_objects = put_torch_tensors(client, tensor_values)
  File "/opt/caoye/v6d/python/vineyard/contrib/ml/torch.py", line 219, in put_torch_tensors
    value = tensor.numpy()
TypeError: Got unsupported ScalarType BFloat16