Apple M1 Max运行TensorFlow训练随机epoch卡顿如何排查

首先明确：你看到的Could not identify NUMA node of platform GPU ID 0是M系列芯片运行tensorflow-metal时的常规无害日志，Apple Silicon架构本身不支持NUMA节点，所有用metal加速的TF环境都会输出这条信息，和训练卡顿没有关联，不需要处理。
你遇到的随机epoch卡死、单step耗时突然飙升到数百秒的问题，是tensorflow-metal 0.5.0搭配tensorflow-macos 2.9.x版本的已知适配bug，不是环境配置错误，也不需要卸载tensorflow-metal放弃GPU加速。这个bug的触发原因是0.5.0版本的metal插件存在统一内存调度死锁问题，当训练过程中显存占用随迭代波动到阈值（你的模型包含多层全连接+Dropout，显存占用波动刚好容易触发这个问题），GPU运算队列就会挂死，表现就是训练进度完全停滞。

按优先级依次尝试，不需要重装整个miniforge环境：

• 版本降级到稳定适配组合（验证有效率最高）
  激活你的conda环境后执行以下命令，将相关包降到M1平台验证过的最稳定版本：

  pip install tensorflow-macos==2.8.0 tensorflow-metal==0.4.0
  

  该版本组合不会损失GPU训练速度，也不存在0.5.0版本的显存调度死锁问题，是目前M1 Max机型跑稠密全连接网络的首选配置。
• 代码侧规避（不想降版本时使用）
  两处代码修改即可绕开bug：
  1. 不要直接将全量训练数据转为tf.constant传入fit，改用tf.data.Dataset封装数据，避免全量数据常驻GPU显存触发调度异常：

     batch_size = 32
     train_ds = tf.data.Dataset.from_tensor_slices((X_train, Y_train))
     train_ds = train_ds.shuffle(1024).batch(batch_size).prefetch(tf.data.AUTOTUNE)
     history = model.fit(train_ds, epochs=100, use_multiprocessing=False)
     

  2. fit调用时传入use_multiprocessing=False，关闭0.5.0版本存在死锁问题的多进程运算调度逻辑。
• 显存上限限制（降版本后仍有偶发卡死时补充使用）
  在训练代码最开头、定义模型前加入以下配置，限制TensorFlow可使用的GPU统一内存上限，64G内存的M1 Max建议设置为40G以内，留足系统运行内存避免显存抢占导致的挂起：

  gpu_devices = tf.config.list_physical_devices('GPU')
  if gpu_devices:
      tf.config.set_logical_device_configuration(
          gpu_devices[0],
          [tf.config.LogicalDeviceConfiguration(memory_limit=40960)]
      )
  

如果以上方案操作后仍存在卡顿，导出你miniforge环境的YAML配置检查即可，重点排查是否存在conda源和pip源混装TensorFlow依赖的问题——如果conda安装了部分TensorFlow底层依赖，pip又单独安装了tensorflow-macos和tensorflow-metal，会出现依赖版本不匹配，也会导致训练随机挂死。

网上流传的卸载tensorflow-metal改用CPU运行的方案确实能解决卡顿，但会损失70%以上的训练速度，完全没有必要，通过版本对齐即可在保留GPU加速的前提下解决问题。

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