Commit 0494c607 authored by Yin, Junqi's avatar Yin, Junqi

add efficient decentralized to sgd

parent ae56ca02
......@@ -40,7 +40,8 @@ class SGD(Optimizer):
cur_rank=self.rank,
world=conf.graph.ranks,
neighbors_info=self.neighbors_info,
aggregator_type="decentralized",
aggregator_type="efficient_decentralized",
graph=conf.graph
)
self.world_aggregator = comm.get_aggregators(
cur_rank=self.rank,
......
......@@ -307,6 +307,10 @@ class EfficientDecentralizedAggregation(Aggregation):
if neighbor_rank != rank
]
self.out_edges, self.in_edges = graph.get_edges()
self.world_size = float(len(self.neighbor_ranks))
for out_edge, in_edge in zip(self.out_edges, self.in_edges):
print("OUT_EDGES: ", out_edge.src, out_edge.dest)
print("IN_EDGES: ", in_edge.src, in_edge.dest)
def _agg(self, data, op, force_wait=True):
"""Aggregate data using `op` operation.
......@@ -324,7 +328,7 @@ class EfficientDecentralizedAggregation(Aggregation):
out_reqs, in_reqs = [], []
for out_edge, in_edge in zip(self.out_edges, self.in_edges):
out_req = dist.broadcast(
tensor=self.in_buffer[self.rank],
tensor=self.in_buffer[out_edge.src],
src=out_edge.src,
group=out_edge.process_group,
async_op=True,
......@@ -338,8 +342,33 @@ class EfficientDecentralizedAggregation(Aggregation):
async_op=True,
)
in_reqs.append(in_req)
return [out_reqs, in_reqs], self.in_buffer
#return [out_reqs, in_reqs], self.in_buffer
# wait until finish.
if force_wait:
self.complete_wait((out_reqs, in_reqs))
# Aggregate local_data
if op == "avg":
output = sum(self.in_buffer.values()) / (self.world_size + 1)
elif op == "weighted":
output = sum(
[
tensor * self.neighbors_info[rank]
for rank, tensor in self.in_buffer.items()
]
)
elif op == "get_raw_sync_data":
output = self.in_buffer
else:
raise NotImplementedError("op {} is not supported yet.".format(op))
return output
else:
if op == "get_raw_sync_data":
return (out_reqs, in_reqs), self.in_buffer
else:
raise NotImplementedError("op {} is not supported yet.".format(op))
def complete_wait(self, reqs):
out_reqs, in_reqs = reqs
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment