Repository: pytorch
Issue: Torch Distributed Pipelinging API datatype mismatch
Labels: oncall: distributed,triaged,module: pipelining

### 🐛 Describe the bug

Hello, I encountered some issues while using torch.distributed.pipelining.



I tested PiPPy/examples/huggingface/pippy_gpt2.py with the default configuration. Because I'm working on full model testings, I added a loss function to complete the backward phase in this example. The modifications refer to the following code https://github.com/pytorch/pytorch/blob/78e40b271b3d394d8d99ff48f3faa0f0f3464c21/test/distributed/pipelining/test_schedule.py#L87-L116. 


Original codes: 
https://github.com/pytorch/PiPPy/blob/main/examples/huggingface/pippy_gpt2.py#L67-L78

My modifications :
```python
-    # Attach to a schedule
-    schedule = ScheduleGPipe(stage, args.chunks)
-
-    # Full batch inputs as in single-worker case
+    loss_fn = torch.nn.CrossEntropyLoss()
+    schedule = ScheduleGPipe(stage, args.chunks,loss_fn=loss_fn)
     inputs = generate_inputs_for_model(
-        model_class, gpt2, model_name, args.batch_size, args.device)
-
-    # Run
+            model_class, gpt2, model_name, args.batch_size, args.device)
+    targets = inputs["input_ids"]
     if args.rank == 0:
-        schedule.step(**inputs)
+            losses = []
+            out = schedule.step(**inputs, target=targets, losses=losses)
     else:
-        out = schedule.step()
+            losses = []
+            out = schedule.step(target=targets, losses=losses)
```

### Q1: Do I use the API correctly?




### Q2

It keeps throwing this error:
```bash
self._maybe_compute_loss(self._stage, output, target_mbs, i)
 [rank0]: File "/home/project/lib/python3.10/site-packages/torch/distributed/pipelining/schedules.py", line 83, in _maybe_compute_loss [rank0]: loss = self._compute_loss(output, target_mbs[mb_index]) # type: ignore[index] 
[rank0]: File "/home/project/lib/python3.10/site-packages/torch/distributed/pipelining/schedules.py", line 192, in _compute_loss 
[rank0]: return self._loss_fn(output, target) # type: ignore[misc] 
[rank0]: File "/home/project/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1566, in _wrapped_call_impl [rank0]: return self._call_impl(*args, **kwargs) 
[rank0]: File "/home/project/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1575, in _call_impl [rank0]: return forward_call(*args, **kwargs) 
[rank0]: File "/home/project/lib/python3.10/site-packages/torch/nn/modules/loss.py", line 538, in forward [rank0]: return F.mse_loss(input, target, reduction=self.reduction) 
[rank0]: File "/home/project/lib/python3.10/site-packages/torch/nn/functional.py", line 3373, in mse_loss [rank0]: if not (target.size() == input.size()): [rank0]: AttributeError: 'tuple' object has no attribute 'size'
```
the `mse_loss` in the last line is another loss function I tested, but they have same errors.

The output of `self._stage.forward_one_chunk(i, arg_mbs[i], kwarg_mbs[i])` is a tuple containing tensors without information like loss and logits, and past key values. However, the `_compute_loss` function expects tensors as inputs but gets tuples.

I have to modify the model and the function of generating input, then compare it with the case of training without parallel to understand the results. My codes are here: https://gist.github.com/YingLiGithub/8e53e1f54c5713f12003cfc4df1945a4
The following modifications in questions 2 and 3 are for my code example.

If I understand correctly, the first item in the output tuple represents the logits. So, I filter this item and pass it to `compute_loss`. Can you tell me if this is the correct way to solve it?

