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

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"""Megatron dataset builder — wires config → IndexedDataset → GPTDataset (+ blend). 

16 

17Config surface (``args.data.*``): 

18 

19- ``train_path``: either a single path-prefix (no suffix) pointing to one 

20 ``.bin``/``.idx`` pair, OR a flat ``"w1 path1 w2 path2 ..."`` string / 

21 list for :class:`BlendableDataset`. 

22- ``max_seq_len``: per-sample token count (input plus implicit label shift). 

23- ``megatron_seed`` (optional, default ``train.seed``): RNG seed for the 

24 document / sample / shuffle indices. 

25- ``pad_token_id`` (optional, default 0): right-pad for the tail sample. 

26- ``eod_token_id`` (optional, default None) + ``eod_mask_loss`` (optional, 

27 default False): mask EOD tokens out of the loss. 

28 

29Sample count: ``num_samples = max_steps * global_batch_size`` so the 

30distributed sampler can drain the dataset without short cycles. Unlike the 

31HF builder, this builder does NOT modify ``base.state.max_steps``: Megatron 

32datasets are typically larger than ``num_samples``, so the configured 

33``max_steps`` is used as-is with no truncation. 

34""" 

35import logging 

36from typing import Any, List, Tuple 

37 

38from hyper_parallel.data.megatron.blendable_dataset import BlendableDataset 

39from hyper_parallel.data.megatron.gpt_dataset import GPTDataset 

40from hyper_parallel.data.megatron.indexed_dataset import IndexedDataset, strip_suffix 

41from hyper_parallel.data.registry import DATASET_REGISTRY 

42 

43 

44logger = logging.getLogger(__name__) 

45 

46 

47def _parse_blend(raw: Any) -> List[Tuple[float, str]]: 

48 """Parse a blend spec into ``[(weight, prefix), ...]``. 

49 

50 Supported shapes: 

51 

52 - ``str`` of whitespace-separated tokens: ``"0.3 /data/a 0.7 /data/b"`` 

53 - ``list[float | str]`` in the same alternating order 

54 - ``list[list]`` of ``[weight, prefix]`` pairs 

55 """ 

56 if isinstance(raw, str): 

57 toks = raw.split() 

58 elif isinstance(raw, (list, tuple)): 

59 # Already in pair form? 

60 if raw and isinstance(raw[0], (list, tuple)): 

61 return [(float(w), str(p)) for w, p in raw] 

62 toks = list(raw) 

63 else: 

64 raise ValueError(f"Unsupported blend spec type: {type(raw)}") 

65 

66 if len(toks) % 2 != 0: 

67 raise ValueError( 

68 f"Blend spec must have an even number of tokens " 

69 f"(alternating weight + prefix); got {len(toks)}: {toks}" 

70 ) 

71 out: List[Tuple[float, str]] = [] 

72 for i in range(0, len(toks), 2): 

73 out.append((float(toks[i]), str(toks[i + 1]))) 

74 return out 

75 

76 

77def _looks_like_blend(train_path: Any) -> bool: 

78 """Return ``True`` when ``train_path`` is a multi-source blend spec. 

79 

80 A blend spec always starts with a numeric weight (``"0.3 /a 0.7 /b"`` or 

81 a list of pairs); a single corpus is a bare path-prefix. Keying off "the 

82 first token parses as a float" (rather than "contains whitespace") lets a 

83 single path that legitimately contains spaces still be read as one source. 

84 """ 

85 if isinstance(train_path, (list, tuple)): 

86 return True 

87 if not isinstance(train_path, str): 

88 return False 

89 toks = train_path.split() 

90 if len(toks) < 2: 

91 return False 

92 try: 

93 float(toks[0]) 

94 except ValueError: 

95 return False 

96 return True 

97 

98 

99@DATASET_REGISTRY.register("megatron") 

100def build_megatron(*, base: Any, args: Any, **_: Any) -> Any: 

101 """Build a Megatron ``.bin``/``.idx`` dataset (single source or blend).""" 

102 del base 

103 data_cfg = args.data 

104 train_cfg = args.train 

105 train_path = data_cfg.train_path 

106 if not train_path: 

107 raise ValueError("data.train_path is required when data.type='megatron'") 

108 

109 seq_length = int(data_cfg.max_seq_len) 

110 # ``megatron_seed`` may be present-but-None on a default DataConfig; fall 

111 # through to ``train.seed`` in that case so the typed config doesn't have 

112 # to invent a sentinel. 

113 megatron_seed = data_cfg.megatron_seed 

114 if megatron_seed is None: 

115 megatron_seed = train_cfg.seed 

116 seed = int(megatron_seed) 

117 pad_token_id = int(data_cfg.pad_token_id) 

118 eod_token_id = data_cfg.eod_token_id 

119 eod_mask_loss = bool(data_cfg.eod_mask_loss) 

120 mmap_bin = bool(data_cfg.mmap_bin_files) 

121 

122 num_samples = int(train_cfg.max_steps * train_cfg.global_batch_size) 

123 if num_samples <= 0: 

124 raise ValueError( 

125 f"num_samples = max_steps * global_batch_size must be > 0; " 

126 f"got max_steps={train_cfg.max_steps}, global_bs={train_cfg.global_batch_size}" 

127 ) 

128 

129 # Single source: simplest path, no blend needed. 

130 if not _looks_like_blend(train_path): 

131 prefix = strip_suffix(train_path) 

132 idx_ds = IndexedDataset(prefix, mmap=mmap_bin) 

133 gpt_ds = GPTDataset( 

134 idx_ds, num_samples=num_samples, seq_length=seq_length, 

135 seed=seed, pad_token_id=pad_token_id, 

136 eod_mask_loss=eod_mask_loss, eod_token_id=eod_token_id, 

137 ) 

138 logger.info( 

139 "Megatron dataset built: prefix=%s docs=%d sequences=%d samples=%d seq_len=%d", 

140 prefix, idx_ds.num_documents, len(idx_ds), len(gpt_ds), seq_length, 

141 ) 

142 return gpt_ds 

143 

144 # Blend: parse, build per-source GPTDataset, wrap in BlendableDataset. 

145 pairs = _parse_blend(train_path) 

146 weights = [w for w, _ in pairs] 

147 sub_datasets = [] 

148 for w, p in pairs: 

149 prefix = strip_suffix(p) 

150 idx_ds = IndexedDataset(prefix, mmap=mmap_bin) 

151 # Each sub-dataset is sized to cover its share of the blend; we ask 

152 # for the full ``num_samples`` here because blending picks indices 

153 # modulo the sub-dataset length. 

154 sub_datasets.append( 

155 GPTDataset( 

156 idx_ds, num_samples=num_samples, seq_length=seq_length, 

157 seed=seed, pad_token_id=pad_token_id, 

158 eod_mask_loss=eod_mask_loss, eod_token_id=eod_token_id, 

159 ) 

160 ) 

161 logger.info( 

162 " blend source w=%g prefix=%s docs=%d sequences=%d", 

163 w, prefix, idx_ds.num_documents, len(idx_ds), 

164 ) 

165 return BlendableDataset(sub_datasets, weights, num_samples=num_samples, seed=seed)