Title: datasets freezes with streaming mode in multiple-gpu
State: open
Labels: description: Something isn't working
name: bug

## Describe the bug
Hi. I am using this dataloader, which is for processing large datasets in streaming mode mentioned in one of examples of huggingface. I am using it to read c4: https://github.com/huggingface/transformers/blob/b48ac1a094e572d6076b46a9e4ed3e0ebe978afc/examples/research_projects/codeparrot/scripts/codeparrot_training.py#L22
During using multi-gpu in accelerator in one node, the code freezes, but works for 1 GPU:   

```
10/16/2022 14:18:46 - INFO - datasets.info - Loading Dataset Infos from /home/jack/.cache/huggingface/modules/datasets_modules/datasets/c4/df532b158939272d032cc63ef19cd5b83e9b4d00c922b833e4cb18b2e9869b01
Steps:   0%|                                                                                                                                                                                                 | 0/400000 [00:00<?, ?it/s]10/16/2022 14:18:47 - INFO - torch.utils.data.dataloader - Shared seed (135290893754684706) sent to store on rank 0
```

# Code to reproduce 
please run this code with `accelerate launch code.py`
```
from accelerate import Accelerator
from accelerate.logging import get_logger
from datasets import load_dataset
from torch.utils.data.dataloader import DataLoader
import torch
from datasets import load_dataset
from transformers import AutoTokenizer
import torch
from accelerate.logging import get_logger
from torch.utils.data import IterableDataset
from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe


logger = get_logger(__name__)


class ConstantLengthDataset(IterableDataset):
    """
    Iterable dataset that returns constant length chunks of tokens from stream of text files.
        Args:
            tokenizer (Tokenizer): The processor used for proccessing the data.
            dataset (dataset.Dataset): Dataset with text files.
            infinite (bool): If True the iterator is reset after dataset reaches end else stops.
            max_seq_length (int): Length of token sequences to return.
            num_of_sequences (int): Number of token sequences to keep in buffer.
            chars_per_token (int): Number of characters per token used to estimate number of tokens in text buffer.
    """

    def __init__(
        self,
        tokenizer,
        dataset,
        infinite=False,
        max_seq_length=1024,
        num_of_sequences=1024,
        chars_per_token=3.6,
    ):
        self.tokenizer = tokenizer
        # self.concat_token_id = tokenizer.bos_token_id
        self.dataset = dataset
        self.max_seq_length = max_seq_length
        self.epoch = 0
        self.infinite = infinite
        self.current_size = 0
        self.max_buffer_size = max_seq_length * chars_per_token * num_of_sequences
        self.content_field = "text"

    def __iter__(self):
        iterator = iter(self.dataset)
        more_examples = True
        while more_examples:
            buffer, buffer_len = [], 0
            while True:
                if buffer_len >= self.max_buffer_size:
                    break
                try:
                    buffer.append(next(iterator)[self.content_field])
                    buffer_len += len(buffer[-1])
                except StopIteration:
                    if self.infinite:
                        iterator = iter(self.dataset)
                        self.epoch += 1
                        logger.info(f"Dataset epoch: {self.epoch}")
                    else:
                        more_examples = False
                        break
            tokenized_inputs = self.tokenizer(buffer, truncation=False)["input_ids"]
            all_token_ids = []
            for tokenized_input in tokenized_inputs:
                all_token_ids.extend(tokenized_input)
            for i in range(0, len(all_token_ids), self.max_seq_length):
                input_ids = all_token_ids[i : i + self.max_seq_length]
                if len(input_ids) == self.max_seq_length:
                    self.current_size += 1
                    yield torch.tensor(input_ids)

    def shuffle(self, buffer_size=1000):
        return ShufflerIterDataPipe(self, buffer_size=buffer_size)


def create_dataloaders(tokenizer, accelerator):
    ds_kwargs = {"streaming": True}
    # In distributed training, the load_dataset function gaurantees that only one process
    # can concurrently download the dataset.
    datasets = load_dataset(
        "c4",
        "en",
        cache_dir="cache_dir",
        **ds_kwargs,
    )
    train_data, valid_data = datasets["train"], datasets["validation"]
    with accelerator.main_process_first():
        train_data = train_data.shuffle(buffer_size=10000, seed=None)
        train_dataset = ConstantLengthDataset(
            tokenizer,
            train_data,
            infinite=True,
            max_seq_length=256,
        )
        valid_dataset = ConstantLengthDataset(
            tokenizer,
            valid_data,
            infinite=False,
            max_seq_length=256,
        )
        train_dataset = train_dataset.shuffle(buffer_size=10000)
    train_dataloader = DataLoader(train_dataset, batch_size=160, shuffle=True)
    eval_dataloader = DataLoader(valid_dataset, batch_size=160)
    return train_dataloader, eval_dataloader


def main():
    # Accelerator.
    logging_dir = "data_save_dir/log"
    accelerator = Accelerator(
        gradient_accumulation_steps=1,
        mixed_precision="bf16",
        log_with="tensorboard",
        logging_dir=logging_dir,
    )
    # We need to initialize the trackers we use, and also store our configuration.
    # The trackers initializes automatically on the main process.
    if accelerator.is_main_process:
        accelerator.init_trackers("test")
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

    # Load datasets and create dataloaders.
    train_dataloader, _ = create_dataloaders(tokenizer, accelerator)

    train_dataloader = accelerator.prepare(train_dataloader)
    for step, batch in enumerate(train_dataloader, start=1):
        print(step)
    accelerator.end_training()


if __name__ == "__main__":
    main()
```

## Results expected
Being able to run the code for streamining datasets with multi-gpu

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 2.5.2
- Platform: linux
- Python version: 3.9.12
- PyArrow version: 9.0.0

@lhoestq  I do not have any idea why this freezing happens, and I removed the streaming mode and this was working fine, so I know this is caused by streaming mode of the dataloader part not working well with multi-gpu setting. Since datasets are large, I hope to keep the streamining mode. I very much appreciate your help. 


Comments:
@lhoestq  I tested the script without accelerator, and I confirm this is due to datasets part as this gets similar results without accelerator.

Hi ! You said it works on 1 GPU but doesn't wortk without accelerator - what's the difference between running on 1 GPU and running without accelerator in your case ?

Hi @lhoestq 
thanks for coming back to me. Sorry for the confusion I made. I meant this works fine on 1 GPU, but on multi-gpu it is freezing. "accelerator" is not an issue as if you adapt the code without accelerator this still gets the same issue.
In order to test it. Please run "accelerate config", then use the setup for multi-gpu in one node.
After that run "accelerate launch code.py" and then you would see the freezing occurs.

Hi @lhoestq 
could you have the chance to reproduce the error by running the minimal example shared?
thanks

I think you need to do `train_dataset = train_dataset.with_format("torch")` to work with the DataLoader in a multiprocessing setup :)

The hang is probably caused by our streamign lib `fsspec` which doesn't work in multiprocessing out of the box - but we made it work with the PyTorch DataLoader when the dataset format is set to "torch"

Hi @lhoestq 
thanks for the response. I added  the line suggested right before calling `with accelerator.main_process_first():` in the code above and I confirm this also freezes. to reproduce it please run "accelerate launch code.py". I was wondering if you could have more suggestions for me? I do not have an idea how to fix this or debug this freezing. many thanks.

Maybe the `fsspec` stuff need to be clearer even before - can you try to run this function at the very beginning of your script ?
```python
import fsspec

def _set_fsspec_for_multiprocess() -> None:
    """
    Clear reference to the loop and thread.
    This is necessary otherwise HTTPFileSystem hangs in the ML training loop.
    Only required for fsspec >= 0.9.0
    See https://github.com/fsspec/gcsfs/issues/379
    """
    fsspec.asyn.iothread[0] = None
    fsspec.asyn.loop[0] = None

_set_fsspec_for_multiprocess()
```

Hi @lhoestq 
thank you. I tried it, I am getting `AttributeError: module 'fsspec' has no attribute 'asyn'`. which version of fsspect do you use?
I am using  
```fsspec                    2022.8.2                 pypi_0    pypi```
thank you.

Hi @lhoestq 
I solved `fsspec` error with this hack for now https://discuss.huggingface.co/t/attributeerror-module-fsspec-has-no-attribute-asyn/19255 but this is still freezing, I greatly appreciate if you could run this script on your side. Many thanks.

