已合并
[feat] fusedsgd #5277
culechan创建于 6月26日
[feat] fusedsgd #5277
已合并
共 4 个文件变更+417-0
| @@ -5569,6 +5569,13 @@ official: | |||
| 5569 | - func: _scaled_grouped_mm_v2(Tensor mat_a, Tensor mat_b, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? offs=None, Tensor? bias=None, ScalarType? out_dtype=None, int[] contraction_dim=[], bool use_fast_accum=False) -> Tensor | 5569 | - func: _scaled_grouped_mm_v2(Tensor mat_a, Tensor mat_b, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? offs=None, Tensor? bias=None, ScalarType? out_dtype=None, int[] contraction_dim=[], bool use_fast_accum=False) -> Tensor |
| 5570 | op_api: [v2.10, newest] | 5570 | op_api: [v2.10, newest] |
| 5571 | 5571 | ||
| 5572 | + - func: _fused_sgd_(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] momentum_buffer_list, *, float weight_decay, float momentum, float lr, float dampening, bool nesterov, bool maximize, bool is_first_step, Tensor? grad_scale=None, Tensor? found_inf=None) -> () | ||
D | |||
| 5573 | + op_api: [v2.7, newest] | ||
| 5574 | + | ||
| 5575 | + - func: _fused_sgd_.tensor_lr(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] momentum_buffer_list, *, float weight_decay, float momentum, Tensor lr, float dampening, bool nesterov, bool maximize, bool is_first_step, Tensor? grad_scale=None, Tensor? found_inf=None) -> () | ||
| 5576 | + op_api: [v2.7, newest] | ||
| 5577 | + device_check: NoCheck | ||
| 5578 | + | ||
| 5572 | autograd: | 5579 | autograd: |
| 5573 | - func: _thnn_fused_gru_cell(Tensor input_gates, Tensor hidden_gates, Tensor hx, Tensor? input_bias=None, Tensor? hidden_bias=None) -> (Tensor, Tensor) | 5580 | - func: _thnn_fused_gru_cell(Tensor input_gates, Tensor hidden_gates, Tensor hx, Tensor? input_bias=None, Tensor? hidden_bias=None) -> (Tensor, Tensor) |
| 5574 | op_api: [v2.5, newest] | 5581 | op_api: [v2.5, newest] |
| @@ -0,0 +1,90 @@ | |||
| 1 | +// Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 2 | +// All rights reserved. | ||
| 3 | +// | ||
| 4 | +// Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +// you may not use this file except in compliance with the License. | ||
| 6 | +// You may obtain a copy of the License at | ||
| 7 | +// | ||
| 8 | +// https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +// | ||
| 10 | +// Unless required by applicable law or agreed to in writing, software | ||
| 11 | +// distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +// See the License for the specific language governing permissions and | ||
| 14 | +// limitations under the License. | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +namespace op_api { | ||
| 22 | + | ||
| 23 | +void _fused_sgd_(at::TensorList params, at::TensorList grads, at::TensorList momentum_buffer_list, | ||
| 24 | + const double weight_decay, const double momentum, const double lr, const double dampening, | ||
| 25 | + const bool nesterov, const bool maximize, const bool is_first_step, | ||
| 26 | + const std::optional<at::Tensor>& grad_scale, | ||
| 27 | + const std::optional<at::Tensor>& found_inf) | ||
| 28 | +{ | ||
| 29 | + if (found_inf.has_value()) { | ||
| 30 | + at::Tensor found_inf_real = found_inf.value(); | ||
| 31 | + if (found_inf_real.item().toFloat() == 1) { | ||
| 32 | + return; | ||
| 33 | + } | ||
| 34 | + } | ||
| 35 | + const float momentum_real = float(momentum); | ||
| 36 | + TORCH_CHECK(momentum_real > 0, "momentum must be positive, but got ", momentum_real); | ||
| 37 | + TORCH_CHECK(at::native::check_fast_path_restrictions({params, grads})); | ||
| 38 | + if (is_first_step && momentum_buffer_list.empty()) { | ||
| 39 | + TORCH_WARN_ONCE( | ||
| 40 | + "`is_first_step` argument has no effect when `momentum_buffer_list` is empty"); | ||
| 41 | + } | ||
| 42 | + const float weight_decay_real = float(weight_decay); | ||
| 43 | + const float lr_real = float(lr); | ||
| 44 | + const float dampening_real = float(dampening); | ||
| 45 | + const at::Tensor grad_scale_real = grad_scale.value_or(at::Tensor()); | ||
| 46 | + EXEC_NPU_CMD(aclnnFusedSgd, params, grads, momentum_buffer_list, grad_scale_real, weight_decay_real, | ||
| 47 | + momentum_real, lr_real, dampening_real, nesterov, maximize, is_first_step); | ||
| 48 | +} | ||
| 49 | + | ||
| 50 | +void _fused_sgd_(at::TensorList params, at::TensorList grads, at::TensorList momentum_buffer_list, | ||
| 51 | + const double weight_decay, const double momentum, const at::Tensor& lr, const double dampening, | ||
| 52 | + const bool nesterov, const bool maximize, const bool is_first_step, | ||
| 53 | + const std::optional<at::Tensor>& grad_scale, | ||
| 54 | + const std::optional<at::Tensor>& found_inf) | ||
| 55 | +{ | ||
| 56 | + TORCH_CHECK(at::native::check_fast_path_restrictions({params, grads})); | ||
| 57 | + if (is_first_step && momentum_buffer_list.empty()) { | ||
| 58 | + TORCH_WARN_ONCE( | ||
| 59 | + "`is_first_step` argument has no effect when `momentum_buffer_list` is empty"); | ||
| 60 | + } | ||
| 61 | + if (grad_scale.has_value()) { | ||
| 62 | + TORCH_CHECK( | ||
| 63 | + grad_scale->device() == params[0].device(), | ||
| 64 | + "grad_scale must be on the same NPU device as the params"); | ||
| 65 | + } | ||
| 66 | + if (found_inf.has_value()) { | ||
| 67 | + TORCH_CHECK( | ||
| 68 | + found_inf->device() == params[0].device(), | ||
| 69 | + "found_inf must be on the same NPU device as the params"); | ||
| 70 | + } | ||
| 71 | + TORCH_CHECK( | ||
| 72 | + lr.device() == params[0].device(), | ||
| 73 | + "lr must be on the same NPU device as the params"); | ||
| 74 | + if (found_inf.has_value()) { | ||
| 75 | + at::Tensor found_inf_real = found_inf.value(); | ||
| 76 | + if (found_inf_real.item().toFloat() == 1) { | ||
| 77 | + return; | ||
| 78 | + } | ||
| 79 | + } | ||
| 80 | + const float momentum_real = float(momentum); | ||
| 81 | + TORCH_CHECK(momentum_real > 0, "momentum must be positive, but got ", momentum_real); | ||
| 82 | + const float weight_decay_real = float(weight_decay); | ||
| 83 | + const float dampening_real = float(dampening); | ||
| 84 | + const float lr_value = lr.item().toFloat(); | ||
| 85 | + const at::Tensor grad_scale_real = grad_scale.value_or(at::Tensor()); | ||
| 86 | + EXEC_NPU_CMD(aclnnFusedSgd, params, grads, momentum_buffer_list, grad_scale_real, weight_decay_real, | ||
| 87 | + momentum_real, lr_value, dampening_real, nesterov, maximize, is_first_step); | ||
| 88 | +} | ||
| 89 | + | ||
| 90 | +} // namespace op_api | ||
| @@ -5044,5 +5044,11 @@ | |||
| 5044 | }, | 5044 | }, |
| 5045 | "func: _scaled_grouped_mm_v2(Tensor mat_a, Tensor mat_b, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? offs=None, Tensor? bias=None, ScalarType? out_dtype=None, int[] contraction_dim=[], bool use_fast_accum=False) -> Tensor": { | 5045 | "func: _scaled_grouped_mm_v2(Tensor mat_a, Tensor mat_b, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? offs=None, Tensor? bias=None, ScalarType? out_dtype=None, int[] contraction_dim=[], bool use_fast_accum=False) -> Tensor": { |
| 5046 | "version": ["v2.10", "newest"] | 5046 | "version": ["v2.10", "newest"] |
| 5047 | + }, | ||
| 5048 | + "func: _fused_sgd_(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] momentum_buffer_list, *, float weight_decay, float momentum, float lr, float dampening, bool nesterov, bool maximize, bool is_first_step, Tensor? grad_scale=None, Tensor? found_inf=None) -> ()": { | ||
| 5049 | + "version": ["v2.7", "newest"] | ||
| 5050 | + }, | ||
| 5051 | + "func: _fused_sgd_.tensor_lr(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] momentum_buffer_list, *, float weight_decay, float momentum, Tensor lr, float dampening, bool nesterov, bool maximize, bool is_first_step, Tensor? grad_scale=None, Tensor? found_inf=None) -> ()": { | ||
| 5052 | + "version": ["v2.7", "newest"] | ||
| 5047 | } | 5053 | } |
| 5048 | } | 5054 | } |
| @@ -0,0 +1,314 @@ | |||
| 1 | +import unittest | ||
| 2 | +import copy | ||
| 3 | +import torch | ||
| 4 | +import torch_npu | ||
| 5 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +def fused_sgd_ref(params, grads, momentum_buffer_list, weight_decay, momentum, lr, | ||
| 9 | + dampening, nesterov, maximize, is_first_step, | ||
| 10 | + grad_scale=None, found_inf=None): | ||
| 11 | + if found_inf is not None and int(found_inf.item()) == 1: | ||
| 12 | + return | ||
| 13 | + grad_scale_val = grad_scale.item() if grad_scale is not None else 1.0 | ||
| 14 | + for i in range(len(params)): | ||
| 15 | + d_p = grads[i].clone() | ||
| 16 | + if grad_scale is not None: | ||
| 17 | + d_p = d_p / grad_scale_val | ||
| 18 | + if weight_decay != 0: | ||
| 19 | + d_p = d_p.add(params[i], alpha=weight_decay) | ||
| 20 | + if momentum != 0: | ||
| 21 | + buf = momentum_buffer_list[i] | ||
| 22 | + if is_first_step: | ||
| 23 | + buf.copy_(d_p) | ||
| 24 | + else: | ||
| 25 | + buf.mul_(momentum).add_(d_p, alpha=1 - dampening) | ||
| 26 | + if nesterov: | ||
| 27 | + d_p = d_p.add(buf, alpha=momentum) | ||
| 28 | + else: | ||
| 29 | + d_p = buf | ||
| 30 | + if maximize: | ||
| 31 | + params[i].add_(d_p, alpha=lr) | ||
| 32 | + else: | ||
| 33 | + params[i].add_(d_p, alpha=-lr) | ||
| 34 | + | ||
| 35 | + | ||
| 36 | +class TestFusedSgd(TestCase): | ||
| 37 | + | ||
| 38 | + def _gen_tensors(self, shapes, dtype=torch.float32, seed=42): | ||
| 39 | + torch.manual_seed(seed) | ||
| 40 | + cpu_params = [torch.randn(s, dtype=dtype) for s in shapes] | ||
| 41 | + cpu_grads = [torch.randn(s, dtype=dtype) for s in shapes] | ||
| 42 | + cpu_bufs = [torch.randn(s, dtype=dtype) for s in shapes] | ||
| 43 | + npu_params = [p.npu() for p in cpu_params] | ||
| 44 | + npu_grads = [g.npu() for g in cpu_grads] | ||
| 45 | + npu_bufs = [b.npu() for b in cpu_bufs] | ||
| 46 | + cpu_params = [p.clone() for p in cpu_params] | ||
| 47 | + cpu_grads = [g.clone() for g in cpu_grads] | ||
| 48 | + cpu_bufs = [b.clone() for b in cpu_bufs] | ||
| 49 | + return cpu_params, cpu_grads, cpu_bufs, npu_params, npu_grads, npu_bufs | ||
| 50 | + | ||
| 51 | + def _run_and_compare(self, shapes, dtype, weight_decay, momentum, lr, | ||
| 52 | + dampening, nesterov, maximize, is_first_step, | ||
| 53 | + grad_scale=None, found_inf=None, prec=None): | ||
| 54 | + cpu_p, cpu_g, cpu_b, npu_p, npu_g, npu_b = self._gen_tensors(shapes, dtype) | ||
| 55 | + | ||
| 56 | + torch._fused_sgd_(cpu_p, cpu_g, cpu_b, weight_decay=weight_decay, | ||
| 57 | + momentum=momentum, lr=lr, dampening=dampening, | ||
| 58 | + nesterov=nesterov, maximize=maximize, | ||
| 59 | + is_first_step=is_first_step, | ||
| 60 | + grad_scale=grad_scale, found_inf=found_inf) | ||
| 61 | + | ||
| 62 | + npu_grad_scale = grad_scale.npu() if grad_scale is not None else None | ||
| 63 | + npu_found_inf = found_inf.npu() if found_inf is not None else None | ||
| 64 | + torch._fused_sgd_(npu_p, npu_g, npu_b, weight_decay=weight_decay, | ||
| 65 | + momentum=momentum, lr=lr, dampening=dampening, | ||
| 66 | + nesterov=nesterov, maximize=maximize, | ||
| 67 | + is_first_step=is_first_step, | ||
| 68 | + grad_scale=npu_grad_scale, found_inf=npu_found_inf) | ||
| 69 | + | ||
| 70 | + for c_p, n_p, c_b, n_b in zip(cpu_p, npu_p, cpu_b, npu_b): | ||
| 71 | + self.assertRtolEqual(c_p, n_p.cpu(), prec=1e-3) | ||
| 72 | + self.assertRtolEqual(c_b, n_b.cpu(), prec=1e-3) | ||
| 73 | + | ||
| 74 | + def test_fused_sgd_first_step(self): | ||
| 75 | + self._run_and_compare( | ||
| 76 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 77 | + weight_decay=0.0, momentum=0.9, lr=0.01, | ||
| 78 | + dampening=0.0, nesterov=False, maximize=False, | ||
| 79 | + is_first_step=True) | ||
| 80 | + | ||
| 81 | + def test_fused_sgd_subsequent_step(self): | ||
| 82 | + self._run_and_compare( | ||
| 83 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 84 | + weight_decay=0.0, momentum=0.9, lr=0.01, | ||
| 85 | + dampening=0.0, nesterov=False, maximize=False, | ||
| 86 | + is_first_step=False) | ||
| 87 | + | ||
| 88 | + def test_fused_sgd_weight_decay(self): | ||
| 89 | + self._run_and_compare( | ||
| 90 | + shapes=[(4, 4), (2, 3, 3)], dtype=torch.float32, | ||
| 91 | + weight_decay=0.01, momentum=0.9, lr=0.01, | ||
| 92 | + dampening=0.0, nesterov=False, maximize=False, | ||
| 93 | + is_first_step=False) | ||
| 94 | + | ||
| 95 | + def test_fused_sgd_nesterov(self): | ||
| 96 | + self._run_and_compare( | ||
| 97 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 98 | + weight_decay=0.0, momentum=0.9, lr=0.01, | ||
| 99 | + dampening=0.0, nesterov=True, maximize=False, | ||
| 100 | + is_first_step=False) | ||
| 101 | + | ||
| 102 | + def test_fused_sgd_maximize(self): | ||
| 103 | + self._run_and_compare( | ||
| 104 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 105 | + weight_decay=0.0, momentum=0.9, lr=0.01, | ||
| 106 | + dampening=0.0, nesterov=False, maximize=True, | ||
| 107 | + is_first_step=False) | ||
| 108 | + | ||
| 109 | + def test_fused_sgd_dampening(self): | ||
| 110 | + self._run_and_compare( | ||
| 111 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 112 | + weight_decay=0.0, momentum=0.9, lr=0.01, | ||
| 113 | + dampening=0.5, nesterov=False, maximize=False, | ||
| 114 | + is_first_step=False) | ||
| 115 | + | ||
| 116 | + def test_fused_sgd_nesterov_with_weight_decay(self): | ||
| 117 | + self._run_and_compare( | ||
| 118 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 119 | + weight_decay=0.001, momentum=0.9, lr=0.01, | ||
| 120 | + dampening=0.0, nesterov=True, maximize=True, | ||
| 121 | + is_first_step=False) | ||
| 122 | + | ||
| 123 | + def test_fused_sgd_found_inf_skip(self): | ||
| 124 | + cpu_p, cpu_g, cpu_b, npu_p, npu_g, npu_b = self._gen_tensors( | ||
| 125 | + [(4, 4), (8, 8)], torch.float32) | ||
| 126 | + cpu_p_before = [p.clone() for p in cpu_p] | ||
| 127 | + cpu_b_before = [b.clone() for b in cpu_b] | ||
| 128 | + npu_p_before = [p.clone() for p in npu_p] | ||
| 129 | + npu_b_before = [b.clone() for b in npu_b] | ||
| 130 | + | ||
| 131 | + found_inf = torch.tensor(1, dtype=torch.int32) | ||
| 132 | + fused_sgd_ref(cpu_p, cpu_g, cpu_b, weight_decay=0.0, momentum=0.9, lr=0.01, dampening= 0.0, | ||
| 133 | + nesterov=False, maximize=False, is_first_step=False, grad_scale=None, found_inf=found_inf) | ||
| 134 | + torch._fused_sgd_(npu_p, npu_g, npu_b, weight_decay=0.0, | ||
| 135 | + momentum=0.9, lr=0.01, dampening=0.0, | ||
| 136 | + nesterov=False, maximize=False, | ||
| 137 | + is_first_step=False, grad_scale=None, | ||
| 138 | + found_inf=found_inf.npu()) | ||
| 139 | + | ||
| 140 | + for c_before, c_after in zip(cpu_p_before, cpu_p): | ||
| 141 | + self.assertRtolEqual(c_before, c_after, prec=1e-2) | ||
| 142 | + for n_before, n_after in zip(npu_p_before, npu_p): | ||
| 143 | + self.assertRtolEqual(n_before.cpu(), n_after.cpu(), prec=1e-2) | ||
| 144 | + | ||
| 145 | + def test_fused_sgd_found_inf_zero(self): | ||
| 146 | + found_inf = torch.tensor(0, dtype=torch.float32) | ||
| 147 | + self._run_and_compare( | ||
| 148 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 149 | + weight_decay=0.0, momentum=0.9, lr=0.01, | ||
| 150 | + dampening=0.0, nesterov=False, maximize=False, | ||
| 151 | + is_first_step=False, found_inf=found_inf) | ||
| 152 | + | ||
| 153 | + def test_fused_sgd_grad_scale(self): | ||
| 154 | + grad_scale = torch.tensor(4.0, dtype=torch.float32) | ||
| 155 | + self._run_and_compare( | ||
| 156 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 157 | + weight_decay=0.0, momentum=0.9, lr=0.01, | ||
| 158 | + dampening=0.0, nesterov=False, maximize=False, | ||
| 159 | + is_first_step=False, grad_scale=grad_scale) | ||
| 160 | + | ||
| 161 | + def test_fused_sgd_grad_scale_with_weight_decay(self): | ||
| 162 | + grad_scale = torch.tensor(2.0, dtype=torch.float32) | ||
| 163 | + self._run_and_compare( | ||
| 164 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 165 | + weight_decay=0.01, momentum=0.9, lr=0.01, | ||
| 166 | + dampening=0.0, nesterov=True, maximize=False, | ||
| 167 | + is_first_step=False, grad_scale=grad_scale) | ||
| 168 | + | ||
| 169 | + def test_fused_sgd_single_tensor(self): | ||
| 170 | + self._run_and_compare( | ||
| 171 | + shapes=[(16, 16)], dtype=torch.float32, | ||
| 172 | + weight_decay=0.01, momentum=0.9, lr=0.05, | ||
| 173 | + dampening=0.1, nesterov=False, maximize=False, | ||
| 174 | + is_first_step=False) | ||
| 175 | + | ||
| 176 | + def test_fused_sgd_multi_tensor(self): | ||
| 177 | + self._run_and_compare( | ||
| 178 | + shapes=[(2, 2), (3, 3), (4, 4), (5, 5), (6, 6)], dtype=torch.float32, | ||
| 179 | + weight_decay=0.001, momentum=0.9, lr=0.01, | ||
| 180 | + dampening=0.0, nesterov=False, maximize=False, | ||
| 181 | + is_first_step=False) | ||
| 182 | + | ||
| 183 | + def test_fused_sgd_momentum_non_positive_error(self): | ||
| 184 | + npu_p = [torch.randn(4, 4, dtype=torch.float32).npu()] | ||
| 185 | + npu_g = [torch.randn(4, 4, dtype=torch.float32).npu()] | ||
| 186 | + npu_b = [torch.randn(4, 4, dtype=torch.float32).npu()] | ||
| 187 | + with self.assertRaises(RuntimeError): | ||
| 188 | + torch._fused_sgd_(npu_p, npu_g, npu_b, weight_decay=0.0, | ||
| 189 | + momentum=0.0, lr=0.01, dampening=0.0, | ||
| 190 | + nesterov=False, maximize=False, | ||
| 191 | + is_first_step=True) | ||
| 192 | + | ||
| 193 | + | ||
| 194 | +class TestFusedSgdTensorLr(TestCase): | ||
| 195 | + | ||
| 196 | + def _gen_tensors(self, shapes, dtype=torch.float32, seed=42): | ||
| 197 | + torch.manual_seed(seed) | ||
| 198 | + cpu_params = [torch.randn(s, dtype=dtype) for s in shapes] | ||
| 199 | + cpu_grads = [torch.randn(s, dtype=dtype) for s in shapes] | ||
| 200 | + cpu_bufs = [torch.randn(s, dtype=dtype) for s in shapes] | ||
| 201 | + npu_params = [p.npu() for p in cpu_params] | ||
| 202 | + npu_grads = [g.npu() for g in cpu_grads] | ||
| 203 | + npu_bufs = [b.npu() for b in cpu_bufs] | ||
| 204 | + cpu_params = [p.clone() for p in cpu_params] | ||
| 205 | + cpu_grads = [g.clone() for g in cpu_grads] | ||
| 206 | + cpu_bufs = [b.clone() for b in cpu_bufs] | ||
| 207 | + return cpu_params, cpu_grads, cpu_bufs, npu_params, npu_grads, npu_bufs | ||
| 208 | + | ||
| 209 | + def _run_and_compare(self, shapes, dtype, weight_decay, momentum, lr_val, | ||
| 210 | + dampening, nesterov, maximize, is_first_step, | ||
| 211 | + grad_scale=None, found_inf=None, prec=None): | ||
| 212 | + cpu_p, cpu_g, cpu_b, npu_p, npu_g, npu_b = self._gen_tensors(shapes, dtype) | ||
| 213 | + | ||
| 214 | + cpu_lr = torch.tensor(lr_val, dtype=torch.float32) | ||
| 215 | + torch._fused_sgd_(cpu_p, cpu_g, cpu_b, weight_decay=weight_decay, momentum=momentum, lr=cpu_lr, | ||
| 216 | + dampening=dampening, nesterov=nesterov, maximize=maximize, is_first_step=is_first_step, | ||
| 217 | + grad_scale=grad_scale, found_inf=found_inf) | ||
| 218 | + | ||
| 219 | + npu_lr = cpu_lr.npu() | ||
| 220 | + npu_grad_scale = grad_scale.npu() if grad_scale is not None else None | ||
| 221 | + npu_found_inf = found_inf.npu() if found_inf is not None else None | ||
| 222 | + torch._fused_sgd_(npu_p, npu_g, npu_b, weight_decay=weight_decay, | ||
| 223 | + momentum=momentum, lr=npu_lr, | ||
| 224 | + dampening=dampening, nesterov=nesterov, | ||
| 225 | + maximize=maximize, is_first_step=is_first_step, | ||
| 226 | + grad_scale=npu_grad_scale, | ||
| 227 | + found_inf=npu_found_inf) | ||
| 228 | + | ||
| 229 | + for c_p, n_p, c_b, n_b in zip(cpu_p, npu_p, cpu_b, npu_b): | ||
| 230 | + self.assertRtolEqual(c_p, n_p.cpu(), prec=1e-3) | ||
| 231 | + self.assertRtolEqual(c_b, n_b.cpu(), prec=1e-3) | ||
| 232 | + | ||
| 233 | + def test_fused_sgd_tensor_lr_first_step(self): | ||
| 234 | + self._run_and_compare( | ||
| 235 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 236 | + weight_decay=0.0, momentum=0.9, lr_val=0.01, | ||
| 237 | + dampening=0.0, nesterov=False, maximize=False, | ||
| 238 | + is_first_step=True) | ||
| 239 | + | ||
| 240 | + def test_fused_sgd_tensor_lr_subsequent_step(self): | ||
| 241 | + self._run_and_compare( | ||
| 242 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 243 | + weight_decay=0.01, momentum=0.9, lr_val=0.01, | ||
| 244 | + dampening=0.0, nesterov=False, maximize=False, | ||
| 245 | + is_first_step=False) | ||
| 246 | + | ||
| 247 | + def test_fused_sgd_tensor_lr_nesterov(self): | ||
| 248 | + self._run_and_compare( | ||
| 249 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 250 | + weight_decay=0.001, momentum=0.9, lr_val=0.05, | ||
| 251 | + dampening=0.0, nesterov=True, maximize=False, | ||
| 252 | + is_first_step=False) | ||
| 253 | + | ||
| 254 | + def test_fused_sgd_tensor_lr_maximize(self): | ||
| 255 | + self._run_and_compare( | ||
| 256 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 257 | + weight_decay=0.0, momentum=0.9, lr_val=0.01, | ||
| 258 | + dampening=0.0, nesterov=False, maximize=True, | ||
| 259 | + is_first_step=False) | ||
| 260 | + | ||
| 261 | + def test_fused_sgd_tensor_lr_found_inf_skip(self): | ||
| 262 | + cpu_p, cpu_g, cpu_b, npu_p, npu_g, npu_b = self._gen_tensors( | ||
| 263 | + [(4, 4), (8, 8)], torch.float32) | ||
| 264 | + cpu_p_before = [p.clone() for p in cpu_p] | ||
| 265 | + npu_p_before = [p.clone() for p in npu_p] | ||
| 266 | + | ||
| 267 | + found_inf = torch.tensor(1, dtype=torch.float32) | ||
| 268 | + torch._fused_sgd_(cpu_p, cpu_g, cpu_b, weight_decay=0.0, momentum=0.9, lr=0.01, dampening=0.0, | ||
| 269 | + nesterov=False, maximize=False, is_first_step=False, grad_scale=None, found_inf=found_inf) | ||
| 270 | + npu_lr = torch.tensor(0.01, dtype=torch.float32).npu() | ||
| 271 | + torch._fused_sgd_(npu_p, npu_g, npu_b, weight_decay=0.0, | ||
| 272 | + momentum=0.9, lr=npu_lr, dampening=0.0, | ||
| 273 | + nesterov=False, maximize=False, | ||
| 274 | + is_first_step=False, grad_scale=None, | ||
| 275 | + found_inf=found_inf.npu()) | ||
| 276 | + | ||
| 277 | + for c_before, c_after in zip(cpu_p_before, cpu_p): | ||
| 278 | + self.assertRtolEqual(c_before, c_after) | ||
| 279 | + for n_before, n_after in zip(npu_p_before, npu_p): | ||
| 280 | + self.assertRtolEqual(n_before.cpu(), n_after.cpu()) | ||
| 281 | + | ||
| 282 | + def test_fused_sgd_tensor_lr_grad_scale(self): | ||
| 283 | + grad_scale = torch.tensor([4.0], dtype=torch.float32) | ||
| 284 | + self._run_and_compare( | ||
| 285 | + shapes=[(4, 4), (8, 8)], dtype=torch.float32, | ||
| 286 | + weight_decay=0.0, momentum=0.9, lr_val=0.01, | ||
| 287 | + dampening=0.0, nesterov=False, maximize=False, | ||
| 288 | + is_first_step=False, grad_scale=grad_scale) | ||
| 289 | + | ||
| 290 | + def test_fused_sgd_tensor_lr_momentum_non_positive_error(self): | ||
| 291 | + npu_p = [torch.randn(4, 4, dtype=torch.float32).npu()] | ||
| 292 | + npu_g = [torch.randn(4, 4, dtype=torch.float32).npu()] | ||
| 293 | + npu_b = [torch.randn(4, 4, dtype=torch.float32).npu()] | ||
| 294 | + npu_lr = torch.tensor(0.01, dtype=torch.float32).npu() | ||
| 295 | + with self.assertRaises(RuntimeError): | ||
| 296 | + torch._fused_sgd_(npu_p, npu_g, npu_b, weight_decay=0.0, | ||
| 297 | + momentum=0.0, lr=npu_lr, | ||
| 298 | + dampening=0.0, nesterov=False, | ||
| 299 | + maximize=False, is_first_step=True) | ||
| 300 | + | ||
| 301 | + def test_fused_sgd_tensor_lr_device_check(self): | ||
| 302 | + npu_p = [torch.randn(4, 4, dtype=torch.float32).npu()] | ||
| 303 | + npu_g = [torch.randn(4, 4, dtype=torch.float32).npu()] | ||
| 304 | + npu_b = [torch.randn(4, 4, dtype=torch.float32).npu()] | ||
| 305 | + cpu_lr = torch.tensor(0.01, dtype=torch.float32) | ||
| 306 | + with self.assertRaises(RuntimeError): | ||
| 307 | + torch._fused_sgd_(npu_p, npu_g, npu_b, weight_decay=0.0, | ||
| 308 | + momentum=0.9, lr=cpu_lr, | ||
| 309 | + dampening=0.0, nesterov=False, | ||
| 310 | + maximize=False, is_first_step=True) | ||
| 311 | + | ||
| 312 | + | ||
| 313 | +if __name__ == "__main__": | ||
| 314 | + run_tests() | ||
需要加meta注册吗