github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-01-29 21:49:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish doc style for ColoTensor (#1457)

Browse Source
This commit is contained in:
Jiarui Fang
2022-08-16 09:21:05 +08:00
committed by GitHub
parent 0dbd61c29b
commit a1476ea882
9 changed files with 197 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/tensor/tensor_spec.py
View File

@@ -7,6 +7,12 @@ from dataclasses import dataclass
@dataclass
class ColoTensorSpec:
""" ColoTensorSpec
A data class for specifications of the `ColoTensor`.
It contains attributes of `ProcessGroup`, `_DistSpec`, `ComputeSpec`.
The latter two attributes are optional. If not set, they are default value is `Replicate()` and `None`.
"""
pg: ProcessGroup
dist_attr: Optional[_DistSpec] = _DistSpec(DistPlacementPattern.REPLICATE)
compute_attr: Optional[ComputeSpec] = None