NVIDIA / NVIDIA/Megatron-LM

[QUESTION]May I allocate and register NCCL userbuffer for activations/gradients transmitted between PP stages?

Open
#1,720 1 comment 2 reactions 0 assignees View on GitHub
community-request module: training question waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

I noted that when use-nccl-ub is enabled, param and grad buffer can be allocated and registered with NCCL user buffer mechanism. That could save SM resources and improve performance. I would like to apply NCCL userbuffer for activations/gradients transmitted between PP stages.

From the current NCCL backend implementation in pytorch code (ProcessGroupNCCL.cpp), it looks P2P NCCL register buffer is not supported.

Image

Please let me know if we have plan to support it.

It would be highly appreciated if any suggestion? Many thanks.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.