Skip to content

Commit

Permalink
Allocate Tensor memory based on Tensor storage id
Browse files Browse the repository at this point in the history 
Summary: Allocate Tensor memory based on Tensor storage id

Differential Revision: D58474695
  • Loading branch information
@shengfukevin @facebook-github-bot
shengfukevin authored and facebook-github-bot committed Jul 19, 2024
1 parent 7c820e5 commit c8d45e0
Show file tree
Hide file tree
Showing 2 changed files with 85 additions and 15 deletions.
5 changes: 5 additions & 0 deletions et_replay/et_replay_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,11 @@ def is_fbgemm(op):

# TODO: Hopefully merge is_fbgemm and skip_op, ignore tid == 2
def skip_op(op):
if (
op.name == "aten::_scaled_dot_product_efficient_attention_backward"
or op.name == "aten::_efficient_attention_backward"
):
return True
# Workaround: skip bounds check indices and other ops under embedding lookup module
return (
not is_fbgemm_forward(op)
Expand Down
95 changes: 80 additions & 15 deletions et_replay/tools/et_replay.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from collections import defaultdict
from datetime import datetime
from functools import reduce
from typing import Dict

import numpy as np
import torch
Expand Down Expand Up @@ -102,6 +103,9 @@ def __init__(self):
# Dict that stores the shapes of a tensor, for the convenience of quickly determining whether
# to create a unique tensor in replay if the id is same but shape is different.
self.tensor_shapes = defaultdict(set)
# Dict that maps tensor storage id to its size, and a map for {device, torch.Tensor}.
# The tensor with the same storage id may located on different devices.
self.tensor_storage_map: Dict[int, []] = defaultdict(set)
# Mark those tensors that occur first as an input in the original et as needing to be instantiated in replay
# at the very beginning.
self.instantiate = set()
Expand Down Expand Up @@ -285,7 +289,7 @@ def reset_registry(self):
None
if v is None
else (
v
v.cpu()
if self.tensor_device[k] == "cpu" or self.cpu
else v.cuda(self.device)
)
Expand All @@ -298,7 +302,9 @@ def reset_registry(self):
None
if v is None
else (
v if k in self.cpu_tensor or self.cpu else v.cuda(self.device)
v.cpu()
if k in self.cpu_tensor or self.cpu
else v.cuda(self.device)
)
)
for k, v in self.tensor_registry_permanent.items()
Expand Down Expand Up @@ -435,7 +441,27 @@ def bfs_traverse(node):
)

def analyze_tensors(self):
def add_storage_tensor(t_id, device):
# t_id is a tupe of (tensor_id, storage_id, offset, number of element,
# number of bytes for each element, device)

# ET does not save the size of the tensor storage, so we iterate over all the
# tensors to find the maximum size of the storage.
storage_id = t_id[1]
if storage_id not in self.tensor_storage_map:
# the storage size for this tensor is the sum of the storage offset and
# the number of elements.
self.tensor_storage_map[storage_id] = [
t_id[2] + t_id[3] * t_id[4],
{},
]
else:
max_size = self.tensor_storage_map[storage_id][0]
max_size = max(max_size, t_id[2] + t_id[3] * t_id[4])
self.tensor_storage_map[storage_id][0] = max_size

def add_unique_tensor(node_name, node_id, t_id, shape, input, device=-1):
add_storage_tensor(t_id, device)
# If we did not see this tensor before, add it as a unique tensor.
if t_id not in self.original_unique_tensors:
self.original_unique_tensors.add(t_id)
Expand Down Expand Up @@ -534,6 +560,8 @@ def add_unique_tensor(node_name, node_id, t_id, shape, input, device=-1):
output_set.add(self.tensors_mapping[(node.id, t_id, False)])

def allocate_tensors(self):
start_ns = time.time_ns()

