Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / init_weights.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"""Empty / on-device weight initialization utilities.""" 

16# pylint: disable=C9006,C9007 

17 

18from collections.abc import Iterator 

19from contextlib import contextmanager 

20 

21import torch 

22 

23from hyper_parallel.core.dtensor._utils import init_on_device as _init_on_device 

24 

25 

26@contextmanager 

27def init_empty_weights(include_buffers: bool = False) -> Iterator[None]: 

28 """Context manager that initializes model parameters (and optionally 

29 buffers) on the meta device. 

30 

31 Models created under this context have no real data, which avoids allocating 

32 any NPU/GPU/CPU memory. This is useful when the model is too large to fit in RAM 

33 and the weights will be loaded from a checkpoint afterwards. 

34 

35 Args: 

36 include_buffers (bool): If ``True``, buffers are also placed on the meta 

37 device. Defaults to ``False``. 

38 

39 Example:: 

40 

41 from hyper_parallel import init_empty_weights 

42 

43 with init_empty_weights(): 

44 model = MyLargeModel() # no memory allocated 

45 

46 Note: 

47 A model initialised under this context has **no weights**. You cannot 

48 call ``model.to(some_device)`` directly; load a checkpoint first. 

49 

50 Deterministic memory filling is suspended only inside this context: 

51 factories may allocate CPU tensors before registration moves them to 

52 meta. Filling that discarded storage defeats empty initialization. 

53 The deterministic algorithm setting is unchanged. 

54 """ 

55 fill_uninitialized_memory = torch.utils.deterministic.fill_uninitialized_memory 

56 try: 

57 torch.utils.deterministic.fill_uninitialized_memory = False 

58 with _init_on_device(torch.device("meta"), include_buffers=include_buffers): 

59 yield 

60 finally: 

61 torch.utils.deterministic.fill_uninitialized_memory = fill_uninitialized_memory 

62 

63 

64@contextmanager 

65def init_on_device(device, include_buffers=False): 

66 """Context manager that initializes model parameters (and optionally buffers) on *device*. 

67 

68 Args: 

69 device: Target device for parameter (and buffer) allocation. 

70 include_buffers (bool): If ``True``, buffers are also placed on *device*. 

71 Defaults to ``False``. 

72 

73 Example:: 

74 

75 from hyper_parallel import init_on_device 

76 import torch 

77 

78 with init_on_device(torch.device("npu")): 

79 model = MyModel() # parameters live on Ascend 

80 """ 

81 with _init_on_device(device, include_buffers=include_buffers): 

82 yield