{"id":2922,"date":"2024-12-01T18:35:22","date_gmt":"2024-12-01T18:35:22","guid":{"rendered":"https:\/\/timallanwheeler.com\/blog\/?p=2922"},"modified":"2024-12-01T18:35:22","modified_gmt":"2024-12-01T18:35:22","slug":"rollouts-on-the-gpu","status":"publish","type":"post","link":"https:\/\/timallanwheeler.com\/blog\/2024\/12\/01\/rollouts-on-the-gpu\/","title":{"rendered":"Rollouts on the GPU"},"content":{"rendered":"\n

Last month I wrote about<\/a> moving the Sokoban policy training code from CPU to GPU, yielding massive speedups<\/mark>.<\/span><\/span>Reduced training time by about 32x for the base case and training time scales much better with model size.<\/span><\/span><\/span><\/span> That significantly shortened both training time and the time it takes to compute basic validation metrics. It has not, unfortunately, significantly changed how long it takes to run rollouts, and relatedly, how long it takes to run beam search.<\/p>\n\n\n\n

The Bottleneck<\/h1>\n\n\n\n

The training that I’ve done so far has all been with teacher forcing<\/a>, which allows all inputs to be passed to the net at once:<\/p>\n\n\n\n

\"\"<\/figure>\n\n\n\n

<\/p>\n\n\n\n

When we do a rollout, we can’t pass everything in at once. We start with our initial state and use the policy to discover where we end up:<\/p>\n\n\n\n

\"\"<\/figure>\n\n\n\n

The problem is that the left-side of that image, the policy call, is happening on the GPU, but the right side, the state advancement, is happening on the CPU. If a rollout involves 62 player steps, then instead of one data transfer step like we have for training, we’re doing 61 transfers! Our bottleneck is all that back-and-forth communication:<\/p>\n\n\n\n

\"\"<\/figure>\n\n\n\n

<\/p>\n\n\n\n

Let’s move everything to the GPU.<\/p>\n\n\n\n

CPU Code<\/h1>\n\n\n\n

So what is currently happening on the CPU?<\/p>\n\n\n\n

At every state, we are:<\/p>\n\n\n\n

    \n
  1. Sampling an action for each board from the action logits<\/li>\n\n\n\n
  2. Applying that action to each board to advance the state<\/li>\n<\/ol>\n\n\n\n

    Sampling from the actions is pretty straightforward to run on the GPU. That’s the bread and butter of transformers and RL in general.<\/p>\n\n\n\n

    # policy_logits are [a\u00d7s\u00d7b] (a=actions, s=sequence length, b = batch size)\npolicy_logits, nsteps_logits = policy(inputs)\n\n# Sample from the logits using the Gumbel-max trick\nsampled_actions = argmax(policy_logits .+ gumbel_noise, dims=1)<\/code><\/pre>\n\n\n\n
    where we use the Gumbel-max trick<\/a> and the Gumble noise is sampled in advance and passed to the GPU like the other inputs:<\/p>\n\n\n\n
    using Distributions.jl\ngumbel_noise = rand(Gumbel(0, 1), size(a, s, b))<\/code><\/pre>\n\n\n\n
    <\/p>\n\n\n\n
    Advancing the board states is more complicated. Here is the CPU method for a single state:<\/p>\n\n\n\n
    function maybe_move!(board::Board, dir::Direction)::Bool\n    \u25a1_player::TileValue=find_player_tile(board)\n    step_fore = get_step_fore(board, dir)\n\n    \u25a1 = \u25a1_player # where the player starts\n    \u25a9 = \u25a1 + step_fore # where the player potentially ends up\n\n    if is_set(board[\u25a9], WALL)\n        return false # We would be walking into a wall\n    end\n\n    if is_set(board[\u25a9], BOX)\n        # We would be walking into a box.\n        # This is only a legal move if we can push the box.\n        \u25f0 = \u25a9 + step_fore # where box ends up\n        if is_set(board[\u25f0],  WALL + BOX)\n            return false # We would be pushing the box into a box or wall\n        end\n\n        # Move the box\n        board[\u25a9] &= ~BOX # Clear the box\n        board[\u25f0] |= BOX # Add the box\n    end\n\n    # At this point we have established this as a legal move.\n    # Finish by moving the player\n    board[\u25a1] &= ~PLAYER # Clear the player\n    board[\u25a9] |= PLAYER # Add the player\n\n    return true\nend<\/code><\/pre>\n\n\n\n
    <\/p>\n\n\n\n
    There are many ways to represent board states. This representation is a simple Matrix{UInt8}<\/code>, so an 8×8 board is just an 8×8 matrix. Each tile is a bitfield with components that can be set for whether that tile has\/is a wall, box, floor, or tile.<\/p>\n\n\n\n
    Moving the player has 3 possible paths:<\/p>\n\n\n\n