mirror of https://github.com/hpcaitech/ColossalAI
33 lines
965 B
Python
33 lines
965 B
Python
from typing import Dict, List, NamedTuple
|
|
|
|
import torch
|
|
|
|
from colossalai.elixir.chunk import ChunkGroup
|
|
|
|
|
|
class ChunkPlan(NamedTuple):
|
|
"""ChunkPlan is a type of configuration used to instruct the initialization of a chunk.
|
|
|
|
args:
|
|
name_list: contains the names of parameters that should be pushed into this chunk
|
|
chunk_size: the size of this chunk
|
|
chunk_dtype: the dtype of this chunk
|
|
kwargs: a dictionary used in __init__ function of Chunk
|
|
"""
|
|
name_list: List[str]
|
|
chunk_size: int
|
|
chunk_dtype: torch.dtype
|
|
kwargs: Dict
|
|
|
|
|
|
class SearchResult(object):
|
|
|
|
def __init__(self,
|
|
chunk_group: ChunkGroup,
|
|
chunk_plans: List[ChunkPlan],
|
|
param_called_per_step: List[List[str]] = None) -> None:
|
|
super().__init__()
|
|
self.chunk_group = chunk_group
|
|
self.param_chunk_plans = chunk_plans
|
|
self.param_called_per_step = param_called_per_step
|