From 91f2011f08e039a9e9ec47cd2e1ea64aab884650 Mon Sep 17 00:00:00 2001 From: Lysandre Date: Fri, 29 Aug 2025 15:48:12 +0200 Subject: [PATCH 1/5] VitPose --- .../modeling_vitpose_backbone.py | 39 ++++++++++++------- 1 file changed, 24 insertions(+), 15 deletions(-) diff --git a/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py b/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py index e9a90cdfcb86..28c3b22a468d 100644 --- a/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py +++ b/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py @@ -230,7 +230,28 @@ def forward(self, hidden_states: torch.Tensor, head_mask: Optional[torch.Tensor] return output -# FIXME: refactor? +class VitPoseNaiveMoe(nn.ModuleList): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + + hidden_features = int(config.hidden_size * config.mlp_ratio) + part_features = config.part_features + + for _ in range(self.num_experts): + self += [nn.Linear(hidden_features, part_features)] + + def forward(self, hidden_state, shared_hidden_state, indices): + expert_hidden_state = torch.zeros_like(hidden_state[:, :, -self.part_features :]) + for i in range(self.num_experts): + selected_index = indices == i + current_hidden_state = self[i](hidden_state) * selected_index + expert_hidden_state = expert_hidden_state + current_hidden_state + + hidden_state = torch.cat([shared_hidden_state, expert_hidden_state], dim=-1) + return hidden_state + class VitPoseBackboneMoeMLP(nn.Module): def __init__(self, config: VitPoseBackboneConfig): super().__init__() @@ -245,28 +266,16 @@ def __init__(self, config: VitPoseBackboneConfig): self.fc1 = nn.Linear(in_features, hidden_features) self.act = ACT2FN[config.hidden_act] self.fc2 = nn.Linear(hidden_features, out_features - part_features) - self.drop = nn.Dropout(config.hidden_dropout_prob) self.num_experts = num_experts - experts = [nn.Linear(hidden_features, part_features) for _ in range(num_experts)] - self.experts = nn.ModuleList(experts) + self.experts = VitPoseNaiveMoe(config) def forward(self, hidden_state: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: - expert_hidden_state = torch.zeros_like(hidden_state[:, :, -self.part_features :]) - hidden_state = self.fc1(hidden_state) hidden_state = self.act(hidden_state) shared_hidden_state = self.fc2(hidden_state) indices = indices.view(-1, 1, 1) - - # to support ddp training - for i in range(self.num_experts): - selected_index = indices == i - current_hidden_state = self.experts[i](hidden_state) * selected_index - expert_hidden_state = expert_hidden_state + current_hidden_state - - hidden_state = torch.cat([shared_hidden_state, expert_hidden_state], dim=-1) - + hidden_state = self.experts(hidden_state, shared_hidden_state, indices) return hidden_state From 50c70954f877d3761ca276272c23f3f062a11540 Mon Sep 17 00:00:00 2001 From: Lysandre Date: Fri, 29 Aug 2025 16:00:10 +0200 Subject: [PATCH 2/5] outside --- .../vitpose_backbone/modeling_vitpose_backbone.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py b/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py index 28c3b22a468d..414ad3c27de8 100644 --- a/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py +++ b/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py @@ -242,15 +242,14 @@ def __init__(self, config): for _ in range(self.num_experts): self += [nn.Linear(hidden_features, part_features)] - def forward(self, hidden_state, shared_hidden_state, indices): + def forward(self, hidden_state, indices): expert_hidden_state = torch.zeros_like(hidden_state[:, :, -self.part_features :]) for i in range(self.num_experts): selected_index = indices == i current_hidden_state = self[i](hidden_state) * selected_index expert_hidden_state = expert_hidden_state + current_hidden_state - hidden_state = torch.cat([shared_hidden_state, expert_hidden_state], dim=-1) - return hidden_state + return expert_hidden_state class VitPoseBackboneMoeMLP(nn.Module): def __init__(self, config: VitPoseBackboneConfig): @@ -275,7 +274,10 @@ def forward(self, hidden_state: torch.Tensor, indices: torch.Tensor) -> torch.Te hidden_state = self.act(hidden_state) shared_hidden_state = self.fc2(hidden_state) indices = indices.view(-1, 1, 1) - hidden_state = self.experts(hidden_state, shared_hidden_state, indices) + + expert_hidden_state = self.experts(hidden_state, shared_hidden_state, indices) + hidden_state = torch.cat([shared_hidden_state, expert_hidden_state], dim=-1) + return hidden_state From 5bedb4e584ae1409d491ff68f73895bb82b09f09 Mon Sep 17 00:00:00 2001 From: Lysandre Date: Fri, 29 Aug 2025 16:14:41 +0200 Subject: [PATCH 3/5] outside --- .../models/vitpose_backbone/modeling_vitpose_backbone.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py b/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py index 414ad3c27de8..8ef8fcdad8e8 100644 --- a/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py +++ b/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py @@ -233,8 +233,7 @@ def forward(self, hidden_states: torch.Tensor, head_mask: Optional[torch.Tensor] class VitPoseNaiveMoe(nn.ModuleList): def __init__(self, config): super().__init__() - self.top_k = config.num_experts_per_tok - self.num_experts = config.num_local_experts + self.num_experts = config.num_experts hidden_features = int(config.hidden_size * config.mlp_ratio) part_features = config.part_features From c89b8b22156ec722586eff0502edc08ccb804faf Mon Sep 17 00:00:00 2001 From: Lysandre Date: Fri, 29 Aug 2025 16:15:14 +0200 Subject: [PATCH 4/5] outside --- .../models/vitpose_backbone/modeling_vitpose_backbone.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py b/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py index 8ef8fcdad8e8..713f27ef0802 100644 --- a/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py +++ b/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py @@ -274,7 +274,7 @@ def forward(self, hidden_state: torch.Tensor, indices: torch.Tensor) -> torch.Te shared_hidden_state = self.fc2(hidden_state) indices = indices.view(-1, 1, 1) - expert_hidden_state = self.experts(hidden_state, shared_hidden_state, indices) + expert_hidden_state = self.experts(hidden_state, indices) hidden_state = torch.cat([shared_hidden_state, expert_hidden_state], dim=-1) return hidden_state From cb16f279a0c8f0334a1f8ccd7d220e90190abfcc Mon Sep 17 00:00:00 2001 From: Lysandre Date: Fri, 29 Aug 2025 16:16:07 +0200 Subject: [PATCH 5/5] fix --- .../models/vitpose_backbone/modeling_vitpose_backbone.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py b/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py index 713f27ef0802..26ac5a2f50d8 100644 --- a/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py +++ b/src/transformers/models/vitpose_backbone/modeling_vitpose_backbone.py @@ -234,6 +234,7 @@ class VitPoseNaiveMoe(nn.ModuleList): def __init__(self, config): super().__init__() self.num_experts = config.num_experts + self.part_features = config.part_features hidden_features = int(config.hidden_size * config.mlp_ratio) part_features = config.part_features