Index

A

accumulate_gradient() (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_hook() (colossalai.engine.Engine 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.collective) all_reduce() (in module colossalai.communication.collective) AMP_TYPE (class in colossalai.amp.amp_type) ApexAMPOptimizer (class in colossalai.amp.apex_amp) AsyncMemoryMonitor (class in colossalai.utils.memory_tracer) AttnMaskType (class in colossalai.kernel.cuda_native.scaled_softmax)

B

backward() (colossalai.amp.apex_amp.ApexAMPOptimizer 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) broadcast() (in module colossalai.communication.collective)

BucketTensorShardStrategy (class in colossalai.zero.shard_utils) build_data_sampler() (in module colossalai.builder.builder) build_dataset() (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.builder) build_hooks() (in module colossalai.builder.builder) build_layer() (in module colossalai.builder.builder) build_loss() (in module colossalai.builder.builder) build_lr_scheduler() (in module colossalai.builder.builder) build_model() (in module colossalai.builder.builder) build_ophooks() (in module colossalai.builder.builder) build_optimizer() (in module colossalai.builder.builder) build_pipeline_model() (in module colossalai.builder.pipeline) build_pipeline_model_from_cfg() (in module colossalai.builder.pipeline) build_schedule() (in module colossalai.builder.builder) build_transform() (in module colossalai.builder.builder)

C

check_sanity() (colossalai.context.parallel_context.ParallelContext method) checkpoint() (in module colossalai.utils.activation_checkpoint) Classifier (class in colossalai.nn.layer.colossalai_layer) Classifier1D (class in colossalai.nn.layer.parallel_1d) Classifier2D (class in colossalai.nn.layer.parallel_2d) Classifier2p5D (class in colossalai.nn.layer.parallel_2p5d) Classifier3D (class in colossalai.nn.layer.parallel_3d) clip_grad_norm() (colossalai.amp.apex_amp.ApexAMPOptimizer method) (colossalai.amp.torch_amp.TorchAMPOptimizer method) clip_grad_norm_fp32() (in module colossalai.utils.common) col_model_deepcopy() (in module colossalai.zero.sharded_model.utils) colossalai module colossalai.amp module colossalai.amp.amp_type module colossalai.amp.apex_amp module colossalai.amp.naive_amp module colossalai.amp.torch_amp module colossalai.builder.builder module colossalai.builder.pipeline 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.moe_context 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_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.ophooks module colossalai.engine.paramhooks 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.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.moe module colossalai.nn.layer.parallel_1d module colossalai.nn.layer.parallel_2d module colossalai.nn.layer.parallel_2p5d module colossalai.nn.layer.parallel_3d module colossalai.nn.layer.parallel_sequence module colossalai.nn.layer.utils module colossalai.nn.layer.vanilla module

colossalai.nn.layer.wrapper module colossalai.nn.loss module colossalai.nn.loss.loss_1d 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.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.colossalai_optimizer module colossalai.nn.optimizer.cpu_adam module colossalai.nn.optimizer.fused_adam module colossalai.nn.optimizer.fused_lamb module colossalai.nn.optimizer.fused_sgd module colossalai.nn.optimizer.hybrid_adam module colossalai.nn.optimizer.lamb module colossalai.nn.optimizer.lars module colossalai.registry module colossalai.registry.registry module colossalai.testing module colossalai.trainer module colossalai.trainer.hooks 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.gradient_accumulation module colossalai.utils.memory_tracer module colossalai.utils.moe module colossalai.utils.multi_tensor_apply module colossalai.utils.multi_tensor_apply.multi_tensor_apply module colossalai.utils.profiler module colossalai.utils.tensor_detector.tensor_detector module colossalai.utils.timer module colossalai.zero module colossalai.zero.init_ctx.init_context module colossalai.zero.shard_utils module colossalai.zero.sharded_model module colossalai.zero.sharded_model.reduce_scatter module colossalai.zero.sharded_model.sharded_model_v2 module colossalai.zero.sharded_model.utils module colossalai.zero.sharded_optim module colossalai.zero.sharded_param module CommProfiler (class in colossalai.utils.profiler) 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) CosineAnnealingLR (class in colossalai.nn.lr_scheduler.cosine) CosineAnnealingWarmupLR (class in colossalai.nn.lr_scheduler.cosine) count_layer_params() (in module colossalai.builder.pipeline) CPUAdam (class in colossalai.nn.optimizer.cpu_adam) 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.data_sampler) debug() (colossalai.logging.DistributedLogger method) DelayerScheduler (class in colossalai.nn.lr_scheduler.delayed) destroy() (colossalai.context.parallel_context.ParallelContext method) disable_existing_loggers() (in module colossalai.logging) DistributedLogger (class in colossalai.logging)

divide() (in module colossalai.nn.layer.utils) Dropout (class in colossalai.nn.layer.colossalai_layer) Dropout1D (class in colossalai.nn.layer.parallel_1d) DropPath (class in colossalai.nn.layer.vanilla) dump_memory_stats() (colossalai.zero.sharded_model.sharded_model_v2.ShardedModelV2 method) (colossalai.zero.sharded_model.ShardedModelV2 method) (colossalai.zero.ShardedModelV2 method)

E

Embedding (class in colossalai.nn.layer.colossalai_layer) Embedding1D (class in colossalai.nn.layer.parallel_1d) Embedding2D (class in colossalai.nn.layer.parallel_2d) Embedding2p5D (class in colossalai.nn.layer.parallel_2p5d) Embedding3D (class in colossalai.nn.layer.parallel_3d) empty_cache() (in module colossalai.utils.cuda)

Engine (class in colossalai.engine) error() (colossalai.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.torch)

F

FFNExperts (class in colossalai.nn.layer.moe) fit() (colossalai.trainer.Trainer method) FlatAnnealingLR (class in colossalai.nn.lr_scheduler.cosine) FlatAnnealingWarmupLR (class in colossalai.nn.lr_scheduler.cosine) flush() (colossalai.zero.sharded_model.reduce_scatter.ReduceScatterBucketer method) forward() (colossalai.nn.loss.loss_1d.VocabParallelCrossEntropyLoss1D method) (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.loss.loss_moe.MoeCrossEntropyLoss method) (colossalai.nn.loss.loss_moe.MoeLoss 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) FP16Optimizer (class in colossalai.amp.naive_amp) free() (colossalai.zero.sharded_model.reduce_scatter.ReduceScatterBucketer method) from_file() (colossalai.context.config.Config static method) FusedAdam (class in colossalai.nn.optimizer.fused_adam) FusedLAMB (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.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.cuda) get_current_mode() (in module colossalai.context.random) get_dataloader() (in module colossalai.utils.data_sampler) get_default_parser() (in module colossalai.initialize) get_dist_logger() (in module colossalai.logging) get_elapsed_time() (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.Timer method) get_history_sum() (colossalai.utils.timer.Timer method) get_info() (colossalai.context.moe_context.MoeContext method) get_instance() (colossalai.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_memory_usage() (colossalai.zero.sharded_optim.ShardedOptimizerV2 method) (colossalai.zero.ShardedOptimizerV2 method) get_module() (colossalai.registry.registry.Registry method) get_moe_epsize_param_dict() (in module colossalai.utils.moe) 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.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) HybridAdam (class in colossalai.nn.optimizer.hybrid_adam)

I

info() (colossalai.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_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_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_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_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_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) 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) 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) 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_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) 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)

K

kaiming_normal_() (in module colossalai.nn.init)

kaiming_uniform_() (in module colossalai.nn.init)

L

