pro训练时爆OOM

你好我跑PRO训练代码会报OOM，我是80G的A800，训练13B的模型，按道理应该不会爆啊
我把batch size设为1，block_size设为100，还是爆了，不知道问题出在哪？
train_hh.sh：
```
export OMP_NUM_THREADS=16
root_dir=..

#stage 23
id=$1
data_path=$2
ranking_len=$3
mkdir -p $root_dir/logs/$id/$ranking_len
# --main_process_port 29534 \
CUDA_VISIBLE_DEVICES=4,5,7 accelerate launch --num_processes 2 --config_file ds_config.yaml --main_process_port=29534 main.py \
--task hh \
--train_file_path $root_dir/data/${data_path} \
--validation_file_path $root_dir/data/hh_dev \
--validation_file_name sampled_dev.json \
--output_dir $root_dir/checkpoints/index_$id/stage_$ranking_len \
--log_path $root_dir/logs/$id/$ranking_len \
--index $id \
--seed 42 \
--temperature 1 \
--sft_weight 0.05 \
--num_train_epochs 2 \
--training_stage_num $ranking_len \
--block_size 100 \
--learning_rate 5e-6 \
--per_device_train_batch_size 1 \
--per_device_eval_batch_size 1 \
--model_name_or_path /mnt/data2/finLLM/models/tigerbot-13b-base \
--do_train \
--do_validation > $root_dir/logs/$id/$ranking_len/train_detail.log 2>&1
```
日志：
```
The following values were not passed to `accelerate launch` and had defaults used instead:
`--dynamo_backend` was set to a value of `'no'`
To avoid this warning pass in values for each of the problematic parameters or run `accelerate config`.
/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/utils/dataclasses.py:541: UserWarning: DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.
warnings.warn("DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.")
/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/utils/dataclasses.py:541: UserWarning: DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.
warnings.warn("DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.")
/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/utils/dataclasses.py:541: UserWarning: DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.
warnings.warn("DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.")

task : hh
do_train : True
do_validation : True
sft_weight : 0.05
index : exp001
seed : 42
temperature : 1.0
training_stage_num : 2
train_file_path : ../data/hh_train_len2
validation_file_path : ../data/hh_dev
validation_file_name : sampled_dev.json
model_name_or_path : /mnt/data2/finLLM/models/tigerbot-13b-base
per_device_train_batch_size : 1
per_device_eval_batch_size : 1
learning_rate : 5e-06
block_size : 100
num_train_epochs : 2
max_train_steps : None
gradient_accumulation_steps : 8
output_dir : ../checkpoints/index_exp001/stage_2
checkpointing_step : 600
log_path : ../logs/exp001/2

Loading checkpoint shards: 0%| | 0/3 [00:00
[2024-03-20 02:14:17,832] [INFO] [logging.py:75:log_dist] [Rank 0] Creating torch.bfloat16 ZeRO stage 2 optimizer
[2024-03-20 02:14:17,832] [INFO] [stage_1_and_2.py:145:__init__] Reduce bucket size 500,000,000
[2024-03-20 02:14:17,832] [INFO] [stage_1_and_2.py:146:__init__] Allgather bucket size 500,000,000
[2024-03-20 02:14:17,832] [INFO] [stage_1_and_2.py:147:__init__] CPU Offload: False
[2024-03-20 02:14:17,832] [INFO] [stage_1_and_2.py:148:__init__] Round robin gradient partitioning: False
Using /home/zhengmingjie/.cache/torch_extensions/py39_cu117 as PyTorch extensions root...
Using /home/zhengmingjie/.cache/torch_extensions/py39_cu117 as PyTorch extensions root...
Using /home/zhengmingjie/.cache/torch_extensions/py39_cu117 as PyTorch extensions root...
Emitting ninja build file /home/zhengmingjie/.cache/torch_extensions/py39_cu117/utils/build.ninja...
Building extension module utils...
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 utils...
Time to load utils op: 0.7705831527709961 seconds
Loading extension module utils...
Time to load utils op: 0.7298614978790283 seconds
Loading extension module utils...
Time to load utils op: 0.7109990119934082 seconds
Rank: 0 partition count [3] and sizes[(4435952640, False)]
Rank: 2 partition count [3] and sizes[(4435952640, False)]
Rank: 1 partition count [3] and sizes[(4435952640, False)]
[2024-03-20 02:15:34,752] [INFO] [utils.py:826:see_memory_usage] Before initializing optimizer states
[2024-03-20 02:15:34,753] [INFO] [utils.py:827:see_memory_usage] MA 41.35 GB Max_MA 49.61 GB CA 49.62 GB Max_CA 50 GB
[2024-03-20 02:15:34,754] [INFO] [utils.py:835:see_memory_usage] CPU Virtual Memory: used = 98.89 GB, percent = 9.8%
Traceback (most recent call last):
File "/mnt/data2/finLLM/DAMO-ConvAI-main/PRO/train/main.py", line 45, in
model = process_manager.train()
File "/mnt/data2/finLLM/DAMO-ConvAI-main/PRO/train/utils/process_manager.py", line 182, in train
model, optimizer, dataset_length = self.init_prepare_train(
File "/mnt/data2/finLLM/DAMO-ConvAI-main/PRO/train/utils/process_manager.py", line 165, in init_prepare_train
model, optimizer, _ = self.accelerator.prepare(
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/accelerator.py", line 1090, in prepare
result = self._prepare_deepspeed(*args)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/accelerator.py", line 1368, in _prepare_deepspeed
engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/__init__.py", line 125, in initialize
engine = DeepSpeedEngine(args=args,
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 336, in __init__
self._configure_optimizer(optimizer, model_parameters)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1292, in _configure_optimizer
self.optimizer = self._configure_zero_optimizer(basic_optimizer)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1542, in _configure_zero_optimizer
optimizer = DeepSpeedZeroOptimizer(
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 524, in __init__
self.initialize_optimizer_states()
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 649, in initialize_optimizer_states
self.optimizer.step()
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/optimizer.py", line 280, in wrapper
out = func(*args, **kwargs)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/optimizer.py", line 33, in _use_grad
ret = func(self, *args, **kwargs)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/adamw.py", line 160, in step
self._init_group(
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/adamw.py", line 114, in _init_group
state["exp_avg"] = torch.zeros_like(
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 16.53 GiB (GPU 0; 79.15 GiB total capacity; 57.88 GiB already allocated; 5.21 GiB free; 57.88 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
Traceback (most recent call last):
File "/mnt/data2/finLLM/DAMO-ConvAI-main/PRO/train/main.py", line 45, in
model = process_manager.train()
File "/mnt/data2/finLLM/DAMO-ConvAI-main/PRO/train/utils/process_manager.py", line 182, in train
model, optimizer, dataset_length = self.init_prepare_train(
File "/mnt/data2/finLLM/DAMO-ConvAI-main/PRO/train/utils/process_manager.py", line 165, in init_prepare_train
model, optimizer, _ = self.accelerator.prepare(
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/accelerator.py", line 1090, in prepare
result = self._prepare_deepspeed(*args)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/accelerator.py", line 1368, in _prepare_deepspeed
engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/__init__.py", line 125, in initialize
engine = DeepSpeedEngine(args=args,
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 336, in __init__
self._configure_optimizer(optimizer, model_parameters)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1292, in _configure_optimizer
self.optimizer = self._configure_zero_optimizer(basic_optimizer)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1542, in _configure_zero_optimizer
optimizer = DeepSpeedZeroOptimizer(
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 524, in __init__
self.initialize_optimizer_states()
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 649, in initialize_optimizer_states
self.optimizer.step()
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/optimizer.py", line 280, in wrapper
out = func(*args, **kwargs)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/optimizer.py", line 33, in _use_grad
ret = func(self, *args, **kwargs)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/adamw.py", line 160, in step
self._init_group(
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/adamw.py", line 118, in _init_group
state["exp_avg_sq"] = torch.zeros_like(
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 16.53 GiB (GPU 1; 79.15 GiB total capacity; 74.40 GiB already allocated; 4.13 GiB free; 74.41 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
Traceback (most recent call last):
File "/mnt/data2/finLLM/DAMO-ConvAI-main/PRO/train/main.py", line 45, in
model = process_manager.train()
File "/mnt/data2/finLLM/DAMO-ConvAI-main/PRO/train/utils/process_manager.py", line 182, in train
model, optimizer, dataset_length = self.init_prepare_train(
File "/mnt/data2/finLLM/DAMO-ConvAI-main/PRO/train/utils/process_manager.py", line 165, in init_prepare_train
model, optimizer, _ = self.accelerator.prepare(
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/accelerator.py", line 1090, in prepare
result = self._prepare_deepspeed(*args)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/accelerator.py", line 1368, in _prepare_deepspeed
engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/__init__.py", line 125, in initialize
engine = DeepSpeedEngine(args=args,
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 336, in __init__
self._configure_optimizer(optimizer, model_parameters)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1292, in _configure_optimizer
self.optimizer = self._configure_zero_optimizer(basic_optimizer)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1542, in _configure_zero_optimizer
optimizer = DeepSpeedZeroOptimizer(
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 524, in __init__
self.initialize_optimizer_states()
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 649, in initialize_optimizer_states
self.optimizer.step()
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/optimizer.py", line 280, in wrapper
out = func(*args, **kwargs)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/optimizer.py", line 33, in _use_grad
ret = func(self, *args, **kwargs)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/adamw.py", line 160, in step
self._init_group(
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/optim/adamw.py", line 118, in _init_group
state["exp_avg_sq"] = torch.zeros_like(
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 16.53 GiB (GPU 2; 79.15 GiB total capacity; 74.40 GiB already allocated; 4.18 GiB free; 74.41 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 92834 closing signal SIGTERM
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 92835 closing signal SIGTERM
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 92833) of binary: /mnt/data2/miniconda3/envs/pro/bin/python
Traceback (most recent call last):
File "/mnt/data2/miniconda3/envs/pro/bin/accelerate", line 8, in
sys.exit(main())
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/commands/accelerate_cli.py", line 45, in main
args.func(args)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/commands/launch.py", line 900, in launch_command
deepspeed_launcher(args)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/accelerate/commands/launch.py", line 643, in deepspeed_launcher
distrib_run.run(args)
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/distributed/run.py", line 785, in run
elastic_launch(
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/distributed/launcher/api.py", line 134, in __call__
return launch_agent(self._config, self._entrypoint, list(args))
File "/mnt/data2/miniconda3/envs/pro/lib/python3.9/site-packages/torch/distributed/launcher/api.py", line 250, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
============================================================
main.py FAILED
------------------------------------------------------------
Failures:

------------------------------------------------------------
Root Cause (first observed failure):
[0]:
time : 2024-03-20_02:15:40
host : oem
rank : 0 (local_rank: 0)
exitcode : 1 (pid: 92833)
error_file:
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================

```