The amount of memory depends on the model, your batch size and the number of tokens in your batches. So in your case you have to either lower your batch size or have batches with a maximum number of tokens that is not too high (like 200 in your example).
It’s not easy to know in advance what will be the amount of GPU used so feel free to test different values if needed