"Trying to create tensor with negative dimension -686245457367107674" when resuming from checkpoint with deepspeed
Author: jonathanasdf
Pull request: no
Labels: name: solved

### System Info

pytorch nightly, latest version of everything

### Who can help?

@pacman100 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```
Attempting to resume from /scratch/tmpk3ri55kw/checkpoint-400
[2023-12-28 09:08:50,967] [INFO] [torch_checkpoint_engine.py:27:load] [Torch] Loading checkpoint from /scratch/tmpk3ri55kw/checkpoint-400/global_step400/zero_pp_rank_0_mp_rank_00_model_states.pt...
[2023-12-28 09:08:50,985] [INFO] [torch_checkpoint_engine.py:29:load] [Torch] Loaded checkpoint from /scratch/tmpk3ri55kw/checkpoint-400/global_step400/zero_pp_rank_0_mp_rank_00_model_states.pt.
[2023-12-28 09:08:50,985] [INFO] [torch_checkpoint_engine.py:27:load] [Torch] Loading checkpoint from /scratch/tmpk3ri55kw/checkpoint-400/global_step400/zero_pp_rank_0_mp_rank_00_model_states.pt...
[2023-12-28 09:08:51,002] [INFO] [torch_checkpoint_engine.py:29:load] [Torch] Loaded checkpoint from /scratch/tmpk3ri55kw/checkpoint-400/global_step400/zero_pp_rank_0_mp_rank_00_model_states.pt.
[2023-12-28 09:08:51,036] [INFO] [torch_checkpoint_engine.py:27:load] [Torch] Loading checkpoint from /scratch/tmpk3ri55kw/checkpoint-400/global_step400/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt...
[2023-12-28 09:08:51,117] [INFO] [torch_checkpoint_engine.py:29:load] [Torch] Loaded checkpoint from /scratch/tmpk3ri55kw/checkpoint-400/global_step400/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt.
[2023-12-28 09:08:51,117] [INFO] [engine.py:2988:_get_all_zero_checkpoint_state_dicts] successfully read 2 ZeRO state_dicts for rank 0
[2023-12-28 09:08:51,159] [INFO] [engine.py:2920:_load_zero_checkpoint] loading 2 zero partition checkpoints for rank 0
Traceback (most recent call last):
  File "/finetune.py", line 630, in <module>
    main()
  File "/finetune.py", line 523, in main
    trainer.train(resume_from_checkpoint=last_checkpoint)
  File "/opt/venv/lib/python3.11/site-packages/transformers/trainer.py", line 1543, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.11/site-packages/transformers/trainer.py", line 1827, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/opt/venv/lib/python3.11/site-packages/accelerate/data_loader.py", line 639, in __iter__
    next_batch, next_batch_info = self._fetch_batches(main_iterator)
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.11/site-packages/accelerate/data_loader.py", line 612, in _fetch_batches
    broadcast_object_list(batch_info)
  File "/opt/venv/lib/python3.11/site-packages/accelerate/utils/operations.py", line 490, in broadcast_object_list
    torch.distributed.broadcast_object_list(object_list, src=from_process)
  File "/opt/venv/lib/python3.11/site-packages/torch/distributed/c10d_logger.py", line 75, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.11/site-packages/torch/distributed/distributed_c10d.py", line 2526, in broadcast_object_list
    object_tensor = torch.empty(  # type: ignore[call-overload]
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate more than 1EB memory.
```

### Expected behavior

This workflow used to work before at some point. Not sure what combinations of things it worked/didn't work for. Training from scratch rather than resuming works fine. It's just resuming from checkpoint that causes this very strange crash.

Comments:
When I try without LoRA I get a different (maybe more useful?) error message

