Skip to content

Support multi node nvlink fabric - #23

Open
lyppg wants to merge 1 commit into
MoonshotAI:masterfrom
bytedance-iaas:jxp/support_multinode_fabric
Open

Support multi node nvlink fabric#23
lyppg wants to merge 1 commit into
MoonshotAI:masterfrom
bytedance-iaas:jxp/support_multinode_fabric

Conversation

@lyppg

@lyppg lyppg commented Jul 31, 2026

Copy link
Copy Markdown

Current MoonEP supports each rank of an EP group lives on the same node. This PR adds support for multi node with nvlink via CUDA fabric handles, valid for every GPU in the same NVLink/IMEX domain, e.g., GB200.

Tests

Run the benchmark on single node and two nodes.

# single node
torchrun --nproc_per_node=4 benchmarks/bench_prefetch.py

Config                   E    B       H     Hp   SMs  Slots   Data(MB)  Worst(us)  BW(GB/s)   CommBW
-----------------------------------------------------------------------------------------------------
slots_1                 32   14    3584   3072    32      1      44.04      34.51   1276.08   638.04
slots_2                 32   14    3584   3072    32      2      88.08      62.30   1413.79   706.90
slots_3                 32   14    3584   3072    32      3     132.12      90.28   1463.48   731.74
slots_5                 32   14    3584   3072    32      5     220.20     146.85   1499.50   749.75
ramp_0_3                32   14    3584   3072    32     12     528.48     343.52   1538.45   769.22
mixed                   32   14    3584   3072    32     12     528.48     344.71   1533.12   766.56
heavy                   32   14    3584   3072    32     17     748.68     484.67   1544.72   772.36
full_3x8                32   14    3584   3072    32     24    1056.96     681.36   1551.26   775.63
dense_B8                32    8    3584   3072    32     15     660.60     429.11   1539.46   769.73
epn4_full               16   14    3584   3072    32     12     528.48     343.51   1538.48   769.24
epn16_mixed             64   16    3584   3072    32     24    1056.96     680.12   1554.09   777.04
thin_7168x128           32   16    7168    128    32     12      44.04      34.85   1263.72   631.86
tall_1024x3072          32   14    1024   3072    32     12     150.99     102.85   1468.05   734.02
tiny_512x512            32   16     512    512    32     17      17.83      17.68   1008.25   504.12

# two GB200 nodes, --node_rank=1 on another node
torchrun --nnodes=2 --node_rank=0 --nproc_per_node=4     --master_addr=$MASTER_ADDR --master_port=$PORT benchmarks/bench_prefetch.py

Config                   E    B       H     Hp   SMs  Slots   Data(MB)  Worst(us)  BW(GB/s)   CommBW
-----------------------------------------------------------------------------------------------------
slots_1                 64   14    3584   3072    32      1      44.04      35.27   1248.76   624.38
slots_2                 64   14    3584   3072    32      2      88.08      63.17   1394.38   697.19
slots_3                 64   14    3584   3072    32      3     132.12      91.21   1448.46   724.23
slots_5                 64   14    3584   3072    32      5     220.20     147.27   1495.18   747.59
ramp_0_3                64   14    3584   3072    32     12     528.48     344.33   1534.80   767.40
mixed                   64   14    3584   3072    32     12     528.48     344.97   1531.97   765.99
heavy                   64   14    3584   3072    32     17     748.68     484.10   1546.55   773.27
full_3x8                64   14    3584   3072    32     24    1056.96     681.24   1551.53   775.77
dense_B8                64    8    3584   3072    32     15     660.60     428.49   1541.70   770.85
epn4_full               32   14    3584   3072    32     12     528.48     344.73   1533.05   766.52
epn16_mixed            128   16    3584   3072    32     24    1056.96     681.67   1550.56   775.28
thin_7168x128           64   16    7168    128    32     12      44.04      35.51   1240.37   620.19
tall_1024x3072          64   14    1024   3072    32     12     150.99     103.75   1455.32   727.66
tiny_512x512            64   16     512    512    32     17      17.83      18.67    954.93   477.46

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant