[RLlib] GPU Memory Leak? Tune + PPO, Policy Server + Client

mannyv · May 24, 2023, 9:07pm

Hi @Denys_Ashikhin and @kourosh,

Here are two posts that might be relevant here.

github.com/ray-project/ray

[RLlib] tf2 framework slight memory leak

opened 03:38AM - 09 Jun 22 UTC

jjyyxx

bug P1 rllib rllib-system

### What happened + What you expected to happen I was hoping that an RLlib tr…aining script could hold a roughly constant amount of memory after at most 1M environment interactions. But the memory usage grows unboundedly, which prohibits long-running training, especially using `ray.tune` to train multiple experiments concurrently. The memory leak happens with `'framework': 'tf2'`. `'framework': 'tf'` seems fine, but our custom algorithm is targeting TensorFlow 2 for the "tracing" feature. The memory leak rate is around 1% RAM every 5M environment interactions (on a 64G RAM machine). It should be noted that we use the nightly Ray version because the fixes in https://github.com/ray-project/ray/pull/15412 are not released yet. Before this fix (e.g. Ray v1.12.1) leaks at a much faster rate of 1% RAM every 1M environment interactions. We tried to debug with the help of `MemoryTrackingCallbacks`, but shows no obvious leaks (This is different from before https://github.com/ray-project/ray/pull/15412, where the callback shows memory allocated for global_timestep is constantly increasing). ![image](https://user-images.githubusercontent.com/13808089/172759068-4cb1859b-dae4-4377-9ca1-8a92d7b54ab2.png) ### Versions / Dependencies Ray: Daily Release at 2022-06-09T02:28:12.704Z from 'https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-2.0.0.dev0-cp39-cp39-manylinux2014_x86_64.whl' TensorFlow: 2.9.1 CPU Python: Python 3.9.13 OS: Ubuntu 18.04.2 LTS ### Reproduction script #### Usage ```bash # Installation conda create -n rllib-leak -c conda-forge --override-channels python=3.9 conda activate rllib-leak pip install tensorflow 'ray[rllib] @ https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-2.0.0.dev0-cp39-cp39-manylinux2014_x86_64.whl' # Run python ppo_leak.py ``` Then TensorBoard's `ray/tune/perf/ram_util_percent` or `ray/tune/custom_metrics/tracemalloc/worker/rss_mean` (if MemoryTrackingCallbacks is used) shows constant memory increment. #### Dummy Env (dummy_env.py) ```python from gym import Env from gym.spaces import Box from gym.wrappers import TimeLimit import numpy as np class DummyEnv(Env): observation_space = Box(-np.inf, np.inf, (70,)) action_space = Box(-1.0, 1.0, (12,)) _dummy_obs = np.random.rand(70) def step(self, action): return DummyEnv._dummy_obs, 0.0, False, {} def reset(self): return DummyEnv._dummy_obs def env_creator(config): return TimeLimit(DummyEnv(), max_episode_steps=1000) ``` #### Script (ppo_leak.py) ```python from ray.rllib.agents.callbacks import MemoryTrackingCallbacks from ray.rllib.agents.ppo import PPOTrainer import ray.tune.registry import dummy_env ray.tune.registry.register_env('dummy', dummy_env.env_creator) agent_config = { # 'callbacks': MemoryTrackingCallbacks, # Env config 'env': 'dummy', # Worker config 'framework': 'tf2', 'eager_tracing': True, 'num_workers': 10, 'num_envs_per_worker': 4, 'num_gpus': 0, 'horizon': 1000, # PPO config 'rollout_fragment_length': 1000, 'train_batch_size': 40000, 'sgd_minibatch_size': 40000, 'num_sgd_iter': 30, 'lr': 3e-4, 'lambda': 0.97, 'gamma': 0.99, 'no_done_at_end': True, # Model config 'model': { '_use_default_native_models': False, # Both True/False leaks at same rate 'fcnet_hiddens': [256, 256], 'fcnet_activation': 'relu', 'vf_share_layers': False, 'free_log_std': True, }, "seed": 1, } trainer = PPOTrainer(agent_config) while True: trainer.train() ``` ### Issue Severity High: It blocks me from completing my task.

Topic		Replies	Views
Memory Leak when training PPO on a single agent environment RLlib	15	1524	December 24, 2022
PPO with PyTorch GPU has a RAM memory leak for Ray 1.6.0 RLlib	5	649	October 5, 2021
Expected RAM usage for PPOTrainer (debugging memory leaks) RLlib	10	887	September 15, 2022
Running out of CUDA Memory - Sample Batch , rollout fragment length is not helping RLlib	2	895	July 28, 2022
Memory leak with PPO and GNN custom model RLlib	2	616	September 17, 2021

[RLlib] GPU Memory Leak? Tune + PPO, Policy Server + Client

Related topics