reduce CPU peak memory in multimodal tensor hashing (#21123)
This commit is contained in:
@@ -1341,23 +1341,29 @@ def tensor_hash(tensor_list) -> int:
|
|||||||
tensor = tensor_list
|
tensor = tensor_list
|
||||||
if isinstance(tensor_list, list):
|
if isinstance(tensor_list, list):
|
||||||
tensor_list = flatten_nested_list(tensor_list)
|
tensor_list = flatten_nested_list(tensor_list)
|
||||||
tensor_list = [
|
tensors = [
|
||||||
x.flatten() if isinstance(x, torch.Tensor) else x for x in tensor_list
|
x.flatten() if isinstance(x, torch.Tensor) else x for x in tensor_list
|
||||||
]
|
]
|
||||||
tensor = torch.concat(tensor_list)
|
# GPU path: concat + triton hash (unchanged)
|
||||||
|
if any(isinstance(t, torch.Tensor) and t.is_cuda for t in tensors):
|
||||||
|
tensor = torch.concat(tensors)
|
||||||
|
return gpu_tensor_hash(tensor.cuda())
|
||||||
|
# CPU path: hash each tensor incrementally without concat
|
||||||
|
hasher = hashlib.sha256()
|
||||||
|
for t in tensors:
|
||||||
|
t = t.detach().contiguous()
|
||||||
|
hasher.update(memoryview(t.view(torch.uint8).numpy()))
|
||||||
|
hash_bytes = hasher.digest()[:8]
|
||||||
|
return int.from_bytes(hash_bytes, byteorder="big", signed=False)
|
||||||
|
|
||||||
|
# Single tensor
|
||||||
if tensor.is_cuda:
|
if tensor.is_cuda:
|
||||||
return gpu_tensor_hash(tensor.cuda())
|
return gpu_tensor_hash(tensor.cuda())
|
||||||
tensor = tensor.detach().contiguous()
|
tensor = tensor.detach().contiguous()
|
||||||
|
hasher = hashlib.sha256()
|
||||||
if tensor.dtype == torch.bfloat16:
|
hasher.update(memoryview(tensor.view(torch.uint8).numpy()))
|
||||||
# memoryview() doesn't support PyTorch's BFloat16 dtype
|
hash_bytes = hasher.digest()[:8]
|
||||||
tensor = tensor.float()
|
return int.from_bytes(hash_bytes, byteorder="big", signed=False)
|
||||||
|
|
||||||
assert isinstance(tensor, torch.Tensor)
|
|
||||||
tensor_cpu = tensor.cpu()
|
|
||||||
|
|
||||||
mv = memoryview(tensor_cpu.numpy())
|
|
||||||
return data_hash(mv.tobytes())
|
|
||||||
|
|
||||||
|
|
||||||
def hash_feature(f):
|
def hash_feature(f):
|
||||||
@@ -1367,8 +1373,10 @@ def hash_feature(f):
|
|||||||
return data_hash(tuple(flatten_nested_list(f)))
|
return data_hash(tuple(flatten_nested_list(f)))
|
||||||
elif isinstance(f, np.ndarray):
|
elif isinstance(f, np.ndarray):
|
||||||
arr = np.ascontiguousarray(f)
|
arr = np.ascontiguousarray(f)
|
||||||
arr_bytes = arr.tobytes()
|
hasher = hashlib.sha256()
|
||||||
return data_hash(arr_bytes)
|
hasher.update(memoryview(arr))
|
||||||
|
hash_bytes = hasher.digest()[:8]
|
||||||
|
return int.from_bytes(hash_bytes, byteorder="big", signed=False)
|
||||||
elif isinstance(f, torch.Tensor):
|
elif isinstance(f, torch.Tensor):
|
||||||
return tensor_hash([f])
|
return tensor_hash([f])
|
||||||
elif isinstance(f, CudaIpcTensorTransportProxy):
|
elif isinstance(f, CudaIpcTensorTransportProxy):
|
||||||
|
|||||||
Reference in New Issue
Block a user