[doc] pipeline doc typos/improvements

[stas00]

This PR fixes a few typos, rewrites a few sentences to hopefully be easier to understand.

And after reading this doc I have a questions:

1. It's not clear from the tutorial why TransformerBlockPipe needs to return mask. And if there is some hidden reason, shouldn't it then return hidden too?

2. Why do you need the vars to be a single tensor or a tuple? This surely could be a problematic limitation since some layers could have forward args which aren't tensors. Perhaps a more dynamic approach could be used instead? e.g. with the automatic data remapper I posted here: rfc: automating the switching of inputs to the device of the params pytorch/pytorch#49961 (comment)

3. wrt Tied Layers - could you please give an example? or at least a pointer to where it can be seen in action?

Please let me know if it'd be better to separate the questions into an Issue and not mix it with the PR. I was just thinking that perhaps the answers could be integrated into the tutorials, hence posted them here.

Thank you!

[stas00]

Is this supposed to communicate that when LayerSpec is used, not only that each GPU allocates only the layers it's assigned to but also that no CPU allocation happens at all?

Or should I have rewritten it to say:

Then each worker will allocate only the layers it's assigned to. So continuing the example from the previous paragraph, a machine with 16 GPUs will need to allocate a total of 1x model size on its CPU, compared to 16x in the LayerSpec example.

I think the two paras have GPU and CPU mixed up.

[ShadenSmith]

That's correct, LayerSpec does not allocate modules until the partitioning is complete. At that point, each rank allocates only the layers assigned to it and they are then moved to the GPU. This lets us build distributed models that wouldn't fit in CPU memory.

[stas00]

So is it better to leave the original modification of this PR, or replace with the more exemplifying:

Then each worker will allocate only the layers it's assigned to. So continuing the example from the previous paragraph, a machine with 16 GPUs will need to allocate a total of 1x model size on its CPU, compared to 16x in the LayerSpec example.

[stas00]

OK, I went ahead with the longer version. Please double check that it's still correct.

[ShadenSmith]

Thanks so much for your ongoing contributions to DeepSpeed!

Good questions!

It's not clear from the tutorial why TransformerBlockPipe needs to return mask. And if there is some hidden reason, shouldn't it then return hidden too?

Pipeline parallelism is somewhat rigid in that the input to each layer's forward() is always the output of the previous layer. When you have a stack of transformers, we need to return the mask in order to provide it to the next layer. We don't return hidden because it's not needed by the next layer. The next layer uses output as its hidden input.

From a usability perspective, I'd love to be able to specify that some inputs are static and globally available (e.g., mask) so that we wouldn't need to forward it down the pipeline. No concrete plans for it yet, though.

Why do you need the vars to be a single tensor or a tuple?

Agreed that this is a limitation. As a starting place, we wanted to match torch.nn.Sequential functionality and that container has the same restriction. I'd like to be more flexible. The primary limitation is the communication code at stage boundaries. I've been working on a more flexible communication manager that would allow something like this and also remove some hacky code in the pipeline engine.

wrt Tied Layers - could you please give an example?

Definitely. I'm trying to release our Megatron GPT implementation asap. The word embeddings are tied because they are used on both the first and last stages.

[stas00]

Thank you for all the other answers in the previous comment, @ShadenSmith.

It's not clear from the tutorial why TransformerBlockPipe needs to return mask. And if there is some hidden reason, shouldn't it then return hidden too?

Pipeline parallelism is somewhat rigid in that the input to each layer's forward() is always the output of the previous layer. When you have a stack of transformers, we need to return the mask in order to provide it to the next layer. We don't return hidden because it's not needed by the next layer. The next layer uses output as its hidden input.

OK, can then this example be expanded to show how the mask is fed to forward on subsequent calls in the pipeline?
https://www.deepspeed.ai/tutorials/pipeline/#inputs-and-outputs
The issue here is that there could be other inputs and it's not clear at all how the user should handle those. i.e. you're saying "The next layer uses output as its hidden input." but this is not generic enough.

I think it'd help if all that code is explicit in the open so that the user can see what's going on and then expand it to support other inputs.

I hope you are can see how the provided example isn't helping with a general case.

We have a huge list of inputs as it can be seen here in just one example of t5:

https://github.com/huggingface/transformers/blob/c60e0e1ee45f4bf1017736b146c51729f120bb83/src/transformers/models/t5/modeling_t5.py#L600-L612

and you can also see that some of them aren't tensors, but Boolean.

I see fairscale has the same limitation:
https://fairscale.readthedocs.io/en/latest/api/nn/pipe.html

Input and output have to be a Tensor or a tuple of tensors. This restriction is applied at partition boundaries too.

