TypeError: Cannot subclass <class 'typing._SpecialForm'> while fine tuning GPT-J

Viewed 170

I am trying to fine tune GPT-J by following this GitHub Repository. When running the training command, I encounter this error:

Traceback (most recent call last):
  File "device_train.py", line 13, in <module>
    from mesh_transformer import util
  File "/home/shreyjain/mesh-transformer-jax/mesh_transformer/util.py", line 36, in <module>
    class ClipByGlobalNormState(OptState):
  File "/usr/lib/python3.8/typing.py", line 317, in __new__
    raise TypeError(f"Cannot subclass {cls!r}")
TypeError: Cannot subclass <class 'typing._SpecialForm'> 

This looks like a source code error but I am not sure. I have also raised an issue on GitHub regarding this. Any help will be appreciated!

1 Answers

I had the same error, this is caused by a wrong Optax version. Fix this by downgrading Optax to 0.0.9 (at least it worked for me!).

pip install optax==0.0.9

Related