Multi-GPU Pretraining is slower than single GPU training?

[leannmlindsey]

Hello, I was able to pretrain a full model using 1 A100 GPU using this config file

composer main.py yamls/main/flex-bert-rope-base.yaml

This took about 4 days for 1 epoch.

When I try to run the same command using 4 A100 GPUs it seems to go into an infinite loop and doesn't progress. I also checked the GPU load with nvidia-smi and there is no GPU activity on any node.

(base) [lindseylm@cn0096 ~]$ nvidia-smi
Mon Mar 31 13:27:33 2025
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.216.03 Driver Version: 535.216.03 CUDA Version: 12.2 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA A100-SXM4-80GB On | 00000000:07:00.0 Off | 0 |
| N/A 30C P0 66W / 400W | 3894MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 1 NVIDIA A100-SXM4-80GB On | 00000000:46:00.0 Off | 0 |
| N/A 31C P0 68W / 400W | 2700MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 2 NVIDIA A100-SXM4-80GB On | 00000000:85:00.0 Off | 0 |
| N/A 32C P0 70W / 400W | 2700MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 3 NVIDIA A100-SXM4-80GB On | 00000000:C7:00.0 Off | 0 |
| N/A 32C P0 68W / 400W | 2628MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=======================================================================================|
| 0 N/A N/A 1081346 C ...eylm/conda/envs/bert24_2/bin/python 2618MiB |
| 0 N/A N/A 1081347 C ...eylm/conda/envs/bert24_2/bin/python 416MiB |
| 0 N/A N/A 1081348 C ...eylm/conda/envs/bert24_2/bin/python 416MiB |
| 0 N/A N/A 1081349 C ...eylm/conda/envs/bert24_2/bin/python 416MiB |
| 1 N/A N/A 1081347 C ...eylm/conda/envs/bert24_2/bin/python 2690MiB |
| 2 N/A N/A 1081348 C ...eylm/conda/envs/bert24_2/bin/python 2690MiB |
| 3 N/A N/A 1081349 C ...eylm/conda/envs/bert24_2/bin/python 2618MiB |
+---------------------------------------------------------------------------------------+

This is the same exact yaml file which runs to completion with no issue on 1 A100 GPU.

I also tried a config file from the pretrain_documentation branch

composer main.py yaml/pretrain/modernbert-base-pretrain.yaml

logfile.txt

This does seem to run properly with no errors and I can see that it is utilizing all of the GPUs, but it doesn't make much forward progress and the rate is extremely slow compared to the single GPU training.

Single GPU training
Train time/batch: 3
Train time/sample: 12288
Train time/batch_in_epoch: 3
Train time/sample_in_epoch: 12288
Train time/token: 6291456
Train time/token_in_epoch: 6291456
Train trainer/device_train_microbatch_size: 128
Train loss/train/total: 8.4557
Train metrics/train/LanguageCrossEntropy: 8.4557
Train metrics/train/MaskedAccuracy: 0.0003
Train time/train: 0.0071
Train time/val: 0.0000
Train time/total: 0.0071
Train lr-DecoupledAdamW/group0: 0.0000
Train lr-DecoupledAdamW/group1: 0.0000

Multi-GPU training
Train trainer/packing_efficiency: 0.9994
Train time/batch: 299
Train time/sample: 593403
Train time/batch_in_epoch: 299
Train time/sample_in_epoch: 593403
Train time/token: 117504980
Train time/token_in_epoch: 117504980
Train trainer/device_train_microbatch_size: 96
Train loss/train/total: 6.9307
Train throughput/batches_per_sec: 0.4220
Train throughput/samples_per_sec: 837.5386
Train throughput/device/batches_per_sec: 0.1055
Train throughput/device/samples_per_sec: 209.3846
Train throughput/tokens_per_sec: 165843.3564
Train throughput/device/tokens_per_sec: 41460.8391
Train time/train: 0.2375
Train time/val: 0.0000
Train time/total: 0.2375
Train lr-StableAdamW/group0: 0.0000
Train lr-StableAdamW/group1: 0.0000
Train gradient_norms/l1_norm: 2351.9883
Train gradient_norms/l2_norm: 0.7196

