fix OOM in BlockManager

[zhyncs]

Motivation

fix OOM in BlockManager

Modification

replace total * ratio with free * ratio in GetBlockCount, exclude the weights memory usage

Checklist

1. Pre-commit or other linting tools are used to fix the potential lint issues.
2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
4. The documentation has been modified accordingly, like docstring or example tutorials.

[zhyncs]

Hi @lzhangzz @lvhan028 May you help review this pr

[lzhangzz]

We need to remind the users when specifying the value by ratio, care must be taken for TP>1 (different ranks may have different free memory sizes).

Any idea on where to add the reminder in the docs ? @lvhan028

[lvhan028]

We need to remind the users when specifying the value by ratio, care must be taken for TP>1 (different ranks may have different free memory sizes).

Any idea on where to add the reminder in the docs ? @lvhan028

I suggest adding some advanced usages in https://github.com/InternLM/lmdeploy/blob/main/docs/en/inference/pipeline.md, showing how to set cache_max_entry_count. Meanwhile, we can update the FAQ for the OOM issue.

[zhyncs]

We need to remind the users when specifying the value by ratio, care must be taken for TP>1 (different ranks may have different free memory sizes).

Any idea on where to add the reminder in the docs ? @lvhan028

Hi @lzhangzz Usually when using a single machine with multiple GPUs for tensor parallelism, hardware specifications are the same, and the weights are also equally divided. May we assume that in most scenarios the free memory of multiple ranks are the same?

[zhyncs]

We need to remind the users when specifying the value by ratio, care must be taken for TP>1 (different ranks may have different free memory sizes).
Any idea on where to add the reminder in the docs ? @lvhan028

I suggest adding some advanced usages in https://github.com/InternLM/lmdeploy/blob/main/docs/en/inference/pipeline.md, showing how to set cache_max_entry_count. Meanwhile, we can update the FAQ for the OOM issue.

Hi @lvhan028 I agree that the usage documentation should also be updated. May this fix be considered to be merged first?

[lvhan028]

Hi, @zhyncs
I still have some concerns.
In the tensor parallelism case, if some specified GPUs are occupied by other programs or zombie processes, free memory will differ among the GPUs. It means the number of k/v cache blocks in each GPU is different, which will bring catastrophe during inference.

[zhyncs]

Hi @lzhangzz @lvhan028 In the latest commit, I have updated the free * ratio with total * ratio - (total - free).

I agreed with @lvhan028's case.

And if we want to completely solve the OOM issue and potential inconsistencies with multi-GPU's k/v cache blocks during TP, we may execute a forward pass with dummy inputs to profile the memory usage of the model like vLLM's profile_run.

[lvhan028]

Hi, @zhyncs
total * ratio - (total - free) might be negative.
I prefer to compute the k/v cache memory size on python side, not C++/cuda side.
@lzhangzz any comments?

[zhyncs]

Hi @lvhan028
I agree that total * ratio - (total - free) might be negative. And we may add assertion for this case. The program should report the error to the user and gracefully exit, prompting the user to provide a correct value for the configuration. Guiding the user through the correct configuration will help prevent any further issues or misunderstandings.
Anyway, the current total * ratio is problematic.
Controlling these parameters and checksums on the Python side is ok.

[zhyncs]

We need to remind the users when specifying the value by ratio, care must be taken for TP>1 (different ranks may have different free memory sizes).
Any idea on where to add the reminder in the docs ? @lvhan028

I suggest adding some advanced usages in https://github.com/InternLM/lmdeploy/blob/main/docs/en/inference/pipeline.md, showing how to set cache_max_entry_count. Meanwhile, we can update the FAQ for the OOM issue.

Hi @lvhan028 @lzhangzz

If it is more difficult for us to decide whether to go with this temporary fix or a more sophisticated one, such as profile_run similar to vLLM, then we can now go ahead and update the docs to temporarily bypass the OOM issue by configuring cache_max_entry_count when the user encounters it, even though there is a certain mind-burden associated with this configuration.

TurbomindEngineConfig cache_max_entry_count
  ->LlamaTritonModel EngineParams cache_max_block_count
    ->LlamaV2 EngineParams cache_max_block_count
      ->LlamaBatch EngineParams cache_max_block_count
        ->SequenceManager block_count
          ->BlockManager block_count
            ->GetBlockCount ratio

[lvhan028]

I agree. But should it be total * ratio <= free?

[zhyncs]

There's no need to have this restriction, it doesn't matter.
e.g. A100 80G, Llama 2 13b Chat, ratio 0.9
total * ratio =80*0.9=72, free =80-26=54, total * ratio - (total - free) =46 (weights: 26g, kv cache: 46g)
In this case, total * ratio is not less than free but it's ok.
We may take total * ratio - (total - free) as a whole.

[lvhan028]

Currently, the ratio definition is that the size of ratio*total memory is allocated for the k/v cache. So, when OOM happens, we suggest decreasing the ratio.
However, the equation total * ratio - (total - free) = free - (1-ratio)*total indicates an increasing ratio when memory is run out.
That's why I suggest FT_CHECK_WITH_INFO(total * ratio < free, ...)

[zhyncs]

