Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_mhc_grad / common.py: 100%

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"""Runtime gradient types and forward cache for HyperMegaMhcGrad.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass 

20 

21import torch 

22 

23from ..mega_mhc_common import MHC_MAPPINGS, MHC_SINKHORN_ITERS, MHC_STREAMS, validate_mega_mhc_tensor 

24 

25 

26MegaMhcGradOutputs = tuple[ 

27 torch.Tensor, 

28 torch.Tensor, 

29 torch.Tensor, 

30 torch.Tensor, 

31 torch.Tensor, 

32 torch.Tensor, 

33 torch.Tensor, 

34 torch.Tensor, 

35 torch.Tensor, 

36] 

37 

38 

39@dataclass(frozen=True) 

40class MegaMhcGradCache: 

41 """Forward tensors required by the native backward kernels.""" 

42 

43 new_residual: torch.Tensor 

44 hc_before_norm: torch.Tensor 

45 inv_rms: torch.Tensor 

46 sum_out: torch.Tensor 

47 norm_out: torch.Tensor 

48 mixed_input: torch.Tensor 

49 rms_rstd: torch.Tensor 

50 

51 

52def validate_mega_mhc_grad_tensors( 

53 gradients: tuple[torch.Tensor, ...], residual: torch.Tensor, cache: MegaMhcGradCache, 

54) -> None: 

55 """Validate upstream gradients and every native forward-cache tensor. 

56 

57 Args: 

58 gradients: Gradients of the five public forward outputs, in output order. 

59 residual: Original residual streams defining the shape and device. 

60 cache: Forward tensors retained for native backward. 

61 """ 

62 if not isinstance(cache, MegaMhcGradCache): 

63 raise TypeError("cache must be a MegaMhcGradCache.") 

64 leading_shape = tuple(residual.shape[:-2]) 

65 hidden_size = residual.shape[-1] 

66 tokens = residual.numel() // (MHC_STREAMS * hidden_size) 

67 gradient_specs = ( 

68 ("grad_new_residual", (*leading_shape, MHC_STREAMS, hidden_size), torch.bfloat16), 

69 ("grad_next_pre_mix", (*leading_shape, MHC_STREAMS), torch.float32), 

70 ("grad_next_post_mix", (*leading_shape, MHC_STREAMS), torch.float32), 

71 ("grad_next_residual_mix", (*leading_shape, MHC_STREAMS, MHC_STREAMS), torch.float32), 

72 ("grad_block_input", (*leading_shape, hidden_size), torch.bfloat16), 

73 ) 

74 if len(gradients) != len(gradient_specs): 

75 raise ValueError(f"Expected five output gradients, got {len(gradients)}.") 

76 for value, (name, shape, dtype) in zip(gradients, gradient_specs): 

77 validate_mega_mhc_tensor(name, value, shape, dtype, residual.device) 

78 cache_specs = ( 

79 ("new_residual", cache.new_residual, (1, tokens, MHC_STREAMS, hidden_size), torch.bfloat16), 

80 ("hc_before_norm", cache.hc_before_norm, (1, tokens, MHC_MAPPINGS), torch.float32), 

81 ("inv_rms", cache.inv_rms, (1, tokens, 1), torch.float32), 

82 ("sum_out", cache.sum_out, (2 * MHC_SINKHORN_ITERS, 1, tokens, MHC_STREAMS), torch.float32), 

83 ("norm_out", cache.norm_out, (2 * MHC_SINKHORN_ITERS, 1, tokens, MHC_STREAMS, MHC_STREAMS), torch.float32), 

84 ("mixed_input", cache.mixed_input, (1, tokens, hidden_size), torch.bfloat16), 

85 ("rms_rstd", cache.rms_rstd, (1, tokens, 1), torch.float32), 

86 ) 

87 for name, value, shape, dtype in cache_specs: 

88 validate_mega_mhc_tensor(f"cache.{name}", value, shape, dtype, residual.device)