```
import fsspec

def _set_fsspec_for_multiprocess() -> None:
    """
    Clear reference to the loop and thread.
    This is necessary otherwise HTTPFileSystem hangs in the ML training loop.
    Only required for fsspec >= 0.9.0
    See https://github.com/fsspec/gcsfs/issues/379
    """
    fsspec.asyn.iothread[0] = None
    fsspec.asyn.loop[0] = None


_set_fsspec_for_multiprocess()

from accelerate import Accelerator
from accelerate.logging import get_logger
from datasets import load_dataset
from torch.utils.data.dataloader import DataLoader
import torch
from datasets import load_dataset
from transformers import AutoTokenizer
import torch
from accelerate.logging import get_logger
from torch.utils.data import IterableDataset
from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe


logger = get_logger(__name__)


class ConstantLengthDataset(IterableDataset):
    """
    Iterable dataset that returns constant length chunks of tokens from stream of text files.
        Args:
            tokenizer (Tokenizer): The processor used for proccessing the data.
            dataset (dataset.Dataset): Dataset with text files.
            infinite (bool): If True the iterator is reset after dataset reaches end else stops.
            max_seq_length (int): Length of token sequences to return.
            num_of_sequences (int): Number of token sequences to keep in buffer.
            chars_per_token (int): Number of characters per token used to estimate number of tokens in text buffer.
    """

    def __init__(
        self,
        tokenizer,
        dataset,
        infinite=False,
        max_seq_length=1024,
        num_of_sequences=1024,
        chars_per_token=3.6,
    ):
        self.tokenizer = tokenizer
        # self.concat_token_id = tokenizer.bos_token_id
        self.dataset = dataset
        self.max_seq_length = max_seq_length
        self.epoch = 0
        self.infinite = infinite
        self.current_size = 0
        self.max_buffer_size = max_seq_length * chars_per_token * num_of_sequences
        self.content_field = "text"

    def __iter__(self):
        iterator = iter(self.dataset)
        more_examples = True
        while more_examples:
            buffer, buffer_len = [], 0
            while True:
                if buffer_len >= self.max_buffer_size:
                    break
                try:
                    buffer.append(next(iterator)[self.content_field])
                    buffer_len += len(buffer[-1])
                except StopIteration:
                    if self.infinite:
                        iterator = iter(self.dataset)
                        self.epoch += 1
                        logger.info(f"Dataset epoch: {self.epoch}")
                    else:
                        more_examples = False
                        break
            tokenized_inputs = self.tokenizer(buffer, truncation=False)["input_ids"]
            all_token_ids = []
            for tokenized_input in tokenized_inputs:
                all_token_ids.extend(tokenized_input)
            for i in range(0, len(all_token_ids), self.max_seq_length):
                input_ids = all_token_ids[i : i + self.max_seq_length]
                if len(input_ids) == self.max_seq_length:
                    self.current_size += 1
                    yield torch.tensor(input_ids)

    def shuffle(self, buffer_size=1000):
        return ShufflerIterDataPipe(self, buffer_size=buffer_size)


def create_dataloaders(tokenizer, accelerator):
    ds_kwargs = {"streaming": True}
    # In distributed training, the load_dataset function gaurantees that only one process
    # can concurrently download the dataset.
    datasets = load_dataset(
        "c4",
        "en",
        cache_dir="cache_dir",
        **ds_kwargs,
    )
    train_data, valid_data = datasets["train"], datasets["validation"]
    with accelerator.main_process_first():
        train_data = train_data.shuffle(buffer_size=10000, seed=None)
        train_dataset = ConstantLengthDataset(
            tokenizer,
            train_data,
            infinite=True,
            max_seq_length=256,
        )
        valid_dataset = ConstantLengthDataset(
            tokenizer,
            valid_data,
            infinite=False,
            max_seq_length=256,
        )
        train_dataset = train_dataset.shuffle(buffer_size=10000)
    train_dataloader = DataLoader(train_dataset, batch_size=160, shuffle=True)
    eval_dataloader = DataLoader(valid_dataset, batch_size=160)
    return train_dataloader, eval_dataloader


def main():
    # Accelerator.
    logging_dir = "data_save_dir/log"
    accelerator = Accelerator(
        gradient_accumulation_steps=1,
        mixed_precision="bf16",
        log_with="tensorboard",
        logging_dir=logging_dir,
    )
    # We need to initialize the trackers we use, and also store our configuration.
    # The trackers initializes automatically on the main process.
    if accelerator.is_main_process:
        accelerator.init_trackers("test")
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

    # Load datasets and create dataloaders.
    train_dataloader, _ = create_dataloaders(tokenizer, accelerator)

    train_dataloader = accelerator.prepare(train_dataloader)
    for step, batch in enumerate(train_dataloader, start=1):
        print(step)
    accelerator.end_training()


if __name__ == "__main__":
    main()
```

