Repository: microsoft/DeepSpeed
File: csrc/cpu/comm/ccl.cpp (C++)

Code before:
                            data.data_ptr(),
                            data.numel(),
                            get_ccl_datatype(data.scalar_type()),
                            get_ccl_reduce_op(op, data),
                            _get_comm_from_group(group),
                            attr)
                 .wait());
}

void inference_all_reduce(torch::Tensor& data)
{
#ifdef DO_PROFILE
    static double total_time = 0.0;
    static double total_time_sq = 0.0;
    static int count = -16;  // warmup
    static double max_time = 0.0;
    static double min_time = DBL_MAX;
    // make sure all rank reach this point before measuring time
    // turn on this if you suspect each rank didn't reach here at the same time (stragger)
    // if (all_ranks_local_p) {
    // barrier_wait(0, world_size);
    //}
    auto start = std::chrono::system_clock::now();
#endif

    auto numel = data.numel();

    int data_size = 0;
    bool data_type_fallback = false;

    switch (data.scalar_type()) {
        case c10::ScalarType::BFloat16: data_size = numel * 2; break;
        case c10::ScalarType::Float: data_size = numel * 4; break;
        default: data_type_fallback = true;
    }

    if (data_type_fallback || !all_ranks_local_p) {
        // fallback to oneccl allreduce
        CCLCHECK(ccl::allreduce(data.data_ptr(),
                                data.data_ptr(),
                                data.numel(),
                                get_ccl_datatype(data.scalar_type()),
                                ccl::reduction::sum,
                                _get_comm_from_group())
                     .wait());
    } else {
        all_reduce_outer_loop(data, numel, data_size);
    }

#ifdef DO_PROFILE
    auto end = std::chrono::system_clock::now();

Review comment (suggestion):
I think it is confusing to drop this assert since `inference_all_reduce` only works for sum even though it is not clear from the name. Users can easily assume that this function is similar to more common ones like nccl allreduce that supports other user-defined reduction ops like MAX. I see two possible solutions:
1. Restore this assert so that users are notified of unsupported ops, and developers can eventually support more ops without changing API.
2. Rename function to something like `inference_all_reduce_sum` to better reflect the intention. 

Code after:
                            data.data_ptr(),
                            data.numel(),
                            get_ccl_datatype(data.scalar_type()),
                            get_ccl_reduce_op(op, data),
                            _get_comm_from_group(group),
                            attr)
                 .wait());
}

void inference_all_reduce(torch::Tensor& data, py::object op)
{
#ifdef DO_PROFILE
    static double total_time = 0.0;
    static double total_time_sq = 0.0;
    static int count = -16;  // warmup
    static double max_time = 0.0;
    static double min_time = DBL_MAX;
    // make sure all rank reach this point before measuring time
    // turn on this if you suspect each rank didn't reach here at the same time (stragger)
    // if (all_ranks_local_p) {
    // barrier_wait(0, world_size);
    //}
    auto start = std::chrono::system_clock::now();
#endif

    static py::object ReduceOp = py::module_::import("deepspeed.comm").attr("ReduceOp");
    static auto ReduceOpSum = (int)py::int_(ReduceOp.attr("SUM").attr("value"));

    assert(py::int_(op.attr("value")) == ReduceOpSum);

    auto numel = data.numel();

    int data_size = 0;
    bool data_type_fallback = false;

    switch (data.scalar_type()) {
        case c10::ScalarType::BFloat16: data_size = numel * 2; break;
        case c10::ScalarType::Float: data_size = numel * 4; break;
        default: data_type_fallback = true;
    }

    if (data_type_fallback || !all_ranks_local_p) {
        // fallback to oneccl allreduce
        CCLCHECK(ccl::allreduce(data.data_ptr(),
                                data.data_ptr(),
                                data.numel(),
                                get_ccl_datatype(data.scalar_type()),
                                get_ccl_reduce_op(op, data),
                                _get_comm_from_group())
                     .wait());
    } else {