Closed LeonEricsson closed 2 months ago
Can you try adding a call to .block_until_ready()
after each pmap
call on one of the outputs? 32 sounds like the size of one of our internal queues that limits how far ahead we can enqueue computations. A possible explanation is that the very first computation you run isn't terminating for some reason, and if you add block_until_ready()
then you should see that behavior right away.
(As to why it isn't terminating? I don't know without a repro, but my guess is it's a bug in your code.)
Thanks for the hasty response.
Note this is how selfplay is called:
while True:
# Log to tensorboard
for k, v in log.items():
writer.add_scalar(f'{k}', v, iteration)
writer.flush()
if iteration >= config.max_num_iters:
break
log = {}
# Selfplay
key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, num_cpu_devices)
env_state, samples = selfplay(
models,
keys,
env_state,
)
samples.obs.block_until_ready()
log.update(
{
'stats/samples_per_sec': SAMPLES_PER_ITERATION / time_sample_generation,
},
)
print(iteration)
iteration += 1
Adding the block does seem to reveal the problem. The selfplay calls are extremely slow, and it takes probably a good 15 minutes until we reach the 32nd iteration. Given your suggestion of a internal queue, this makes a lot of sense. I've failed to consider the asynchronous dispatch, and was fooled by the print statements!
Thank you again.
Description
I've got a self-play (not important) function that I would like to execute on the CPU, but I'm running into an issue where the execution mysteriously freezes on the 32nd iteration. Now, when I say freeze or halt, I actually don't know if this is a complete freeze or just a immense slow down. The first 31 iterations run in less than a second, and I've waited for 5 minutes without the 32nd iteration completing. Anyway, let me walk through the code and what I've found so far.
This is the selfplay function that is called externally N times, and halts to a stop on exactly the 32nd call (every time):
I've identified the
modified_gumbel_muzero_policy
function as the culprit (swapping it out for a dummy initialization fixes the issue). My initial thought was that there was some complex logic inmodified_gumbel_muzero_policy
that for some (unknown) reason was really resource intense. So I started replacing all the logic step-by-step with dummy initializers to see when the issue fixed itself, and I ended up with this:There's hardly anything left?! I've been able to remove almost all the logic from this function and everything still freezes on the 32nd selfplay call.
Swapping out
action
andaction_weights
to dummy initializers does fix the problem:This runs.
What is going on here? Does anything have any intuition as to why this is happening?
I'm enabling CPU executing across cores using:
System info (python version, jaxlib version, accelerator, etc.)