Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / preset_pt.py: 23%

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"""``preset_pt`` dataset for replaying pre-tokenized ``.pt`` batches. 

16 

17Each entry is either a stacked batch dict ``{key: (B, S)-Tensor}`` or a list 

18of per-rank dicts. The loader flattens both forms into per-sample rows so the 

19standard ``DataLoader`` can batch them again with the trainer collator. 

20""" 

21import logging 

22from typing import Any, Dict, List 

23 

24import torch 

25from torch.utils.data import Dataset 

26 

27from hyper_parallel.data.registry import DATASET_REGISTRY 

28 

29 

30logger = logging.getLogger(__name__) 

31 

32 

33class PresetPtDataset(Dataset): 

34 """Wrap a pre-expanded list of per-sample dicts.""" 

35 

36 def __init__(self, samples: List[Dict[str, Any]]) -> None: 

37 self.samples = samples 

38 

39 def __len__(self) -> int: 

40 return len(self.samples) 

41 

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

43 return self.samples[idx] 

44 

45 

46_OPTIONAL_2D_FIELDS = ("attention_mask", "mm_token_type_ids") 

47 

48_PIXEL_PAIRS = (("pixel_values", "image_grid_thw"), 

49 ("pixel_values_videos", "video_grid_thw")) 

50 

51 

52def _split_pixel_block( 

53 b: Dict[str, Any], i: int, batch_size: int, pix_key: str, grid_key: str, 

54) -> Dict[str, Any]: 

55 """Slice the ``(pix_key, grid_key)`` rows owned by sample ``i``.""" 

56 pv = b[pix_key] 

57 thw = b[grid_key] 

58 grids_per_sample = thw.shape[0] // batch_size if thw.dim() == 2 else 0 

59 if grids_per_sample == 0: 

60 return {} 

61 thw_i = thw[i * grids_per_sample:(i + 1) * grids_per_sample].clone() 

62 pv_count = int(thw_i.prod(dim=-1).sum().item()) 

63 offset = sum( 

64 int(thw[j].prod(dim=-1).sum().item()) 

65 for j in range(i * grids_per_sample) 

66 ) 

67 return { 

68 pix_key: pv[offset:offset + pv_count].clone(), 

69 grid_key: thw_i, 

70 } 

71 

72 

73def _slice_position_ids(position_ids: Any, sample_idx: int, batch_size: int) -> Any: 

74 """Return one sample's position ids from an LF/HF stacked batch. 

75 

76 Text Qwen3.5 batches may carry either ``[B, S]`` plain position ids or 

77 mRoPE ``[R, B, S]`` ids (``R`` is 3 or 4 in Transformers). Preserve the 

78 rotary-rank axis and slice only the batch axis so the trainer can rebuild 

79 the original stacked shape in its collate function. 

80 """ 

81 if position_ids.dim() == 2 and position_ids.shape[0] == batch_size: 

82 return position_ids[sample_idx].clone() 

83 if position_ids.dim() == 3 and position_ids.shape[1] == batch_size: 

84 return position_ids[:, sample_idx].clone() 

85 if position_ids.dim() == 3 and position_ids.shape[0] == batch_size: 

86 return position_ids[sample_idx].clone() 

87 raise ValueError( 

88 "preset_pt position_ids must be [B, S], [R, B, S], or [B, R, S]; " 

89 f"got shape={tuple(position_ids.shape)} with batch_size={batch_size}." 

90 ) 

91 

92 

93def _expand_batch(b: Dict[str, Any], *, vl: bool) -> List[Dict[str, Any]]: 

94 """Split a stacked LM/VL batch dict into per-sample dicts. 

95 

96 LM samples carry only ``input_ids`` / ``labels`` (plus optionally 

97 ``attention_mask``). VL samples additionally carry ``mm_token_type_ids`` 

98 and ``(pixel_values, image_grid_thw)`` / ``(pixel_values_videos, 

99 video_grid_thw)`` pairs sliced according to the per-sample grid product. 

100 """ 

101 ids = b["input_ids"] 

102 labels = b["labels"] 

103 batch_size = ids.shape[0] 

104 out: List[Dict[str, Any]] = [] 

105 for i in range(batch_size): 

106 rec: Dict[str, Any] = { 

107 "input_ids": ids[i].clone(), 

108 "labels": labels[i].clone(), 

109 } 

110 if "num_items_in_batch" in b: 

111 rec["num_items_in_batch"] = int((rec["labels"] != -100).sum().item()) 

112 for k in _OPTIONAL_2D_FIELDS: 

113 v = b.get(k) 

114 if v is not None and v.dim() == 2: 

115 rec[k] = v[i].clone() 

116 position_ids = b.get("position_ids") 

117 if position_ids is not None: 

118 rec["position_ids"] = _slice_position_ids(position_ids, i, batch_size) 

119 if vl: 

120 for pix_key, grid_key in _PIXEL_PAIRS: 

121 if b.get(pix_key) is not None and b.get(grid_key) is not None: 

122 rec.update(_split_pixel_block(b, i, batch_size, pix_key, grid_key)) 

123 out.append(rec) 

124 return out 

125 

126 

127def _is_vl(batch_entry: Any) -> bool: 

128 """Heuristic: VL batches always carry pixel data (image or video).""" 

129 if isinstance(batch_entry, list): 

130 batch_entry = batch_entry[0] if batch_entry else None 

131 return isinstance(batch_entry, dict) and any( 

132 pix_key in batch_entry for pix_key, _ in _PIXEL_PAIRS) 

133 

134 

135@DATASET_REGISTRY.register("preset_pt") 

136def build_preset_pt(*, base: Any, args: Any, **_: Any) -> PresetPtDataset: 

137 """Build the preset replay dataset. 

138 

139 Auto-detects whether the .pt holds VL batches (pixel_values present) 

140 or plain LM batches by inspecting the first entry, and dispatches the 

141 matching per-sample expansion. 

142 """ 

143 train_path = args.data.train_path 

144 if not train_path: 

145 raise ValueError("data.train_path is required when data.type='preset_pt'") 

146 batches = torch.load(train_path, map_location="cpu", weights_only=False) 

147 if not isinstance(batches, list) or not batches: 

148 raise ValueError(f"preset_pt expects List, got {type(batches)}") 

149 

150 vl = _is_vl(batches[0]) 

151 per_sample: List[Dict[str, Any]] = [] 

152 for b in batches: 

153 if isinstance(b, list): 

154 for br in b: 

155 per_sample.extend(_expand_batch(br, vl=vl)) 

156 else: 

157 per_sample.extend(_expand_batch(b, vl=vl)) 

158 

159 ds = PresetPtDataset(per_sample) 

160 if args.train.max_steps: 

161 base.state.max_steps = int(args.train.max_steps) 

162 logger.info( 

163 "preset_pt dataset (%s): %d samples loaded from %s", 

164 "vl" if vl else "lm", len(per_sample), train_path, 

165 ) 

166 return ds