Repository: transformers
Issue: Error occurs when using model.generate with Gemma2 in ZeRO3 environment
Labels: bug

### System Info

> pip install git+https://github.com/huggingface/transformers.git@a6256ec0982fee2c57cc41237bff7e64ed4dcda9
> pip install deepspeed==0.15.4
> pip install torch==2.5.1+cu121

<details>
<summary>huggingface-env</summary>

transformers: 4.48.0.dev0
huggingface_hub version: 0.27.0
Platform: Linux-5.15.0-124-generic-x86_64-with-glibc2.35
Python version: 3.10.12
Running in iPython ?: No
Running in notebook ?: No
Running in Google Colab ?: No
Running in Google Colab Enterprise ?: No
Token path ?: /root/.cache/huggingface/token
Has saved token ?: True
Who am I ?: jp1924
Configured git credential helpers: 
FastAI: N/A
Tensorflow: N/A
Torch: 2.5.1+cu121
Jinja2: 3.1.5
Graphviz: N/A
keras: N/A
Pydot: N/A
Pillow: 10.4.0
hf_transfer: N/A
gradio: N/A
tensorboard: N/A
numpy: 1.26.4
pydantic: 2.10.4
aiohttp: 3.11.11
ENDPOINT: https://huggingface.co
HF_HUB_CACHE: /root/.cache/huggingface/hub
HF_ASSETS_CACHE: /root/.cache/huggingface/assets
HF_TOKEN_PATH: /root/.cache/huggingface/token
HF_STORED_TOKENS_PATH: /root/.cache/huggingface/stored_tokens
HF_HUB_OFFLINE: False
HF_HUB_DISABLE_TELEMETRY: False
HF_HUB_DISABLE_PROGRESS_BARS: None
HF_HUB_DISABLE_SYMLINKS_WARNING: False
HF_HUB_DISABLE_EXPERIMENTAL_WARNING: False
HF_HUB_DISABLE_IMPLICIT_TOKEN: False
HF_HUB_ENABLE_HF_TRANSFER: False
HF_HUB_ETAG_TIMEOUT: 10
HF_HUB_DOWNLOAD_TIMEOUT: 10

</details>

<details>
<summary>deepspeed-env</summary>

 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
 [WARNING]  FP Quantizer is using an untested triton version (3.1.0), only 2.3.(0, 1) and 3.0.0 are known to be compatible with these kernels
fp_quantizer ........... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
 [WARNING]  gds requires the dev libaio .so object and headers but these were not found.
 [WARNING]  gds: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
gds .................... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.5
 [WARNING]  using untested triton version (3.1.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]


DeepSpeed general environment info:
torch install path ............... ['/usr/local/lib/python3.10/dist-packages/torch']
torch version .................... 2.5.1+cu121
deepspeed install path ........... ['/usr/local/lib/python3.10/dist-packages/deepspeed']
deepspeed info ................... 0.15.4, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 12.1
deepspeed wheel compiled w. ...... torch 0.0, cuda 0.0
shared memory (/dev/shm) size .... 126.00 GB

</details>

<details>
<summary>torch-env</summary>

PyTorch version: 2.5.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.31.2
Libc version: glibc-2.35

Python version: 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-124-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100 80GB PCIe
GPU 1: NVIDIA A100 80GB PCIe
GPU 2: NVIDIA A100 80GB PCIe
GPU 3: NVIDIA A100 80GB PCIe

