Repository: vllm-project/vllm

Issue:
[Bug]: problems using long context, Command-R
### Your current environment

```
Collecting environment information...
PyTorch version: 2.1.2+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.29.0
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.105.1.el7.x86_64-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
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: 550.54.15
cuDNN version: Could not collect
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:       48 bits physical, 48 bits virtual
Byte Order:          Little Endian
CPU(s):              64
On-line CPU(s) list: 0-63
Vendor ID:           AuthenticAMD
Model name:          AMD EPYC 7763 64-Core Processor
CPU family:          25
Model:               1
Thread(s) per core:  1
Core(s) per socket:  64
Socket(s):           1
Stepping:            1
Frequency boost:     enabled
CPU max MHz:         2450.0000
CPU min MHz:         1500.0000
BogoMIPS:            4900.49
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc art rep_good nopl nonstop_tsc extd_apicid aperfmperf eagerfpu pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_l2 cpb cat_l3 cdp_l3 invpcid_single hw_pstate sme ssbd rsb_ctxsw ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif umip pku ospke vaes vpclmulqdq overflow_recov succor smca
Virtualization:      AMD-V
L1d cache:           2 MiB (64 instances)
L1i cache:           2 MiB (64 instances)
L2 cache:            32 MiB (64 instances)
L3 cache:            256 MiB (8 instances)
NUMA node(s):        1
NUMA node0 CPU(s):   0-63

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.1.2
[pip3] triton==2.1.0
[conda] Could not collectROCM Version: Could not collect
Neuron SDK Version: N/A
vLLM Version: 0.4.0
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
GPU0	GPU1	GPU2	GPU3	CPU Affinity	NUMA Affinity	GPU NUMA ID
GPU0	 X 	PIX	SYS	SYS	0-63	0		N/A
GPU1	PIX	 X 	SYS	SYS	0-63	0		N/A
GPU2	SYS	SYS	 X 	PIX	0-63	0		N/A
GPU3	SYS	SYS	PIX	 X 	0-63	0		N/A

Legend:

  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinks

```


### 🐛 Describe the bug

Build: stock```vllm/vllm-openai:v0.4.0.post1``` container image.

I'm receiving bad output from Command-R with total prompt+generation sequence length >8192 tokens, while in theory the model supports 128k via rope scaling.   Vanilla transformers responds appropriately.  Any suggestions?

| Tot Seq Length | Env | HF | vLLM | 
| --- | --- | --- | --- |
| <=8k | HF | ok | ok |
| >8k | HF | ok | bad |

HF output >8k:
```
Your favorite word is "chicken"! This is a fun and playful word, often bringing to mind images (8212 total tokens)
```

vLLM output >8k:
```
Your ignore ignore\n\n\n\n ignore\n\nconst\n � � � �  (8212 total tokens)
```

### Example code:
```
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer, AutoModelForCausalLM
import sys
import os
import gc
import torch

os.environ['TOKENIZERS_PARALLELISM']="false"

model_id = "CohereForAI/c4ai-command-r-v01"
tokenizer = AutoTokenizer.from_pretrained(model_id)

target_seq_len=8192+20  # push us a bit past 8k
max_new_tokens=20
base_prompt_len=22 # Template+"my favorite word..."+"What is my favorite word"

num_extra_tokens=target_seq_len - base_prompt_len - max_new_tokens 

# Pad with nonsense "Blah " tokens to the desired length
messages=[
    {"role": "user", "content": "My favorite word is \"Chicken\".\n" +
                                "Blah "*num_extra_tokens +
                                "\nWhat is my favorite word?\n" 
    }
    ]

input_ids  = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")

print(f"Tokens (incl. template): {input_ids.shape[1]}")

#############################
## HF
tokenizer = AutoTokenizer.from_pretrained(model_id)
hf_model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)

gen_tokens = hf_model.generate( input_ids, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.3)

gen_text = tokenizer.decode(gen_tokens[0, input_ids.shape[1]:])

tot_tokens=gen_tokens.shape[1]
print(f"HF output:{gen_text} ({tot_tokens} total tokens)")

del hf_model ; gc.collect() ; torch.cuda.empty_cache()

#############################
## vLLM

vllm_model = LLM(model=model_id, tokenizer=model_id, tensor_parallel_size=4, max_model_len=target_seq_len)

samplingparams = SamplingParams(max_tokens=max_new_tokens, temperature=0.3, skip_special_tokens=False)

output = vllm_model.generate(prompt_token_ids=[ input_ids[0].tolist() ], sampling_params=samplingparams)

tot_tokens=target_seq_len-max_new_tokens+len(output[0].outputs[0].token_ids)
print(f"vLLM output: {output[0].outputs[0]} ({tot_tokens} total tokens)")
```