Maybe we can consider revising the definition of ratio. Using the ratio*total-sized memory for k/v cache can impose a significant mental burden on users. For example, when running Llama 2 7b chat fp16 on an A30 24g device, the weights occupy 14g, and the default ratio parameter is set to 0.5. If not configured, it could result in OOM. Even when configured manually, the user would need to subtract the weights from the total memory, make estimations, and perform calculations to determine an appropriate ratio. It's not user-friendly.
The ratio can convey the following meanings, eliminating the need for users to consider model weights.

weights: total - free
k/v cache: total * ratio - (total - free)
other: (1 - ratio) * total

[zhyncs]

Maybe the ratio definition could be the size of ratio*total - MODEL_WEIGHTS memory is allocated for the k/v cache.

[lzhangzz]

With kv = total * ratio - (total - free), the semantic of ratio becomes ratio = kv% + used%, which is the ratio of kv cache plus the ratio of weight.

However, it's inconsistent with the option name cache_max_entry_count, which suggesting we are configuring the amount kv-cache blocks.

And it does not solve the possible inconsistency in TP mode. Anything depends on free suffers the same problem.

I suggest we still use kv = free * ratio as it's consistent with the current option.

To deal with the inconsistency, we may

1. Find the mininal free of all ranks, which requires additional communication and synchronization.
2. Do nothing for now and flag a warning that the user should check for consistency.

[zhyncs]

With kv = total * ratio - (total - free), the semantic of ratio becomes ratio = kv% + used%, which is the ratio of kv cache plus the ratio of weight.

However, it's inconsistent with the option name cache_max_entry_count, which suggesting we are configuring the amount kv-cache blocks.

And it does not solve the possible inconsistency in TP mode. Anything depends on free suffers the same problem.

I suggest we still use kv = free * ratio as it's consistent with the current option.

To deal with the inconsistency, we may

1. Find the mininal free of all ranks, which requires additional communication and synchronization.
2. Do nothing for now and flag a warning that the user should check for consistency.

@lzhangzz Ok I agree. The first commit is just free * ratio. In order to use cache_max_entry_count as before, I will revert to this and add warning for the TP consistency. And I think this check is suitable before initializing AsyncEngine. @AllentDan Do you have any suggestion?

[zhyncs]

from py3nvml.py3nvml import (
    nvmlInit,
    nvmlDeviceGetCount,
    nvmlDeviceGetHandleByIndex,
    nvmlDeviceGetMemoryInfo,
    nvmlShutdown,
)
import os


def get_device_ids():
    devices = os.getenv('CUDA_VISIBLE_DEVICES', '')
    return list(map(int, devices.split(','))) if devices else []


def compare_individual_gpu_memory(device_ids):
    try:
        nvmlInit()
        device_count = nvmlDeviceGetCount()
        total_mem = []
        free_mem = []

        for i in range(device_count):
            if device_ids and i not in device_ids:
                continue
            handle = nvmlDeviceGetHandleByIndex(i)
            mem_info = nvmlDeviceGetMemoryInfo(handle)
            total_mem.append(mem_info.total)
            free_mem.append(mem_info.free)

        all_total_equal = total_mem.count(total_mem[0]) == len(total_mem)
        all_free_equal = free_mem.count(free_mem[0]) == len(free_mem)

        print(f"Total memory equal across GPUs: {all_total_equal}")
        print(f"Free memory equal across GPUs: {all_free_equal}")

        nvmlShutdown()

    except Exception as e:
        print(f"An exception occurred: {e}")


device_ids = get_device_ids()
compare_individual_gpu_memory(device_ids)

It works when TP > 1. Some code snippets like this @AllentDan

[AllentDan]

It would be better that the check and logging a suggestion ratio could be done in Turbomind. Not only AsyncEngine used Turbomind.

[zhyncs]

It would be better that the check and logging a suggestion ratio could be done in Turbomind. Not only AsyncEngine used Turbomind.

ok

[AllentDan]

As for python package to get the memory information, I would suggest using torch.cuda.mem_get_info. This function requires no extra third party package.

[zhyncs]

Hi @lzhangzz @AllentDan May you help review the latest commit? Thanks.

[zhyncs]

Hi @lvhan028 May you help merge this pr? Thanks.

[lzhangzz]

fix lint issues please

[zhyncs]

fix lint issues please

ok

[lvhan028]

@zhyncs please resolve linting error

pip install pre-commit
cd lmdeploy
pre-commit install .
pre-commit run --all-files

[zhyncs]

@zhyncs please resolve linting error

pip install pre-commit
cd lmdeploy
pre-commit install .
pre-commit run --all-files

[username@hostname lmdeploy]# pre-commit run --all-files
flake8...................................................................Passed
isort....................................................................Passed
yapf.....................................................................Passed
trim trailing whitespace.................................................Passed
check yaml...............................................................Passed
fix end of files.........................................................Passed
fix requirements.txt.....................................................Passed
fix double quoted strings................................................Passed
check for merge conflicts................................................Passed
fix python encoding pragma...............................................Passed
mixed line ending........................................................Passed
mdformat.................................................................Passed
codespell................................................................Passed
docformatter.............................................................Passed
check copyright..........................................................Passed

done

[lzhangzz]

hold on, just found some problems in _compare_individual_gpu_memory

[zhyncs]

Hi @lzhangzz May you help review the latest commit? Thanks.