Repository navigation
[BUG]: Cannot use pipeline and gemini at the same time #3403
Description
🐛 Describe the bug
I previously attempted to submit a similar issue on #3383, but some of my imprecise expressions may cause unnecessary misunderstandings, which could increase the cost of understanding this issue. So here is a clearer version.
Hi~ We are trying to use ColoInitContext and GeminiDDP to make it possible to use both Gemini and pipeline parallel at the same time.But when different ranks tried to transmit the forward results, the code got stuck.The following is a simple reproduction based on the official example:
train.py
import os
import torch
from titans.model.vit.vit import _create_vit_model
from tqdm import tqdm
import colossalai
from colossalai.context import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.logging import get_dist_logger
from colossalai.nn import CrossEntropyLoss
from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer
from colossalai.nn.parallel import GeminiDDP
from colossalai.pipeline.pipelinable import PipelinableContext
from colossalai.utils import is_using_pp, get_current_device
from colossalai.utils.model.colo_init_context import ColoInitContext
from colossalai.tensor import ProcessGroup
import torch.distributed as dist
class DummyDataloader():
def __init__(self, length, batch_size):
self.length = length
self.batch_size = batch_size
def generate(self):
data = torch.rand(self.batch_size, 3, 224, 224)
label = torch.randint(low=0, high=10, size=(self.batch_size,))
return data, label
def __iter__(self):
self.step = 0
return self
def __next__(self):
if self.step < self.length:
self.step += 1
return self.generate()
else:
raise StopIteration
def __len__(self):
return self.length
def main():
# launch from torch
parser = colossalai.get_default_parser()
args = parser.parse_args()
colossalai.launch_from_torch(config=args.config)
# get logger
logger = get_dist_logger()
logger.info("initialized distributed environment", ranks=[0])
if hasattr(gpc.config, 'LOG_PATH'):
if gpc.get_global_rank() == 0:
log_path = gpc.config.LOG_PATH
if not os.path.exists(log_path):
os.mkdir(log_path)
logger.log_to_file(log_path)
# create model
model_kwargs = dict(img_size=gpc.config.IMG_SIZE,
patch_size=gpc.config.PATCH_SIZE,
hidden_size=gpc.config.HIDDEN_SIZE,
depth=gpc.config.DEPTH,
num_heads=gpc.config.NUM_HEADS,
mlp_ratio=gpc.config.MLP_RATIO,
num_classes=10,
init_method='jax',
checkpoint=gpc.config.CHECKPOINT)
pipelinable = PipelinableContext()
pg = ProcessGroup(ranks = [dist.get_rank()])
with ColoInitContext(device=torch.device('cuda'), dtype=torch.half, default_pg=pg):
with pipelinable:
model = _create_vit_model(**model_kwargs)
pipelinable.to_layer_list()
pipelinable.policy = "uniform"
model = pipelinable.partition(1, gpc.pipeline_parallel_size, gpc.get_local_rank(ParallelMode.PIPELINE))
model = GeminiDDP(model, device=get_current_device(), placement_policy='cpu', pin_memory=True)
# use synthetic dataset
# we train for 10 steps and eval for 5 steps per epoch
train_dataloader = DummyDataloader(length=10, batch_size=gpc.config.BATCH_SIZE)
test_dataloader = DummyDataloader(length=5, batch_size=gpc.config.BATCH_SIZE)
# create loss function
criterion = CrossEntropyLoss(label_smoothing=0.1)
# create optimizer
optimizer = GeminiAdamOptimizer(model, lr=gpc.config.LEARNING_RATE)
# initialize
engine, train_dataloader, test_dataloader, _ = colossalai.initialize(model=model,
optimizer=optimizer,
criterion=criterion,
train_dataloader=train_dataloader,
test_dataloader=test_dataloader)
data_iter = iter(train_dataloader)
logger.info("Engine is built", ranks=[0])
for epoch in range(gpc.config.NUM_EPOCHS):
# training
engine.zero_grad()
_, _, loss = engine.execute_schedule(data_iter, return_output_label=False)
engine.step()
if dist.get_rank() == logger_rank:
logger.info(f"step: {step}, loss: {loss.item()}")
if __name__ == '__main__':
main()
config.py
from colossalai.amp import AMP_TYPE
# hyperparameters
# BATCH_SIZE is as per GPU
# global batch size = BATCH_SIZE x data parallel size
BATCH_SIZE = 4
LEARNING_RATE = 3e-3
WEIGHT_DECAY = 0.3
NUM_EPOCHS = 2
WARMUP_EPOCHS = 1
# model config
IMG_SIZE = 224
PATCH_SIZE = 16
HIDDEN_SIZE = 16
DEPTH = 5
NUM_HEADS = 4
MLP_RATIO = 2
NUM_CLASSES = 10
CHECKPOINT = False
SEQ_LENGTH = (IMG_SIZE // PATCH_SIZE)**2 + 1 # add 1 for cls token
# parallel setting
TENSOR_PARALLEL_SIZE = 1
TENSOR_PARALLEL_MODE = '1d'
parallel = dict(
pipeline=4,
tensor=dict(mode=TENSOR_PARALLEL_MODE, size=TENSOR_PARALLEL_SIZE),
)
# fp16 = dict(mode=AMP_TYPE.NAIVE)
clip_grad_norm = 1.0
# pipeline config
NUM_MICRO_BATCHES = parallel['pipeline']
Script startup command
colossalai run --nproc_per_node 4 train.py --config config.py
The following is the corresponding log for the above code:
[03/31/23 13:04:17] INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/context/paral
lel_context.py:521 set_device
[03/31/23 13:04:17] INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/context/paral
lel_context.py:521 set_device
INFO colossalai - colossalai - INFO: process rank 3 is
bound to device 3
[03/31/23 13:04:17] INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/context/paral
lel_context.py:521 set_device
INFO colossalai - colossalai - INFO: process rank 0 is
bound to device 0
INFO colossalai - colossalai - INFO: process rank 2 is
bound to device 2
[03/31/23 13:04:17] INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/context/paral
lel_context.py:521 set_device
INFO colossalai - colossalai - INFO: process rank 1 is
bound to device 1
[03/31/23 13:04:19] INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/context/paral
lel_context.py:557 set_seed
[03/31/23 13:04:19] INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/context/paral
lel_context.py:557 set_seed
[03/31/23 13:04:19] INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/context/paral
lel_context.py:557 set_seed
[03/31/23 13:04:19] INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/context/paral
lel_context.py:557 set_seed
INFO colossalai - colossalai - INFO: initialized seed on
rank 3, numpy: 1024, python random: 1024,
ParallelMode.DATA: 1024, ParallelMode.TENSOR:
4096,the default parallel seed is
ParallelMode.DATA.
INFO colossalai - colossalai - INFO: initialized seed on
rank 0, numpy: 1024, python random: 1024,
ParallelMode.DATA: 1024, ParallelMode.TENSOR:
1024,the default parallel seed is
ParallelMode.DATA.
INFO colossalai - colossalai - INFO: initialized seed on
rank 2, numpy: 1024, python random: 1024,
ParallelMode.DATA: 1024, ParallelMode.TENSOR:
3072,the default parallel seed is
ParallelMode.DATA.
INFO colossalai - colossalai - INFO: initialized seed on
rank 1, numpy: 1024, python random: 1024,
ParallelMode.DATA: 1024, ParallelMode.TENSOR:
2048,the default parallel seed is
ParallelMode.DATA.
INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/initialize.py
:116 launch
INFO colossalai - colossalai - INFO: Distributed
environment is initialized, data parallel size: 1,
pipeline parallel size: 4, tensor parallel size: 1
INFO colossalai - colossalai - INFO: train.py:53 main
INFO colossalai - colossalai - INFO: initialized
distributed environment
[03/31/23 13:04:19] INFO colossalai - ProcessGroup - INFO:
/workspace/code/ColossalAI/colossalai/tensor/proces
s_group.py:22 log_pg_init
INFO colossalai - ProcessGroup - INFO: Pytorch
ProcessGroup Init:
backend: nccl
ranks: [0]
INFO colossalai - ProcessGroup - INFO:
/workspace/code/ColossalAI/colossalai/tensor/proces
s_group.py:22 log_pg_init
INFO colossalai - ProcessGroup - INFO: Pytorch
ProcessGroup Init:
backend: nccl
ranks: [1]
INFO colossalai - ProcessGroup - INFO:
/workspace/code/ColossalAI/colossalai/tensor/proces
s_group.py:22 log_pg_init
INFO colossalai - ProcessGroup - INFO: Pytorch
ProcessGroup Init:
backend: nccl
ranks: [2]
INFO colossalai - ProcessGroup - INFO:
/workspace/code/ColossalAI/colossalai/tensor/proces
s_group.py:22 log_pg_init
INFO colossalai - ProcessGroup - INFO: Pytorch
ProcessGroup Init:
backend: nccl
ranks: [3]
INFO colossalai - ProcessGroup - INFO:
/workspace/code/ColossalAI/colossalai/tensor/proces
s_group.py:22 log_pg_init
INFO colossalai - ProcessGroup - INFO: Pytorch
ProcessGroup Init:
backend: nccl
ranks: [0, 1, 2, 3]
searching chunk configuration is completed in 0.00 s.
used number: 0.02 MB, wasted number: 0.00 MB
total wasted percentage is 0.00%
False
False
False
[extension] Compiling or loading the JIT-built fused_optim kernel during runtime now
Detected CUDA files, patching ldflags
Emitting ninja build file /root/.cache/colossalai/torch_extensions/torch1.13_cu11.7/build.ninja...
Building extension module fused_optim...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module fused_optim...
Loading extension module fused_optim...
Loading extension module fused_optim...
Loading extension module fused_optim...
[extension] Time to compile or load fused_optim op: 0.4035046100616455 seconds
[03/31/23 13:04:21] INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/initialize.py
:265 initialize
INFO colossalai - colossalai - INFO:
========== Your Config ========
{'BATCH_SIZE': 4,
'CHECKPOINT': False,
'DEPTH': 5,
'HIDDEN_SIZE': 16,
'IMG_SIZE': 224,
'LEARNING_RATE': 0.003,
'MLP_RATIO': 2,
'NUM_CHUNKS': 2,
'NUM_CLASSES': 10,
'NUM_EPOCHS': 2,
'NUM_HEADS': 4,
'NUM_MICRO_BATCHES': 4,
'PATCH_SIZE': 16,
'SEQ_LENGTH': 197,
'TENSOR_PARALLEL_MODE': '1d',
'TENSOR_PARALLEL_SIZE': 1,
'WARMUP_EPOCHS': 1,
'WEIGHT_DECAY': 0.3,
'clip_grad_norm': 1.0,
'parallel': {'pipeline': 4, 'tensor': {'mode':
'1d', 'size': 1}}}
================================
INFO colossalai - colossalai - INFO:
/workspace/code/ColossalAI/colossalai/initialize.py
:277 initialize
INFO colossalai - colossalai - INFO: cuDNN benchmark =
False, deterministic = False
WARNING colossalai - colossalai - WARNING:
/workspace/code/ColossalAI/colossalai/initialize.py
:442 initialize
WARNING colossalai - colossalai - WARNING: No PyTorch DDP
or gradient handler is set up, please make sure you
do not need to all-reduce the gradients after a
training step.
INFO colossalai - colossalai - INFO: train.py:102 main
INFO colossalai - colossalai - INFO: Engine is built
By tracing the code, we found that processes are stuck here.
How should I solve this issue? Does Colossal AI have better debugging tools?
Some of our guesses about this issue
We found a similar discussion in PyTorch.
Furthermore, we use top and find the CPU usage of two processes was very high, but the GPU memory usage is not high:
PID USER PR NI VIRT RES SHR S %CPU %MEM TIME+ COMMAND
35597 root 20 0 26.6g 2.3g 679720 R 101.3 0.5 1:53.95 python3.8
35598 root 20 0 25.3g 1.6g 565736 R 101.0 0.3 1:52.47 python3.8
Is it possible that some async communication was carried out by Gemini and Pipeline which resulted in a deadlock?
Environment
docker image based on nvidia/cuda:11.7.1-devel-ubuntu20.04.Installing PyTorch and ColossalAI in this image will give you an environment similar to ours.
The container startup command we use is:
sudo docker run -itd --ipc=host --network=host --gpus all --hostname colossalai --name colossalai --cap-add=SYS_PTRACE ${image_name} /bin/bash
CUDA 11.7.1
torch 1.13.1
colossalai 0.2.8
Reacted by binmakeswell
Activity
1SAA commentedon Apr 7, 2023
Hi @liuzeming-yuxi
Thanks for your feedback. Currently, Gemini is not compatible with pipeline parallellism. The reason is that Gemini is not compatible with gradient accumulation which is necessary in pipeline parallelism. You can try ZeRO-1 and ZeRO2 with pipeline parallelism instead.
JazyJiang commentedon Oct 22, 2026
Hi, thanks for reporting this, and sorry it went unanswered for so long. We're going through older issues and closing the ones that no longer apply to the current codebase. The code this issue refers to (
colossalai/communication/p2p.py,colossalai.core) has since been refactored or moved tocolossalai/legacy, so the problem as described no longer applies to the current main branch. We're marking it as stale and will close it in 14 days. If you still see this problem with the latest ColossalAI on PyTorch 2.x, please comment here (or open a new issue) with your environment and a minimal reproduction, and we'll take a look.JazyJiang commentedon Nov 5, 2026
Closing this issue since there has been no response in the 14 days after it was marked as stale. If you still run into this problem with the latest ColossalAI on PyTorch 2.x, please open a new issue using the bug report template, with your environment and a minimal reproduction.