Distributed predictions are generated repeatedly
Author: Sakurakdx
Pull request: no

### System Info

- `transformers` version: 4.35.0.dev0
- Platform: Linux-5.4.143.bsk.8-amd64-x86_64-with-glibc2.28
- Python version: 3.10.13
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.23.0
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: DEEPSPEED
        - use_cpu: False
        - debug: False
        - num_processes: 1
        - machine_rank: 0
        - num_machines: 1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - deepspeed_config: {'deepspeed_config_file': 'none', 'zero3_init_flag': False}
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
        - dynamo_config: {'dynamo_backend': 'INDUCTOR', 'dynamo_mode': 'default', 'dynamo_use_dynamic': False, 'dynamo_use_fullgraph': False}
- PyTorch version (GPU?): 2.1.0+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker @younesbelkada @muellerz muller @pacman100 

### Information

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

### Tasks

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

### Reproduction

Loading Lora fine-tuned LLAMA for prediction, when I predict on multiple GPUs, there are duplicate predictions, but on a single GPU this problem does not occur.

Mytrainer:
```python
class CausalFtTrainer(Seq2SeqTrainer):
    # https://discuss.huggingface.co/t/trainer-vs-seq2seqtrainer/3145/2
    def __init__(
        self,
        params: Dict[str, Any],
        model: torch.nn.Module,
        train_dataset: Optional[Dataset] = None,
        eval_dataset: Optional[Dataset] = None,
        data_collator: Optional[Callable] = None,
        tokenizer: Optional[PreTrainedTokenizerBase] = None,
        callback: Optional[LLMCallback] = None,
        compute_metrics: Optional[Callable] = None,
        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
    ) -> None:
        # TODO: change the base class to Trainer
        self.combine_causal_loss_factor = params.pop("combine_causal_loss_factor", 0)

        self.params = params

        if train_dataset and not eval_dataset:
            # generation_config should be a generation_config instance or None,
            # not a dict, so if eval_dataset is None, set this key to None
            self.params["generation_config"] = None
        if eval_dataset:
            # https://huggingface.co/docs/transformers/main/main_classes/text_generation
            generation_config = GenerationConfig(**self.params.pop("generation_config"))
            self.params["generation_config"] = generation_config

        self.arguments = CausalFtTrainingArguments(**self.params)

        super().__init__(
            model=model,
            args=self.arguments,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            data_collator=data_collator,
            tokenizer=tokenizer,
            callbacks=[LLMCallback] if not callback else None,
            compute_metrics=compute_metrics,
            optimizers=optimizers,
        )


    def predict(
        self,
        test_dataset: Dataset,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "test",
        **gen_kwargs,
    ) -> PredictionOutput:
        data_size = len(test_dataset)
        # for arnold
        if os.environ.get("ARNOLD_WORKER_GPU", None):
            gpu_number = int(os.environ.get("ARNOLD_WORKER_GPU", 1)) * int(os.environ.get("ARNOLD_WORKER_NUM", 1))
        else:
            gpu_number = torch.cuda.device_count()
        self.total_step = math.ceil(data_size // (self.params["per_device_eval_batch_size"] * gpu_number))
        self.current_step = 0
        return super().predict(test_dataset, ignore_keys, metric_key_prefix, **gen_kwargs)

    def prediction_step(
        self, model: torch.nn.Module, inputs: Dict[str, torch.Tensor], prediction_loss_only: bool, ignore_keys: Optional[List[str]] = None, **kwargs
    ) -> Tuple[float, torch.Tensor, torch.Tensor]:
        y_true = inputs.pop("ground_truth_labels", None)

        if not self.args.predict_with_generate or prediction_loss_only:
            return super().prediction_step(model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys)

        has_labels = "labels" in inputs
        inputs = self._prepare_inputs(inputs)

        gen_kwargs = self._gen_kwargs.copy()

        if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None:
            gen_kwargs["max_length"] = self.model.config.max_length
        gen_kwargs["num_beams"] = gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.model.config.num_beams
        default_synced_gpus = True if is_deepspeed_zero3_enabled() else False
        gen_kwargs["synced_gpus"] = gen_kwargs["synced_gpus"] if gen_kwargs.get("synced_gpus") is not None else default_synced_gpus

        # If the `decoder_input_ids` was created from `labels`, evict the former, so that the model can freely generate
        # (otherwise, it would continue generating from the padded `decoder_input_ids`)
        if "labels" in inputs and "decoder_input_ids" in inputs and inputs["labels"].shape == inputs["decoder_input_ids"].shape:
            inputs = {k: v for k, v in inputs.items() if k != "decoder_input_ids"}

        # extra params for generation_config
        if not isinstance(self.tokenizer, CasterTokenizer):
            self.params["generation_config"].bos_token_id = self.tokenizer.bos_token_id
        self.params["generation_config"].eos_token_id = self.tokenizer.eos_token_id
        self.params["generation_config"].pad_token_id = self.tokenizer.pad_token_id
        gen_kwargs["generation_config"] = self.params["generation_config"]

        scores = None
        if not getattr(self.params["generation_config"], "output_scores", False):
            generated_tokens = self.model.generate(**inputs, **gen_kwargs)
        else:
            # set return_dict_in_generate to True in order to get the output dict
            self.params["generation_config"].return_dict_in_generate = True
            outputs = self.model.generate(**inputs, **gen_kwargs)
            generated_tokens, scores = outputs["sequences"], outputs["scores"]
        # generated_tokens:         [batch_size, seq_len]
        # scores:                   ([1, vocab_size], [1, vocab_size], ...)
        generated_str_length = generated_tokens.size(0)

        generated_token_probs = [[] for _ in range(generated_str_length)]
        if scores:
            generated_token_length = len(scores)
            for index_ in range(-generated_token_length, 0):
                probs = (
                    torch.nn.functional.softmax(scores[index_], dim=1).type(torch.float32).cpu().numpy()
                )  # bf16 should change type to torch.float32, so use torch.float32 to convert
                for _ in range(generated_str_length):
                    index_probs = generated_tokens[_][index_]
                    if self.args.care_tokens:
                        index_probs = self.tokenizer.encode(self.args.care_tokens)
                    generated_token_probs[_].append(probs[_][index_probs])
            generated_token_probs = np.array(generated_token_probs).tolist()
            del scores
        # Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop
        # TODO: remove this hack when the legacy code that initializes generation_config from a model config is
        # removed in https://github.com/huggingface/transformers/blob/98d88b23f54e5a23e741833f1e973fdf600cc2c5/src/transformers/generation/utils.py#L1183
        if self.model.generation_config._from_model_config:
            self.model.generation_config._from_model_config = False

        # Retrieves GenerationConfig from model.generation_config
        gen_config = self.model.generation_config
        # in case the batch is shorter than max length, the output should be padded
        if generated_tokens.shape[-1] < gen_config.max_length:
            generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_config.max_length)
        elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
            generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_config.max_new_tokens + 1)

        with torch.no_grad():
            if has_labels:
                with self.compute_loss_context_manager():
                    outputs = model(**inputs)
                if self.label_smoother is not None:
                    loss = self.label_smoother(outputs, inputs["labels"]).mean().detach()
                else:
                    loss = (outputs["loss"] if isinstance(outputs, dict) else outputs[0]).mean().detach()
            else:
                loss = None

        if self.args.prediction_loss_only:
            return loss, None, None

        if has_labels:
            labels = inputs["labels"]
            if labels.shape[-1] < gen_config.max_length:
                labels = self._pad_tensors_to_max_len(labels, gen_config.max_length)
            elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1:
                labels = self._pad_tensors_to_max_len(labels, gen_config.max_new_tokens + 1)
        else:
            labels = None

        # process the generated tokens
        if dist.get_rank() == 0 and not self.args.do_eval:
            if self.current_step % self.params["logging_steps"] == 0:
                ground_truth_ = self.decode([y_true[0]]) if y_true is not None else None
                output = self.decode([generated_tokens[0]])
                logger.info(f"[GPU_{dist.get_rank()}, {self.current_step}/{self.total_step}]: \n ground_truth: {ground_truth_[0]} \n output: {output[0]}")
            self.current_step += 1
        self._output_generate_results(inputs, generated_tokens, y_true, generated_token_probs)
        # generated_tokens = torch.tensor([1]).cuda(f"cuda:{dist.get_rank()}")

        return loss, generated_tokens, labels
```

