From 82d0ed3cf37250b54dffa302ef8abf98a1b3ce75 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sun, 28 Apr 2024 21:55:21 -0400 Subject: [PATCH] cap default dataset wikipedia max_workers to 32 (#4345) 64 on tinybox OOM --- extra/datasets/wikipedia.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/extra/datasets/wikipedia.py b/extra/datasets/wikipedia.py index 32e3513339..df422be7a6 100644 --- a/extra/datasets/wikipedia.py +++ b/extra/datasets/wikipedia.py @@ -402,7 +402,7 @@ if __name__ == "__main__": elif sys.argv[1] == "pre-train": os.makedirs(BASEDIR / "train", exist_ok=True) if sys.argv[2] == "all": # Use all 500 parts for training generation - process_map(process_part, [part for part in range(500)], max_workers=getenv('NUM_WORKERS', os.cpu_count()), chunksize=1) + process_map(process_part, [part for part in range(500)], max_workers=getenv('NUM_WORKERS', min(os.cpu_count(), 32)), chunksize=1) else: # Use a specific part for training generation part = int(sys.argv[2]) os.makedirs(BASEDIR / "train" / str(part), exist_ok=True)