[BUG] Training speed is slower and consumes more memory than FSDP

Describe the bug
I am working on language model fine-tuning tasks using data parallel techniques. Other than deepspeed, I also use Fully Sharded Data Parallel from HuggingFace (https://huggingface.co/docs/accelerate/usage_guides/fsdp). Since the Zero optimizer of deepspeed has a lot of implementation in common with FSDP, I would expect it has similar performance compared to FSDP. But the benchmark results do not support this:

Falcon-7B 4k context length fine-tuning, 8xA100 40GB GPUs.

On the table, FSDP and Zero-3 seems to have similar inference speed, but Zero-3 consumes more memory. If offloading is enabled, their memory footprint is the same but zero-3 is significantly slower.

To Reproduce
Detailed code is attached DeepSpeed-Falcon-7B.zip
Steps to reproduce the behavior:

1. Download Falcon-7B pretrained weights:
   python download_model.py tiiuae/falcon-7b
2. Change context length of Falcon-7B to 4K by modify the 4th line of ./models/tiiuae_falcon-7b/tokenizer_config.json to "model_max_length": 4096,
3. Launch training using this command and wikitext dataset will be automatically downloaded as toy dataset:
   torchrun --nproc_per_node=8 --master_port=3045 run_clm.py falcon-7b-dp.json.

Expected behavior
Expect Deepspeed's Zero-3 optimizer has equal or even better performance than FSDP. Also expect parameter and optimizer offloading to be faster. (now it's 2.9s versus 13.7s)

ds_report output

[2023-07-07 20:34:33,705] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.0
 [WARNING]  using untested triton version (2.0.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/usr/local/lib/python3.8/dist-packages/torch']
torch version .................... 2.0.1+cu118
deepspeed install path ........... ['/usr/local/lib/python3.8/dist-packages/deepspeed']
deepspeed info ................... 0.9.5, unknown, unknown
torch cuda version ............... 11.8
torch hip version ................ None
nvcc version ..................... 11.8
deepspeed wheel compiled w. ...... torch 2.0, cuda 11.8

System info (please complete the following information):

• OS: Ubuntu 22.04
• GPU count and types: 8x A100 40G GPUs, one machine
• Interconnects (if applicable): None
• Python version: 3.8.10
• Any other relevant info about your setup: None

Launcher context
I use torchrun to launch single node training with multiple GPUs.
torchrun --nproc_per_node=8 --master_port=3045 run_clm.py configs/falcon-7b-dp.json

Docker context
I am using off-the-shelf docker image huggingface/transformers-pytorch-gpu:latest and install deepspeed using pip install deepspeed

Start with the attached DeepSpeed-Falcon-7B.zip and reproduce the command in run_clm.py using configs/falcon-7b-dp.json, then review the ds_report output and benchmark settings. Compare ZeRO stages and offload modes against the reported FSDP results; done means identifying and addressing the performance or memory gap with repeatable benchmark evidence.