SigLIP memory usage

[alexisdrakopoulos]

I am currently running with the following:

python -m open_clip_train.main
--save-frequency 1
--zeroshot-frequency 1
--dataset-type webdataset
--train-data 'tar_files/dataset_{000..042}.tar'
--train-num-samples 1280000
--warmup 10000
--batch-size=192
--lr=1e-3
--wd=0
--epochs=200
--workers=4
--model ViT-B-16-SigLIP
--pretrained webli
--siglip

Images should be 224x224 but I can't use a batch size above 200ish without getting OOM on a 24gb VRAM machine. How did the original authors run 32k? that would require 3840gb of memory wouldn't it? They claimed to run it on 4 TPUv4s which is only 128gb.

They write "For ex-ample, with four TPU-v4 chips, we could fit a batch size of 4096". That's over 5 times the batch size that I can fit.

[alexisdrakopoulos]

Nevermind I needed
--grad-checkpointing