Index

A

accumulate_gradient() (in module colossalai.utils) (in module colossalai.utils.gradient_accumulation) Accuracy2D (class in colossalai.nn.metric.accuracy_2d) Accuracy2p5D (class in colossalai.nn.metric.accuracy_2p5d) Accuracy3D (class in colossalai.nn.metric.accuracy_3d) AccuracyHook (class in colossalai.trainer.hooks) add_global_rank() (colossalai.context.parallel_context.ParallelContext method) add_group() (colossalai.context.parallel_context.ParallelContext method) add_local_rank() (colossalai.context.parallel_context.ParallelContext method) add_ranks_in_group() (colossalai.context.parallel_context.ParallelContext method) add_seed() (colossalai.context.random.seed_manager.SeedManager method) (in module colossalai.context.random) add_world_size() (colossalai.context.parallel_context.ParallelContext method) after_hook_is_attached() (colossalai.trainer.hooks.BaseHook method) after_test() (colossalai.trainer.hooks.BaseHook method) (colossalai.trainer.hooks.LogMemoryByEpochHook method)

after_test_epoch() (colossalai.trainer.hooks.BaseHook method) (colossalai.trainer.hooks.LogTimingByEpochHook method) after_test_iter() (colossalai.trainer.hooks.BaseHook method) after_train() (colossalai.trainer.hooks.BaseHook method) after_train_epoch() (colossalai.trainer.hooks.BaseHook method) (colossalai.trainer.hooks.LogMemoryByEpochHook method) (colossalai.trainer.hooks.LogTimingByEpochHook method) (colossalai.trainer.hooks.SaveCheckpointHook method) after_train_iter() (colossalai.trainer.hooks.BaseHook method) all_gather() (in module colossalai.communication) (in module colossalai.communication.collective) AMP_TYPE (class in colossalai.amp.amp_type) ApexAMPOptimizer (class in colossalai.amp.apex_amp) (class in colossalai.amp.apex_amp.apex_amp) AttnMaskType (class in colossalai.kernel.cuda_native.scaled_softmax)

B

backward() (colossalai.amp.apex_amp.apex_amp.ApexAMPOptimizer method) (colossalai.amp.apex_amp.ApexAMPOptimizer method) (colossalai.amp.torch_amp.torch_amp.TorchAMPOptimizer method) (colossalai.amp.torch_amp.TorchAMPOptimizer method) (colossalai.engine.Engine method) backward_by_grad() (colossalai.engine.Engine method) backward_step() (colossalai.engine.schedule.PipelineSchedule method) BaseGradientHandler (class in colossalai.engine.gradient_handler) BaseHook (class in colossalai.trainer.hooks) BaseSchedule (class in colossalai.engine.schedule) before_test() (colossalai.trainer.hooks.BaseHook method) before_test_epoch() (colossalai.trainer.hooks.BaseHook method) before_test_iter() (colossalai.trainer.hooks.BaseHook method) before_train() (colossalai.trainer.hooks.BaseHook method) (colossalai.trainer.hooks.LoadCheckpointHook method) (colossalai.trainer.hooks.LogMemoryByEpochHook method) before_train_epoch() (colossalai.trainer.hooks.BaseHook method) before_train_iter() (colossalai.trainer.hooks.BaseHook method) build_data_sampler() (in module colossalai.builder) (in module colossalai.builder.builder) build_dataset() (in module colossalai.builder) (in module colossalai.builder.builder) build_from_config() (in module colossalai.builder.builder) build_from_registry() (in module colossalai.builder.builder) build_gradient_handler() (in module colossalai.builder) (in module colossalai.builder.builder)

build_hooks() (in module colossalai.builder) (in module colossalai.builder.builder) build_layer() (in module colossalai.builder) (in module colossalai.builder.builder) build_loss() (in module colossalai.builder) (in module colossalai.builder.builder) build_lr_scheduler() (in module colossalai.builder) (in module colossalai.builder.builder) build_model() (in module colossalai.builder) (in module colossalai.builder.builder) build_ophooks() (in module colossalai.builder) (in module colossalai.builder.builder) build_optimizer() (in module colossalai.builder) (in module colossalai.builder.builder) build_pipeline_model() (in module colossalai.builder) (in module colossalai.builder.pipeline) build_pipeline_model_from_cfg() (in module colossalai.builder) (in module colossalai.builder.pipeline) build_schedule() (in module colossalai.builder) (in module colossalai.builder.builder) build_transform() (in module colossalai.builder) (in module colossalai.builder.builder) bytes_to_GB() (in module colossalai.utils.memory) bytes_to_MB() (in module colossalai.utils.memory)

C

