JAX runs slow on GPU compared to CPU

[chkda]

Hi folks,
I have a piece of code that I am trying to train an agent. I dont know if this is the right place for a question like this. The code runs much faster on my local CPU machine compared to V100 on Colab. I am not sure if I am doing something wrong or the code has some memory bottlenecks. My code is pasted below. Any help will be appreciated.

import jax
import jax.numpy as jnp
import flax
import flax.linen as nn
import optax
from flax.training.train_state import TrainState
from brax import envs
from brax.io import html
from brax.envs.wrappers.training import wrap
from brax.envs.wrappers.gym import GymWrapper, VectorGymWrapper
from brax.experimental.barkour_v0 import barkour_joystick

import os
import time
import wandb
import random
import numpy as np

from torch.utils.tensorboard import SummaryWriter

class ReplayBuffer(object):
    def __init__(self, size):
        self._storage = []
        self._maxsize = size
        self._next_idx = 0

    def __len__(self):
        return len(self._storage)

    def add(self, obs_t, action, reward, obs_tp1, done):
        data = (obs_t, action, reward, obs_tp1, done)

        if self._next_idx >= len(self._storage):
            self._storage.append(data)
        else:
            self._storage[self._next_idx] = data
        self._next_idx = (self._next_idx + 1) % self._maxsize

    def _encode_sample(self, idxes):
        obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
        for i in idxes:
            data = self._storage[i]
            obs_t, action, reward, obs_tp1, done = data
            obses_t.append(np.array(obs_t, copy=False))
            actions.append(np.array(action, copy=False))
            rewards.append(reward)
            obses_tp1.append(np.array(obs_tp1, copy=False))
            dones.append(done)
        return np.array(obses_t), np.array(actions), np.array(rewards), np.array(obses_tp1), np.array(dones)

    def sample(self, batch_size):
        idxes = [random.randint(0, len(self._storage) - 1) for _ in range(batch_size)]
        return self._encode_sample(idxes)

class TrainigConfig:

    def __init__(
        self, seed, total_timesteps, learning_rate, buffer_size,
        gamma, tau, batch_size, exploration_noise, learning_starts,
        policy_frequency, noise_clip, num_envs
    ):
        self.seed = seed
        self.total_timesteps = total_timesteps
        self.learning_rate = learning_rate
        self.buffer_size = buffer_size
        self.gamma = gamma
        self.tau = tau
        self.batch_size = batch_size
        self.exploration_noise = exploration_noise
        self.learning_starts = learning_starts
        self.policy_frequency = policy_frequency
        self.noise_clip = noise_clip
        self.num_envs = num_envs


class QNetwork(nn.Module):

    @nn.compact
    def __call__(self, obs, actions):
        combined_input = jnp.concatenate([obs, actions], axis=-1)
        x = nn.Dense(256)(combined_input)
        x = nn.relu(x)
        x = nn.Dense(256)(x)
        x = nn.relu(x)
        x = nn.Dense(1)(x)

        return x

class Actor(nn.Module):

    action_shape: tuple

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(256)(x)
        x = nn.relu(x)
        x = nn.Dense(256)(x)
        x = nn.relu(x)
        x = nn.Dense(self.action_shape[1])(x)
        x = jnp.tanh(x)

        return x 

class TrainState(TrainState):
    target_params: flax.core.FrozenDict


