EC2 AMI新实例首次加载PyTorch模型耗时过长问题咨询

• EBS卷冷初始化：自定义AMI对应的EBS快照，新实例挂载后首次访问未读取过的数据块时，需要从AWS后端的快照存储拉取数据，大体积的模型权重文件首次读取会严重拖慢速度，后续读取从本地EBS缓存获取，因此速度恢复正常。
• PyTorch CUDA算子编译缓存未保留：PyTorch运行时会对用到的CUDA算子做即时编译（JIT），编译产物默认存在~/.cache/torch/kernels、~/.cache/torch/transformers（若用到HuggingFace相关组件）、~/.nv/ComputeCache路径下，制作AMI时未将这些缓存目录打包进镜像，新实例首次运行会重新编译所有用到的CUDA算子，耗时极长。
• 第三方库字节码缓存不全：仅保留了自有业务代码的__pycache__，torch、numpy等第三方库的字节码首次访问时也需要生成，不过这部分耗时占比极低。

• 制作AMI前预热全量数据和缓存
  生成AMI前在母实例执行以下操作，确保所有需要用到的资源都提前生成、所有数据块都被读取过：
  1. 完整运行一次业务代码，确保CUDA算子编译完成、模型权重全量读取到磁盘缓存
  2. 固定缓存路径避免权限问题：将~/.cache/torch、~/.nv目录移动到全局可访问路径比如/opt/下，在systemd服务配置中提前配置环境变量TORCH_CACHE_DIR=/opt/torch_cache、CUDA_CACHE_PATH=/opt/nv_cache，确保新实例启动时直接读取预生成的缓存
  3. 制作AMI前执行sudo sync确保所有磁盘写入落地
• 开启EBS快速快照还原（FSR）
  对制作AMI的源EBS快照开启快速快照还原，可直接消除EBS冷启动的块读取延迟，适合对伸缩启动速度要求高的场景。
• 优化模型加载逻辑
  调整加载参数减少IO开销，参考优化后的代码：

  # Import libraries
  import torch
  import numpy as np
  import os
  
  # 提前指定和AMI制作阶段一致的缓存路径
  os.environ['TORCH_CACHE_DIR'] = '/opt/torch_cache'
  os.environ['CUDA_CACHE_PATH'] = '/opt/nv_cache'
  
  # Import personal models
  from model1 import model1
  from model2 import model2
  
  # 初始化模型
  model1_object = model1()
  model2_object = model2()
  
  # 开启内存映射加载，跳过CPU到GPU的额外拷贝，减少磁盘IO耗时
  model1_object.load_state_dict(torch.load("model1.pth", map_location="cuda", mmap=True))
  model2_object.load_state_dict(torch.load("model2.pth", map_location="cuda", mmap=True))
  

• 调整systemd启动时机
  在systemd服务配置中添加依赖，等待EBS卷挂载完成、NVIDIA驱动初始化完成后再启动业务代码，避免依赖未就绪导致的额外耗时。

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