### Output:
```
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.

No chat template is defined for this tokenizer - using the default template for the CohereTokenizerFast class. If the default is not appropriate for your model, please set `tokenizer.chat_template` to an appropriate template. See https://huggingface.co/docs/transformers/main/chat_templating for more information.

Tokens (incl. template): 8192
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 15/15 [02:05<00:00,  8.34s/it]
/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py:2756: UserWarning: Specified kernel cache directory could not be created! This disables kernel caching. Specified directory is /tmp/xdg-cache/torch/kernels. This warning will appear only once per process. (Triggered internally at ../aten/src/ATen/native/cuda/jit_utils.cpp:1442.)
  next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
This is a friendly reminder - the current text generation call will exceed the model's predefined maximum length (8192). Depending on the model, you may observe exceptions, performance degradation, or nothing at all.
HF output:Your favorite word is "chicken"! This is a fun and playful word, often bringing to mind images (8212 total tokens)
2024-04-06 22:19:02,193	INFO worker.py:1752 -- Started a local Ray instance.
INFO 04-06 22:19:03 llm_engine.py:81] Initializing an LLM engine (v0.4.0.post1) with config: model='CohereForAI/c4ai-command-r-v01', speculative_config=None, tokenizer='CohereForAI/c4ai-command-r-v01', tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=8212, download_dir=None, load_format=auto, tensor_parallel_size=4, disable_custom_all_reduce=True, quantization=None, enforce_eager=False, kv_cache_dtype=auto, device_config=cuda, seed=0)
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
INFO 04-06 22:19:11 selector.py:16] Using FlashAttention backend.
(RayWorkerVllm pid=772) INFO 04-06 22:19:13 selector.py:16] Using FlashAttention backend.
INFO 04-06 22:19:13 pynccl_utils.py:45] vLLM is using nccl==2.18.1
(RayWorkerVllm pid=772) INFO 04-06 22:19:13 pynccl_utils.py:45] vLLM is using nccl==2.18.1
INFO 04-06 22:19:17 weight_utils.py:177] Using model weights format ['*.safetensors']
(RayWorkerVllm pid=772) INFO 04-06 22:19:18 weight_utils.py:177] Using model weights format ['*.safetensors']
(RayWorkerVllm pid=920) INFO 04-06 22:19:13 selector.py:16] Using FlashAttention backend. [repeated 2x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/ray-logging.html#log-deduplication for more options.)
(RayWorkerVllm pid=920) INFO 04-06 22:19:13 pynccl_utils.py:45] vLLM is using nccl==2.18.1 [repeated 2x across cluster]
INFO 04-06 22:20:12 model_runner.py:104] Loading model weights took 16.2916 GB
(RayWorkerVllm pid=772) INFO 04-06 22:20:13 model_runner.py:104] Loading model weights took 16.2916 GB
(RayWorkerVllm pid=920) INFO 04-06 22:19:18 weight_utils.py:177] Using model weights format ['*.safetensors'] [repeated 2x across cluster]
INFO 04-06 22:20:17 ray_gpu_executor.py:244] # GPU blocks: 10620, # CPU blocks: 819
INFO 04-06 22:20:20 model_runner.py:791] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.
INFO 04-06 22:20:20 model_runner.py:795] CUDA graphs can take additional 1~3 GiB memory per GPU. If you are running out of memory, consider decreasing `gpu_memory_utilization` or enforcing eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.
(RayWorkerVllm pid=772) INFO 04-06 22:20:20 model_runner.py:791] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.
(RayWorkerVllm pid=772) INFO 04-06 22:20:20 model_runner.py:795] CUDA graphs can take additional 1~3 GiB memory per GPU. If you are running out of memory, consider decreasing `gpu_memory_utilization` or enforcing eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.
(RayWorkerVllm pid=920) INFO 04-06 22:20:13 model_runner.py:104] Loading model weights took 16.2916 GB [repeated 2x across cluster]
INFO 04-06 22:20:31 model_runner.py:867] Graph capturing finished in 10 secs.
Processed prompts:   0%|                                                                                                                                                                | 0/1 [00:00<?, ?it/s](RayWorkerVllm pid=772) INFO 04-06 22:20:31 model_runner.py:867] Graph capturing finished in 10 secs.
(RayWorkerVllm pid=920) INFO 04-06 22:20:20 model_runner.py:791] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. [repeated 2x across cluster]
(RayWorkerVllm pid=920) INFO 04-06 22:20:20 model_runner.py:795] CUDA graphs can take additional 1~3 GiB memory per GPU. If you are running out of memory, consider decreasing `gpu_memory_utilization` or enforcing eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage. [repeated 2x across cluster]
Processed prompts: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.76s/it]
vLLM output: CompletionOutput(index=0, text='Your ignore ignore\n\n\n\n ignore\n\nconst\n � � � � ', token_ids=[17700, 24434, 24434, 206, 206, 206, 206, 24434, 206, 206, 4399, 206, 199861, 199861, 199861, 199861, 228, 199861, 199861, 199861], cumulative_logprob=-16.78925259313553, logprobs=None, finish_reason=length, stop_reason=None) (8212 total tokens)
(RayWorkerVllm pid=920) INFO 04-06 22:20:31 model_runner.py:867] Graph capturing finished in 10 secs. [repeated 2x across cluster]
```