Are you using `Pytorch 1.11`? Otherwise the script freezes because of the shuffling in this line: 
```
        return ShufflerIterDataPipe(self, buffer_size=buffer_size)
```
`ShufflerIterDataPipe` behavior must have changed for newer Pytorch versions. But this doesn't change whether you're using streaming or not in `datasets`, so probably not the same issue, but something to try.

> Are you using `Pytorch 1.11`? Otherwise the script freezes because of the shuffling in this line:
> 
> ```
>         return ShufflerIterDataPipe(self, buffer_size=buffer_size)
> ```
> 
> `ShufflerIterDataPipe` behavior must have changed for newer Pytorch versions. But this doesn't change whether you're using streaming or not in `datasets`, so probably not the same issue, but something to try.

I met the same issue for pytorch 1.12 and 1.13, is there a way to work around for this function for newer pytorch versions?

I met the same problem, is there any updates?

Can you share more details ? this issue is quite old now and libs have evolved, it can be worth opening a new issue with your specific setup details and code

> Can you share more details ? this issue is quite old now and libs have evolved, it can be worth opening a new issue with your specific setup details and code

In dataset codes, I use `train_dataset = load_dataset("parquet", data_files=path_train_files,  keep_in_memory=False, streaming = True, split='train')`.
In model codes, I use `torch.distributed.all_gather` to gather tensors from all GPUs.

In multi-GPU mode, If I set the `streaming = False` in dataset, training is fine with multiple epochs, If I do not gather tensors from multiple GPUs, the training is also fine with mutiple epochs. However, If I set `streaming = True` together with `torch.distributed.all_gather`, the training hang after one epoch.

After debug, I guess the hang might be caused by unevenly distributed batch data to different GPUs? But set `"dataloader_drop_last": true` did not help. Someone said we should use `DistributedSampler` when training with multiple GPUs, but how do I set the args in `TrainingArguments`? There is no such args in `TrainingArguments`.


> After debug, I guess the hang might be caused by unevenly distributed batch data to different GPUs?

This is likely the case, though DistributedSampler can't help here since it's ioly useful for Dataset objects, not IterableDataset which is what you get in streaming mode.

We discussed this exact issue in https://github.com/huggingface/datasets/issues/6623, and the apparent best fix is to use the Join context manager, which correctly finishes the processes that run out of batches without freezing the other processes

I posted an example of how to use it here : https://github.com/huggingface/datasets/issues/6623#issuecomment-3376944099, and a direction that may work with the Trainer here: https://github.com/huggingface/datasets/issues/6623#issuecomment-3612944796

Please let me know what works best for you, as we will probably want to document this better, or maybe improve the Trainer directly for this case

@lhoestq Thank you for your quick reply.

I tried your solution like this: https://github.com/huggingface/datasets/issues/6623#issuecomment-3612944796.

But I get error `[rank1]: AttributeError: 'AcceleratedOptimizer' object has no attribute 'join_hook'`, how can I resolve it?

I am using the following packages: 

- python                    3.10.12              h955ad1f_0    defaults
- transformers              4.57.3                   pypi_0    pypi
- datasets                  3.5.0                    pypi_0    pypi
- accelerate                1.12.0                   pypi_0    pypi

