Upgrade PJRT CUDA Plugin to version 0.2.3, adding NCCL support for correct sharding.

This commit is contained in:
Tarry Singh 2023-04-12 15:47:06 +00:00
parent 8e43a45a3c
commit 833ff5f28d

View File

@ -182,8 +182,8 @@ cc_import(
http_archive( http_archive(
name = "libpjrt_cuda", name = "libpjrt_cuda",
build_file = "libpjrt_cuda.BUILD.bazel", build_file = "libpjrt_cuda.BUILD.bazel",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v0.2.2/pjrt-cuda_linux-amd64.tar.gz", url = "https://github.com/zml/pjrt-artifacts/releases/download/v0.2.3/pjrt-cuda_linux-amd64.tar.gz",
sha256 = "45e91e8649bcccc43900f90d6dcbf0cfe87d3d2ee76f1763f41263d2ed44d31b", sha256 = "14f39ffef0c9ac529b1a8957750b0b5f5d2f6d310c0d997436051c53a9eb1618",
) )
return mctx.extension_metadata( return mctx.extension_metadata(