Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / mpipe / sampler.py: 41%

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"""Per-PP-rank disk-sharded sampler for MPipe + ``data.load: single``. 

16 

17Framework-agnostic: only manipulates integer indices, never touches sample 

18data. 

19""" 

20from typing import Any, Dict, Iterable, Iterator 

21 

22 

23class PPRankOwnedSampler: 

24 """Filter a DP-sharded base sampler to indices for ``owned_micros`` only. 

25 

26 Used with MPipe + ``data.load: "single"`` to shard disk reads across PP 

27 ranks. The base sampler emits a deterministic stream of global indices 

28 (same seed on every rank); within each pipeline step the next 

29 ``pp_micro_batch_num * micro_batch_size`` indices are laid out 

30 micro-major, and this wrapper yields only the slots in ``owned_micros``. 

31 """ 

32 

33 def __init__(self, base_sampler: Iterable[int], owned_micros: Iterable[int], 

34 micro_batch_size: int, pp_micro_batch_num: int) -> None: 

35 """Wrap ``base_sampler``, keeping only this rank's micro-batch slots. 

36 

37 Args: 

38 base_sampler (Iterable[int]): Deterministic global-index stream. 

39 owned_micros (Iterable[int]): Micro indices this rank loads. 

40 micro_batch_size (int): Samples per micro-batch. 

41 pp_micro_batch_num (int): Micro-batches per pipeline step. 

42 """ 

43 self.base = base_sampler 

44 self.owned_micros = sorted(owned_micros) 

45 self.micro_bs = int(micro_batch_size) 

46 self.step_size = int(pp_micro_batch_num) * self.micro_bs 

47 

48 def __iter__(self) -> Iterator[int]: 

49 """Yield only the owned slots of each micro-major step chunk.""" 

50 chunk = [] 

51 for idx in self.base: 

52 chunk.append(idx) 

53 if len(chunk) == self.step_size: 

54 for m in self.owned_micros: 

55 yield from chunk[m * self.micro_bs : (m + 1) * self.micro_bs] 

56 chunk = [] 

57 

58 def __len__(self) -> int: 

59 """Number of indices yielded per epoch (trailing partial step dropped).""" 

60 full_chunks = len(self.base) // self.step_size 

61 return full_chunks * len(self.owned_micros) * self.micro_bs 

62 

63 def set_epoch(self, epoch: int) -> None: 

64 """Forward the epoch to the base sampler when it supports reseeding.""" 

65 if hasattr(self.base, "set_epoch"): 

66 self.base.set_epoch(epoch) 

67 

68 def state_dict(self) -> Dict[str, Any]: 

69 """Return the base sampler's state, or an empty dict if it has none.""" 

70 if hasattr(self.base, "state_dict"): 

71 return {"base": self.base.state_dict()} 

72 return {} 

73 

74 def load_state_dict(self, state: Dict[str, Any]) -> None: 

75 """Restore the base sampler's state from :meth:`state_dict` output.""" 

76 if "base" in state and hasattr(self.base, "load_state_dict"): 

77 self.base.load_state_dict(state["base"]) 

78 

79 

80def mpipe_owned_micros(pp_size: int, micro_batch_num: int, pp_rank: int, 

81 mode: str = "full") -> frozenset: 

82 """Micro indices PP rank ``pp_rank`` loads under MPipe (``NT = min(PP, M)``). 

83 

84 ``mode="full"`` (default): round-robin — rank ``i < NT`` owns 

85 ``{m | m % NT == i}``. ``mode="min"``: rank 0 owns ``{0, NT..M-1}``, 

86 ranks ``1..NT-1`` own ``{rank}``. Ranks ``i >= NT`` own nothing; for 

87 ``M <= NT`` both modes collapse to ``{rank}`` per rank. 

88 """ 

89 nt = min(pp_size, micro_batch_num) 

90 if pp_rank >= nt: 

91 return frozenset() 

92 if mode == "full": 

93 return frozenset(m for m in range(micro_batch_num) if m % nt == pp_rank) 

94 if mode == "min": 

95 if pp_rank == 0: 

96 return frozenset({0, *range(nt, micro_batch_num)}) 

97 return frozenset({pp_rank}) 

98 raise ValueError(f"mpipe transpose overflow mode must be 'full' or 'min', got {mode!r}")