Is this the case where the pipeline implementation tries to protect itself from having an object as input, since it won't know how to .to() such input to the correct device. But it shouldn't be a problem with a simple non-tensor variable or a list/dict of tensors and/or simple variables/list/dict/tuple. Am I missing something? As I mentioned I wrote a helper function that handles any such types transparently:

def recursive_to(device, item):
    """
    Switch any tensors found in `item` to `device`.
    Currently can handle a single tensor, or any of the nested list, tuple and dict structures.
    """

    if torch.is_tensor(item):
        return item.to(device)

    elif isinstance(item, list):
        for i, x in enumerate(item):
            item[i] = recursive_to(device, x)
        return item

    elif isinstance(item, tuple):
        return tuple(recursive_to(device, list(item)))

    elif isinstance(item, dict):
        for k, v in item.items():
            item[k] = recursive_to(device, v)
        return item

    else:
        return item

I do understand that we need each stage in the pipeline to expect self + input and return a single output, so that it stacks up. The key is that if the input or output is a tuple it should be able to contain non-tensor elements I believe.

[stas00]

After studying this more I finally wrapped my head around it.

I propose the following change to https://www.deepspeed.ai/tutorials/pipeline/#inputs-and-outputs to make it crystal clear from the get going:

1. Rewrite the code as:

before:

class TransformerBlock(nn.Module)
    ...
    def forward(self, hidden, mask):
        hidden = self.compute(hidden, mask)
        return hidden

after:

class TransformerBlockPipe(TransformerBlock)
    def forward(self, inputs):
        hidden, mask = inputs
        hidden = super().forward(hidden, mask)
        return (hidden, mask)

and add a note:

It's clear to see how the output of the previous stage in the pipeline becomes the input to next stage, since the previous stage in the pipeline now returns a tuple of (hidden, mask) and the next stage expects besides self, a tuple of (hidden, mask).

2. One more piece is missing from that section of the tutorial. Since the original code expected just one non-tuple variable returning from TransformerBlock.forward you have to modify the outmost pre-pipe caller from:

hidden = self.transformer_block(hidden, mask)

to:

hidden, _ = self.transformer_block_pipe(hidden, mask)

because otherwise hidden will contain a tuple, and not the original single variable.

Please let me know if you feel this makes things easier to understand and I will expand this PR.

[stas00]

And as I'm trying to port t5 transformers to Pipe I see why there is a requirement of all inputs to be Tensors - because they have to be spliced into micro-batches.

1. If I got it right - they shouldn't just be tensors, they should be tensors with first dimension of batch size. I actually had a tensor that was tracking things, but it's a no-go, since it's of a wrong shape. So this pre-requisite can and probably should be clarified.

2. Surely, there should be a way to handle booleans here. So the input parts that are of batch size will be spliced into micro-batches, but parts that are on/off switches can be probably left alone - hmm, but then it might be difficult to put them back into outputs correctly - perhaps the inputs should have an optional leave_me_alone_inputs which should be replicated as is on micro-batch splicing? So the sig would be:

def forward(self, inputs, leave_me_alone_inputs):
    ...
    return (inputs, leave_me_alone_inputs)

?
This would also allow passing random other inputs that don't need to be micro-batch-sliced.

3. Same issue happens to optional input parts, which could be None or a tensor. The Pipe needs to slice the tensor, but leave None alone.

Currently in transformers t5 I have 15 distinct input parts, some tensors of BS, some None, some Bool, some tensors that aren't of BS.

I think I have an idea how to work around this by creating a Pipe wrapper class at run time and initialize it with a late __init__ to overcome the need to pass flags and states via forward. I will probably have to come up with some closure to keep global state of the pipe. I will keep you posted.

[ShadenSmith]

Hi @stas00 , thanks a ton for your contibutions and comments. Unfortunately I lost my first reply in a power outage. Here are a few quick comments and clarifications and I'll double back over the day.

DeepSpeed does not slice inputs into micro-batches, we just rely on the data loader to do that. We already have the concept of micro-batches for gradient accumulation and build off of that instead. Instead, we just next() the provided data iterator gradient_accumulation_steps number of times.

The forward() input limitations are partly a trade-off of pipeline parallelism, and partly engineering limitations. Pipeline modules use this forward implementation implicitly:

def forward(self, inputs):
    x = inputs
    for layer in self.layers:
        x = layer(x)
    return x

This contract lets us split self.layers and place them on different machines. There are two consequences here:

