Fix issues of attention.core_attention.softmax_offset is None (#330)

yueshen2016 · web-flow · commit d94fc1b44e37 · 2025-09-17T20:25:38.000Z
Signed-off-by: Yue &lt;yueshen@nvidia.com&gt;
diff --git a/modelopt/torch/export/plugins/megatron_importer.py b/modelopt/torch/export/plugins/megatron_importer.py
@@ -512,7 +512,7 @@ def _import_state_dict(self):
                             self.rules["k_layernorm"](attention.k_layernorm, layer_id)
                         self.rules["linear_qkv"](attention.linear_qkv, layer_id)
                         self.rules["linear_proj"](attention.linear_proj, layer_id)
-                        if hasattr(attention.core_attention, "softmax_offset"):
+                        if getattr(attention.core_attention, "softmax_offset", None) is not None:
                             self.rules["softmax_offset"](
                                 attention.core_attention.softmax_offset, layer_id
                             )

Original file line number	Diff line number	Diff line change
`@@ -512,7 +512,7 @@ def _import_state_dict(self):`
`512`	`512`	`self.rules["k_layernorm"](attention.k_layernorm, layer_id)`
`513`	`513`	`self.rules["linear_qkv"](attention.linear_qkv, layer_id)`
`514`	`514`	`self.rules["linear_proj"](attention.linear_proj, layer_id)`
`515`		`- if hasattr(attention.core_attention, "softmax_offset"):`
	`515`	`+ if getattr(attention.core_attention, "softmax_offset", None) is not None:`
`516`	`516`	`self.rules["softmax_offset"](`
`517`	`517`	`attention.core_attention.softmax_offset, layer_id`
`518`	`518`	`)`