Commit 6e21c2fb authored by Yin, Junqi's avatar Yin, Junqi
Browse files

add before_sync print for ring/torus

parent dc292b93
......@@ -98,6 +98,12 @@ class SGD(Optimizer):
n_bits = get_n_bits(flatten_grads.buffer)
with kargs["timer"]("sync.apply_grad", epoch=self.conf.epoch_):
if self.conf.print_grad:
self.param_groups, self.state, apply_grad_to_model=True,
dryrun = False, logger=self.conf.logger
self.param_groups, self.state, apply_grad_to_model=True
......@@ -51,12 +51,14 @@ def apply_gradient(param_groups, state, apply_grad_to_model=True,
d_p = buf
if apply_grad_to_model:
if dryrun:
pnew = torch.add(, d_p, alpha=-group["lr"])
name = group["name"]
if dryrun: # complete graph
pnew = torch.add(, d_p, alpha=-group["lr"])
logger.log(f"parameter: {name} param_norm: {pnew.norm()}")
else:, alpha=-group["lr"])
if logger is not None: # ring or torus
logger.log(f"parameter: {name} param_norm: {p.norm()}")
else: = d_p
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