Support DP attention with GPT-OSS (#9359)

2025-08-20 16:36:31 -07:00
parent d4bce29721
commit c10b8e6a0f
2 changed files with 6 additions and 5 deletions
@@ -1091,7 +1091,7 @@ class GptOssForCausalLM(nn.Module):
                    if name in params_dict.keys():
                        param = params_dict[name]
                        if "sinks" in name:
-                            start = tp_rank * param.numel()
+                            start = get_attention_tp_rank() * param.numel()
                            param.data.copy_(
                                loaded_weight[start : start + param.numel()]
                            )