[FEAT] [non-CUDA]: Support alternative implementation for  `constraints.positive_definite.check`
Author: tjtanaa
Pull request: no
Labels: name: Feature request

### Feature request

Could there be an alternative implementation for
```
/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py:2470: in _init_added_embeddings_weights_with_mean
    is_covariance_psd = constraints.positive_definite.check(epsilon * covariance).all()
```

the `torch.linalg.cholesky` only exists for CUDA in pytorch.

### Motivation

To support vision language embedding model (llava model) on vLLM for ROCm.

When I am trying to enable vision_language embedding model support on vLLM for ROCm, I encounter this issue.

```
tests/models/embedding/vision_language/test_llava_next.py:134: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/models/embedding/vision_language/test_llava_next.py:63: in _run_test
    hf_model.model.resize_token_embeddings(
/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py:2109: in resize_token_embeddings
    model_embeds = self._resize_token_embeddings(new_num_tokens, pad_to_multiple_of, mean_resizing)
/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py:2134: in _resize_token_embeddings
    new_embeddings = self._get_resized_embeddings(
/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py:2291: in _get_resized_embeddings
    self._init_added_embeddings_weights_with_mean(
/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py:2470: in _init_added_embeddings_weights_with_mean
    is_covariance_psd = constraints.positive_definite.check(epsilon * covariance).all()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = PositiveDefinite()
value = tensor([[ 8.4661e-14, -9.3146e-17,  5.4274e-16,  ..., -1.2541e-16,
          8.1008e-16,  2.6355e-16],
        [-9.314...       [ 2.6355e-16, -5.6042e-16,  5.1984e-16,  ..., -1.9993e-16,
         -2.7124e-16,  8.5429e-14]], device='cuda:0')

    def check(self, value):
        sym_check = super().check(value)
        if not sym_check.all():
            return sym_check
>       return torch.linalg.cholesky_ex(value).info.eq(0)
E       RuntimeError: Calling torch.linalg.cholesky on a CUDA tensor requires compiling PyTorch with MAGMA. Please use PyTorch built with MAGMA support.
```

the `torch.linalg.cholesky` only exists for CUDA in pytorch.

### Your contribution

By helping to test on AMD GPUs with the fixed and providing feedback.

Comments:
Hi @tjtanaa, yes, we should definitely add a workaround here! This function is only called when token embeddings are resized, rather than on every forward pass, so I think it's okay to do it on CPU. @ivarflakstad is there a way we can check if we're running on ROCm so we can send the tensors to CPU before running this check?

