[fix] Set deepspeed's fp16 auto_cast to false

[maxreciprocate]

Fix of #238

https://github.com/huggingface/accelerate/blob/4e5cc0c6b9c28fc569bfc70dd0508225d187fbf9/src/accelerate/utils/dataclasses.py#L527

https://wandb.ai/sorry/trlx/reports/-fix-Set-deepspeed-s-fp16-auto_cast-to-false-279--VmlldzozNDkzMzkx

Thanks @ZHAOTING for finding this fix

[cat-state]

Thanks @reciprocated and @ZHAOTING . Does this increase memory usage? Or are the model weights still being cast to fp16 somewhere else? I see in the report that the bf16 and fp16 runs have about the same but theres a run which seems to have more memory use than the baseline:

[maxreciprocate]

No, it doesn't increase the memory usage. For the baseline I took by accident a prior bf16 run, while on the fix is a run without accelerate, so the difference here is between bf16 and fp32 (it's not possible to make a fp16 run on the main). Just in case I made a few more runs with fp32, zero-fp32, zero-fp16, zero-bf16:
https://wandb.ai/sorry/trlx/reports/Set-deepspeed-s-fp16-auto_cast-to-false-279--VmlldzozNDk0OTMz

[cat-state]

Ah okay, got it. This LGTM!

[Jiaxin-Wen]

it's not possible to make a fp16 run on the main

Hi, I have some related problems.
Specifically, I am running the example of summarize_rlhf. I observe that "mixed_preicision" is set to "no" in summarize_rlhf/configs/default_accelerate_config.yaml, while "fp16" is enabled in summarize_rlhf/configs/ds_config_trlx_gptj_summarize.json.

First, without any modification, the code will fail with the following error:

ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored: ['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device', 
'zero3_save_16bit_model', 'mixed_precision'].

And after I delete "mixed_precision" in summarize_rlhf/configs/default_accelerate_config.yaml, the code will again with the following error:

│ /data/wenjiaxin/home/trlx/trlx/trainer/nn/ppo_models.py:220 in loss                              │                                                                                                                               
│                                                                                                  │                                                                                                                               
│    217 │   │   │   │   value_loss=vf_loss.item(),                                                │                                                                                                                               
│    218 │   │   │   ),                                                                            │                                                                                                                               
│    219 │   │   │   values=dict(                                                                  │                                                                                                                               
│ ❱  220 │   │   │   │   get_tensor_stats(values, mask, n),                                        │                                                                                                                               
│    221 │   │   │   │   values_error=torch.sum(((values - returns) * mask) ** 2) / n,             │                                                                                                                               
│    222 │   │   │   │   clipfrac=vf_clipfrac,                                                     │                                                                                                                               
│    223 │   │   │   ),                                                                            │                                                                                                                               
│                                                                                                  │                                                                                                                               
│ /data/wenjiaxin/home/trlx/trlx/utils/modeling.py:242 in get_tensor_stats                         │                                                                                                                               
│                                                                                                  │                                                                                                                               
│   239 │   mean = (xs * mask).sum() / n                                                           │                                                                                                                               
│   240 │   return dict(                                                                           │                                                                                                                               
│   241 │   │   mean=mean,                                                                         │                                                                                                                               
│ ❱ 242 │   │   min=torch.where(mask.bool(), xs, np.inf).min(),                                    │                                                                                                                               
│   243 │   │   max=torch.where(mask.bool(), xs, -np.inf).max(),                                   │                                                                                                                               
│   244 │   │   std=torch.sqrt(((xs - mean) * mask).pow(2).sum() / n),                             │                                                                                                                               
│   245 │   )                                                                                      │                                                                                                                               
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯                                                                                                                               
RuntimeError: expected scalar type c10::Half but found double   

[maxreciprocate]

Hi! @XWwwwww, I'm able to reproduce the first part, it's related to the newest accelerate release which started to raise errors in this case, however I'm not getting the runtime error you're having. Maybe it's related to some particular versioning (at least that's what a quick google search suggests). Are you having the same problem for other examples as well?

[Jiaxin-Wen]

Hi! @reciprocated
For the second problem, I fix it as follows:

tmp = np.inf * torch.ones_like(xs)
min=torch.where(mask.bool(), xs, tmp).min(),
max=torch.where(mask.bool(), xs, -tmp).max(),

and my torch version is 1.10.1

[PhungVanDuy]

Can you try to upgrade your pytorch version? ref: #264 (comment)