Repository: huggingface/trl

Issue:
Multi-GPU RuntimeError: Tensors must be CUDA and dense
I'm encountering a runtime error in my code from using 8 GPUs and seeking assistance to resolve it. The error message states: "RuntimeError: Tensors must be CUDA and dense." This error happens in the evaluation step. According to this https://github.com/Lightning-AI/lightning/discussions/2529, it seems like I have to move the model/metric to the device, but I couldn't find a way to resolve this through RewardTrainer. Here is the relevant section of the traceback:
```
Traceback (most recent call last):
  File "/nobackup/jirayu/llama-hh-rlhf/reward_model.py", line 147, in <module>
  File "/nobackup/jirayu/llama-hh-rlhf/reward_model.py", line 130, in run_training
    
  File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 1645, in train
    return inner_training_loop(
  File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 2020, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 2321, in _maybe_log_save_evaluate
    metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
  File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 3053, in evaluate
    output = eval_loop(
  File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 3272, in evaluation_loop
    labels = self._nested_gather(labels)
  File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 3388, in _nested_gather
    tensors = distributed_concat(tensors)
  File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer_pt_utils.py", line 197, in distributed_concat
    dist.all_gather(output_tensors, tensor)
  File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 1451, in wrapper
    return func(*args, **kwargs)
  File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 2448, in all_gather
    work = default_pg.allgather([tensor_list], [tensor])
RuntimeError: Tensors must be CUDA and dense
```

Code:
```python
import argparse
import os

import torch
from accelerate import Accelerator
from datasets import load_dataset
from peft import LoraConfig
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, logging, set_seed
from trl import RewardTrainer


os.environ["WANDB_PROJECT"] = "llama-hh-rlhf"
os.environ["WANDB_RUN_NAME"] = "llama-7b-rm"


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", type=str, default="./llama-7b-sft")

    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument("--num_train_epochs", type=int, default=1)

    parser.add_argument("--learning_rate", type=float, default=3e-5)
    parser.add_argument("--lr_scheduler_type", type=str, default="linear")
    parser.add_argument("--num_warmup_steps", type=int, default=0)
    parser.add_argument("--weight_decay", type=float, default=0.05)

    parser.add_argument("--lora_r", type=int, default=16)
    parser.add_argument("--lora_alpha", type=int, default=32)
    parser.add_argument("--lora_dropout", type=float, default=0.1)

    parser.add_argument("--fp16", action="store_true", default=False)
    parser.add_argument("--bf16", action="store_true", default=True)
    parser.add_argument("--gradient_checkpointing", action="store_true", default=True)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--output_dir", type=str, default="./llama-7b-rm-adapter")
    parser.add_argument("--log_freq", default=1, type=int)
    parser.add_argument("--eval_freq", default=50, type=int)

    return parser.parse_args()


def preprocess_function(examples, tokenizer):
    tokenized_chosen = tokenizer(examples["chosen"], truncation=True)
    tokenized_rejected = tokenizer(examples["rejected"], truncation=True)
    return {
        "input_ids_chosen": tokenized_chosen["input_ids"],
        "attention_mask_chosen": tokenized_chosen["attention_mask"],
        "input_ids_rejected": tokenized_rejected["input_ids"],
        "attention_mask_rejected": tokenized_rejected["attention_mask"],
    }


def run_training(args):
    print("Loading model...")
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_path, num_labels=1, torch_dtype=torch.bfloat16,
        device_map={"": Accelerator().process_index}
    )
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path, use_auth_token=True, torch_dtype=torch.bfloat16,
        device_map={"": Accelerator().process_index}
    )
    tokenizer.pad_token = tokenizer.eos_token
    model.config.pad_token_id = tokenizer.eos_token_id
    model.config.use_cache = not args.gradient_checkpointing

    print("Loading dataset...")
    dataset = load_dataset("Anthropic/hh-rlhf")
    train_dataset = dataset["train"]
    eval_dataset = dataset["test"]

    num_proc = 24
    original_columns = train_dataset.column_names
    train_dataset = train_dataset.map(
        preprocess_function,
        batched=True,
        num_proc=num_proc,
        remove_columns=original_columns,
        fn_kwargs={"tokenizer": tokenizer}
    )
    eval_dataset = eval_dataset.map(
        preprocess_function,
        batched=True,
        num_proc=num_proc,
        remove_columns=original_columns,
        fn_kwargs={"tokenizer": tokenizer}
    )

    print("Setting up training...")
    peft_config = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        inference_mode=False,
        task_type="SEQ_CLS",
    )
    training_args = TrainingArguments(
        output_dir=args.output_dir,
        dataloader_drop_last=True,
        num_train_epochs=args.num_train_epochs,
        evaluation_strategy="steps",
        eval_steps=args.eval_freq,
        logging_steps=args.log_freq,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        learning_rate=args.learning_rate,
        lr_scheduler_type=args.lr_scheduler_type,
        warmup_steps=args.num_warmup_steps,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        gradient_checkpointing=args.gradient_checkpointing,
        fp16=args.fp16,
        bf16=args.bf16,
        weight_decay=args.weight_decay,
        report_to="wandb",
    )
    trainer = RewardTrainer(
        model=model,
        args=training_args,
        tokenizer=tokenizer,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        peft_config=peft_config,
    )

    print("Training...")
    trainer.train()

    print("Saving model...")
    model = trainer.model

    print(f"Saving to {args.output_dir}")
    model.save_pretrained(args.output_dir)


if __name__ == "__main__":
    args = get_args()

    set_seed(args.seed)
    os.makedirs(args.output_dir, exist_ok=True)

    logging.set_verbosity_error()

    run_training(args)
```

