[BUG] 10x Performance Degradation for Stage 3 vs Stage 2

[yukw777]

Describe the bug

This is a follow up from https://github.com/microsoft/DeepSpeed/issues/1069#issuecomment-1518294721. @tjruwase

I'm in the process of reproducing Alpaca-LoRA, and I've noticed a 10x performance degradation when using Stage 3 compared to Stage 2, which I believe is higher than usual.

To Reproduce

1. Clone https://github.com/yukw777/stanford_alpaca and set it up, e.g., install Python dependencies in requirements.txt.
2. Train Alpaca-LoRA-7B using Stage 2 by running the following command. Note that you can modify this code block to initialize a random LLaMA model.

echo '{
    "bf16": {
        "enabled": "auto"
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "none",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}' > ds_zero_2.json

torchrun --nproc_per_node=4 train.py \
--model_name_or_path <path_to_converted_llama_weights> \
--output_dir <path_to_output_dir> \
--num_train_epochs 10 \
--learning_rate 3e-4 \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 1 \
--bf16 True \
--use_lora True \
--warmup_steps 100 \
--evaluation_strategy steps \
--eval_steps 200 \
--save_strategy steps \
--save_steps 200 \
--save_total_limit 3 \
--group_by_length True \
--logging_steps 10 \
--deepspeed ds_zero_2.json

3. Train Alpaca-LoRA-7B using Stage 3 by running the following command.

echo '{
    "bf16": {
        "enabled": "auto"
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "none",
            "pin_memory": true
        },
        "offload_param": {
            "device": "none",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}' > ds_zero_3.json

torchrun --nproc_per_node=4 train.py \
--model_name_or_path /nfs/turbo/coe-chaijy/pre-trained-weights/LLaMA-hf/7B \
--output_dir output/Alpaca-LoRA/7B \
--num_train_epochs 10 \
--learning_rate 3e-4 \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 1 \
--bf16 True \
--use_lora True \
--warmup_steps 100 \
--evaluation_strategy steps \
--eval_steps 200 \
--save_strategy steps \
--save_steps 200 --save_total_limit 3 \
--group_by_length True \
--logging_steps 10 \
--deepspeed ds_zero_3.json

4. You'll notice that training using Stage 3 is about 10x slower than Stage 2. In my setup, Stage 3 took 2.72s per iteration, while Stage 2 took about 0.29s per iteration (3.49 iterations per second).

Expected behavior I'm not sure what'd be the correct expected behavior, but I'm hoping it's stage 3 being faster than 10x Stage 2.

ds_report output

$ ds_report
--------------------------------------------------
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 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]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/kpyu/stanford_alpaca/.venv/lib/python3.10/site-packages/torch']
torch version .................... 2.0.0+cu117
deepspeed install path ........... ['/home/kpyu/stanford_alpaca/.venv/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.9.0, b7abafb, main
torch cuda version ............... 11.7
torch hip version ................ None
nvcc version ..................... 11.8
deepspeed wheel compiled w. ...... torch 0.0, cuda 0.0

Screenshots N/A

System info (please complete the following information):

• OS: [e.g. Ubuntu 18.04] Red Hat Enterprise Linux release 8.4 (Ootpa)
• GPU count and types [e.g. two machines with x8 A100s each] One machine with 4 A40 48GB GPUs
• Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB] One machine with 4 A40 48GB GPUs connected via PCIe.
• Python version 3.10
• Any other relevant info about your setup I'm using a Slurm cluster.

Launcher context torchrun

Docker context N/A

Additional context I thought it might be related to the fact that I have some frozen layers (i.e. LoRA), so I tried to reproduce it without LoRA, but unfortunately, I got CUDA OOM for training Alpaca-7B with Stage 2, so I couldn't perform an apples-to-apples comparison. You can run the training script without LoRA by simply setting --use_lora to False.