```python
class JoinCallback(TrainerCallback):
    def __init__(self, trainer):
        self.trainer = trainer

    def on_train_begin(self, args, state, control, **kwargs):
        self.join = Join([self.trainer.model_wrapped, self.trainer.optimizer])
        self.join.__enter__()

    def on_train_end(self, args, state, control, **kwargs):
        self.join.__exit__(None, None, None)

trainer = Trainer(...)
join_callback = JoinCallback(trainer)
trainer.add_callback(join_callback)

trainer.train()
```

you can try passing the model alone to Join():

```
self.join = Join([self.trainer.model_wrapped])
```

@lhoestq Thank you for your quick reply.
The training still hung after using `Join`. Here are the trace logs:
Can we just drop the last batch to avoid the hanging problem? Why `"dataloader_drop_last": true` did not work?
```
[rank0]:[E1216 00:16:14.707919603 ProcessGroupNCCL.cpp:685] [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=6285, OpType=ALLGATHER, NumelIn=1024, NumelOut=2048, Timeout(ms)=1800000) ran for 1800032 milliseconds before timing out.
[rank1]:[E1216 00:16:14.707943871 ProcessGroupNCCL.cpp:685] [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=6286, OpType=ALLGATHER, NumelIn=896, NumelOut=1792, Timeout(ms)=1800000) ran for 1800069 milliseconds before timing out.
[rank0]:[E1216 00:16:14.708070370 ProcessGroupNCCL.cpp:2252] [PG ID 0 PG GUID 0(default_pg) Rank 0]  failure detected by watchdog at work sequence id: 6285 PG status: last enqueued work: 6286, last completed work: 6284
[rank0]:[E1216 00:16:14.708079288 ProcessGroupNCCL.cpp:732] Stack trace of the failed collective not found, potentially because FlightRecorder is disabled. You can enable it by setting TORCH_NCCL_TRACE_BUFFER_SIZE to a non-zero value.
[rank1]:[E1216 00:16:14.708075071 ProcessGroupNCCL.cpp:2252] [PG ID 0 PG GUID 0(default_pg) Rank 1]  failure detected by watchdog at work sequence id: 6286 PG status: last enqueued work: 6286, last completed work: 6285
[rank1]:[E1216 00:16:14.708086915 ProcessGroupNCCL.cpp:732] Stack trace of the failed collective not found, potentially because FlightRecorder is disabled. You can enable it by setting TORCH_NCCL_TRACE_BUFFER_SIZE to a non-zero value.
[rank0]:[E1216 00:16:14.708110648 ProcessGroupNCCL.cpp:2584] [PG ID 0 PG GUID 0(default_pg) Rank 0] First PG on this rank to signal dumping.
[rank1]:[E1216 00:16:14.708158143 ProcessGroupNCCL.cpp:2584] [PG ID 0 PG GUID 0(default_pg) Rank 1] First PG on this rank to signal dumping.
[rank0]:[E1216 00:16:15.397595260 ProcessGroupNCCL.cpp:1870] [PG ID 0 PG GUID 0(default_pg) Rank 0] Received a dump signal due to a collective timeout from this local rank and we will try our best to dump the debug info. Last enqueued NCCL work: 6286, last completed NCCL work: 6284.This is most likely caused by incorrect usages of collectives, e.g., wrong sizes used across ranks, the order of collectives is not same for all ranks or the scheduled collective, for some reason, didn't run. Additionally, this can be caused by GIL deadlock or other reasons such as network errors or bugs in the communications library (e.g. NCCL), etc. 
[rank1]:[E1216 00:16:15.397615519 ProcessGroupNCCL.cpp:1870] [PG ID 0 PG GUID 0(default_pg) Rank 1] Received a dump signal due to a collective timeout from this local rank and we will try our best to dump the debug info. Last enqueued NCCL work: 6286, last completed NCCL work: 6285.This is most likely caused by incorrect usages of collectives, e.g., wrong sizes used across ranks, the order of collectives is not same for all ranks or the scheduled collective, for some reason, didn't run. Additionally, this can be caused by GIL deadlock or other reasons such as network errors or bugs in the communications library (e.g. NCCL), etc. 
[rank1]:[E1216 00:16:15.397882866 ProcessGroupNCCL.cpp:1589] [PG ID 0 PG GUID 0(default_pg) Rank 1] ProcessGroupNCCL preparing to dump debug info. Include stack trace: 1
[rank0]:[E1216 00:16:15.397927611 ProcessGroupNCCL.cpp:1589] [PG ID 0 PG GUID 0(default_pg) Rank 0] ProcessGroupNCCL preparing to dump debug info. Include stack trace: 1
[rank0]:[E1216 00:17:14.708222861 ProcessGroupNCCL.cpp:746] [Rank 0] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[rank0]:[E1216 00:17:14.708262370 ProcessGroupNCCL.cpp:760] [Rank 0] To avoid data inconsistency, we are taking the entire process down.
[rank1]:[E1216 00:17:14.708268735 ProcessGroupNCCL.cpp:746] [Rank 1] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[rank1]:[E1216 00:17:14.708295064 ProcessGroupNCCL.cpp:760] [Rank 1] To avoid data inconsistency, we are taking the entire process down.
[rank0]:[E1216 00:17:14.709239054 ProcessGroupNCCL.cpp:2068] [PG ID 0 PG GUID 0(default_pg) Rank 0] Process group watchdog thread terminated with exception: [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=6285, OpType=ALLGATHER, NumelIn=1024, NumelOut=2048, Timeout(ms)=1800000) ran for 1800032 milliseconds before timing out.
Exception raised from checkTimeout at /pytorch/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:688 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0x80 (0x7fa68bfdeeb0 in /usr/local/miniconda3/lib/python3.10/site-packages/torch/lib/libc10.so)
frame #1: c10d::ProcessGroupNCCL::WorkNCCL::checkTimeout(std::optional<std::chrono::duration<long, std::ratio<1l, 1000l> > >) + 0x247 (0x7fa68cea0147 in /usr/local/miniconda3/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #2: c10d::ProcessGroupNCCL::Watchdog::runLoop() + 0x1591 (0x7fa68cea3b61 in /usr/local/miniconda3/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #3: c10d::ProcessGroupNCCL::Watchdog::run() + 0xd2 (0x7fa68cea4ec2 in /usr/local/miniconda3/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #4: <unknown function> + 0xdbbf4 (0x7fa67065abf4 in /usr/local/miniconda3/bin/../lib/libstdc++.so.6)
frame #5: <unknown function> + 0x8937a (0x7fa6f7ad237a in /lib64/libc.so.6)
frame #6: clone + 0x40 (0x7fa6f7b52110 in /lib64/libc.so.6)

[rank1]:[E1216 00:17:14.709266281 ProcessGroupNCCL.cpp:2068] [PG ID 0 PG GUID 0(default_pg) Rank 1] Process group watchdog thread terminated with exception: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=6286, OpType=ALLGATHER, NumelIn=896, NumelOut=1792, Timeout(ms)=1800000) ran for 1800069 milliseconds before timing out.
Exception raised from checkTimeout at /pytorch/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:688 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0x80 (0x7f3b8eb7eeb0 in /usr/local/miniconda3/lib/python3.10/site-packages/torch/lib/libc10.so)
frame #1: c10d::ProcessGroupNCCL::WorkNCCL::checkTimeout(std::optional<std::chrono::duration<long, std::ratio<1l, 1000l> > >) + 0x247 (0x7f3b320a0147 in /usr/local/miniconda3/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #2: c10d::ProcessGroupNCCL::Watchdog::runLoop() + 0x1591 (0x7f3b320a3b61 in /usr/local/miniconda3/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #3: c10d::ProcessGroupNCCL::Watchdog::run() + 0xd2 (0x7f3b320a4ec2 in /usr/local/miniconda3/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #4: <unknown function> + 0xdbbf4 (0x7f3b15a5abf4 in /usr/local/miniconda3/bin/../lib/libstdc++.so.6)
frame #5: <unknown function> + 0x8937a (0x7f3b9cdbd37a in /lib64/libc.so.6)
frame #6: clone + 0x40 (0x7f3b9ce3d110 in /lib64/libc.so.6)
```

by the way, I also use `deepspeed  0.16.9  pypi_0  pypi`.

I'm not familiar with accelerate enough, maybe @SunMarc you had a chance to look a bit into this ?

> Can we just drop the last batch to avoid the hanging problem? Why "dataloader_drop_last": true did not work?

Because some nodes may have N extra batches, not just 1