Turning off shuffling for datasets to decrease initialization time

lukebaumann · lukebaumann · commit 42109b9090ca · 2025-07-14T22:30:21.000Z
diff --git a/axlearn/experiments/text/gpt/common.py b/axlearn/experiments/text/gpt/common.py
@@ -504,7 +504,8 @@ def mixture_train_input_source(
                 config_for_function(input_tf_data.tfds_dataset).set(
                     dataset_name=component.name,
                     split=component.split,
-                    train_shuffle_buffer_size=64 * component.shuffle_buffer_size,
+                    # train_shuffle_buffer_size=64 * component.shuffle_buffer_size,
+                    train_shuffle_buffer_size=0,
                     read_config=tfds_read_config(),
                 )
             )

Original file line number	Diff line number	Diff line change
`@@ -504,7 +504,8 @@ def mixture_train_input_source(`
`504`	`504`	`config_for_function(input_tf_data.tfds_dataset).set(`
`505`	`505`	`dataset_name=component.name,`
`506`	`506`	`split=component.split,`
`507`		`- train_shuffle_buffer_size=64 * component.shuffle_buffer_size,`
	`507`	`+ # train_shuffle_buffer_size=64 * component.shuffle_buffer_size,`
	`508`	`+ train_shuffle_buffer_size=0,`
`508`	`509`	`read_config=tfds_read_config(),`
`509`	`510`	`)`
`510`	`511`	`)`