Pytorch MPS超参调优时触发'mps.scatter_nd'张量形状错误

问题背景 #

使用NeverGrad为SASRec模型搭建超参数调优流程时，首次训练可正常完成，但第二次训练启动时触发MPS设备报错：

MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":49:0)): error: 'mps.scatter_nd' op invalid input tensor shape: updates tensor shape and data tensor shape must match along inner dimensions
/AppleInternal/Library/BuildRoots/97f6331a-ba75-11ed-a4bc-863efbbaf80d/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphExecutable.mm:1710: failed assertion `Error: MLIR pass manager failed'

仅MPS设备出现该问题，CPU设备运行完全正常，模型本身无逻辑错误。

核心原因与解决方案 #

1. 清理MPS残留资源 #

MPS设备在训练结束后可能残留未释放的张量或计算图状态，导致下一次训练时张量形状不匹配。在main函数末尾添加资源清理逻辑：

# 清理MPS缓存
torch.mps.empty_cache()
# 将模型移回CPU，彻底释放MPS设备资源
model.to('cpu')

同时在objective函数中，每次调用main后也执行清理：

def objective(hyperparameters):
    # ... 原有参数赋值代码 ...
    hit_10 = main(args)
    # 训练后立即清理MPS资源
    torch.mps.empty_cache()
    return -hit_10

2. 避免全局args对象的状态污染 #

当前代码直接修改全局args对象的属性，不同trial间可能存在参数残留。改为创建args副本进行操作：

def objective(hyperparameters):
    # 创建args副本，隔离不同trial的参数状态
    trial_args = argparse.Namespace(**vars(args))
    
    trial_args.dropout_rate = hyperparameters["dropout_rate"]
    trial_args.maxlen = int(hyperparameters["maxlen"])
    trial_args.lr = hyperparameters["lr"]
    trial_args.hidden_units = hyperparameters["hidden_units"]

    n_heads_max = trial_args.hidden_units // 2
    var_heads = [i for i in range(1, n_heads_max + 1) if trial_args.hidden_units % i == 0]
    trial_args.num_heads = hyperparameters["num_heads"] % len(var_heads)
    trial_args.num_heads = var_heads[trial_args.num_heads]

    trial_args.num_blocks = int(hyperparameters["num_blocks"])

    hit_10 = main(trial_args)
    torch.mps.empty_cache()
    return -hit_10

3. 修复WarpSampler的多线程残留 #

WarpSampler的多线程可能在MPS设备下残留未终止的线程资源，修改main函数中的sampler关闭逻辑：

# 在main函数末尾
sampler.close()
# 等待所有线程完全终止
import time
time.sleep(1)
torch.mps.empty_cache()

4. 禁用WandB设备同步（可选） #

WandB的自动设备跟踪可能保留MPS张量引用，导致资源无法释放。初始化WandB时添加禁用配置：

wandb.init(
    project='SASRec_RND', 
    config=args, 
    settings=wandb.Settings(_disable_device_sync=True)
)

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