Closed qGentry closed 1 month ago
related XLA issue: https://github.com/openxla/xla/issues/14332
Thanks for raising this.
I think it's an XLA:GPU issue, and we don't have any way to fix it from JAX.
That said, the hard-to-parsae error may be something we can get traction on from JAX... can you say a bit more about what would've helped in the error message? We attach Python source information to the HLO program, but it's up to XLA to raise errors that reference it... from JAX we could've at least told you which jitted function raised the compiler error, but I'm not sure if we have other information to provide...
I think we should close this in favor of the XLA issue. Looks like it just got assigned yesterday!
Description
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.
System info (python version, jaxlib version, accelerator, etc.)