使用torch.compile时Amazon SageMaker训练任务失败求助

环境配置 #

• 实例：Amazon SageMaker 单台ml.g4dn.xlarge
• 依赖：PyTorch 2.0、Python 3.10
• 模型：ViT视觉Transformer

问题表现 #

未启用模型编译时，训练任务全程正常完成；添加torch.compile()后，训练流程已完整执行（已打印Training completed!，所有检查点、模型产物均成功存入指定S3存储桶），但任务最终失败，报错内容：

free(): invalid pointer
Aborted (core dumped)

简化训练脚本 #

#### Training script ###
model = get_model()

# 添加此行后任务会失败
# 训练完成后执行该语句
model = torch.compile(model)

# PyTorch Lightning训练逻辑
run_training()

# 标记训练脚本执行完毕
print('Training completed!')

额外信息 #

• 任务失败状态：任务状态显示为失败，错误类型标注为算法错误
• 训练日志：日志末尾输出上述free(): invalid pointer和Aborted (core dumped)的崩溃信息

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