Fix incorrect LoRA weight loading for fused gate_up_proj (#6734)

This commit is contained in:
Lifu Huang
2025-05-31 13:41:44 -07:00
committed by GitHub
parent 888cb175a6
commit 094fbdacd5
4 changed files with 29 additions and 14 deletions
+16 -9
View File
@@ -296,23 +296,30 @@ class Idefics2VisionTransformer(nn.Module):
def compute_cu_seqlens(
self,
tgt_sizes: Optional[torch.Tensor] = None,
atch_attention_mask: Optional[torch.BoolTensor] = None,
input_embeds: Optional[torch.Tensor] = None,
) -> torch.Tensor:
# shape: (batch_size,)
if tgt_sizes is not None:
patch_len = tgt_sizes[:, 0] * tgt_sizes[:, 1]
seqlen = tgt_sizes[:, 0] * tgt_sizes[:, 1]
elif input_embeds is not None:
seqlen = torch.full(
size=(input_embeds.shape[0],),
fill_value=input_embeds.shape[1],
dtype=torch.int32,
device=input_embeds.device,
)
else:
patch_len = atch_attention_mask[:, :, 0].sum(dim=1) * atch_attention_mask[
:, 0, :
].sum(dim=1)
raise ValueError(
"Either `tgt_sizes` or `input_embeds` must be provided to compute cu_seqlens."
)
cu_seqlens = torch.cat(
[
torch.tensor([0], device=patch_len.device, dtype=torch.int32),
torch.cumsum(patch_len, dim=0, dtype=torch.int32),
torch.tensor([0], device=seqlen.device, dtype=torch.int32),
torch.cumsum(seqlen, dim=0, dtype=torch.int32),
],
dim=0,
).to(patch_len.device)
).to(seqlen.device)
return cu_seqlens
def forward(
@@ -326,7 +333,7 @@ class Idefics2VisionTransformer(nn.Module):
patch_attention_mask=patch_attention_mask,
tgt_sizes=tgt_sizes,
)
cu_seqlens = self.compute_cu_seqlens(tgt_sizes, patch_attention_mask)
cu_seqlens = self.compute_cu_seqlens(tgt_sizes, hidden_states)
encoder_outputs = self.encoder(
hidden_states,
cu_seqlens=cu_seqlens,