[tensor] design DistSpec and DistSpecManager for ColoTensor #934
Conversation
|
All tests passed. |
| from .op_wrapper import ( | ||
| colo_op_impl,) | ||
| from .colo_tensor import ColoTensor | ||
| from .colo_parameter import ColoParameter | ||
| from .utils import convert_parameter, named_params_with_colotensor | ||
| from ._ops import * | ||
| from .optim.colo_optimizer import ColoOptimizer | ||
| from . import dist_spec |
There was a problem hiding this comment.
why not name it as DistSpec?
There was a problem hiding this comment.
Or DistSignature? to be consistent with our design doc?
There was a problem hiding this comment.
It's a package not class. I just use package to implement the "namespace" in C++.
There was a problem hiding this comment.
I've updated this part in our design doc. I think current implementation is simpler and it's easy to comprehend for users.
| from .op_wrapper import ( | ||
| colo_op_impl,) | ||
| from .colo_tensor import ColoTensor | ||
| from .colo_parameter import ColoParameter | ||
| from .utils import convert_parameter, named_params_with_colotensor | ||
| from ._ops import * | ||
| from .optim.colo_optimizer import ColoOptimizer | ||
| from . import dist_spec |
There was a problem hiding this comment.
Or DistSignature? to be consistent with our design doc?
colossalai/tensor/colo_tensor.py
Outdated
|
|
||
| def has_spec(self) -> bool: | ||
| return self._shard_spec is not None and self._shard_spec.num_action > 0 | ||
| return self._spec is not None and self._spec.num_action > 0 |
There was a problem hiding this comment.
we'd better set _spec as a dummy one by default.
I think making it an optional variable may bring some trouble. According to our experiments in ShardTensor design.
There was a problem hiding this comment.
Oh, yes. Actually our default value of _spec is not None. This code snippet is outdated and I forgot to update.
| from .op_wrapper import ( | ||
| colo_op_impl,) | ||
| from .colo_tensor import ColoTensor | ||
| from .colo_parameter import ColoParameter | ||
| from .utils import convert_parameter, named_params_with_colotensor | ||
| from ._ops import * | ||
| from .optim.colo_optimizer import ColoOptimizer | ||
| from . import dist_spec |
test_model.py