Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ORT training support stage3 #1439

Closed
wants to merge 7 commits into from
Closed

ORT training support stage3 #1439

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
a7f4372
enable stage3
pengwa Oct 8, 2023
a81a59a
fix
pengwa Oct 8, 2023
6c623f8
fix the check
pengwa Oct 21, 2023
f5d241e
Merge branch 'main' of https://github.com/huggingface/optimum into pe…
pengwa Oct 21, 2023
06a429a
refine the comment
pengwa Oct 21, 2023
c259bba
fix make stype
pengwa Oct 26, 2023
d5c7971
Merge branch 'main' of https://github.com/huggingface/optimum into pe…
pengwa Oct 26, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions optimum/onnxruntime/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -459,6 +459,14 @@ def _inner_training_loop(

# Wrap the model with `ORTModule`
logger.info("Wrap ORTModule for ONNX Runtime training.")

from onnxruntime import version as ort_version

ort_support_stage3 = version.parse(ort_version) >= version.parse("1.17.0")
os.environ["ORTMODULE_ENABLE_ZERO_STAGE3"] = str(
int(ort_support_stage3 and self.is_deepspeed_enabled and is_deepspeed_zero3_enabled())
)

model = ORTModule(self.model)
self.model_wrapped = model
self.model = model
Expand All @@ -469,9 +477,9 @@ def _inner_training_loop(
self._created_lr_scheduler = False

if self.is_deepspeed_enabled:
if is_deepspeed_zero3_enabled():
if is_deepspeed_zero3_enabled() and not ort_support_stage3:
raise NotImplementedError(
"`ORTTrainer` does not support ZeRO stage 3 for the moment. Please use DeepSpeed stage 1 or 2 instead."
"`ORTTrainer` does not support ZeRO stage 3 for the moment. Please use DeepSpeed stage 1 or 2 instead or consider to update ONNX Runtime to 1.17 or later."
)
if args.bf16:
warnings.warn(
Expand Down
Loading