Yep 👍 
In certain edge cases we already have specific checks for ROCm, like [here](https://github.com/huggingface/transformers/blob/071a161d3e38f56dbda2743b979f0afeed2cd4f1/src/transformers/modeling_utils.py#L2106).

@Rocketknight1 @ivarflakstad  I have tried running the model on CPU, however it seems the Pytorch version in the vLLM image is not compiled for CPU. I am getting similar error message:

```
tests/models/embedding/vision_language/test_llava_next.py:114: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/models/embedding/vision_language/test_llava_next.py:71: in _run_test
    hf_model.model.resize_token_embeddings(
/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py:2109: in resize_token_embeddings
    model_embeds = self._resize_token_embeddings(new_num_tokens, pad_to_multiple_of, mean_resizing)
/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py:2134: in _resize_token_embeddings
    new_embeddings = self._get_resized_embeddings(
/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py:2291: in _get_resized_embeddings
    self._init_added_embeddings_weights_with_mean(
/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py:2470: in _init_added_embeddings_weights_with_mean
    is_covariance_psd = constraints.positive_definite.check(epsilon * covariance).all()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = PositiveDefinite()
value = tensor([[ 8.4661e-14, -9.3147e-17,  5.4274e-16,  ..., -1.2541e-16,
          8.1008e-16,  2.6355e-16],
        [-9.314..., -2.7124e-16],
        [ 2.6355e-16, -5.6042e-16,  5.1984e-16,  ..., -1.9993e-16,
         -2.7124e-16,  8.5429e-14]])

    def check(self, value):
        sym_check = super().check(value)
        if not sym_check.all():
            return sym_check
>       return torch.linalg.cholesky_ex(value).info.eq(0)
E       RuntimeError: Calling torch.linalg.cholesky on a CPU tensor requires compiling PyTorch with LAPACK. Please use PyTorch built with LAPACK support.

```

So, is there other possible approach that do not call `torch.linalg.cholesky_ex` function?

this is going to end up with me writing a cholesky decomposition by hand isn't it :sob:

@tjtanaa can you check if `torch.linalg.eigvalsh` or `torch.linalg.eigh` work for that image?

@Rocketknight1 Unfortunately no 😭

```
>>> import torch
>>> A = torch.tensor([[1.0, 2.0], [2.0, 3.0]])
>>> eigenvalues = torch.linalg.eigvalsh(A)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Calling torch.linalg.eigh or eigvalsh on a CPU tensor requires compiling PyTorch with LAPACK. Please use PyTorch built with LAPACK support.
```

```
>>> import torch
>>> A = torch.tensor([[1.0, 2.0], [2.0, 3.0]])
>>> eigenvalues, eigenvectors = torch.linalg.eigh(A)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Calling torch.linalg.eigh or eigvalsh on a CPU tensor requires compiling PyTorch with LAPACK. Please use PyTorch built with LAPACK support.
```

This is quite difficult - it's hard for us to support an environment where we can't run linear algebra operations on CPU or on GPU! The only solutions I can see at this point are:

1) Updating the environment so that it can correctly use `torch.linalg` somehow
2) Someone (me) manually writes a semi-efficient cholesky decomposition using elementary operations

The problem with 2) is that even if it works, the next line uses `torch.distributions.multivariate_normal.MultivariateNormal`, which I suspect also makes some calls to `torch.linalg`.

@tjtanaa can you try replacing the line `is_covariance_psd = constraints.positive_definite.check(epsilon * covariance).all()` with `is_covariance_psd = True` and telling me if the code works for you?

If you don't know how to do that: Clone the repo from Github, then install it with `pip install -e .` 

After this, you can edit files in the repo and they will be reflected in your `transformers` installation.

@Rocketknight1 Yup it fails.


```diff
                precision_matrix.shape[:-2], loc.shape[:-1]                                                                       [0/1977]
            )                                                                                                                             
            self.precision_matrix = precision_matrix.expand(batch_shape + (-1, -1))                                                       
        self.loc = loc.expand(batch_shape + (-1,))                                                                                        
                                                                                                                                          
        event_shape = self.loc.shape[-1:]                                                                                                 
        super().__init__(batch_shape, event_shape, validate_args=validate_args)                                                           
                                                                                                                                          
        if scale_tril is not None:                                                                                                        
            self._unbroadcasted_scale_tril = scale_tril                                                                                   
        elif covariance_matrix is not None:                                                                                               
>           self._unbroadcasted_scale_tril = torch.linalg.cholesky(covariance_matrix)                                                     
E           RuntimeError: Calling torch.linalg.cholesky on a CUDA tensor requires compiling PyTorch with MAGMA. Please use PyTorch built w
ith MAGMA support.

+ /usr/local/lib/python3.12/dist-packages/torch/distributions/multivariate_normal.py:186: RuntimeError
======================================================== short test summary info =========================================================
FAILED tests/models/embedding/vision_language/test_llava_next.py::test_models_text[half-royokong/e5-v] - RuntimeError: Calling torch.linal
g.cholesky on a CUDA tensor requires compiling PyTorch with MAGMA. Please use PyTorch built with MAG...
FAILED tests/models/embedding/vision_language/test_llava_next.py::test_models_image[half-royokong/e5-v] - RuntimeError: Calling torch.lina
lg.cholesky on a CUDA tensor requires compiling PyTorch with MAGMA. Please use PyTorch built with MAG...
====================================================== 2 failed in 87.58s (0:01:27) ======================================================
```

@Rocketknight1 I dig down a bit further. 

