topk casting

Author: rebel-thkim

examples/text2text-generation/run_gpt_oss.py

@@ -66,7 +66,7 @@ def main(
             rbln_tensor_parallel_size=tensor_parallel_size,
             rbln_kvcache_partition_len=kvcache_partition_len,
             config=target_config,
-            dtype=torch.float32,
+            # dtype=torch.float32,
         )
         model.save_pretrained(os.path.basename(model_id))
     else:

src/optimum/rbln/transformers/models/gemma3/modeling_gemma3.py

@@ -97,11 +97,9 @@ def redirect(func):
 
     def can_generate(self):
         return True
-    
-    
+
     @classmethod
     def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
-        
         with no_init_weights():
             model_cls_name = model.model.language_model.__class__.__name__
             causal_model_cls_name = model_cls_name.replace("TextModel", "ForCausalLM")

src/optimum/rbln/transformers/models/gpt_oss/gpt_oss_architecture.py

@@ -68,11 +68,17 @@ def __init__(self, model):
         self.weight = model.weight
         self.bias = model.bias
 
+    def casted_top_K(self, router_logits, hidden_states):
+        logits = router_logits.to(torch.float32)
+        router_top_value, router_indices = torch.topk(logits, self.top_k, dim=-1)
+
+        return router_top_value.to(hidden_states.dtype), router_indices
+
     def forward(self, hidden_states):
         hidden_states = hidden_states.reshape(-1, self.hidden_dim)
         router_logits = F.linear(hidden_states, self.weight, self.bias)  # (seq_len, num_experts)
-        router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (seq_len, top_k)
-        router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype)
+        router_top_value, router_indices = self.casted_top_K(router_logits, hidden_states)
+        router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=hidden_states.dtype)
         router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
 
         return router_scores, router_indices
@@ -120,7 +126,9 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
         hidden_states = hidden_states.repeat(num_experts, 1)
         hidden_states = hidden_states.view(num_experts, -1, self.hidden_size)
 
-        gate_up = torch.bmm(hidden_states, self.gate_up_proj.to(hidden_states.dtype)) + self.gate_up_proj_bias[..., None, :].to(hidden_states.dtype)
+        gate_up = torch.bmm(hidden_states, self.gate_up_proj.to(hidden_states.dtype)) + self.gate_up_proj_bias[
+            ..., None, :
+        ].to(hidden_states.dtype)
         gate, up = gate_up[..., ::2], gate_up[..., 1::2]
         gate = gate.clamp(min=None, max=self.limit)
         up = up.clamp(min=-self.limit, max=self.limit)

src/optimum/rbln/transformers/models/grounding_dino/modeling_grounding_dino.py

@@ -204,8 +204,7 @@ def save_torch_artifacts(
         save_dict["bbox_embed"] = model.bbox_embed.state_dict()
 
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
-        
-    
+
     @classmethod
     def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
         model.encoder = model.model.encoder

src/optimum/rbln/transformers/models/siglip/modeling_siglip.py

@@ -66,7 +66,9 @@ class RBLNSiglipVisionModel(RBLNModel):
     _tp_support = False
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNSiglipVisionModelConfig) -> torch.nn.Module:
+    def _wrap_model_if_needed(
+        cls, model: torch.nn.Module, rbln_config: RBLNSiglipVisionModelConfig
+    ) -> torch.nn.Module:
         wrapper_cfg = {
             "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
             "output_hidden_states": rbln_config.output_hidden_states,