colossalai.builder

colossalai.builder.build_gradient_handler(config, model, optimizer)

Returns a gradient handler object of BaseGradientHandler constructed from config, model and optimizer.

Parameters

• config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the GRADIENT_HANDLER.

• model (nn.Module) – A model containing parameters for the gradient handler

• optimizer (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)

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)

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.

• colossalai.builder.builder
  • build_from_config()
  • build_from_registry()
  • build_gradient_handler()