[Bug]: OOM occurs when running THUDM/glm-4-9b-chat-1m, both with dense and minference mode.

Describe the bug

When using THUDM/glm-4-9b-chat-1m, the model runs into OOM errors in both dense and minference modes.

In dense mode, the model runs successfully only when I comment out the following line in new_patch:

def new_patch(model, config):
    if model.__class__.__name__ == "ChatGLMForConditionalGeneration":
        # model = patch_glm_4_1m(model, config)  # ← causes OOM
        return model

Any suggestions? Thanks in advance.

Steps to reproduce

python ../infinitebench.py --task kv_retrieval --max_seq_length 160000 --attn minference --num_eval_examples 1

The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
0it [00:19, ?it/s]
Traceback (most recent call last):
  File "/home/eval/infinite_bench/scripts/../infinitebench.py", line 295, in <module>
    pred = get_pred(
  File "/home/eval/infinite_bench/scripts/../infinitebench.py", line 111, in get_pred
    outputs = model.generate(**input_tensors, generation_config=generation_config)
  File "/home/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/home/.local/lib/python3.10/site-packages/transformers/generation/utils.py", line 2215, in generate
    result = self._sample(
  File "/home/.local/lib/python3.10/site-packages/transformers/generation/utils.py", line 3206, in _sample
    outputs = self(**model_inputs, return_dict=True)
  File "/home/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/MInference/minference/utils.py", line 99, in new_forward
    return original_forward(*args, **kwargs)
  File "/home/model_hub/huggingface/modules/transformers_modules/THUDM/glm-4-9b-chat-1m/0aa722c7e0745dd21453427dd44c257dd253304f/modeling_chatglm.py", line 1000, in forward
    transformer_outputs = self.transformer(
  File "/home/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/model_hub/huggingface/modules/transformers_modules/THUDM/glm-4-9b-chat-1m/0aa722c7e0745dd21453427dd44c257dd253304f/modeling_chatglm.py", line 898, in forward
    hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
  File "/home/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/MInference/minference/patch.py", line 960, in <lambda>
    lambda self, *args, **kwargs: transformer_forward(self, *args, **kwargs)
  File "/home/MInference/minference/utils.py", line 443, in transformer_forward
    layer_ret = layer(
  File "/home/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/model_hub/huggingface/modules/transformers_modules/THUDM/glm-4-9b-chat-1m/0aa722c7e0745dd21453427dd44c257dd253304f/modeling_chatglm.py", line 651, in forward
    mlp_output = self.mlp(layernorm_output)
  File "/home/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/model_hub/huggingface/modules/transformers_modules/THUDM/glm-4-9b-chat-1m/0aa722c7e0745dd21453427dd44c257dd253304f/modeling_chatglm.py", line 585, in forward
    intermediate_parallel = self.activation_func(intermediate_parallel)
  File "/home/model_hub/huggingface/modules/transformers_modules/THUDM/glm-4-9b-chat-1m/0aa722c7e0745dd21453427dd44c257dd253304f/modeling_chatglm.py", line 569, in swiglu
    return F.silu(x[0]) * x[1]
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.51 GiB. GPU 0 has a total capacity of 79.15 GiB of which 3.07 GiB is free. Process 338657 has 76.07 GiB memory in use. Of the allocated memory 65.20 GiB is allocated by PyTorch, and 10.38 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Hardware: A100 with 80G

Expected Behavior

No response

Logs

No response

Additional Information

No response