快捷方式

測試說明

測試(位於 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/test/ 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 將被自動檢測並用於測試。要在具有 GPU 的機器上以僅 CPU 模式執行測試和基準測試,必須在環境中設定 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/test/ 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/bench/ directory
cd bench

python tbe_training_benchmark.py device

文件

訪問 PyTorch 的全面開發者文件

檢視文件

教程

獲取針對初學者和高階開發者的深度教程

檢視教程

資源

查詢開發資源並獲得問題解答

檢視資源