在Amazon SageMaker部署StarCoder时遭遇CUDA内存不足问题

问题概述 #

使用ml.g5.8xlarge实例在Amazon SageMaker部署bigcode/starcoder模型时，出现两类错误：ShardCannotStart分片启动失败，以及torch.cuda.OutOfMemoryError显存不足。

错误日志 #

错误1：ShardCannotStart #

Error: ShardCannotStart
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 155, in serve
    asyncio.run(serve_inner(model_id, revision, sharded, quantize, trust_remote_code))
  File "/opt/conda/lib/python3.9/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 647, in run_until_complete
    return future.result()
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 124, in serve_inner
    model = get_model(model_id, revision, sharded, quantize, trust_remote_code)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/__init__.py", line 134, in get_model
    return santacoder_cls(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_santacoder.py", line 62, in __init__
    self.load_weights(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_santacoder.py", line 96, in load_weights
    value = value.to(device if quantize is None else "cpu").to(dtype)

错误2：CUDA显存不足 #

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 288.00 MiB (GPU 0; 22.20 GiB total capacity; 19.72 GiB already allocated; 143.12 MiB free; 21.11 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

原部署代码 #

import json
import sagemaker
import boto3
from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

# Hub Model configuration
hub = {
    'HF_MODEL_ID':'bigcode/starcoder',
    'SM_NUM_GPUS': json.dumps(1),
    'HF_API_TOKEN': '<TOKEN>'
}


# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    image_uri=get_huggingface_llm_image_uri("huggingface",version="0.8.2"),
    env=hub,
    role=role, 
)

# deploy model to SageMaker Inference
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.8xlarge",
    container_startup_health_check_timeout=400,
    endpoint_name="Starcoder"
   )
  
# send request
predictor.predict({
    "inputs": "def print_hello_world():",
})

解决方案 #

1. 启用模型量化 #

StarCoder是15B参数模型，全精度加载需约60GB显存，远超ml.g5.8xlarge单GPU的22.2GiB可用显存。通过启用4bit量化可将显存占用降至约7.5GB，满足实例配置要求。

在hub环境变量中添加量化配置：

'QUANTIZE': 'bitsandbytes'  # 启用4bit量化

2. 优化CUDA内存分配 #

添加内存碎片优化配置，减少显存碎片化导致的OOM：

'PYTORCH_CUDA_ALLOC_CONF': 'max_split_size_mb:512'

3. 升级HF LLM镜像版本（可选） #

旧版本镜像可能存在内存管理bug，建议升级至最新兼容版本（如1.1.0）：

image_uri=get_huggingface_llm_image_uri("huggingface",version="1.1.0")

修改后的完整部署代码 #

import json
import sagemaker
import boto3
from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

# Hub Model configuration
hub = {
    'HF_MODEL_ID':'bigcode/starcoder',
    'SM_NUM_GPUS': json.dumps(1),
    'HF_API_TOKEN': '<TOKEN>',
    'QUANTIZE': 'bitsandbytes',
    'PYTORCH_CUDA_ALLOC_CONF': 'max_split_size_mb:512'
}


# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    image_uri=get_huggingface_llm_image_uri("huggingface",version="1.1.0"),
    env=hub,
    role=role, 
)

# deploy model to SageMaker Inference
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.8xlarge",
    container_startup_health_check_timeout=400,
    endpoint_name="Starcoder"
   )
  
# send request
predictor.predict({
    "inputs": "def print_hello_world():",
})

说明 #

ShardCannotStart错误本质是模型加载时显存不足导致的启动失败，解决OOM问题后该错误会自动消失。

内容的提问来源于stack exchange，提问作者texnoob