nemo_automodel.components.moe.uccl_ep

UCCL-EP integration for expert parallelism.

UCCL-EP (https://github.com/uccl-project/uccl/tree/main/ep) has the same interface and functionality as DeepEP, and enables GPU-initiated communication for MoE models across heterogeneous GPUs and NICs.

Vendored files (from https://github.com/uccl-project/uccl/tree/main/ep/deep_ep_wrapper): _buffer.py <- deep_ep_wrapper/deep_ep/buffer.py _utils.py <- deep_ep_wrapper/deep_ep/utils.py

Usage: Set dispatcher: uccl_ep in the model backend config.

Submodules

• nemo_automodel.components.moe.uccl_ep._buffer
• nemo_automodel.components.moe.uccl_ep._utils
• nemo_automodel.components.moe.uccl_ep.buffer

Package Contents

Classes

API

class nemo_automodel.components.moe.uccl_ep.buffer.UCCLBuffer(
    group,
    num_nvl_bytes: int = 0,
    num_rdma_bytes: int = 0,
    low_latency_mode: bool = False,
    num_qps_per_rank: int = 24,
    allow_nvlink_for_low_latency_mode: bool = True,
    allow_mnnvl: bool = False,
    explicitly_destroy: bool = False,
    is_intranode: bool = False
)