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)

config = TrainigConfig(
    seed = 1,
    total_timesteps = 500000,
    learning_rate = 0.00003,
    buffer_size = 10000,
    gamma = 0.99,
    tau = 0.005,
    batch_size = 1024,
    exploration_noise = 0.1,
    learning_starts = 25000,
    policy_frequency = 2,
    noise_clip = 0.5,
    num_envs=4,
)

train(config,use_wandb=False)

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

[jakevdp]

Hi - thanks for the question! Your code defines a number of functions, but doesn't run them, and it's not clear how to call the train() function to reproduce the results you're seeing. Can you edit your question to add this information?

[chkda]

Hey, I apologise for missing that part. I have added it now. Please have a look

[jakevdp]

Thanks – I don't know much about the brax and wandb libraries, and what computations these are doing under the hood.

To diagnose this, I would start by looking at whether the slowdown is being driven by runtime or by compile time. If it's compile time, then it may be that the GPU compiler is just doing more optimizations and therefore taking more time (but then the GPU runtime would probably be faster than CPU)

If it's runtime, then I would guess that one of the libraries you are calling is using workflows with heavy reliance on scan or while_loop, which can be more efficient on CPU than GPU, because CPU is purpose-built for that kind of sequential workflow.

Another possibility to check: because this is an iterative solver with a convergence criterion based on accuracy, it's possible that this is driven by relative numeric accuracy of computations on the different devices. You could check this by seeing how many iterations it takes on CPU vs. GPU. If the GPU iterations are much more than CPU iterations, you should look at adjusting e.g. the matmul precision settings.

Hope that helps!