Distributed Concat (Nested Gather - HF Trainer): Support for Mapping/Dict
Author: prajwal967
Pull request: no

### Feature request

* The request is with respect to the [distributed_concat](https://github.com/huggingface/transformers/blob/b9af152efb748b1bff8f6fe0130e62ebb8e11a53/src/transformers/trainer_pt_utils.py#L188) function call in the [nested gather](https://github.com/huggingface/transformers/blob/b9af152efb748b1bff8f6fe0130e62ebb8e11a53/src/transformers/trainer.py#L3182) function in the HuggingFace Trainer.

* The feature request: Addition of an additional **IF** statement to check if the **input (tensor)** is of type **Mapping/Dict**.

* The [distributed_concat](https://github.com/huggingface/transformers/blob/b9af152efb748b1bff8f6fe0130e62ebb8e11a53/src/transformers/trainer_pt_utils.py#L188) function currently checks for the **List/Tuple** type, but does not check for the **Mapping/Dict** type.

```
  def distributed_concat(tensor: Any, num_total_examples: Optional[int] = None) -> Any:
      try:
          if isinstance(tensor, (tuple, list)):
              return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
          tensor = atleast_1d(tensor).contiguous()
          output_tensors = [tensor.clone() for _ in range(dist.get_world_size())]
          dist.all_gather(output_tensors, tensor)
          concat = torch.cat(output_tensors, dim=0)
  
          # truncate the dummy elements added by SequentialDistributedSampler
          if num_total_examples is not None:
              concat = concat[:num_total_examples]
          return concat
      except AssertionError:
          raise AssertionError("Not currently using distributed training")
  ```

### Motivation

* The [nested_xla_mesh_reduce](https://github.com/huggingface/transformers/blob/b9af152efb748b1bff8f6fe0130e62ebb8e11a53/src/transformers/trainer_pt_utils.py#L171) & [smp_gather](https://github.com/huggingface/transformers/blob/b9af152efb748b1bff8f6fe0130e62ebb8e11a53/src/transformers/trainer_pt_utils.py#L1084) functions have an **IF** statement to check if the input is of type **Mapping/Dict** and then iterates over the items. The code runs without throwing any errors.

* For example in [nested_xla_mesh_reduce](https://github.com/huggingface/transformers/blob/b9af152efb748b1bff8f6fe0130e62ebb8e11a53/src/transformers/trainer_pt_utils.py#L171):

  ```
  def nested_xla_mesh_reduce(tensors, name):
      if is_torch_tpu_available():
          import torch_xla.core.xla_model as xm
  
          if isinstance(tensors, (list, tuple)):
              return type(tensors)(nested_xla_mesh_reduce(t, f"{name}_{i}") for i, t in enumerate(tensors))
          if isinstance(tensors, Mapping):
              return type(tensors)(
                  {k: nested_xla_mesh_reduce(t, f"{name}_{i}") for i, (k, t) in enumerate(tensors.items())}
              )
  
          tensors = atleast_1d(tensors)
          return xm.mesh_reduce(name, tensors, torch.cat)
      else:
          raise ImportError("Torch xla must be installed to use `nested_xla_mesh_reduce`")
  ```


* The [distributed_concat](https://github.com/huggingface/transformers/blob/b9af152efb748b1bff8f6fe0130e62ebb8e11a53/src/transformers/trainer_pt_utils.py#L188) function currently checks for the **List/Tuple** type, but does not check for the **Mapping/Dict** type. 

  ```
  def distributed_concat(tensor: Any, num_total_examples: Optional[int] = None) -> Any:
      try:
          if isinstance(tensor, (tuple, list)):
              return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
          tensor = atleast_1d(tensor).contiguous()
          output_tensors = [tensor.clone() for _ in range(dist.get_world_size())]
          dist.all_gather(output_tensors, tensor)
          concat = torch.cat(output_tensors, dim=0)
  
          # truncate the dummy elements added by SequentialDistributedSampler
          if num_total_examples is not None:
              concat = concat[:num_total_examples]
          return concat
      except AssertionError:
          raise AssertionError("Not currently using distributed training")
  ```
  
* The same code and data (which contains labels of type **Mapping/Dict**) run without any errors when using torch/cuda/python. The nested_xla_mesh_reduce function is called and no errors are thrown.

* The same code however throws an error when I run it using deepspeed, since the [distributed_concat](https://github.com/huggingface/transformers/blob/b9af152efb748b1bff8f6fe0130e62ebb8e11a53/src/transformers/trainer_pt_utils.py#L188) function is called, and the items of the **Mapping/Dict** object are not appropriately gathered and converted.

* The fix seems to be simple, I currently got the code to work by adding the following code to the [distributed_concat](https://github.com/huggingface/transformers/blob/b9af152efb748b1bff8f6fe0130e62ebb8e11a53/src/transformers/trainer_pt_utils.py#L188) function:

  ```
  if isinstance(tensor, Mapping):
    return type(tensor)(
      {k: distributed_concat(t, num_total_examples) for k, t in tensor.items()}
     )
  ```

### Your contribution

I could submit a **PR** with the following lines of code added to the [distributed_concat](https://github.com/huggingface/transformers/blob/b9af152efb748b1bff8f6fe0130e62ebb8e11a53/src/transformers/trainer_pt_utils.py#L188) function:

  ```
  if isinstance(tensor, Mapping):
    return type(tensor)(
      {k: distributed_concat(t, num_total_examples) for k, t in tensor.items()}
     )
  ```
  
  **The entire function would look something like this:**
  ```
def distributed_concat(tensor: Any, num_total_examples: Optional[int] = None) -> Any:
    try:
        if isinstance(tensor, (tuple, list)):
            return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
        if isinstance(tensor, Mapping):
            return type(tensor)(
                {k: distributed_concat(t, num_total_examples) for k, t in tensor.items()}
            )
        
        tensor = atleast_1d(tensor).contiguous()
        output_tensors = [tensor.clone() for _ in range(dist.get_world_size())]
        dist.all_gather(output_tensors, tensor)
        concat = torch.cat(output_tensors, dim=0)

        # truncate the dummy elements added by SequentialDistributedSampler
        if num_total_examples is not None:
            concat = concat[:num_total_examples]
        return concat
    except AssertionError:
        raise AssertionError("Not currently using distributed training")
```

Comments:
Happy to review a PR if you want to make one :-) 

Sure, will do! Thanks!

I just raised a PR (https://github.com/huggingface/transformers/pull/21500). I haven't raised a PR before for HuggingFace. Let me know if I need to change or fix anything or if I forgot to do something. Thanks!