def train(config: TrainigConfig, use_wandb: bool):
    env_id = "barkour_v0_joystick"
    run_name = f"{env_id}__{config.seed}__{int(time.time())}"
    if use_wandb:
        wandb.init(
        project="ddpg-barkour-benchmark",
        config={
            "env": env_id,
            "seed": config.seed,
            "timesteps": config.total_timesteps,
            "lr": config.learning_rate,
            "buffer_size": config.buffer_size,
            "gamma": config.gamma,
            "tau": config.tau,
            "batch_size": config.batch_size,
            "exploration_noise": config.exploration_noise,
            "learning_starts": config.learning_starts,
            "policy_frequency": config.policy_frequency,
            "noise_clip": config.noise_clip,
        },
        sync_tensorboard=True,
        monitor_gym=True,
        name=run_name
        )
    
    writer = SummaryWriter(f"runs/{run_name}")

    random.seed(config.seed)
    np.random.seed(config.seed)
    rng = jax.random.PRNGKey(config.seed)
    rng, eval_rng = jax.random.split(rng)

    barkour_env = envs.create('barkour_v0_joystick', backend='generalized', batch_size=config.num_envs)
    env = VectorGymWrapper(env=barkour_env, seed=config.seed)

    barkour_eval_env = envs.create('barkour_v0_joystick', backend='generalized')
    eval_env = GymWrapper(env=barkour_eval_env, seed=config.seed)
    # jit_eval_env_reset = jax.jit(barkour_eval_env.reset)
    # jit_eval_env_step = jax.jit(barkour_eval_env.step)
    
    obs_shape = env.observation_space.shape
    action_shape = env.action_space.shape
    actor = Actor(action_shape=action_shape)
    critic = QNetwork()

    dummy_obs = jnp.ones(obs_shape, dtype=jnp.float32)  
    dummy_action = jnp.ones(action_shape, dtype=jnp.float32)  

    actor_params = actor.init(rng, dummy_obs)
    critic_params = critic.init(rng, dummy_obs, dummy_action)

    actor_optimiser = optax.adam(config.learning_rate)
    critic_optimiser = optax.adam(config.learning_rate)

    actor_state = TrainState.create(apply_fn=actor.apply, 
                                    params=actor_params, 
                                    target_params=actor_params, 
                                    tx=actor_optimiser)

    critic_state = TrainState.create(apply_fn=critic.apply, 
                                    params=critic_params, 
                                    target_params=critic_params, 
                                    tx=critic_optimiser)

    actor.apply = jax.jit(actor.apply)
    critic.apply = jax.jit(critic.apply)

    @jax.jit
    def actor_update(actor_state, critic_state, states):
        def loss_fn(actor_params, critic_params):
            actions = actor.apply(actor_params, states)
            return -(critic.apply(critic_params, states, actions)).mean()
        
        loss_value, grads = jax.value_and_grad(loss_fn)(actor_state.params, critic_state.params)
        actor_state = actor_state.apply_gradients(grads=grads)
        return loss_value, actor_state
    
    @jax.jit
    def critic_update(actor_state, critic_state, states, actions, rewards, next_states, dones, gamma):
        next_state_actions = actor.apply(actor_state.target_params, next_states)
        next_target = critic.apply(critic_state.params, next_states, next_state_actions)
        next_q_value = rewards + (1 - dones) * gamma * next_target

        def mse_loss(params):
            q_pred = critic.apply(params, states, actions)
            return ((q_pred - next_q_value)**2).mean(), q_pred

        (loss, q_pred), grads = jax.value_and_grad(mse_loss, has_aux=True)(critic_state.params)
        critic_state = critic_state.apply_gradients(grads=grads)
        return loss, q_pred, critic_state


    rb = ReplayBuffer(config.buffer_size)
    obs = env.reset()

    start_time = time.time()

    for global_step in range(config.total_timesteps):

        if global_step < config.learning_starts:
            action = jax.random.uniform(rng, shape=action_shape, minval=-1.0, maxval=1.0)
        else:
            action = actor.apply(actor_state.params, obs)
            noise = config.exploration_noise *jax.random.normal(rng, action_shape)
            action = jnp.clip(action + noise, -action_shape[0], action_shape[0])

        next_obs, reward, done, info = env.step(action)
        for i in range(config.num_envs):
            rb.add(obs[i], action[i], reward[i], next_obs[i], done[i])

        obs = next_obs

        if global_step > config.learning_starts:
            states, actions, rewards, next_states, dones = rb.sample(config.batch_size)
            critic_loss, old_val, critic_state = critic_update(
                actor_state,
                critic_state,
                states,
                actions,
                rewards,
                next_states,
                dones,
                config.gamma
            )

            if global_step % config.policy_frequency == 0:
                actor_loss, actor_state = actor_update(
                    actor_state,
                    critic_state,
                    states
                )

                actor_state = actor_state.replace(
                    target_params=optax.incremental_update(actor_state.params, actor_state.target_params, config.tau)
                )

                critic_state = critic_state.replace(
                    target_params=optax.incremental_update(critic_state.params, critic_state.target_params, config.tau)
                )

            if global_step % 100 == 0:
                writer.add_scalar("losses/actor_loss", jax.device_get(actor_loss), global_step)
                writer.add_scalar("losses/critic_loss", jax.device_get(critic_loss), global_step)
                writer.add_scalar("charts/SPS",int(global_step / (time.time() - start_time)),global_step)

        if global_step % 100000 == 0:
            eval_state = eval_env.reset()
            rollout = []
            for i in range(5000):
                action = actor.apply(actor_state.params, eval_state)
                eval_env.step(action)
                rollout.append(eval_env._state.pipeline_state)

            filename = "eval-" + str(global_step) + ".html"
            with open(filename, "w") as f:
                f.write(html.render(barkour_eval_env.sys, rollout))

        print("Global step:", global_step)

Also attaching CPU and GPU config

CPU

Architecture:            x86_64
  CPU op-mode(s):        32-bit, 64-bit
  Address sizes:         39 bits physical, 48 bits virtual
  Byte Order:            Little Endian
CPU(s):                  8
  On-line CPU(s) list:   0-7
Vendor ID:               GenuineIntel
  Model name:            Intel(R) Core(TM) i5-8265U CPU @ 1.60GHz
    CPU family:          6
    Model:               142
    Thread(s) per core:  2
    Core(s) per socket:  4
    Socket(s):           1
    Stepping:            12
    CPU max MHz:         3900.0000
    CPU min MHz:         400.0000

GPU

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla V100-SXM2-16GB           Off | 00000000:00:04.0 Off |                    0 |
| N/A   35C    P0              50W / 300W |  12534MiB / 16384MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
+---------------------------------------------------------------------------------------+

Corresponding jax version

jax                              0.4.23
jaxlib                           0.4.23+cuda12.cudnn89

[chiamp]

I unfortunately don't have access to a GPU so I can't benchmark this, but what batch size are you using? From personal experience, too small of a batch size is inefficient since the data needs to be loaded and unloaded from GPU memory. To take advantage of a GPU, I usually increase the batch size to improve efficiency.

[chkda]

Hey thanks for taking time to answer. I checked as per your advice, found out that replay buffer was the problem probably because of lot of GPU to CPU operations I guess. Now I have optimised it. Its working good. Thanks for your help

[8bitmp3]

Thank you @chkda and @chiamp This is good to know