"ValueError: Predictions and/or references don't match the expected format." error
Author: ashokmon-aws
Pull request: no
Labels: name: bug

### System Info

Ran on trn1.2xlarge instance with Ubuntu 22

transformers version: 4.44.0
Platform: Linux/UNIX
Python version: 3.10.12
Accelerate version: 1.7.0
PyTorch version (GPU?): 2.7.0.2.8.6734+ac864f72

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. Install dependencies:
```
env TOKENIZERS_PARALLELISM=True
pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
pip install -U "protobuf<4" "transformers==4.53.*" "datasets<=3.6.0" "accelerate==1.7.*" scikit-learn evaluate
git clone https://github.com/huggingface/transformers --branch v4.52.3
```

2. Set the parameters:
```
model_name = "camembert-base"
env_var_options = "XLA_USE_BF16=1 NEURON_CC_FLAGS=\"--model-type=transformer\""
num_workers = 2
task_name = "sst2"
batch_size = 8
max_seq_length = 128
learning_rate = 2e-05
dataset_name =  "glue"
num_train_epochs = 5
model_base_name = model_name
```

3. Compile model with neuron_parallel_compile along with `--eval_do_concat_batches=False` 
```
import subprocess
print("Compile model")
COMPILE_CMD = f"""{env_var_options} neuron_parallel_compile \
torchrun --nproc_per_node={num_workers} \
transformers/examples/pytorch/text-classification/run_glue.py \
--model_name_or_path {model_name} \
--task_name {task_name} \
--dataset_name {dataset_name} \
--do_train \
--max_seq_length {max_seq_length} \
--per_device_train_batch_size {batch_size} \
--learning_rate {learning_rate} \
--max_train_samples 128 \
--overwrite_output_dir \
--eval_do_concat_batches=False \
--output_dir {model_base_name}-{task_name}-{batch_size}bs"""

print(f'Running command: \n{COMPILE_CMD}')
if subprocess.check_call(COMPILE_CMD,shell=True):
   print("There was an error with the compilation command")
else:
   print("Compilation Success!!!")
```

4. Fine tune the model
```
print("Train model")
RUN_CMD = f"""{env_var_options} torchrun --nproc_per_node={num_workers} \
transformers/examples/pytorch/text-classification/run_glue.py \
--model_name_or_path {model_name} \
--task_name {task_name} \
--do_train \
--do_eval \
--max_seq_length {max_seq_length} \
--dataset_name {dataset_name} \
--per_device_train_batch_size {batch_size} \
--learning_rate {learning_rate} \
--num_train_epochs {num_train_epochs} \
--overwrite_output_dir \
--eval_do_concat_batches=False \
--output_dir {model_base_name}-{task_name}-{num_workers}w-{batch_size}bs"""

print(f'Running command: \n{RUN_CMD}')
if subprocess.check_call(RUN_CMD,shell=True):
   print("There was an error with the fine-tune command")
else:
   print("Fine-tune Successful!!!")
```

