`"histogram_cpu" not implemented for 'BFloat16'` when using deepspeed and reporting to wandb
## Environment info
- `transformers` version: 4.18.0
- Platform: Linux-5.13.0-20-generic-x86_64-with-glibc2.10
- Python version: 3.8.5
- Huggingface_hub version: 0.5.1
- PyTorch version (GPU?): 1.11.0+cu113 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: Deepspeed

### Who can help
 @stas00

## Information

Model I am using (Bert, XLNet ...): bart-large

The problem arises when using:
* [X] the official example scripts: (give details below)
* [X] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: (give details below)

## To reproduce

I'm using a training script adapted from the run_summarization.py example with a model using bart-large architecture and a custom tokenizer. I'm working locally on my workstation with two RTX 3090s. I had been training using deepspeed and fp16, but I saw that the latest transformers update added bf16 support to the deepspeed integration, so I wanted to try that in order to reduce the constant overflow errors I had been getting.

But when using deepspeed, bf16, and reporting to wandb, my training crashes. 

I'm able to reproduce the error using the example scripts:
```
deepspeed run_summarization.py \
      --model_name_or_path facebook/bart-large \
      --dataset_name cnn_dailymail --dataset_config_name 3.0.0 \
      --do_train --per_device_train_batch_size 4 --bf16  \
      --overwrite_output_dir --output_dir models/text_summarization \
      --deepspeed config/deepspeed_config-zero2-bf16.json 
```

with the deepspeed config being:
```
{
  "bf16": {
     "enabled": true
   },

  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": "auto",
      "weight_decay": "auto"
    }
  },

  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto"
    }
  },

  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true,
    "cpu_offload": false
  },

  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "steps_per_print": 2000,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false
}
```

After 500 steps (when saving the first checkpoint), wandb throws this error:
`RuntimeError: "histogram_cpu" not implemented for 'BFloat16'`

The error doesn't occur if I run the same script without deepspeed. And no other error gets thrown if I use deepspeed and don't report to wandb.

[A very similar issue was reported to wandb last month](https://github.com/wandb/client/issues/3332). The wandb people say it's an issue with pytorch and not wandb, but since everything is working without deepspeed, maybe there's something different about how the deepspeed integration is reporting to wandb?

## Expected behavior

The training should continue without crashing, and should report as much info to wandb as possible (not sure if there are limits to that introduced by bf16 )

Hey @jncasey, thanks for the report.

> since everything is working without deepspeed, maybe there's something different about how the deepspeed integration is reporting to wandb?

Could you check that the script still breaks when you set `"bf16"` to `"enabled": false` in the DeepSpeed config? Also, a full traceback would be helpful!<|||||>Updating the DeepSpeed config to set `"bf16" { "enabled": false }` raises this error, as one might expect:
```
ValueError: Please correct the following DeepSpeed config values that mismatch TrainingArguments values:
- ds bf16.enabled=False vs hf bf16|bf16_full_eval=True
```
Setting `"bf16" { "enabled": "auto" }` yields the original error I reported.

Here's the full traceback:
```
  0%|▎                                                      | 499/107670 [03:39<13:13:26,  2.25it/s]Traceback (most recent call last):
  File "./bin/run_summarization.py", line 706, in <module>
    main()
  File "./bin/run_summarization.py", line 625, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/transformers/trainer.py", line 1422, in train
    tr_loss_step = self.training_step(model, inputs)
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/transformers/trainer.py", line 2027, in training_step
    loss = self.deepspeed.backward(loss)
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1667, in backward
    self.optimizer.backward(loss)
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1921, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 53, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/torch/_tensor.py", line 363, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/wandb/wandb_torch.py", line 266, in <lambda>
    handle = var.register_hook(lambda grad: _callback(grad, log_track))
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/wandb/wandb_torch.py", line 264, in _callback
    self.log_tensor_stats(grad.data, name)
  File "/opt/miniconda3/envs/hf/lib/python3.8/site-packages/wandb/wandb_torch.py", line 215, in log_tensor_stats
    tensor = flat.histc(bins=self._num_bins, min=tmin, max=tmax)
RuntimeError: "histogram_cpu" not implemented for 'BFloat16'
```<|||||>Please help me understand this issue - why are we discussing a problem in wandb at `transformers`? Clearly wandb can't handle bf16 inputs in at least one code path - what does it have to do with deepspeed or transformers?

The next logical step is to either have `wandb` workaround bf16 if it can't handle it and not use `histogram_cpu` for bf16 inputs or ask pytorch to implement it for bf16/cpu.<|||||>Hi Stas,

My thought was that since the same script can report to wandb using bf16 when not using DeepSpeed, there might be something different in how the DeepSpeed integration handles the reporting, and it might be possible to avoid the problem the way the non-DeepSpeed run does. 

But I admit I'm out of my depth here, so maybe m thinking is flawed and there's nothing to be done on the transformers side.<|||||>The difference is that normally w/o deepspeed you're using bf16/amp, which keeps the model and activations in fp32 and downcasts them to bf16 when needed.

Deepspeed doesn't use amp and uses a different approach where it keeps the model and activations in the half precision mode from the get going (fp16 or bf16) (but keeps a fp32 weights copy in its optimizer), and so it trips wandb's code which doesn't expect bf16 tensors.

It's not something that can be changed in Deepspeed or the HF integration.

Possible solutions:

For example, if the bf16 input < 64k, wandb could make a copy and safely convert it to fp16 and run `histogram_cpu` on it (I assume it supports fp16), if not perhaps it could do the processing on gpu.

But the simplest solution is to request pytorch to implement `histogram_cpu` for `BFloat16` inputs, which you or the wandb folks can ask via a "feature request" at https://github.com/pytorch/pytorch/issues/new/choose but which of course will take time. So possibly both solutions could be used together.



<|||||>Got it! Thanks for the super clear explanation.<|||||>Late update: This seems to be fixed with the release of pytorch 1.12<|||||>Super! Thank you for the update, @jncasey <|||||>I'm still having this bug with torch 0.13.0+cu116 (`RuntimeError: "histogram_cpu" not implemented for 'Byte'`)<|||||>> I'm still having this bug with torch 0.13.0+cu116 (`RuntimeError: "histogram_cpu" not implemented for 'Byte'`)

@julien-blanchon Did you solve it? 

I am getting this error when reporting to wandb using `torch==2.0.1`, `wandb==0.15.5` and `accelerate==0.20.3`.<|||||>No, I'm still waiting for a patch here: https://github.com/pytorch/pytorch/issues/75667