github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 20:10:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[checkpointio] hotfix torch 2.0 compatibility (#4824)

Browse Source
This commit is contained in:
Hongxin Liu
2023-10-07 10:45:52 +08:00
committed by GitHub
parent ad23460cf8
commit cb3a25a062
3 changed files with 26 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/checkpoint_io/utils.py
View File

@@ -9,6 +9,7 @@ from typing import Iterator, List, Mapping, Optional, OrderedDict, Tuple
import torch
import torch.nn as nn
from packaging.version import Version
from torch.optim import Optimizer
from colossalai.tensor.d_tensor import (
@@ -663,7 +664,10 @@ def sharded_optimizer_loading_epilogue(optimizer: Optimizer):
"""
# Do the cleaning up as in src code of Pytorch.
optimizer._hook_for_profile() # To support multiprocessing pickle/unpickle.
if Version(torch.__version__) >= Version("2.0.0"):
optimizer._patch_step_function() # To support multiprocessing pickle/unpickle
else:
optimizer._hook_for_profile() # To support multiprocessing pickle/unpickle.
optimizer.defaults.setdefault("differentiable", False)