Lamb (class in colossalai.nn.optimizer.lamb) LambdaLR (class in colossalai.nn.lr_scheduler.torch) LambdaWrapper (class in colossalai.nn.layer.wrapper) lap() (colossalai.utils.timer.Timer method) Lars (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) LayerNorm2D (class in colossalai.nn.layer.parallel_2d) LayerNorm2p5D (class in colossalai.nn.layer.parallel_2p5d) LayerNorm3D (class in colossalai.nn.layer.parallel_3d) Linear (class in colossalai.nn.layer.colossalai_layer) Linear1D (class in colossalai.nn.layer.parallel_1d) Linear1D_Col (class in colossalai.nn.layer.parallel_1d)

Linear1D_Row (class in colossalai.nn.layer.parallel_1d) Linear2D (class in colossalai.nn.layer.parallel_2d) Linear2p5D (class in colossalai.nn.layer.parallel_2p5d) Linear3D (class in colossalai.nn.layer.parallel_3d) LinearWarmupLR (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) 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

MemProfiler (class in colossalai.utils.profiler) 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.naive_amp colossalai.amp.torch_amp colossalai.builder.builder colossalai.builder.pipeline colossalai.communication.collective colossalai.communication.p2p colossalai.communication.ring colossalai.communication.utils colossalai.context colossalai.context.config colossalai.context.moe_context 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_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.ophooks colossalai.engine.paramhooks 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.nn colossalai.nn.init colossalai.nn.layer colossalai.nn.layer.base_layer colossalai.nn.layer.colossalai_layer colossalai.nn.layer.moe colossalai.nn.layer.parallel_1d colossalai.nn.layer.parallel_2d colossalai.nn.layer.parallel_2p5d colossalai.nn.layer.parallel_3d colossalai.nn.layer.parallel_sequence colossalai.nn.layer.utils colossalai.nn.layer.vanilla colossalai.nn.layer.wrapper colossalai.nn.loss colossalai.nn.loss.loss_1d colossalai.nn.loss.loss_2d colossalai.nn.loss.loss_2p5d colossalai.nn.loss.loss_3d colossalai.nn.loss.loss_moe 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_optimizer colossalai.nn.optimizer.cpu_adam colossalai.nn.optimizer.fused_adam colossalai.nn.optimizer.fused_lamb colossalai.nn.optimizer.fused_sgd colossalai.nn.optimizer.hybrid_adam colossalai.nn.optimizer.lamb colossalai.nn.optimizer.lars colossalai.registry colossalai.registry.registry colossalai.testing colossalai.trainer colossalai.trainer.hooks colossalai.utils.activation_checkpoint colossalai.utils.checkpointing colossalai.utils.common colossalai.utils.cuda colossalai.utils.data_sampler colossalai.utils.gradient_accumulation colossalai.utils.memory_tracer colossalai.utils.moe colossalai.utils.multi_tensor_apply colossalai.utils.multi_tensor_apply.multi_tensor_apply colossalai.utils.profiler colossalai.utils.tensor_detector.tensor_detector colossalai.utils.timer colossalai.zero colossalai.zero.init_ctx.init_context colossalai.zero.shard_utils colossalai.zero.sharded_model colossalai.zero.sharded_model.reduce_scatter colossalai.zero.sharded_model.sharded_model_v2 colossalai.zero.sharded_model.utils colossalai.zero.sharded_optim colossalai.zero.sharded_param

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

N

NaiveAMPOptimizer (class in colossalai.amp.naive_amp) NonPipelineSchedule (class in colossalai.engine.schedule)

normal_() (in module colossalai.nn.init) NormalNoiseGenerator (class in colossalai.nn.layer.moe)

O

OneCycleLR (class in colossalai.nn.lr_scheduler.onecycle) ones_() (in module colossalai.nn.init)

ophooks (colossalai.engine.Engine property) optimizer (colossalai.engine.Engine property)

P

ParallelContext (class in colossalai.context.parallel_context) ParallelMode (class in colossalai.context.parallel_mode) parameterize() (in module colossalai.testing) PatchEmbedding (class in colossalai.nn.layer.colossalai_layer) PatchEmbedding2D (class in colossalai.nn.layer.parallel_2d) PatchEmbedding2p5D (class in colossalai.nn.layer.parallel_2p5d) PatchEmbedding3D (class in colossalai.nn.layer.parallel_3d) PcieProfiler (class in colossalai.utils.profiler) PipelineSchedule (class in colossalai.engine.schedule)

