5 ms·

The abstractions provided by JAX for parallelism are beautiful. JAX is an absolute master-class in programming-interface design and a lesson in the power of pr

The abstractions provided by JAX for parallelism are beautiful. JAX is an absolute master-class in programming-interface design and a lesson in the power of providing composable primitive operations and FP inspired design. An astounding amount of complexity is hidden from the user behind primitives like pmap, and the power is exposed in such a simple interface.

6gvONxR4sf7o 4y ago
That's true, and is a massive part of what I love about JAX, but they also form barriers in weird parts of your code, preventing standard introspection tools, which is the single thing I hate about JAX. The errors are amazingly opaque.
- mattjjatgoogle 4y ago
  If you have any particular examples in mind, and time to share them on https://github.com/google/jax/issues https://github.com/google/jax/issues, we'd love to try to improve them. Improving error messages is a priority. About introspection tools, at least for runtime value debugging there is to some extent a fundamental challenge: since jax.jit stages computation out of Python (though jax.grad and jax.vmap don't), it means standard Python runtime value inspection tools, like printing and pdb, can't work under a jax.jit as the values aren't available as the Python code is executing. You can always remove the jax.jit while debugging (or use `with jax.disable_jit(): ...`), but that's not always convenient, and we need jax.jit for good performance. We recently added some runtime value debugging tools which work even with jax.jit-staged-out code (even in automatically parallelized code!), though they're not the standard introspection tools: see `jax.debug.print` and `jax.debug.breakpoint` on https://jax.readthedocs.io/en/latest/debugging/index.html https://jax.readthedocs.io/en/latest/debugging/index.html and https://jax.readthedocs.io/en/latest/debugging/print_breakpoint.html https://jax.readthedocs.io/en/latest/debugging/print_breakpo.... If you were thinking about other kinds of introspection tooling, I'd love to hear about it!
  - 6gvONxR4sf7o 4y ago
    > with jax.disable_jit(): ... That's handy, and I hadn't seen it before, thanks. It's been a bit, but I think the most frustrating errors were around mapping pytrees (like this issue https://github.com/google/jax/issues/9928 https://github.com/google/jax/issues/9928). I'm not sure the exact solution, but the axis juggling and specifications were where I remember a lot of pain, and the docs (though extensive) were unclear. At times it feels like improvements are punted on in the hopes that xmap eventually fixes everything (and xmap has been in experimental for far longer than I expected). Also the barriers where I couldn't disable jit. IIRC pmap automatically jits, so there was no way to avoid staging that part out. When it came to doing some complex jax.lax.ppermute, it felt more difficult than it needed to be to debug. Next time I encounter something particularly opaque, I'll share on the github issue tracker.
    - mattjjatgoogle 4y ago
      Thanks for taking the time to explain these. > It's been a bit, but I think the most frustrating errors were around mapping pytrees (like this issue https://github.com/google/jax/issues/9928 https://github.com/google/jax/issues/9928). We've improved some of these pytree error messages but it seems that vmap one is still not great. Thanks for the ping on it. > Also the barriers where I couldn't disable jit. IIRC pmap automatically jits, so there was no way to avoid staging that part out. That was indeed a longstanding issue in pmap's implementation. And since people came to expect jit to be "built in" to pmap, it wasn't easy to revise. However, we recently (https://github.com/google/jax/pull/11854 https://github.com/google/jax/pull/11854) made `jax.disable_jit()` work with pmap, in the sense that it makes pmap execute eagerly, so that you can print/pdb/etc to your heart's content. (The pmap successor, shard_map (https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html https://jax.readthedocs.io/en/latest/jep/14273-shard-map.htm...), is eager by default. Also it has uniformly good error messages from the start!) > Next time I encounter something particularly opaque, I'll share on the github issue tracker. Thank you for the constructive feedback!
      - 6gvONxR4sf7o 4y ago
        Thanks! One last thing, since I have your ear. The function transformation aspects of jax seem to make their way into downstream libraries like haiku, resulting in a lot of "magic" that can be difficult to examine and debug. Are there any utils you made to make jax's own transformations more transparent, which you think might be helpful to third party transformations? Higher order functions are difficult in general, and it would be fantastic to have core patterns or tools for breaking them open.
mattjjatgoogle 4y ago
Thanks for the kind words! We've been doing a lot more work in this direction too, for both compiler-based automatic parallelization [0] and a work-in-progress pmap successor for 'manual' parallelism (per-device code and explicit collectives) [1] which composes seamlessly with the compiler-based stuff. [0] https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html https://jax.readthedocs.io/en/latest/notebooks/Distributed_a... [1] https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html https://jax.readthedocs.io/en/latest/jep/14273-shard-map.htm...
alfalfasprout 4y ago
Agreed. Though keep in mind they built on a lot of failed attempts at doing the same to get here.
hummus_bae 4y ago
[dead]

6gvONxR4sf7o 4y ago

That's true, and is a massive part of what I love about JAX, but they also form barriers in weird parts of your code, preventing standard introspection tools, which is the single thing I hate about JAX. The errors are amazingly opaque.

mattjjatgoogle 4y ago
If you have any particular examples in mind, and time to share them on https://github.com/google/jax/issues https://github.com/google/jax/issues, we'd love to try to improve them. Improving error messages is a priority. About introspection tools, at least for runtime value debugging there is to some extent a fundamental challenge: since jax.jit stages computation out of Python (though jax.grad and jax.vmap don't), it means standard Python runtime value inspection tools, like printing and pdb, can't work under a jax.jit as the values aren't available as the Python code is executing. You can always remove the jax.jit while debugging (or use `with jax.disable_jit(): ...`), but that's not always convenient, and we need jax.jit for good performance. We recently added some runtime value debugging tools which work even with jax.jit-staged-out code (even in automatically parallelized code!), though they're not the standard introspection tools: see `jax.debug.print` and `jax.debug.breakpoint` on https://jax.readthedocs.io/en/latest/debugging/index.html https://jax.readthedocs.io/en/latest/debugging/index.html and https://jax.readthedocs.io/en/latest/debugging/print_breakpoint.html https://jax.readthedocs.io/en/latest/debugging/print_breakpo.... If you were thinking about other kinds of introspection tooling, I'd love to hear about it!
- 6gvONxR4sf7o 4y ago
  > with jax.disable_jit(): ... That's handy, and I hadn't seen it before, thanks. It's been a bit, but I think the most frustrating errors were around mapping pytrees (like this issue https://github.com/google/jax/issues/9928 https://github.com/google/jax/issues/9928). I'm not sure the exact solution, but the axis juggling and specifications were where I remember a lot of pain, and the docs (though extensive) were unclear. At times it feels like improvements are punted on in the hopes that xmap eventually fixes everything (and xmap has been in experimental for far longer than I expected). Also the barriers where I couldn't disable jit. IIRC pmap automatically jits, so there was no way to avoid staging that part out. When it came to doing some complex jax.lax.ppermute, it felt more difficult than it needed to be to debug. Next time I encounter something particularly opaque, I'll share on the github issue tracker.
  - mattjjatgoogle 4y ago
    Thanks for taking the time to explain these. > It's been a bit, but I think the most frustrating errors were around mapping pytrees (like this issue https://github.com/google/jax/issues/9928 https://github.com/google/jax/issues/9928). We've improved some of these pytree error messages but it seems that vmap one is still not great. Thanks for the ping on it. > Also the barriers where I couldn't disable jit. IIRC pmap automatically jits, so there was no way to avoid staging that part out. That was indeed a longstanding issue in pmap's implementation. And since people came to expect jit to be "built in" to pmap, it wasn't easy to revise. However, we recently (https://github.com/google/jax/pull/11854 https://github.com/google/jax/pull/11854) made `jax.disable_jit()` work with pmap, in the sense that it makes pmap execute eagerly, so that you can print/pdb/etc to your heart's content. (The pmap successor, shard_map (https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html https://jax.readthedocs.io/en/latest/jep/14273-shard-map.htm...), is eager by default. Also it has uniformly good error messages from the start!) > Next time I encounter something particularly opaque, I'll share on the github issue tracker. Thank you for the constructive feedback!
    - 6gvONxR4sf7o 4y ago
      Thanks! One last thing, since I have your ear. The function transformation aspects of jax seem to make their way into downstream libraries like haiku, resulting in a lot of "magic" that can be difficult to examine and debug. Are there any utils you made to make jax's own transformations more transparent, which you think might be helpful to third party transformations? Higher order functions are difficult in general, and it would be fantastic to have core patterns or tools for breaking them open.

mattjjatgoogle 4y ago

Thanks for the kind words! We've been doing a lot more work in this direction too, for both compiler-based automatic parallelization [0] and a work-in-progress pmap successor for 'manual' parallelism (per-device code and explicit collectives) [1] which composes seamlessly with the compiler-based stuff. [0] https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html https://jax.readthedocs.io/en/latest/notebooks/Distributed_a... [1] https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html https://jax.readthedocs.io/en/latest/jep/14273-shard-map.htm...

alfalfasprout 4y ago

Agreed. Though keep in mind they built on a lot of failed attempts at doing the same to get here.

hummus_bae 4y ago

[dead]