| 236 | # =>> Reference: https://github.com/haotian-liu/LLaVA/blob/main/llava/train/llava_trainer.py#L60 |
| 237 | # =>> Reference: https://github.com/huggingface/transformers/blob/main/src/transformers/trainer_pt_utils.py#L603 |
| 238 | class SplitModalitySampler(Sampler): |
| 239 | def __init__( |
| 240 | self, |
| 241 | dataset: Dataset, |
| 242 | modality_lengths: List[Tuple[bool, int]], |
| 243 | global_batch_size: int, |
| 244 | num_replicas: Optional[int] = None, |
| 245 | rank: Optional[int] = None, |
| 246 | seed: int = 0, |
| 247 | drop_last: bool = False, |
| 248 | ) -> None: |
| 249 | super().__init__() |
| 250 | self.num_replicas = num_replicas if num_replicas is not None else dist.get_world_size() |
| 251 | self.rank = rank if rank is not None else dist.get_rank() |
| 252 | self.seed, self.epoch = seed, 0 |
| 253 | |
| 254 | # Custom Parameters |
| 255 | self.dataset, self.modality_lengths, self.drop_last = dataset, modality_lengths, drop_last |
| 256 | self.global_batch_size = global_batch_size |
| 257 | |
| 258 | # For our purposes, `drop_last` is always False! |
| 259 | assert not self.drop_last, "SplitModalitySampler must set `drop_last = False`!" |
| 260 | self.total_size = math.ceil(len(self.dataset) / self.global_batch_size) * self.global_batch_size |
| 261 | self.num_samples = self.total_size // self.num_replicas |
| 262 | |
| 263 | @staticmethod |
| 264 | def reindex_batch(batch_idxs: List[int], idx2lengths: List[int], n_buckets: int) -> List[List[int]]: |
| 265 | """Re-indexes a batch in a way that is conducive to DistributedSampler + grouping by seqlen per rank.""" |
| 266 | assert len(batch_idxs) % n_buckets == 0, "Batch length is not divisible by `num_replicas`!" |
| 267 | |
| 268 | # Establish initial buckets, capacities, and max number of elements per bucket |
| 269 | n_examples_per_bucket = len(batch_idxs) // n_buckets |
| 270 | bucket_indices = [[] for _ in range(n_buckets)] |
| 271 | bucket_lengths = [0 for _ in range(n_buckets)] |
| 272 | |
| 273 | # Note that `batch_idxs` is already sorted by corresponding length (in descending order) |
| 274 | for idx in batch_idxs: |
| 275 | shortest_bucket_idx = bucket_lengths.index(min(bucket_lengths)) |
| 276 | bucket_indices[shortest_bucket_idx].append(idx) |
| 277 | |
| 278 | # Update `bucket_lengths` --> set length to infinity if at capacity! |
| 279 | bucket_lengths[shortest_bucket_idx] += idx2lengths[idx] |
| 280 | if len(bucket_indices[shortest_bucket_idx]) == n_examples_per_bucket: |
| 281 | bucket_lengths[shortest_bucket_idx] = float("inf") |
| 282 | |
| 283 | return bucket_indices |
| 284 | |
| 285 | def get_modality_and_length_grouped_indices(self, generator: torch.Generator) -> List[int]: |
| 286 | """ |
| 287 | Returns a list of indices so that each slice of `global_batch_size` consecutive indices corresponds to elements |
| 288 | of the same modality with each sub-sequence of `per_replica_batch_size` (the batch size each unique device sees |
| 289 | during distributed training) is roughly grouped by sequence length (for training efficiency). |
| 290 | """ |
| 291 | multimodal_indices, multimodal_lengths = zip( |
| 292 | *[(idx, length) for idx, (is_multimodal, length) in enumerate(self.modality_lengths) if is_multimodal] |
| 293 | ) |
| 294 | |
| 295 | # Handle Special Case --> no "unimodal" inputs |