Pipeline parallelism hello world

[dionhaefner]

I'm experimenting with pipeline parallelism where subsequent computations are executed on different devices.

I've not been able to jit a simple function that takes parameters living on different devices. I'm assuming this needs some explicit sharding information but I got confused by the tutorials which seem to be written for the more advanced case of sharding individual axes (instead of entire arrays).

Example code:

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=2"

import jax
import jax.numpy as jnp

devices = jax.devices("cpu")

@jax.jit
def pipeline_test(x, t1, t2):
    x = jax.device_put(x, devices[0])
    x = x + t1
    x = jax.device_put(x, devices[1])
    x = x + t2
    return x

inp = jnp.ones(100)
t1 = jax.device_put(inp, devices[0])
t2 = jax.device_put(inp, devices[1])
pipeline_test(inp, t1, t2)

Error:

ValueError: Received incompatible devices for jitted computation. Got argument x of pipeline_test with shape float32[100] and device ids [0] on platform CPU and argument t2 of pipeline_test with shape float32[100] and device ids [1] on platform CPU

[aniquetahir]

I think you have to rely on something like alpa for now:
https://github.com/alpa-projects/alpa

[dionhaefner]

Thanks, that's what I figured. Alpa looks unmaintained unfortunately...

[mjsML]

Thank you for your question, indeed there is no ergonomic recipe on how to do this today, however a reference implementation you might want to use is implemented in Praxis , I don’t know enough about your use case, do mind sharing more details?

[dionhaefner]

I'm trying to solve memory issues by computing different parts of the pipeline on different devices (without replicating parameters). I also need to backpropagate through it so I can't split it into multiple jit blocks.

[aniquetahir]

One workaround would be to use tensor parallelism on the parameters. This is not as efficient as pipeline parallelism in terms of the reserved memory but it worked out for a model I wanted to run on reduced memory per gpu.

[aniquetahir]

Another solution would be to write the implementation of sequential layers as separate jit functions and then write the gpipe code to connect them. This would be much more programming intensive.

[MoFHeka]

@mjsML pipeline parallelism in praxis seems not 1F1B.