Discussion:
Hi @agt, could you try to change `max_position_embeddings` equal to `model_max_length` in `config.json` file?
Thanks @esmeetu , no change with ```max_position_embeddings```==```model_max_length```==```131072```.   

 ~~It appears the model config doesn't enable rope scaling, while in the Transformers implementation it may happen implicitly.~~

~~A crude patch to ```models/commandr.py``` setting ```rope_scaling``` (with max_position_embeddings back at 8192) seems to yield acceptable output for longer sequence lengths up to 32k, but I'm not sure if that's a reasonable direction to take, or how that factor ought to be calculated/set, or whether higher scaling factors (e.g. *16=128k seq_len) ought to work.~~

```
class CohereModel(nn.Module):
    def __init__(
        self,
        config: CohereConfig,
        linear_method: Optional[LinearMethodBase] = None,
    ):
        # Force vLLM to adopt rope scaling
        config.rope_scaling = { "type" : "dynamic", "factor": 4.0 }
        super().__init__()
```

~~@ saurabhdash2512 , would you be able to weigh in?~~

Transformers impl. isn't using scaling - forcing it here doesn't seem correct

Resolving patch:
diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py
--- a/vllm/model_executor/models/commandr.py
+++ b/vllm/model_executor/models/commandr.py
@@ -140,7 +140,9 @@ def __init__(
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
-        self.max_position_embeddings = config.max_position_embeddings
+        self.max_position_embeddings = getattr(
+            config, "model_max_length", None) or getattr(
+                config, "max_position_embeddings", 8192)
         self.rope_theta = config.rope_theta
         self.rope_scaling = getattr(config, "rope_scaling", None)
         self.use_qk_norm = getattr(config, "use_qk_norm", False)


Test patch: