colossalai.builder

colossalai.builder.build_schedule(config)

Returns a schedule of colossalai.engine.schedule.BaseSchedule.

Parameters: config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
Returns: An object of colossalai.engine.schedule.BaseSchedule
Return type: colossalai.engine.schedule.BaseSchedule

colossalai.builder.build_lr_scheduler(config, optimizer)

Returns a learning rate scheduler object of torch.optim.lr_scheduler constructed from config, optimizer, total_steps and num_steps_per_epoch.

Parameters

config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
optimizer (torch.optim.Optimizer) – An optimizer object containing parameters for the learning rate scheduler

Returns

An object of torch.optim.lr_scheduler

Return type

torch.optim.lr_scheduler

colossalai.builder.build_model(config)

Returns a model object of nn.Module constructed from config.

Parameters: config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
Returns: An object of torch.nn.Module
Return type: torch.nn.Module

colossalai.builder.build_optimizer(config, model)

Returns an optimizer object of torch.optim.Optimizer constructed from config, ‘model’ and ‘params’.

Parameters

config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
model (nn.Module) – A model containing parameters for the optimizer

Returns

An object of torch.optim.Optimizer

Return type

torch.optim.Optimizer

colossalai.builder.build_layer(config)

Returns a layer object of nn.Module constructed from config.

Parameters: config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
Returns: An object of torch.nn.Module
Return type: torch.nn.Module

colossalai.builder.build_loss(config)

Returns a loss function object of torch.autograd.Function constructed from config.

Parameters: config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
Returns: An object of torch.nn.modules.loss._Loss
Return type: torch.nn.modules.loss._Loss

colossalai.builder.build_hooks(config, trainer)

Returns a hook object of BaseHook constructed from config and trainer.

Parameters

config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
trainer (Trainer) – A Trainer object containing parameters for the hook

Returns

An object of colossalai.trainer.hooks.BaseHook

Return type

colossalai.trainer.hooks.BaseHook

colossalai.builder.build_dataset(config)

Returns a dataset object of torch.utils.data.Dataset constructed from config.

Parameters: config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
Returns: An object of torch.utils.data.Dataset
Return type: torch.utils.data.Dataset

colossalai.builder.build_transform(config)

Returns a transformation object of torchvision.transforms constructed from config.

Parameters: config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
Returns: An object of torchvision.transforms
Return type: torchvision.transforms

colossalai.builder.build_data_sampler(config, dataset)

Returns a data sampler object of colossalai.nn.data.sampler.BaseSampler constructed from config.

Parameters

config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
dataset (torch.utils.data.Dataset) – An object of torch.utils.data.Dataset containing information used in the construction of the return object

Returns

An object of colossalai.utils.data_sampler.BaseSampler

Return type

colossalai.utils.data_sampler.BaseSampler

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 return object
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

Return type

colossalai.engine.BaseGradientHandler

colossalai.builder.build_pipeline_model(layers, num_chunks=1, verbose=False)

An intializer to split the model into different stages for pipeline parallelism. Note that layer must be torch.nn.Sequential.

Parameters

layers (torch.nn.Sequential) – Layers of model
num_chunks (int, optional) – The number of chunks you want to have on the current stage. This value should be 1 in most cases unless you are using virutal pipeline parallelism.
verbose (bool, optional) – Whether to print the logs

colossalai.builder.build_pipeline_model_from_cfg(config, num_chunks=1, partition_method='parameter', verbose=False)

An intializer to split the model into different stages for pipeline parallelism.

An example for the model config is shown below. The class VisionTransformerFromConfig should inherit colossalai.nn.model.ModelFromConfig to allow this initializer to build model from a sequence of layer configurations.

model_config = dict(: type=’VisionTransformerFromConfig’, embedding_cfg=dict(…), …

)

Parameters

config (dict) – Configuration of the model
num_chunks (int, optional) – The number of chunks you want to have on the current stage. This value should be 1 in most cases unless you are using virutal pipeline parallelism.
partition_method (str, optional) – This parameter determines how you want to split your model layers into stages, you can set it as ‘layer’ or ‘parameter’
verbose (bool, optional) – Whether to print the logs

colossalai.builder.build_ophooks(config)

Returns a hook object of BaseOpHook constructed from config.

Parameters: config (dict or colossalai.context.Config) – A python dict or a colossalai.context.Config object containing information used in the construction of the return object
Returns: An object of colossalai.trainer.hooks.BaseOpHook
Return type: colossalai.trainer.hooks.BaseOpHook