Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / special_handlers.py: 80%

1# Copyright 2025-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"""special_handlers: special-parameter handler registry (planner Phase 6, 05 §6.4.6). 

16 

17``SPECIAL_HANDLERS`` maps handler names to ``callable(module, param_name, mesh)``; 

18``_SPECIAL_HANDLER_PATTERNS`` maps lowercase FQN substrings to handler names; 

19``_collect_special_handlers`` is the Phase 6 collection pass (SPECIAL-role 

20parameters → handler name). 

21""" 

22 

23from typing import Callable, Dict 

24 

25from torch import nn 

26 

27from hyper_parallel.core.dtensor.dtensor import distribute_tensor 

28from hyper_parallel.core.dtensor.placement_types import Shard 

29from hyper_parallel.distributed.tensor_parallel.param_role import ( 

30 ParamRole, 

31 _match_any, 

32) 

33 

34 

35def _shard_gated_delta(module, param_name, mesh): 

36 """Custom TP sharding skeleton for gated_delta modules (SSM/Mamba-style 

37 modules, 05 §6.4.6). 

38 

39 Shards along the SSM head structure rather than standard 

40 colwise/rowwise. Skeleton implementation: structural recognition plus 

41 a standard Shard(0) fallback; the head-aligned fine-grained sharding is 

42 left to be completed when a concrete model is onboarded. 

43 """ 

44 param = getattr(module, param_name, None) 

45 if param is None: 

46 return 

47 sharded = distribute_tensor(param.data, mesh, [Shard(0)]) 

48 module.register_parameter(param_name, nn.Parameter(sharded)) 

49 

50 

51# {handler_name: callable(module, param_name, mesh)} — Phase B special parameter handlers. 

52SPECIAL_HANDLERS: Dict[str, Callable] = { 

53 "gated_delta_tp_shard": _shard_gated_delta, 

54} 

55 

56# planner-side pattern → handler_name mapping (lowercase fqn substring match). 

57_SPECIAL_HANDLER_PATTERNS: Dict[str, str] = { 

58 "gated_delta": "gated_delta_tp_shard", 

59 "a_log": "gated_delta_tp_shard", 

60 "dt_bias": "gated_delta_tp_shard", 

61} 

62 

63 

64def _collect_special_handlers( 

65 param_roles: Dict[str, ParamRole], 

66 special_handler_patterns: Dict[str, str], 

67) -> Dict[str, str]: 

68 """SPECIAL-role parameters → handler name (unregistered patterns 

69 fall back to "default").""" 

70 result: Dict[str, str] = {} 

71 for fqn, role in param_roles.items(): 

72 if role != ParamRole.SPECIAL: 

73 continue 

74 handler_name = "default" 

75 for pattern, hname in special_handler_patterns.items(): 

76 if _match_any(fqn.lower(), [pattern.lower()]): 

77 handler_name = hname 

78 break 

79 result[fqn] = handler_name 

80 return result