calc_loss() (colossalai.engine.Engine method) check_sanity() (colossalai.context.parallel_context.ParallelContext method) checkpoint() (in module colossalai.utils) (in module colossalai.utils.activation_checkpoint) Classifier (class in colossalai.nn.layer.colossalai_layer) (class in colossalai.nn.layer.colossalai_layer.linear) Classifier1D (class in colossalai.nn.layer.parallel_1d) (class in colossalai.nn.layer.parallel_1d.layers) Classifier2D (class in colossalai.nn.layer.parallel_2d) (class in colossalai.nn.layer.parallel_2d.layers) Classifier2p5D (class in colossalai.nn.layer.parallel_2p5d) (class in colossalai.nn.layer.parallel_2p5d.layers) Classifier3D (class in colossalai.nn.layer.parallel_3d) (class in colossalai.nn.layer.parallel_3d.layers) clip_grad_norm() (colossalai.amp.apex_amp.apex_amp.ApexAMPOptimizer method) (colossalai.amp.apex_amp.ApexAMPOptimizer method) (colossalai.amp.torch_amp.torch_amp.TorchAMPOptimizer method) (colossalai.amp.torch_amp.TorchAMPOptimizer method) clip_grad_norm_fp32() (in module colossalai.utils) (in module colossalai.utils.common) colossalai module colossalai.amp module colossalai.amp.amp_type module colossalai.amp.apex_amp module colossalai.amp.apex_amp.apex_amp module colossalai.amp.naive_amp module colossalai.amp.naive_amp.naive_amp module colossalai.amp.torch_amp module colossalai.amp.torch_amp.torch_amp module colossalai.builder module colossalai.builder.builder module colossalai.builder.pipeline module colossalai.communication module colossalai.communication.collective module colossalai.communication.p2p module colossalai.communication.ring module colossalai.communication.utils module colossalai.context module colossalai.context.config module colossalai.context.parallel_context module colossalai.context.parallel_mode module colossalai.context.process_group_initializer module colossalai.context.process_group_initializer.initializer_1d module colossalai.context.process_group_initializer.initializer_2d module colossalai.context.process_group_initializer.initializer_2p5d module colossalai.context.process_group_initializer.initializer_3d module colossalai.context.process_group_initializer.initializer_data module colossalai.context.process_group_initializer.initializer_model module colossalai.context.process_group_initializer.initializer_moe module colossalai.context.process_group_initializer.initializer_pipeline module colossalai.context.process_group_initializer.initializer_sequence module colossalai.context.process_group_initializer.initializer_tensor module colossalai.context.process_group_initializer.process_group_initializer module colossalai.context.random module colossalai.context.random.seed_manager module colossalai.engine module colossalai.engine.gradient_handler module colossalai.engine.schedule module colossalai.initialize module colossalai.kernel module colossalai.kernel.cuda_native module colossalai.kernel.cuda_native.layer_norm module colossalai.kernel.cuda_native.multihead_attention module colossalai.kernel.cuda_native.scaled_softmax module colossalai.kernel.jit module colossalai.kernel.jit.bias_dropout_add module colossalai.kernel.jit.bias_gelu module colossalai.kernel.jit.option module colossalai.logging module colossalai.logging.logging module colossalai.nn module colossalai.nn.init module colossalai.nn.layer module colossalai.nn.layer.base_layer module colossalai.nn.layer.colossalai_layer module colossalai.nn.layer.colossalai_layer.dropout module colossalai.nn.layer.colossalai_layer.embedding module colossalai.nn.layer.colossalai_layer.linear module colossalai.nn.layer.colossalai_layer.normalization module colossalai.nn.layer.moe module colossalai.nn.layer.moe.layers module

colossalai.nn.layer.parallel_1d module colossalai.nn.layer.parallel_1d.layers module colossalai.nn.layer.parallel_2d module colossalai.nn.layer.parallel_2d.layers module colossalai.nn.layer.parallel_2p5d module colossalai.nn.layer.parallel_2p5d.layers module colossalai.nn.layer.parallel_3d module colossalai.nn.layer.parallel_3d.layers module colossalai.nn.layer.parallel_sequence module colossalai.nn.layer.parallel_sequence.layers module colossalai.nn.layer.utils module colossalai.nn.layer.utils.common module colossalai.nn.layer.vanilla module colossalai.nn.layer.vanilla.layers module colossalai.nn.layer.wrapper module colossalai.nn.layer.wrapper.lambda_wrapper module colossalai.nn.layer.wrapper.pipeline_wrapper module colossalai.nn.loss module colossalai.nn.loss.loss_2d module colossalai.nn.loss.loss_2p5d module colossalai.nn.loss.loss_3d module colossalai.nn.loss.loss_moe module colossalai.nn.lr_scheduler module colossalai.nn.lr_scheduler.cosine module colossalai.nn.lr_scheduler.delayed module colossalai.nn.lr_scheduler.linear module colossalai.nn.lr_scheduler.multistep module colossalai.nn.lr_scheduler.onecycle module colossalai.nn.lr_scheduler.poly module colossalai.nn.lr_scheduler.torch module colossalai.nn.metric module colossalai.nn.metric.accuracy_2d module colossalai.nn.metric.accuracy_2p5d module colossalai.nn.metric.accuracy_3d module colossalai.nn.model module colossalai.nn.model.model_from_config module colossalai.nn.optimizer module colossalai.nn.optimizer.colossalai_optimizer module colossalai.nn.optimizer.fused_adam module colossalai.nn.optimizer.fused_lamb module colossalai.nn.optimizer.fused_sgd module colossalai.nn.optimizer.lamb module colossalai.nn.optimizer.lars module colossalai.registry module colossalai.registry.registry module colossalai.trainer module colossalai.trainer.hooks module colossalai.utils module colossalai.utils.activation_checkpoint module colossalai.utils.checkpointing module colossalai.utils.common module colossalai.utils.cuda module colossalai.utils.data_sampler module colossalai.utils.data_sampler.base_sampler module colossalai.utils.data_sampler.data_parallel_sampler module colossalai.utils.gradient_accumulation module colossalai.utils.memory module colossalai.utils.multi_tensor_apply module colossalai.utils.multi_tensor_apply.multi_tensor_apply module colossalai.utils.timer module colossalai.zero module Config (class in colossalai.context.config) (class in colossalai.kernel.cuda_native.multihead_attention) convert_to_amp() (in module colossalai.amp) convert_to_apex_amp() (in module colossalai.amp.apex_amp) convert_to_naive_amp() (in module colossalai.amp.naive_amp) convert_to_torch_amp() (in module colossalai.amp.torch_amp) convert_to_zero() (in module colossalai.zero) CosineAnnealingLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.cosine) CosineAnnealingWarmupLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.cosine) count_layer_params() (in module colossalai.builder.pipeline) criterion (colossalai.engine.Engine property) CrossEntropyLoss2D (class in colossalai.nn.loss.loss_2d) CrossEntropyLoss2p5D (class in colossalai.nn.loss.loss_2p5d) CrossEntropyLoss3D (class in colossalai.nn.loss.loss_3d) cur_epoch (colossalai.trainer.Trainer property) cur_step (colossalai.trainer.Trainer property)

