Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / tensor_parallel / __init__.py: 93%

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# pylint: disable=undefined-all-variable 

16"""Tensor parallel declarative APIs (parallel styles and module parallelization).""" 

17from importlib import import_module as _import_module # pylint: disable=invalid-name 

18 

19from hyper_parallel.core.tensor_parallel.api import parallelize_module 

20from hyper_parallel.core.tensor_parallel.style import ( 

21 ColwiseParallel, 

22 NoParallel, 

23 ParallelStyle, 

24 PrepareModuleInput, 

25 PrepareModuleInputOutput, 

26 PrepareModuleOutput, 

27 RowwiseParallel, 

28 SequenceParallel, 

29) 

30from hyper_parallel.core.tensor_parallel.loss_parallel import ( 

31 loss_parallel, 

32 is_loss_parallel_active, 

33) 

34 

35# MC2 fused kernels import torch at module load; keep them off the eager path so 

36# MindSpore-only environments can import this package without torch installed. 

37_LAZY_EXPORTS = { 

38 "MC2Linear": ".mc2", 

39 "MC2ColwiseParallel": ".mc2_style", 

40 "MC2RowwiseParallel": ".mc2_style", 

41} 

42 

43 

44def __getattr__(name): # pylint: disable=invalid-name 

45 """Lazily import MC2 symbols that require torch.""" 

46 if name not in _LAZY_EXPORTS: 

47 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 

48 module = _import_module(_LAZY_EXPORTS[name], __name__) 

49 value = getattr(module, name) 

50 globals()[name] = value 

51 return value 

52 

53 

54def __dir__(): # pylint: disable=invalid-name 

55 """Include lazy MC2 exports in ``dir()``.""" 

56 return sorted(set(globals()) | set(_LAZY_EXPORTS)) 

57 

58 

59__all__ = [ 

60 # Parallel styles 

61 "ColwiseParallel", 

62 "MC2ColwiseParallel", 

63 "MC2RowwiseParallel", 

64 "NoParallel", 

65 "ParallelStyle", 

66 "PrepareModuleInput", 

67 "PrepareModuleInputOutput", 

68 "PrepareModuleOutput", 

69 "RowwiseParallel", 

70 "SequenceParallel", 

71 # MC2 fused linear 

72 "MC2Linear", 

73 # Module parallelization 

74 "parallelize_module", 

75 # Loss parallel 

76 "loss_parallel", 

77 "is_loss_parallel_active", 

78]