deepspeed resume from ckpt fixes and adding support for deepspeed optimizer and HF scheduler #25863

pacman100 · 2023-08-30T11:41:16Z

What does this PR do?

Add support for deep speed optimizer and HF scheduler. Should be merged after Add support for deepspeed optimizer and custom scheduler accelerate#1909
Fixing the lr scheduler not being saved by passing them to the DeepSpeed engine for all schedulers that are instances of LRScheduler. Should be merged after Add support for deepspeed optimizer and custom scheduler accelerate#1909

Below we will run the 4 combinations of optimizer and schedulers for the run_glue.py transformers example
Initial setup:

cd transformers
export CUDA_VISISBLE_DEVICES=0,1
export TASK_NAME=mrpc

a. HF Optimizer + HF Scheduler Case:

i. ds config ds_config_z3_hf_optim_hf_scheduler.json:

{
  "fp16": {
    "enabled": "auto",
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },

  "bf16": {
    "enabled": "auto"
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "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
}

ii. command to run:

torchrun --nnodes 1 --nproc-per-node 2 ./examples/pytorch/text-classification/run_glue.py --model_name_or_path bert-base-cased  --task_name $TASK_NAME  --do_train  --do_eval  --max_seq_length 128  --per_device_train_batch_size 16  --learning_rate 5e-5  --num_train_epochs 3  --output_dir /tmp/$TASK_NAME/ --overwrite_output_dir --deepspeed ~/transformers/tests/deepspeed/ds_config_z3_hf_optim_hf_scheduler.json --lr_scheduler_type cosine --save_strategy "epoch" --evaluation_strategy "epoch" --logging_steps 1

Kill the process after epoch 1. run the above command with --resume_from_checkpoint as below:

torchrun --nnodes 1 --nproc-per-node 2 ./examples/pytorch/text-classification/run_glue.py --model_name_or_path bert-base-cased  --task_name $TASK_NAME  --do_train  --do_eval  --max_seq_length 128  --per_device_train_batch_size 16  --learning_rate 5e-5  --num_train_epochs 3  --output_dir /tmp/$TASK_NAME/ --overwrite_output_dir --deepspeed ~/transformers/tests/deepspeed/ds_config_z3_hf_optim_hf_scheduler.json --lr_scheduler_type cosine --save_strategy "epoch" --evaluation_strategy "epoch" --logging_steps 1 --resume_from_checkpoint /tmp/$TASK_NAME/checkpoint-115/

iii. Plots of loss and learning rate:

a. DS Optimizer + DS Scheduler Case:
i. ds config ds_config_z3_ds_optim_ds_scheduler.json:

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "bf16": {
        "enabled": "auto"
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
	    "total_num_steps": "auto", 
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },

    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "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
}

rest of the steps as above. Plots:

c. HF Optimizer + DS Scheduler Case:
i. ds config ds_config_z3_hf_optim_ds_scheduler.json:

{
  "fp16": {
    "enabled": "auto",
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },

  "bf16": {
    "enabled": "auto"
  },

  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
	  "total_num_steps": "auto", 
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto"
    }
  },

  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "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
}

rest of the steps as above. Plots:

c. DS Optimizer + HF Scheduler Case:
i. ds config ds_config_z3_ds_optim_hf_scheduler.json:

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "bf16": {
        "enabled": "auto"
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "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
}

rest of the steps as above. Plots:

Adding tests to check the resume from ckpt is working properly with DeepSpeed.

HuggingFaceDocBuilderDev · 2023-08-30T12:08:40Z

The documentation is not available anymore as the PR was closed or merged.

… scenario

BenjaminBossan

Thanks a lot for tackling the issue you described and conducting (+ showing) the experiments you ran to prove that it works. Personally, I miss the experience with deepspeed required to understand the bigger picture, so I cannot provide a full on review, only some small comments.

tests/deepspeed/test_deepspeed.py

src/transformers/integrations/deepspeed.py

muellerzr

Looks good to me, thanks! Let's definitely keep an eye out for pickle problems, and be prepared to move that to a util if needed

