Deepspeed reward training hangs at end of training with Dataset.from_list
State: open | Author: Matt00n (NONE)

There seems to be a weird interaction between Deepspeed, the Dataset.from_list method and trl's RewardTrainer. On a multi-GPU setup (10 A100s), training always hangs at the very end of training until it times out. The training itself works fine until the end of training and running the same script with Deepspeed on a single GPU works without hangig. The issue persisted across a wide range of Deepspeed configs and training arguments. The issue went away when storing the exact same dataset as a JSON and using `dataset = load_dataset("json", ...)`. Here is my training script:

```python
import pickle
import os
import random
import warnings

import torch
from datasets import load_dataset, Dataset
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from trl import RewardConfig, RewardTrainer, ModelConfig

####################################### Reward model #################################################

# Explicitly set arguments
model_name_or_path = "Qwen/Qwen2.5-1.5B"
output_dir = "Qwen2-0.5B-Reward-LoRA"
per_device_train_batch_size = 2
num_train_epochs = 5
gradient_checkpointing = True
learning_rate = 1.0e-4
logging_steps = 25
eval_strategy = "steps"
eval_steps = 50
max_length = 2048
torch_dtype = "auto"
trust_remote_code = False

model_args = ModelConfig(
    model_name_or_path=model_name_or_path,
    model_revision=None,
    trust_remote_code=trust_remote_code,
    torch_dtype=torch_dtype,
    lora_task_type="SEQ_CLS", # Make sure task type is seq_cls
)

training_args = RewardConfig(
    output_dir=output_dir,
    per_device_train_batch_size=per_device_train_batch_size,
    num_train_epochs=num_train_epochs,
    gradient_checkpointing=gradient_checkpointing,
    learning_rate=learning_rate,
    logging_steps=logging_steps,
    eval_strategy=eval_strategy,
    eval_steps=eval_steps,
    max_length=max_length,
    gradient_checkpointing_kwargs=dict(use_reentrant=False),
    center_rewards_coefficient = 0.01,
    fp16=False,
    bf16=True,
    save_strategy="no",
    dataloader_num_workers=0,
    # deepspeed="./configs/deepspeed_config.json",
)

################
# Model & Tokenizer
################

model_kwargs = dict(
    revision=model_args.model_revision,
    use_cache=False if training_args.gradient_checkpointing else True,
    torch_dtype=model_args.torch_dtype,
)
tokenizer = AutoTokenizer.from_pretrained(
    model_args.model_name_or_path, use_fast=True
)
model = AutoModelForSequenceClassification.from_pretrained(
    model_args.model_name_or_path, num_labels=1, trust_remote_code=model_args.trust_remote_code, **model_kwargs
)
# Align padding tokens between tokenizer and model
model.config.pad_token_id = tokenizer.pad_token_id

# If post-training a base model, use ChatML as the default template
if tokenizer.chat_template is None:
    model, tokenizer = setup_chat_format(model, tokenizer)

if model_args.use_peft and model_args.lora_task_type != "SEQ_CLS":
    warnings.warn(
        "You are using a `task_type` that is different than `SEQ_CLS` for PEFT. This will lead to silent bugs"
        " Make sure to pass --lora_task_type SEQ_CLS when using this script with PEFT.",
        UserWarning,
    )

##############
# Load dataset
##############


with open('./prefs.pkl', 'rb') as fh:
    loaded_data = pickle.load(fh)

random.shuffle(loaded_data)

dataset = []
for a_wins, a, b in loaded_data:
    if a_wins == 0:
        a, b = b, a
    dataset.append({'chosen': a, 'rejected': b})


dataset = Dataset.from_list(dataset)

# Split the dataset into training and evaluation sets
train_eval_split = dataset.train_test_split(test_size=0.15, shuffle=True, seed=42)

# Access the training and evaluation datasets
train_dataset = train_eval_split['train']
eval_dataset = train_eval_split['test']


##########
# Training
##########


trainer = RewardTrainer(
    model=model,
    processing_class=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()
```

Replacing `dataset = Dataset.from_list(dataset)` with 
```python
with open('./prefs.json', 'w') as fh:
    json.dump(dataset, fh)
dataset = load_dataset("json", data_files="./prefs.json", split='train')
```
resolves the issue.

Comments: Hi ! How big is the dataset ? if you load it using `from_list`, the dataset lives in memory and has to be copied to every gpu process, which can be slow.

It's fasted if you load it from JSON files from disk, because in that case the dataset in converted to Arrow and loaded from disk using memory mapping. Memory mapping allows to quickly reload the dataset in other processes.

Maybe we can change `from_list` and other methods to always use the disk though, instead of loading in memory, WDYT ?