Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / batching / build_collate_fn.py: 43%

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"""Build final batch collators for Indexed and Online sources.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Callable, Mapping, Sequence 

20from dataclasses import dataclass 

21from typing import Any 

22 

23import torch 

24from torch.utils.data import default_collate 

25 

26from hyper_parallel.data.constants import IGNORE_INDEX, OMNI_CONCAT_FIELDS 

27 

28 

29def _get_sequence_parallel_size(mesh_context: Any | None) -> int: 

30 """Return the unified sequence-sharding degree used for batch alignment.""" 

31 if mesh_context is None: 

32 return 1 

33 

34 cp_size = int(getattr(mesh_context, "cp_size", 1)) 

35 tp_size = int(getattr(mesh_context, "tp_size", 1)) 

36 sequence_parallel = bool(getattr(mesh_context, "sequence_parallel", False)) 

37 if cp_size <= 0 or tp_size <= 0: 

38 raise ValueError("mesh_context tp_size and cp_size must be positive") 

39 

40 return cp_size * (tp_size if sequence_parallel else 1) 

41 

42 

43@dataclass 

44class TextPackingCollator: 

45 """Pack unpadded Online text samples and emit ``cu_seq_lens``. 

46 

47 Only ``input_ids`` and pre-shifted ``labels`` are packed here. Unified 

48 ``get_batch`` constructs loss, position, mask, and CP runtime fields. 

49 """ 

50 

51 sequence_parallel_size: int = 1 

52 

53 def __call__(self, model_samples: Sequence[Mapping[str, Any]]) -> Mapping[str, Any]: 

54 """Pack samples into one ``[1, packed_length]`` forward-backward batch. 

55 

56 Args: 

57 model_samples: Online samples selected for one forward-backward step. 

58 

59 Returns: 

60 Packed ``input_ids`` and ``labels`` plus int32 ``cu_seq_lens``. 

61 

62 Raises: 

63 ValueError: If no samples are provided. 

64 """ 

65 if not model_samples: 

66 raise ValueError("model_samples must contain at least one Online sample") 

67 

68 packed_batch = {} 

69 for field in ("input_ids", "labels"): 

70 values = [model_sample[field] for model_sample in model_samples] 

71 packed_batch[field] = torch.cat(values, dim=-1).unsqueeze(0) 

72 

73 input_ids = packed_batch.get("input_ids") 

74 if input_ids is None or packed_batch.get("labels") is None: 

75 raise ValueError("packed batch requires input_ids and labels") 

76 

77 packed_seq_len = input_ids.shape[-1] 

78 pad_len = (-packed_seq_len) % self.sequence_parallel_size 

79 if pad_len: 

80 input_padding = input_ids.new_zeros((1, pad_len)) 

81 label_padding = packed_batch["labels"].new_full((1, pad_len), IGNORE_INDEX) 

82 packed_batch["input_ids"] = torch.cat((input_ids, input_padding), dim=-1) 

83 packed_batch["labels"] = torch.cat((packed_batch["labels"], label_padding), dim=-1) 

84 

85 seq_lens = model_samples[0]["input_ids"].new_tensor( 

86 [model_sample["input_ids"].shape[-1] for model_sample in model_samples] 

87 ) 

88 zero = seq_lens.new_zeros(1) 

89 seq_ends = seq_lens.cumsum(dim=0) 

90 if pad_len: 

91 # Represent the alignment tail as one synthetic packed sequence so 

92 # attention metadata covers every physical Q/KV token. Its labels 

93 # remain IGNORE_INDEX and therefore do not contribute to the loss. 

94 padded_end = seq_ends[-1:] + pad_len 

95 seq_ends = torch.cat((seq_ends, padded_end)) 

96 cu_seq_lens = torch.cat((zero, seq_ends)) 

97 cu_seq_lens = cu_seq_lens.to(torch.int32) 

98 packed_batch["cu_seq_lens"] = cu_seq_lens 

99 

100 return packed_batch 

101 

102 

103class OmniCollator: 

104 """Batch already prepared Omni samples without performing sequence packing.""" 

105 

106 def __call__(self, model_samples: Sequence[Mapping[str, Any]]) -> Mapping[str, Any]: 

107 """Stack token fields and concatenate variable-length modality fields.""" 

108 if not model_samples: 

109 raise ValueError("model_samples must contain at least one Omni sample") 

110 

111 sample_fields = set() 

112 for model_sample in model_samples: 

113 sample_fields.update(model_sample) 

114 

115 batch = {} 

116 for field in sample_fields: 

117 values = [] 

118 for model_sample in model_samples: 

119 if field in model_sample: 

120 values.append(model_sample[field]) 

121 

122 if field in OMNI_CONCAT_FIELDS: 

123 batch[field] = torch.cat(values, dim=0) 

124 else: 

125 batch[field] = default_collate(values) 

126 

127 return batch 

128 

129 

130def build_default_collate_fn() -> Callable[[list[Any]], Any]: 

131 """Build ordinary PyTorch collation for already prepared samples. 

132 

133 Returns: 

134 PyTorch default collation. 

135 """ 

136 collate_fn = default_collate 

137 

138 return collate_fn 

139 

140 

141# Indexed samples are already prepared and use the same ordinary collation. 

142build_indexed_collate_fn = build_default_collate_fn 

143 

144 

145def build_omni_collate_fn() -> OmniCollator: 

146 """Build ordinary collation for encoded or finally packed Omni samples.""" 

147 collate_fn = OmniCollator() 

148 return collate_fn 

149 

150 

151def build_online_text_collate_fn(mesh_context: Any | None = None) -> TextPackingCollator: 

152 """Build Online text collation shared by fixed N and dynamic K batching. 

153 

154 Packing concatenates ``input_ids`` and ``labels`` and emits ``cu_seq_lens``. 

155 Tail padding is derived from the runtime topology rather than user config. 

156 The unified alignment size is ``cp_size * tp_size`` when TP sequence 

157 parallelism is enabled, otherwise it is ``cp_size``. Real sample boundaries 

158 remain unchanged and one synthetic final boundary covers the physical 

159 padding segment. 

160 

161 Args: 

162 mesh_context: Runtime TP/CP topology injected by the Trainer. 

163 

164 Returns: 

165 A collator producing one forward-backward batch. 

166 """ 

167 sequence_parallel_size = _get_sequence_parallel_size(mesh_context) 

168 collate_fn = TextPackingCollator(sequence_parallel_size=sequence_parallel_size) 

169 

170 return collate_fn