Running command: `accelerate launch reward_model.py`

accelerate config:
```
compute_environment: LOCAL_MACHINE
distributed_type: MULTI_GPU
downcast_bf16: 'no'
gpu_ids: all
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```



Discussion:
Hi @Top34051 
Thanks for the issue, indeed we did not tested the RM in a distributed setup. I think the issues comes from the fact that the default `compute_accuracy` returns a dict of np arrays and you are trying to all_gather thzm.

Can you try to replace the compute accuracy metrics by:
```python
def compute_accuracy(eval_pred) -> Dict[str, float]:
    predictions, labels = eval_pred
    # Here, predictions is rewards_chosen and rewards_rejected.
    # We want to see how much of the time rewards_chosen > rewards_rejected.
    predictions = np.argmax(predictions, axis=1)

    accuracy = np.array(predictions == labels, dtype=float).mean().item()
    return {"accuracy": torch.Tensor(accuracy)}
```
and pass it to the `RewardTrainer`'s init `compute_metrics=compute_accuracy` . Let me know how it goes
Hi @younesbelkada 

Thanks for the suggestion. However, replacing the compute accuracy metric still gives the CUDA must be dense error. 

Error:
```
File "/nobackup/jirayu/llama-hh-rlhf/reward.py", line 186, in <module>
  run_training(args)
File "/nobackup/jirayu/llama-hh-rlhf/reward.py", line 174, in run_training
  trainer.train()
File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 1645, in train
  return inner_training_loop(
File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 2020, in _inner_training_loop
  self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 2321, in _maybe_log_save_evaluate
  metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 3053, in evaluate
  output = eval_loop(
File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 3272, in evaluation_loop
  labels = self._nested_gather(labels)
File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer.py", line 3388, in _nested_gather
  tensors = distributed_concat(tensors)
File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/transformers/trainer_pt_utils.py", line 197, in distributed_concat
  dist.all_gather(output_tensors, tensor)
File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 1451, in wrapper
  return func(*args, **kwargs)
File "/nobackup/jirayu/miniconda3/envs/llm/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 2448, in all_gather
  work = default_pg.allgather([tensor_list], [tensor])
RuntimeError: Tensors must be CUDA and dense
```

