colossalai.builder
- colossalai.builder.build_gradient_handler(config, model, optimizer)[source]
Returns a gradient handler object of
BaseGradientHandler
constructed from config, model and optimizer.- Parameters
config (dict or
colossalai.context.Config
) – A python dict or acolossalai.context.Config
object containing information used in the construction of theGRADIENT_HANDLER
.model (
nn.Module
) – A model containing parameters for the gradient handleroptimizer (
torch.optim.Optimizer
) – An optimizer object containing parameters for the gradient handler
- Returns
An object of
colossalai.engine.BaseGradientHandler
- colossalai.builder.build_from_config(module, config)[source]
Returns an object of
module
constructed from config.- Parameters
module – A python or user-defined class
config – A python dict containing information used in the construction of the return object
Returns: An
object
of interest- Raises
AssertionError – Raises an AssertionError if module is not a class
- colossalai.builder.build_from_registry(config, registry)[source]
Returns an object constructed from config, the type of the object is specified by registry.
Note
the config is used to construct the return object such as LAYERS, OPTIMIZERS and other support types in registry. The config should contain all required parameters of corresponding object. The details of support types in registry and the mod_type in config could be found in registry.
- Parameters
config (dict or
colossalai.context.colossalai.context.Config
) – information used in the construction of the return object.registry (
Registry
) – A registry specifying the type of the return object
- Returns
A Python object specified by registry.
- Raises
Exception – Raises an Exception if an error occurred when building from registry.