mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-08 12:30:42 +00:00
[autoparallel] update CommSpec (#1667)
This commit is contained in:
YuliangLiu0306
GitHub
@@ -11,355 +11,9 @@ import math
|
||||
from functools import reduce
|
||||
import operator
|
||||
from torch.distributed import ReduceOp
|
||||
from .comm_spec import *
|
||||
|
||||
__all__ = [
|
||||
'CollectiveCommPattern', 'CommSpec', 'ShapeConsistencyManager', 'ShapeConsistencyOptions',
|
||||
'set_shape_consistency_options'
|
||||
]
|
||||
|
||||
|
||||
def _all_gather(tensor, comm_spec):
|
||||
'''
|
||||
Implement all gather operation on device mesh based on information provided by comm_spec.
|
||||
'''
|
||||
process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis]
|
||||
for rank_list, process_group in process_groups_list:
|
||||
if dist.get_rank() in rank_list:
|
||||
tensor_list = [
|
||||
torch.zeros(tensor.shape, dtype=tensor.dtype, device=tensor.device)
|
||||
for _ in range(comm_spec.device_mesh.mesh_shape[comm_spec.logical_process_axis])
|
||||
]
|
||||
tensor = tensor
|
||||
group = process_group
|
||||
dist.all_gather(tensor_list, tensor, group=group)
|
||||
output = torch.cat(tuple(tensor_list), comm_spec.gather_dim).contiguous()
|
||||
return output
|
||||
|
||||
|
||||
def _split(tensor, comm_spec):
|
||||
'''
|
||||
Implement shard operation on device mesh based on information provided by comm_spec.
|
||||
'''
|
||||
process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis]
|
||||
for rank_list, _ in process_groups_list:
|
||||
if dist.get_rank() in rank_list:
|
||||
tensor = tensor
|
||||
dim = comm_spec.shard_dim
|
||||
length = tensor.shape[comm_spec.shard_dim] // len(rank_list)
|
||||
start = length * rank_list.index(dist.get_rank())
|
||||
output = torch.narrow(tensor, dim, start, length)
|
||||
return output
|
||||
|
||||
|
||||
def _all_to_all(tensor, comm_spec):
|
||||
'''
|
||||
Implement all to all operation on device mesh based on information provided by comm_spec.
|
||||
'''
|
||||
process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis]
|
||||
for rank_list, process_group in process_groups_list:
|
||||
if dist.get_rank() in rank_list:
|
||||
new_shape = list(tensor.shape)
|
||||
new_shape[comm_spec.shard_dim] = new_shape[comm_spec.shard_dim] // len(rank_list)
|
||||
new_shape = torch.Size(new_shape)
|
||||
output_tensor_list = [
|
||||
torch.zeros(new_shape, dtype=tensor.dtype, device=tensor.device) for _ in range(len(rank_list))
|
||||
]
|
||||
dim = comm_spec.shard_dim
|
||||
length = tensor.shape[comm_spec.shard_dim] // len(rank_list)
|
||||
input_tensor_list = [
|
||||
torch.narrow(tensor, dim, length * i, length).contiguous() for i in range(len(rank_list))
|
||||
]
|
||||
group = process_group
|
||||
dist.all_to_all(output_tensor_list, input_tensor_list, group)
|
||||
output = torch.cat(tuple(output_tensor_list), comm_spec.gather_dim).contiguous()
|
||||
return output
|
||||
|
||||
|
||||
def _all_reduce(tensor, comm_spec):
|
||||
'''
|
||||
Implement all reduce operation on device mesh based on information provided by comm_spec.
|
||||
'''
|
||||
process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis]
|
||||
for rank_list, process_group in process_groups_list:
|
||||
if dist.get_rank() in rank_list:
|
||||
dist.all_reduce(tensor, op=ReduceOp.SUM, group=process_group)
|
||||
return tensor
|
||||
|
||||
|
||||
class _ReduceGrad(torch.autograd.Function):
|
||||
"""
|
||||
A customized communication operation which forward is an identity operation,
|
||||
backward is all_reduce operation.
|
||||
|
||||
Args:
|
||||
input_: input matrix.
|
||||
comm_spec: comm_spec will give information like process group, rank list, etc.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def symbolic(graph, input_):
|
||||
return input_
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, input_, comm_spec):
|
||||
ctx.comm_spec = comm_spec
|
||||
return input_
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
return _all_reduce(grad_output, ctx.comm_spec), None
|
||||
|
||||
|
||||
class _ReduceInput(torch.autograd.Function):
|
||||
"""
|
||||
A customized communication operation which forward is all_reduce operation,
|
||||
backward is an identity operation.
|
||||
|
||||
Args:
|
||||
input_: input matrix.
|
||||
comm_spec: comm_spec will give information like process group, rank list, etc.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def symbolic(graph, input_):
|
||||
return _all_reduce(input_)
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, input_, comm_spec):
|
||||
return _all_reduce(input_, comm_spec)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
return grad_output, None
|
||||
|
||||
|
||||
class _SplitForwardGatherBackward(torch.autograd.Function):
|
||||
"""
|
||||
A customized communication operation which forward is split operation,
|
||||
backward is an all gather operation.
|
||||
|
||||
Args:
|
||||
input_: input matrix.
|
||||
comm_spec: comm_spec will give information like process group, rank list, etc.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def symbolic(graph, input_):
|
||||
return _split(input_)
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, input_, comm_spec):
|
||||
ctx.comm_spec = comm_spec
|
||||
return _split(input_, comm_spec)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
return _all_gather(grad_output, ctx.comm_spec), None
|
||||
|
||||
|
||||
class _GatherForwardSplitBackward(torch.autograd.Function):
|
||||
"""
|
||||
A customized communication operation which forward is an all gather operation,
|
||||
backward is split operation.
|
||||
|
||||
Args:
|
||||
input_: input matrix.
|
||||
comm_spec: comm_spec will give information like process group, rank list, etc.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def symbolic(graph, input_):
|
||||
return _all_gather(input_)
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, input_, comm_spec):
|
||||
ctx.comm_spec = comm_spec
|
||||
return _all_gather(input_, comm_spec)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
return _split(grad_output, ctx.comm_spec), None
|
||||
|
||||
|
||||
class _AllToAll(torch.autograd.Function):
|
||||
"""
|
||||
A customized communication operation which forward is an all to all operation,
|
||||
backward is an all to all operation.
|
||||
|
||||
Args:
|
||||
input_: input matrix.
|
||||
comm_spec: comm_spec will give information like process group, rank list, etc.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def symbolic(graph, input_):
|
||||
return _all_to_all(input_)
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, input_, comm_spec):
|
||||
output = _all_to_all(input_, comm_spec)
|
||||
comm_spec_for_backward = CommSpec(comm_pattern=comm_spec.comm_pattern,
|
||||
sharding_spec=comm_spec.sharding_spec,
|
||||
gather_dim=comm_spec.shard_dim,
|
||||
shard_dim=comm_spec.gather_dim,
|
||||
logical_process_axis=comm_spec.logical_process_axis)
|
||||
ctx.comm_spec = comm_spec_for_backward
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_outputs):
|
||||
return _all_to_all(grad_outputs, ctx.comm_spec), None
|
||||
|
||||
|
||||
def reduce_grad(input_, comm_spec):
|
||||
return _ReduceGrad.apply(input_, comm_spec)
|
||||
|
||||
|
||||
def reduce_input(input_, comm_spec):
|
||||
return _ReduceInput.apply(input_, comm_spec)
|
||||
|
||||
|
||||
def split_forward_gather_backward(input_, comm_spec):
|
||||
return _SplitForwardGatherBackward.apply(input_, comm_spec)
|
||||
|
||||
|
||||
def gather_forward_split_backward(input_, comm_spec):
|
||||
return _GatherForwardSplitBackward.apply(input_, comm_spec)
|
||||
|
||||
|
||||
def all_to_all(input_, comm_spec):
|
||||
return _AllToAll.apply(input_, comm_spec)
|
||||
|
||||
|
||||
class CollectiveCommPattern(Enum):
|
||||
GATHER_FWD_SPLIT_BWD = 'gather_fwd_split_bwd'
|
||||
ALL2ALL_FWD_ALL2ALL_BWD = 'all2all_fwd_all2all_bwd'
|
||||
SPLIT_FWD_GATHER_BWD = 'split_fwd_gather_bwd'
|
||||
ALLREDUCE_FWD_IDENTITY_BWD = 'all_reduce_fwd_identity_bwd'
|
||||
IDENTITY_FWD_ALLREDUCE_BWD = 'identity_fwd_all_reduce_bwd'
|
||||
|
||||
|
||||
class CommSpec:
|
||||
'''
|
||||
Communication spec is used to record the communication action. It has two main functions:
|
||||
1. Compute the communication cost which will be used in auto parallel solver.
|
||||
2. Convert the communication spec to real action which will be used in runtime.
|
||||
It contains comm_pattern to determine the
|
||||
communication method, sharding_spec to determine the communication size, gather_dim and shard_dim
|
||||
to determine the buffer shape, and logical_process_axis
|
||||
|
||||
Argument:
|
||||
comm_pattern(CollectiveCommPattern): decribe the communication method used in this spec.
|
||||
sharding_spec(ShardingSpec): This is sharding spec of the tensor which will join the communication action.
|
||||
gather_dim(int, Optional): The gather_dim of the tensor will be gathered.
|
||||
shard_dim(int, Optional): The shard_dim of the tensor will be sharded.
|
||||
logical_process_axis(Union(int, List[int]), Optional): The mesh_dim to implement the communication action.
|
||||
'''
|
||||
|
||||
def __init__(self,
|
||||
comm_pattern,
|
||||
sharding_spec,
|
||||
gather_dim=None,
|
||||
shard_dim=None,
|
||||
logical_process_axis=None,
|
||||
forward_only=False):
|
||||
self.comm_pattern = comm_pattern
|
||||
self.sharding_spec = sharding_spec
|
||||
self.gather_dim = gather_dim
|
||||
self.shard_dim = shard_dim
|
||||
self.logical_process_axis = logical_process_axis
|
||||
self.forward_only = forward_only
|
||||
if isinstance(self.logical_process_axis, list):
|
||||
self.device_mesh = self.sharding_spec.device_mesh.flatten_device_mesh
|
||||
self.logical_process_axis = 0
|
||||
else:
|
||||
self.device_mesh = self.sharding_spec.device_mesh
|
||||
|
||||
def __repr__(self):
|
||||
res_list = ["CommSpec:("]
|
||||
if self.comm_pattern == CollectiveCommPattern.GATHER_FWD_SPLIT_BWD:
|
||||
res_list.append(f"comm_pattern:GATHER_FWD_SPLIT_BWD, ")
|
||||
res_list.append(f"gather_dim:{self.gather_dim}, ")
|
||||
res_list.append(f"logical_process_axis:{self.logical_process_axis})")
|
||||
elif self.comm_pattern == CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD:
|
||||
res_list.append(f"comm_pattern:ALL2ALL_FWD_ALL2ALL_BWD, ")
|
||||
res_list.append(f"gather_dim:{self.gather_dim}, ")
|
||||
res_list.append(f"shard_dim:{self.shard_dim}, ")
|
||||
res_list.append(f"logical_process_axis: {self.logical_process_axis})")
|
||||
elif self.comm_pattern == CollectiveCommPattern.SPLIT_FWD_GATHER_BWD:
|
||||
res_list.append(f"comm_pattern:SPLIT_FWD_GATHER_BWD, ")
|
||||
res_list.append(f"shard_dim:{self.shard_dim}, ")
|
||||
res_list.append(f"logical_process_axis:{self.logical_process_axis})")
|
||||
elif self.comm_pattern == CollectiveCommPattern.ALLREDUCE_FWD_IDENTITY_BWD:
|
||||
res_list.append(f"comm_pattern:ALLREDUCE_FWD_IDENTITY_BWD, ")
|
||||
res_list.append(f"logical_process_axis:{self.logical_process_axis})")
|
||||
elif self.comm_pattern == CollectiveCommPattern.IDENTITY_FWD_ALLREDUCE_BWD:
|
||||
res_list.append(f"comm_pattern:IDENTITY_FWD_ALLREDUCE_BWD, ")
|
||||
res_list.append(f"logical_process_axis:{self.logical_process_axis})")
|
||||
|
||||
return ''.join(res_list)
|
||||
|
||||
def get_comm_cost(self):
|
||||
'''
|
||||
For all_gather, all2all, and all_reduce operation, the formula provided in DeviceMesh with alpha-beta model is used to
|
||||
compute the communication cost.
|
||||
For shard operation, it is an on-chip operation, so the communication cost is zero.
|
||||
'''
|
||||
comm_size = reduce(operator.mul, self.sharding_spec.get_sharded_shape_per_device(), 1)
|
||||
if self.comm_pattern == CollectiveCommPattern.GATHER_FWD_SPLIT_BWD:
|
||||
forward_communication_cost = self.device_mesh.all_gather_cost(comm_size, self.logical_process_axis)
|
||||
# give a tiny cost to shard
|
||||
backward_communication_cost = 10
|
||||
|
||||
if self.comm_pattern == CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD:
|
||||
forward_communication_cost = self.device_mesh.all_to_all_cost(comm_size, self.logical_process_axis)
|
||||
# grad should have same shape as input tensor
|
||||
# all to all operation has same logical process axis as forward.
|
||||
backward_communication_cost = self.device_mesh.all_to_all_cost(comm_size, self.logical_process_axis)
|
||||
|
||||
if self.comm_pattern == CollectiveCommPattern.ALLREDUCE_FWD_IDENTITY_BWD:
|
||||
forward_communication_cost = self.device_mesh.all_reduce_cost(comm_size, self.logical_process_axis)
|
||||
backward_communication_cost = 0
|
||||
|
||||
if self.comm_pattern == CollectiveCommPattern.IDENTITY_FWD_ALLREDUCE_BWD:
|
||||
forward_communication_cost = 0
|
||||
backward_communication_cost = self.device_mesh.all_reduce_cost(comm_size, self.logical_process_axis)
|
||||
|
||||
if self.comm_pattern == CollectiveCommPattern.SPLIT_FWD_GATHER_BWD:
|
||||
# give a tiny cost to shard
|
||||
forward_communication_cost = 10
|
||||
backward_communication_cost = self.device_mesh.all_gather_cost(comm_size, self.logical_process_axis)
|
||||
try:
|
||||
if self.forward_only:
|
||||
total_communication_cost = forward_communication_cost
|
||||
else:
|
||||
total_communication_cost = forward_communication_cost + backward_communication_cost
|
||||
except:
|
||||
raise RuntimeError(f"Could not find a matching CollectiveCommPattern for {self.comm_pattern}.")
|
||||
|
||||
return total_communication_cost
|
||||
|
||||
def covert_spec_to_action(self, tensor):
|
||||
'''
|
||||
Convert CommSpec into runtime action, implement real collection communication to target tensor.
|
||||
The collection communication action is directed by the CommSpec.
|
||||
|
||||
Argument:
|
||||
tensor(torch.Tensor): Tensor stored in each device, which could be different in different ranks.
|
||||
'''
|
||||
if self.comm_pattern in pattern_to_func_dict:
|
||||
tensor.data = pattern_to_func_dict[self.comm_pattern](tensor, self)
|
||||
else:
|
||||
tensor.data = tensor
|
||||
|
||||
|
||||
pattern_to_func_dict = {
|
||||
CollectiveCommPattern.GATHER_FWD_SPLIT_BWD: gather_forward_split_backward,
|
||||
CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD: all_to_all,
|
||||
CollectiveCommPattern.SPLIT_FWD_GATHER_BWD: split_forward_gather_backward,
|
||||
CollectiveCommPattern.ALLREDUCE_FWD_IDENTITY_BWD: reduce_input,
|
||||
CollectiveCommPattern.IDENTITY_FWD_ALLREDUCE_BWD: reduce_grad,
|
||||
}
|
||||
__all__ = ['ShapeConsistencyManager', 'ShapeConsistencyOptions', 'set_shape_consistency_options']
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -406,7 +60,7 @@ class ShapeConsistencyManager(metaclass=SingletonMeta):
|
||||
assert isinstance(value, bool)
|
||||
self._forward_only = value
|
||||
|
||||
def get_all_all_gather_spec(self, source_spec, orig_cost):
|
||||
def get_all_all_gather_spec(self, source_spec, orig_cost_dict):
|
||||
'''
|
||||
Get all valid sharding specs from source_spec with single all-gather operation, and
|
||||
accumulate commucation cost on origin cost which will finally be used in auto sharding solver.
|
||||
@@ -463,16 +117,18 @@ class ShapeConsistencyManager(metaclass=SingletonMeta):
|
||||
forward_only=self.forward_only)
|
||||
|
||||
# compute the communication cost with CommSpec
|
||||
cost = comm_spec.get_comm_cost()
|
||||
cost_dict = comm_spec.get_comm_cost()
|
||||
|
||||
# generate new sharding spec
|
||||
new_sharding_spec = ShardingSpec(source_spec.device_mesh,
|
||||
source_spec.entire_shape,
|
||||
dim_partition_dict=new_dim_partition_dict)
|
||||
valid_spec_dict[new_sharding_spec] = (comm_spec, orig_cost + cost)
|
||||
for phase, cost in cost_dict.items():
|
||||
cost_dict[phase] = cost + orig_cost_dict[phase]
|
||||
valid_spec_dict[new_sharding_spec] = (comm_spec, cost_dict)
|
||||
return valid_spec_dict
|
||||
|
||||
def get_all_all_to_all_spec(self, source_spec, orig_cost):
|
||||
def get_all_all_to_all_spec(self, source_spec, orig_cost_dict):
|
||||
'''
|
||||
Get all valid sharding specs from source_spec with single all-to-all operation, and
|
||||
accumulate commucation cost on origin cost which will finally be used in auto sharding solver.
|
||||
@@ -552,7 +208,7 @@ class ShapeConsistencyManager(metaclass=SingletonMeta):
|
||||
forward_only=self.forward_only)
|
||||
|
||||
# compute the communication cost with CommSpec
|
||||
cost = comm_spec.get_comm_cost()
|
||||
cost_dict = comm_spec.get_comm_cost()
|
||||
new_dim_partition_dict = deepcopy(source_spec.dim_partition_dict)
|
||||
|
||||
# We won't add empty list into dim_partition_dict
|
||||
@@ -570,10 +226,12 @@ class ShapeConsistencyManager(metaclass=SingletonMeta):
|
||||
new_sharding_spec = ShardingSpec(source_spec.device_mesh,
|
||||
source_spec.entire_shape,
|
||||
dim_partition_dict=new_dim_partition_dict)
|
||||
valid_spec_dict[new_sharding_spec] = (comm_spec, orig_cost + cost)
|
||||
for phase, cost in cost_dict.items():
|
||||
cost_dict[phase] = cost + orig_cost_dict[phase]
|
||||
valid_spec_dict[new_sharding_spec] = (comm_spec, cost_dict)
|
||||
return valid_spec_dict
|
||||
|
||||
def get_all_shard_spec(self, source_spec, orig_cost):
|
||||
def get_all_shard_spec(self, source_spec, orig_cost_dict):
|
||||
'''
|
||||
Get all valid sharding specs from source_spec with single shard operation, and
|
||||
accumulate commucation cost on origin cost which will finally be used in auto sharding solver.
|
||||
@@ -639,16 +297,18 @@ class ShapeConsistencyManager(metaclass=SingletonMeta):
|
||||
forward_only=self.forward_only)
|
||||
|
||||
# compute the communication cost with CommSpec
|
||||
cost = comm_spec.get_comm_cost()
|
||||
cost_dict = comm_spec.get_comm_cost()
|
||||
|
||||
# generate new sharding spec
|
||||
new_sharding_spec = ShardingSpec(source_spec.device_mesh,
|
||||
source_spec.entire_shape,
|
||||
dim_partition_dict=new_dim_partition_dict)
|
||||
valid_spec_dict[new_sharding_spec] = (comm_spec, orig_cost + cost)
|
||||
for phase, cost in cost_dict.items():
|
||||
cost_dict[phase] = cost + orig_cost_dict[phase]
|
||||
valid_spec_dict[new_sharding_spec] = (comm_spec, cost_dict)
|
||||
return valid_spec_dict
|
||||
|
||||
def get_all_one_step_transform_spec(self, source_spec, orig_cost):
|
||||
def get_all_one_step_transform_spec(self, source_spec, orig_cost_dict):
|
||||
'''
|
||||
Get all valid sharding specs from source_spec with one step transform, and
|
||||
accumulate commucation cost on origin cost which will finally be used in auto sharding solver.
|
||||
@@ -665,9 +325,9 @@ class ShapeConsistencyManager(metaclass=SingletonMeta):
|
||||
valid_spec_dict(Dict[ShardingSpec, float]): all valid sharding specs from source_spec with single all-to-all operation.
|
||||
'''
|
||||
valid_spec_dict = {}
|
||||
valid_spec_dict.update(self.get_all_all_gather_spec(source_spec, orig_cost))
|
||||
valid_spec_dict.update(self.get_all_all_to_all_spec(source_spec, orig_cost))
|
||||
valid_spec_dict.update(self.get_all_shard_spec(source_spec, orig_cost))
|
||||
valid_spec_dict.update(self.get_all_all_gather_spec(source_spec, orig_cost_dict))
|
||||
valid_spec_dict.update(self.get_all_all_to_all_spec(source_spec, orig_cost_dict))
|
||||
valid_spec_dict.update(self.get_all_shard_spec(source_spec, orig_cost_dict))
|
||||
return valid_spec_dict
|
||||
|
||||
def shape_consistency(self, source_spec, target_spec):
|
||||
@@ -730,7 +390,7 @@ class ShapeConsistencyManager(metaclass=SingletonMeta):
|
||||
total_cost: 12294.402000000002
|
||||
'''
|
||||
MAX_TRANSFORM_STEPS = 20
|
||||
total_cost = 0
|
||||
total_cost_dict = {'forward': 0, 'backward': 0, 'total': 0}
|
||||
total_steps = 0
|
||||
transform_path = []
|
||||
comm_action_sequence = []
|
||||
@@ -740,35 +400,37 @@ class ShapeConsistencyManager(metaclass=SingletonMeta):
|
||||
# We do nothing if the sharding spec is all the same.
|
||||
if source_spec.sharding_sequence_difference(target_spec) == 0:
|
||||
self.cached_spec_pairs_transform_path[spec_pairs] = (transform_path, comm_action_sequence)
|
||||
return (transform_path, comm_action_sequence, total_cost)
|
||||
return (transform_path, comm_action_sequence, total_cost_dict)
|
||||
|
||||
temp_sharding_spec = source_spec
|
||||
transform_path.append(temp_sharding_spec)
|
||||
# To avoid dead loop, the loop will break after MAX_TRANSFORM_STEPS transforms
|
||||
while total_steps <= MAX_TRANSFORM_STEPS:
|
||||
valid_transform_spec_dict = self.get_all_one_step_transform_spec(temp_sharding_spec, total_cost)
|
||||
valid_transform_spec_dict = self.get_all_one_step_transform_spec(temp_sharding_spec, total_cost_dict)
|
||||
best_difference_score = math.inf
|
||||
|
||||
for sharding_spec, info_pairs in valid_transform_spec_dict.items():
|
||||
comm_spec, cost = info_pairs
|
||||
comm_spec, cost_dict = info_pairs
|
||||
spec_difference = sharding_spec.sharding_sequence_difference(target_spec)
|
||||
|
||||
if spec_difference == 0:
|
||||
total_cost += cost
|
||||
for phase, cost in total_cost_dict.items():
|
||||
total_cost_dict[phase] = cost + cost_dict[phase]
|
||||
transform_path.append(sharding_spec)
|
||||
comm_action_sequence.append(comm_spec)
|
||||
self.cached_spec_pairs_transform_path[spec_pairs] = (transform_path, comm_action_sequence)
|
||||
return (transform_path, comm_action_sequence, total_cost)
|
||||
return (transform_path, comm_action_sequence, total_cost_dict)
|
||||
|
||||
if spec_difference < best_difference_score:
|
||||
temp_sharding_spec = sharding_spec
|
||||
temp_cost = cost
|
||||
temp_cost_dict = cost_dict
|
||||
temp_comm_spec = comm_spec
|
||||
best_difference_score = spec_difference
|
||||
|
||||
transform_path.append(temp_sharding_spec)
|
||||
comm_action_sequence.append(temp_comm_spec)
|
||||
total_cost += temp_cost
|
||||
for phase, cost in total_cost_dict.items():
|
||||
total_cost_dict[phase] = cost + temp_cost_dict[phase]
|
||||
total_steps += 1
|
||||
|
||||
raise RuntimeError(f"Could not find a valid transform path with in {MAX_TRANSFORM_STEPS} steps.")
|
||||
|
||||
Reference in New Issue
Block a user