Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / megatron / blendable_dataset.py: 22%

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"""Blend multiple datasets with per-source weights. 

16 

17Given ``N`` sub-datasets and ``N`` non-negative weights ``w_i``, a sample at 

18global index ``g`` is read from sub-dataset ``dataset_index[g]`` at local 

19index ``dataset_sample_index[g]``. Both indices are precomputed once (and 

20kept in memory) using a deterministic deficit scheduler, so the same 

21``(weights, num_samples)`` tuple always yields the same blend. ``seed`` is 

22accepted only for API compatibility. 

23 

24Use this to mix corpora — e.g. 30 % code, 70 % web — while keeping the 

25downstream :class:`GPTDataset` API intact (each sub-dataset can be a 

26GPTDataset itself, so blending stacks cleanly). 

27""" 

28import logging 

29from typing import Any, Dict, List, Sequence 

30 

31import numpy as np 

32from torch.utils.data import Dataset 

33 

34 

35logger = logging.getLogger(__name__) 

36 

37 

38def _build_blend_indices( 

39 dataset_sizes: Sequence[int], 

40 weights: Sequence[float], 

41 num_samples: int, 

42 seed: int, # pylint: disable=W0613 

43) -> tuple: 

44 """Return ``(dataset_index, dataset_sample_index)`` for a weighted blend. 

45 

46 At every global slot ``g``, pick the source whose realised fraction lags 

47 its target the most (i.e. ``target * (g+1) - realised`` is maximal). 

48 This guarantees the realised blend tracks ``weights`` exactly up to a 

49 one-sample rounding gap, with no run-to-run variance. 

50 

51 ``seed`` is accepted for API parity but unused; the algorithm is 

52 deterministic in ``(weights, num_samples)``. 

53 """ 

54 if len(dataset_sizes) != len(weights): 

55 raise ValueError( 

56 f"dataset_sizes ({len(dataset_sizes)}) and weights ({len(weights)}) " 

57 f"must have the same length" 

58 ) 

59 if any(w < 0 for w in weights): 

60 raise ValueError(f"weights must be non-negative, got {list(weights)}") 

61 total_w = float(sum(weights)) 

62 if total_w <= 0: 

63 raise ValueError("weights must contain at least one non-zero value") 

64 norm_w = np.asarray([w / total_w for w in weights], dtype=np.float64) 

65 sizes = np.asarray([int(s) for s in dataset_sizes], dtype=np.int64) 

66 # A zero-length source with a non-zero weight would be selected by the 

67 # greatest-error rule and then hit ``counters[d] % sizes[d]`` → divide by 

68 # zero. Reject up-front (a source contributing samples must have samples). 

69 empty_weighted = [i for i, (s, w) in enumerate(zip(sizes, weights)) if s == 0 and w > 0] 

70 if empty_weighted: 

71 raise ValueError( 

72 f"BlendableDataset sources {empty_weighted} have a non-zero weight " 

73 f"but zero length; an empty source cannot contribute samples" 

74 ) 

75 

76 dataset_index = np.zeros(num_samples, dtype=np.int32) 

77 dataset_sample_index = np.zeros(num_samples, dtype=np.int64) 

78 realised = np.zeros(len(weights), dtype=np.float64) 

79 counters = np.zeros(len(weights), dtype=np.int64) 

80 for g in range(num_samples): 

81 # error[d] = how far behind source d's realised fraction is — 

82 # picking the argmax keeps every source within 1 sample of its 

83 # target share at every prefix (Megatron's invariant). 

84 error = norm_w * (g + 1) - realised 

85 d = int(np.argmax(error)) 

86 dataset_index[g] = d 

87 # Modulo wrap so a small sub-dataset reused inside a long blend 

88 # cycles through its samples in order instead of indexing OOB. 

89 dataset_sample_index[g] = counters[d] % sizes[d] 

90 realised[d] += 1.0 

91 counters[d] += 1 

92 return dataset_index, dataset_sample_index 

93 

94 

95class BlendableDataset(Dataset): 

96 """Weighted blend over a list of datasets. 

97 

98 Args: 

99 datasets: Concrete ``torch.utils.data.Dataset`` objects; each must 

100 implement ``__len__`` and ``__getitem__``. Typically these are 

101 :class:`GPTDataset` instances. 

102 weights: One non-negative weight per dataset (auto-normalised). 

103 num_samples: Total length exposed by the blend. 

104 seed: Accepted for API compatibility; index construction is 

105 deterministic and does not use random sampling. 

106 

107 Note: 

108 At every prefix, each source's realised count tracks its weighted 

109 target within the scheduler's one-sample rounding bound. 

110 """ 

111 

112 def __init__( 

113 self, 

114 datasets: List[Dataset], 

115 weights: Sequence[float], 

116 num_samples: int, 

117 seed: int = 1234, 

118 ) -> None: 

119 if not datasets: 

120 raise ValueError("BlendableDataset requires at least one dataset") 

121 self.datasets = datasets 

122 sizes = [len(d) for d in datasets] 

123 self.num_samples = int(num_samples) 

124 self.dataset_index, self.dataset_sample_index = _build_blend_indices( 

125 sizes, weights, self.num_samples, seed, 

126 ) 

127 logger.info( 

128 "BlendableDataset built: %d sub-datasets, weights=%s, total samples=%d", 

129 len(datasets), list(weights), self.num_samples, 

130 ) 

131 

132 def __len__(self) -> int: 

133 return self.num_samples 

134 

135 def __getitem__(self, idx: int) -> Dict[str, Any]: 

136 d = int(self.dataset_index[idx]) 

137 s = int(self.dataset_sample_index[idx]) 

138 return self.datasets[d][s]