Deepspeed reward training hangs at end of training with Dataset.from_list

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.