From 599aee7ddbb610c9a9127192425b452fcd9e786e Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 5 Aug 2024 19:17:41 +0530 Subject: [PATCH] fix: fixed sharding in dataset --- training.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/training.py b/training.py index f9e787a..e8b9f31 100644 --- a/training.py +++ b/training.py @@ -300,7 +300,7 @@ def get_dataset_grain( shuffle=True, seed=seed, num_epochs=num_epochs, - shard_options=pygrain.NoSharding(), + shard_options=pygrain.ShardByJaxProcess(), ) transformations = [ @@ -1093,9 +1093,9 @@ def main(args): main(args) """ -JAX_TRACEBACK_FILTERING=off python3 training.py --dataset=laiona_coco --dataset_path='/home/mrwhite0racle/gcs_mount/arrayrecord/laion-aesthetics-12m+mscoco-2017'\ - --epochs=40 --batch_size=128 \ +python3 training.py --dataset=laiona_coco --dataset_path='/home/mrwhite0racle/gcs_mount/arrayrecord/laion-aesthetics-12m+mscoco-2017'\ + --epochs=40 --batch_size=64 \ --learning_rate=2.7e-4 --num_res_blocks=3 \ - --use_self_and_cross=False --dtype=float32 --precision=high --attention_heads=16 \ - --experiment_name='batch 128 multi-host laiona_coco' + --use_self_and_cross=False --dtype=bfloat16 --precision=high --attention_heads=16\ + --experiment_name='batch 64 v4-16 host laiona_coco'" """ \ No newline at end of file