feat: attention backward pass Phase A (residual passthrough)
Add backward_residual() method to GpuAttention that implements gradient flow through the residual connection for frozen attention weights. This is Phase A of the attention backward implementation. The residual connection (output = input + attention(input)) ensures that gradients flow through unchanged: d_input = d_output. With frozen weights, we don't compute attention weight gradients, making this a literal no-op when the input and output buffers alias (as they do in apply_iqn_trunk_gradient). Phase B will add attention path gradients when weights are unfrozen.
This commit is contained in:
@@ -161,6 +161,27 @@ impl GpuAttention {
|
||||
Ok(&self.output_buf)
|
||||
}
|
||||
|
||||
/// Backward pass (Phase A: residual passthrough).
|
||||
///
|
||||
/// For frozen attention weights, the gradient flows entirely through
|
||||
/// the residual connection: d_input = d_output. This is a no-op when
|
||||
/// d_input and d_output are the same buffer (which they are in
|
||||
/// apply_iqn_trunk_gradient where bw_d_h_s2 is used for both).
|
||||
///
|
||||
/// Phase B (unfrozen weights) would add attention path gradients
|
||||
/// and accumulate weight gradients via atomicAdd.
|
||||
pub fn backward_residual(
|
||||
&self,
|
||||
_d_output: &CudaSlice<f32>,
|
||||
_d_input: &mut CudaSlice<f32>,
|
||||
_batch_size: usize,
|
||||
) -> Result<(), MLError> {
|
||||
// Phase A: d_input = d_output through residual path.
|
||||
// When d_input and d_output alias (same bw_d_h_s2 buffer),
|
||||
// this is a literal no-op — the gradient is already in place.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get the output buffer reference.
|
||||
pub fn output(&self) -> &CudaSlice<f32> {
|
||||
&self.output_buf
|
||||
|
||||
Reference in New Issue
Block a user