| 123456789101112131415161718 |
- # EXPERIMENT: arx RDMA all-reduce for decode-sized TP all-reduces (see arx.py),
- # and with VLLM_ARXBIG=1 arxbig.cu's all-gather for prefill-sized SP gathers.
- # docker compose -f compose/glm53.yaml -f compose/arx.yaml up -d
- services:
- glm53:
- volumes:
- - ../experimental/arx/arx.py:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/arx.py:ro
- - ../experimental/arx/arx_vllm.cu:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/arx_vllm.cu:ro
- - ../experimental/arx/arxbig.cu:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/arxbig.cu:ro
- - ../experimental/arx/cuda_communicator.py:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/cuda_communicator.py:ro
- environment:
- - VLLM_ARX_ALLREDUCE=1
- - VLLM_ARXBIG=${VLLM_ARXBIG:-1}
- - VLLM_ARXBIG_RS=${VLLM_ARXBIG_RS:-1}
- - VLLM_ARXBIG_AG=${VLLM_ARXBIG_AG:-0}
- # Decode all-reduces above 256 KB and up to this size run arx's two-shot mode
- # (reduce-scatter + all-gather over RDMA); 0 sends them to NCCL.
- - VLLM_ARX_TWO_SHOT_MAX_KB=${VLLM_ARX_TWO_SHOT_MAX_KB:-2048}
|
