wehub-resource-sync
1284 lines
50 KiB
Python
1284 lines
50 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License
|
|
|
|
import copy
|
|
from collections import defaultdict
|
|
|
|
from paddle.distributed.passes.pass_base import PassContext
|
|
from paddle.framework import IrGraph, core, set_flags
|
|
|
|
from ..process_mesh import ProcessMesh
|
|
from .dist_op import DistributedOperator
|
|
from .dist_tensor import DistributedTensor
|
|
from .utils import (
|
|
__no_shape_var_type__,
|
|
_copy_dist_attr_to_cpp,
|
|
is_loss_grad_op,
|
|
)
|
|
|
|
# There always exists a default context for user. And user can set it to another one.
|
|
_g_default_distributed_context = None
|
|
|
|
|
|
def get_default_distributed_context():
|
|
global _g_default_distributed_context
|
|
if _g_default_distributed_context is None:
|
|
dist_context = DistributedContext()
|
|
set_default_distributed_context(dist_context)
|
|
return _g_default_distributed_context
|
|
|
|
|
|
def set_default_distributed_context(dist_context):
|
|
global _g_default_distributed_context
|
|
_g_default_distributed_context = dist_context
|
|
|
|
|
|
def _node_id(node):
|
|
return (node.node.graph_id(), node.node.id())
|
|
|
|
|
|
class DistributedContext:
|
|
"""
|
|
DistributedContext is used to collect related distributed information for program and graph.
|
|
One auto-parallel run should use its own DistributedContext to avoid interfering other run.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
serial_main_prog=None,
|
|
serial_startup_prog=None,
|
|
serial_optimizer=None,
|
|
serial_loss=None,
|
|
feed_vars={},
|
|
fetch_vars={},
|
|
cluster=None,
|
|
strategy=None,
|
|
json_config=None,
|
|
):
|
|
# Data members related to original programs (unchanged)
|
|
self._original_serial_main_program = serial_main_prog
|
|
self._original_serial_startup_program = serial_startup_prog
|
|
self._original_serial_optimizer = serial_optimizer
|
|
self._original_serial_loss = serial_loss
|
|
self._original_serial_feed_vars = feed_vars
|
|
self._original_serial_fetch_vars = fetch_vars
|
|
|
|
# Data members related to programs (changed)
|
|
self._serial_main_program = None
|
|
self._serial_startup_program = None
|
|
self._serial_loss = None
|
|
self._serial_optimizer = None
|
|
self._serial_feed_vars = {}
|
|
self._serial_fetch_vars = {}
|
|
self._lr_optimizer = None # record the optimizer holding lr_scheduler
|
|
|
|
# Data members related to the program
|
|
self._dist_tensors_for_program = {}
|
|
self._dist_ops_for_program = {}
|
|
|
|
# Data members related to the graph
|
|
self._serial_graph = None
|
|
self._dist_tensors_for_graph = {}
|
|
self._dist_ops_for_graph = {}
|
|
self._node_id_to_tensor_id = {}
|
|
self._node_id_to_op_id = {}
|
|
|
|
# Data members related to the distributed programs
|
|
# Distributed programs
|
|
self._dist_main_programs = {}
|
|
self._dist_startup_programs = {}
|
|
self._dist_op_context = DistributedOperatorContext()
|
|
self._process_meshes = []
|
|
|
|
self._cluster = cluster
|
|
self._strategy = strategy
|
|
|
|
# Pass Context
|
|
self._pass_context = PassContext()
|
|
self._block_state = BlockState()
|
|
|
|
# Other data members
|
|
self._serial_ordered_tensor_nodes = []
|
|
self._serial_ordered_op_nodes = []
|
|
self._serial_ordered_nodes = []
|
|
# self._tensor_id_to_tensor_node_ids = {}
|
|
|
|
self._is_initialized = False
|
|
# TODO: need a better way to remove the following flag
|
|
self._need_copy_dist_attr_to_graph = False
|
|
self._backup_pass_context_stack = []
|
|
self._backup_block_state_stack = []
|
|
self._backup_dist_tensors_for_program_stack = []
|
|
self._backup_dist_ops_for_program_stack = []
|
|
self._backup_serial_main_program_stack = []
|
|
self._backup_serial_startup_program_stack = []
|
|
|
|
# flag whether scale gradient with dp size
|
|
self._gradient_scale = True
|
|
|
|
# whether use allreduce_avg to scale gradient, i.e., allreduce_sum + scale -> allreduce_avg
|
|
self._gradient_scale_using_allreduce_avg = False
|
|
|
|
# A flag indicates whether the used parallelism is data parallel
|
|
self._data_parallel = False
|
|
|
|
# record upstream and downstream of cur rank
|
|
self._up_down_streams = UpDownStream()
|
|
|
|
self._json_config = json_config
|
|
|
|
# record vpp chunk size
|
|
self._num_model_chunks = 0
|
|
|
|
@property
|
|
def serial_main_program(self):
|
|
return self._serial_main_program
|
|
|
|
@property
|
|
def serial_startup_program(self):
|
|
return self._serial_startup_program
|
|
|
|
@property
|
|
def serial_loss(self):
|
|
return self._serial_loss
|
|
|
|
@property
|
|
def serial_optimizer(self):
|
|
return self._serial_optimizer
|
|
|
|
@property
|
|
def serial_feed_vars(self):
|
|
return self._serial_feed_vars
|
|
|
|
@property
|
|
def serial_fetch_vars(self):
|
|
return self._serial_fetch_vars
|
|
|
|
@property
|
|
def dist_main_programs(self):
|
|
return self._dist_main_programs
|
|
|
|
@property
|
|
def dist_startup_programs(self):
|
|
return self._dist_startup_programs
|
|
|
|
@property
|
|
def cluster(self):
|
|
return self._cluster
|
|
|
|
@property
|
|
def strategy(self):
|
|
return self._strategy
|
|
|
|
@property
|
|
def serial_graph(self):
|
|
return self._serial_graph
|
|
|
|
@property
|
|
def serial_ordered_nodes(self):
|
|
return self._serial_ordered_nodes
|
|
|
|
@property
|
|
def process_meshes(self):
|
|
return self._process_meshes
|
|
|
|
@process_meshes.setter
|
|
def process_meshes(self, val):
|
|
self._process_meshes = val
|
|
|
|
@property
|
|
def pass_context(self):
|
|
return self._pass_context
|
|
|
|
@property
|
|
def dist_op_context(self):
|
|
return self._dist_op_context
|
|
|
|
@property
|
|
def block_state(self):
|
|
return self._block_state
|
|
|
|
@property
|
|
def has_annotation(self):
|
|
return len(self._dist_tensors_for_program) or len(
|
|
self._dist_ops_for_program
|
|
)
|
|
|
|
@property
|
|
def gradient_scale(self):
|
|
return self._gradient_scale
|
|
|
|
@gradient_scale.setter
|
|
def gradient_scale(self, gs):
|
|
self._gradient_scale = gs
|
|
|
|
@property
|
|
def gradient_scale_using_allreduce_avg(self):
|
|
return self._gradient_scale_using_allreduce_avg
|
|
|
|
@gradient_scale_using_allreduce_avg.setter
|
|
def gradient_scale_using_allreduce_avg(
|
|
self, gradient_scale_using_allreduce_avg
|
|
):
|
|
self._gradient_scale_using_allreduce_avg = (
|
|
gradient_scale_using_allreduce_avg
|
|
)
|
|
|
|
@property
|
|
def data_parallel(self):
|
|
return self._data_parallel
|
|
|
|
@property
|
|
def up_down_streams(self):
|
|
return self._up_down_streams
|
|
|
|
@data_parallel.setter
|
|
def data_parallel(self, dp):
|
|
self._data_parallel = dp
|
|
|
|
def _backup_serial_info(self, mode):
|
|
self._backup_serial_main_program_stack.append(
|
|
self._serial_main_program.clone()
|
|
)
|
|
self._backup_serial_startup_program_stack.append(
|
|
self._serial_startup_program.clone()
|
|
)
|
|
self._backup_pass_context_stack.append(
|
|
copy.deepcopy(self._pass_context)
|
|
)
|
|
self._backup_block_state_stack.append(copy.deepcopy(self._block_state))
|
|
|
|
def _backup_dist_info(self, mode):
|
|
self._backup_dist_tensors_for_program_stack.append(
|
|
copy.deepcopy(self._dist_tensors_for_program)
|
|
)
|
|
self._backup_dist_ops_for_program_stack.append(
|
|
copy.deepcopy(self._dist_ops_for_program)
|
|
)
|
|
|
|
def _backup(self, serial=True, serial_mode=None, dist=True, dist_mode=None):
|
|
# Use this function carefully
|
|
if serial:
|
|
self._backup_serial_info(serial_mode)
|
|
if dist:
|
|
self._backup_dist_info(dist_mode)
|
|
|
|
def _restore_serial_loss(self):
|
|
if self._original_serial_loss:
|
|
if isinstance(self._original_serial_loss, list):
|
|
if len(self._original_serial_loss) == 1:
|
|
loss = self._original_serial_loss[0]
|
|
block_idx = loss.block.idx
|
|
var_name = loss.name
|
|
var = self._serial_main_program.blocks[
|
|
block_idx
|
|
]._var_recursive(var_name)
|
|
self._serial_loss = var
|
|
elif len(self._original_serial_loss) == 0:
|
|
self._serial_loss = []
|
|
else:
|
|
raise ValueError("multi loss vars are not supported.")
|
|
else:
|
|
block_idx = self._original_serial_loss.block.idx
|
|
var_name = self._original_serial_loss.name
|
|
var = self._serial_main_program.blocks[
|
|
block_idx
|
|
]._var_recursive(var_name)
|
|
self._serial_loss = var
|
|
|
|
def _restore_serial_feed_vars(self):
|
|
for key, var_list in self._original_serial_feed_vars.items():
|
|
new_var_list = []
|
|
for var in var_list:
|
|
block_idx = var.block.idx
|
|
var_name = var.name
|
|
var = self._serial_main_program.blocks[
|
|
block_idx
|
|
]._var_recursive(var_name)
|
|
new_var_list.append(var)
|
|
self._serial_feed_vars[key] = new_var_list
|
|
|
|
def _restore_serial_fetch_vars(self):
|
|
for key, var_list in self._original_serial_fetch_vars.items():
|
|
new_var_list = []
|
|
# metrics is a list of list
|
|
if key == "metrics":
|
|
for inner_var_list in var_list:
|
|
new_inner_var_list = []
|
|
for var in inner_var_list:
|
|
block_idx = var.block.idx
|
|
var_name = var.name
|
|
var = self._serial_main_program.blocks[
|
|
block_idx
|
|
]._var_recursive(var_name)
|
|
new_inner_var_list.append(var)
|
|
new_var_list.append(new_inner_var_list)
|
|
else:
|
|
for var in var_list:
|
|
block_idx = var.block.idx
|
|
var_name = var.name
|
|
var = self._serial_main_program.blocks[
|
|
block_idx
|
|
]._var_recursive(var_name)
|
|
new_var_list.append(var)
|
|
self._serial_fetch_vars[key] = new_var_list
|
|
|
|
def _restore_serial_info(self, mode="to_backup"):
|
|
if mode == "to_backup":
|
|
self._serial_main_program = (
|
|
self._backup_serial_main_program_stack.pop()
|
|
)
|
|
self._serial_startup_program = (
|
|
self._backup_serial_startup_program_stack.pop()
|
|
)
|
|
elif mode == "to_original":
|
|
assert self._original_serial_main_program is not None
|
|
assert self._original_serial_startup_program is not None
|
|
self._serial_main_program = (
|
|
self._original_serial_main_program.clone()
|
|
)
|
|
self._serial_startup_program = (
|
|
self._original_serial_startup_program.clone()
|
|
)
|
|
|
|
self._restore_serial_loss()
|
|
self._restore_serial_feed_vars()
|
|
self._restore_serial_fetch_vars()
|
|
self._serial_optimizer = self._original_serial_optimizer
|
|
self._pass_context = self._backup_pass_context_stack.pop()
|
|
self._block_state = self._backup_block_state_stack.pop()
|
|
|
|
def _restore_dist_info(self, mode="to_backup"):
|
|
if mode == "to_backup":
|
|
self._dist_tensors_for_program = (
|
|
self._backup_dist_tensors_for_program_stack.pop()
|
|
)
|
|
self._dist_ops_for_program = (
|
|
self._backup_dist_ops_for_program_stack.pop()
|
|
)
|
|
elif mode == "to_original":
|
|
assert self._original_dist_tensors_for_program
|
|
assert self._original_dist_ops_for_program
|
|
self._dist_tensors_for_program = copy.deepcopy(
|
|
self._original_dist_tensors_for_program
|
|
)
|
|
self._dist_ops_for_program = copy.deepcopy(
|
|
self._original_dist_ops_for_program
|
|
)
|
|
elif mode == "to_default":
|
|
new_tensors_ids = []
|
|
for (
|
|
tensor_id,
|
|
dist_tensor,
|
|
) in self._dist_tensors_for_program.items():
|
|
if tensor_id in self._tensors_ids:
|
|
dist_tensor.dist_attr.reset()
|
|
else:
|
|
new_tensors_ids.append(tensor_id)
|
|
for tensor_id in new_tensors_ids:
|
|
self._dist_tensors_for_program.pop(tensor_id)
|
|
new_ops_ids = []
|
|
for op_id, dist_op in self._dist_ops_for_program.items():
|
|
if op_id in self._ops_ids:
|
|
dist_op.dist_attr.reset()
|
|
else:
|
|
new_ops_ids.append(op_id)
|
|
for op_id in new_ops_ids:
|
|
self._dist_ops_for_program.pop(op_id)
|
|
else:
|
|
new_tensors_ids = []
|
|
for (
|
|
tensor_id,
|
|
dist_tensor,
|
|
) in self._dist_tensors_for_program.items():
|
|
new_tensors_ids.append(tensor_id)
|
|
for tensor_id in new_tensors_ids:
|
|
self._dist_tensors_for_program.pop(tensor_id)
|
|
new_ops_ids = []
|
|
for op_id, dist_op in self._dist_ops_for_program.items():
|
|
new_ops_ids.append(op_id)
|
|
for op_id in new_ops_ids:
|
|
self._dist_ops_for_program.pop(op_id)
|
|
self._dist_main_programs = {}
|
|
self._dist_startup_programs = {}
|
|
self._dist_op_context = DistributedOperatorContext()
|
|
self._need_copy_dist_attr_to_graph = True
|
|
self._process_meshes = []
|
|
|
|
def _restore(
|
|
self,
|
|
serial=True,
|
|
serial_mode="to_backup",
|
|
dist=True,
|
|
dist_mode="to_backup",
|
|
):
|
|
# Use this function carefully
|
|
if serial:
|
|
self._restore_serial_info(serial_mode)
|
|
if dist:
|
|
self._restore_dist_info(dist_mode)
|
|
|
|
def initialize(self, with_graph=True, with_cpp=False, no_default=False):
|
|
if not self._is_initialized:
|
|
if not self._serial_main_program:
|
|
if self._original_serial_main_program:
|
|
self._serial_main_program = (
|
|
self._original_serial_main_program.clone()
|
|
)
|
|
if not self._serial_startup_program:
|
|
if self._original_serial_startup_program:
|
|
self._serial_startup_program = (
|
|
self._original_serial_startup_program.clone()
|
|
)
|
|
if not self._serial_loss:
|
|
self._restore_serial_loss()
|
|
if not self._serial_optimizer:
|
|
self._serial_optimizer = self._original_serial_optimizer
|
|
if not self._serial_feed_vars:
|
|
self._restore_serial_feed_vars()
|
|
if not self._serial_fetch_vars:
|
|
self._restore_serial_fetch_vars()
|
|
|
|
self._init_dist_attr_for_program(no_default)
|
|
# Backup the original distributed information for later restore
|
|
self._original_dist_tensors_for_program = copy.deepcopy(
|
|
self._dist_tensors_for_program
|
|
)
|
|
self._original_dist_ops_for_program = copy.deepcopy(
|
|
self._dist_ops_for_program
|
|
)
|
|
self._tensors_ids = list(self._dist_tensors_for_program.keys())
|
|
self._ops_ids = list(self._dist_ops_for_program.keys())
|
|
self._is_initialized = True
|
|
|
|
# TODO: This will be removed in the future
|
|
if with_cpp:
|
|
_copy_dist_attr_to_cpp(self)
|
|
|
|
if with_graph:
|
|
set_flags({"FLAGS_convert_all_blocks": True})
|
|
self._serial_graph = IrGraph(
|
|
core.Graph(self._serial_main_program.desc)
|
|
)
|
|
self._init_dist_attr_for_graph()
|
|
self._need_copy_dist_attr_to_graph = False
|
|
|
|
if self._need_copy_dist_attr_to_graph and with_graph:
|
|
self.copy_dist_attr_from_program_to_graph()
|
|
|
|
def add_process_mesh(self, process_mesh):
|
|
assert isinstance(process_mesh, (ProcessMesh, core.ProcessMesh)), (
|
|
'The type of dim_mapping must be ProcessMesh.'
|
|
)
|
|
if process_mesh not in self.process_meshes:
|
|
self._process_meshes.append(process_mesh)
|
|
|
|
def add_dist_tensor_for_program(self, dist_tensor):
|
|
inner_serial_tensor = dist_tensor.serial_tensor
|
|
inner_serial_tensor_id = inner_serial_tensor.desc.original_id()
|
|
self._dist_tensors_for_program[inner_serial_tensor_id] = dist_tensor
|
|
|
|
def add_dist_op_for_program(self, dist_op):
|
|
inner_serial_op = dist_op.serial_op
|
|
inner_serial_op_id = inner_serial_op.desc.original_id()
|
|
self._dist_ops_for_program[inner_serial_op_id] = dist_op
|
|
|
|
def get_dist_tensor_for_program(self, serial_tensor):
|
|
serial_tensor_id = serial_tensor.desc.id()
|
|
dist_tensor = self._dist_tensors_for_program.get(serial_tensor_id, None)
|
|
if dist_tensor:
|
|
return dist_tensor
|
|
else:
|
|
serial_tensor_id = serial_tensor.desc.original_id()
|
|
dist_tensor = self._dist_tensors_for_program.get(
|
|
serial_tensor_id, None
|
|
)
|
|
if dist_tensor:
|
|
return dist_tensor
|
|
else:
|
|
return None
|
|
|
|
def get_dist_tensor_for_graph(self, serial_tensor_node):
|
|
serial_tensor_node_id = _node_id(serial_tensor_node)
|
|
return self._dist_tensors_for_graph.get(serial_tensor_node_id, None)
|
|
|
|
def get_dist_op_for_program(self, serial_op):
|
|
serial_op_id = serial_op.desc.id()
|
|
dist_op = self._dist_ops_for_program.get(serial_op_id, None)
|
|
if dist_op:
|
|
return dist_op
|
|
else:
|
|
serial_op_id = serial_op.desc.original_id()
|
|
dist_op = self._dist_ops_for_program.get(serial_op_id, None)
|
|
if dist_op:
|
|
return dist_op
|
|
else:
|
|
return None
|
|
|
|
def del_dist_op_for_program(self, serial_tensor):
|
|
serial_tensor_id = serial_tensor.desc.id()
|
|
if self._dist_ops_for_program.get(serial_tensor_id, None):
|
|
del self._dist_ops_for_program[serial_tensor_id]
|
|
|
|
def get_dist_op_for_graph(self, serial_op_node):
|
|
serial_op_node_id = _node_id(serial_op_node)
|
|
return self._dist_ops_for_graph.get(serial_op_node_id, None)
|
|
|
|
def get_tensor_dist_attr_for_program(self, serial_tensor):
|
|
serial_tensor_id = serial_tensor.desc.id()
|
|
dist_tensor = self._dist_tensors_for_program.get(serial_tensor_id, None)
|
|
if dist_tensor:
|
|
return dist_tensor.dist_attr
|
|
else:
|
|
serial_tensor_id = serial_tensor.desc.original_id()
|
|
dist_tensor = self._dist_tensors_for_program.get(
|
|
serial_tensor_id, None
|
|
)
|
|
if dist_tensor:
|
|
return dist_tensor.dist_attr
|
|
else:
|
|
return None
|
|
|
|
def get_tensor_dist_attr_for_program_with_id(self, tensor_id):
|
|
dist_tensor = self._dist_tensors_for_program.get(tensor_id, None)
|
|
if dist_tensor:
|
|
return dist_tensor.dist_attr
|
|
else:
|
|
return None
|
|
|
|
def set_tensor_dist_attr_for_program(self, serial_tensor, dist_attr):
|
|
dist_tensor = DistributedTensor(serial_tensor, dist_attr)
|
|
self.add_dist_tensor_for_program(dist_tensor)
|
|
|
|
def get_tensor_dist_attr_for_graph(self, serial_tensor_node):
|
|
serial_tensor_node_id = _node_id(serial_tensor_node)
|
|
dist_tensor = self._dist_tensors_for_graph.get(
|
|
serial_tensor_node_id, None
|
|
)
|
|
if dist_tensor:
|
|
return dist_tensor.dist_attr
|
|
else:
|
|
return None
|
|
|
|
def get_op_dist_attr_for_program(self, serial_op):
|
|
serial_op_id = serial_op.desc.id()
|
|
dist_op = self._dist_ops_for_program.get(serial_op_id, None)
|
|
if dist_op:
|
|
return dist_op.dist_attr
|
|
else:
|
|
serial_op_id = serial_op.desc.original_id()
|
|
dist_op = self._dist_ops_for_program.get(serial_op_id, None)
|
|
if dist_op:
|
|
return dist_op.dist_attr
|
|
else:
|
|
return None
|
|
|
|
def get_op_dist_attr_for_program_with_id(self, op_id):
|
|
dist_op = self._dist_ops_for_program.get(op_id, None)
|
|
if dist_op:
|
|
return dist_op.dist_attr
|
|
else:
|
|
return None
|
|
|
|
def set_op_dist_attr_for_program(self, serial_op, dist_attr):
|
|
dist_op = DistributedOperator(serial_op, dist_attr)
|
|
self.add_dist_op_for_program(dist_op)
|
|
|
|
def get_op_dist_attr_for_graph(self, serial_op_node):
|
|
serial_op_node_id = _node_id(serial_op_node)
|
|
dist_op = self._dist_ops_for_graph.get(serial_op_node_id, None)
|
|
if dist_op:
|
|
return dist_op.dist_attr
|
|
else:
|
|
return None
|
|
|
|
def get_dist_attr_for_graph(self, serial_node):
|
|
if serial_node.is_var() and serial_node.var() is not None:
|
|
serial_tensor_node_id = _node_id(serial_node)
|
|
dist_tensor = self._dist_tensors_for_graph.get(
|
|
serial_tensor_node_id, None
|
|
)
|
|
if dist_tensor:
|
|
return dist_tensor.dist_attr
|
|
else:
|
|
return None
|
|
if serial_node.is_op() and serial_node.op() is not None:
|
|
serial_op_node_id = _node_id(serial_node)
|
|
dist_op = self._dist_ops_for_graph.get(serial_op_node_id, None)
|
|
if dist_op:
|
|
return dist_op.dist_attr
|
|
else:
|
|
return None
|
|
return None
|
|
|
|
def _init_dist_attr_for_program(self, no_default=False):
|
|
# Copy the dist tensors and dist ops annotated by users from the default context
|
|
if not no_default:
|
|
default_ctx = get_default_distributed_context()
|
|
self._process_meshes = copy.deepcopy(default_ctx.process_meshes)
|
|
else:
|
|
default_ctx = self
|
|
# Copy the data parallel flag from the default context
|
|
self._data_parallel = default_ctx.data_parallel
|
|
for block in self._serial_main_program.blocks:
|
|
for tensor in block.vars.values():
|
|
# Copy the distributed tensors in the default context
|
|
default_dist_tensor = default_ctx.get_dist_tensor_for_program(
|
|
tensor
|
|
)
|
|
if default_dist_tensor and default_ctx is not self:
|
|
dist_tensor = DistributedTensor(tensor)
|
|
dist_tensor.dist_attr = copy.deepcopy(
|
|
default_dist_tensor.dist_attr
|
|
)
|
|
self.add_dist_tensor_for_program(dist_tensor)
|
|
current_dist_tensor = self.get_dist_tensor_for_program(tensor)
|
|
if current_dist_tensor is None:
|
|
dist_tensor = DistributedTensor(tensor)
|
|
self.add_dist_tensor_for_program(dist_tensor)
|
|
for op in block.ops:
|
|
# Copy the distributed operators in the default context
|
|
default_dist_op = default_ctx.get_dist_op_for_program(op)
|
|
if default_dist_op and default_ctx is not self:
|
|
dist_op = DistributedOperator(op)
|
|
dist_op.dist_attr = copy.deepcopy(default_dist_op.dist_attr)
|
|
self.add_dist_op_for_program(dist_op)
|
|
current_dist_op = self.get_dist_op_for_program(op)
|
|
if current_dist_op is None:
|
|
dist_op = DistributedOperator(op)
|
|
self.add_dist_op_for_program(dist_op)
|
|
self._original_dist_tensors_for_program = copy.deepcopy(
|
|
self._dist_tensors_for_program
|
|
)
|
|
self._original_dist_ops_for_program = copy.deepcopy(
|
|
self._dist_ops_for_program
|
|
)
|
|
|
|
def _order_nodes_by_program_order(self):
|
|
serial_ordered_tensor_nodes = []
|
|
serial_ordered_op_nodes = []
|
|
all_nodes = []
|
|
visited = {}
|
|
for idx, graph in enumerate(self._serial_graph.all_sub_graphs()):
|
|
for node in graph.all_nodes():
|
|
all_nodes.append(node)
|
|
for node in all_nodes:
|
|
if node.is_var() and node.var() is not None:
|
|
serial_ordered_tensor_nodes.append(node)
|
|
visited[_node_id(node)] = False
|
|
if node.is_op() and node.op() is not None:
|
|
serial_ordered_op_nodes.append(node)
|
|
serial_ordered_tensor_nodes.sort(
|
|
key=lambda node: node.node.original_desc_id()
|
|
)
|
|
serial_ordered_op_nodes.sort(
|
|
key=lambda node: node.node.original_desc_id()
|
|
)
|
|
num_nodes_before = len(serial_ordered_tensor_nodes) + len(
|
|
serial_ordered_op_nodes
|
|
)
|
|
|
|
new_serial_ordered_tensor_nodes = []
|
|
new_serial_ordered_op_nodes = []
|
|
new_serial_ordered_nodes = []
|
|
for op_node in serial_ordered_op_nodes:
|
|
tensor_nodes = []
|
|
for tensor_node in op_node.inputs:
|
|
if (
|
|
tensor_node.is_var()
|
|
and tensor_node.var() is not None
|
|
and not visited[_node_id(tensor_node)]
|
|
):
|
|
tensor_nodes.append(tensor_node)
|
|
new_serial_ordered_tensor_nodes.append(tensor_node)
|
|
visited[_node_id(tensor_node)] = True
|
|
|
|
tensor_nodes.sort(key=lambda node: node.node.original_desc_id())
|
|
new_serial_ordered_nodes.extend(tensor_nodes)
|
|
new_serial_ordered_nodes.append(op_node)
|
|
new_serial_ordered_op_nodes.append(op_node)
|
|
tensor_nodes = []
|
|
for tensor_node in op_node.outputs:
|
|
if (
|
|
tensor_node.is_var()
|
|
and tensor_node.var() is not None
|
|
and not visited[_node_id(tensor_node)]
|
|
):
|
|
tensor_nodes.append(tensor_node)
|
|
new_serial_ordered_tensor_nodes.append(tensor_node)
|
|
visited[_node_id(tensor_node)] = True
|
|
tensor_nodes.sort(key=lambda node: node.node.original_desc_id())
|
|
new_serial_ordered_nodes.extend(tensor_nodes)
|
|
new_serial_ordered_tensor_nodes.sort(
|
|
key=lambda node: node.node.original_desc_id()
|
|
)
|
|
new_serial_ordered_op_nodes.sort(
|
|
key=lambda node: node.node.original_desc_id()
|
|
)
|
|
self._serial_ordered_tensor_nodes = new_serial_ordered_tensor_nodes
|
|
self._serial_ordered_op_nodes = new_serial_ordered_op_nodes
|
|
self._serial_ordered_nodes = new_serial_ordered_nodes
|
|
assert len(self._serial_ordered_nodes) == len(
|
|
self._serial_ordered_tensor_nodes
|
|
) + len(self._serial_ordered_op_nodes)
|
|
# graph_id -> tensor->name -> node_lists
|
|
self._tensor_nodes_with_same_name = defaultdict(dict)
|
|
for idx, node in enumerate(self._serial_ordered_nodes):
|
|
if node.is_var() and node.var() is not None:
|
|
graph_id = node.node.graph_id()
|
|
tensor_name = node.var().name()
|
|
if (
|
|
self._tensor_nodes_with_same_name[graph_id].get(
|
|
tensor_name, None
|
|
)
|
|
is None
|
|
):
|
|
self._tensor_nodes_with_same_name[graph_id][
|
|
tensor_name
|
|
] = []
|
|
self._tensor_nodes_with_same_name[graph_id][tensor_name].append(
|
|
(idx, node)
|
|
)
|
|
|
|
self._serial_orphan_tensor_nodes = []
|
|
for tensor_node in serial_ordered_tensor_nodes:
|
|
if not visited[_node_id(tensor_node)]:
|
|
self._serial_orphan_tensor_nodes.append(tensor_node)
|
|
if len(self._serial_ordered_nodes) != num_nodes_before:
|
|
print(
|
|
"WARNING: there are some orphan tensors or ops which are not used in the execution."
|
|
)
|
|
|
|
def _init_dist_attr_for_graph(self):
|
|
# Convert program to graph and initialize the distributed attributes
|
|
self._order_nodes_by_program_order()
|
|
self._tensor_original_id_to_id = {}
|
|
self._op_original_id_to_id = {}
|
|
for tensor_id, tensor in self._dist_tensors_for_program.items():
|
|
original_id = tensor.serial_tensor.desc.original_id()
|
|
self._tensor_original_id_to_id[original_id] = tensor_id
|
|
for op_id, op in self._dist_ops_for_program.items():
|
|
original_id = op.serial_op.desc.original_id()
|
|
self._op_original_id_to_id[original_id] = op_id
|
|
for node in self.serial_ordered_nodes:
|
|
if node.is_var() and node.var() is not None:
|
|
dist_tensor = None
|
|
tensor_id = node.node.original_desc_id()
|
|
cur_dist_tensor = self._dist_tensors_for_program.get(
|
|
tensor_id, None
|
|
)
|
|
if cur_dist_tensor is not None:
|
|
cur_tensor_id = tensor_id
|
|
else:
|
|
cur_tensor_id = self._tensor_original_id_to_id[tensor_id]
|
|
cur_dist_tensor = self._dist_tensors_for_program.get(
|
|
cur_tensor_id, None
|
|
)
|
|
dist_tensor = cur_dist_tensor
|
|
self._node_id_to_tensor_id[_node_id(node)] = cur_tensor_id
|
|
assert dist_tensor is not None, (
|
|
"Tensor must have a distributed tensor after the initialization for program."
|
|
)
|
|
serial_tensor_node_id = _node_id(node)
|
|
new_dist_tensor = DistributedTensor(
|
|
dist_tensor.serial_tensor, dist_tensor.dist_attr
|
|
)
|
|
self._dist_tensors_for_graph[serial_tensor_node_id] = (
|
|
new_dist_tensor
|
|
)
|
|
if node.is_op() and node.op() is not None:
|
|
dist_op = None
|
|
op_id = node.node.original_desc_id()
|
|
cur_dist_op = self._dist_ops_for_program.get(op_id, None)
|
|
if cur_dist_op is not None:
|
|
cur_op_id = op_id
|
|
else:
|
|
cur_op_id = self._op_original_id_to_id[op_id]
|
|
cur_dist_op = self._dist_ops_for_program.get(
|
|
cur_op_id, None
|
|
)
|
|
dist_op = cur_dist_op
|
|
self._node_id_to_op_id[_node_id(node)] = cur_op_id
|
|
assert dist_op is not None, (
|
|
"Operator must have a distributed operator after the initialization for program."
|
|
)
|
|
serial_op_node_id = _node_id(node)
|
|
new_dist_op = DistributedOperator(
|
|
dist_op.serial_op, dist_op.dist_attr
|
|
)
|
|
self._dist_ops_for_graph[serial_op_node_id] = new_dist_op
|
|
|
|
def clear_dist_info_for_program(self):
|
|
self._dist_tensors_for_program.clear()
|
|
self._dist_ops_for_program.clear()
|
|
|
|
def clear_dist_info_for_graph(self):
|
|
self._dist_tensors_for_graph.clear()
|
|
self._dist_ops_for_graph.clear()
|
|
|
|
def copy_dist_attr_from_program_to_graph(self):
|
|
for node in self.serial_ordered_nodes:
|
|
if node.is_var() and node.var() is not None:
|
|
dist_tensor = None
|
|
tensor_id = node.node.original_desc_id()
|
|
cur_dist_tensor = self._dist_tensors_for_program.get(
|
|
tensor_id, None
|
|
)
|
|
if cur_dist_tensor is not None:
|
|
cur_tensor_id = tensor_id
|
|
else:
|
|
cur_tensor_id = self._tensor_original_id_to_id[tensor_id]
|
|
cur_dist_tensor = self._dist_tensors_for_program.get(
|
|
cur_tensor_id, None
|
|
)
|
|
dist_tensor = cur_dist_tensor
|
|
assert dist_tensor is not None, (
|
|
"Tensor must have a distributed tensor after the initialization for program."
|
|
)
|
|
serial_tensor_node_id = _node_id(node)
|
|
new_dist_tensor = DistributedTensor(
|
|
dist_tensor.serial_tensor, dist_tensor.dist_attr
|
|
)
|
|
self._dist_tensors_for_graph[serial_tensor_node_id] = (
|
|
new_dist_tensor
|
|
)
|
|
if node.is_op() and node.op() is not None:
|
|
dist_op = None
|
|
op_id = node.node.original_desc_id()
|
|
cur_dist_op = self._dist_ops_for_program.get(op_id, None)
|
|
if cur_dist_op is not None:
|
|
cur_op_id = op_id
|
|
else:
|
|
cur_op_id = self._op_original_id_to_id[op_id]
|
|
cur_dist_op = self._dist_ops_for_program.get(
|
|
cur_op_id, None
|
|
)
|
|
dist_op = cur_dist_op
|
|
assert dist_op is not None, (
|
|
"Operator must have a distributed operator after the initialization for program."
|
|
)
|
|
serial_op_node_id = _node_id(node)
|
|
new_dist_op = DistributedOperator(
|
|
dist_op.serial_op, dist_op.dist_attr
|
|
)
|
|
self._dist_ops_for_graph[serial_op_node_id] = new_dist_op
|
|
|
|
def copy_dist_attr_from_graph_to_program(self):
|
|
assert self._is_initialized, (
|
|
"Both program and graph must be initialized."
|
|
)
|
|
updated_tensors = {}
|
|
all_nodes = self._serial_ordered_nodes
|
|
process_meshes = [self.process_meshes[0]]
|
|
for node in all_nodes:
|
|
if node.is_var() and node.var() is not None:
|
|
tensor_id = self._node_id_to_tensor_id[_node_id(node)]
|
|
updated = updated_tensors.get(tensor_id, False)
|
|
# If a var has multiples var nodes in graph, only use the first one for now
|
|
if not updated:
|
|
tensor_dist_attr_for_graph = (
|
|
self.get_tensor_dist_attr_for_graph(node)
|
|
)
|
|
dist_tensor_for_program = self._dist_tensors_for_program[
|
|
tensor_id
|
|
]
|
|
dist_tensor_for_program.dist_attr = (
|
|
tensor_dist_attr_for_graph
|
|
)
|
|
updated_tensors[tensor_id] = True
|
|
process_mesh = tensor_dist_attr_for_graph.process_mesh
|
|
if process_mesh not in process_meshes:
|
|
process_meshes.append(process_mesh)
|
|
if node.is_op() and node.op() is not None:
|
|
op_id = self._node_id_to_op_id[_node_id(node)]
|
|
op_dist_attr_for_graph = self.get_op_dist_attr_for_graph(node)
|
|
dist_op_for_program = self._dist_ops_for_program[op_id]
|
|
dist_op_for_program.dist_attr = op_dist_attr_for_graph
|
|
process_mesh = op_dist_attr_for_graph.process_mesh
|
|
if process_mesh not in process_meshes:
|
|
process_meshes.append(process_mesh)
|
|
# NOTE(zhaoyingli):
|
|
# The order of process_meshes is execution order of the ops,
|
|
# which will help pipeline strategy to get pp_rank info.
|
|
self.process_meshes = copy.deepcopy(process_meshes)
|
|
# TODO: the completion algorithm will skipped orphan tensors,
|
|
# here we just set there process_mesh to the first one.
|
|
for orphan_node in self._serial_orphan_tensor_nodes:
|
|
serial_tensor_id = orphan_node.var().id()
|
|
dist_tensor = self._dist_tensors_for_program.get(
|
|
serial_tensor_id, None
|
|
)
|
|
if not dist_tensor:
|
|
serial_tensor_id = orphan_node.var().original_id()
|
|
dist_tensor = self._dist_tensors_for_program.get(
|
|
serial_tensor_id, None
|
|
)
|
|
dist_tensor.dist_attr.process_mesh = self.process_meshes[0]
|
|
|
|
def amend_dist_attr_for_program(self):
|
|
for dist_tensor in self._dist_tensors_for_program.values():
|
|
serial_tensor = dist_tensor.serial_tensor
|
|
dist_attr = dist_tensor.dist_attr
|
|
if serial_tensor.type in __no_shape_var_type__:
|
|
tensor_shape = []
|
|
else:
|
|
tensor_shape = serial_tensor.shape
|
|
dims_mapping = dist_attr.dims_mapping
|
|
process_mesh_shape = dist_attr.process_mesh.shape
|
|
process_mesh_processes = dist_attr.process_mesh.process_ids
|
|
# If the dimension of tensor is less than the sharding dimension of process mesh,
|
|
# we just amend the dimension mapping to -1. (Is this really OK?)
|
|
for i in range(len(tensor_shape)):
|
|
if (
|
|
dims_mapping[i] != -1
|
|
and tensor_shape[i] > 0
|
|
and process_mesh_shape[dims_mapping[i]] > tensor_shape[i]
|
|
):
|
|
dims_mapping[i] = -1
|
|
if dims_mapping[i] != -1 and len(process_mesh_processes) == 1:
|
|
dims_mapping[i] = -1
|
|
dist_attr.dims_mapping = dims_mapping
|
|
|
|
for dist_op in self._dist_ops_for_program.values():
|
|
serial_op = dist_op.serial_op
|
|
dist_attr = dist_op.dist_attr
|
|
process_mesh_shape = dist_attr.process_mesh.shape
|
|
process_mesh_processes = dist_attr.process_mesh.process_ids
|
|
for arg_name in serial_op.input_arg_names:
|
|
if dist_op.get_serial_input(arg_name) is None:
|
|
tensor_shape = []
|
|
else:
|
|
if (
|
|
dist_op.get_serial_input(arg_name).type
|
|
in __no_shape_var_type__
|
|
):
|
|
tensor_shape = []
|
|
else:
|
|
tensor_shape = dist_op.get_serial_input(arg_name).shape
|
|
dims_mapping = dist_attr.get_input_dims_mapping(arg_name)
|
|
# If the dimension of tensor is less than the sharding dimension of process mesh,
|
|
# we just amend the dimension mapping to -1. (Is this really OK?)
|
|
for i in range(len(tensor_shape)):
|
|
if (
|
|
dims_mapping[i] != -1
|
|
and tensor_shape[i] > 0
|
|
and process_mesh_shape[dims_mapping[i]]
|
|
> tensor_shape[i]
|
|
):
|
|
dims_mapping[i] = -1
|
|
if (
|
|
dims_mapping[i] != -1
|
|
and len(process_mesh_processes) == 1
|
|
):
|
|
dims_mapping[i] = -1
|
|
dist_attr.set_input_dims_mapping(arg_name, dims_mapping)
|
|
for arg_name in serial_op.output_arg_names:
|
|
if (
|
|
dist_op.get_serial_output(arg_name).type
|
|
in __no_shape_var_type__
|
|
):
|
|
tensor_shape = []
|
|
else:
|
|
tensor_shape = dist_op.get_serial_output(arg_name).shape
|
|
dims_mapping = dist_attr.get_output_dims_mapping(arg_name)
|
|
# If the dimension of tensor is less than the sharding dimension of process mesh,
|
|
# we just amend the dimension mapping to -1. (Is this really OK?)
|
|
for i in range(len(tensor_shape)):
|
|
if (
|
|
dims_mapping[i] != -1
|
|
and tensor_shape[i] > 0
|
|
and process_mesh_shape[dims_mapping[i]]
|
|
> tensor_shape[i]
|
|
):
|
|
dims_mapping[i] = -1
|
|
if (
|
|
dims_mapping[i] != -1
|
|
and len(process_mesh_processes) == 1
|
|
):
|
|
dims_mapping[i] = -1
|
|
dist_attr.set_output_dims_mapping(arg_name, dims_mapping)
|
|
if (
|
|
len(process_mesh_processes) == 1
|
|
and dist_op.serial_op.type != "dropout"
|
|
):
|
|
dist_op.dist_attr.impl_type = "default"
|
|
dist_op.dist_attr.impl_idx = 0
|
|
|
|
def validate_dist_attr_for_program(self):
|
|
if not self._is_initialized:
|
|
raise AssertionError(
|
|
"Program must be initialized before validating its distributed attributes"
|
|
)
|
|
for block in self.serial_main_program.blocks:
|
|
for tensor in block.vars.values():
|
|
dist_tensor = self.get_dist_tensor_for_program(tensor)
|
|
assert dist_tensor is not None, (
|
|
f"Tensor {dist_tensor.serial_tensor.name} does not have a distributed attribute."
|
|
)
|
|
if (dist_tensor is not None) and (
|
|
not dist_tensor.validate_dist_attr()
|
|
):
|
|
raise AssertionError(
|
|
f"Tensor {dist_tensor.serial_tensor.name} (id: {dist_tensor.serial_tensor.desc.id()}, original_id: {dist_tensor.serial_tensor.desc.original_id()}) has a wrong distributed attributes {dist_tensor.dist_attr}."
|
|
)
|
|
for op in block.ops:
|
|
dist_op = self.get_dist_op_for_program(op)
|
|
assert dist_op is not None, (
|
|
f"Operator {dist_op.serial_op.type} does not have a distributed attribute."
|
|
)
|
|
if (dist_op is not None) and (not dist_op.validate_dist_attr()):
|
|
raise AssertionError(
|
|
f"Operator {dist_op.serial_op.type} (id: {dist_op.serial_op.desc.id()}, original_id: {dist_op.serial_op.desc.original_id()}) has a wrong distributed attributes {dist_op.dist_attr} ."
|
|
)
|
|
if (
|
|
op.has_attr("op_namescope")
|
|
and 'auto_parallel/rc_' in op.attr("op_namescope")
|
|
and not self.strategy.recompute.enable
|
|
):
|
|
self.strategy.recompute.enable = True
|
|
return True
|
|
|
|
def __deepcopy__(self, memo):
|
|
cls = self.__class__
|
|
result = cls.__new__(cls)
|
|
memo[id(self)] = result
|
|
for k, v in self.__dict__.items():
|
|
if k in [
|
|
"_original_serial_main_program",
|
|
"_original_serial_startup_program",
|
|
"_serial_main_program",
|
|
"_serial_startup_program",
|
|
"_serial_graph",
|
|
"_dist_main_programs",
|
|
"_dist_startup_programs",
|
|
"_serial_ordered_nodes",
|
|
"_serial_ordered_tensor_nodes",
|
|
"_serial_ordered_op_nodes",
|
|
"_original_serial_loss",
|
|
"_original_serial_feed_vars",
|
|
"_original_serial_fetch_vars",
|
|
"_serial_loss",
|
|
"_serial_feed_vars",
|
|
"_serial_fetch_vars",
|
|
"_serial_optimizer",
|
|
"_backup_serial_main_program_stack",
|
|
"_backup_serial_startup_program_stack",
|
|
"_pass_context",
|
|
"_tensor_nodes_with_same_name",
|
|
]:
|
|
setattr(result, k, v)
|
|
else:
|
|
setattr(result, k, copy.deepcopy(v, memo))
|
|
|
|
# update dist tensor's dist_context
|
|
for key in result._dist_tensors_for_program.keys():
|
|
result._dist_tensors_for_program[key]._dist_context = result
|
|
return result
|
|
|
|
|
|
class DistributedOperatorContext:
|
|
"""
|
|
DistributedOperatorContext is used to create a dist op desc in Program.
|
|
Every time to create a new dist op, the context should be updated for it accordingly.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._dst_main_program = None
|
|
self._main_block = None
|
|
self._dst_startup_program = None
|
|
self._startup_block = None
|
|
self._cur_src_op = None
|
|
self._cur_dist_attr = None
|
|
self.grad_op_id_to_op_id = {}
|
|
self.grad_var_to_var = defaultdict(dict)
|
|
self._work_block = None
|
|
self.already_init_sync_vars = set()
|
|
self.varname_mapping = None
|
|
self.rank_id = None
|
|
# NOTE Support correct parallelism for high-order differential model.
|
|
# by default exceed_backward_init_op is False and it means we are in Forward phase; After exceed_backward_init_op = True,
|
|
# it means we are in Backward phase.
|
|
# And the final solution should be revise high-order differential logic for these two phases in future.
|
|
self._exceed_backward_init_op = False
|
|
|
|
def __deepcopy__(self, memo):
|
|
cls = self.__class__
|
|
result = cls.__new__(cls)
|
|
memo[id(self)] = result
|
|
for k, v in self.__dict__.items():
|
|
if k in [
|
|
"_dst_main_program",
|
|
"_dst_startup_program",
|
|
"_cur_src_op",
|
|
"_work_block",
|
|
"_main_block",
|
|
"_startup_block",
|
|
]:
|
|
setattr(result, k, v)
|
|
else:
|
|
setattr(result, k, copy.deepcopy(v, memo))
|
|
return result
|
|
|
|
@property
|
|
def dst_main_program(self):
|
|
return self._dst_main_program
|
|
|
|
@dst_main_program.setter
|
|
def dst_main_program(self, prog):
|
|
self._dst_main_program = prog
|
|
self._main_block = prog.blocks[0]
|
|
|
|
@property
|
|
def main_block(self):
|
|
return self._main_block
|
|
|
|
@property
|
|
def dst_startup_program(self):
|
|
return self._dst_startup_program
|
|
|
|
@dst_startup_program.setter
|
|
def dst_startup_program(self, prog):
|
|
self._dst_startup_program = prog
|
|
self._startup_block = prog.blocks[0]
|
|
|
|
@property
|
|
def startup_block(self):
|
|
return self._startup_block
|
|
|
|
@property
|
|
def work_block(self):
|
|
assert self._work_block is not None
|
|
return self._work_block
|
|
|
|
@work_block.setter
|
|
def work_block(self, block):
|
|
assert block is not None
|
|
self._work_block = block
|
|
|
|
@property
|
|
def cur_src_op(self):
|
|
assert self._cur_src_op is not None
|
|
return self._cur_src_op
|
|
|
|
def in_backward_phase(self):
|
|
return self._exceed_backward_init_op
|
|
|
|
def prepare_context(self, src_op):
|
|
self._cur_src_op = src_op
|
|
|
|
if is_loss_grad_op(src_op):
|
|
self._exceed_backward_init_op = True
|
|
|
|
# build input varname mapping
|
|
kinputs = {}
|
|
for input_name in src_op.desc.input_names():
|
|
varnames = []
|
|
for varname in src_op.desc.input(input_name):
|
|
assert varname in self.varname_mapping[src_op.block.idx]
|
|
varnames.append(self.varname_mapping[src_op.block.idx][varname])
|
|
kinputs[input_name] = varnames
|
|
|
|
# build output varname mapping
|
|
koutputs = {}
|
|
for output_name in src_op.desc.output_names():
|
|
varnames = []
|
|
for varname in src_op.desc.output(output_name):
|
|
assert varname in self.varname_mapping[src_op.block.idx]
|
|
varnames.append(self.varname_mapping[src_op.block.idx][varname])
|
|
koutputs[output_name] = varnames
|
|
|
|
return kinputs, koutputs
|
|
|
|
|
|
class BlockState:
|
|
def __init__(self):
|
|
self.nblock = 0
|
|
self.forward_indices = []
|
|
self.backward_indices = []
|
|
self.backward_to_forward_index_map = {}
|
|
|
|
def parse_forward_blocks(self, program):
|
|
program._roll_to_global_block()
|
|
assert program.current_block_idx == 0
|
|
|
|
for idx, block in enumerate(program.blocks):
|
|
assert idx == block.idx, "index doesn't match"
|
|
assert block.forward_block_idx == -1, (
|
|
f"forward_block_idx of forward block [{idx}] is not [{block.forward_block_idx}]"
|
|
)
|
|
self.forward_indices.append(idx)
|
|
self.nblock += 1
|
|
|
|
assert self.nblock >= 1
|
|
|
|
def parse_backward_blocks(self, program):
|
|
assert 0 in self.forward_indices, (
|
|
f"forward block idx are{self.forward_indices}"
|
|
)
|
|
self.backward_to_forward_index_map[0] = 0
|
|
|
|
for idx, block in enumerate(program.blocks):
|
|
if idx < len(self.forward_indices):
|
|
continue
|
|
|
|
assert idx == block.idx, "index doesn't match"
|
|
assert block.forward_block_idx in self.forward_indices
|
|
self.backward_indices.append(idx)
|
|
self.backward_to_forward_index_map[idx] = block.forward_block_idx
|
|
self.nblock += 1
|
|
|
|
assert self.nblock == len(program.blocks)
|
|
|
|
|
|
class UpDownStream:
|
|
def __init__(self):
|
|
self._ups = {}
|
|
self._downs = {}
|
|
|
|
def add_up_stream(self, rank, up_stream):
|
|
ups = self._ups.get(rank, None)
|
|
if not ups:
|
|
self._ups[rank] = [up_stream]
|
|
elif up_stream != -1:
|
|
ups = list(filter(lambda a: a != -1, ups))
|
|
ups.append(up_stream)
|
|
self._ups[rank] = ups
|
|
|
|
def add_down_stream(self, rank, down_stream):
|
|
downs = self._downs.get(rank, None)
|
|
if not downs:
|
|
self._downs[rank] = [down_stream]
|
|
elif down_stream != -1:
|
|
downs = list(filter(lambda a: a != -1, downs))
|
|
downs.append(down_stream)
|
|
self._downs[rank] = downs
|
|
|
|
def add_pair_stream(self, up, down):
|
|
self.add_up_stream(up, -1)
|
|
self.add_up_stream(down, up)
|
|
self.add_down_stream(up, down)
|
|
self.add_down_stream(down, -1)
|
|
|
|
def ups(self, rank):
|
|
ups = self._ups.get(rank, None)
|
|
if not ups:
|
|
return None
|
|
return list(set(ups))
|
|
|
|
def downs(self, rank):
|
|
downs = self._downs.get(rank, None)
|
|
if not downs:
|
|
return None
|
|
return list(set(downs))
|