resolve comments

Author: LiuzcEECS

tests/data/test_dynamic_batching_dataset.py

@@ -16,6 +16,7 @@
         - test_dynamic_batching_dataset_no_shuffle
 """
 
+import argparse
 import os
 import subprocess
 import sys
@@ -366,8 +367,8 @@ def build_command(shuffle=True, save_by_idx=True):
         "--train.rmpad=false",
         "--train.rmpad_with_pos_ids=true",
         "--train.dyn_bsz=true",
-        "--train.dyn_bsz_in_worker_loop=false",
-        f"--train.dyn_bsz_dataset_save_by_idx={str(save_by_idx).lower()}",
+        "--dyn_bsz_in_worker_loop=false",
+        f"--save_by_idx={str(save_by_idx).lower()}",
         "--train.seed=42",
     ]
     return command
@@ -403,6 +404,12 @@ def main_distributed_test():
 
 def _run_distributed_test():
     """Internal function that runs the actual distributed test."""
+    _parser = argparse.ArgumentParser()
+    _parser.add_argument("--save_by_idx", type=lambda x: x.lower() == "true", default=True)
+    _parser.add_argument("--dyn_bsz_in_worker_loop", type=lambda x: x.lower() == "true", default=True)
+    test_args, remaining_argv = _parser.parse_known_args()
+    sys.argv = [sys.argv[0]] + remaining_argv
+
     args = parse_args(Arguments)
     world_size = int(os.environ["WORLD_SIZE"])
     rank = int(os.environ["RANK"])
@@ -452,11 +459,11 @@ def _run_distributed_test():
         train_steps=train_steps,
         rmpad=args.train.rmpad,
         dyn_bsz=args.train.dyn_bsz,
-        dyn_bsz_in_worker_loop=args.train.dyn_bsz_in_worker_loop,
+        dyn_bsz_in_worker_loop=test_args.dyn_bsz_in_worker_loop,
         bsz_warmup_ratio=args.train.bsz_warmup_ratio,
         rmpad_with_pos_ids=args.train.rmpad_with_pos_ids,
         dyn_bsz_buffer_size=READY_FOR_MICRO_BATCH_THRESHOLD,
-        dyn_bsz_dataset_save_by_idx=args.train.dyn_bsz_dataset_save_by_idx,
+        dyn_bsz_dataset_save_by_idx=test_args.save_by_idx,
         num_workers=2,
         drop_last=False,
         pin_memory=args.data.pin_memory,
@@ -504,12 +511,12 @@ def _run_distributed_test():
 
             # Print batch info for debugging
             """
-            logger.info(f"[rank{rank}] epoch:{epoch} step:{local_step} global_step:{global_step} num_micro_batches:{len(micro_batches)}")
+            logger.error(f"[rank{rank}] epoch:{epoch} step:{local_step} global_step:{global_step} num_micro_batches:{len(micro_batches)} dataset_iter: {dataloader.dataset._data_iter}")
             for micro_idx, micro_batch in enumerate(micro_batches):
                 # Extract sample indices from input_ids (each sample has all same values)
                 input_ids = micro_batch["input_ids"].squeeze(0)  # Remove batch dim
                 input_ids = set(input_ids.tolist())
-                logger.info(f"[rank{rank}] epoch:{epoch} step:{local_step} global_step:{global_step} micro_batch[{micro_idx}]: {input_ids}")
+                logger.error(f"[rank{rank}] epoch:{epoch} step:{local_step} global_step:{global_step} micro_batch[{micro_idx}]: {input_ids}")
             """
 
             if epoch > save_epoch or (epoch == save_epoch and local_step > save_step):

veomni/arguments/arguments_types.py

@@ -396,22 +396,10 @@ class TrainingArguments:
         default="worker",
         metadata={"help": "Use main process or worker process to run dynamic batch size."},
     )
-    dyn_bsz_in_worker_loop: bool = field(
-        default=True,
-        metadata={
-            "help": "Whether the dynamic batch construction is in DataLoader's worker loop or in Dataset's iterator."
-        },
-    )
     dyn_bsz_buffer_size: int = field(
         default=200,
         metadata={"help": "Buffer size for dynamic batch size."},
     )
-    dyn_bsz_dataset_save_by_idx: bool = field(
-        default=True,
-        metadata={
-            "help": "When dyn_bsz_in_worker_loop is False, it is to decide whether to save buffer by index for checkpointing in DynamicBatchingSizeDataset."
-        },
-    )
     bsz_warmup_ratio: float = field(
         default=0,
         metadata={"help": "Ratio of batch size warmup steps."},
@@ -718,13 +706,8 @@ def __post_init__(self):
         # for:
         #   - DynamicBatchingSizeDataset and StatefulDataLoader
         #   - StreamingDataset and StreamingDataLoader
-        if (self.rmpad or self.rmpad_with_pos_ids) and self.dyn_bsz:
-            if self.dyn_bsz_in_worker_loop:
-                self.dataloader_batch_size = 1
-            else:
-                self.dataloader_batch_size = self.global_batch_size // (
-                    self.micro_batch_size * self.data_parallel_size
-                )
+        if (self.rmpad or self.rmpad_with_pos_ids) and self.dyn_bsz and self.dyn_bsz_runtime == "worker":
+            self.dataloader_batch_size = 1
         else:
             self.dataloader_batch_size = self.global_batch_size // self.data_parallel_size  # = micro bsz * grad accu
 

veomni/data/data_loader.py

@@ -135,6 +135,7 @@ def build_native_dataloader(
             )
             collate_fn = UnpackDataCollator()
         else:
+            dataloader_batch_size = num_micro_batch
             dataset = DynamicBatchingSizeDataset(
                 dataset=dataset,
                 micro_batch_seq_length=token_micro_bsz,