timmoon10
commented
5 days ago I see the callstack includes c10d::ProcessGroupMPI::createProcessGroupMPI, which implies the error is happening when PyTorch is initializing MPI. The most likely culprit is at
https://github.com/NVIDIA/Megatron-LM/blob/0bc3547702464501feefeb5523b7a17e591b21fa/megatron/training/initialize.py#L207. This looks fine to me though, so we'll need to debug further.
The first thing to confirm is that MPI is properly configured on your system. Try running the following script with torchrun:
import torch
torch.distributed.init_process_group(backend="nccl")
torch.distributed.new_group(backend="mpi")If that works, my suspicion falls on the Userbuffers initialization at https://github.com/NVIDIA/Megatron-LM/blob/0bc3547702464501feefeb5523b7a17e591b21fa/megatron/training/initialize.py#L209. Could there be some race condition where running initialize_ub on one process causes MPI initialization to fail on another process? One thing to try is commenting out the call to initialize_ub in Megatron-LM and seeing if it gets past that point (it'll probably error out during the first forward pass). Pinging @denera. It may also be helpful to take a look at https://github.com/NVIDIA/TransformerEngine/issues/827.
denera
XLzed
When use megatron-lm with transformer_engine, the training core dump with arguments --tp-comm-overlap.
environment
container
FROM nvcr.io/nvidia/pytorch:24.03-py3 RUN pip install --no-cache-dir tinydb flask sentencepiece git+https://github.com/NVIDIA/TransformerEngine.git@maincode version
cmd
NCCL_IB_GID_INDEX=3 NCCL_NET_GDR_LEVEL=SYS torchrun --nproc_per_node 8 --nnodes 1 --node_rank 0 --master_addr localhost --master_port 6000 pretrain_gpt.py --seq-length 4096 --num-layers 32 --hidden-size 4096 --ffn-hidden-size 11008 --num-attention-heads 32 --max-position-embeddings 4096 --disable-bias-linear --init-method-std 0.01 --attention-dropout 0.0 --hidden-dropout 0.0 --normalization RMSNorm --position-embedding-type rope --swiglu --untie-embeddings-and-output-weights --no-masked-softmax-fusion --no-position-embedding --tokenizer-type Llama2Tokenizer --tokenizer-model /workspace/train_perf/train_scripts/data_and_tokenizer/llama2/tokenizer/tokenizer.model --data-path /workspace/train_perf/train_scripts/data_and_tokenizer/llama2/preprocessed_data/CC_text_document --split 97,2,1 --micro-batch-size 1 --global-batch-size 8 --lr 1e-4 --train-iters 20 --lr-decay-iters 15 --lr-warmup-iters 5 --lr-decay-style cosine --min-lr 1.0e-5 --weight-decay 0.1 --clip-grad 1.0 --bf16 --tensor-model-parallel-size 2 --pipeline-model-parallel-size 1 --sequence-parallel --use-distributed-optimizer --use-flash-attn --use-mcore-models --overlap-grad-reduce --overlap-param-gather --log-interval 1 --save-interval 10000 --eval-interval 1000 --eval-iters 1 --log-throughput --no-load-optim --no-load-rng --tp-comm-overlaplog
`setting number of micro-batches to constant 2