From bd32cd0e2bb83156dd2e5b89002fb02caf4ef224 Mon Sep 17 00:00:00 2001 From: lit Date: Sun, 18 Jan 2026 19:27:10 -0800 Subject: [PATCH] Add backward_dw to GDN --- megatron/core/ssm/gated_delta_net.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 70e749724dc..7b1149a781d 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -508,6 +508,19 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_gr return sharded_state_dict + def backward_dw(self): + """Execute weight gradient computation for all linear layers.""" + self._backward_in_proj() + self._backward_out_proj() + + def _backward_in_proj(self): + """Computes weight gradients of input projection layer.""" + self.in_proj.backward_dw() + + def _backward_out_proj(self): + """Computes weight gradients of output projection layer.""" + self.out_proj.backward_dw() + def _split_tensor_factory( orig_sh_ten: ShardedTensor, split_sections: List[int], split_names: List[str], split_dim: int