Fix data type in PartialDistributedGradientTape (#3985)

Signed-off-by: Chen Zhu <[email protected]>
horovod · Sep 18, 2023 · 7e4d993 · 7e4d993
1 parent 8724eff
commit 7e4d993
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/horovod/_keras/__init__.py b/horovod/_keras/__init__.py
@@ -193,7 +193,7 @@ def __filtered_reduce_grads(grads, vars):
  if v_ref in self._local_vars and v2g[v_ref] is not None:
  grad = v2g[v_ref]
  if isinstance(grad, tf.IndexedSlices):
- grad = tf.IndexedSlices(grad.values / horovod_size, grad.indices, grad.dense_shape)
+ grad = tf.IndexedSlices(grad.values / float(horovod_size), grad.indices, grad.dense_shape)
  else:
  grad /= float(horovod_size)
  v2g[v_ref] = grad
@@ -209,7 +209,7 @@ def __filtered_reduce_grads(grads, vars):
  if v in self._local_vars and v2g[v] is not None:
  grad = v2g[v]
  if isinstance(grad, tf.IndexedSlices):
- grad = tf.IndexedSlices(grad.values / horovod_size, grad.indices, grad.dense_shape)
+ grad = tf.IndexedSlices(grad.values / float(horovod_size), grad.indices, grad.dense_shape)
  else:
  grad /= float(horovod_size)
  v2g[v] = grad

diff --git a/horovod/tensorflow/__init__.py b/horovod/tensorflow/__init__.py
@@ -744,7 +744,7 @@ def _filtered_reduce_grads(grads, vars):
  if v_ref in self._local_vars and v2g[v_ref]:
  grad = v2g[v_ref]
  if isinstance(grad, tf.IndexedSlices):
- grad = tf.IndexedSlices(grad.values / horovod_size, grad.indices, grad.dense_shape)
+ grad = tf.IndexedSlices(grad.values / float(horovod_size), grad.indices, grad.dense_shape)
  else:
  grad /= float(horovod_size)
  v2g[v_ref] = grad
@@ -760,7 +760,7 @@ def _filtered_reduce_grads(grads, vars):
  if v in self._local_vars and v2g[v]:
  grad = v2g[v]
  if isinstance(grad, tf.IndexedSlices):
- grad = tf.IndexedSlices(grad.values / horovod_size, grad.indices, grad.dense_shape)
+ grad = tf.IndexedSlices(grad.values / float(horovod_size), grad.indices, grad.dense_shape)
  else:
  grad /= float(horovod_size)
  v2g[v] = grad
@@ -1086,7 +1086,7 @@ def gradient(self, target, sources, output_gradients=None, use_generic_names=Fal
  if s_ref in self._local_sources and s2g[s_ref] is not None:
  grad = s2g[s_ref]
  if isinstance(grad, tf.IndexedSlices):
- grad = tf.IndexedSlices(grad.values / horovod_size, grad.indices, grad.dense_shape)
+ grad = tf.IndexedSlices(grad.values / float(horovod_size), grad.indices, grad.dense_shape)
  else:
  grad /= float(horovod_size)
  s2g[s_ref] = grad
@@ -1102,7 +1102,7 @@ def gradient(self, target, sources, output_gradients=None, use_generic_names=Fal
  if s in self._local_sources and s2g[s] is not None:
  grad = s2g[s]
  if isinstance(grad, tf.IndexedSlices):
- grad = tf.IndexedSlices(grad.values / horovod_size, grad.indices, grad.dense_shape)
+ grad = tf.IndexedSlices(grad.values / float(horovod_size), grad.indices, grad.dense_shape)
  else:
  grad /= float(horovod_size)
  s2g[s] = grad
@@ -1237,4 +1237,4 @@ def PartialDistributedGradientTape(gradtape, device_dense='', device_sparse='',
  num_groups, groups, process_set, scale_local_gradients)
  for var in local_vars:
  _tape.register_local_source(var)
- return _tape
+ return _tape