Skip to content

Commit ba8b85f

Browse files
committed
add version select
1 parent f07f6e5 commit ba8b85f

File tree

1 file changed

+9
-0
lines changed

1 file changed

+9
-0
lines changed

src/xccl/ProcessGroupXCCL.cpp

+9
Original file line numberDiff line numberDiff line change
@@ -973,6 +973,15 @@ c10::intrusive_ptr<Work> ProcessGroupXCCL::allreduce(
973973
xcclReduceOp,
974974
comm,
975975
ccl::create_stream(stream.queue()));
976+
#if !defined(XCCL_HAS_AVG)
977+
if (opts.reduceOp == ReduceOp::AVG) {
978+
auto divisor = getSize();
979+
c10::StreamGuard guard(stream);
980+
c10::xpu::XPUCachingAllocator::recordStream(
981+
output.storage().data_ptr(), stream);
982+
output.div_(divisor);
983+
}
984+
#endif
976985
return;
977986
},
978987
OpType::ALLREDUCE,

0 commit comments

Comments
 (0)