8000 clear locally accumulated gradient by assigning with zeros_like to avoid infinite gradient not correctly cleared up by yundai424 · Pull Request #3505 · horovod/horovod · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

clear locally accumulated gradient by assigning with zeros_like to avoid infinite gradient not correctly cleared up #3505

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 15, 2022
Merged

clear locally accumulated gradient by assigning with zeros_like to avoid infinite gradient not correctly cleared up #3505

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions horovod/tensorflow/gradient_aggregation_eager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,8 +119,8 @@ def _allreduce_helper(self, grads, vars):
def _clear_vars(self):
self.counter.assign(0)
for idx in self.locally_aggregated_grads.keys():
self.locally_aggregated_grads[idx].assign_add(
-1 * self.locally_aggregated_grads[idx])
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this a breaking change for existing user code?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

inf - inf is NaN, so semantically it would definitely be an improvement to assign zero here.

However, I suspect that tf.assign_add() has been used here originally to avoid an intermediate extra memory allocation for the result of tf.zeros_like(). In the past I've seen a similar effect actually cause perceivable memory waste.

I wonder if this is still the case or if recent releases of TensorFlow can optimize x.assign(tf.zeros_like(x)) appropriately..

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tgaddair Is there a reason of introducing tf.assign_add() originally?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@maxhgerlach i had the same doubt initially, but according to this, it seems like even the old assign_add created a new buffer. I guess ultimately the proper way is to use the in-place c++ api which is a bit involved than the scope of the pr.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would assume:

a.assign_add(-1 * a)

consumes the same amount of memory as:

a.assign(tf.zeros_like(a))

since both likely create temporaries, as mentioned by @Tixxx.

Is there a TF API to just set all values of a tensor to a scalar value?

Copy link
Contributor Author
@yundai424 yundai424 Apr 8, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should go into this route when assigning zero to local gradient. What happens afterwards will be the same for Assign and AssignAdd, the only difference is for assign the update is params.device(d) = update; and for assignAdd it's params.device(d) += update;. In this sense it looks like pretty much in place (w/ a buffer for the update). My understanding could be wrong :)

self.locally_aggregated_grads[idx].assign(
tf.zeros_like(self.locally_aggregated_grads[idx]))

def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):
def increment_optimizer_iteration():
Expand Down
27 changes: 27 additions & 0 deletions test/parallel/test_tensorflow2_keras.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,33 @@ def test_sparse_as_dense_with_grad_aggregation(self):
aggregation_counter = opt._agg_helper.counter.numpy()
assert aggregation_counter == training_steps % backward_passes_per_step

def test_grad_aggregation_with_inf_grad(self):
backward_passes_per_step = 2
step_count = tf.Variable(0, trainable=False, dtype=tf.int32)
opt = tf.optimizers.SGD()
opt = hvd.DistributedOptimizer(
opt,
backward_passes_per_step=backward_passes_per_step,
sparse_as_dense=True
)
x = tf.Variable(0.)
var = [x]

def loss():
step_count.assign_add(1)
return tf.cond(
pred=tf.greater(step_count, 1),
true_fn=lambda: x,
false_fn=lambda: x * float('inf')
)
for _ in range(2 * backward_passes_per_step):
# in the first aggregation cycle the gradient is infinite,
# and it should be cleaned up to zero after apply_gradients
# and doesn't affect the 2nd aggregation cycle
grads_and_vars = opt._compute_gradients(loss=loss, var_list=var)
opt.apply_gradients(grads_and_vars)
assert tf.math.is_finite(grads_and_vars[0][0])

def test_from_config(self):
opt = keras.optimizers.Adam()
hopt = hvd.DistributedOptimizer(opt)
Expand Down
0