```
  File "/opt/venv/lib/python3.11/site-packages/transformers/trainer.py", line 1827, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/opt/venv/lib/python3.11/site-packages/accelerate/data_loader.py", line 646, in __iter__
    next_batch, next_batch_info = self._fetch_batches(main_iterator)
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.11/site-packages/accelerate/data_loader.py", line 619, in _fetch_batches
    broadcast_object_list(batch_info)
  File "/opt/venv/lib/python3.11/site-packages/accelerate/utils/operations.py", line 491, in broadcast_object_list
    torch.distributed.broadcast_object_list(object_list, src=from_process)
  File "/opt/venv/lib/python3.11/site-packages/torch/distributed/c10d_logger.py", line 75, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.11/site-packages/torch/distributed/distributed_c10d.py", line 2526, in broadcast_object_list
    object_tensor = torch.empty(  # type: ignore[call-overload]
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Trying to create tensor with negative dimension -686245457367107674: [-686245457367107674]
```

@pacman100 

I have a repro with 2xa100

```python
import logging
import os
import shutil
import tempfile
import datasets
import torch
import torch.distributed
import transformers

output = 'test'
os.makedirs(output, exist_ok=True)
local_output = tempfile.mkdtemp()
training_args = transformers.TrainingArguments(
    local_output, 
    per_device_train_batch_size=1, 
    save_steps=2, 
    max_steps=3, 
    deepspeed={'zero_optimization': {'stage': 3}, 'train_micro_batch_size_per_gpu': 'auto'}, 
    report_to=[])
if training_args.local_rank <= 0:
  logging.basicConfig(level=logging.INFO, format='%(levelname)-8s | %(message)s')
  transformers.logging.set_verbosity(logging.INFO)
  # Copy from remote (eg. cloud storage) to local. Here we use two local dirs as an example.
  shutil.copytree(output, local_output, dirs_exist_ok=True)
if torch.distributed.is_initialized():
  torch.distributed.barrier()

last_checkpoint = transformers.trainer_utils.get_last_checkpoint(local_output)
dataset = {'input_ids': torch.tensor([0]), 'labels': torch.tensor([0]), 'attention_mask': torch.tensor([1])}
train_dataset = datasets.IterableDataset.from_generator(lambda: (dataset for _ in range(100)))
model = transformers.AutoModelForCausalLM.from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v0.6')
trainer = transformers.Trainer(args=training_args, model=model, train_dataset=train_dataset)
trainer.train(resume_from_checkpoint=last_checkpoint)
if training_args.local_rank <= 0:
  shutil.copytree(local_output, output, dirs_exist_ok=True)
```

requirements.txt:
```
--pre --extra-index-url https://download.pytorch.org/whl/nightly/cu118
accelerate==0.25.0
datasets==2.15.0
deepspeed==0.12.6
torch==2.3.0.dev20231229+cu118
transformers==4.36.2
```

In a new folder, save as test.py. Run with `deepspeed test.py`. First run will initialize and run successfully and save checkpoint. Second run will try to load the checkpoint and crash.

This only crashes with deepspeed with multiple GPUs. Running with `python3 test.py` or `deepspeed --num_gpus=1 test.py` works fine. Additionally for some reason the copy from `output` to `local_output` is somehow important, if you remove `local_output` and all the `copytree` calls and just use `output` everywhere the bug also does not happen.

Additionally, if you use a Dataset instead of IterableDataset then it doesn't crash but just hangs forever while loading.
```
dataset = {'input_ids': torch.tensor([[0]]), 'labels': torch.tensor([[0]]), 'attention_mask': torch.tensor([[1]])}
train_dataset = datasets.Dataset.from_dict(dataset)
```

Bump

@pacman100 this is still failing with newest versions

updated requirements.txt
```
--pre --extra-index-url https://download.pytorch.org/whl/nightly/cu121
accelerate==0.26.1
datasets==2.16.1
deepspeed==0.13.0
torch==2.3.0.dev20240122+cu121
transformers==4.37.0
```

Updated title to be more relevant and trying another ping

cc @SunMarc if you can have a look as well 😉 