D

DataParallelGradientHandler (class in colossalai.engine.gradient_handler) DataParallelSampler (class in colossalai.utils) (class in colossalai.utils.data_sampler) (class in colossalai.utils.data_sampler.data_parallel_sampler) debug() (colossalai.logging.DistributedLogger method) (colossalai.logging.logging.DistributedLogger method) DelayerScheduler (class in colossalai.nn.lr_scheduler.delayed) destroy() (colossalai.context.parallel_context.ParallelContext method) DistributedLogger (class in colossalai.logging) (class in colossalai.logging.logging)

divide() (in module colossalai.nn.layer.utils) (in module colossalai.nn.layer.utils.common) drop_path() (in module colossalai.nn.layer.vanilla.layers) Dropout (class in colossalai.nn.layer.colossalai_layer) (class in colossalai.nn.layer.colossalai_layer.dropout) Dropout1D (class in colossalai.nn.layer.parallel_1d) (class in colossalai.nn.layer.parallel_1d.layers) DropPath (class in colossalai.nn.layer.vanilla) (class in colossalai.nn.layer.vanilla.layers)

E

Embedding (class in colossalai.nn.layer.colossalai_layer) (class in colossalai.nn.layer.colossalai_layer.embedding) Embedding1D (class in colossalai.nn.layer.parallel_1d) (class in colossalai.nn.layer.parallel_1d.layers) Embedding2D (class in colossalai.nn.layer.parallel_2d) (class in colossalai.nn.layer.parallel_2d.layers) Embedding2p5D (class in colossalai.nn.layer.parallel_2p5d) (class in colossalai.nn.layer.parallel_2p5d.layers) Embedding3D (class in colossalai.nn.layer.parallel_3d) (class in colossalai.nn.layer.parallel_3d.layers)

empty_cache() (in module colossalai.utils) (in module colossalai.utils.cuda) Engine (class in colossalai.engine) error() (colossalai.logging.DistributedLogger method) (colossalai.logging.logging.DistributedLogger method) eval() (colossalai.engine.Engine method) evaluate() (colossalai.trainer.Trainer method) Experts (class in colossalai.nn.layer.moe) ExponentialLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.torch)

F

fit() (colossalai.trainer.Trainer method) FlatAnnealingLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.cosine) FlatAnnealingWarmupLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.cosine) forward() (colossalai.nn.loss.loss_2d.CrossEntropyLoss2D method) (colossalai.nn.loss.loss_2d.VocabParallelCrossEntropyLoss2D method) (colossalai.nn.loss.loss_2p5d.CrossEntropyLoss2p5D method) (colossalai.nn.loss.loss_2p5d.VocabParallelCrossEntropyLoss2p5D method) (colossalai.nn.loss.loss_3d.CrossEntropyLoss3D method) (colossalai.nn.loss.loss_3d.VocabParallelCrossEntropyLoss3D method) (colossalai.nn.metric.accuracy_2d.Accuracy2D method) (colossalai.nn.metric.accuracy_2p5d.Accuracy2p5D method) (colossalai.nn.metric.accuracy_3d.Accuracy3D method)

forward_backward_step() (colossalai.engine.schedule.BaseSchedule method) (colossalai.engine.schedule.NonPipelineSchedule method) (colossalai.engine.schedule.PipelineSchedule method) forward_step() (colossalai.engine.schedule.PipelineSchedule method) from_file() (colossalai.context.config.Config static method) FusedAdam (class in colossalai.nn.optimizer) (class in colossalai.nn.optimizer.fused_adam) FusedLAMB (class in colossalai.nn.optimizer) (class in colossalai.nn.optimizer.fused_lamb) FusedScaleMaskSoftmax (class in colossalai.kernel) (class in colossalai.kernel.cuda_native.scaled_softmax) FusedSGD (class in colossalai.nn.optimizer) (class in colossalai.nn.optimizer.fused_sgd)