• Inputs like booleans, etc. still need to "flow" into the layer unless they are default valued keyward args. That's where something like the TransformerBlockPipe subclass comes in, we can wrap the functionality of an existing module and just specialize the call into forward with any additional args. Within the pipeline, there's no call site since it's done implicitly by the engine.
• Successive layers could reside on different machines, and so any inputs must be communicated by NCCL. This is why we are limited to tensor inputs/outputs today. I agree that your recursive to() would help in a non-distributed environment, but we need a bit more in order to send/recv (also, the receiving processes do not know the shape/type of data to receive without an initial handshake). Maybe we could something like pickle more general data structures and send them as ByteTensors ?

[stas00]

• Inputs like booleans, etc. still need to "flow" into the layer unless they are default valued keyward args. That's where something like the TransformerBlockPipe subclass comes in, we can wrap the functionality of an existing module and just specialize the call into forward with any additional args. Within the pipeline, there's no call site since it's done implicitly by the engine.

Yes, this one I have already figured out and basically created a wrapper class to which I pass the flow True/False args via its __init__. So we are good here. Thank you!

• Successive layers could reside on different machines, and so any inputs must be communicated by NCCL. This is why we are limited to tensor inputs/outputs today. I agree that your recursive to() would help in a non-distributed environment, but we need a bit more in order to send/recv (also, the receiving processes do not know the shape/type of data to receive without an initial handshake). Maybe we could something like pickle more general data structures and send them as ByteTensors ?

I think what would help you to help me is to show you exactly what I'm dealing with:

1. I have a few simple things like:

   None, which may become a tensor inside the stages of the pipe - but it shouldn't be sliced or reconstructed on the way out.

   There really should be a None type of tensor

2. Then I have a medium complexity things like:

   (). which gets filled out by each stage - e.g. it accumulates all hidden_states, attentions, etc - I have 5 of these in t5

   This one is difficult since now instead of accumulating say 6 values (for a stack of 6 blocks), I end up with 12, 18 or more depending on the number of chunks and so I'm attempting to re-construct them once we are out of the pipeline. Not simple but doable.

   Same here - it shouldn't be sliced or reconstructed on the way out.

   What I'm attempting at the moment is a closure, that I pass to the wrapper's __init__ and grow these accumulators via the closure.

3. And the one I'm fighting with at the moment is this:

   past_key_values which is tuple of tuples of tensors, and the batch dimension is not even on the top level.

   This one needs to be sliced but on the inside of a tuple, but since it's a tuple of a tuple I can't pass it with input - I'm trying to pass it to wrapper's __init__ instead and to do my own slicing from inside the wrapper, by deriving the size of the microbatch and counting each invocation (and hoping they are not async).

   I'm thinking of adding some encoding logic where I convert the tuple of a tuple into a new stack of tensors, pushing the batch dimension to the front so that it get sliced correctly and then inside forward pick out the section that is relevant to that block. pretty nuts.

The original logic runs a simple loop, picks the right segment for each block in the loop from the first tuple, then takes chunks of that segment via the nested tuple and eventually gets a 2 tensor slice which it adds to hidden_state in self-attention. here is the structure I need to pass around:

In this case we have BS=3, 6 blocks, and each block has 4 past key values stored. As you can see the batch dimension (bs=3 here) that PP needs to slice on is buried 3 levels. So I need to invert that if I want to pass it to PP as part of input.

Also note that none of these have an issue of being switched to the right device, they are all tensors and simple variables, but aren't necessarily w/o a nested python structure around them (tuple or a tuple of tuples).

Hope this helps to understand what I'm dealing with and I hope you have some practical ideas.

Thank you!

Maybe we could something like pickle more general data structures and send them as ByteTensors ?

I haven't thought of this - thank you!

Oh, I think in fairscale/pytorch pipeline it does slice - i.e. expects the inputs to be of batch size in the first dimension. Let me double check.

DeepSpeed does not slice inputs into micro-batches, we just rely on the data loader to do that.

Hmm, ok, then I don't think this would work at all in the specific case I'm attempting - I'm currently trying to the use the newly added Pipeline in pytorch which I understand is coming from fairscale. The thing is again trying to make all of transformers t5 layers to be fully a pipeline would be very difficult, so what I'm trying right now is a hybrid approach of using the pipe in 2 places specifically where we have stacks. So it looks something like:

T5Model
  Encoder -> T5Stack -> Pipe
  Decoder -> T5Stack -> Pipe

so 2 pipes!

and it sounds that this approach won't work with DeepSpeed pipeline.

Am I correct?

[stas00]

let's merge this as is?

[g-karthik]

@stas00 @ShadenSmith I think there's a typo in this sentence, it should be "compared to 16x in the Sequential example"

[stas00]

Good catch, @g-karthik

How about:

So, comparing to the
example from the previous paragraph, a machine with 16 GPUs will need to allocate a
total of 1x model size on its CPU and not 16x.

[g-karthik]

Sounds fine to me!

[stas00]

#888