ByteDance-Seed / ByteDance-Seed/Triton-distributed
test_gemm_rs performance issue
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 172
- PR merge metrics
- No merged PRs in 30d
Description
I run on a single node with 8*H100 GPUs with 4 nics (4\*cx7-4\*400gb/s).
`bash ./launch.sh ./python/triton_dist/test/nvidia/test_gemm_rs.py 8192 8192 29568`
I found that the performance of Triton-Distributed is far inferior to PyTorch. Is there any configuration error?
commit: b10fa2e7a7e561431687a2d71bfaa09f882fa122
M, N, K (8192, 8192, 295568) are the dimensions mentioned above, which will have errors and do not match the given data.

The performance of `ag_gemm` is improved, but (8192,53248, 16384) will hang.
`bash ./launch.sh ./python/triton_dist/test/nvidia/test_ag_gemm.py --case perf --shape_id LLaMA-3.1-405B`
The startup script is as follows (launch.sh):
```
#!/bin/bash
export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
export CUDA_LAUNCH_BLOCKING=${CUDA_LAUNCH_BLOCKING:-0}
export TORCH_CPP_LOG_LEVEL=1
export NVSHMEM_DISABLE_CUDA_VMM=1
export NCCL_IB_HCA=^=mlx5_bond_0
export NCCL_SOCKET_IFNAME=bond0
export NCCL_DEBUG=INFO
export NCCL_IB_GID_INDEX=3
SCRIPT_DIR=/xx/triton-distributed/Triton-distributed
NVSHMEM_ROOT=${SCRIPT_DIR}/3rdparty/nvshmem/build/install
OMPI_BUILD=${SCRIPT_DIR}/shmem/rocshmem_bind/ompi_build/install/ompi
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${NVSHMEM_ROOT}/lib:${OMPI_BUILD}/lib
export TRITON_CACHE_DIR=${SCRIPT_DIR}/triton_cache
export NVSHMEM_HOME=${NVSHMEM_ROOT}
export PYTHONPATH=$PYTHONPATH:${SCRIPT_DIR}/python
mkdir -p ${SCRIPT_DIR}/triton_cache
export NVSHMEM_SYMMETRIC_SIZE=${NVSHMEM_SYMMETRIC_SIZE:-1000000000}
export LD_LIBRARY_PATH=${NVSHMEM_DIR}/lib:$LD_LIBRARY_PATH
export NVSHMEM_DISABLE_CUDA_VMM=${NVSHMEM_DISABLE_CUDA_VMM:-1} # moving from cpp to shell
export NVSHMEM_BOOTSTRAP=UID
export NVSHMEM_IFNAME=bond0
export NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME=bond0
export NVSHMEM_DEBUG=INFO
export NVSHMEM_ENABLE_NIC_PE_MAPPING=1
export NVSHMEM_HCA_PE_MAPPING="mlx5_0:1,mlx5_1:1,mlx5_4:1,mlx5_5:1"
nproc_per_node=${ARNOLD_WORKER_GPU:=$(nvidia-smi --list-gpus | wc -l)}
nnodes=${ARNOLD_WORKER_NUM:=1}
node_rank=${ARNOLD_ID:=0}
master_addr=${ARNOLD_WORKER_0_HOST:="xx.xxx.xx.x"}
if [ -z ${ARNOLD_WORKER_0_PORT} ]; then
master_port="23456"
else
master_port=$(echo "$ARNOLD_WORKER_0_PORT" | cut -d "," -f 1)
fi
additional_args="--rdzv_endpoint=${master_addr}:${master_port}"
CMD="torchrun \
--node_rank=${node_rank} \
--nproc_per_node=${nproc_per_node} \
--nnodes=${nnodes} \
${DIST_TRITON_EXTRA_TORCHRUN_ARGS} \
${additional_args} \
${DIST_TRITON_EXTRA_TORCHRUN_ARGS} \
$@"
echo ${CMD}
${CMD}
ret=$?
exit $ret
```
Contributor guide
Research direction
Start by reproducing the commands in launch.sh for python/triton_dist/test/nvidia/test_gemm_rs.py and test_ag_gemm.py at commit b10fa2e7a7e561431687a2d71bfaa09f882fa122. Compare the reported dimensions, performance, and hang behavior with the supplied environment settings. Done means the configuration or implementation cause is identified and the affected test cases behave as expected.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- bash, python
- Domain
- distributed-systems, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100