Error at scatter num_items_in_batch in ddp/dp
Author: MVPavan
Pull request: no
Labels: name: bug

### System Info

Name: transformers
Version: 4.46.2
Python 3.11.10

### Who can help?

Hi @muellerzr @SunMarc, thanks in advance.

Trainer in compute loss adds num_items_in_batch as tensor scalar to kwargs, in DDP when kwargs are scattered it throws following error:

`RuntimeError: chunk expects at least a 1-dimensional tensor`

reference issue: [TRL Issue](https://github.com/huggingface/trl/issues/2338)

### 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

Using two gpu machine. Code to replicate the issue:

```
from datasets import load_dataset
dataset = load_dataset(path="HuggingFaceTB/smoltalk", name="everyday-conversations")
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer, setup_chat_format
import torch

device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available() else "cpu"
)

model_name = "HuggingFaceTB/SmolLM2-135M"

model = AutoModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path=model_name
).to(device)
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)
model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)

finetune_name = "SmolLM2-FT-Lora-test"
finetune_tags = ["smol-course", "module_1"]

from peft import LoraConfig

rank_dimension = 6
lora_alpha = 8
lora_dropout = 0.05

peft_config = LoraConfig(
    r=rank_dimension,  # Rank dimension - typically between 4-32
    lora_alpha=lora_alpha,  # LoRA scaling factor - typically 2x rank
    lora_dropout=lora_dropout,  # Dropout probability for LoRA layers
    bias="none",  # Bias type for LoRA. the corresponding biases will be updated during training.
    target_modules="all-linear",  # Which modules to apply LoRA to
    task_type="CAUSAL_LM",  # Task type for model architecture
)
args = SFTConfig(
    # Output settings
    output_dir=finetune_name,  # Directory to save model checkpoints
    # Training duration
    num_train_epochs=5,  # Number of training epochs
    # Batch size settings
    per_device_train_batch_size=2,  # Batch size per GPU
    gradient_accumulation_steps=2,  # Accumulate gradients for larger effective batch
    # Memory optimization
    gradient_checkpointing=True,  # Trade compute for memory savings
    gradient_checkpointing_kwargs = {"use_reentrant": False},
    # Optimizer settings
    optim="adamw_torch_fused",  # Use fused AdamW for efficiency
    learning_rate=2e-4,  # Learning rate (QLoRA paper)
    max_grad_norm=0.3,  # Gradient clipping threshold
    # Learning rate schedule
    warmup_ratio=0.03,  # Portion of steps for warmup
    lr_scheduler_type="constant",  # Keep learning rate constant after warmup
    # Logging and saving
    logging_steps=10,  # Log metrics every N steps
    save_strategy="epoch",  # Save checkpoint every epoch
    # Precision settings
    bf16=True,  # Use bfloat16 precision
    # Integration settings
    push_to_hub=False,  # Don't push to HuggingFace Hub
    report_to="none",  # Disable external logging
)

max_seq_length = 1512  # max sequence length for model and packing of the dataset


trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset["train"],
    peft_config=peft_config,  # LoRA configuration
    max_seq_length=max_seq_length,  # Maximum sequence length
    tokenizer=tokenizer,
    packing=True,  # Enable input packing for efficiency
    dataset_kwargs={
        "add_special_tokens": False,  # Special tokens handled by template
        "append_concat_token": False,  # No additional separator needed
    },
)

trainer.train()
```

### Expected behavior

`RuntimeError: chunk expects at least a 1-dimensional tensor`

Comments:
Solution can be to make the tensor scalar to tensor of 1 dim:

In trainer.py -> compute_loss function
```
    if self.model_accepts_loss_kwargs:
        loss_kwargs = {}
        if num_items_in_batch is not None:
            #NOTE: To handle mutigpu scatter, num_items_in_batch is required to have at least 1 dimension
            # So tensor scalar is unsqueezed, This is a workaround for the current implementation of scatter in PyTorch
            if isinstance(num_items_in_batch, torch.Tensor) and self.args.n_gpu > 1:
                if len(num_items_in_batch.size()) == 0:
                    num_items_in_batch = num_items_in_batch.unsqueeze(0).repeat(self.args.n_gpu)
            loss_kwargs["num_items_in_batch"] = num_items_in_batch
        inputs = {**inputs, **loss_kwargs}
```

This will also make loss 1 dim tensor, however in ddp during gather loss is anyway converted from scalar tensor to tensor of 1dim, so overall no effect at the end.

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.

Just be the sure, is the issue fixed by the PR linked above @MVPavan ?