Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix nccl regression on PyTorch 2.3 upgrade #2099

Merged
merged 5 commits into from
Jul 8, 2024
Merged

Fix nccl regression on PyTorch 2.3 upgrade #2099

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2502ce4
fix nccl issue
fxmarty Jun 20, 2024
a76b6f4
add note in dockerfile
fxmarty Jun 20, 2024
27a3792
use v2.22.3 that also fixes @samsamoa's repro
fxmarty Jun 20, 2024
62a1ddb
poetry actually can't handle the conflict between torch and nccl
fxmarty Jun 20, 2024
a1695ce
set LD_PRELOAD
fxmarty Jun 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,9 @@ RUN cargo build --profile release-opt
# Adapted from: https://github.com/pytorch/pytorch/blob/master/Dockerfile
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 as pytorch-install

# NOTE: When updating PyTorch version, beware to remove `pip install nvidia-nccl-cu12==2.22.3` below in the Dockerfile. Context: https://github.com/huggingface/text-generation-inference/pull/2099
ARG PYTORCH_VERSION=2.3.0

ARG PYTHON_VERSION=3.10
# Keep in sync with `server/pyproject.toml
ARG CUDA_VERSION=12.1
Expand Down Expand Up @@ -232,7 +234,10 @@ COPY server/Makefile server/Makefile
RUN cd server && \
make gen-server && \
pip install -r requirements_cuda.txt && \
pip install ".[bnb, accelerate, quantize, peft, outlines]" --no-cache-dir
pip install ".[bnb, accelerate, quantize, peft, outlines]" --no-cache-dir && \
pip install nvidia-nccl-cu12==2.22.3
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would have liked to use pyproject.toml for that, but poetry disapproves of conflict handling python-poetry/poetry#697 (comment)


ENV LD_PRELOAD=/opt/conda/lib/python3.10/site-packages/nvidia/nccl/lib/libnccl.so.2
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need to preload?

Copy link
Contributor Author

@fxmarty fxmarty Jul 8, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Otherwise, the shared object is not used. The current base docker image of TGI is nvidia/cuda:12.1.0-base-ubuntu22.04, where there is no libnccl.so anywhere and it is not loaded by pytorch either, although we have /opt/conda/lib/libcudart.so.12.1.105 etc. COPY --from=pytorch-install /opt/conda /opt/conda does not seem to copy any libnccl.so. Weird.


# Deps before the binaries
# The binaries change on every build given we burn the SHA into them
Expand Down
2 changes: 1 addition & 1 deletion server/Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,5 +34,5 @@ run-dev:
SAFETENSORS_FAST_GPU=1 python -m torch.distributed.run --nproc_per_node=2 text_generation_server/cli.py serve bigscience/bloom-560m --sharded

export-requirements:
poetry export -o requirements_cuda.txt --without-hashes
poetry export -o requirements_cuda.txt --without-hashes --with cuda
poetry export -o requirements_rocm.txt --without-hashes
Loading