[megatron] fix pp mla (#4904)

Jintao-Huang · Jintao-Huang · commit 086a9f94758f · 2025-07-11T10:06:43.000+08:00
diff --git a/swift/megatron/model/gpt_model.py b/swift/megatron/model/gpt_model.py
@@ -79,7 +79,7 @@ def __init__(
                 use_cpu_initialization=config.use_cpu_initialization,
             )
             # save memory
-            for i in range(config.num_layers):
+            for i in range(len(self.decoder.layers)):
                 if hasattr(self.decoder.layers[i].self_attention, 'rotary_pos_emb'):
                     del self.decoder.layers[i].self_attention.rotary_pos_emb
         self.attention_scaling = 1.

Original file line number	Diff line number	Diff line change
`@@ -79,7 +79,7 @@ def __init__(`
`79`	`79`	`use_cpu_initialization=config.use_cpu_initialization,`
`80`	`80`	`)`
`81`	`81`	`# save memory`
`82`		`- for i in range(config.num_layers):`
	`82`	`+ for i in range(len(self.decoder.layers)):`
`83`	`83`	`if hasattr(self.decoder.layers[i].self_attention, 'rotary_pos_emb'):`
`84`	`84`	`del self.decoder.layers[i].self_attention.rotary_pos_emb`
`85`	`85`	`self.attention_scaling = 1.`