Deepspeed script changes to align with transformers>=4.30.0 for LoRA training
diff --git a/docs/training.md b/docs/training.md
index df16922a47..628557ad5a 100644
--- a/docs/training.md
+++ b/docs/training.md
@@ -30,7 +30,7 @@ torchrun --nproc_per_node=4 --master_port=9778 fastchat/train/train_flant5.py \
 After training, please use our post-processing [function](https://github.com/lm-sys/FastChat/blob/55051ad0f23fef5eeecbda14a2e3e128ffcb2a98/fastchat/utils.py#L166-L185) to update the saved model weight. Additional discussions can be found [here](https://github.com/lm-sys/FastChat/issues/643).
 
 ### Fine-tuning using (Q)LoRA
-You can use the following command to train Vicuna-7B using QLoRA using ZeRO2. Note that ZeRO3 is not currently supported with QLoRA but ZeRO3 does support LoRA, which has a reference configuraiton under `playground/deepspeed_config_s3.json`.
+You can use the following command to train Vicuna-7B using QLoRA using ZeRO2. Note that ZeRO3 is not currently supported with QLoRA but ZeRO3 does support LoRA, which has a reference configuraiton under playground/deepspeed_config_s3.json. To use QLoRA, you must have bitsandbytes>=0.39.0 and transformers>=4.30.0 installed.
 ```bash
 deepspeed train_lora.py \
     --model_name_or_path ~/model_weights/llama-7b  \
diff --git a/fastchat/train/train_lora.py b/fastchat/train/train_lora.py
index 67a71f2e62..9481fa788b 100644
--- a/fastchat/train/train_lora.py
+++ b/fastchat/train/train_lora.py
@@ -103,13 +103,14 @@ def train():
     ) = parser.parse_args_into_dataclasses()
 
     device_map = None
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    ddp = world_size != 1
     if lora_args.q_lora:
-        world_size = int(os.environ.get("WORLD_SIZE", 1))
-        device_map = (
-            {"": int(os.environ.get("LOCAL_RANK") or 0)} if world_size != 1 else None
-        )
+        device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)} if ddp else None
         if len(training_args.fsdp) > 0 or deepspeed.is_deepspeed_zero3_enabled():
-            logging.warn("FSDP and ZeRO3 are both currently incompatible with QLoRA.")
+            logging.warning(
+                "FSDP and ZeRO3 are both currently incompatible with QLoRA."
+            )
 
     compute_dtype = (
         torch.float16
@@ -143,7 +144,7 @@ def train():
         model = prepare_model_for_kbit_training(
             model, use_gradient_checkpointing=training_args.gradient_checkpointing
         )
-        if torch.cuda.device_count() > 1:
+        if not ddp and torch.cuda.device_count() > 1:
             # keeps Trainer from trying its own DataParallelism when more than 1 gpu is available
             model.is_parallelizable = True
             model.model_parallel = True
@@ -178,7 +179,7 @@ def train():
     trainer.save_state()
 
     # check if zero3 mode enabled
-    if trainer.hf_deepspeed_config_orig.is_zero3():
+    if deepspeed.is_deepspeed_zero3_enabled():
         # use deepspeed engine internal function to gather state dict
         # state_dict_zero3 contains whole parameters of base and lora adapters
         # we will not extract lora parameters since peft save_pretrained will do that
diff --git a/playground/deepspeed_config_s2.json b/playground/deepspeed_config_s2.json
index aa350bb02c..3113cd3655 100644
--- a/playground/deepspeed_config_s2.json
+++ b/playground/deepspeed_config_s2.json
@@ -1,21 +1,12 @@
 {
-    "zero_optimization": {
-      "stage": 2,
-      "offload_optimizer": {
-              "device": "cpu"
-      },
-      "contiguous_gradients": true,
-      "overlap_comm": true
+  "zero_optimization": {
+    "stage": 2,
+    "offload_optimizer": {
+      "device": "cpu"
     },
-    "optimizer": {
-      "type": "AdamW",
-      "params": {
-          "lr": "auto",
-          "betas": "auto",
-          "eps": "auto",
-          "weight_decay": "auto"
-      }
+    "contiguous_gradients": true,
+    "overlap_comm": true
   },
-    "train_micro_batch_size_per_gpu": "auto",
-    "gradient_accumulation_steps": "auto"
-  }
\ No newline at end of file
+  "train_micro_batch_size_per_gpu": "auto",
+  "gradient_accumulation_steps": "auto"
+}
\ No newline at end of file
diff --git a/playground/deepspeed_config_s3.json b/playground/deepspeed_config_s3.json
index 629bfd8ade..07f4b16a66 100644
--- a/playground/deepspeed_config_s3.json
+++ b/playground/deepspeed_config_s3.json
@@ -1,4 +1,12 @@
 {
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
     "zero_optimization": {
         "stage": 3,
         "offload_optimizer": {
@@ -11,24 +19,14 @@
         },
         "overlap_comm": true,
         "contiguous_gradients": true,
-        "sub_group_size": "auto",
-        "reduce_bucket_size": "auto",
-        "stage3_prefetch_bucket_size": "auto",
-        "stage3_param_persistence_threshold": "auto",
-        "stage3_max_live_parameters": "auto",
-        "stage3_max_reuse_distance": "auto",
+        "stage3_max_live_parameters" : 1e9,
+        "stage3_max_reuse_distance" : 1e9,
+        "stage3_prefetch_bucket_size" : 5e8,
+        "stage3_param_persistence_threshold" : 1e6,
+        "sub_group_size" : 1e12,
         "stage3_gather_16bit_weights_on_model_save": true
     },
-    "optimizer": {
-        "type": "AdamW",
-        "params": {
-            "lr": "auto",
-            "betas": "auto",
-            "eps": "auto",
-            "weight_decay": "auto"
-        }
-    },
     "train_batch_size": "auto",
-    "micro_batch_size_per_gpu": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
     "gradient_accumulation_steps": "auto"
 }
\ No newline at end of file

## Why are these changes needed?
Since QLoRA relies on at least transformers version==4.30.0 and bitsandbytes==0.39.0, we propose some changes to the deepspeed and `train_lora` scripts to fix it.

We align deepspeed and trainer with transformers>=4.30.0 by doing the following: 
1. we must remove optimizer from the deepspeed config. Check [here](https://github.com/huggingface/transformers/issues/24359) for more information about supported combinations of optimizer and scheduler.
2. We also add the tag for `fp16` in the script to fix an input tensor mismatch [issue](https://github.com/microsoft/DeepSpeed/issues/3654).

This was tested on a distributed setup of 8 12GB GPUs with LLaMA-7B taking up ~5GB VRAM and training with effective training batch size of 1 taking around 10GB VRAM.

## Related issue number (if applicable)
#1741

## Checks

- [x] I've run `format.sh` to lint the changes in this PR.
- [x] I've included any doc changes needed.
- [ ] I've made sure the relevant tests are passing (if applicable).
