Deepspeed reward training hangs at end of training with Dataset.from_list
Opened by Matt00n (open)

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 ?

Thanks for raising this! As lhoestq mentioned, the root cause seems to be that `Dataset.from_list()` creates an in-memory dataset, which causes issues with DeepSpeed across multiple GPUs due to the cost of copying that memory to all processes.

Using `load_dataset("json", ...)` works because Hugging Face datasets then convert the data to Apache Arrow and use **memory mapping**, which avoids this copying overhead.

Possible improvement could be to add an option like `use_disk=True` to `Dataset.from_list()` to allow users to write to Arrow + memory-map the dataset, enabling compatibility with multi-process settings like DeepSpeed, while keeping the current fast behavior by default.

Would love to hear if this direction sounds acceptable before attempting a PR.