I tried using PyTorch from ROCm/pytroch: `docker pull rocm/pytorch:rocm6.3.4_ubuntu24.04_py3.12_pytorch_release_2.4.0`

The PyTorch in this Docker supports these `torch.linalg` functions. So, there is something wrong with the PyTorch found in the vLLM `Dockerfile.rocm_base`. I will look the direction to enable the build for these `torch.linalg` issue.

Thank you very much for your time @Rocketknight1 and @ivarflakstad . You guys are amazing! Responses are quick and helpful. 

Great!
For extra internet points perhaps you could open an issue in vllm an reference this issue? :)

> [@Rocketknight1](https://github.com/Rocketknight1) I dig down a bit further.
> 
> I tried using PyTorch from ROCm/pytroch: `docker pull rocm/pytorch:rocm6.3.4_ubuntu24.04_py3.12_pytorch_release_2.4.0`
> 
> The PyTorch in this Docker supports these `torch.linalg` functions. So, there is something wrong with the PyTorch found in the vLLM `Dockerfile.rocm_base`. I will look the direction to enable the build for these `torch.linalg` issue.
> 
> Thank you very much for your time [@Rocketknight1](https://github.com/Rocketknight1) and [@ivarflakstad](https://github.com/ivarflakstad) . You guys are amazing! Responses are quick and helpful.

Same issues here when trying to resize the tokenizer. Do we see a quick solution after 5 months later? :)

> Same issues here when trying to resize the tokenizer. Do we see a quick solution after 5 months later? :)

Is the issue in vllm for you as well?

> > Same issues here when trying to resize the tokenizer. Do we see a quick solution after 5 months later? :)
> 
> Is the issue in vllm for you as well?

Hi,

Our team @yanran-tang did a bit of exploration on this.

We are using a few different docker images to test whether this issue happened.

Firstly, according to some simple search, the issue seems to be about whether MAGMA is used when making the image.

And we tried:
1. An verl image made by AMD staff (@yushengsu-thu) at: https://verl.readthedocs.io/en/latest/amd_tutorial/amd_build_dockerfile_page.html#, or the docker image in: https://hub.docker.com/layers/yushengsuthu/verl/verl-0.4.1_ubuntu-22.04_rocm6.3.4-numa-patch_vllm0.8.5_sglang0.4.6.post4/images/sha256-a5faf39725f411576b5c7e760bba994803acafef3b1828e7d92a0b6711cd0f03 (**failed, no MAGMA is listed**)
2. The official verl at: https://hub.docker.com/layers/rocm/verl/verl-0.3.0.post0_rocm6.2_vllm0.6.3/images/sha256-cbe423803fd7850448b22444176bee06f4dcf22cd3c94c27732752d3a39b04b2 (**passed, MAGMA is listed in the image setup**)
3. The official pytorch (the one used by @tjtanaa above. Not totally the same image, pt 2.4 vs 2.6, but from the same uploader) at: https://hub.docker.com/layers/rocm/pytorch/rocm6.4.2_ubuntu24.04_py3.12_pytorch_release_2.6.0/images/sha256-6a287591500b4048a9556c1ecc92bc411fd3d552f6c8233bc399f18eb803e8d6 (**passed, MAGMA is listed in the image setup**)
4. The official vllm at: https://hub.docker.com/layers/rocm/vllm/latest/images/sha256-4a429705fa95a58f6d20aceab43b1b76fa769d57f32d5d28bd3f4e030e2a78ea (**failed, no MAGMA is listed**)

So we believe this issue could be in whether the following is used when making the image:
```
COPY file:d98119f5af1b9faf43e5508c7dff2c997ec9066336c0b7117755c698c92006fb in install_rocm_magma.sh
ENV MAGMA_HOME=/opt/rocm/magma
```

The native Pytorch should already support MAGMA: https://docs.pytorch.org/docs/stable/backends.html#torch.backends.cuda.preferred_linalg_library

@RuihongQiu  Thank you so much for summarizing the situation, and you are right on spot. In vLLM docker image, the Pytorch are built from scratch and it doesn't include the compilation of the MAGMA dependencies.