G

gather_split_1d_tensor() (in module colossalai.communication.utils) get_checkpoint_path() (in module colossalai.utils.checkpointing) get_current_device() (in module colossalai.utils) (in module colossalai.utils.cuda) get_current_mode() (in module colossalai.context.random) get_dataloader() (in module colossalai.utils) (in module colossalai.utils.data_sampler) (in module colossalai.utils.data_sampler.data_parallel_sampler) get_default_parser() (in module colossalai.initialize) get_dist_logger() (in module colossalai.logging) get_elapsed_time() (colossalai.utils.Timer method) (colossalai.utils.timer.Timer method) get_global_rank() (colossalai.context.parallel_context.ParallelContext method) get_group() (colossalai.context.parallel_context.ParallelContext method) get_history_mean() (colossalai.utils.Timer method) (colossalai.utils.timer.Timer method) get_history_sum() (colossalai.utils.Timer method) (colossalai.utils.timer.Timer method)

get_instance() (colossalai.logging.DistributedLogger static method) (colossalai.logging.logging.DistributedLogger static method) get_latest_checkpoint_path() (in module colossalai.utils.checkpointing) get_latest_checkpoint_pattern() (in module colossalai.utils.checkpointing) get_local_rank() (colossalai.context.parallel_context.ParallelContext method) get_module() (colossalai.registry.registry.Registry method) get_next_global_rank() (colossalai.context.parallel_context.ParallelContext method) get_prev_global_rank() (colossalai.context.parallel_context.ParallelContext method) get_ranks_in_group() (colossalai.context.parallel_context.ParallelContext method) get_seeds() (in module colossalai.context.random) get_states() (in module colossalai.context.random) get_timer() (colossalai.utils.MultiTimer method) (colossalai.utils.timer.MultiTimer method) get_world_size() (colossalai.context.parallel_context.ParallelContext method) GradAccumDataloader (class in colossalai.utils.gradient_accumulation) GradAccumGradientHandler (class in colossalai.utils.gradient_accumulation) GradAccumLrSchedulerByStep (class in colossalai.utils.gradient_accumulation) GradAccumOptimizer (class in colossalai.utils.gradient_accumulation)

H

handle_gradient() (colossalai.engine.gradient_handler.BaseGradientHandler method) (colossalai.engine.gradient_handler.DataParallelGradientHandler method) (colossalai.engine.gradient_handler.MoeGradientHandler method) (colossalai.engine.gradient_handler.PipelineSharedModuleGradientHandler method) (colossalai.engine.gradient_handler.SequenceParallelGradientHandler method) (colossalai.engine.gradient_handler.ZeROGradientHandler method)

has() (colossalai.registry.registry.Registry method)

I

info() (colossalai.logging.DistributedLogger method) (colossalai.logging.logging.DistributedLogger method) init_dist_group() (colossalai.context.process_group_initializer.Initializer_1D method) (colossalai.context.process_group_initializer.initializer_1d.Initializer_1D method) (colossalai.context.process_group_initializer.Initializer_2D method) (colossalai.context.process_group_initializer.initializer_2d.Initializer_2D method) (colossalai.context.process_group_initializer.initializer_2d.Initializer_2D_Col method) (colossalai.context.process_group_initializer.initializer_2d.Initializer_2D_Row method) (colossalai.context.process_group_initializer.Initializer_2p5D method) (colossalai.context.process_group_initializer.initializer_2p5d.Initializer_2p5D method) (colossalai.context.process_group_initializer.initializer_2p5d.Initializer_2p5D_Col method) (colossalai.context.process_group_initializer.initializer_2p5d.Initializer_2p5D_Dep method) (colossalai.context.process_group_initializer.initializer_2p5d.Initializer_2p5D_ROW method) (colossalai.context.process_group_initializer.initializer_2p5d.Initializer_2p5D_XZ method) (colossalai.context.process_group_initializer.Initializer_3D method) (colossalai.context.process_group_initializer.initializer_3d.Initializer_3D method) (colossalai.context.process_group_initializer.initializer_3d.Initializer_3D_Input method) (colossalai.context.process_group_initializer.initializer_3d.Initializer_3D_Output method) (colossalai.context.process_group_initializer.initializer_3d.Initializer_3D_Weight method) (colossalai.context.process_group_initializer.Initializer_Data method) (colossalai.context.process_group_initializer.initializer_data.Initializer_Data method) (colossalai.context.process_group_initializer.Initializer_Model method) (colossalai.context.process_group_initializer.initializer_model.Initializer_Model method) (colossalai.context.process_group_initializer.Initializer_Moe method) (colossalai.context.process_group_initializer.initializer_moe.Initializer_Moe method) (colossalai.context.process_group_initializer.initializer_moe.Initializer_Moedata method) (colossalai.context.process_group_initializer.initializer_moe.Initializer_Moemodel method) (colossalai.context.process_group_initializer.Initializer_Pipeline method) (colossalai.context.process_group_initializer.initializer_pipeline.Initializer_Pipeline method) (colossalai.context.process_group_initializer.Initializer_Sequence method) (colossalai.context.process_group_initializer.initializer_sequence.Initializer_Sequence method) (colossalai.context.process_group_initializer.initializer_sequence.Initializer_Sequence_DP method) (colossalai.context.process_group_initializer.Initializer_Tensor method) (colossalai.context.process_group_initializer.initializer_tensor.Initializer_Tensor method) init_global_dist() (colossalai.context.parallel_context.ParallelContext method) init_parallel_groups() (colossalai.context.parallel_context.ParallelContext method)

