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 contextlib import contextmanager 

19 

20import torch 

21 

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

23 

24 

25@contextmanager 

26def init_empty_weights(include_buffers=False): 

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

28 buffers) on the meta device. 

29 

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

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

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

33 

34 Args: 

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

36 device. Defaults to ``False``. 

37 

38 Example:: 

39 

40 from hyper_parallel import init_empty_weights 

41 

42 with init_empty_weights(): 

43 model = MyLargeModel() # no memory allocated 

44 

45 Note: 

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

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

48 """ 

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

50 yield 

51 

52 

53@contextmanager 

54def init_on_device(device, include_buffers=False): 

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

56 

57 Args: 

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

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

60 Defaults to ``False``. 

61 

62 Example:: 

63 

64 from hyper_parallel import init_on_device 

65 import torch 

66 

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

68 model = MyModel() # parameters live on Ascend 

69 """ 

70 with _init_on_device(device, include_buffers=include_buffers): 

71 yield