AttributeError: module 'torch.utils' has no attribute 'checkpoint' in gpt-neox/gpt-neox

[kinoc]

while running " deepspeed train_enwik8.py --deepspeed --deepspeed_config ./configs/deepspeed_zero2.json" I got the error AttributeError: module 'torch.utils' has no attribute 'checkpoint' in gpt-neox/gpt-neox
seems similar to problem described in pyTorch attributeerror-module-torch-utils-has-no-attribute-checkpoint

quick-fix: adding
"from torch.utils.checkpoint import checkpoint"
to top of file seems to fix the problem

[StellaAthena]

Are you sure you have the versions specified in requirements.txt installed? I just ran this without any issue.

[joshlk]

I get the same error with torch v1.6.0 or v1.7.0. What version are you using?

[StellaAthena]

~/gpt-neox# python3
Python 3.8.5 (default, Sep  4 2020, 07:30:14) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> c = torch.utils.checkpoint.checkpoint
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: module 'torch.utils' has no attribute 'checkpoint'
>>> from torch.utils.checkpoint import checkpoint
>>> checkpoint
<function checkpoint at 0x7fcb633aca60>
>>> 

Well then. No idea how this is even possible, but that's the error it seems. We're about to roll something out, I'll make sure this fix is included.

[kinoc]

Are you sure you have the versions specified in requirements.txt installed? I just ran this without any issue.

einops 0.3.0
torch 1.7.1
torchcontrib 0.0.2
tqdm 4.56.0
transformers 4.2.2
ftfy 5.8
lm-dataformat 0.0.19

The one stand out is :
tensorflow-gpu 1.14.0

It runs fine with the line locally, and not when I remove it.