…imizer and HF scheduler (huggingface#25863) * Add support for deepspeed optimizer and HF scheduler * fix bug * fix the import * fix issue with deepspeed scheduler saving for hf optim + hf scheduler scenario * fix loading of hf scheduler when loading deepspeed checkpoint * fix import of `DeepSpeedSchedulerWrapper` * add tests * add the comment and skip the failing tests * address comment

…imizer and HF scheduler (#25863) * Add support for deepspeed optimizer and HF scheduler * fix bug * fix the import * fix issue with deepspeed scheduler saving for hf optim + hf scheduler scenario * fix loading of hf scheduler when loading deepspeed checkpoint * fix import of `DeepSpeedSchedulerWrapper` * add tests * add the comment and skip the failing tests * address comment

…imizer and HF scheduler (huggingface#25863) * Add support for deepspeed optimizer and HF scheduler * fix bug * fix the import * fix issue with deepspeed scheduler saving for hf optim + hf scheduler scenario * fix loading of hf scheduler when loading deepspeed checkpoint * fix import of `DeepSpeedSchedulerWrapper` * add tests * add the comment and skip the failing tests * address comment

sunyclj · 2023-12-21T02:16:26Z

DS Optimizer + HF Scheduler
Use the configuration in DS Optimizer + HF Scheduler ,
error : ValueError: You cannot create a DummyScheduler without specifying a scheduler in the config file.
It may have something to do with the version of accelerate. my version of accelerate is 0.22.0，what's your version, please.

Add support for deepspeed optimizer and HF scheduler

397dcad

pacman100 added 2 commits August 31, 2023 15:06

fix bug

c19567a

fix the import

2bcc4e6

pacman100 mentioned this pull request Aug 31, 2023

Deepspeed lr_schedule is None #25865

Closed

4 tasks

pacman100 changed the title ~~Add support for deepspeed optimizer and HF scheduler~~ deepspeed resume from ckpt fixes and adding support for deepspeed optimizer and HF scheduler Aug 31, 2023

pacman100 mentioned this pull request Aug 31, 2023

ValueError: Found optimizer configured in the DeepSpeed config, but no scheduler. Please configure a scheduler in the DeepSpeed config. #24359

Closed

pacman100 added 7 commits September 1, 2023 15:25

Merge branch 'main' into smangrul/ds-fixes-and-tests

c4d0530

Merge branch 'main' into smangrul/ds-fixes-and-tests

51d720d

fix issue with deepspeed scheduler saving for hf optim + hf scheduler…

224892a

… scenario

fix loading of hf scheduler when loading deepspeed checkpoint

577150f

fix import of DeepSpeedSchedulerWrapper

2cb4088

add tests

d77abfa

add the comment and skip the failing tests

2d509d1

pacman100 marked this pull request as ready for review September 4, 2023 14:42

pacman100 requested review from muellerzr and BenjaminBossan September 4, 2023 14:44

pacman100 added 2 commits September 5, 2023 11:04

Merge branch 'main' into smangrul/ds-fixes-and-tests

83971a1

Merge branch 'main' into smangrul/ds-fixes-and-tests

44892b7

BenjaminBossan approved these changes Sep 5, 2023

View reviewed changes

tests/deepspeed/test_deepspeed.py Outdated Show resolved Hide resolved

src/transformers/integrations/deepspeed.py Show resolved Hide resolved

address comment

76ae056

muellerzr approved these changes Sep 5, 2023

View reviewed changes

pacman100 merged commit 6bc517c into main Sep 5, 2023

pacman100 deleted the smangrul/ds-fixes-and-tests branch September 5, 2023 17:01

pacman100 mentioned this pull request Sep 7, 2023

fix the deepspeed tests #26021

Merged

VeryLazyBoy mentioned this pull request Oct 16, 2023

Support custom scheduler in deepspeed training #26831

Merged

5 tasks

jklj077 mentioned this pull request Dec 20, 2023

[BUG] <title>请问多GPU微调训练，学习率为什么不变呢？ QwenLM/Qwen#672

Closed

2 tasks

dwyatte mentioned this pull request Dec 23, 2023

Update docs around mixing hf scheduler with deepspeed optimizer #28223

Merged

5 tasks

VeryLazyBoy mentioned this pull request Dec 24, 2023

Skip saving frozen parameters if using peft model with deepspeed #26503

Closed

5 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

deepspeed resume from ckpt fixes and adding support for deepspeed optimizer and HF scheduler #25863

deepspeed resume from ckpt fixes and adding support for deepspeed optimizer and HF scheduler #25863

pacman100 commented Aug 30, 2023 •

edited

Loading

HuggingFaceDocBuilderDev commented Aug 30, 2023 •

edited

Loading

BenjaminBossan left a comment

muellerzr left a comment

sunyclj commented Dec 21, 2023

deepspeed resume from ckpt fixes and adding support for deepspeed optimizer and HF scheduler #25863

deepspeed resume from ckpt fixes and adding support for deepspeed optimizer and HF scheduler #25863

Conversation

pacman100 commented Aug 30, 2023 • edited Loading

What does this PR do?

HuggingFaceDocBuilderDev commented Aug 30, 2023 • edited Loading

BenjaminBossan left a comment

Choose a reason for hiding this comment

muellerzr left a comment

Choose a reason for hiding this comment

sunyclj commented Dec 21, 2023

pacman100 commented Aug 30, 2023 •

edited

Loading

HuggingFaceDocBuilderDev commented Aug 30, 2023 •

edited

Loading