Apple M1 Pro上AWS SageMaker本地训练/调试的正确配置

问题背景 #

计划在AWS SageMaker上运行RL训练任务，先在Apple M1 Pro设备上通过SageMaker本地模式做测试，再提交到云端付费实例。此前本地训练用TensorFlow-metal和TensorFlow-macos，但SageMaker的framework_version无法指定这些版本，且不确定适用于Linux GPU的local_gpu参数是否适配Apple Silicon架构，运行代码时卡在登录步骤，急需适配M1 Pro的配置、Docker镜像或示例。

当前尝试的代码 #

## main.py
import boto3
import sagemaker
import os
import keras
import numpy as np
from keras.datasets import fashion_mnist
from sagemaker.tensorflow import TensorFlow

sess = sagemaker.Session() 
# role = '<arn:aws:iam::0000000000000:role/CFN-SM-IM-Lambda-Catalog-sk-SageMakerExecutionRole-BlaBlaBla>' # 替换为你的IAM角色

(x_train, y_train), (x_val, y_val) = fashion_mnist.load_data()

os.makedirs("./data", exist_ok = True)

np.savez('./data/training', image=x_train, label=y_train)
np.savez('./data/validation', image=x_val, label=y_val)

# 本地数据路径，也可以用S3 URI
training_input_path   = 'file://data/training.npz'
validation_input_path = 'file://data/validation.npz'

# 本地模型输出路径，也可以用S3 URI
output_path           = 'file:///tmp/model/'

tf_estimator = TensorFlow(entry_point='mnist_keras_tf.py',
                          role=role,
                          instance_count=1, 
                          instance_type='local_gpu',   # 原配置尝试使用本地GPU
                          framework_version='2.1.0',
                          py_version='py3',
                          hyperparameters={'epochs': 1},
                          output_path=output_path
                         )

tf_estimator.fit({'training': training_input_path, 'validation': validation_input_path})

登录命令 #

aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com

适配方案建议 #

1. 关于local_gpu参数的兼容性 #

SageMaker官方的local_gpu参数是针对Linux系统的CUDA GPU设计的，Apple M1 Pro采用ARM64架构+Metal加速框架，官方镜像不支持该参数，因此不能直接使用local_gpu，需要自定义适配ARM64的Docker镜像。

2. 构建适配M1 Pro的自定义Docker镜像 #

基于Apple官方提供的TensorFlow-macos ARM64镜像，添加SageMaker训练工具依赖，示例Dockerfile如下：

# 基于Apple官方TensorFlow-macos ARM64镜像
FROM apple/tensorflow:2.15.0-macos-arm64

# 安装SageMaker训练工具包
RUN pip install sagemaker-training-toolkit

# 安装TensorFlow-metal加速插件（版本需与TensorFlow对应）
RUN pip install tensorflow-metal==1.1.0

构建镜像命令：

docker build -t m1-sagemaker-tf:latest .

3. 修改SageMaker代码配置 #

放弃官方TensorFlow Estimator，改用通用Estimator指定自定义镜像，实例类型改用local（自定义镜像已集成Metal加速），修改后的代码片段：

from sagemaker.estimator import Estimator

# 指定自定义构建的镜像标签
custom_image = 'm1-sagemaker-tf:latest'

estimator = Estimator(
    image_uri=custom_image,
    role=role,
    instance_count=1,
    instance_type='local',
    hyperparameters={'epochs': 1},
    output_path=output_path
)

estimator.fit({'training': training_input_path, 'validation': validation_input_path})

4. 解决登录卡住问题 #

官方ECR的763104351884仓库中的镜像均为AMD64架构，在M1 Pro上拉取时需通过Rosetta转译，容易出现卡顿或失败。使用上述自定义ARM64镜像可跳过这一步，直接使用本地镜像运行训练。

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