for node in self.sorted_nodes:
if node.name == "record_param_comms" and (
self.compute_only or self.args.separate
Expand All @@ -557,7 +585,9 @@ def allocate_tensors(self):
self.args.alpha,
)
for idx, (data_type, t_id, shape) in enumerate(get_input_tensors(node)):
device = self.device
if self.tensor_with_device:
device = t_id[5]
t_id = tuple(list(t_id)[:5])
replay_t_id = self.tensors_mapping[(node.id, t_id, True)]
if (
Expand All @@ -578,14 +608,21 @@ def allocate_tensors(self):
self.unchangeable_intermediate_tensors.add(replay_t_id)
else:
if data_type == "Tensor(signed char)":
dtype, rng = TORCH_DTYPES_RNG["signed char"]
dtype, _ = TORCH_DTYPES_RNG["signed char"]
else:
dtype, rng = TORCH_DTYPES_RNG[
dtype, _ = TORCH_DTYPES_RNG[
data_type.lstrip("Tensor(").rstrip(")")
]
self.tensor_registry_permanent[replay_t_id] = rng(shape).to(
dtype

tensor = self.get_tensor_from_storage(
t_id[1], # storage_id
t_id[2], # offset
t_id[4], # number of bytes per element
device,
shape,
dtype,
)
self.tensor_registry_permanent[replay_t_id] = tensor
if node.name == "aten::embedding_bag":
self.unchangeable_intermediate_tensors.add(replay_t_id)
if node.name == "aten::pin_memory" and idx == 0:
Expand Down Expand Up @@ -615,6 +652,8 @@ def allocate_tensors(self):
][i] = (i * nnz)
######

print(f"Tensor allocation time: {(time.time_ns() - start_ns) / 1000000.0} ms")

def build_func(self, node):
if is_fbgemm_forward(node):
if self.cpu:
Expand Down Expand Up @@ -1013,6 +1052,36 @@ def _generate_run_ops_str(override):
print(code_str, file=f)
exec(code_str)

def get_tensor_from_storage(
self, storage_id, data_offset, elem_bytes, device, shape, data_type
):
assert storage_id in self.tensor_storage_map

tensor_data = self.tensor_storage_map[storage_id]
device = torch.device(device)
if device not in tensor_data[1]:
if data_type in [torch.half, torch.float32, torch.float64, torch.bfloat16]:
storage_tensor = torch.rand(
(tensor_data[0] // elem_bytes), dtype=data_type, device=device
)
else:
storage_tensor = torch.ones(
(tensor_data[0] // elem_bytes), dtype=data_type, device=device
)
tensor_data[1][device] = storage_tensor
else:
storage_tensor = tensor_data[1][device]
x = torch.empty(0, dtype=data_type)
if device != torch.device("cpu"):
x = x.cuda(torch.device(device))
x = x.set_(
storage_tensor.untyped_storage(),
storage_offset=data_offset,
size=shape,
)

return x

def get_inputs(self, node):
try:
if is_fbgemm_forward(node):
Expand Down Expand Up @@ -1165,13 +1234,6 @@ def run_op(self, node, iter):
for (_, t_id, _), output in zip(get_output_tensors(node), outputs):
if self.tensor_with_device:
t_id = tuple(list(t_id)[:5])
# if output.isnan().any():
# print(
# node.id,
# t_id,
# output,
# inputs,
# )
if (
t_id in self.dependency_permanent
and self.tensors_mapping[(node.id, t_id, False)]
Expand Down Expand Up @@ -1283,8 +1345,7 @@ def benchTime(self):
self.preprocess_graph()
if self.generator:
return
print("Start to execution: ")
time.sleep(2)
print("Start execution: ")

total_time = 0.0
event_1 = torch.cuda.Event(enable_timing=True)
Expand All @@ -1299,6 +1360,8 @@ def benchTime(self):
qps_print_interval = 10

prev_iter = self.numWarmupIters

start_mem = torch.cuda.memory_allocated(self.device)
if self.profile_replay:
try:
from aiplatform.monitoring.unitrace.upload_manifold import (
Expand Down Expand Up @@ -1400,6 +1463,8 @@ def benchTime(self):
# self.reset_registry()
benchmark_result["execution finished"] = True
print("Execution finished!")
end_mem = torch.cuda.memory_allocated(self.device)
print(f"GPU Memory allocated: {(end_mem - start_mem) / 1024.0 / 1024.0} MB")

if self.profile_memory:
print("Allocated GPU memory(B):")
Expand Down

0 comments on commit c8d45e0

Please sign in to comment.