Skip to content

Reduction backward passes force a device sync via grad.item() #22

Description

@CanReader

src/autograd/ops/reduce.rs:17 and :32:

vec![Tensor::full(&self.shape, grad.item()).to(grad.device())]

grad.item() downloads from the GPU, the full tensor is then built on the host and
uploaded again. That's every sum() and mean() backward, so every loss, doing a
synchronising round trip plus an allocation the size of the whole input.

CrossEntropyBackward (src/autograd/ops/loss.rs:20) and
WeightedCrossEntropyBackward do the same thing.

fastnn_cuda_fill is already declared in ffi.rs and implemented in kernels.cu
and never called. Either that, or a mul_scalar against a ones buffer, keeps this
on device.

Metadata

Metadata

Assignees

No one assigned

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions