Enables ZeRO-3 inference

[jeffra]

This enables a ZeRO-3 inference support, this means no optimizer is specified. This supports ZeRO-3 w. multiple GPUs and also supports ZeRO-3 w. parameter CPU offload.

This enables initial functional support, we have not fully evaluated performance and all memory reduction scenarios. More PRs to come I imagine :)

[stas00]

It definitely uses less GPU memory w/ offload configured as compared to not.

2 gpus, w/ offload:

export BS=16; rm -r output_dir; PYTHONPATH=src USE_TF=0 deepspeed --num_gpus=2 examples/pytorch/translation/run_translation.py --model_name_or_path t5-small --output_dir output_dir --adam_eps 1e-06 --evaluation_strategy=steps --do_eval --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 500 --max_source_length 128 --max_target_length 128  --overwrite_output_dir --per_device_eval_batch_size $BS --predict_with_generate --sortish_sampler --source_lang en --target_lang ro --dataset_name wmt16 --dataset_config "ro-en" --source_prefix "translate English to Romanian: " --val_max_target_length 128 --warmup_steps 50 --max_eval_samples 50 --deepspeed tests/deepspeed/ds_config_zero3.json  --fp16 --skip_memory_metrics 0

  before_init_mem_cpu       =     5552MB
  before_init_mem_gpu       =       32MB
  eval_bleu                 =    28.1872
  eval_gen_len              =      34.88
  eval_loss                 =     3.6001
  eval_mem_cpu_alloc_delta  =     2175MB
  eval_mem_cpu_peaked_delta =        0MB
  eval_mem_gpu_alloc_delta  =        0MB
  eval_mem_gpu_peaked_delta =      264MB
  eval_runtime              = 0:00:16.08
  eval_samples              =         50
  eval_samples_per_second   =      3.108
  eval_steps_per_second     =      0.124
  init_mem_cpu_alloc_delta  =        5MB
  init_mem_cpu_peaked_delta =        0MB
  init_mem_gpu_alloc_delta  =        0MB
  init_mem_gpu_peaked_delta =        0MB

w/o offload (same cmd and config but changed "device": "none",)

  before_init_mem_cpu       =     5428MB
  before_init_mem_gpu       =      106MB
  eval_bleu                 =    28.1872
  eval_gen_len              =      34.88
  eval_loss                 =     3.6001
  eval_mem_cpu_alloc_delta  =      668MB
  eval_mem_cpu_peaked_delta =        0MB
  eval_mem_gpu_alloc_delta  =      332MB
  eval_mem_gpu_peaked_delta =      264MB
  eval_runtime              = 0:00:30.00
  eval_samples              =         50
  eval_samples_per_second   =      1.666
  eval_steps_per_second     =      0.067
  init_mem_cpu_alloc_delta  =        3MB
  init_mem_cpu_peaked_delta =        0MB
  init_mem_gpu_alloc_delta  =        0MB
  init_mem_gpu_peaked_delta =        0MB

the change was:

--- a/tests/deepspeed/ds_config_zero3.json
+++ b/tests/deepspeed/ds_config_zero3.json
@@ -30,11 +30,11 @@
     "zero_optimization": {
         "stage": 3,
         "offload_optimizer": {
-            "device": "cpu",
+            "device": "none",
             "pin_memory": true
         },
         "offload_param": {
-            "device": "cpu",
+            "device": "none",
             "pin_memory": true
         },
         "overlap_comm": true,