Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / runtime / bandwidth.py: 32%

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"""Hardware bandwidth profiling utilities.""" 

16 

17import torch 

18 

19from hyper_parallel.auto_parallel.hyper_offload.runtime.timer import DeviceTimer 

20 

21 

22def profile_transfer_bandwidth() -> tuple[float, float]: 

23 """Profile host-to-device and device-to-host transfer bandwidth. 

24 

25 Returns: 

26 A tuple of (d2h_bandwidth_gbps, h2d_bandwidth_gbps). 

27 

28 """ 

29 if not torch.accelerator.is_available(): 

30 raise RuntimeError( 

31 "Cannot profile transfer bandwidth: no accelerator available" 

32 ) 

33 

34 device = torch.accelerator.current_accelerator() 

35 size = 16 * 1024 * 1024 

36 src = torch.empty(size, dtype=torch.uint8, device=device) 

37 host = torch.empty(size, dtype=torch.uint8, pin_memory=True) 

38 timer = DeviceTimer() 

39 

40 # D2H bandwidth 

41 timer.start() 

42 host.copy_(src, non_blocking=True) 

43 d2h_ms = timer.stop() 

44 d2h = size / (d2h_ms / 1000.0) / 1024**3 

45 

46 # H2D bandwidth 

47 timer.start() 

48 src.copy_(host, non_blocking=True) 

49 h2d_ms = timer.stop() 

50 h2d = size / (h2d_ms / 1000.0) / 1024**3 

51 

52 return max(d2h, 1.0), max(h2d, 1.0)