github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-03 10:06:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish comments for Chunk class (#2116)

Browse Source
This commit is contained in:
Jiarui Fang
2022-12-12 15:39:31 +08:00
committed by GitHub
parent 09d69e1c25
commit e99edfcb51
6 changed files with 95 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
colossalai/tensor/param_op_hook.py
View File

@@ -9,10 +9,11 @@ from colossalai.tensor.tensor_spec import ColoTensorSpec
class ColoParamOpHook(ABC):
"""Hook which is triggered by each operation when operands contain ColoParameter.
"""
Hook which is triggered by each operation when operands contain ColoParameter.
To customize it, you must inherit this abstract class, and implement ``pre_forward``,
``post_forward``, ``pre_backward`` and ``post_backward``. These four methods take a list
of ColoParameter.
``post_forward``, ``pre_backward`` and ``post_backward``.
These four methods apply a list of ColoParameter as input args.
"""
@abstractmethod
@@ -33,7 +34,8 @@ class ColoParamOpHook(ABC):
class ColoParamOpHookManager:
"""Manage your param op hooks. It only has static methods.
"""
Manage your param op hooks. It only has static methods.
The only static method you should call is ``use_hooks(*hooks)``.
"""
hooks: Tuple[ColoParamOpHook, ...] = tuple()