Sagemaker部署PyTorch模型GPU推理时卡在模型调用阶段

1. TorchScript模型GPU兼容性问题 #

• 问题根源：在CPU环境下trace生成的TorchScript模型，未适配GPU运行逻辑，直接在GPU环境加载时可能触发静默崩溃（TorchScript底层错误不会抛出Python可捕获异常，直接终止进程）。
• 解决办法：
  • 在GPU环境重新trace模型，确保trace时模型和输入都在CUDA设备上：

    import torch
    model = YourModel().to("cuda")
    dummy_input = torch.randn(1, 3, 224, 224).to("cuda")  # 匹配实际输入维度
    traced_model = torch.jit.trace(model, dummy_input)
    traced_model.save("traced_model_gpu.pt")
    

  • 部署前在GPU环境本地测试：加载模型并运行推理，确认无崩溃情况。

2. 自定义Transform函数设备不匹配 #

• 问题根源：预处理后的张量仍留在CPU，传入GPU模型时触发底层设备不兼容错误，且该错误未被Python层捕获，导致进程挂掉。
• 解决办法：
  • 修改transform函数，强制将输入张量转移到CUDA：

    def transform_fn(model, request_body, content_type, accept_type):
        # 原有预处理逻辑（解析、转换为张量等）
        input_tensor = ...
        input_tensor = input_tensor.to("cuda")  # 新增设备转移
        with torch.no_grad():
            output = model(input_tensor)
        # 后续输出处理逻辑
    

  • 检查预处理流程中所有张量的设备，避免CPU张量流入GPU模型。

3. 模型加载未显式指定GPU #

• 问题根源：SageMaker PyTorch推理容器默认可能不会自动将TorchScript模型转移到GPU，导致推理时设备不匹配。
• 解决办法：
  • 添加自定义模型加载脚本model.py，放在模型归档根目录：

    import torch
    
    def model_fn(model_dir):
        model = torch.jit.load(f"{model_dir}/traced_model.pt")
        model = model.to("cuda")  # 显式转移到GPU
        model.eval()
        return model
    

  • 重新打包模型并部署，容器会自动执行该脚本加载模型。

4. GPU显存不足 #

• 问题根源：ml.g4dn.xlarge的GPU显存为16GB，若模型或输入批量过大，可能触发OOM，但部分情况下不会输出明确日志，直接终止进程。
• 解决办法：
  • 减小请求的批量大小，测试是否能正常推理。
  • 在本地GPU环境用torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()监控显存使用，确认是否存在OOM。
  • 若显存确实不足，升级实例类型（如ml.g4dn.2xlarge）或对模型做量化优化。

5. 日志捕获不全 #

• 问题根源：底层CUDA或PyTorch错误可能输出到stderr，未被CloudWatch日志系统收集，即使设置log_level=DEBUG也无法看到。
• 解决办法：
  • 自定义容器启动脚本，将stderr重定向到stdout：修改启动命令为python -m sagemaker_pytorch_serving_container.serve 2>&1，确保所有日志被捕获。
  • 直接登录到SageMaker实例，查看/var/log/sagemaker/目录下的本地日志，可能找到更详细的错误栈。

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