[BUG]: ValueError: mutable default <class 'colossalai.legacy.tensor.distspec._DistSpec'> for field dist_attr is not allowed: use default_factory
🐛 Describe the bug
ValueError Traceback (most recent call last) Cell In[2], line 1 ----> 1 from colossalai.booster import Booster
File ~/.local/lib/python3.11/site-packages/colossalai/booster/init.py:2 1 from .accelerator import Accelerator ----> 2 from .booster import Booster 3 from .plugin import Plugin
File ~/.local/lib/python3.11/site-packages/colossalai/booster/booster.py:17 15 from .accelerator import Accelerator 16 from .mixed_precision import MixedPrecision, mixed_precision_factory ---> 17 from .plugin import Plugin 18 from .plugin.pp_plugin_base import PipelinePluginBase 20 all = ["Booster"]
File ~/.local/lib/python3.11/site-packages/colossalai/booster/plugin/init.py:1 ----> 1 from .gemini_plugin import GeminiPlugin 2 from .hybrid_parallel_plugin import HybridParallelPlugin 3 from .low_level_zero_plugin import LowLevelZeroPlugin
File ~/.local/lib/python3.11/site-packages/colossalai/booster/plugin/gemini_plugin.py:31 29 from colossalai.interface import ModelWrapper, OptimizerWrapper 30 from colossalai.shardformer import ShardConfig, ShardFormer ---> 31 from colossalai.zero import GeminiDDP, GeminiOptimizer 32 from colossalai.zero.gemini.memory_tracer import MemStats 34 from .dp_plugin_base import DPPluginBase
File ~/.local/lib/python3.11/site-packages/colossalai/zero/init.py:1 ----> 1 from .gemini import GeminiAdamOptimizer, GeminiDDP, GeminiOptimizer, get_static_torch_model 2 from .low_level import LowLevelZeroOptimizer 3 from .wrapper import zero_model_wrapper, zero_optim_wrapper
File ~/.local/lib/python3.11/site-packages/colossalai/zero/gemini/init.py:2 1 from .chunk import ChunkManager, TensorInfo, TensorState, search_chunk_configuration ----> 2 from .gemini_ddp import GeminiDDP 3 from .gemini_mgr import GeminiManager 4 from .gemini_optimizer import GeminiAdamOptimizer, GeminiOptimizer
File ~/.local/lib/python3.11/site-packages/colossalai/zero/gemini/gemini_ddp.py:34 31 from colossalai.utils import _cast_float, free_storage, is_ddp_ignored 33 from .chunk import Chunk, ChunkManager, TensorState, init_chunk_manager ---> 34 from .gemini_hook import GeminiZeROHook 35 from .gemini_mgr import GeminiManager 36 from .memory_tracer import MemStats, OrderedParamGenerator
File ~/.local/lib/python3.11/site-packages/colossalai/zero/gemini/gemini_hook.py:11 9 from colossalai.utils import is_ddp_ignored 10 from colossalai.zero.gemini import TensorState ---> 11 from colossalai.zero.gemini.gemini_mgr import GeminiManager 14 class TrainingPhase(Enum): 15 FORWARD = 0
File ~/.local/lib/python3.11/site-packages/colossalai/zero/gemini/gemini_mgr.py:9 7 from .chunk import Chunk, ChunkManager 8 from .memory_tracer import ChunkMemStatsCollector, MemStats ----> 9 from .placement_policy import PlacementPolicyFactory 12 class GeminiManager: 13 """ 14 Stateful Tensor Manager, inspired from PatrickStar 15 (...) 24 memstats (MemStats, optional): a mem stats collected by a runtime mem tracer. if None then GeminiManager will collect it during a warmup iteration. 25 """
File ~/.local/lib/python3.11/site-packages/colossalai/zero/gemini/placement_policy.py:10 7 import torch 9 from colossalai.accelerator import get_accelerator ---> 10 from colossalai.legacy.utils.memory import colo_device_memory_capacity 11 from colossalai.zero.gemini.chunk import Chunk 13 from .chunk import Chunk, ChunkManager
File ~/.local/lib/python3.11/site-packages/colossalai/legacy/init.py:1 ----> 1 from .initialize import ( 2 get_default_parser, 3 initialize, 4 launch, 5 launch_from_openmpi, 6 launch_from_slurm, 7 launch_from_torch, 8 ) 10 all = [ 11 "launch", 12 "launch_from_openmpi", (...) 16 "get_default_parser", 17 ]
File ~/.local/lib/python3.11/site-packages/colossalai/legacy/initialize.py:21 19 from colossalai.context import Config, ConfigException 20 from colossalai.interface import OptimizerWrapper ---> 21 from colossalai.legacy.amp import AMP_TYPE, convert_to_amp 22 from colossalai.legacy.amp.naive_amp import NaiveAMPModel 23 from colossalai.legacy.builder.builder import build_gradient_handler
File ~/.local/lib/python3.11/site-packages/colossalai/legacy/amp/init.py:11 8 from colossalai.context import Config 10 from .amp_type import AMP_TYPE ---> 11 from .apex_amp import convert_to_apex_amp 12 from .naive_amp import convert_to_naive_amp 13 from .torch_amp import convert_to_torch_amp
File ~/.local/lib/python3.11/site-packages/colossalai/legacy/amp/apex_amp/init.py:4
1 import torch.nn as nn
2 from torch.optim import Optimizer
----> 4 from .apex_amp import ApexAMPOptimizer
7 def convert_to_apex_amp(model: nn.Module, optimizer: Optimizer, amp_config):
8 r"""A helper function to wrap training components with Apex AMP modules
9
10 Args:
(...)
34 More details about amp_config refer to amp_config <[https://nvidia.github.io/apex/amp.html?highlight=apex%20amp>.](https://nvidia.github.io/apex/amp.html?highlight=apex%20amp%3E`.%3C/span%3E)
35 """
File ~/.local/lib/python3.11/site-packages/colossalai/legacy/amp/apex_amp/apex_amp.py:14 11 from torch import Tensor 13 from colossalai.interface import OptimizerWrapper ---> 14 from colossalai.legacy.utils import clip_grad_norm_fp32 17 class ApexAMPOptimizer(OptimizerWrapper): 18 """A wrapper class for APEX optimizer and it implements apex-specific backward and clip_grad_norm 19 methods 20 """
File ~/.local/lib/python3.11/site-packages/colossalai/legacy/utils/init.py:1 ----> 1 from .checkpointing import load_checkpoint, save_checkpoint 2 from .common import ( 3 clip_grad_norm_fp32, 4 copy_tensor_parallel_attributes, (...) 16 sync_model_param, 17 ) 18 from .data_sampler import DataParallelSampler, get_dataloader
File ~/.local/lib/python3.11/site-packages/colossalai/legacy/utils/checkpointing.py:16 13 except ImportError: 14 _EXTRA_STATE_KEY_SUFFIX = "_extra_state" ---> 16 from .common import is_using_pp 18 all = ["save_checkpoint", "load_checkpoint"] 21 def broadcast_state_dict(state_dict, parallel_mode):
File ~/.local/lib/python3.11/site-packages/colossalai/legacy/utils/common.py:16 14 from colossalai.legacy.core import global_context as gpc 15 from colossalai.legacy.global_variables import tensor_parallel_env as env ---> 16 from colossalai.legacy.tensor import ProcessGroup 17 from colossalai.tensor import ColoParameter 18 from colossalai.utils.multi_tensor_apply import multi_tensor_applier
File ~/.local/lib/python3.11/site-packages/colossalai/legacy/tensor/init.py:6 4 from .distspec import ReplicaSpec, ShardSpec 5 from .process_group import ProcessGroup ----> 6 from .tensor_spec import ColoTensorSpec 8 all = [ 9 "ComputePattern", 10 "ComputeSpec", (...) 16 "ReplicaSpec", 17 ]
File ~/.local/lib/python3.11/site-packages/colossalai/legacy/tensor/tensor_spec.py:10
5 from colossalai.legacy.tensor.process_group import ProcessGroup
7 from .compute_spec import ComputeSpec
---> 10 @dataclass
11 class ColoTensorSpec:
12 """ColoTensorSpec
13
14 A data class for specifications of the ColoTensor.
15 It contains attributes of ProcessGroup, _DistSpec, ComputeSpec.
16 The latter two attributes are optional. If not set, they are default value is Replicate() and None.
17 """
19 pg: ProcessGroup
File /usr/local/lib/python3.11/dataclasses.py:1230, in dataclass(cls, init, repr, eq, order, unsafe_hash, frozen, match_args, kw_only, slots, weakref_slot) 1227 return wrap 1229 # We're called as @dataclass without parens. -> 1230 return wrap(cls)
File /usr/local/lib/python3.11/dataclasses.py:1220, in dataclass.
File /usr/local/lib/python3.11/dataclasses.py:958, in _process_class(cls, init, repr, eq, order, unsafe_hash, frozen, match_args, kw_only, slots, weakref_slot) 955 kw_only = True 956 else: 957 # Otherwise it's a field of some type. --> 958 cls_fields.append(_get_field(cls, name, type, kw_only)) 960 for f in cls_fields: 961 fields[f.name] = f
File /usr/local/lib/python3.11/dataclasses.py:815, in _get_field(cls, a_name, a_type, default_kw_only) 811 # For real fields, disallow mutable defaults. Use unhashable as a proxy 812 # indicator for mutability. Read the hash attribute from the class, 813 # not the instance. 814 if f._field_type is _FIELD and f.default.class.hash is None: --> 815 raise ValueError(f'mutable default {type(f.default)} for field ' 816 f'{f.name} is not allowed: use default_factory') 818 return f
ValueError: mutable default <class 'colossalai.legacy.tensor.distspec._DistSpec'> for field dist_attr is not allowed: use default_factory
Environment
python3.11,cuda12.2,nvidia-cudnn-cu12==8.9.2.2,nvidia-nccl-cu12==2.18.1; colossalai 0.3.6, PyTorch 2.1.1
Hi, this issue has been fixed. https://github.com/hpcaitech/ColossalAI/pull/5440/files