Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / _precision.py: 94%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Preserve full-precision control parameters during ordinary model dtype conversion.""" 

16 

17from collections.abc import Callable, Iterable 

18 

19import torch 

20 

21 

22def preserve_fp32_parameters( 

23 module: torch.nn.Module, transform: Callable[[torch.Tensor], torch.Tensor], names: Iterable[str], 

24) -> Callable[[torch.Tensor], torch.Tensor]: 

25 """Wrap Module._apply without rounding protected Parameters or their existing gradients. 

26 

27 Args: 

28 module: Module holding the protected control Parameters. 

29 transform: Requested parameter device or dtype transformation. 

30 names: Names of control Parameters which must retain FP32 values. 

31 """ 

32 protected = set() 

33 for name in names: 

34 parameter = getattr(module, name, None) 

35 if parameter is not None: 

36 protected.add(id(parameter)) 

37 if parameter.grad is not None: 

38 protected.add(id(parameter.grad)) 

39 

40 def apply(tensor: torch.Tensor) -> torch.Tensor: 

41 """Retain source FP32 values while following the requested target device. 

42 

43 Args: 

44 tensor: Parameter or gradient being transformed. 

45 """ 

46 converted = transform(tensor) 

47 if id(tensor) not in protected or converted.dtype == torch.float32: 

48 return converted 

49 if tensor.is_meta: 

50 return converted.to(dtype=torch.float32) 

51 return tensor.to(device=converted.device, dtype=torch.float32) 

52 

53 return apply