Scaling up Batch Size and GPU Usage to Accelerate Training

[yoyowang0109]

Hi, and thank you for your support so far.

After setting tokens_per_batch to 8192, the training runs smoothly. However, based on the source code, it seems that each batch is constrained to a size of 1, which also implies that training is limited to using only a single GPU — otherwise, the following assertion error is triggered:

This restriction results in a large number of batches being processed. For example, pre-training on the MIMIC-IV dataset currently takes approximately 42–45 days to complete.

Given that Appendix C of your paper mentions using “24 Intel Xeon 2.70GHz CPU cores and 8 Nvidia V100 GPUs,” I’m wondering:

1. Does the current version of the code support multi-GPU training and larger batch sizes?
2. If not directly, are there recommended changes or workarounds to enable this?

Any guidance or suggestions would be greatly appreciated.
Thank you again for your time and work on this project.

[EthanSteinberg]

@yoyowang0109 The code uses dynamic batching, so that as many patients as possible are put in the 8192 tokens per batch.

So if you increase the number of tokens, that will increase the number of patients per batch.

This version of the code doesn't support multiple GPUs. The 8 V100s were primarily used to speed up hyperprameter tuning runs and related. 42 days seems a bit much. Is that for 100 epochs? I wouldn't recommend 100 epochs. 10 should be sufficient.

In theory, getting this code to work on multiple GPUs shouldn't be too big of an issue, it would just require several minor tweaks.

I don't know if I have any particular plans to implement that soon, but feel free to do so yourself and we can leave this issue open (since it is something that should be implemented).

[yoyowang0109]

@EthanSteinberg Thank you for your elaboration. Here, I explain how I calculate the running time.

After converting the batches, I obtained 26,965 batches from the MIMIC-IV dataset. With 4 training epochs, this results in a total of 26,965 * 4 = 107,860 steps. Currently, after running the program for 4 days, it has completed approximately 10,000 steps. Assuming the current processing speed remains constant, the entire training would take around 40 days to finish.

Following your suggestion, I plan to increase the number of tokens per batch to reduce the total number of batches and speed up training. As long as memory usage remains manageable, is there a practical upper limit to the number of tokens per batch?

Here are the current hyperparameters I used:
{
"frac_test": 0.15,
"min_fraction": 1e-9,
"tokens_per_batch": 256,
"num_proc": 4,
"learning_rate": 1e-05,
"vocab_size": 65536,
"num_tasks": 2048,
"num_bins": 4,
"final_layer_size": 32,
"n_layer": 12,
"hidden_size": 768,
"per_device_train_batch_size": 1,
"per_device_eval_batch_size": 1,
"num_train_epochs": 4,
"eval_steps": 20,
"logging_steps": 20
}

According to a comment in the source code:

“Normally min_fraction should be set higher (e.g., 1e-4), but a smaller value is needed to get enough codes.”

Does the min_fraction parameter affect training time significantly? I'm open to any suggestions for tuning the above hyperparameters to improve performance.

Thank you for your time and help. It seems GPU capacity may not be the main bottleneck here, and I'm doing my best to work through the lengthy training process.

[EthanSteinberg]

@yoyowang0109 Hmm. Not sure why things would be running so slowly. I don't think min_fraction would change much, it doesn't exclude that many tokens. Maybe increase the number of processes? Might be a slow CPU job bottlenecking? I would use nvidia-smi to find where the bottleneck is.