init_runner_states() (colossalai.trainer.hooks.BaseHook method) initialize() (in module colossalai.initialize) Initializer_1D (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.initializer_1d) Initializer_2D (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.initializer_2d) Initializer_2D_Col (class in colossalai.context.process_group_initializer.initializer_2d) Initializer_2D_Row (class in colossalai.context.process_group_initializer.initializer_2d) Initializer_2p5D (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.initializer_2p5d) Initializer_2p5D_Col (class in colossalai.context.process_group_initializer.initializer_2p5d) Initializer_2p5D_Dep (class in colossalai.context.process_group_initializer.initializer_2p5d) Initializer_2p5D_ROW (class in colossalai.context.process_group_initializer.initializer_2p5d) Initializer_2p5D_XZ (class in colossalai.context.process_group_initializer.initializer_2p5d) Initializer_3D (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.initializer_3d) Initializer_3D_Input (class in colossalai.context.process_group_initializer.initializer_3d) Initializer_3D_Output (class in colossalai.context.process_group_initializer.initializer_3d) Initializer_3D_Weight (class in colossalai.context.process_group_initializer.initializer_3d) Initializer_Data (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.initializer_data) Initializer_Model (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.initializer_model) Initializer_Moe (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.initializer_moe) Initializer_Moedata (class in colossalai.context.process_group_initializer.initializer_moe) Initializer_Moemodel (class in colossalai.context.process_group_initializer.initializer_moe) Initializer_Pipeline (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.initializer_pipeline) Initializer_Sequence (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.initializer_sequence) Initializer_Sequence_DP (class in colossalai.context.process_group_initializer.initializer_sequence) Initializer_Tensor (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.initializer_tensor) is_first_rank() (colossalai.context.parallel_context.ParallelContext method) is_initialized() (colossalai.context.parallel_context.ParallelContext method) is_last_rank() (colossalai.context.parallel_context.ParallelContext method)

L

Lamb (class in colossalai.nn.optimizer) (class in colossalai.nn.optimizer.lamb) LambdaLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.torch) LambdaWrapper (class in colossalai.nn.layer.wrapper) (class in colossalai.nn.layer.wrapper.lambda_wrapper) lap() (colossalai.utils.Timer method) (colossalai.utils.timer.Timer method) Lars (class in colossalai.nn.optimizer) (class in colossalai.nn.optimizer.lars) launch() (in module colossalai.initialize) launch_from_openmpi() (in module colossalai.initialize) launch_from_slurm() (in module colossalai.initialize) launch_from_torch() (in module colossalai.initialize) LayerNorm (class in colossalai.nn.layer.colossalai_layer) (class in colossalai.nn.layer.colossalai_layer.normalization) LayerNorm2D (class in colossalai.nn.layer.parallel_2d) (class in colossalai.nn.layer.parallel_2d.layers) LayerNorm2p5D (class in colossalai.nn.layer.parallel_2p5d) (class in colossalai.nn.layer.parallel_2p5d.layers) LayerNorm3D (class in colossalai.nn.layer.parallel_3d) (class in colossalai.nn.layer.parallel_3d.layers) Linear (class in colossalai.nn.layer.colossalai_layer) (class in colossalai.nn.layer.colossalai_layer.linear) Linear1D (class in colossalai.nn.layer.parallel_1d) (class in colossalai.nn.layer.parallel_1d.layers)

Linear1D_Col (class in colossalai.nn.layer.parallel_1d) (class in colossalai.nn.layer.parallel_1d.layers) Linear1D_Row (class in colossalai.nn.layer.parallel_1d) (class in colossalai.nn.layer.parallel_1d.layers) Linear2D (class in colossalai.nn.layer.parallel_2d) (class in colossalai.nn.layer.parallel_2d.layers) Linear2p5D (class in colossalai.nn.layer.parallel_2p5d) (class in colossalai.nn.layer.parallel_2p5d.layers) Linear3D (class in colossalai.nn.layer.parallel_3d) (class in colossalai.nn.layer.parallel_3d.layers) LinearWarmupLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.linear) load_batch() (colossalai.engine.schedule.BaseSchedule method) load_checkpoint() (in module colossalai.utils.checkpointing) load_config() (colossalai.context.parallel_context.ParallelContext method) LoadCheckpointHook (class in colossalai.trainer.hooks) log_to_file() (colossalai.logging.DistributedLogger method) (colossalai.logging.logging.DistributedLogger method) LogMemoryByEpochHook (class in colossalai.trainer.hooks) LogMetricByEpochHook (class in colossalai.trainer.hooks) LogMetricByStepHook (class in colossalai.trainer.hooks) LogTimingByEpochHook (class in colossalai.trainer.hooks) LossHook (class in colossalai.trainer.hooks) LRSchedulerHook (class in colossalai.trainer.hooks)

M

MetricHook (class in colossalai.trainer.hooks) model (colossalai.engine.Engine property) module colossalai colossalai.amp colossalai.amp.amp_type colossalai.amp.apex_amp colossalai.amp.apex_amp.apex_amp colossalai.amp.naive_amp colossalai.amp.naive_amp.naive_amp colossalai.amp.torch_amp colossalai.amp.torch_amp.torch_amp colossalai.builder colossalai.builder.builder colossalai.builder.pipeline colossalai.communication colossalai.communication.collective colossalai.communication.p2p colossalai.communication.ring colossalai.communication.utils colossalai.context colossalai.context.config colossalai.context.parallel_context colossalai.context.parallel_mode colossalai.context.process_group_initializer colossalai.context.process_group_initializer.initializer_1d colossalai.context.process_group_initializer.initializer_2d colossalai.context.process_group_initializer.initializer_2p5d colossalai.context.process_group_initializer.initializer_3d colossalai.context.process_group_initializer.initializer_data colossalai.context.process_group_initializer.initializer_model colossalai.context.process_group_initializer.initializer_moe colossalai.context.process_group_initializer.initializer_pipeline colossalai.context.process_group_initializer.initializer_sequence colossalai.context.process_group_initializer.initializer_tensor colossalai.context.process_group_initializer.process_group_initializer colossalai.context.random colossalai.context.random.seed_manager colossalai.engine colossalai.engine.gradient_handler colossalai.engine.schedule colossalai.initialize colossalai.kernel colossalai.kernel.cuda_native colossalai.kernel.cuda_native.layer_norm colossalai.kernel.cuda_native.multihead_attention colossalai.kernel.cuda_native.scaled_softmax colossalai.kernel.jit colossalai.kernel.jit.bias_dropout_add colossalai.kernel.jit.bias_gelu colossalai.kernel.jit.option colossalai.logging colossalai.logging.logging colossalai.nn colossalai.nn.init colossalai.nn.layer colossalai.nn.layer.base_layer colossalai.nn.layer.colossalai_layer colossalai.nn.layer.colossalai_layer.dropout colossalai.nn.layer.colossalai_layer.embedding colossalai.nn.layer.colossalai_layer.linear colossalai.nn.layer.colossalai_layer.normalization colossalai.nn.layer.moe colossalai.nn.layer.moe.layers colossalai.nn.layer.parallel_1d colossalai.nn.layer.parallel_1d.layers colossalai.nn.layer.parallel_2d colossalai.nn.layer.parallel_2d.layers colossalai.nn.layer.parallel_2p5d colossalai.nn.layer.parallel_2p5d.layers colossalai.nn.layer.parallel_3d colossalai.nn.layer.parallel_3d.layers colossalai.nn.layer.parallel_sequence colossalai.nn.layer.parallel_sequence.layers colossalai.nn.layer.utils colossalai.nn.layer.utils.common colossalai.nn.layer.vanilla colossalai.nn.layer.vanilla.layers colossalai.nn.layer.wrapper colossalai.nn.layer.wrapper.lambda_wrapper colossalai.nn.layer.wrapper.pipeline_wrapper colossalai.nn.loss colossalai.nn.loss.loss_2d colossalai.nn.loss.loss_2p5d colossalai.nn.loss.loss_3d colossalai.nn.loss.loss_moe colossalai.nn.lr_scheduler colossalai.nn.lr_scheduler.cosine colossalai.nn.lr_scheduler.delayed colossalai.nn.lr_scheduler.linear colossalai.nn.lr_scheduler.multistep colossalai.nn.lr_scheduler.onecycle colossalai.nn.lr_scheduler.poly colossalai.nn.lr_scheduler.torch colossalai.nn.metric colossalai.nn.metric.accuracy_2d colossalai.nn.metric.accuracy_2p5d colossalai.nn.metric.accuracy_3d colossalai.nn.model colossalai.nn.model.model_from_config colossalai.nn.optimizer colossalai.nn.optimizer.colossalai_optimizer colossalai.nn.optimizer.fused_adam colossalai.nn.optimizer.fused_lamb colossalai.nn.optimizer.fused_sgd colossalai.nn.optimizer.lamb colossalai.nn.optimizer.lars colossalai.registry colossalai.registry.registry colossalai.trainer colossalai.trainer.hooks colossalai.utils colossalai.utils.activation_checkpoint colossalai.utils.checkpointing colossalai.utils.common colossalai.utils.cuda colossalai.utils.data_sampler colossalai.utils.data_sampler.base_sampler colossalai.utils.data_sampler.data_parallel_sampler colossalai.utils.gradient_accumulation colossalai.utils.memory colossalai.utils.multi_tensor_apply colossalai.utils.multi_tensor_apply.multi_tensor_apply colossalai.utils.timer colossalai.zero

