FP32 Mode for ZeRO about deepspeed

[saareliad]

Hi, what is the status for this?
I want to compare ZeRo in FP32 to my system.

There are several use-cases for FP32 where it can use less memory compared to mixed precision since in mixed precision we still keep master parameters. (so for example given 44GB model in FP32, mixed precision will hold 66GB just for the model).
if the activations are small enough than mixed precision may not be good choice (at least memory wise).

[samyam]

@LvanderGoten, ZeRO should work with FP32.

Please disable the fp16 flag via

"fp16": {
"enabled": false
}
and also turn on the fp32_allreduce flag.

Since we implemented ZeRO with fp16 in mind, some of the variable names may not make sense for fp32, and there will be some redundant data copies and memory allocation, but you should still be able to see up to 2x memory reduction with ZeRO in fp32.

Please give it a try and let us know if you run into any issues.

[LvanderGoten]

I get the following:

Traceback (most recent call last):
  File "identification/generative/CPGAN/train_deepspeed.py", line 186, in <module>
    main()
  File "identification/generative/CPGAN/train_deepspeed.py", line 179, in main
    save_every_n_steps=args.save_every_n_steps)
  File "/home/deepspeed/Code/identification/generative/CPGAN/orchestrator_msggan_deepspeed.py", line 250, in train
    training_data=data)
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/__init__.py", line 95, in initialize
    collate_fn=collate_fn)
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/pt/deepspeed_light.py", line 126, in __init__
    self._configure_with_arguments(args, mpu)
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/pt/deepspeed_light.py", line 324, in _configure_with_arguments
    self._config = DeepSpeedConfig(args.deepspeed_config, mpu)
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/pt/deepspeed_config.py", line 243, in __init__
    self._do_sanity_check()
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/pt/deepspeed_config.py", line 359, in _do_sanity_check
    self._do_error_check()
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/pt/deepspeed_config.py", line 379, in _do_error_check
    assert self.fp16_enabled, "DeepSpeedConfig: ZeRO is only supported if fp16 is enabled"
AssertionError: DeepSpeedConfig: ZeRO is only supported if fp16 is enabled

Using this config:

{
  "train_batch_size": 8,
  "gradient_accumulation_steps": 1,
  "steps_per_print": 1,
  "zero_optimization": true,
  "fp32_allreduce": true,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.0001
    }
  },

  "fp16": {
    "enabled": false
  }
}

[jeffra]

We haven't tested ZeRO with fp32 so there may be a few issues along the way here :) If you want to try it out I would comment out the these lines to disable this check: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/pt/deepspeed_config.py#L378-L379

[jeffra]

Feel free to re-open if you run into other issues or let us know if this solved your problem.