Updated code:
```python
import argparse
import numpy as np
import os

import torch
from accelerate import Accelerator
from datasets import load_dataset
from peft import LoraConfig, TaskType
from tqdm import tqdm
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    logging,
    set_seed
)
from trl import RewardTrainer


os.environ["WANDB_PROJECT"] = "llama-hh-rlhf"
os.environ["WANDB_RUN_NAME"] = "llama-7b-rm"


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", type=str, default="./llama-7b-sft")
    parser.add_argument("--output_dir", type=str, default="./llama-7b-rm-adapter")

    parser.add_argument("--num_train_epochs", type=int, default=1)
    parser.add_argument("--per_device_train_batch_size", type=int, default=4)
    parser.add_argument("--per_device_eval_batch_size", type=int, default=4)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument("--gradient_checkpointing", action="store_true", default=True)

    parser.add_argument("--fp16", action="store_true", default=False)
    parser.add_argument("--bf16", action="store_true", default=True)
    parser.add_argument("--max_seq_length", type=int, default=512)

    parser.add_argument("--learning_rate", type=float, default=3e-5)
    parser.add_argument("--lr_scheduler_type", type=str, default="linear")
    parser.add_argument("--weight_decay", type=float, default=0.05)
    parser.add_argument("--warmup_ratio", type=float, default=0.0)

    parser.add_argument("--lora_r", type=int, default=16)
    parser.add_argument("--lora_alpha", type=int, default=32)
    parser.add_argument("--lora_dropout", type=float, default=0.1)

    parser.add_argument("--eval_steps", type=int, default=10)
    parser.add_argument("--logging_steps", type=int, default=1)
    parser.add_argument("--save_steps", type=int, default=2000)
    parser.add_argument("--seed", type=int, default=0)

    return parser.parse_args()


def get_dataset(args, tokenizer):
    dataset = load_dataset("Anthropic/hh-rlhf", split="train")
    dataset = dataset.train_test_split(test_size=0.05, seed=args.seed)
    train_dataset = dataset["train"]
    eval_dataset = dataset["test"]

    def preprocess_function(examples):
        tokenized_chosen = tokenizer(examples["chosen"], truncation=True)
        tokenized_rejected = tokenizer(examples["rejected"], truncation=True)
        return {
            "input_ids_chosen": tokenized_chosen["input_ids"],
            "attention_mask_chosen": tokenized_chosen["attention_mask"],
            "input_ids_rejected": tokenized_rejected["input_ids"],
            "attention_mask_rejected": tokenized_rejected["attention_mask"],
        }

    num_proc = 24
    original_columns = train_dataset.column_names

    train_dataset = train_dataset.map(
        preprocess_function,
        batched=True,
        num_proc=num_proc,
        remove_columns=original_columns
    )
    train_dataset = train_dataset.filter(
        lambda x: len(x["input_ids_chosen"]) <= args.max_seq_length and len(x["input_ids_rejected"]) <= args.max_seq_length
    )

    eval_dataset = eval_dataset.map(
        preprocess_function,
        batched=True,
        num_proc=num_proc,
        remove_columns=original_columns
    )
    eval_dataset = eval_dataset.filter(
        lambda x: len(x["input_ids_chosen"]) <= args.max_seq_length and len(x["input_ids_rejected"]) <= args.max_seq_length
    )

    return train_dataset, eval_dataset


def print_trainable_parameters(model):
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            trainable_params += param.numel()
    print(f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}")


def compute_accuracy(eval_pred):
    predictions, labels = eval_pred
    # Here, predictions is rewards_chosen and rewards_rejected.
    # We want to see how much of the time rewards_chosen > rewards_rejected.
    predictions = np.argmax(predictions, axis=1)

    accuracy = np.array(predictions == labels, dtype=float).mean().item()
    return {"accuracy": torch.Tensor(accuracy)}


def run_training(args):
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_path, num_labels=1, 
        load_in_8bit=False,
        device_map={"": Accelerator().process_index},
        torch_dtype=torch.bfloat16,
    )

    tokenizer = AutoTokenizer.from_pretrained(args.model_path, use_auth_token=True)
    tokenizer.pad_token = tokenizer.eos_token

    train_dataset, eval_dataset = get_dataset(args, tokenizer)

    model.config.pad_token_id = tokenizer.eos_token_id
    model.config.use_cache = not args.gradient_checkpointing

    peft_config = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        inference_mode=False,
        task_type=TaskType.SEQ_CLS,
    )

    trainer = RewardTrainer(
        model=model,
        args=TrainingArguments(
            output_dir=args.output_dir,
            num_train_epochs=args.num_train_epochs,
            per_device_train_batch_size=args.per_device_train_batch_size,
            per_device_eval_batch_size=args.per_device_eval_batch_size,
            gradient_accumulation_steps=args.gradient_accumulation_steps,
            gradient_checkpointing=args.gradient_checkpointing,
            fp16=args.fp16,
            bf16=args.bf16,
            learning_rate=args.learning_rate,
            weight_decay=args.weight_decay,
            lr_scheduler_type=args.lr_scheduler_type,
            warmup_ratio=args.warmup_ratio,
            label_names=[],
            evaluation_strategy="steps",
            eval_steps=args.eval_steps,
            logging_steps=args.logging_steps,
            save_steps=args.save_steps,
            report_to="wandb",
        ),
        tokenizer=tokenizer,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_accuracy,
        peft_config=peft_config,
    )

    print_trainable_parameters(trainer.model)

    print("Start training")
    trainer.train()

    print("Save last checkpoint of the model")
    trainer.model.save_pretrained(args.output_dir)


if __name__ == "__main__":
    args = get_args()

    set_seed(args.seed)
    os.makedirs(args.output_dir, exist_ok=True)

    run_training(args)
```
I tracked this down a bit. The tensors that are failing are the `labels` that are returned in the `def prediction_step` of `reward_trainer.py`.

Setting `labels = None` gets around the issue, but isn't really a solution.

```python
def prediction_step(
        self,
        model: Union[PreTrainedModel, nn.Module],
        inputs: Dict[str, Union[torch.Tensor, Any]],
        prediction_loss_only: bool,
        ignore_keys: Optional[List[str]] = None,
    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
        inputs = self._prepare_inputs(inputs)
        if ignore_keys is None:
            if hasattr(self.model, "config"):
                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
            else:
                ignore_keys = []

        with torch.no_grad():
            loss, logits_dict = self.compute_loss(model, inputs, return_outputs=True)

        if prediction_loss_only:
            return (loss, None, None)

        loss = loss.detach()
        logits = tuple(v for k, v in logits_dict.items() if k not in ignore_keys)
        logits = nested_detach(logits)
        # Stack accepted against rejected, mean over logits
        # and softmax to get preferences between accepted and rejected to sum to 1
        logits = torch.stack(logits).mean(dim=2).softmax(dim=0).T

        # PROBLEM IS HERE
        # labels = torch.zeros(logits.shape[0])
        labels = None

        return loss, logits, labels
```

Resolving patch:
diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py
--- a/trl/trainer/reward_trainer.py
+++ b/trl/trainer/reward_trainer.py
@@ -206,5 +206,6 @@ def prediction_step(
         logits = torch.stack(logits).mean(dim=2).softmax(dim=0).T
 
         labels = torch.zeros(logits.shape[0])
+        labels = self._prepare_inputs(labels)
 
         return loss, logits, labels


Test patch: