"megatron/git@developer.sourcefind.cn:OpenDAS/megatron-lm.git" did not exist on "1005e9d34d1c04472a388a434add424176fec46d"
Unverified Commit c31db6d6 authored by Andrey Talman's avatar Andrey Talman Committed by GitHub
Browse files

Using conda metapackage when installing cuda 11.6 and 11.7 (#6371)

* Testing conda metapackage

* Testing

* Testing pytorch-cuda

* Testing

* Testing

* Testing

* Testing

* Testing

* Using new metapackage

* Testing

* Simplifing logic

* Further simplifying
parent 6e535db2
......@@ -22,11 +22,10 @@ else
echo "Using CUDA $CUDA_VERSION as determined by CU_VERSION"
version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
cuda_toolkit_pckg="cudatoolkit"
if [[ "$CU_VERSION" == cu116 ]]; then
cuda_toolkit_pckg="cuda"
cudatoolkit="nvidia::cudatoolkit=${version}"
if [[ "$version" == "11.6" || "$version" == "11.7" ]]; then
cudatoolkit=" pytorch-cuda=${version}"
fi
cudatoolkit="nvidia::${cuda_toolkit_pckg}=${version}"
fi
case "$(uname -s)" in
......
......@@ -261,8 +261,11 @@ setup_conda_cudatoolkit_constraint() {
export CONDA_BUILD_VARIANT="cpu"
else
case "$CU_VERSION" in
cu117)
export CONDA_CUDATOOLKIT_CONSTRAINT="- pytorch-cuda=11.7 # [not osx]"
;;
cu116)
export CONDA_CUDATOOLKIT_CONSTRAINT="- cuda >=11.6,<11.7 # [not osx]"
export CONDA_CUDATOOLKIT_CONSTRAINT="- pytorch-cuda=11.6 # [not osx]"
;;
cu113)
export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.3,<11.4 # [not osx]"
......@@ -290,8 +293,11 @@ setup_conda_cudatoolkit_plain_constraint() {
export CMAKE_USE_CUDA=0
else
case "$CU_VERSION" in
cu117)
export CONDA_CUDATOOLKIT_CONSTRAINT="pytorch-cuda=11.7"
;;
cu116)
export CONDA_CUDATOOLKIT_CONSTRAINT="cuda=11.6"
export CONDA_CUDATOOLKIT_CONSTRAINT="pytorch-cuda=11.6"
;;
cu113)
export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=11.3"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment