Ray Tune训练模型时节点内存持续上涨导致试验运行报错

问题现象 #

使用Ray Tune做深度学习模型超参数搜索时，节点内存占用持续上涨最终导致试验运行报错，运行时状态输出如下：

== Status ==
Current time: 2022-06-16 13:27:43 (running for 00:09:14.60)
Memory usage on this node: 26.0/62.8 GiB
Using AsyncHyperBand: num_stopped=0
Bracket: Iter 4.000: None | Iter 2.000: None | Iter 1.000: None
Resources requested: 8.0/8 CPUs, 1.0/1 GPUs, 0.0/51.37 GiB heap, 0.0/0.93 GiB objects (0.0/1.0 accelerator_type:A40)
Result logdir: /app/ray_experiment_0
Number of trials: 3/20 (1 PENDING, 2 RUNNING)
+--------------+----------+----------------+-----------------+--------------+
| Trial name   | status   | loc            |   learning_rate |   batch_size |
|--------------|----------|----------------|-----------------+--------------|
| run_cf921dd8 | RUNNING  | 172.17.0.3:402 |       0.0374603 |           64 |
| run_d20c6f50 | RUNNING  | 172.17.0.3:437 |       0.0950719 |           64 |
| run_d20e37cc | PENDING  |                |       0.0732021 |           64 |
+--------------+----------+----------------+-----------------+--------------+

尝试设置ray.init(object_store_memory = 10**9)后问题没有解决，使用的核心逻辑为贝叶斯优化采样超参数+ASHAScheduler提前终止低表现试验，核心代码如下：

def grid_search(config):

    # 提前终止表现不佳的试验
    scheduler = ASHAScheduler(
        max_t=5,
        grace_period=1,
        reduction_factor=2)

    # 贝叶斯优化智能采样超参数
    algo = BayesOptSearch(random_search_steps=4, mode="min")

    reporter = CLIReporter(
        parameter_columns=["learning_rate", "batch_size"],
        metric_columns=["loss", "mean_accuracy", "training_iteration"])

    resources_per_trial = {"cpu": config["n_cpu_per_trials"], "gpu": config["n_gpu_per_trials"]}

    trainable = tune.with_parameters(run)

    analysis = tune.run(trainable,
        resources_per_trial=resources_per_trial,
        metric="loss",
        mode="min",
        config=config,
        num_samples=config["n_sampling"], # 超参数空间采样次数
        scheduler=scheduler,
        progress_reporter=reporter,
        name=config["name_experiment"],
        local_dir="/app/.",
        search_alg=algo)

    print("搜索到的最优超参数为: ", analysis.best_config)

可行解决方案 #

• 先排查训练函数自身的内存泄漏：80%以上的同类问题和Ray框架无关，是传入的run训练函数存在长期引用的未释放对象，比如每轮迭代把loss、预测张量、模型权重存在全局列表/类属性中未清理，或是深度学习框架的计算图缓存未清空。以PyTorch为例，可以在每轮验证结束后调用torch.cuda.empty_cache()，同时不要把不需要持久化的张量长期放在迭代作用域之外。
• 调整Tune默认的结果缓存配置：默认状态下Tune会将所有trial的全量指标、检查点都存在内存中供后续结果分析调用，试验样本量上来后会持续占用内存。可以在tune.run中增加配置：设置keep_checkpoints_num=1、checkpoint_score_attr="loss"，自动删除表现较差trial的检查点，仅保留最优trial的检查点；同时配置sync_config=ray.train.SyncConfig(sync_artifacts=False)，关闭不必要的工件内存缓存。
• 修正大对象传参逻辑：不要把初始化后的大模型、全量数据集这类大体积对象通过config参数传入Tune，这类对象会被每个trial重复序列化拷贝，产生大量冗余内存占用。大对象提前用ray.put存入对象存储，仅传对象引用给trial即可。
• 合理设置对象存储内存：之前设置的1GB对象存储内存远低于深度学习任务的常规需求，会导致Ray频繁将内存对象溢出转存到磁盘，转存过程的缓存进程反而会占满内存。一般建议将对象存储内存设置为节点总内存的20%-30%，比如示例中62.8GiB总内存的节点，设置为20GiB左右即可。
• 开启自动清理+主动垃圾回收：在tune.run中配置trial_graceful_cleanup=True，trial结束后自动清理关联的进程和对象；同时可以在训练函数每轮迭代结束后手动调用import gc; gc.collect()，及时清理没有引用的残余对象。
• 升级Ray版本：旧版本（2.0之前）的BayesOptSearch实现存在已知的内存泄漏问题，会全量缓存所有历史trial的观测结果不做裁剪，升级到2.0及以上的稳定版本即可修复这类框架层面的内存问题。

调试时可以在运行过程中执行ray memory命令，查看对象存储中存活对象的引用来源，快速定位是训练逻辑中的用户态对象泄漏，还是框架本身的对象残留问题。

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