Hello, Thank you @jonathanasdf for the minimal reproducer. I can replicate it but this seems unrelated with DeepSpeed, the error stems from resuming dataloader cc @muellerzr . Also, when not using DeepSpeed I notice a complete hang:
```
Traceback (most recent call last):
  File "/raid/sourab/temp/issues/transformers/issue_28270.py", line 33, in <module>
    trainer.train(resume_from_checkpoint=last_checkpoint)
  File "/raid/sourab/transformers/src/transformers/trainer.py", line 1561, in train
    return inner_training_loop(
  File "/raid/sourab/transformers/src/transformers/trainer.py", line 1862, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/raid/sourab/accelerate/src/accelerate/data_loader.py", line 654, in __iter__
    next_batch, next_batch_info = self._fetch_batches(main_iterator)
  File "/raid/sourab/accelerate/src/accelerate/data_loader.py", line 627, in _fetch_batches
    broadcast_object_list(batch_info)
  File "/raid/sourab/accelerate/src/accelerate/utils/operations.py", line 506, in broadcast_object_list
    torch.distributed.broadcast_object_list(object_list, src=from_process)
  File "/raid/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/torch/distributed/c10d_logger.py", line 72, in wrapper
    return func(*args, **kwargs)
  File "/raid/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 2422, in broadcast_object_list
    object_tensor = torch.empty(  # type: ignore[call-overload]
RuntimeError: Trying to create tensor with negative dimension -868913248929906688: [-868913248929906688]
[2024-02-06 09:00:56,748] [INFO] [launch.py:315:sigkill_handler] Killing subprocess 4105240
[2024-02-06 09:00:56,983] [INFO] [launch.py:315:sigkill_handler] Killing subprocess 4105241
[2024-02-06 09:00:56,984] [ERROR] [launch.py:321:sigkill_handler] ['/raid/sourab/miniconda3/envs/hf/bin/python', '-u', 'issue_28270.py', '--local_rank=1'] exits with return code = 1
```

Also, when directly using the output directory without the copying logic between temp dir and output dir, everything works as expected.

Hi @pacman100, thank you for looking into it.

I don't see how the code you added could be it. The if statement you added is in a block that is already under another `if training_args.local_rank <= 0:`

oh, yes, didn't notice it was already in the `if training_args.local_rank <= 0:`. I ran the below code and it worked fine:

```
import logging
import os
import shutil
import tempfile
import datasets
import torch
import torch.distributed
import transformers

output = 'test'
os.makedirs(output, exist_ok=True)
local_output = "test-2"
os.makedirs(local_output, exist_ok=True)
training_args = transformers.TrainingArguments(
    local_output, 
    per_device_train_batch_size=1, 
    save_steps=2, 
    max_steps=3, 
    deepspeed={'zero_optimization': {'stage': 3}, 'train_micro_batch_size_per_gpu': 'auto'}, 
    report_to=[])
if training_args.local_rank <= 0:
  logging.basicConfig(level=logging.INFO, format='%(levelname)-8s | %(message)s')
  transformers.logging.set_verbosity(logging.INFO)
  # Copy from remote (eg. cloud storage) to local. Here we use two local dirs as an example.
  shutil.copytree(output, local_output, dirs_exist_ok=True)
if torch.distributed.is_initialized():
  torch.distributed.barrier()

last_checkpoint = transformers.trainer_utils.get_last_checkpoint(local_output)
dataset = {'input_ids': torch.tensor([0]), 'labels': torch.tensor([0]), 'attention_mask': torch.tensor([1])}
train_dataset = datasets.IterableDataset.from_generator(lambda: (dataset for _ in range(100)))
model = transformers.AutoModelForCausalLM.from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v0.6')
trainer = transformers.Trainer(args=training_args, model=model, train_dataset=train_dataset)
trainer.train(resume_from_checkpoint=last_checkpoint)
if training_args.local_rank <= 0:
  shutil.copytree(local_output, output, dirs_exist_ok=True)
```

