bugfix: remove out large samples from the multi pack batch sampler

larger entries, longer than max_batch_len, may be present in the output of the sampler which gets filtered out in the collator in rare instances, this can lead to empty batches which can be problematic for the training loop Signed-off-by: ChanderG <[email protected]>
instructlab · Jun 19, 2024 · f8ae22b · f8ae22b
1 parent 1cb03c5
commit f8ae22b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/instructlab/training/multipack_sampler.py b/src/instructlab/training/multipack_sampler.py
@@ -282,6 +282,9 @@ def generate_batches(self, set_stats=False):
             len(self.lengths)
         )
 
+        # remove indices where the entries are longer than batch max length
+        indices = indices[self.lengths[indices] <= self.batch_max_length]
+
         lengths = self.lengths[indices]
         lengths_cumsum = np.cumsum(lengths)