add back packing efficiency estimate so epochs and multi-gpu works pr…

…operly (axolotl-ai-cloud#1697)
quangnd2911 · Jun 8, 2024 · ed8ef65 · ed8ef65
1 parent 00ac302
commit ed8ef65
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py
@@ -387,6 +387,7 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
             return MultipackBatchSampler(
                 RandomSampler(self.train_dataset),
                 lengths=get_dataset_lengths(self.train_dataset),
+                packing_efficiency_estimate=self.args.sample_packing_efficiency,
                 batch_max_len=batch_max_len,
                 batch_size=batch_size,
                 group_size=self.args.sample_packing_group_size,
@@ -412,6 +413,7 @@ def _get_eval_sampler(
             return MultipackBatchSampler(
                 SequentialSampler(eval_dataset),
                 lengths=get_dataset_lengths(self.eval_dataset),
+                packing_efficiency_estimate=self.args.sample_packing_efficiency,
                 batch_max_len=batch_max_len,
                 batch_size=batch_size,
                 group_size=self.args.sample_packing_group_size,