main function
```python
predictor = CausalFtTrainer(
    params=trainer_params,
    model=causal_model,
    eval_dataset=dataset,
    tokenizer=tokenizer,
    data_collator=data_collator,
)

predictor.predict(dataset, **predict_params)
```
Another problem is that when running on multiple cards, it will get stuck in the pred_gather function or report an error.

![image](https://github.com/huggingface/transformers/assets/48399040/e53b0f72-f96e-432d-a620-d221f6b39e72)
![image](https://github.com/huggingface/transformers/assets/48399040/6ad105f2-1ac8-4e15-a077-4f67f3c5a653)


### Expected behavior

![image](https://github.com/huggingface/transformers/assets/48399040/e53b0f72-f96e-432d-a620-d221f6b39e72)
![image](https://github.com/huggingface/transformers/assets/48399040/6ad105f2-1ac8-4e15-a077-4f67f3c5a653)


Comments:
cc @SunMarc if you have some time have a look! 

Hi @Sakurakdx, were you able to fix the issue ? If not, could you provide a minimal reproducer ? It will help me a lot to debug the issue I see that you are using your own trainer. Does this issue also happen on HF Trainer ? Thanks. 

@SunMarc, I didn't actually solve the problem; I just did some post-processing after the prediction was completed. The trainer I used was inherited from Hugging Face's trainer, and theoretically, it's a bug in the native trainer. I have already left the project, so I'm unable to provide a reproducer, sorry.

Thanks for your reply and sorry to not have answered earlier. I'm closing this issue if you are fine if that. 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.