[Issue]: One-rank avg reduction missing elements

[kwen2501]

How is this issue impacting you?

Data corruption

Share Your Debug Logs

This is an instantiation of a PyTorch issue.

More details here: pytorch/pytorch#168092

Steps to Reproduce the Issue

No response

NCCL Version

NCCL version 2.27.5+cuda12.9

Your platform details

H100

Error Message & Behavior

Last 8 elements of the buffer seems un-processed:

msg=f"Last entries are: {output[-10:]} vs {input_tensor[-10:]}",

AssertionError: Last entries are: tensor([-1.1953,  0.2295,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,
         0.0000,  0.0000], device='cuda:0', dtype=torch.bfloat16) vs tensor([-1.1953,  0.2295,  1.6016,  0.2891,  1.9844, -0.5703, -0.1904,  0.6172,
         1.4141,  1.2578], device='cuda:0', dtype=torch.bfloat16)

[d4l3k]

This seems to fix it:

--- a/src/device/onerank.cu
+++ b/src/device/onerank.cu
@@ -22,8 +22,8 @@
 
     // each block/channel gets a roughly equal segment of 16 byte packs
     constexpr int EltPerPack = 16/sizeof(T);
-    intptr_t i0 = (bid+0)*alignUp(nElts/bn, EltPerPack);
-    intptr_t i1 = (bid+1)*alignUp(nElts/bn, EltPerPack);
+    intptr_t i0 = (bid+0)*alignUp(divUp(nElts, bn), EltPerPack);
+    intptr_t i1 = (bid+1)*alignUp(divUp(nElts, bn), EltPerPack);
     i0 = min(i0, nElts);
     i1 = min(i1, nElts);
     src = (T*)src + i0;

https://github.com/NVIDIA/nccl/blob/master/src/device/onerank.cu

Thanks to @pavanbalaji for root causing this

[kiskra-nvidia]

@jbachan It looks like you touched this code last 😃. Could you confirm the above fix?

[jbachan]

Yes looks good

[teojgo]

The fix is going to be included in the next release.