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 – layers of model

  • num_chunks (int) – 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) – 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) – 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) – this parameter determines how you want to split your model layers into stages, you can set it as ‘layer’ or ‘parameter’

  • verbose (bool) – whether to print the logs