I modified the code, function:  _maybe_compute_loss 
[schedules.py](https://github.com/pytorch/pytorch/blob/main/torch/distributed/pipelining/schedules.py#L80-L83)

```python
def _maybe_compute_loss(self, stage, output, target_mbs, mb_index):
       if stage.is_last and self._has_backward:
           test1=output[0].view(-1, output[0].size(-1))
           test2=target_mbs[mb_index].view(-1)
           loss = self._compute_loss(test1, test2) # type: ignore[index]
           self._internal_losses.append(loss)
```


### Q3: 

After making these changes, I was able to run the backward process, but it returned the following error after several stages.
```bash
[rank0]: Traceback (most recent call last): [rank0]: File "/home/project/lib/python3.10/site-packages/torch/distributed/pipelining/_backward.py", line 66, in stage_backward [rank0]: extract_tensors_with_grads(stage_output, output_grads) [rank0]: File "/home/project/lib/python3.10/site-packages/torch/distributed/pipelining/_backward.py", line 52, in extract_tensors_with_grads [rank0]: assert len(output_val) == len(grad_val) 
[rank0]: AssertionError [rank0]: The above exception was the direct cause of the following exception: 
[rank0]: Traceback (most recent call last): [rank0]: File "/home/SimValidation/pipe/pygpt2-training-profiler_singleBatch_pipev2.py", line 182, in <module> 
[rank0]: run(args) [rank0]: File "/home/SimValidation/pipe/pygpt2-training-profiler_singleBatch_pipev2.py", line 131, in run [rank0]: out = schedule.step(**inputs, target=targets, losses=losses) 
[rank0]: File "/home/project/lib/python3.10/site-packages/torch/distributed/pipelining/schedules.py", line 341, in step 
[rank0]: self._step_microbatches(args_split, kwargs_split, targets_split, losses) 
[rank0]: File "/home/project/lib/python3.10/site-packages/torch/distributed/pipelining/schedules.py", line 419, in _step_microbatches 
[rank0]: self._stage.backward_one_chunk(i, loss=loss) [rank0]: File "/home/project/lib/python3.10/site-packages/torch/distributed/pipelining/stage.py", line 602, in backward_one_chunk 
[rank0]: self.grads_input = self.backward_maybe_with_nosync(bwd_kwargs) 
[rank0]: File "/home/project/lib/python3.10/site-packages/torch/distributed/pipelining/stage.py", line 493, in backward_maybe_with_nosync [rank0]: grads_input = stage_backward(**bwd_kwargs) [rank0]: File "/home/project/lib/python3.10/site-packages/torch/distributed/pipelining/_backward.py", line 101, in stage_backward 
[rank0]: raise RuntimeError(exc_msg) from e 
[rank0]: RuntimeError: [rank0]: Failed to run stage backward: 
[rank0]: Stage output: ('Tensor(torch.Size([1, 128, 768]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)', 'Tensor(torch.Size([1, 12, 128, 64]), grad=True, dtype=torch.float32)') [rank0]: Output gradient: ('Tensor(torch.Size([1, 128, 768]), grad=False, dtype=torch.float32)',) [rank0]: Input: ['Tensor(torch.Size([1, 128]), grad=False, dtype=torch.int64)'] [rank0]: [rank0]:
```
I guess this issue relates to the data format. So I modify the following function  `function backward_one_chunk`
https://github.com/pytorch/pytorch/blob/78e40b271b3d394d8d99ff48f3faa0f0f3464c21/torch/distributed/pipelining/stage.py#L574-L592

My modification: 
```python
# Compute backward
       if self.is_last:
           # Last stage computes gradients from loss and has no gradients from
           # next stage
           bwd_kwargs = {
               "stage_output": loss,
               "output_grads": None,
               "input_values": input_values,
           }
       else:
           # Otherwise, receive gradients from next stage
           grads_output = self._retrieve_recv_grads(bwd_chunk_id)
           # If an input to the pipeline requires gradient,
           # `torch.autograd.backward` will accumulate the gradient into the
           # `.grad` field of such input
           bwd_kwargs = {
               "stage_output": stage_output[0], 
               "output_grads": grads_output[0], 
               "input_values": input_values,
           }
```
Can you tell me if this is the correct way to solve it?

### Q4:  
After the above modifications, it runs successfully when the chunk size is not 1. I saw a comment in the argparser `the division result (i.e. chunk size) must be 1`. Could you tell me why it has such limitations?

Original codes:
https://github.com/pytorch/PiPPy/blob/main/examples/huggingface/pippy_gpt2.py#L92-L96

### Q5: 

Since torch.distributed.pipelining  is designed for GPipe, having a larger chunk number under the same batch size should increase GPU utilization by reducing idle time. Therefore, I expect the total processing time to decrease.

However, when I ran training experiments with 1 chunk and 4 chunks for 2 GPUs using PyTorch profiler to measure GPU time, I observed significantly longer times, particularly with NCCL time. Is this behavior expected?

Result is:
when BS32
1 GPU without pipeline: 1022ms
2 GPUs 1 chunk: GPU0, 5778ms;  GPU1, 1712ms
2 GPU 4 chunks: GPU0, 22830ms, GPU1, 4697ms


My codes are here: https://gist.github.com/YingLiGithub/8e53e1f54c5713f12003cfc4df1945a4




### Versions

nightly

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @chauhang @d4l3k