[BUG]: GPT2-10b Error at self.cpu_shard = torch.empty(self.shard_size, dtype=self.dtype, pin_memory=self.pin_memory) #3322

🐛 Describe the bug

使用gpt2 gemini example例子在A100上测试，运行gpt-xl正常，运行gpt-10b出现错误，run_gemini.sh配置为

set -x
export DISTPLAN=${DISTPLAN:-"CAI_Gemini"}
export GPUNUM=${GPUNUM:-2}
export TPDEGREE=${TPDEGREE:-2}
export PLACEMENT=${PLACEMENT:-"cpu"}
export USE_SHARD_INIT=${USE_SHARD_INIT:-True}
export BATCH_SIZE=${BATCH_SIZE:-1}
export MODEL_TYPE=${MODEL_TYPE:-"gpt2_10b"}
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}
2>&1 | tee ./gemini_logs/${MODEL_TYPE}${DISTPLAN}gpu${GPUNUM}bs${BATCH_SIZE}tp${TPDEGREE}${PLACEMENT}.log

报错日志为

environmental variable OMP_NUM_THREADS is set to 192.
-------> key: config, value: None
-------> key: host, value: None
-------> key: port, value: None
-------> key: world_size, value: None
-------> key: rank, value: None
-------> key: local_rank, value: None
-------> key: backend, value: nccl
-------> key: distplan, value: CAI_Gemini
-------> key: tp_degree, value: 1
-------> key: placement, value: cuda
-------> key: shardinit, value: True
-------> key: batch_size, value: 1
-------> key: model_type, value: gpt2_10b
-------> key: train_step, value: 20
environmental variable OMP_NUM_THREADS is set to 192.
-------> key: config, value: None
-------> key: host, value: None
-------> key: port, value: None
-------> key: world_size, value: None
-------> key: rank, value: None
-------> key: local_rank, value: None
-------> key: backend, value: nccl
-------> key: distplan, value: CAI_Gemini
-------> key: tp_degree, value: 1
-------> key: placement, value: cuda
-------> key: shardinit, value: True
-------> key: batch_size, value: 1
-------> key: model_type, value: gpt2_10b
-------> key: train_step, value: 20
[03/29/23 17:18:44] INFO colossalai - colossalai - INFO: /usr/local/anaconda3/lib/python3.8/site-packages/colossalai/context/parallel_context.py:521 set_device
INFO colossalai - colossalai - INFO: process rank 1 is bound to device 1
[03/29/23 17:18:44] INFO colossalai - colossalai - INFO: /usr/local/anaconda3/lib/python3.8/site-packages/colossalai/context/parallel_context.py:521 set_device
INFO colossalai - colossalai - INFO: process rank 0 is bound to device 0
[03/29/23 17:18:48] INFO colossalai - colossalai - INFO: /usr/local/anaconda3/lib/python3.8/site-packages/colossalai/context/parallel_context.py:557 set_seed
[03/29/23 17:18:48] INFO colossalai - colossalai - INFO: /usr/local/anaconda3/lib/python3.8/site-packages/colossalai/context/parallel_context.py:557 set_seed
INFO colossalai - colossalai - INFO: initialized seed on rank 1, 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 0, numpy: 1024, python random: 1024, ParallelMode.DATA: 1024, ParallelMode.TENSOR: 1024,the default parallel seed is
ParallelMode.DATA.
INFO colossalai - colossalai - INFO: /usr/local/anaconda3/lib/python3.8/site-packages/colossalai/initialize.py:116 launch
INFO colossalai - colossalai - INFO: Distributed environment is initialized, data parallel size: 2, pipeline parallel size: 1, tensor parallel size: 1
INFO colossalai - colossalai - INFO: ./train_gpt2.py:224 main
INFO colossalai - colossalai - INFO: gpt2_10b, CAI_Gemini, batch size 1
args.distplan: CAI_Gemini, world_size: 2, shard_pg： ProcessGroup(ranks=[0, 1],
rank=0, dp=1, tp=2), default_dist_spec: DistSpec(dims=(-1,), num_partitions=(2,), placement=DistPlacementPattern.SHARD)
args.distplan: CAI_Gemini, world_size: 2, shard_pg： ProcessGroup(ranks=[0, 1],
rank=1, dp=1, tp=2), default_dist_spec: DistSpec(dims=(-1,), num_partitions=(2,), placement=DistPlacementPattern.SHARD)
------------>>>> GPT2LMHeadModel bert_param: 5.138485248
------------>>>> GPTLMModel bert_param: 5.241421824
------------>>>> GPT2LMHeadModel bert_param: 5.138485248
------------>>>> GPTLMModel bert_param: 5.241421824
False
Emitting ninja build file /root/.cache/colossalai/torch_extensions/torch1.12_cu11.3/build.ninja...
Building extension module cpu_adam...
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 cpu_adam...
False
False
[extension] Compiling or loading the JIT-built cpu_adam kernel during runtime now
Detected CUDA files, patching ldflags
Emitting ninja build file /root/.cache/colossalai/torch_extensions/torch1.12_cu11.3/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 cpu_adam...
[extension] Time to compile or load cpu_adam op: 0.26808643341064453 seconds
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.12_cu11.3/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...
[extension] Time to compile or load fused_optim op: 0.43032264709472656 seconds
searching chunk configuration is completed in 1.20 s.
used number: 9997.22 MB, wasted number: 62.40 MB
total wasted percentage is 0.62%
Traceback (most recent call last):
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/gemini/chunk/manager.py", line 63, in register_tensor
chunk_group[-1].append_tensor(tensor)
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/gemini/chunk/chunk.py", line 264, in append_tensor
raise ChunkFullError
colossalai.gemini.chunk.chunk.ChunkFullError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "./train_gpt2.py", line 422, in
Traceback (most recent call last):
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/gemini/chunk/manager.py", line 63, in register_tensor
main()
File "./train_gpt2.py", line 326, in main
model = zero_model_wrapper(model, zero_stage, gemini_config)
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/nn/parallel/zero_wrapper.py", line 43, in zero_model_wrapper
chunk_group[-1].append_tensor(tensor)
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/gemini/chunk/chunk.py", line 264, in append_tensor
wrapped_model = GeminiDDP(model, **gemini_config)
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/nn/parallel/gemini_parallel.py", line 63, in init
super().init(module, gemini_manager, pin_memory, force_outputs_fp32, strict_ddp_mode)
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/nn/parallel/data_parallel.py", line 242, in init
raise ChunkFullError
colossalai.gemini.chunk.chunk.ChunkFullError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "./train_gpt2.py", line 422, in
self._init_chunks(param_order=param_order,
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/nn/parallel/data_parallel.py", line 665, in _init_chunks
main()
File "./train_gpt2.py", line 326, in main
self.chunk_manager.register_tensor(tensor=p,
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/gemini/chunk/manager.py", line 71, in register_tensor
self.__close_one_chunk(chunk_group[-1])
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/gemini/chunk/manager.py", line 231, in __close_one_chunk
model = zero_model_wrapper(model, zero_stage, gemini_config)
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/nn/parallel/zero_wrapper.py", line 43, in zero_model_wrapper
chunk.close_chunk()wrapped_model = GeminiDDP(model, **gemini_config)
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/nn/parallel/gemini_parallel.py", line 63, in init

File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/gemini/chunk/chunk.py", line 302, in close_chunk
super().init(module, gemini_manager, pin_memory, force_outputs_fp32, strict_ddp_mode)
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/nn/parallel/data_parallel.py", line 242, in init
self.cpu_shard = torch.empty(self.shard_size, dtype=self.dtype, pin_memory=self.pin_memory)
RuntimeError: CUDA error: invalid argument
self._init_chunks(param_order=param_order,
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/nn/parallel/data_parallel.py", line 665, in _init_chunks
self.chunk_manager.register_tensor(tensor=p,
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/gemini/chunk/manager.py", line 71, in register_tensor
self.__close_one_chunk(chunk_group[-1])
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/gemini/chunk/manager.py", line 231, in __close_one_chunk
chunk.close_chunk()
File "/usr/local/anaconda3/lib/python3.8/site-packages/colossalai/gemini/chunk/chunk.py", line 302, in close_chunk
self.cpu_shard = torch.empty(self.shard_size, dtype=self.dtype, pin_memory=self.pin_memory)
RuntimeError: CUDA error: invalid argument
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 3757 closing signal SIGTERM
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 3756) of binary: /usr/local/anaconda3/bin/python
Traceback (most recent call last):
File "/usr/local/anaconda3/bin/torchrun", line 8, in
sys.exit(main())
File "/usr/local/anaconda3/lib/python3.8/site-packages/torch/distributed/elastic/multiprocessing/errors/init.py", line 345, in wrapper
return f(*args, **kwargs)
File "/usr/local/anaconda3/lib/python3.8/site-packages/torch/distributed/run.py", line 761, in main
run(args)
File "/usr/local/anaconda3/lib/python3.8/site-packages/torch/distributed/run.py", line 752, in run
elastic_launch(
File "/usr/local/anaconda3/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 131, in call
return launch_agent(self._config, self._entrypoint, list(args))
File "/usr/local/anaconda3/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 245, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:

Environment

cuda 11.3
pytorch 1.12