bugfixes

src/axolotl/utils/samplers/multipack.py

@@ -148,8 +148,11 @@ def generate_batches(self, set_stats=False):
         )
 
         batches = [
-            [[indices[b_idx] for b_idx in batch] for _ in range(self.batch_size)]
-            for batch in batches
+            [
+                [indices[b_idx] for b_idx in batch]
+                for batch in batches[i : i + self.batch_size]
+            ]
+            for i in range(0, len(batches), self.batch_size)
         ]
 
         # statistics
@@ -192,7 +195,7 @@ def _len_est(self):
                     0.99
                     * lengths_sum_per_device
                     / self.packing_efficiency_estimate
-                    // self.batch_max_len
+                    // (self.batch_max_len * self.batch_size)
                 )
                 - 1
             ),

src/axolotl/utils/trainer.py

@@ -255,7 +255,7 @@ def calculate_total_num_steps(cfg, train_dataset, update=True):
                 train_dataset.remove_columns(["length"]),
                 batch_sampler=sampler,
             )
-            data_loader_len = len(data_loader)
+            data_loader_len = len(data_loader) // batch_size
             actual_eff = sampler.efficiency()
             LOG.debug(f"data_loader_len: {data_loader_len}", main_process_only=True)
             # FIXME: is there a bug here somewhere? the total num steps depends