测试说明¶

测试（在 fbgemm_gpu/test/ 目录中）和基准测试（在 fbgemm_gpu/bench/ 目录中）提供了有关如何使用 FBGEMM_GPU 的良好示例。

设置 FBGEMM_GPU 测试环境¶

在构建/安装 FBGEMM_GPU 软件包后，环境可用，还需要安装其他软件包才能使测试正确运行。

# !! Run inside the Conda environment !!

# From the /fbgemm_gpu/ directory
python -m pip install -r requirements.txt

运行 FBGEMM_GPU 测试¶

在构建/安装 FBGEMM_GPU 软件包后运行测试。

# !! Run inside the Conda environment !!

# From the /fbgemm_gpu/ directory
cd test

python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning split_table_batched_embeddings_test.py
python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning quantize_ops_test.py
python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning sparse_ops_test.py
python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning split_embedding_inference_converter_test.py

使用 CUDA 变体进行测试¶

对于 FBGEMM_GPU CUDA 软件包，将自动检测 GPU 并用于测试。要在仅支持 CPU 的模式下在支持 GPU 的机器上运行测试和基准测试，必须在环境中设置 CUDA_VISIBLE_DEVICES=-1。

# !! Run inside the Conda environment !!

# Enable for running in CPU-only mode (when on a GPU-capable machine)
export CUDA_VISIBLE_DEVICES=-1

# Enable for debugging failed kernel executions
export CUDA_LAUNCH_BLOCKING=1

# For operators involving NCCL, if the rpath is not set up correctly for
# libnccl.so.2, LD_LIBRARY_PATH will need to be updated.
export LD_LIBRARY_PATH="/path/to/nccl/lib:${LD_LIBRARY_PATH}"

python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning split_table_batched_embeddings_test.py

使用 ROCm 变体进行测试¶

对于 ROCm 机器，需要使用在环境中设置 FBGEMM_TEST_WITH_ROCM=1 来启用针对 ROCm GPU 的测试。

# !! Run inside the Conda environment !!

# From the /fbgemm_gpu/ directory
cd test

export FBGEMM_TEST_WITH_ROCM=1
# Enable for debugging failed kernel executions
export HIP_LAUNCH_BLOCKING=1

python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning split_table_batched_embeddings_test.py

运行 FBGEMM_GPU 基准测试¶

要运行基准测试。

# !! Run inside the Conda environment !!

# From the /fbgemm_gpu/ directory
cd bench

python split_table_batched_embeddings_benchmark.py uvm