est total tokens, fix field loop

src/axolotl/utils/trainer.py

@@ -155,6 +155,16 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
     if cfg.sample_packing:
         train_dataset = train_dataset.map(add_position_ids)
         eval_dataset = eval_dataset.map(add_position_ids)
+        if cfg.sample_packing_eff_est:
+            total_num_tokens = sum(len(s["input_ids"]) for s in train_dataset)
+            total_num_steps = math.ceil(
+                total_num_tokens
+                / cfg.sample_packing_eff_est
+                / 2048
+                * cfg.num_epochs
+                / cfg.batch_size
+            )
+        else:
             sampler = DistributedSampler(
                 train_dataset,
                 num_replicas=1,
@@ -176,7 +186,10 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
             LOG.info(f"data_loader_len: {data_loader_len}")
             total_num_steps = int(
                 math.ceil(
-                data_loader_len * cfg.micro_batch_size * cfg.num_epochs / cfg.batch_size
+                    data_loader_len
+                    * cfg.micro_batch_size
+                    * cfg.num_epochs
+                    / cfg.batch_size
                 )
             )
     else: