fix:impot error

fastfold/distributed/comm.py

@@ -4,8 +4,8 @@ import torch
 import torch.distributed as dist
 from torch import Tensor
 
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 from .core import ensure_divisibility
 

fastfold/distributed/comm_async.py

@@ -5,8 +5,8 @@ import torch
 import torch.distributed as dist
 from torch import Tensor
 
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 from .comm import _split, divide
 

fastfold/model/fastnn/evoformer.py

@@ -4,8 +4,8 @@ from functools import partial
 import torch
 import torch.nn as nn
 
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 from fastfold.model.fastnn import MSACore, OutProductMean, PairCore
 from fastfold.model.fastnn.ops import Linear

fastfold/model/fastnn/msa.py

@@ -18,8 +18,8 @@ from typing import Optional, Tuple
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from fastfold.model.fastnn.kernel import LayerNorm, bias_dropout_add
 from fastfold.model.fastnn.ops import (ChunkMSARowAttentionWithPairBias, ChunkTransition, 
                                        SelfAttention, GlobalAttention, Transition, 

fastfold/model/fastnn/template.py

@@ -18,8 +18,8 @@ from typing import Optional, List
 import torch
 import torch.nn as nn
 
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 from fastfold.model.nn.primitives import Attention
 from fastfold.utils.checkpointing import checkpoint_blocks

tests/test_fastnn/test_msa_att_col.py

@@ -10,8 +10,8 @@ from fastfold.model.fastnn.ops import set_chunk_size
 from fastfold.model.hub import AlphaFold
 from fastfold.utils.inject_fastnn import inject_fastnn
 from fastfold.utils.import_weights import import_jax_weights_
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from fastfold.distributed.comm import gather, scatter, row_to_col
 from fastfold.utils.test_utils import get_param_path
 

tests/test_fastnn/test_msa_att_row.py

@@ -10,8 +10,8 @@ from fastfold.model.fastnn.ops import set_chunk_size
 from fastfold.model.hub import AlphaFold
 from fastfold.utils.inject_fastnn import inject_fastnn
 from fastfold.utils.import_weights import import_jax_weights_
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from fastfold.utils.test_utils import get_param_path
 from fastfold.distributed.comm import gather, scatter
 

tests/test_fastnn/test_msa_global_att_col.py

@@ -10,8 +10,8 @@ from fastfold.model.fastnn.ops import set_chunk_size
 from fastfold.model.hub import AlphaFold
 from fastfold.utils.inject_fastnn import inject_fastnn
 from fastfold.utils.import_weights import import_jax_weights_
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from fastfold.distributed.comm import gather, scatter, row_to_col
 from fastfold.utils.test_utils import get_param_path
 

tests/test_fastnn/test_out_product_mean.py

@@ -11,8 +11,8 @@ from fastfold.model.hub import AlphaFold
 from fastfold.utils.inject_fastnn import inject_fastnn
 from fastfold.utils.import_weights import import_jax_weights_
 from fastfold.config import model_config
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from fastfold.utils.test_utils import get_param_path
 from fastfold.distributed import scatter, row_to_col
 from fastfold.distributed.comm import gather, scatter