Fix Gpt3 MultiHeadAttention out projection dimension

hsuan-lun-chiang · hsuan-lun-chiang · commit ebef575c4390 · 2025-12-11T03:07:13.000Z
diff --git a/src/MaxText/layers/gpt3.py b/src/MaxText/layers/gpt3.py
@@ -246,7 +246,7 @@ def __init__(
       self.key = self.create_projection_layer(feature_dim, (self.num_heads, self.head_dim), ("embed", "heads", "kv"))
       self.value = self.create_projection_layer(feature_dim, (self.num_heads, self.head_dim), ("embed", "heads", "kv"))
     self.out = self.create_projection_layer(
-        (self.num_heads, self.head_dim), self.num_heads * self.head_dim, ("heads", "kv", "embed"), axis=(-2, -1)
+        (self.num_heads, self.head_dim), feature_dim[-1], ("heads", "kv", "embed"), axis=(-2, -1)
     )
     self.attention_op = AttentionOp(
         config=config,

Original file line number	Diff line number	Diff line change
`@@ -246,7 +246,7 @@ def __init__(`
`246`	`246`	`self.key = self.create_projection_layer(feature_dim, (self.num_heads, self.head_dim), ("embed", "heads", "kv"))`
`247`	`247`	`self.value = self.create_projection_layer(feature_dim, (self.num_heads, self.head_dim), ("embed", "heads", "kv"))`
`248`	`248`	`self.out = self.create_projection_layer(`
`249`		`- (self.num_heads, self.head_dim), self.num_heads * self.head_dim, ("heads", "kv", "embed"), axis=(-2, -1)`
	`249`	`+ (self.num_heads, self.head_dim), feature_dim[-1], ("heads", "kv", "embed"), axis=(-2, -1)`
`250`	`250`	`)`
`251`	`251`	`self.attention_op = AttentionOp(`
`252`	`252`	`config=config,`