in blocksparse/matmul.py [0:0]
def group_param_grads(param_grad, group_size=8):
assert group_size <= 8
# backward walk param grad to find BlocksparseMatmulDW ops
# this should only hit BlocksparseMatmulDWs, BlocksparseMatmulDGs, AddNs or FloatCasts
ops = get_parents(param_grad, "BlocksparseMatmulDW")
if len(ops) <= 1:
return param_grad
# this sorting is dependent on the op names being correctly ordered.
ops.sort(key=lambda op: op.name.split('/')[-1], reverse=True)
# for x in ops:
# print(x.name)
# print("")
# exit()
segment_size = len(ops)
if ops[0].get_attr("gate_grad") and len(ops[0].inputs) == 4:
gate_count = dict()
max_count = 0
for op in ops:
gate = op.inputs[3]
count = gate_count.get(gate, 0) + 1
gate_count[gate] = count
max_count = max(max_count, count)
for count in gate_count.values():
if count != max_count:
raise ValueError("Non-uniform gate broadcasting detected.")
segment_size = max_count
if group_size > segment_size:
group_size = segment_size
else:
assert segment_size % group_size == 0
# nothing to rewrite here.
if segment_size == 1:
return param_grad
# use the parent scope for the new ops
scope = ops[-1].name.split('/')
scope = '/'.join(scope[0:-1])
# we're going to be using absolute names, so clear name_scope
with tf.name_scope(None):
dw = None
dws = list()
offset = 0
seg_cnt = 0
while offset < len(ops):
xs = [op.inputs[0] for op in ops[offset:offset+group_size] ]
gs = [op.inputs[1] for op in ops[offset:offset+group_size] ]
# Get the corresponding activation grad op for the last param grad op in the group
bprop = None
for consumer in gs[-1].consumers():
if consumer.type == "BlocksparseMatmulDX":
bprop = consumer
break
assert bprop is not None
# get attributes of first op in group
up = ops[offset]
blocks = up.get_attr("blocks")
bsize = up.get_attr("bsize")
axis = up.get_attr("axis")
gated_dw = up.get_attr("gated_dw")
gate_grad = up.get_attr("gate_grad")
C = up.get_attr("C")
K = up.get_attr("K")
bench = up.get_attr("bench") // len(xs)
lut = up.inputs[2]
name = "%s/matmul_concat_updat_%03d" % (scope, offset)
gate = [up.inputs[3]] if len(up.inputs) > 3 else []
# The first op needs to allocate a new dw tensor
if dw is None:
dw = blocksparse_matmul_dw(
xs, gs, lut, gate, gated_dw=gated_dw,
gate_grad=gate_grad, blocks=blocks, bsize=bsize, axis=axis,
C=C, K=K, bench=bench, name=name)
# subsequent ops can just accumulate in place
else:
dw = blocksparse_matmul_dwa(
xs, gs, lut, dw, gate, gated_dw=gated_dw,
gate_grad=gate_grad, blocks=blocks, bsize=bsize, axis=axis,
C=C, K=K, bench=bench, name=name)
# force the dw op before any more time steps are processed
bprop._add_control_input(dw.op)
seg_cnt += group_size
offset += group_size
if gate_grad and seg_cnt >= segment_size:
seg_cnt = 0
dws.append(dw)
dw = None
if gate_grad:
for i, dw in enumerate(dws):
# for op in ops[i*group_size:(i+1)*group_size]:
# print(op.name)
# print()
dw_op = ops[i*segment_size:(i+1)*segment_size][-1]
dws[i] = group_dg_grads(dw_op, dw, scope)
# add up final dw values in groups of 4 for good mix of perforamnce and memory use
dw = ew.add_n8_op(dws[0:4]) if len(dws) > 1 else dws[0]
for i in range(4, len(dws), 4):
dw = ew.add_n8_op(dws[i:i+4] + [dw])
# splice in these grad op types sitting on top of the param
if param_grad.op.type in ("Cast", "FloatCast", "L2NormalizeGradCK", "L2NormalizeGainGradCK"):
param_grad.op._update_input(0, dw)
dw = param_grad
elif param_grad.op.type not in ("AddN", "AddN8", "BlocksparseMatmulDW","BlocksparseMatmulDG"):
raise ValueError("Unexpected grad op type:", param_grad.op.type, param_grad.op.name)
return dw