From a809259e719537ac65bfb839db762ffa7afea4c1 Mon Sep 17 00:00:00 2001 From: Pedro Rodriguez <par@meta.com> Date: Fri, 24 Jan 2025 10:57:20 -0800 Subject: [PATCH] Use load_async flag to not start MP iterator (#33) Summary: Test Plan: --- bytelatent/args.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/bytelatent/args.py b/bytelatent/args.py index b9144c6..a332c89 100644 --- a/bytelatent/args.py +++ b/bytelatent/args.py @@ -150,11 +150,13 @@ class DataloaderArgs(BaseModel): enable_byte_ngrams=self.enable_byte_ngrams, ) packing_iterator = PackingIterator(sampling_iterator, packing_args=packing_args) - mp_iterator = MultiprocessIterator( - packing_iterator, n_batches_to_prefetch=self.prefetch_size - ) - - return mp_iterator + if self.load_async: + mp_iterator = MultiprocessIterator( + packing_iterator, n_batches_to_prefetch=self.prefetch_size + ) + return mp_iterator + else: + return packing_iterator class TrainArgs(BaseModel):