I would really like to be able to use multi-GPU for training, as this is why I was using your model. I am running this on Biowulf at the NIH.

[leannmlindsey]

Update: I switched to the pretraining_documentation branch and I also realized that though the config is labeled bert-base, it has 12 layers and 22 hidden layers? while the one I trained on 1 A100 using the flex-bert-base config is listed as having 12 layers and 12 hidden layers.

I am now training on 4 GPU and it looks like it is progressing so far.

[NohTow]

Hello,

though the config is labeled bert-base, it has 12 layers and 22 hidden layers? while the one I trained on 1 A100 using the flex-bert-base config is listed as having 12.

Yes, our ModernBERT-base has not the same shape has the original BERT-base model, you can see more information on this in the Model Design section of the paper where we give all the information and provide insights on how we optimized shapes. Note that this is the case for a lot of base-sized models, the base more refer to the number of parameters (which itself, could be a bit variable depending on the models, e.g refer more to <150M models).

The flex-bert-rope-base.yaml is outdated and the configurations to use are the ones of the pretraining_documentation branch, I am deeply sorry that this is not more clear and still not merged, I'll ask again for us to merge this branch.

As a side note, did you make sure flash-attn is installed? Should speed up the training by quite a margin.

[leannmlindsey]

Yes, Flash Attention 2 is installed (I am training on 4xA100).

Just to be clear, I am now able to train on 4xA100 GPU using the ModernBert configs from the pretraining branch. But I was only able to do that when I switched the full codebase to the pretraining branch. There is some problem with using the configs from that branch in the main branch, it appears to be running (ie no errors and all GPUs are in use) but it makes extremely slow progress (ie it is faster to train on 1 GPU).

I did reduce the hidden layers to 12, and that is what I have been training overnight. Now, with your explanation above, It seems that I need to move that back to 22 and hopefully it will still be able to train.

I have limits on how many GPUs I can allocate at one time, but I do think that I could get 8, or possible 16 (If I waited in the queue a long time). For ModernBert large, how many GPUs do you think you need to make reasonable progress (ie train in less than 1 week).

[NohTow]

There is some problem with using the configs from that branch in the main branch, it appears to be running (ie no errors and all GPUs are in use) but it makes extremely slow progress (ie it is faster to train on 1 GPU).

AFAIK, there is no change to the code at all in this branch, so it seems a bit odd.
As we are using gradient accumulation through the MBS parameter, maybe having less GPUs makes it a bit slow to perform step, but it should not be that slow.

I have limits on how many GPUs I can allocate at one time, but I do think that I could get 8, or possible 16 (If I waited in the queue a long time). For ModernBert large, how many GPUs do you think you need to make reasonable progress (ie train in less than 1 week).

The training times of our models are detailed in the paper. On 8 GPUs (H100), large took a bit more than 21 days for the full training, 17.5 for the stable phase only. Using more GPUs in a infiniband setup should scale reasonably well from our experiments, but in the end it also depends on the number of tokens you will be training for.

[ahxxm]

I noticied that disabling sequence packing makes both small(reduced layers, on single GPU) and base/large(multiple 4 or 8 gpus) faster, according to nvidia-smi dmon, both mem% and sm% won't have periodical drop 1%-10% util rates anymore.

But I haven't really located the root cause, might be any or a combo of these:

• sequence packing loader: the preloading looks mostly good to me, except time.sleep a little suspicious, when preloaded(packing_prefetch_factor==5) data are being consumed too fast, this could cause problem
• tokenizer parallelism disabled: can't keep up with required amount of packed data? feel less likely, previous benchmarks showed that it can tokenize several thousands of samples per second
• disk bottleneck: my training data are on regular AWS gp3, although significantly less than 1.7T tokens

In my setup, keeping other configs the same, training Large with 4 A10G is much faster, estimate hours 320h -> 220h.