MoeCrossEntropyLoss (class in colossalai.nn.loss.loss_moe) MoeGradientHandler (class in colossalai.engine.gradient_handler) MoeLayer (class in colossalai.nn.layer.moe) (class in colossalai.nn.layer.moe.layers) MoeLoss (class in colossalai.nn.loss.loss_moe) MultiHeadAttention (class in colossalai.kernel) (class in colossalai.kernel.cuda_native.multihead_attention) MultiplicativeLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.torch) MultiStepLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.multistep) MultiStepWarmupLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.multistep) MultiTensorApply (class in colossalai.utils.multi_tensor_apply.multi_tensor_apply) MultiTimer (class in colossalai.utils) (class in colossalai.utils.timer)

N

NaiveAMPModel (class in colossalai.amp.naive_amp.naive_amp) NaiveAMPOptimizer (class in colossalai.amp.naive_amp) (class in colossalai.amp.naive_amp.naive_amp)

NonPipelineSchedule (class in colossalai.engine.schedule) NormalNoiseGenerator (class in colossalai.nn.layer.moe)

O

OneCycleLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.onecycle)

optimizer (colossalai.engine.Engine property)

P

ParallelContext (class in colossalai.context.parallel_context) ParallelMode (class in colossalai.context.parallel_mode) PatchEmbedding (class in colossalai.nn.layer.colossalai_layer) (class in colossalai.nn.layer.colossalai_layer.embedding) PatchEmbedding2D (class in colossalai.nn.layer.parallel_2d) (class in colossalai.nn.layer.parallel_2d.layers) PatchEmbedding2p5D (class in colossalai.nn.layer.parallel_2p5d) (class in colossalai.nn.layer.parallel_2p5d.layers) PatchEmbedding3D (class in colossalai.nn.layer.parallel_3d) (class in colossalai.nn.layer.parallel_3d.layers) PipelineSchedule (class in colossalai.engine.schedule)

PipelineSharedModuleGradientHandler (class in colossalai.engine.gradient_handler) PolynomialLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.poly) PolynomialWarmupLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.poly) pre_processing() (colossalai.engine.schedule.BaseSchedule method) predict() (colossalai.trainer.Trainer method) print_rank_0() (in module colossalai.utils) (in module colossalai.utils.common) ProcessGroupInitializer (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.process_group_initializer)

R

recv_backward() (in module colossalai.communication) (in module colossalai.communication.p2p) recv_forward() (in module colossalai.communication) (in module colossalai.communication.p2p) recv_tensor_meta() (in module colossalai.communication) (in module colossalai.communication.utils) reduce_by_batch_2d() (in module colossalai.nn.layer.parallel_2d) reduce_by_batch_2p5d() (in module colossalai.nn.layer.parallel_2p5d) reduce_by_batch_3d() (in module colossalai.nn.layer.parallel_3d) reduce_scatter() (in module colossalai.communication) (in module colossalai.communication.collective)

register_module() (colossalai.registry.registry.Registry method) Registry (class in colossalai.registry.registry) report_memory_usage() (in module colossalai.utils) (in module colossalai.utils.memory) reset() (colossalai.utils.MultiTimer method) (colossalai.utils.Timer method) (colossalai.utils.timer.MultiTimer method) (colossalai.utils.timer.Timer method) ring_forward() (in module colossalai.communication) (in module colossalai.communication.ring) RingAV (class in colossalai.nn.layer.parallel_sequence) RingQK (class in colossalai.nn.layer.parallel_sequence)

S

save_checkpoint() (in module colossalai.utils.checkpointing) SaveCheckpointHook (class in colossalai.trainer.hooks) ScaledMaskedSoftmax (class in colossalai.kernel.cuda_native.scaled_softmax) ScaledUpperTriangMaskedSoftmax (class in colossalai.kernel.cuda_native.scaled_softmax) seed() (in module colossalai.context.random) SeedManager (class in colossalai.context.random.seed_manager) send_backward() (in module colossalai.communication) (in module colossalai.communication.p2p) send_backward_recv_backward() (in module colossalai.communication) (in module colossalai.communication.p2p) send_backward_recv_forward() (in module colossalai.communication) (in module colossalai.communication.p2p) send_forward() (in module colossalai.communication) (in module colossalai.communication.p2p) send_forward_backward_recv_forward_backward() (in module colossalai.communication) (in module colossalai.communication.p2p) send_forward_recv_backward() (in module colossalai.communication) (in module colossalai.communication.p2p) send_forward_recv_forward() (in module colossalai.communication) (in module colossalai.communication.p2p) send_tensor_meta() (in module colossalai.communication) (in module colossalai.communication.utils) SequenceParallelGradientHandler (class in colossalai.engine.gradient_handler) set_device() (colossalai.context.parallel_context.ParallelContext method) set_epoch() (colossalai.utils.data_sampler.data_parallel_sampler.DataParallelSampler method) (colossalai.utils.data_sampler.DataParallelSampler method) (colossalai.utils.DataParallelSampler method) set_jit_fusion_options() (in module colossalai.kernel.jit) (in module colossalai.kernel.jit.option) set_level() (colossalai.logging.DistributedLogger method) (colossalai.logging.logging.DistributedLogger method) set_mode() (colossalai.context.random.seed_manager.SeedManager method) (in module colossalai.context.random) set_seed() (colossalai.context.parallel_context.ParallelContext method) set_seed_states() (in module colossalai.context.random)

