🐛 Describe the bug

以下为1卡与单机4卡的对比，显存占用并无差异，开启了tensor_parallelize，求问原因

单机4卡运行ColossalAI/examples/language/gpt/gemini 目录下的run_gemini.sh，bs为48，日志显示占用前向占用显存最大10G，但是Nvidai-smi显示占用20G，脚本设置如下

`
set -x
export DISTPLAN=${DISTPLAN:-"CAI_Gemini"}

export GPUNUM=${GPUNUM:-4}
export TPDEGREE=${TPDEGREE:-4}
export PLACEMENT=${PLACEMENT:-"cpu"}
export USE_SHARD_INIT=${USE_SHARD_INIT:-False}
export BATCH_SIZE=${BATCH_SIZE:-48}
export MODEL_TYPE=${MODEL_TYPE:-"gpt2_xl"}
export TRAIN_STEP=${TRAIN_STEP:-20}

echo -e "BATCH_SIZE: ${BATCH_SIZE}"
echo -e "TRAIN_STEP: ${TRAIN_STEP}"

if [ ${USE_SHARD_INIT} = "True" ]; then
USE_SHARD_INIT="--shardinit"
else
USE_SHARD_INIT=""
fi

mkdir -p gemini_logs

torchrun --standalone --nproc_per_node=${GPUNUM} ./train_gpt2.py
--tp_degree=${TPDEGREE}
--model_type=${MODEL_TYPE}
--batch_size=${BATCH_SIZE}
--placement=${PLACEMENT}
${USE_SHARD_INIT}
--distplan=${DISTPLAN}
--train_step=${TRAIN_STEP}
--shardinit
2>&1 | tee ./gemini_logs/${MODEL_TYPE}${DISTPLAN}gpu${GPUNUM}bs${BATCH_SIZE}tp${TPDEGREE}${PLACEMENT}.log
`

运行日志如下：

[03/28/23 18:39:16] INFO colossalai - colossalai - INFO: ./train_gpt2.py:349 main
INFO colossalai - colossalai - INFO: the size of testing model size is 469.9M.
Iter (loss=X.XXX): 0%| | 0/53 [00:00<?, ?it/s][03/28/23 18:39:16] INFO colossalai - colossalai - INFO: ./train_gpt2.py:334 main
INFO colossalai - colossalai - INFO: After init optim, GPU memory usage: 12.03 MB, CPU memory usage: 4872.61 MB
[03/28/23 18:39:16] INFO colossalai - colossalai - INFO: ./train_gpt2.py:349 main
INFO colossalai - colossalai - INFO: the size of testing model size is 469.9M.
[03/28/23 18:39:16] INFO colossalai - colossalai - INFO: ./train_gpt2.py:349 main
Iter (loss=X.XXX): 0%| | 0/53 [00:00<?, ?it/s] INFO colossalai - colossalai - INFO: the size of testing model size is 469.9M.
INFO colossalai - colossalai - INFO: ./train_gpt2.py:349 main
INFO colossalai - colossalai - INFO: the size of testing model size is 469.9M.
INFO colossalai - colossalai - INFO: ./train_gpt2.py:350 main
INFO colossalai - colossalai - INFO: After init model, GPU memory usage: 12.03 MB, CPU memory usage: 4872.61 MB
Iter (loss=X.XXX): 0%| | 0/53 [00:00<?, ?it/s]use_cache=True is incompatible with gradient checkpointing. Setting use_cache=False...
use_cache=True is incompatible with gradient checkpointing. Setting use_cache=False...
use_cache=True is incompatible with gradient checkpointing. Setting use_cache=False...
use_cache=True is incompatible with gradient checkpointing. Setting use_cache=False...
Iter (loss=11.195): 0%| | 0/53 [00:22<?, ?it/s][03/28/23 18:39:38] INFO colossalai - colossalai - INFO: ./train_gpt2.py:380 train_step
INFO colossalai - colossalai - INFO: [1/20] Forward GPU memory usage: 7880.17 MB, CPU memory usage: 5169.12 MB
Iter (loss= nan): 0%| | 0/53 [00:27<?, ?it/s]
[03/28/23 18:40:37] INFO colossalai - colossalai - INFO: ./train_gpt2.py:391 train_step
INFO colossalai - colossalai - INFO: [1/20] Backward GPU memory usage: 2129.11 MB, CPU memory usage: 5171.02 MB

Nvida-smi显示如下：

单机1卡运行ColossalAI/examples/language/gpt/gemini 目录下的run_gemini.sh，除GPUNUM=1，TPDEGREE=1，其他设置和单机4卡脚本相同，运行日志如下：

total wasted percentage is 0.40%
[03/28/23 18:45:10] INFO colossalai - colossalai - INFO: ./train_gpt2.py:334 main
INFO colossalai - colossalai - INFO: After init optim, GPU memory usage: 12.03 MB, CPU memory usage: 17072.59 MB
[03/28/23 18:45:11] INFO colossalai - colossalai - INFO: ./train_gpt2.py:349 main
INFO colossalai - colossalai - INFO: the size of testing model size is 1.6B.
INFO colossalai - colossalai - INFO: ./train_gpt2.py:350 main
INFO colossalai - colossalai - INFO: After init model, GPU memory usage: 12.03 MB, CPU memory usage: 17072.59 MB
Iter (loss=X.XXX): 0%| | 0/209 [00:00<?, ?it/s]use_cache=True is incompatible with gradient checkpointing. Setting use_cache=False...
Iter (loss=11.188): 0%| | 0/209 [00:12<?, ?it/s][03/28/23 18:45:23] INFO colossalai - colossalai - INFO: ./train_gpt2.py:380 train_step
INFO colossalai - colossalai - INFO: [1/20] Forward GPU memory usage: 6996.68 MB, CPU memory usage: 17369.62 MB
[03/28/23 18:45:57] INFO colossalai - colossalai - INFO: ./train_gpt2.py:391 train_step
INFO colossalai - colossalai - INFO: [1/20] Backward GPU memory usage: 1862.32 MB, CPU memory usage: 17371.45 MB
INFO colossalai - colossalai - INFO: /opt/conda/lib/python3.7/site-packages/colossalai/nn/optimizer/zero_optimizer.py:218 step
INFO colossalai - colossalai - INFO: Found overflow. Skip step
INFO colossalai - colossalai - INFO: ./train_gpt2.py:397 train_step
INFO colossalai - colossalai - INFO: [1/20] Optimizer step GPU memory usage: 1862.32 MB, CPU memory usage: 17371.46 MB
INFO colossalai - colossalai - INFO: ./train_gpt2.py:402 train_step
INFO colossalai - colossalai - INFO: [1/20] Loss:11.188, Step time: 46.644s, TFLOPS: 13.802, FWD time: 12.461s, BWD time:
34.166s, OPTIM time: 0.017s

Nvidia-smi显示如下：

Environment

pip install --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple \
transformers==4.18.0
torchmetrics==0.8.2
torch==1.12.1+cu113
torchvision==0.13.1+cu113
torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113

colossalai==0.2.7