Skip to content

Commit 289903c

Browse files
Fix for CI hand - deepspeedai#7851
Signed-off-by: Phalani Paladugu <mailofphalani@gmail.com>
1 parent 1ea484b commit 289903c

1 file changed

Lines changed: 5 additions & 1 deletion

File tree

deepspeed/comm/torch.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,11 @@ def all_reduce(self, tensor, op=torch.distributed.ReduceOp.SUM, group=None, asyn
167167
return torch.distributed.all_reduce(tensor=tensor, op=op, group=group, async_op=async_op)
168168

169169
def inference_all_reduce(self, tensor, op, group=None):
170-
if not hasattr(torch.ops, 'deepspeed') or not hasattr(torch.ops.deepspeed, 'inference_all_reduce_'):
170+
use_ds_op = hasattr(torch.ops, 'deepspeed') and hasattr(torch.ops.deepspeed, 'inference_all_reduce_')
171+
world_size = torch.distributed.get_world_size(group=group)
172+
if world_size <= 1:
173+
return tensor
174+
if not use_ds_op:
171175
op = self._reduce_op(op)
172176
return torch.distributed.all_reduce(tensor=tensor, op=op, group=group, async_op=False)
173177
else:

0 commit comments

Comments
 (0)