Skip to content

Commit

Permalink
Revert "Feat graph load to new device (#10335)"
Browse files Browse the repository at this point in the history 
This reverts commit dea3f43.
  • Loading branch information
@jackalcooper
jackalcooper committed Dec 20, 2023
1 parent 21b51a6 commit 29b5c13
Show file tree
Hide file tree
Showing 6 changed files with 35 additions and 271 deletions.
4 changes: 2 additions & 2 deletions oneflow/core/graph/stream_id.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@ namespace oneflow {

// StreamId encoding (bits)
// | reserved | node_index | device_type | device_index | stream_index |
// | -- 18 -- | ----- 19 ----- | ---- 5 ---- | ----- 7 ----- | |
// | -- 21 -- | ----- 19 ----- | ---- 5 ---- | ----- 7 ----- | |
// | | DeviceId | |
// | | ------------------- 31 --------------------- | ---- 15 ---- |
// | | ------------------- 31 --------------------- | ---- 12 ---- |
// | StreamId |
// | -------------------------------- 64 ---------------------------------- |

Expand Down
4 changes: 2 additions & 2 deletions oneflow/core/graph/task_id.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@ namespace oneflow {

// TaskId encoding (maybe extended to 128 bits in future)
// | rank | device_type | device_index | |
// | ----------- 16 ----------- | ---- 5 ---- | ----- 7 ----- | |
// | ----------- 19 ----------- | ---- 5 ---- | ----- 7 ----- | |
// | DeviceId | stream_index | |
// | ------------------------- 31 --------------------------- | ---- 15 ---- | |
// | ------------------------- 31 --------------------------- | ---- 12 ---- | |
// | StreamId | task_index |
// | -------------------------------- 43 ----------------------------------- | --- 21 --- |
// | TaskId |
Expand Down
24 changes: 1 addition & 23 deletions python/oneflow/nn/graph/cache.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@

from oneflow.framework.args_tree import ArgsTree
from oneflow.framework.tensor import Tensor
import oneflow
import oneflow as flow


class LRUCache(object):
Expand Down Expand Up @@ -134,28 +134,6 @@ def runtime_state_dict(
destination[state_dict["graph_name"]] = state_dict
return destination

@staticmethod
def runtime_state_dict_to(
state_dict: Union[
Dict[str, Union[Dict[str, Tensor], str]],
Dict[str, Dict[str, Union[Dict[str, Tensor], str]]],
],
device: str,
) -> Union[
Dict[str, Union[Dict[str, Tensor], str]],
Dict[str, Dict[str, Union[Dict[str, Tensor], str]]],
]:
destination = OrderedDict()
destination._metadata = OrderedDict()
for (key, sub_state_dict) in state_dict.items():
dest_sub_state_dict = oneflow.nn.Graph.runtime_state_dict_to(
sub_state_dict, device
)
dest_sub_state_dict["cache_order"] = sub_state_dict["cache_order"]
dest_sub_state_dict["cache_key"] = sub_state_dict["cache_key"]
destination[key] = dest_sub_state_dict
return destination

def _init_and_get_a_graph_in_cache(self, cache_key):
self._base_graph._print(
0,
Expand Down
78 changes: 13 additions & 65 deletions python/oneflow/nn/graph/graph.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,9 +52,6 @@
GraphIR,
seq_to_func_return,
sys_exc_error_msg,
_rsd_sub_destination_to,
_job_to,
_plan_to,
)
from oneflow.framework.args_tree import ArgsTree
from oneflow.nn.modules.module import Module
Expand Down Expand Up @@ -1072,35 +1069,34 @@ def _fill_sub_destination(dest_dict, name_list, tensor_tuple):
assert len(tensor_tuple) == len(name_list)
for name_idx in range(len(name_list)):
tensor_item = tensor_tuple[name_idx]
device_str = ":".join(
(tensor_item.device.type, str(tensor_item.device.index))
)
dest_dict[name_list[name_idx]] = (tensor_item, device_str)
dest_dict[name_list[name_idx]] = (tensor_item, tensor_item.device.type)

# This is original outputs is needed to build output buffer.
tuple_idx = -1

def gen_index_in_tuple(item):
def gen_index_in_tuple(eager_out):
nonlocal tuple_idx
if isinstance(item, Tensor):
tuple_idx += 1
return "_OFTPI" + str(tuple_idx) # oneflow tuple index
else:
return item
tuple_idx += 1
return "_OFTPI" + str(tuple_idx) # oneflow tuple index

inputs_sub_destination = OrderedDict()
_fill_sub_destination(
inputs_sub_destination, self._input_op_names, self._inputs_tensor_tuple
)

_eager_inputs_args, _eager_inputs_kwargs = self.__map_io_lite(
gen_index_in_tuple, *self.inputs_original[0], **self.inputs_original[1],
_eager_inputs_args, _eager_inputs_kwargs = self.__map_io(
"input",
gen_index_in_tuple,
*self.inputs_original[0],
**self.inputs_original[1],
)
destination["inputs"] = inputs_sub_destination
destination["inputs_original"] = (_eager_inputs_args, _eager_inputs_kwargs)

tuple_idx = -1
_eager_outputs, _ = self.__map_io_lite(gen_index_in_tuple, *self._eager_outputs)
_eager_outputs, _ = self.__map_io(
"output", gen_index_in_tuple, *self._eager_outputs
)
destination["outputs_original"] = _eager_outputs
assert len(self._outputs_tensor_tuple) == tuple_idx + 1
outputs_sub_destination = OrderedDict()
Expand Down Expand Up @@ -1150,7 +1146,7 @@ def load_runtime_state_dict(
Dict[str, Dict[str, Union[Dict[str, Tensor], str]]],
],
*,
warmup_with_run: bool = True,
warmup_with_run: bool = False,
) -> None:
if self._run_with_cache == True:
return self._dynamic_input_graph_cache.load_runtime_state_dict(
Expand Down Expand Up @@ -1297,7 +1293,6 @@ def get_tensor_in_tuple(tensor_tuple, map_item):
self.__run(
*_eager_inputs_args, **_eager_inputs_kwargs
) # pre-run to warm up
oneflow._oneflow_internal.eager.Sync()
build_graph_end = time.perf_counter()
self.__print(
0,
Expand All @@ -1309,53 +1304,6 @@ def get_tensor_in_tuple(tensor_tuple, map_item):
+ "\n",
)

@staticmethod
def runtime_state_dict_to(
state_dict: Union[
Dict[str, Union[Dict[str, Tensor], str]],
Dict[str, Dict[str, Union[Dict[str, Tensor], str]]],
],
device: str,
) -> Union[
Dict[str, Union[Dict[str, Tensor], str]],
Dict[str, Dict[str, Union[Dict[str, Tensor], str]]],
]:
if "job_id" not in state_dict:
from oneflow.nn.graph.cache import GraphCache

return GraphCache.runtime_state_dict_to(state_dict, device)

dest_device = oneflow.device(device)
assert dest_device.type == "cuda", "device must be cuda."

destination = OrderedDict()
destination._metadata = OrderedDict()
destination["oneflow_version"] = state_dict["oneflow_version"]
destination["graph_name"] = state_dict["graph_name"]
destination["job_id"] = state_dict["job_id"]
destination["inputs"] = _rsd_sub_destination_to(state_dict["inputs"], device)
destination["inputs_original"] = state_dict["inputs_original"]
destination["outputs"] = _rsd_sub_destination_to(state_dict["outputs"], device)
destination["outputs_original"] = state_dict["outputs_original"]
destination["oneflow_with_eager_tensor"] = state_dict[
"oneflow_with_eager_tensor"
]
if "states" in state_dict:
destination["states"] = _rsd_sub_destination_to(
state_dict["states"], device
)
destination["exe_plan"] = _plan_to(state_dict["exe_plan"], dest_device)
if "forward_graph" in state_dict:
forward_graph = deepcopy(state_dict["forward_graph"])
_job_to(forward_graph, dest_device)
destination["forward_graph"] = forward_graph
if "compile_graph" in state_dict:
compile_graph = deepcopy(state_dict["compile_graph"])
_job_to(compile_graph, dest_device)
destination["compile_graph"] = compile_graph
destination["id_state"] = state_dict["id_state"]
return destination

def build_graph(self, *args, **kwargs):
# Build graph
try:
Expand Down
119 changes: 0 additions & 119 deletions python/oneflow/nn/graph/util.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,12 @@
import sys
from string import Template
from typing import Callable, Dict, Union, List, Tuple, Optional
from collections import OrderedDict

from google.protobuf import text_format
from google.protobuf.message import Message

import oneflow
import oneflow.core.job.job_pb2 as job_pb
import oneflow.core.job.plan_pb2 as plan_pb
import oneflow.core.common.device_type_pb2 as device_type
import oneflow.core.operator.op_conf_pb2 as op_conf_util
from oneflow.framework.tensor import Tensor

Expand Down Expand Up @@ -311,119 +308,3 @@ def seq_to_func_return(seq, need_unpack=False):
if need_unpack:
return seq[0]
return seq


def _rsd_sub_destination_to(origin_dict, dest_device_str):
dest_dict = OrderedDict()
for k, v in origin_dict.items():
tensor_item, device_str = v
dest_dict[k] = (
tensor_item.to(device=oneflow.device(dest_device_str), copy=True),
dest_device_str,
)
return dest_dict


def _parallel_conf_to(parallel_conf, dest_device):
if parallel_conf.device_tag == "cuda":
assert len(parallel_conf.device_name) == 1
parallel_conf.device_name[0] = "@0:" + str(dest_device.index)


def _mem_case_to(mem_case, dest_device):
if mem_case.device_type == device_type.DeviceType.kCUDA:
mem_case.device_id = dest_device.index
if (
mem_case.HasField("pinned_device_type")
and mem_case.pinned_device_type == device_type.DeviceType.kCUDA
):
mem_case.pinned_device_id = dest_device.index


def _job_to(job, dest_device):
for pg in job.placement.placement_group:
_parallel_conf_to(pg.parallel_conf, dest_device)
for bpg in job.placement.blob_placement_group:
_parallel_conf_to(bpg.parallel_conf, dest_device)


def _modify_bits(original_num, k, j, new_num):
if k > j:
return original_num
mask = ((1 << (j - k + 1)) - 1) << k
cleared_num = original_num & ~mask
modified_num = cleared_num | ((new_num & ((1 << (j - k + 1)) - 1)) << k)
return modified_num


def _get_bits(original_num, k, j):
mask = ((1 << (j - k + 1)) - 1) << k
cleared_num = (original_num & mask) >> k

return cleared_num


def _task_id_to(task_id, dest_device):
if _get_bits(task_id, 43, 48) == 2:
new_id = _modify_bits(task_id, 36, 43, dest_device.index)

return new_id
else:
return task_id


def _thrd_id_to(thrd_id, dest_device):
if _get_bits(thrd_id, 22, 27) == 2:
new_id = _modify_bits(thrd_id, 15, 22, dest_device.index)
return new_id
else:
return thrd_id


def _plan_to(plan_str, dest_device):
plan = plan_pb.Plan()
plan.ParseFromString(plan_str)
for task in plan.task:
task.task_id = _task_id_to(task.task_id, dest_device)
task.thrd_id = _thrd_id_to(task.thrd_id, dest_device)
for node in task.exec_sequence.exec_node:
_parallel_conf_to(
node.kernel_conf.op_attribute.parallel_conf_signature.op_parallel_conf,
dest_device,
)
for name, regst in task.produced_regst_desc.items():
regst.producer_task_id = _task_id_to(regst.producer_task_id, dest_device)
for c_task_id_idx in range(len(regst.consumer_task_id)):
regst.consumer_task_id[c_task_id_idx] = _task_id_to(
regst.consumer_task_id[c_task_id_idx], dest_device
)
_mem_case_to(regst.mem_case, dest_device)
for mem_block in plan.block_chunk_list.mem_block:
_mem_case_to(mem_block.mem_case, dest_device)
mem_block.thrd_id_hint = _thrd_id_to(mem_block.thrd_id_hint, dest_device)
for chunk in plan.block_chunk_list.chunk:
_mem_case_to(chunk.mem_case, dest_device)

new_ctrl_regst_desc_id2producer_task_id = {}
for (
regst_desc_id,
producer_task_id,
) in plan.ctrl_regst_desc_info.ctrl_regst_desc_id2producer_task_id.items():
new_ctrl_regst_desc_id2producer_task_id[regst_desc_id] = _task_id_to(
producer_task_id, dest_device
)
for (
regst_desc_id,
producer_task_id,
) in new_ctrl_regst_desc_id2producer_task_id.items():
plan.ctrl_regst_desc_info.ctrl_regst_desc_id2producer_task_id[
regst_desc_id
] = producer_task_id

for job_id, op_attr_tab in plan.job_id2op_attribute_ref_table.items():
for _, op_attr in op_attr_tab.op_name2op_attribute.items():
_parallel_conf_to(
op_attr.parallel_conf_signature.op_parallel_conf, dest_device
)

return plan.SerializeToString()
Loading

0 comments on commit 29b5c13

Please sign in to comment.