torch-native pipeline parallelism for big models

[muellerzr]

Example use:

import torch
from accelerate.inference import prepare_pippy
from accelerate.utils import set_seed
from transformers import T5ForConditionalGeneration, T5Config

set_seed(42)

config = T5Config()
model = T5ForConditionalGeneration(config)
model.eval()

# Create example inputs for the model
input = torch.randint(
    low=0,
    high=config.vocab_size,
    size=(2, 1024),  # bs x seq_len
    device="cpu",
    dtype=torch.int64,
    requires_grad=False,
)

example_inputs = {"input_ids": input, "decoder_input_ids": input}

model = prepare_pippy(model, example_kwargs=example_inputs)

args = (
    example_inputs["input_ids"].to("cuda:0"),
    example_inputs["decoder_input_ids"].to("cuda:0")
)
with torch.no_grad():
    output = model(*args)

Speed up:

Using 2x4090's in full precision

Bert

	Accelerate/Sequential	PiPPy + Accelerate
First batch	0.2137s	0.3119s
Average of 5 batches	0.0099s	0.0062s

GPT2

	Accelerate/Sequential	PiPPy + Accelerate
First batch	0.1959s	0.4189s
Average of 5 batches	0.0205s	0.0126s

T5

	Accelerate/Sequential	PiPPy + Accelerate
First batch	0.2789s	0.3809s
Average of 5 batches	0.0198s	0.0166s

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

Very cool API ! I like the design and how easily it is to use. I left a few comments around the split_points mainly.

[SunMarc]

nice !

[SunMarc]

We can definitely generate a balanced device_map for pippy exclusively "device_map = "balanced_pippy" if the current balanced option is not the best for that. However, I think it would be great if the user can use other options like "sequential". I didn't try but what happens when we only fill 2 gpus out of the 4 available (possible sequential case) ?

[SunMarc]

Just a thought about how to handle the split points.

• 1. We only expose device_map with predefined options ("sequential", "balanced_pippy")
• 2. We let the user use a custom device_map. For the custom case, it can be complicated since the user needs to be careful about the order (OrderedDict()) and he needs to attribute the gpu in a sequential manner because of that split_points.append(next(k for k, v in device_map.items() if v == i)). So that can be quite complicated.
• 3. We let the user let his own split points List[str].
     I think that 1) is a must. between 2) and 3), I prefer 3) since it is easier for the user.

[muellerzr]

Agreed to do 1 and 3

[SunMarc]

The API is in good shape ! Let's document the main functions a bit and we can merge it. I left a few comments but nothing blocking.

[SunMarc]

it would be great to have a sanity check, to make sure that we indeed have self.num_processes split points when we are generating the split_points + when the user manually pass them

[kwen2501]

Thanks a lot for the integration effort!
LGTM!

[SunMarc]

Thx for iterating ! LGTM

[kwen2501]

Thanks for writing the doc so quick! Looks good to me!

[muellerzr]

cc @MKhalusova for the docs!

[MKhalusova]

Nice work! I left a few comments to polish things in the docs a bit.

[muellerzr]

Final comment before merging, things that still need to be done in a latter PR at some point (but okay not being in the first iteration of this joint effort):

1. Specify balanced_pippy device map and allow a sequential device_map when making the pipeline via prepare_pippy
2. Look into supporting model.generate() through an alternative hook into the model forward if possible
3. Make sure all outputs end up on the CPU so users don't need to check at the end and we can call them via a .gather
4. Migrate the pippy-device-map-playground examples over to here as part of our examples folder

(I'll be doing 3& 4 this week as a follow-up prior to release)