Skip to content

Commit bbb2304

Browse files
authored
Fix issue of attention.core_attention is None (#334)
Signed-off-by: Yue <[email protected]>
1 parent 146d1d9 commit bbb2304

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

modelopt/torch/export/unified_export_megatron.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1085,7 +1085,10 @@ def _get_state_dict(self):
10851085
self.rules["k_layernorm"](layer.self_attention.k_layernorm, layer_id)
10861086
self.rules["linear_qkv"](layer.self_attention.linear_qkv, layer_id)
10871087
self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id)
1088-
if hasattr(layer.self_attention.core_attention, "softmax_offset"):
1088+
if (
1089+
getattr(layer.self_attention.core_attention, "softmax_offset", None)
1090+
is not None
1091+
):
10891092
self.rules["softmax_offset"](
10901093
layer.self_attention.core_attention.softmax_offset, layer_id
10911094
)

0 commit comments

Comments
 (0)