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.buffer
• nemo_automodel.components.moe.uccl_ep._utils

Package Contents

Data

__all__

API

nemo_automodel.components.moe.uccl_ep.__all__

[‘UCCLBuffer’]