Error:
```
Traceback (most recent call last):
[rank1]:   File "/home/ubuntu/aws-neuron-samples/torch-neuronx/training/hf_text_classification/transformers/examples/pytorch/text-classification/run_glue.py", line 626, in <module>
[rank1]:     main()
[rank1]:   File "/home/ubuntu/aws-neuron-samples/torch-neuronx/training/hf_text_classification/transformers/examples/pytorch/text-classification/run_glue.py", line 564, in main
[rank1]:     metrics = trainer.evaluate(eval_dataset=eval_dataset)
[rank1]:   File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/transformers/trainer.py", line 4199, in evaluate
[rank1]:     output = eval_loop(
[rank1]:   File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/transformers/trainer.py", line 4489, in evaluation_loop
[rank1]:     metrics = self.compute_metrics(
[rank1]:   File "/home/ubuntu/aws-neuron-samples/torch-neuronx/training/hf_text_classification/transformers/examples/pytorch/text-classification/run_glue.py", line 498, in compute_metrics
[rank1]:     result = metric.compute(predictions=preds, references=p.label_ids)
[rank1]:   File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/evaluate/module.py", line 455, in compute
[rank1]:     self.add_batch(**inputs)
[rank1]:   File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/evaluate/module.py", line 546, in add_batch
[rank1]:     raise ValueError(error_msg) from None
[rank1]: ValueError: Predictions and/or references don't match the expected format.
[rank1]: Expected format: {'predictions': Value(dtype='int64', id=None), 'references': Value(dtype='int64', id=None)},
[rank1]: Input predictions: [[ 1  0]
[rank1]:  [ 3  7]
[rank1]:  [ 2  0]
[rank1]:  [ 6  3]
[rank1]:  [10  3]
[rank1]:  [ 6 14]
[rank1]:  [14  7]
[rank1]:  [ 0  5]
[rank1]:  [ 2  1]
[rank1]:  [11  7]
[rank1]:  [ 0 10]
[rank1]:  [ 6 15]
[rank1]:  [11 14]
[rank1]:  [13 14]
[rank1]:  [ 2  7]
[rank1]:  [ 2 10]
[rank1]:  [ 2  7]
[rank1]:  [ 0  3]
[rank1]:  [ 8  2]
[rank1]:  [15  0]
[rank1]:  [ 9  0]
[rank1]:  [ 1  0]
[rank1]:  [ 9  1]
[rank1]:  [12  1]
[rank1]:  [ 8 12]
[rank1]:  [ 8 15]
[rank1]:  [ 3 11]
[rank1]:  [ 4 11]
[rank1]:  [14  7]
[rank1]:  [ 5 15]
[rank1]:  [13 12]
[rank1]:  [ 6  8]
[rank1]:  [ 0  2]
[rank1]:  [10  7]
[rank1]:  [ 5 13]
[rank1]:  [ 2  5]
[rank1]:  [14  1]
[rank1]:  [ 6  1]
[rank1]:  [ 7  8]
[rank1]:  [ 5  4]
[rank1]:  [ 2 15]
[rank1]:  [ 6  2]
[rank1]:  [ 4  2]
[rank1]:  [ 5  3]
[rank1]:  [11 10]
[rank1]:  [ 2  0]
[rank1]:  [12  8]
[rank1]:  [13  6]
[rank1]:  [ 8  4]
[rank1]:  [11 15]
[rank1]:  [14  7]
[rank1]:  [12  1]
[rank1]:  [ 1  5]
[rank1]:  [ 6  0]
[rank1]:  [ 0  8]],
[rank1]: Input references: [array([1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1]), array([1, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 1, 0]), array([1, 0, 0, 0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1]), ..., array([0, 0, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1]), array([1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0]), array([0, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0])]
[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/ubuntu/aws-neuron-samples/torch-neuronx/training/hf_text_classification/transformers/examples/pytorch/text-classification/run_glue.py", line 626, in <module>
[rank0]:     main()
[rank0]:   File "/home/ubuntu/aws-neuron-samples/torch-neuronx/training/hf_text_classification/transformers/examples/pytorch/text-classification/run_glue.py", line 564, in main
[rank0]:     metrics = trainer.evaluate(eval_dataset=eval_dataset)
[rank0]:   File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/transformers/trainer.py", line 4199, in evaluate
[rank0]:     output = eval_loop(
[rank0]:   File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/transformers/trainer.py", line 4489, in evaluation_loop
[rank0]:     metrics = self.compute_metrics(
[rank0]:   File "/home/ubuntu/aws-neuron-samples/torch-neuronx/training/hf_text_classification/transformers/examples/pytorch/text-classification/run_glue.py", line 498, in compute_metrics
[rank0]:     result = metric.compute(predictions=preds, references=p.label_ids)
[rank0]:   File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/evaluate/module.py", line 455, in compute
[rank0]:     self.add_batch(**inputs)
[rank0]:   File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/evaluate/module.py", line 546, in add_batch
[rank0]:     raise ValueError(error_msg) from None
[rank0]: ValueError: Predictions and/or references don't match the expected format.
[rank0]: Expected format: {'predictions': Value(dtype='int64', id=None), 'references': Value(dtype='int64', id=None)},
[rank0]: Input predictions: [[ 1  0]
[rank0]:  [ 3  7]
[rank0]:  [ 2  0]
[rank0]:  [ 6  3]
[rank0]:  [10  3]
[rank0]:  [ 6 14]
[rank0]:  [14  7]
[rank0]:  [ 0  5]
[rank0]:  [ 2  1]
[rank0]:  [11  7]
[rank0]:  [ 0 10]
[rank0]:  [ 6 15]
[rank0]:  [11 14]
[rank0]:  [13 14]
[rank0]:  [ 2  7]
[rank0]:  [ 2 10]
[rank0]:  [ 2  7]
[rank0]:  [ 0  3]
[rank0]:  [ 8  2]
[rank0]:  [15  0]
[rank0]:  [ 9  0]
[rank0]:  [ 1  0]
[rank0]:  [ 9  1]
[rank0]:  [12  1]
[rank0]:  [ 8 12]
[rank0]:  [ 8 15]
[rank0]:  [ 3 11]
[rank0]:  [ 4 11]
[rank0]:  [14  7]
[rank0]:  [ 5 15]
[rank0]:  [13 12]
[rank0]:  [ 6  8]
[rank0]:  [ 0  2]
[rank0]:  [10  7]
[rank0]:  [ 5 13]
[rank0]:  [ 2  5]
[rank0]:  [14  1]
[rank0]:  [ 6  1]
[rank0]:  [ 7  8]
[rank0]:  [ 5  4]
[rank0]:  [ 2 15]
[rank0]:  [ 6  2]
[rank0]:  [ 4  2]
[rank0]:  [ 5  3]
[rank0]:  [11 10]
[rank0]:  [ 2  0]
[rank0]:  [12  8]
[rank0]:  [13  6]
[rank0]:  [ 8  4]
[rank0]:  [11 15]
[rank0]:  [14  7]
[rank0]:  [12  1]
[rank0]:  [ 1  5]
[rank0]:  [ 6  0]
[rank0]:  [ 0  8]],
[rank0]: Input references: [array([1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1]), array([1, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 1, 0]), array([1, 0, 0, 0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1]), ..., array([0, 0, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1]), array([1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0]), array([0, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0])]
W0718 19:39:23.555000 107154 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 107163 closing signal SIGTERM
E0718 19:39:24.170000 107154 torch/distributed/elastic/multiprocessing/api.py:874] failed (exitcode: 1) local_rank: 1 (pid: 107164) of binary: /home/ubuntu/aws_neuron_venv_pytorch/bin/python
Traceback (most recent call last):
  File "/home/ubuntu/aws_neuron_venv_pytorch/bin/torchrun", line 8, in <module>
    sys.exit(main())
  File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
    return f(*args, **kwargs)
  File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/torch/distributed/run.py", line 892, in main
    run(args)
  File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/torch/distributed/run.py", line 883, in run
    elastic_launch(
  File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 139, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 270, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
transformers/examples/pytorch/text-classification/run_glue.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2025-07-18_19:39:23
  host      : ip-172-31-43-136.us-west-2.compute.internal
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 107164)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
---------------------------------------------------------------------------
CalledProcessError                        Traceback (most recent call last)
Cell In[7], line 18
      2 RUN_CMD = f"""{env_var_options} torchrun --nproc_per_node={num_workers} \
      3 transformers/examples/pytorch/text-classification/run_glue.py \
      4 --model_name_or_path {model_name} \
   (...)
     14 --eval_do_concat_batches=False \
     15 --output_dir {model_base_name}-{task_name}-{num_workers}w-{batch_size}bs"""
     17 print(f'Running command: \n{RUN_CMD}')
---> 18 if subprocess.check_call(RUN_CMD,shell=True):
     19    print("There was an error with the fine-tune command")
     20 else:

File /usr/lib/python3.10/subprocess.py:369, in check_call(*popenargs, **kwargs)
    367     if cmd is None:
    368         cmd = popenargs[0]
--> 369     raise CalledProcessError(retcode, cmd)
    370 return 0

CalledProcessError: Command 'XLA_USE_BF16=1 NEURON_CC_FLAGS="--model-type=transformer" torchrun --nproc_per_node=2 transformers/examples/pytorch/text-classification/run_glue.py --model_name_or_path camembert-base --task_name sst2 --do_train --do_eval --max_seq_length 128 --dataset_name glue --per_device_train_batch_size 8 --learning_rate 2e-05 --num_train_epochs 5 --overwrite_output_dir --eval_do_concat_batches=False --output_dir camembert-base-sst2-2w-8bs' returned non-zero exit status 1.
```


### Expected behavior

Training should complete successfully

Comments:
Test passes using the below config
```
pip install -U "protobuf<4" "transformers==4.53.*" "datasets<=3.6.0" "accelerate==1.7.*" scikit-learn evaluate
git clone https://github.com/huggingface/transformers --branch v4.52.3
```

```
***** eval metrics *****
  epoch                   =        5.0
  eval_accuracy           =     0.8841
  eval_loss               =     0.4748
  eval_runtime            = 0:00:03.40
  eval_samples            =        872
  eval_samples_per_second =     129.41
  eval_steps_per_second   =      8.235
Fine-tune Successful!!!
```