Nvidia driver version: 560.35.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 57 bits virtual
Byte Order:                           Little Endian
CPU(s):                               48
On-line CPU(s) list:                  0-47
Vendor ID:                            GenuineIntel
Model name:                           Intel(R) Xeon(R) Silver 4310 CPU @ 2.10GHz
CPU family:                           6
Model:                                106
Thread(s) per core:                   2
Core(s) per socket:                   12
Socket(s):                            2
Stepping:                             6
CPU max MHz:                          3300.0000
CPU min MHz:                          800.0000
BogoMIPS:                             4200.00
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 invpcid_single ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect wbnoinvd dtherm ida arat pln pts avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid fsrm md_clear pconfig flush_l1d arch_capabilities
Virtualization:                       VT-x
L1d cache:                            1.1 MiB (24 instances)
L1i cache:                            768 KiB (24 instances)
L2 cache:                             30 MiB (24 instances)
L3 cache:                             36 MiB (2 instances)
NUMA node(s):                         2
NUMA node0 CPU(s):                    0-11,24-35
NUMA node1 CPU(s):                    12-23,36-47
Vulnerability Gather data sampling:   Mitigation; Microcode
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI SW loop, KVM SW loop
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.1.3.1
[pip3] nvidia-cuda-cupti-cu12==12.1.105
[pip3] nvidia-cuda-nvrtc-cu12==12.1.105
[pip3] nvidia-cuda-runtime-cu12==12.1.105
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.0.2.54
[pip3] nvidia-curand-cu12==10.3.2.106
[pip3] nvidia-cusolver-cu12==11.4.5.107
[pip3] nvidia-cusparse-cu12==12.1.0.106
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.1.105
[pip3] torch==2.5.1+cu121
[pip3] torchvision==0.20.1+cu121
[pip3] triton==3.1.0
[conda] Could not collect
</details>

### Who can help?

@muellerzr 
@gante

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
import os

from accelerate import Accelerator

from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    Gemma2ForCausalLM,
    GenerationConfig,
    TextGenerationPipeline,
    TrainingArguments,
)
from transformers.integrations import is_deepspeed_zero3_enabled


CHAT_TEMPLATE = "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{{ bos_token }}{% for message in messages %}{{ '<start_of_turn>' }}{% if message.role == 'user' %}{{ '### User:\\n' }}{% if message.content is not string %}{% for content in message.content %}{% if content.type == 'image' %}{{ '<img>' }}{% elif content.type == 'text' %}{{ content.text }}{% else %}{# Do nothing #}{% endif %}{% endfor %}{% else %}{{ message.content }}{% endif %}{{ '\\n\\n' }}{% elif message.role == 'system' %}{{ '### System:\\n' }}{% if message.content is not string %}{% for content in message.content %}{% if content.type == 'image' %}{{ '<img>' }}{% elif content.type == 'text' %}{{ content.text }}{% else %}{# Do nothing #}{% endif %}{% endfor %}{% else %}{{ message.content }}{% endif %}{{ '\\n\\n' }}{% elif message.role == 'assistant' %}{{ '### Assistant:\\n' }}{% if message.content is not string %}{% for content in message.content %}{% if content.type == 'text' %}{{ content.text }}{% else %}{# Do nothing #}{% endif %}{% endfor %}{% else %}{{ message.content }}{% endif %}{% else %}{{ '' }}{% endif %}{{ '<end_of_turn>' }}{% endfor %}{% if not add_generation_prompt %}{{ eos_token }}{% elif add_generation_prompt %}{{ '<start_of_turn>' }}{{ '### Assistant:\\n' }}{% else %}{# Do nothing #}{% endif %}"


def main(args: TrainingArguments) -> None:
    if not is_deepspeed_zero3_enabled():
        raise ValueError("This script is only compatible with DeepSpeed Zero3.")

    accelerator = Accelerator(deepspeed_plugin=args.deepspeed_plugin)

    model_name = "google/gemma-2-9b"
    config = AutoConfig.from_pretrained(model_name, attn_implementation="eager")
    pipeline = TextGenerationPipeline(
        model=AutoModelForCausalLM.from_pretrained(model_name, config=config),
        tokenizer=AutoTokenizer.from_pretrained(model_name, chat_template=CHAT_TEMPLATE),
        device=args.device,
        batch_size=args.per_device_eval_batch_size,
    )
    pipeline.model = accelerator.prepare(pipeline.model)

    match args.local_rank:
        case 0:
            max_new_tokens = 10
        case 1:
            max_new_tokens = 2

    generate_kwargs = {
        "generation_config": GenerationConfig(
            max_new_tokens=max_new_tokens,
            cache_implementation="hybrid",
            use_cache=True,
        ),
        "synced_gpus": True,
    }
    conversations = [{"role": "user", "content": "hi"}]
    pipeline(conversations, return_full_text=True, **generate_kwargs)


if "__main__" in __name__:
    args = TrainingArguments(
        output_dir=".",
        per_device_eval_batch_size=1,
        deepspeed={
            "bf16": {"enabled": "auto"},
            "zero_optimization": {
                "stage": 3,
                "allgather_partitions": True,
                "reduce_bucket_size": "auto",
                "stage3_prefetch_bucket_size": "auto",
                "stage3_param_persistence_threshold": "auto",
                "stage3_max_live_parameters": 0,
                "stage3_max_reuse_distance": 0,
                "allgather_bucket_size": 0,
                "overlap_comm": False,
                "reduce_scatter": True,
                "contiguous_gradients": True,
                "stage3_gather_16bit_weights_on_model_save": True,
                "offload_optimizer": {"device": "cpu", "pin_memory": True},
            },
            "gradient_accumulation_steps": "auto",
            "gradient_clipping": "auto",
            "steps_per_print": 2000,
            "train_batch_size": "auto",
            "train_micro_batch_size_per_gpu": "auto",
            "wall_clock_breakdown": False,
        },
        bf16=True,
        local_rank=os.getenv("LOCAL_RANK", -1),
    )
    main(args)
# pip install git+https://github.com/huggingface/transformers.git@a6256ec0982fee2c57cc41237bff7e64ed4dcda9
```

`deepspeed --num_gpus=2 main.py`

### Expected behavior

I'm getting the following error when using model.generate with Gemma2 in ZeRO3 environment with hybrid cache:

```python
[rank1]:   File "/root/main.py", line 86, in <module>
[rank1]:     main(args)
[rank1]:   File "/root/main.py", line 52, in main
[rank1]:     pipeline(conversations, return_full_text=True, **generate_kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/text_generation.py", line 278, in __call__
[rank1]:     return super().__call__(Chat(text_inputs), **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/base.py", line 1362, in __call__
[rank1]:     return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/base.py", line 1369, in run_single
[rank1]:     model_outputs = self.forward(model_inputs, **forward_params)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/base.py", line 1269, in forward
[rank1]:     model_outputs = self._forward(model_inputs, **forward_params)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/text_generation.py", line 383, in _forward
[rank1]:     generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
[rank1]:   File "/home/jp/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank1]:     return func(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py", line 2254, in generate
[rank1]:     result = self._sample(
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py", line 3246, in _sample
[rank1]:     model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/transformers/models/gemma2/modeling_gemma2.py", line 904, in prepare_inputs_for_generation
[rank1]:     input_ids = input_ids[:, cache_position]
[rank1]: RuntimeError: CUDA error: device-side assert triggered
[rank1]: Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

I've only debugged this with the Gemma2 model so far. I don't know how it behaves with other models.
The issue occurs due to shape mismatch between input_ids and cache_position or attention_mask.

When `sync_gpus` and `this_peer_finished` are true, before skipping the addition of generated tokens to input_ids,
it always goes through `self._update_model_kwargs_for_generation`, which causes an error by adding 1 to cache_position or attention_mask.

Here's my proposed solution:
```diff
            if "attention_mask" in model_kwargs:
                attention_mask = model_kwargs["attention_mask"]
                model_kwargs["attention_mask"] = torch.cat(
-                    [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
+                    [attention_mask, attention_mask.new_ones((attention_mask.shape[0], num_new_tokens))], dim=-1
                )
```

```diff
        if model_kwargs.get("use_cache", True):
            model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
        else:
            past_positions = model_kwargs.pop("cache_position")
            new_positions = torch.arange(
-                past_positions[-1] + 1, past_positions[-1] + num_new_tokens + 1, dtype=past_positions.dtype
+                past_positions[-1] + num_new_tokens, past_positions[-1] + num_new_tokens + 1, dtype=past_positions.dtype
            ).to(past_positions.device)
            model_kwargs["cache_position"] = torch.cat((past_positions, new_positions))
        return model_kwargs
```

And then:

```diff
            model_kwargs = self._update_model_kwargs_for_generation(
                outputs,
                model_kwargs,
                is_encoder_decoder=self.config.is_encoder_decoder,
+                num_new_tokens = 0 if this_peer_finished else 1,
            )
```

I think this needs to be added. This way, when `this_peer_finished` is True, 
it won't touch irrelevant things like attention_mask or cache_position.