Method reduceGrad

csrc/includes/deepcompile.h:502–552

    }

    virtual at::Tensor reduceGrad(at::Tensor grad_tensor, long ds_id)
    {
        int world_size = process_group_->getSize();
        const DSParam& param = param_registry_->getParam(ds_id);
        const auto expected_grad_dtype = param.getExpectedGradDtype();
        // Match PyTorch's leaf grad accumulation dtype before bucket selection:
        // https://docs.pytorch.org/docs/main/generated/torch.sparse.semi_structured.SparseSemiStructuredTensorCUSPARSELT.html#torch.sparse.semi_structured.SparseSemiStructuredTensorCUSPARSELT.grad_dtype
        if (expected_grad_dtype.has_value() &&
            grad_tensor.scalar_type() != expected_grad_dtype.value()) {
            grad_tensor = grad_tensor.to(expected_grad_dtype.value());
        }
        const auto scalar_type = grad_tensor.scalar_type();
        std::shared_ptr<ReduceBucket> reduce_bucket = reduce_buckets_->getBuffer(scalar_type);

        auto comp_stream = at::cuda::getCurrentCUDAStream();

        if (reduce_bucket->shouldFlush(grad_tensor.numel())) {
            int rank = process_group_->getRank();

            flushReduceBucket(scalar_type);

            // reduce_bucket is swapped in flushReduceBucket if double buffering is enabled
            reduce_bucket = reduce_buckets_->getBuffer(scalar_type);
        }

        if (grad_tensor.numel() > reduce_bucket->getSize()) {
            // extend buckets
            at::cuda::stream_synchronize(rs_stream_);
            reduce_bucket->reserve(grad_tensor.numel());
        }

        at::Tensor reduce_in_buffer = reduce_bucket->allocate(grad_tensor.numel());

        // This ensures the order of reduce_scatter -> copy
        // Without this block, copy may start while reduce_scatter is still running
        reduce_buckets_->getEvent(scalar_type)->block(comp_stream);
        auto copy_src = grad_tensor.contiguous().view({-1}).detach();
        // keep references to copy src
        reduce_tasks_[scalar_type].emplace_back(ds_id, copy_src, reduce_in_buffer);

        // computation must be done before copy
        rs_comp_done_events_[ds_id]->record(comp_stream);
        rs_comp_done_events_[ds_id]->block(copy_stream_);
        {
            at::cuda::CUDAStreamGuard guard(copy_stream_);
            reduce_in_buffer.copy_(copy_src, true);
            rs_copy_done_events_[ds_id]->record(copy_stream_);
        }

        return at::Tensor();
    }

    bool hasParam(long ds_id) const { return hasKey(has_acc_grad_, ds_id); }