set_state() (colossalai.context.random.seed_manager.SeedManager method) set_to_cuda() (in module colossalai.utils) (in module colossalai.utils.cuda) split_batch_3d() (in module colossalai.nn.layer.parallel_3d) split_tensor_2d() (in module colossalai.nn.layer.parallel_2d) split_tensor_2p5d() (in module colossalai.nn.layer.parallel_2p5d) split_tensor_3d() (in module colossalai.nn.layer.parallel_3d) split_tensor_into_1d_equal_chunks() (in module colossalai.communication.utils) start() (colossalai.utils.MultiTimer method) (colossalai.utils.Timer method) (colossalai.utils.timer.MultiTimer method) (colossalai.utils.timer.Timer method) step() (colossalai.amp.torch_amp.torch_amp.TorchAMPOptimizer method) (colossalai.amp.torch_amp.TorchAMPOptimizer method) (colossalai.engine.Engine method) (colossalai.nn.optimizer.fused_adam.FusedAdam method) (colossalai.nn.optimizer.fused_lamb.FusedLAMB method) (colossalai.nn.optimizer.fused_sgd.FusedSGD method) (colossalai.nn.optimizer.FusedAdam method) (colossalai.nn.optimizer.FusedLAMB method) (colossalai.nn.optimizer.FusedSGD method) (colossalai.nn.optimizer.Lamb method) (colossalai.nn.optimizer.lamb.Lamb method) (colossalai.nn.optimizer.Lars method) (colossalai.nn.optimizer.lars.Lars method) StepLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.torch) stop() (colossalai.utils.MultiTimer method) (colossalai.utils.Timer method) (colossalai.utils.timer.MultiTimer method) (colossalai.utils.timer.Timer method) sync_model_param() (in module colossalai.utils) (in module colossalai.utils.common) synchronize() (in module colossalai.utils) (in module colossalai.utils.cuda)

T

TensorboardHook (class in colossalai.trainer.hooks) ThroughputHook (class in colossalai.trainer.hooks) Timer (class in colossalai.utils) (class in colossalai.utils.timer) Top1Router (class in colossalai.nn.layer.moe) (class in colossalai.nn.layer.moe.layers) Top2Router (class in colossalai.nn.layer.moe) (class in colossalai.nn.layer.moe.layers) TorchAMPLoss (class in colossalai.amp.torch_amp) (class in colossalai.amp.torch_amp.torch_amp)

TorchAMPModel (class in colossalai.amp.torch_amp) (class in colossalai.amp.torch_amp.torch_amp) TorchAMPOptimizer (class in colossalai.amp.torch_amp) (class in colossalai.amp.torch_amp.torch_amp) train() (colossalai.engine.Engine method) Trainer (class in colossalai.trainer) TransformerSelfAttentionRing (class in colossalai.nn.layer.parallel_sequence) (class in colossalai.nn.layer.parallel_sequence.layers)

V

VanillaClassifier (class in colossalai.nn.layer.vanilla) (class in colossalai.nn.layer.vanilla.layers) VanillaPatchEmbedding (class in colossalai.nn.layer.vanilla) (class in colossalai.nn.layer.vanilla.layers) VocabParallelClassifier1D (class in colossalai.nn.layer.parallel_1d) (class in colossalai.nn.layer.parallel_1d.layers) VocabParallelClassifier2D (class in colossalai.nn.layer.parallel_2d) (class in colossalai.nn.layer.parallel_2d.layers) VocabParallelClassifier2p5D (class in colossalai.nn.layer.parallel_2p5d) (class in colossalai.nn.layer.parallel_2p5d.layers) VocabParallelClassifier3D (class in colossalai.nn.layer.parallel_3d) (class in colossalai.nn.layer.parallel_3d.layers)

VocabParallelCrossEntropyLoss2D (class in colossalai.nn.loss.loss_2d) VocabParallelCrossEntropyLoss2p5D (class in colossalai.nn.loss.loss_2p5d) VocabParallelCrossEntropyLoss3D (class in colossalai.nn.loss.loss_3d) VocabParallelEmbedding1D (class in colossalai.nn.layer.parallel_1d) (class in colossalai.nn.layer.parallel_1d.layers) VocabParallelEmbedding2D (class in colossalai.nn.layer.parallel_2d) (class in colossalai.nn.layer.parallel_2d.layers) VocabParallelEmbedding2p5D (class in colossalai.nn.layer.parallel_2p5d) (class in colossalai.nn.layer.parallel_2p5d.layers) VocabParallelEmbedding3D (class in colossalai.nn.layer.parallel_3d) (class in colossalai.nn.layer.parallel_3d.layers)

W

WarmupDelayerScheduler (class in colossalai.nn.lr_scheduler.delayed) WarmupScheduler (class in colossalai.nn.lr_scheduler.delayed) warning() (colossalai.logging.DistributedLogger method) (colossalai.logging.logging.DistributedLogger method)

with_seed() (in module colossalai.context.random) WrappedDropout (class in colossalai.nn.layer.vanilla) (class in colossalai.nn.layer.vanilla.layers) WrappedDropPath (class in colossalai.nn.layer.vanilla) (class in colossalai.nn.layer.vanilla.layers)

Z

zero_grad() (colossalai.engine.Engine method)

ZeROGradientHandler (class in colossalai.engine.gradient_handler)