Skip to content

Commit 45b37b7

Browse files
author
Aleksandr Malyshev
committed
AI day rocm base
Signed-off-by: Aleksandr Malyshev <maleksan@amd.com>
1 parent c806f1d commit 45b37b7

1 file changed

Lines changed: 36 additions & 36 deletions

File tree

‎docker/Dockerfile.rocm_base‎

Lines changed: 36 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -69,15 +69,15 @@ RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cyth
6969
# && make package
7070
# RUN mkdir -p /app/install && cp /app/hipBLASLt/build/release/*.deb /app/hipBLAS-common/build/*.deb /app/install
7171

72-
# FROM base AS build_triton
73-
# ARG TRITON_BRANCH
74-
# ARG TRITON_REPO
75-
# RUN git clone ${TRITON_REPO}
76-
# RUN cd triton \
77-
# && git checkout ${TRITON_BRANCH} \
78-
# # && cd python \
79-
# && python3 setup.py bdist_wheel --dist-dir=python/dist
80-
# RUN mkdir -p /app/install && cp /app/triton/python/dist/*.whl /app/install
72+
FROM base AS build_triton
73+
ARG TRITON_BRANCH
74+
ARG TRITON_REPO
75+
RUN git clone ${TRITON_REPO}
76+
RUN cd triton \
77+
&& git checkout ${TRITON_BRANCH} \
78+
# && cd python \
79+
&& python3 setup.py bdist_wheel --dist-dir=python/dist
80+
RUN mkdir -p /app/install && cp /app/triton/python/dist/*.whl /app/install
8181

8282
FROM base AS build_amdsmi
8383
RUN cd /opt/rocm/share/amd_smi \
@@ -91,24 +91,24 @@ ARG PYTORCH_REPO
9191
ARG PYTORCH_VISION_REPO
9292
ARG FA_BRANCH
9393
ARG FA_REPO
94-
# RUN git clone ${PYTORCH_REPO} pytorch
95-
# RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \
96-
# pip install -r requirements.txt && git submodule update --init --recursive \
97-
# && python3 tools/amd_build/build_amd.py \
98-
# && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \
99-
# && pip install dist/*.whl
100-
# RUN git clone ${PYTORCH_VISION_REPO} vision
101-
# RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
102-
# && python3 setup.py bdist_wheel --dist-dir=dist \
103-
# && pip install dist/*.whl
94+
RUN git clone ${PYTORCH_REPO} pytorch
95+
RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \
96+
pip install -r requirements.txt && git submodule update --init --recursive \
97+
&& python3 tools/amd_build/build_amd.py \
98+
&& CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \
99+
&& pip install dist/*.whl
100+
RUN git clone ${PYTORCH_VISION_REPO} vision
101+
RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
102+
&& python3 setup.py bdist_wheel --dist-dir=dist \
103+
&& pip install dist/*.whl
104104
RUN git clone ${FA_REPO}
105105
RUN cd flash-attention \
106106
&& git checkout ${FA_BRANCH} \
107107
&& git submodule update --init \
108108
&& GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist
109-
# RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
110-
# && cp /app/vision/dist/*.whl /app/install \
111-
# && cp /app/flash-attention/dist/*.whl /app/install
109+
RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
110+
&& cp /app/vision/dist/*.whl /app/install \
111+
&& cp /app/flash-attention/dist/*.whl /app/install
112112
RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
113113

114114
FROM base AS build_aiter
@@ -128,12 +128,12 @@ FROM base AS debs
128128
RUN mkdir /app/debs
129129
# RUN --mount=type=bind,from=build_hipblaslt,src=/app/install/,target=/install \
130130
# cp /install/*.deb /app/debs
131-
# RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
132-
# cp /install/*.whl /app/debs
131+
RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
132+
cp /install/*.whl /app/debs
133133
RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
134134
cp /install/*.whl /app/debs
135-
# RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
136-
# cp /install/*.whl /app/debs
135+
RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
136+
cp /install/*.whl /app/debs
137137
RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
138138
cp /install/*.whl /app/debs
139139

@@ -143,12 +143,12 @@ FROM base AS final
143143
# && perl -p -i -e 's/, hipblas-common-dev \([^)]*?\), /, /g' /var/lib/dpkg/status \
144144
# && perl -p -i -e 's/, hipblaslt-dev \([^)]*?\), /, /g' /var/lib/dpkg/status \
145145
# && perl -p -i -e 's/, hipblaslt \([^)]*?\), /, /g' /var/lib/dpkg/status
146-
# RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
147-
# pip install /install/*.whl
146+
RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
147+
pip install /install/*.whl
148148
RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
149149
pip install /install/*.whl
150-
# RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
151-
# pip install /install/*.whl
150+
RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
151+
pip install /install/*.whl
152152
RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
153153
pip install /install/*.whl
154154

@@ -170,12 +170,12 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
170170
# && echo "HIPBLAS_COMMON_BRANCH: ${HIPBLAS_COMMON_BRANCH}" >> /app/versions.txt \
171171
# && echo "HIPBLASLT_BRANCH: ${HIPBLASLT_BRANCH}" >> /app/versions.txt \
172172
# && echo "LEGACY_HIPBLASLT_OPTION: ${LEGACY_HIPBLASLT_OPTION}" >> /app/versions.txt \
173-
# && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
174-
# && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
175-
# && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \
176-
# && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \
177-
# && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \
178-
# && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \
173+
&& echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
174+
&& echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
175+
&& echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \
176+
&& echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \
177+
&& echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \
178+
&& echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \
179179
&& echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
180180
&& echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
181181
&& echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt

0 commit comments

Comments
 (0)