Dynamic batch size

[dxoigmn]

The batch size is a critical component of llmart's efficiency. Having a device (CPU/GPU) with more memory (RAM/VRAM) means being able to process more swaps (samples) in parallel.

Right now the user must pass per_device_bs=N if they want LLMart to process a batch of N samples in parallel.

HF provides find_executable_batch_size, that enables us to programmatically find a suitable batch size. Your job is to figure out where to put this call in LLMart such that it doesn't significantly slow down the basic attack. That is, accelerate launch -m llmart model=llama3-8b-instruct data=basic per_device_bs=64 should run just about as fast as ts -nfG1 python -m llmart model=llama3-8b-instruct data=basic per_device_bs=auto.

[sushma-4]

I am thinking enclosing the make_closure function in the given wrapper will suffice since the per_device_bs within the train function is really only being used by the make_closure function

[sebszyller]

@sushma-4 best way to validate it is to make a PR with the change, a performance comparison table, and a way to easily reproduce said table