So, it seems to be the temp directory related issue? Because if I change the code to 

```diff
import logging
import os
import shutil
import tempfile
import datasets
import torch
import torch.distributed
import transformers

output = 'test'
os.makedirs(output, exist_ok=True)
+ # local_output = "test-2"
+ # os.makedirs(local_output, exist_ok=True)
+ local_output = tempfile.mkdtemp()
+ print(local_output)
training_args = transformers.TrainingArguments(
    local_output, 
    per_device_train_batch_size=1, 
    save_steps=2, 
    max_steps=3, 
    deepspeed={'zero_optimization': {'stage': 3}, 'train_micro_batch_size_per_gpu': 'auto'}, 
    report_to=[])
if training_args.local_rank <= 0:
  logging.basicConfig(level=logging.INFO, format='%(levelname)-8s | %(message)s')
  transformers.logging.set_verbosity(logging.INFO)
  # Copy from remote (eg. cloud storage) to local. Here we use two local dirs as an example.
  shutil.copytree(output, local_output, dirs_exist_ok=True)
if torch.distributed.is_initialized():
  torch.distributed.barrier()

last_checkpoint = transformers.trainer_utils.get_last_checkpoint(local_output)
dataset = {'input_ids': torch.tensor([0]), 'labels': torch.tensor([0]), 'attention_mask': torch.tensor([1])}
train_dataset = datasets.IterableDataset.from_generator(lambda: (dataset for _ in range(100)))
model = transformers.AutoModelForCausalLM.from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v0.6')
trainer = transformers.Trainer(args=training_args, model=model, train_dataset=train_dataset)
trainer.train(resume_from_checkpoint=last_checkpoint)
if training_args.local_rank <= 0:
  shutil.copytree(local_output, output, dirs_exist_ok=True)
```

I see 2 different directories being created by each process instead of all processes using the same directory:
```
/tmp/tmpvmspza_x
/tmp/tmpto8s4e6w
[2024-02-06 10:18:40,312] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)
```

Ah! that makes sense. The part in my code that makes the temp directories was last changed in Dec 2023, and back then it worked fine. So I guess there might be some change in the deepspeed launcher at some point that caused each process to create a different temp directory name...

Or maybe at some point process 0 was in charge of orchestrating the checkpoint loading and the checkpoint dir from process 0 was broadcast to all processes and now each process tries to load its own shard and so sees the different folder.

From a user code standpoint, because using torchrun/deepspeed causes the script to be executed multiple times, I'm not sure how I could define a shared temporary folder name to be used by all processes (other than creating a launch.sh script that creates the temp folder and passes it as an input arg which is kinda ugly). I'm not sure if anyone has thoughts here on how to do it. Probably if trainer.train() is called with different values for resume_from_checkpoint, there should be an error? Maybe we should be calling it with `resume_from_checkpoint=last_checkpoint if local_rank == 0 else None`?

This issue happens in resuming dataloader part as mentioned  https://github.com/huggingface/transformers/issues/28270#issuecomment-1928988185. Conventionally, users need to provide the same output directory to the Trainer.

I agree that was the bug in my code and happy to close this issue. Just wanted to see what we can do here to improve the user experience (eg. raise error when output directories are not the same, provide a function that can generate the same random string across processes, etc)

The specific error message with the dataloader is a red herring. The real issue is rank 0 getting `resume_from_checkpoint=something` while rank 1 getting `resume_from_checkpoint=None` means rank 0 is blocked on checkpoint loading while rank 1 attempts to start training. 

Adding a `dist.barrier()` here: https://github.com/huggingface/transformers/blob/6529a5b5c13210b41bcd87c555c72696cd7083a5/src/transformers/trainer.py#L1740 results in more expected behavior which is that it just hangs

Anyways it's pretty clear what's going on so I'll close this now. Thank you for helping figure it out.