burn_optim/grad_clipping/
base.rs

1use burn_core as burn;
2
3use burn::{
4    config::Config,
5    tensor::{DType, FloatDType, Tensor},
6};
7
8/// Gradient Clipping provides a way to mitigate exploding gradients
9#[derive(Config, Debug)]
10pub enum GradientClippingConfig {
11    /// Clip the gradient by value.
12    Value(f32),
13
14    /// Clip the gradient by norm.
15    Norm(f32),
16}
17
18impl GradientClippingConfig {
19    /// Initialize the gradient clipping.
20    ///
21    /// # Returns
22    ///
23    /// The gradient clipping.
24    pub fn init(&self) -> GradientClipping {
25        match self {
26            GradientClippingConfig::Value(val) => GradientClipping::Value(*val),
27            GradientClippingConfig::Norm(val) => GradientClipping::Norm(*val),
28        }
29    }
30}
31
32/// Gradient Clipping provides a way to mitigate exploding gradients
33/// by clipping every component of the gradient by value or by norm during
34/// backpropagation.
35#[derive(Clone)]
36pub enum GradientClipping {
37    /// Clip the gradient by value.
38    Value(f32),
39
40    /// Clip the gradient by norm.
41    Norm(f32),
42}
43
44impl GradientClipping {
45    /// Clip the gradient.
46    ///
47    /// # Arguments
48    ///
49    /// * `grad` - The gradient to clip.
50    ///
51    /// # Returns
52    ///
53    /// The clipped gradient.
54    pub fn clip_gradient<const D: usize>(&self, grad: Tensor<D>) -> Tensor<D> {
55        match self {
56            GradientClipping::Value(threshold) => self.clip_by_value(grad, *threshold),
57            GradientClipping::Norm(max_norm) => self.clip_by_norm(grad, *max_norm),
58        }
59    }
60
61    fn clip_by_value<const D: usize>(&self, grad: Tensor<D>, threshold: f32) -> Tensor<D> {
62        let greater_mask = grad.clone().greater_scalar(threshold);
63        let lower_mask = grad.clone().lower_scalar(-threshold);
64
65        let clipped_grad = grad.mask_fill(greater_mask, threshold);
66
67        clipped_grad.mask_fill(lower_mask, -threshold)
68    }
69
70    fn clip_by_norm<const D: usize>(&self, grad: Tensor<D>, threshold: f32) -> Tensor<D> {
71        // Compute the norm in F32: in F16 the sum of squares overflows once the norm
72        // exceeds ~256, and BF16 loses too much precision accumulating it.
73        let dtype = grad.dtype();
74        let grad = match dtype {
75            DType::F16 | DType::BF16 => grad.cast(FloatDType::F32),
76            _ => grad,
77        };
78
79        let norm = Self::l2_norm(grad.clone());
80        let min_positive = grad
81            .dtype()
82            .finfo()
83            .unwrap_or(FloatDType::F32.finfo())
84            .min_positive;
85        let clip_coef = threshold / norm.add_scalar(min_positive);
86        let clip_coef_clamped = clip_coef.clamp_max(1.0);
87        grad.mul(clip_coef_clamped.unsqueeze()).cast(dtype)
88    }
89
90    fn l2_norm<const D: usize>(tensor: Tensor<D>) -> Tensor<1> {
91        let squared = tensor.square();
92        let sum = squared.sum();
93        sum.sqrt()
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use burn::tensor::Tensor;
101
102    #[test]
103    fn test_clip_by_value() {
104        let gradient: Tensor<2> = Tensor::from_floats(
105            [
106                [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310],
107                [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883],
108            ],
109            &Default::default(),
110        );
111
112        let clipped_gradient = GradientClipping::Value(0.5).clip_gradient(gradient);
113        let clipped_gradient_data = clipped_gradient.into_data();
114
115        for value in clipped_gradient_data.iter::<f32>() {
116            assert!(value <= 0.5);
117        }
118    }
119
120    #[test]
121    fn test_clip_by_norm() {
122        let gradient: Tensor<2> = Tensor::from_floats(
123            [
124                [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310],
125                [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883],
126            ],
127            &Default::default(),
128        );
129
130        let clipped_gradient = GradientClipping::Norm(2.2).clip_gradient(gradient);
131        let clipped_gradient_data = clipped_gradient.into_data();
132
133        for value in clipped_gradient_data.iter::<f32>() {
134            assert!(value <= 0.88);
135        }
136    }
137    #[test]
138    fn test_clip_by_norm_no_clipping() {
139        let gradient: Tensor<2> = Tensor::from_floats(
140            [[0.3, 0.4, 0.5, 0.2], [0.1, 0.6, 0.3, 0.4]],
141            &Default::default(),
142        );
143
144        let clipped_gradient = GradientClipping::Norm(2.2).clip_gradient(gradient.clone());
145
146        clipped_gradient
147            .into_data()
148            .assert_eq(&gradient.into_data(), true);
149    }
150
151    #[test]
152    fn test_clip_by_norm_f16_does_not_overflow() {
153        let gradient = Tensor::<1>::from_floats([300.0], &Default::default()).cast(FloatDType::F16);
154
155        let clipped_gradient = GradientClipping::Norm(1.0).clip_gradient(gradient);
156
157        assert_eq!(clipped_gradient.dtype(), DType::F16);
158        let actual = clipped_gradient.into_scalar::<f32>();
159        assert!((actual - 1.0).abs() < 1e-3, "actual={actual}");
160    }
161
162    #[test]
163    fn test_clip_by_norm_f16_norm_above_f16_max() {
164        let gradient =
165            Tensor::<1>::from_floats([60000.0, 60000.0], &Default::default()).cast(FloatDType::F16);
166
167        let clipped_gradient = GradientClipping::Norm(1.0).clip_gradient(gradient);
168
169        let expected = core::f32::consts::FRAC_1_SQRT_2;
170        for value in clipped_gradient.into_data().iter::<f32>() {
171            assert!((value - expected).abs() < 1e-3, "value={value}");
172        }
173    }
174}
burn_optim/grad_clipping/base.rs

burn_optim/grad_clipping/
base.rs
