Updating multiple networks simultaneously on CPU

username: ervteng
cooked: <p>Currently working on moving an implementation of a couple popular reinforcement learning algorithms from TensorFlow to PyTorch, and the PyTorch code is noticeably slower (up to 50%). We believe it’s because some algorithms, such as Soft-Actor Critic, have multiple semi-independent neural networks (e.g. Policy, Value, Q Function) that during the update step need to be evaluated, their losses computed, and backpropagation performed on them. In TensorFlow these could all be computed with a single <code>session.run</code> call, and they would be parallelized across multiple cores. However, in PyTorch we are limited to evaluating and updating each network sequentially.</p>
<p>Note that our networks are quite small and we don’t expect much benefit in running on GPU. However, the asynchronous execution in CUDA seems to alleviate a lot of the performance gap between TensorFlow, but at the end of the day we still need to support CPU.</p>
<p>Was wondering what the “right” way to do this for PyTorch on CPU. I’ve played around with using Python threading to evaluate the networks (PyTorch does release the GIL, right?) as well as changing the <code>num_threads</code> and <code>KMP_BLOCKTIME</code> settings, with some success, but am still at the 50% performance gap. Guidance would be much appreciated, thanks!</p>

username: iffiX
cooked: <p>Show your implementation.</p>

username: ervteng
cooked: <p>Here is the pseudo-code of the model update (soft actor-critic):</p>
<pre><code class="lang-auto"># Sample observations, dones, rewards from experience replay buffer
observations, next_observations, dones, rewards = sample_batch()

# Evaluate current policy on sampled observations
(
    sampled_actions,
    log_probs,
    entropies,
    sampled_values,
) = policy.sample_actions(observations)

# Evaluate Q networks on observations and actions
q1p_out, q2p_out = value_network(observations, sampled_actions)
q1_out, q2_out = value_network(observations, actions)

# Evaluate target network on next observations
with torch.no_grad():
    target_values = target_network(next_observations)

# Evaluate losses
q1_loss, q2_loss = sac_q_loss(q1_out, q2_out, target_values, dones, rewards)
value_loss = sac_value_loss(log_probs, sampled_values, q1p_out, q2p_out)
policy_loss = sac_policy_loss(log_probs, q1p_out)
entropy_loss = sac_entropy_loss(log_probs)

total_value_loss = q1_loss + q2_loss + value_loss

# Backprop and weights update
policy_optimizer.zero_grad()
policy_loss.backward()
policy_optimizer.step()

value_optimizer.zero_grad()
total_value_loss.backward()
value_optimizer.step()

entropy_optimizer.zero_grad()
entropy_loss.backward()
entropy_optimizer.step()
</code></pre>
<p>I’m looking for a way to parallelize the <code>value_network</code>, <code>sample_actions</code>, and <code>target_network</code> evaluations, as well as the three <code>backward()</code> calls, as these components consume the most amount of time. I’ve tried running each in different Python threads, but probably a combination of the GIL and the context-switching overhead seems to negate any possible speedup.</p>

username: iffiX
cooked: <p>If your model is small, you can consider about using <code>torch.jit._fork</code> and <code>torch.jit._wait</code>, they can work around the GIL problem, but also requires to convert your models to <code>JitScriptModule</code>.</p>
<p>While pytorch does release GIL, it could still affect model performance if your model consists of many small linear layers like 100(in)-&gt;100(out).</p>

username: ervteng
cooked: <p>The <code>jit._fork</code> seems promising, thanks! It wasn’t clear to me from the documentation, but does the fork and wait have to be called within a <code>JitScriptModule</code> or can it be called <em>on</em> functions from <code>JitScriptModule</code>s in regular Python code?</p>

username: iffiX
cooked: <p>Compile a model to <code>JitScriptModule</code> like this:</p>
<pre><code class="lang-auto">jit_actor = t.jit.script(acc)
</code></pre>
<p>and call it like this:</p>
<pre><code class="lang-auto">future = [t.jit._fork(model, *args_list) for model, args_list in zip(models, args_lists)]
</code></pre>
<p>Then asynchronously wait for results:</p>
<pre><code class="lang-auto">result = [t.jit._wait(fut) for fut in future]
</code></pre>
<p>These APIs work outside of your JIT module.</p>