Skip to content

Fix/janus pro flash attention #152

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

Open
vadseshu wants to merge 2 commits into
base: develop
Choose a base branch
from
Open

Fix/janus pro flash attention #152

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
549e9ec
Added gfx950a to pyt_mochi_inference.ubuntu dockerfile to enable fa
vadseshu1 Apr 16, 2026
eb07c6c
Changed FA steps and fixing runtime failures wrt transformers-5.x
vadseshu Apr 22, 2026
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
47 changes: 31 additions & 16 deletions docker/pyt_janus_pro_inference.ubuntu.amd.Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,27 +43,42 @@ ENV HIP_ARCHITECTURES=${MAD_SYSTEM_GPU_ARCHITECTURE}
RUN echo HIP_ARCHITECTURES = ${HIP_ARCHITECTURES}

# Install flash attention
ARG BUILD_FA="1"
ARG FA_BRANCH="v3.0.0.r1-cktile"
ARG FA_REPO="https://github.com/ROCm/flash-attention.git"
RUN if [ "$BUILD_FA" = "1" ]; then \
cd ${APP_DIR} \
&& pip uninstall -y flash-attention \
&& rm -rf flash-attention \
&& git clone ${FA_REPO} \
&& cd flash-attention \
&& git checkout ${FA_BRANCH} \
&& git submodule update --init \
&& GPU_ARCHS=${HIP_ARCHITECTURES} python3 setup.py bdist_wheel --dist-dir=dist \
&& pip install dist/*.whl \
&& python -c "import flash_attn; print(f'Flash Attention version == {flash_attn.__version__}')"; \
fi
#ARG BUILD_FA="1"
#ARG FA_BRANCH="v3.0.0.r1-cktile"
#ARG FA_REPO="https://github.com/ROCm/flash-attention.git"
#RUN if [ "$BUILD_FA" = "1" ]; then \
# cd ${APP_DIR} \
# && pip uninstall -y flash-attention \
# && rm -rf flash-attention \
# && git clone ${FA_REPO} \
# && cd flash-attention \
# && git checkout ${FA_BRANCH} \
# && git submodule update --init \
# && GPU_ARCHS=${HIP_ARCHITECTURES} python3 setup.py bdist_wheel --dist-dir=dist \
# && pip install dist/*.whl \
# && python -c "import flash_attn; print(f'Flash Attention version == {flash_attn.__version__}')"; \
# fi
Comment on lines 45 to +60
Copy link

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

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

The large commented-out flash-attention build block (BUILD_FA/FA_BRANCH/FA_REPO) looks like dead code after the new install path was added. To keep the Dockerfile maintainable, either remove the commented block or re-enable it behind a build ARG so there is a single authoritative installation path.

Copilot uses AI. Check for mistakes.

# Install Janus-pro
# install flash attention
ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"

RUN git clone https://github.com/ROCm/flash-attention.git &&\
cd flash-attention &&\
python setup.py install
Comment on lines +64 to +67
Copy link

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

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

The flash-attention install is now unpinned (no branch/SHA) and skips git submodule update --init, and python setup.py install builds from the live HEAD. This makes builds non-reproducible and is likely to break (other Dockerfiles in this repo pin a ref and init submodules before building a wheel). Consider restoring pinning (FA_BRANCH/FA_SHA), initializing submodules, and installing via a built wheel (or pip install .) with explicit GPU_ARCHS=${HIP_ARCHITECTURES} to avoid long/incorrect builds.

Suggested change
RUN git clone https://github.com/ROCm/flash-attention.git &&\
cd flash-attention &&\
python setup.py install
ARG FA_REPO="https://github.com/ROCm/flash-attention.git"
ARG FA_BRANCH="v3.0.0.r1-cktile"
ARG FA_SHA=""
RUN cd ${APP_DIR} \
&& rm -rf flash-attention \
&& git clone --branch ${FA_BRANCH} ${FA_REPO} flash-attention \
&& cd flash-attention \
&& if [ -n "${FA_SHA}" ]; then git checkout ${FA_SHA}; fi \
&& git submodule update --init --recursive \
&& GPU_ARCHS=${HIP_ARCHITECTURES} python3 -m pip wheel . --no-build-isolation --wheel-dir=dist \
&& pip install dist/*.whl \
&& python -c "import flash_attn; print(f'Flash Attention version == {flash_attn.__version__}')"

Copilot uses AI. Check for mistakes.

# Install Janus-pro (patches for Hugging Face transformers 5.x):
# - @dataclass rejects mutable default dict on PretrainedConfig subclasses; use default_factory
# - PreTrainedModel.post_init() sets all_tied_weights_keys (required by from_pretrained / loading)
RUN cd ${APP_DIR} \
&& git clone https://github.com/deepseek-ai/Janus.git \
&& cd Janus \
&& sed -i 's/^torch==/# torch==/' requirements.txt \
&& for f in janus/models/modeling_vlm.py janus/janusflow/models/modeling_vlm.py; do \
sed -i '/^from transformers.configuration_utils import PretrainedConfig/a from dataclasses import field' "$f" \
&& sed -i 's/params: AttrDict = {}/params: AttrDict = field(default_factory=dict)/g' "$f"; \
done \
&& sed -i '/self\.language_model = LlamaForCausalLM(language_config)/a\ self.post_init()' janus/models/modeling_vlm.py \
&& sed -i '/self\.vision_gen_dec_aligner = nn.Linear(2048, 768, bias=True)/a\ self.post_init()' janus/janusflow/models/modeling_vlm.py \
Comment on lines +78 to +81
Copy link

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

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

The Janus/transformers 5.x compatibility changes are applied via sed against specific source lines. sed will still exit successfully even if the patterns stop matching upstream, which can silently produce an unpatched image and reintroduce the runtime errors this PR is addressing. Consider pinning Janus to a known commit and/or adding post-sed assertions (e.g., grep for default_factory=dict and self.post_init()) so the build fails if the patch didn’t apply.

Suggested change
&& sed -i 's/params: AttrDict = {}/params: AttrDict = field(default_factory=dict)/g' "$f"; \
done \
&& sed -i '/self\.language_model = LlamaForCausalLM(language_config)/a\ self.post_init()' janus/models/modeling_vlm.py \
&& sed -i '/self\.vision_gen_dec_aligner = nn.Linear(2048, 768, bias=True)/a\ self.post_init()' janus/janusflow/models/modeling_vlm.py \
&& sed -i 's/params: AttrDict = {}/params: AttrDict = field(default_factory=dict)/g' "$f" \
&& grep -q 'from dataclasses import field' "$f" \
&& grep -q 'params: AttrDict = field(default_factory=dict)' "$f"; \
done \
&& sed -i '/self\.language_model = LlamaForCausalLM(language_config)/a\ self.post_init()' janus/models/modeling_vlm.py \
&& grep -q 'self.post_init()' janus/models/modeling_vlm.py \
&& sed -i '/self\.vision_gen_dec_aligner = nn.Linear(2048, 768, bias=True)/a\ self.post_init()' janus/janusflow/models/modeling_vlm.py \
&& grep -q 'self.post_init()' janus/janusflow/models/modeling_vlm.py \

Copilot uses AI. Check for mistakes.
&& pip install -e . \
&& pip install datasets \
&& cd ..
Expand Down
2 changes: 1 addition & 1 deletion docker/pyt_mochi_inference.ubuntu.amd.Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ RUN mkdir -p $WORKSPACE_DIR
WORKDIR $WORKSPACE_DIR

ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx1100;gfx1101;gfx1200;gfx1201
ARG PYTORCH_ROCM_ARCH=gfx950;gfx90a;gfx942;gfx1100;gfx1101;gfx1200;gfx1201
RUN git clone ${FA_REPO}
RUN cd flash-attention \
&& git submodule update --init \
Expand Down