diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 8023f53056e..5db3154f552 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -1337,6 +1337,17 @@ def _qkv_down_projection(self, hidden_states): ) return q_compressed, kv_combined + def backward_dw(self) -> NoReturn: + """Execute weight gradient computation.""" + self.linear_kv_up_proj.backward_dw() + self.linear_qkv_down_proj.backward_dw() + self.linear_q_up_proj.backward_dw() + self._backward_output_proj() + + def set_for_recompute_input_layernorm(self): + """Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4.""" + set_save_original_input(self.linear_qkv_down_proj) + def sharded_state_dict(self, prefix: str = "", sharded_offsets: tuple = (), metadata=None): """Return a sharded state dict compatible with pre-fusion checkpoints.""" sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata) diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index 7c6fc5a46da..863a4e23d6d 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/tests/unit_tests/transformer/test_multi_latent_attention.py @@ -1768,6 +1768,36 @@ def test_backward_pass(self): assert hidden_states.grad is not None +def test_fused_mla_training_hooks_use_fused_down_projection(monkeypatch): + """Training hooks should use fused q/kv down projection attributes.""" + + class LinearWithDelayedWgrad: + def __init__(self, name): + self.name = name + + def backward_dw(self): + calls.append(self.name) + + calls = [] + fused = FusedMLASelfAttention.__new__(FusedMLASelfAttention) + fused.linear_kv_up_proj = LinearWithDelayedWgrad("kv_up") + fused.linear_qkv_down_proj = LinearWithDelayedWgrad("qkv_down") + fused.linear_q_up_proj = LinearWithDelayedWgrad("q_up") + fused.linear_proj = LinearWithDelayedWgrad("out") + + fused.backward_dw() + + assert calls == ["kv_up", "qkv_down", "q_up", "out"] + + saved_inputs = [] + mla_module = __import__(FusedMLASelfAttention.__module__, fromlist=["set_save_original_input"]) + monkeypatch.setattr(mla_module, "set_save_original_input", saved_inputs.append) + + fused.set_for_recompute_input_layernorm() + + assert saved_inputs == [fused.linear_qkv_down_proj] + + class TestFusedMLALoadFromStateDict: @pytest.fixture(scope='function', autouse=True)