Error when importing batch_encode_sequences

[yangkky]

Bug description

After installing according to the instructions, I try:

from esm.esmfold.v1.misc import batch_encode_sequences

And get:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>                                                                                                                                  
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/esm/esmfold/v1/misc.py", line 12, in <module>
    from openfold.np import residue_constants              
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/openfold/__init__.py", line 1, in <module>
    from . import model
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/openfold/model/__init__.py", line 11, in <module>
    _modules = [(m, importlib.import_module("." + m, __name__)) for m in __all__]
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/openfold/model/__init__.py", line 11, in <listcomp>
    _modules = [(m, importlib.import_module("." + m, __name__)) for m in __all__]
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/importlib/__init__.py", line 127, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/openfold/model/evoformer.py", line 22, in <module>
    from openfold.model.primitives import Linear, LayerNorm
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/openfold/model/primitives.py", line 23, in <module>
    import deepspeed
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/deepspeed/__init__.py", line 16, in <module>
    from .runtime.engine import DeepSpeedEngine, DeepSpeedOptimizerCallable, DeepSpeedSchedulerCallable
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 25, in <module>
    from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/deepspeed/runtime/zero/__init__.py", line 6, in <module>
    from .partition_parameters import ZeroParamType
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 25, in <module>
    from ..config import DeepSpeedConfig
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/deepspeed/runtime/config.py", line 24, in <module>
    from .zero.config import DeepSpeedZeroConfig
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/deepspeed/runtime/zero/config.py", line 10, in <module>
    from .offload_config import get_offload_param_config, get_default_offload_param_config, \
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/deepspeed/runtime/zero/offload_config.py", line 8, in <module>
    from .utils import logger
  File "/home/kevyan/miniconda3/envs/esm-test/lib/python3.9/site-packages/deepspeed/runtime/zero/utils.py", line 36, in <module>
    ZERO_SUPPORTED_OPTIMIZERS.append(apex.optimizers.FusedAdam)
AttributeError: module 'apex' has no attribute 'optimizers'

Reproduction steps

Here's the conda environment to reproduce:

name: esm-test
channels:
 - conda-forge
 - bioconda
 - pytorch
dependencies:
 - conda-forge::python=3.9
 - conda-forge::pip
 - conda-forge::cudatoolkit==11.3.*
 - pytorch::pytorch=1.12.*
 - pip:
   - fair-esm[esmfold]
   - git+https://github.com/NVIDIA/dllogger.git
   - git+https://github.com/aqlaboratory/openfold.git@4b41059694619831a7db195b7e0988fc4ff3a307

Expected behavior

Error-free import