-
Notifications
You must be signed in to change notification settings - Fork 397
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[XLA:GPU] Check failed in collective_pipeliner when using gradient accumulation with non-unrolled loop #14332
Comments
related JAX issue: |
Actually, this problem persists even with |
For the compilation error in collective_pipeliner, can you try with xla_gpu_run_post_layout_collective_pipeliner=false ? |
Hi Filipp, could you try TJ's suggestion and update this issue with the results and any errors if applicable? |
Hi guys, looks like this flag was added in very recent commit and has not been added to JAX latest release (0.4.30). I'll wait for JAX 0.4.31 to test it. Thank you! |
Hi, I have following setup:
I'm using following flags:
This works correctly and indeed hide layers' weights all-gather and gradient reduce-scatter behind computations.
Problems are starting to arise when I try to use gradient accumulation in this setup. It is implemented like this:
When I set gradient accumulation factor (num_minibatches_in_batch in this snippet) to value greater than 1, I'm getting following error during compilation:
Here is --xla_dump_to result:
xla_dump.tgz
One important fact here is that if I set
unroll
in jax.lax.fori_loop to True, then there is no compilation error and everything works. But this obviously leads to additional memory usage proportional to gradient accumulation factor so this hack doesn't seem to be viable.My hypothesis is that when using
--xla_gpu_enable_while_loop_double_buffering=true
with pipelined collectives and latency hiding scheduler, XLA compiler tries to double buffer this fori_loop which is actually undesired behavior.Basically, there are two problems:
I've tested this on JAX 0.4.29 and 0.4.30.
The text was updated successfully, but these errors were encountered: