From 6fb70a82dfc5af0a33546b131d936eb8bb0afb73 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Thu, 13 Oct 2022 13:03:01 +0800 Subject: [PATCH] [ORTModule] Update Supported DeepSpeed Version for FP16_Optimizer (#13305) Update supported deepspeed highest version from 0.7.1 to 0.7.3 for FP16_Optimizer. Also add version info to warning log. --- .../python/training/optim/_ds_modifier.py | 20 +++++++++++++++---- .../python/orttraining_test_ortmodule_api.py | 1 + 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/orttraining/orttraining/python/training/optim/_ds_modifier.py b/orttraining/orttraining/python/training/optim/_ds_modifier.py index 3d75ff3242..9e3470d33c 100644 --- a/orttraining/orttraining/python/training/optim/_ds_modifier.py +++ b/orttraining/orttraining/python/training/optim/_ds_modifier.py @@ -12,10 +12,10 @@ import types import warnings -from distutils.version import LooseVersion import torch from numpy import inf +from packaging.version import Version from ._modifier import FP16OptimizerModifier, check_overflow, check_overflow_for_grads from ._multi_tensor_apply import MultiTensorApply @@ -30,9 +30,21 @@ class DeepSpeedZeROModifier(FP16OptimizerModifier): def can_be_modified(self): import deepspeed - ds_version = LooseVersion(deepspeed.__version__) - if ds_version > LooseVersion("0.7.1") or ds_version < LooseVersion("0.4.0"): - warnings.warn("Skip modifying optimizer because of unsupported DeepSpeed version.", UserWarning) + # This modifier relies on the implementation of has_overflow_serial, get_grad_norm_direct, + # and has_overflow_partitioned_grads_serial + # in https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage_1_and_2.py. + # Everytime if we want to update this version supporting list to a newer version, + # we need to check if the implementation of these functions are changed. + # An easy way to check is to check the history of this file, if there is no change during the update, + # it's safe to update the version supporting list. Otherwise, or the file is moved or renamed, + # we need to check the implementation of these functions in detail. + ds_version = Version(deepspeed.__version__) + if ds_version > Version("0.7.3") or ds_version < Version("0.4.0"): + warnings.warn( + "Skip modifying optimizer because of unsupported DeepSpeed version {}, " + "supported version: 0.4.0 - 0.7.3.".format(deepspeed.__version__), + UserWarning, + ) return False return self.check_requirements( diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index f6e0d1f57a..cef9c75276 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -1318,6 +1318,7 @@ def test_gradient_correctness_reducesum(dim, keepdim): loss.backward() return prediction + torch.manual_seed(2333) for _ in range(10): pt_input = torch.rand((N, D, H), device=device, requires_grad=True) ort_input = copy.deepcopy(pt_input)