PipelineSharedModuleGradientHandler (class in colossalai.engine.gradient_handler) PolynomialLR (class in colossalai.nn.lr_scheduler.poly) PolynomialWarmupLR (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.common) ProcessGroupInitializer (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.process_group_initializer) ProfilerContext (class in colossalai.utils.profiler)

R

recv_backward() (in module colossalai.communication.p2p) recv_forward() (in module colossalai.communication.p2p) recv_tensor_meta() (in module colossalai.communication.utils) reduce() (in module colossalai.communication.collective) 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.collective) reduce_scatter_async() (colossalai.zero.sharded_model.reduce_scatter.ReduceScatterBucketer method) ReduceScatterBucketer (class in colossalai.zero.sharded_model.reduce_scatter)

register_module() (colossalai.registry.registry.Registry method) register_ophooks_recursively() (in module colossalai.engine.ophooks) Registry (class in colossalai.registry.registry) remove_hook() (colossalai.engine.Engine method) rerun_on_exception() (in module colossalai.testing) reset() (colossalai.utils.timer.MultiTimer method) (colossalai.utils.timer.Timer method) ring_forward() (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.p2p) send_backward_recv_backward() (in module colossalai.communication.p2p) send_backward_recv_forward() (in module colossalai.communication.p2p) send_forward() (in module colossalai.communication.p2p) send_forward_backward_recv_forward_backward() (in module colossalai.communication.p2p) send_forward_recv_backward() (in module colossalai.communication.p2p) send_forward_recv_forward() (in module colossalai.communication.p2p) send_tensor_meta() (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.DataParallelSampler method) set_jit_fusion_options() (in module colossalai.kernel.jit) (in module colossalai.kernel.jit.option) set_level() (colossalai.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.cuda) ShardedModelV2 (class in colossalai.zero) (class in colossalai.zero.sharded_model) (class in colossalai.zero.sharded_model.sharded_model_v2) ShardedOptimizerV2 (class in colossalai.zero) (class in colossalai.zero.sharded_optim) 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.timer.MultiTimer method) (colossalai.utils.timer.Timer method) step() (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.lamb.Lamb method) (colossalai.nn.optimizer.lars.Lars method) StepLR (class in colossalai.nn.lr_scheduler.torch) stop() (colossalai.utils.timer.MultiTimer method) (colossalai.utils.timer.Timer method) sync_model_param() (in module colossalai.utils.common) sync_moe_model_param() (in module colossalai.utils.moe) synchronize() (in module colossalai.utils.cuda)

T

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

TorchAMPModel (class in colossalai.amp.torch_amp) TorchAMPOptimizer (class in colossalai.amp.torch_amp) TPExperts (class in colossalai.nn.layer.moe) train() (colossalai.engine.Engine method) Trainer (class in colossalai.trainer) TransformerSelfAttentionRing (class in colossalai.nn.layer.parallel_sequence) trunc_normal_() (in module colossalai.nn.init)

U

uniform_() (in module colossalai.nn.init)

UniformNoiseGenerator (class in colossalai.nn.layer.moe)

V

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

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) VocabParallelEmbedding2D (class in colossalai.nn.layer.parallel_2d) VocabParallelEmbedding2p5D (class in colossalai.nn.layer.parallel_2p5d) VocabParallelEmbedding3D (class in colossalai.nn.layer.parallel_3d)

W

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

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

X

xavier_normal_() (in module colossalai.nn.init)

xavier_uniform_() (in module colossalai.nn.init)

Z

zero_grad() (colossalai.engine.Engine method) ZeroContextConfig (class in colossalai.zero.init_ctx.init_context)

ZeROGradientHandler (class in colossalai.engine.gradient_handler) ZeroInitContext (class in colossalai.zero.init_ctx.init_context) zeros_() (in module colossalai.nn.init)