[BUG]Tensor parallelism produces divergent loss with non-divisible vocabulary sizes
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
Tensor parallelism produces significantly divergent loss values compared to single GPU training when using vocabulary sizes that cannot be evenly divided by the tensor parallel size. After 2000 training steps, the absolute error between single GPU and TP=2 reaches 47.137 (over 50% relative error), indicating a systematic numerical issue rather than expected floating-point variations.
**To Reproduce**
You can compare loss of exp_single.sh & exp_tensor_parallel.sh
exp_single.sh
```
export CUBLAS_WORKSPACE_CONFIG=:4096:8
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
python megatron-lm/pretrain_gpt.py \
--tensor-model-parallel-size 1 \
--pipeline-model-parallel-size 1 \
--num-layers 6 \
--hidden-size 128 \
--num-attention-heads 4 \
--seq-length 256 \
--max-position-embeddings 256 \
--vocab-size 50257 \
--make-vocab-size-divisible-by 1 \
--micro-batch-size 1 \
--global-batch-size 1 \
--train-samples 2000 \
--lr-decay-samples 2000 \
--lr-warmup-samples 0 \
--seed 1234 \
--lr 0.3 \
--min-lr 0.3 \
--lr-decay-style constant \
--weight-decay 0.0 \
--clip-grad 1.0 \
--optimizer adam \
--adam-beta1 0.9 \
--adam-beta2 0.95 \
--adam-eps 1e-08 \
--attention-dropout 0.0 \
--hidden-dropout 0.0 \
--attention-softmax-in-fp32 \
--init-method-std 0.5 \
--use-cpu-initialization \
--untie-embeddings-and-output-weights \
--log-interval 5 \
--eval-iters 10 \
--eval-interval 100 \
--save-interval 2000 \
--distributed-backend nccl \
--data-path ./minimal_data/test_data_text_document \
--vocab-file ./vocab/gpt2-vocab.json \
--merge-file ./vocab/gpt2-merges.txt \
--tokenizer-type GPT2BPETokenizer \
--split 949,50,1 \
--save ./checkpoints/divergence_6_single \
--tensorboard-dir ./logs/divergence_6_single/tb \
--no-load-optim \
--no-load-rng 2>&1 | tee ./logs/divergence_6_single/training.log
```
exp_tensor_parallel.sh
```
export CUBLAS_WORKSPACE_CONFIG=:4096:8
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
torchrun --nnodes=1 --nproc-per-node=2 --master-port=6004 megatron-lm/pretrain_gpt.py \
--tensor-model-parallel-size 2 \
--pipeline-model-parallel-size 1 \
--num-layers 6 \
--hidden-size 128 \
--num-attention-heads 4 \
--seq-length 256 \
--max-position-embeddings 256 \
--vocab-size 50257 \
--make-vocab-size-divisible-by 1 \
--micro-batch-size 1 \
--global-batch-size 1 \
--train-samples 2000 \
--lr-decay-samples 2000 \
--lr-warmup-samples 0 \
--seed 1234 \
--lr 0.3 \
--min-lr 0.3 \
--lr-decay-style constant \
--weight-decay 0.0 \
--clip-grad 1.0 \
--optimizer adam \
--adam-beta1 0.9 \
--adam-beta2 0.95 \
--adam-eps 1e-08 \
--attention-dropout 0.0 \
--hidden-dropout 0.0 \
--attention-softmax-in-fp32 \
--init-method-std 0.5 \
--use-cpu-initialization \
--untie-embeddings-and-output-weights \
--log-interval 5 \
--eval-iters 10 \
--eval-interval 100 \
--save-interval 2000 \
--distributed-backend nccl \
--data-path ./minimal_data/test_data_text_document \
--vocab-file ./vocab/gpt2-vocab.json \
--merge-file ./vocab/gpt2-merges.txt \
--tokenizer-type GPT2BPETokenizer \
--split 949,50,1 \
--save ./checkpoints/divergence_6_tp \
--tensorboard-dir ./logs/divergence_6_tp/tb \
--no-load-optim \
--no-load-rng 2>&1 | tee ./logs/divergence_6_tp/training.log
```
**Expected behavior**
Loss values should remain consistent between single GPU and tensor parallel execution, with only minor numerical differences due to floating-point operation ordering.
**Stack trace/logs**
None
**Environment (please complete the following information):**
- Commit ID:75e2efd10fd324819a29c5ce581fdf0f3c37b29f
- PyTorch version: 2.2.0a0+81ea7a4
- CUDA version: 12.3
- NCCL version: 2.19.3
- Container: nvcr.io/nvidia/pytorch:23.12-py3
- GPUs: 2x NVIDIA GPUs
**Proposed fix**
The issue appears to stem from vocabulary parallelism when vocab_size % tensor_parallel_size != 0. The vocabulary is split unevenly (e.g., 50257 tokens split as 25129/25128 between GPUs), causing systematic differences in loss computation. Potential fixes:
- Add proper handling for uneven vocabulary splits in the output layer
- Add a warning when vocabulary size is not divisible by TP size
Contributor guide
Assessment
This issue has not been assessed yet.