Mixtral OutOfMemoryError with 2 GPUs

[thomasgauthier]

I'm trying to run Mixtral (Mixtral Hermes) with two 48GB GPUs but it seems that sglang server is not using my second GPU.

CUDA_VISIBLE_DEVICES="0,1" python -m sglang.launch_server --model-path /workspace/model --port 30000

errors out with

Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Process Process-1:
Traceback (most recent call last):
router init state: Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/manager.py", line 68, in start_router_process
    model_client = ModelRpcClient(server_args, port_args)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_rpc.py", line 448, in __init__
    self.model_server.exposed_init_model(0, server_args, port_args)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_rpc.py", line 54, in exposed_init_model
    self.model_runner = ModelRunner(
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_runner.py", line 229, in __init__
    self.load_model()
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_runner.py", line 272, in load_model
    model = model_class(
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/models/mixtral.py", line 322, in __init__
    self.model = MixtralModel(config, linear_method)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/models/mixtral.py", line 285, in __init__
    [
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/models/mixtral.py", line 286, in <listcomp>
    MixtralDecoderLayer(config, i, linear_method=linear_method)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/models/mixtral.py", line 238, in __init__
    self.block_sparse_moe = MixtralMoE(config=config, linear_method=linear_method)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/models/mixtral.py", line 99, in __init__
    [
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/models/mixtral.py", line 100, in <listcomp>
    MixtralMLP(
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/models/mixtral.py", line 55, in __init__
    self.w2 = ReplicatedLinear(
  File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/layers/linear.py", line 108, in __init__
    self.linear_weights = self.linear_method.create_weights(
  File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/layers/linear.py", line 55, in create_weights
    weight = Parameter(torch.empty(output_size_per_partition,
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_device.py", line 77, in __torch_function__
    return func(*args, **kwargs)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 112.00 MiB. GPU 0 has a total capacty of 44.35 GiB of which 9.38 MiB is free. Process 2776126 has 44.33 GiB memory in use. Of the allocated memory 44.02 GiB is allocated by PyTorch, and 14.25 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

detoken init state: init ok

I had the same error on Modal with 2 x 80GB GPUs.

Thanks for the support

[hnyls2002]

You have to add arguments --tp-size 2 to indicate the server to run the tensor parallelism on 2 GPUs.

[thomasgauthier]

Ok thank you that seems to get the second GPU to work. However I still get an error :

python -m sglang.launch_server --tp-size 2 --model-path /workspace/model --port 30000
server started on [0.0.0.0]:10004
server started on [0.0.0.0]:10005
accepted ('127.0.0.1', 47232) with fd 28
welcome ('127.0.0.1', 47232)
accepted ('127.0.0.1', 48754) with fd 24
welcome ('127.0.0.1', 48754)
router init state: Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/manager.py", line 68, in start_router_process
    model_client = ModelRpcClient(server_args, port_args)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_rpc.py", line 469, in __init__
    rets = [obtain(x) for x in executor.map(init_model, range(tp_size))]
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_rpc.py", line 469, in <listcomp>
    rets = [obtain(x) for x in executor.map(init_model, range(tp_size))]
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 621, in result_iterator
    yield _result_or_cancel(fs.pop())
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 319, in _result_or_cancel
    return fut.result(timeout)
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 458, in result
    return self.__get_result()
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result
    raise self._exception
  File "/usr/lib/python3.10/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_rpc.py", line 467, in init_model
    return self.model_servers[i].init_model(i, server_args, port_args)
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/netref.py", line 240, in __call__
    return syncreq(_self, consts.HANDLE_CALL, args, kwargs)
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/netref.py", line 63, in syncreq
    return conn.sync_request(handler, proxy, *args)
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/protocol.py", line 718, in sync_request
    return _async_res.value
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/async_.py", line 108, in value
    raise self._obj
_get_exception_class.<locals>.Derived: Trying to create tensor with negative dimension -789: [-789]

========= Remote Traceback (1) =========
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/protocol.py", line 359, in _dispatch_request
    res = self._HANDLERS[handler](self, *args)
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/protocol.py", line 837, in _handle_call
    return obj(*args, **dict(kwargs))
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_rpc.py", line 54, in exposed_init_model
    self.model_runner = ModelRunner(
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_runner.py", line 230, in __init__
    self.init_memory_pool(total_gpu_memory)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_runner.py", line 300, in init_memory_pool
    self.req_to_token_pool = ReqToTokenPool(
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/memory_pool.py", line 11, in __init__
    self.mem_state = torch.ones((size,), dtype=torch.bool, device="cuda")
RuntimeError: Trying to create tensor with negative dimension -789: [-789]


detoken init state: init ok
goodbye ('127.0.0.1', 47232)
goodbye ('127.0.0.1', 48754)

I'm not sure if it's related to multi GPU or not

[hnyls2002]

Ok thank you that seems to get the second GPU to work. However I still get an error :

python -m sglang.launch_server --tp-size 2 --model-path /workspace/model --port 30000
server started on [0.0.0.0]:10004
server started on [0.0.0.0]:10005
accepted ('127.0.0.1', 47232) with fd 28
welcome ('127.0.0.1', 47232)
accepted ('127.0.0.1', 48754) with fd 24
welcome ('127.0.0.1', 48754)
router init state: Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/manager.py", line 68, in start_router_process
    model_client = ModelRpcClient(server_args, port_args)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_rpc.py", line 469, in __init__
    rets = [obtain(x) for x in executor.map(init_model, range(tp_size))]
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_rpc.py", line 469, in <listcomp>
    rets = [obtain(x) for x in executor.map(init_model, range(tp_size))]
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 621, in result_iterator
    yield _result_or_cancel(fs.pop())
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 319, in _result_or_cancel
    return fut.result(timeout)
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 458, in result
    return self.__get_result()
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result
    raise self._exception
  File "/usr/lib/python3.10/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_rpc.py", line 467, in init_model
    return self.model_servers[i].init_model(i, server_args, port_args)
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/netref.py", line 240, in __call__
    return syncreq(_self, consts.HANDLE_CALL, args, kwargs)
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/netref.py", line 63, in syncreq
    return conn.sync_request(handler, proxy, *args)
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/protocol.py", line 718, in sync_request
    return _async_res.value
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/async_.py", line 108, in value
    raise self._obj
_get_exception_class.<locals>.Derived: Trying to create tensor with negative dimension -789: [-789]

========= Remote Traceback (1) =========
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/protocol.py", line 359, in _dispatch_request
    res = self._HANDLERS[handler](self, *args)
  File "/usr/local/lib/python3.10/dist-packages/rpyc/core/protocol.py", line 837, in _handle_call
    return obj(*args, **dict(kwargs))
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_rpc.py", line 54, in exposed_init_model
    self.model_runner = ModelRunner(
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_runner.py", line 230, in __init__
    self.init_memory_pool(total_gpu_memory)
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/managers/router/model_runner.py", line 300, in init_memory_pool
    self.req_to_token_pool = ReqToTokenPool(
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/memory_pool.py", line 11, in __init__
    self.mem_state = torch.ones((size,), dtype=torch.bool, device="cuda")
RuntimeError: Trying to create tensor with negative dimension -789: [-789]


detoken init state: init ok
goodbye ('127.0.0.1', 47232)
goodbye ('127.0.0.1', 48754)

I'm not sure if it's related to multi GPU or not

When meeting this error, what is the machine? I think 48G x 2 is not enough for a 8 x 7B model,

[thomasgauthier]

Yes thank you that seems to be the issue. Got it working on modal with 2x80GB and --tp-size 2

code

import sys

from modal import Image, Stub, gpu, method


MODEL_DIR = "/model"
BASE_MODEL = "NousResearch/Nous-Hermes-2-Mixtral-8x7B-SFT"

GPU_CONFIG = gpu.A100(memory=80, count=2)

def download_model_to_folder():
    import os
    from huggingface_hub import snapshot_download
    from transformers.utils import move_cache

    os.makedirs(MODEL_DIR, exist_ok=True)

    snapshot_download(
        BASE_MODEL,
        local_dir=MODEL_DIR,
        ignore_patterns="*.pt",  # Using safetensors
    )
    move_cache()



app_image = (
    Image.micromamba()
    .apt_install("build-essential")
    .apt_install("procps")    
    .micromamba_install(
        "cudatoolkit=11.8",
        "cuda-nvcc",
        channels=["conda-forge", "nvidia"],
    )
    .pip_install(
        "transformers",
        "hf_transfer",
        "huggingface_hub[hf_transfer]",
        "sglang[all]",
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
    .run_function(download_model_to_folder)
)

stub = Stub("sglang-mixtral")

@stub.function(
    gpu=GPU_CONFIG,
    timeout=60 * 10,
    container_idle_timeout=60 * 10,
    allow_concurrent_inputs=10,
    image=app_image,
)
def test():
    # taken from https://modal.com/docs/examples/vllm_mixtral
    if GPU_CONFIG.count > 1:
        # Patch issue from https://github.com/vllm-project/vllm/issues/1116
        import ray

        ray.shutdown()
        ray.init(num_gpus=GPU_CONFIG.count)

    from sglang import function, system, user, assistant, gen, set_default_backend, Runtime

    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained('NousResearch/Nous-Hermes-2-Mixtral-8x7B-SFT', use_fast=True)

    # taken from https://modal.com/docs/examples/vllm_mixtral
    if GPU_CONFIG.count > 1:
        import subprocess

        RAY_CORE_PIN_OVERRIDE = "cpuid=0 ; for pid in $(ps xo '%p %c' | grep ray:: | awk '{print $1;}') ; do taskset -cp $cpuid $pid ; cpuid=$(($cpuid + 1)) ; done"
        subprocess.call(RAY_CORE_PIN_OVERRIDE, shell=True)

    runtime = Runtime(model_path=MODEL_DIR, tp_size=2)
    set_default_backend(runtime)

    @function
    def text_qa(s, question):
        s += "Q: " + question + "\n"
        s += "A:" + gen("answer", stop="\n")

    states = text_qa.run_batch(
        [
            {"question": "What is the capital of the United Kingdom?"},
            {"question": "What is the capital of France?"},
            {"question": "What is the capital of Japan?"},
        ],
        progress_bar=True
    )

    print(states)
    runtime.shutdown()

@stub.local_entrypoint()
def main():
    test.remote()