Skip to content

Hacky tensor parallel for Vision Models #791

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

Draft
wants to merge 3 commits into
base: dev
Choose a base branch
from
Draft

Hacky tensor parallel for Vision Models #791

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
764e257
Pixtral TP architecture.py
Ph0rk0z May 25, 2025
a85c4fb
Bypass Torch Restrictions for inference mode.
Ph0rk0z May 25, 2025
4430729
architecture.py - Tested qwen2 VL 72b TP and it works.
Ph0rk0z May 25, 2025
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
2 changes: 2 additions & 0 deletions exllamav2/architecture.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -319,6 +319,7 @@ class Params:
layer_keys_llama_mlp
self.lm.expect_keys += \
expect_keys_llama
self.lm.supports_tp = True

self.vt_prefix = "vision_tower."
self.vt.keys.update({
Expand Down Expand Up @@ -478,6 +479,7 @@ class Params:
self.lm.attention_bias_qkv = True
self.lm.mrope = True
self.lm.rope_freq_half = True
self.lm.supports_tp = True

self.vt_prefix = "visual."
if arch_string == "Qwen2VLForConditionalGeneration":
Expand Down
22 changes: 15 additions & 7 deletions exllamav2/tensor_p.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -351,16 +351,24 @@ def allgather(
return bc_tensors


def copy_pinned(
self,
buffer: int,
inputs: torch.Tensor
):
# def copy_pinned(
# self,
# buffer: int,
# inputs: torch.Tensor
# ):
# pt = self.pinned_temp[buffer][:inputs.numel()]
# pt = pt.view(inputs.shape)
# pt.copy_(inputs)
# return pt

def copy_pinned(self, buffer: int, inputs: torch.Tensor):
pt = self.pinned_temp[buffer][:inputs.numel()]
pt = pt.view(inputs.shape)
pt.copy_(inputs)
return pt

# Bypass PyTorch entirely - direct memory copy
import ctypes
ctypes.memmove(pt.data_ptr(), inputs.data_ptr(), inputs.numel() * inputs.element_size())
return pt

def add_residual(
self,
Expand Down