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) acquire() (colossalai.logging.logging.Handler method) 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) addHandler() (colossalai.logging.logging.Logger method) addLevelName() (in module colossalai.logging.logging) 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) AsyncMemoryMonitor (class in colossalai.utils.memory_tracer) (class in colossalai.utils.memory_tracer.async_memtracer) 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) basicConfig() (in module colossalai.logging.logging) 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) BucketTensorShardStrategy (class in colossalai.zero.shard_utils) (class in colossalai.zero.shard_utils.bucket_tensor_shard_strategy) BufferingFormatter (class in colossalai.logging.logging) 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)

C

calc_loss() (colossalai.engine.Engine method) callHandlers() (colossalai.logging.logging.Logger method) captureWarnings() (in module colossalai.logging.logging) 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) close() (colossalai.logging.logging.FileHandler method) (colossalai.logging.logging.Handler method) 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.apex_amp.apex_amp module colossalai.amp.naive_amp module colossalai.amp.naive_amp.grad_scaler module colossalai.amp.naive_amp.grad_scaler.base_grad_scaler module colossalai.amp.naive_amp.grad_scaler.constant_grad_scaler module colossalai.amp.naive_amp.grad_scaler.dynamic_grad_scaler 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.constants 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.core module colossalai.engine module colossalai.engine.gradient_handler module colossalai.engine.gradient_handler.utils module colossalai.engine.ophooks module colossalai.engine.ophooks.zero_hook module colossalai.engine.paramhooks module colossalai.engine.schedule module colossalai.global_variables 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.experts module colossalai.nn.layer.moe.layers module colossalai.nn.layer.moe.utils 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_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 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.cpu_adam 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.testing module colossalai.testing.comparison module colossalai.testing.utils 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_tracer module colossalai.utils.memory_tracer.async_memtracer module colossalai.utils.memory_tracer.memstats_collector module colossalai.utils.memory_tracer.model_data_memtracer module colossalai.utils.memory_tracer.test_memstats_collector 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.profiler.comm_profiler module colossalai.utils.profiler.pcie_profiler module colossalai.utils.profiler.prof_utils module colossalai.utils.tensor_detector module colossalai.utils.tensor_detector.tensor_detector module colossalai.utils.timer module colossalai.zero module colossalai.zero.init_ctx module colossalai.zero.init_ctx.init_context module colossalai.zero.shard_utils module colossalai.zero.shard_utils.base_shard_strategy module colossalai.zero.shard_utils.bucket_tensor_shard_strategy module colossalai.zero.shard_utils.tensor_shard_strategy module colossalai.zero.sharded_model module colossalai.zero.sharded_model.param_manager module colossalai.zero.sharded_model.reduce_scatter module colossalai.zero.sharded_model.sharded_grad module colossalai.zero.sharded_model.sharded_model module colossalai.zero.sharded_model.sharded_model_v2 module colossalai.zero.sharded_model.utils module colossalai.zero.sharded_optim module colossalai.zero.sharded_optim.bookkeeping module colossalai.zero.sharded_optim.bookkeeping.base_store module colossalai.zero.sharded_optim.bookkeeping.bucket_store module colossalai.zero.sharded_optim.bookkeeping.gradient_store module colossalai.zero.sharded_optim.bookkeeping.parameter_store module colossalai.zero.sharded_optim.bookkeeping.tensor_bucket module colossalai.zero.sharded_optim.sharded_optim module colossalai.zero.sharded_optim.sharded_optim_v2 module colossalai.zero.sharded_param module colossalai.zero.sharded_param.sharded_param module colossalai.zero.sharded_param.sharded_tensor module CommEvent (class in colossalai.utils.profiler.comm_profiler) CommHandler (class in colossalai.utils.profiler.comm_profiler) CommProfiler (class in colossalai.utils.profiler.comm_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) convert_to_zero() (in module colossalai.zero) converter() (colossalai.logging.logging.Formatter method) 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) createLock() (colossalai.logging.logging.Handler method) criterion (colossalai.engine.Engine property) critical() (colossalai.logging.logging.Logger method) (colossalai.logging.logging.LoggerAdapter method) (in module colossalai.logging.logging) 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.Logger method) (colossalai.logging.logging.LoggerAdapter method) (in module colossalai.logging.logging) DelayerScheduler (class in colossalai.nn.lr_scheduler.delayed) destroy() (colossalai.context.parallel_context.ParallelContext method) disable() (in module colossalai.logging.logging)

disable_existing_loggers() (in module colossalai.logging) DistributedLogger (class in colossalai.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) emit() (colossalai.logging.logging.FileHandler method) (colossalai.logging.logging.Handler method) (colossalai.logging.logging.NullHandler method) (colossalai.logging.logging.StreamHandler method) empty_cache() (in module colossalai.utils) (in module colossalai.utils.cuda)

Engine (class in colossalai.engine) error() (colossalai.logging.DistributedLogger method) (colossalai.logging.logging.Logger method) (colossalai.logging.logging.LoggerAdapter method) (in module colossalai.logging.logging) eval() (colossalai.engine.Engine method) evaluate() (colossalai.trainer.Trainer method) exception() (colossalai.logging.logging.Logger method) (colossalai.logging.logging.LoggerAdapter method) (in module colossalai.logging.logging) Experts (class in colossalai.nn.layer.moe) (class in colossalai.nn.layer.moe.experts) ExponentialLR (class in colossalai.nn.lr_scheduler) (class in colossalai.nn.lr_scheduler.torch)

F

fatal() (colossalai.logging.logging.Logger method) (in module colossalai.logging.logging) FFNExperts (class in colossalai.nn.layer.moe) (class in colossalai.nn.layer.moe.experts) FileHandler (class in colossalai.logging.logging) Filter (class in colossalai.logging.logging) filter() (colossalai.logging.logging.Filter method) findCaller() (colossalai.logging.logging.Logger method) 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) flush() (colossalai.logging.logging.Handler method) (colossalai.logging.logging.StreamHandler method) (colossalai.zero.sharded_model.reduce_scatter.ReduceScatterBucketer method) format() (colossalai.logging.logging.BufferingFormatter method) (colossalai.logging.logging.Formatter method) (colossalai.logging.logging.Handler method) formatException() (colossalai.logging.logging.Formatter method) formatFooter() (colossalai.logging.logging.BufferingFormatter method) formatHeader() (colossalai.logging.logging.BufferingFormatter method) formatStack() (colossalai.logging.logging.Formatter method) Formatter (class in colossalai.logging.logging) formatTime() (colossalai.logging.logging.Formatter 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.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) FP32LinearGate (class in colossalai.nn.layer.moe.layers) free() (colossalai.zero.sharded_model.reduce_scatter.ReduceScatterBucketer 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_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_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.MultiTimer method) (colossalai.utils.timer.MultiTimer method) get_world_size() (colossalai.context.parallel_context.ParallelContext method) getChild() (colossalai.logging.logging.Logger method) getEffectiveLevel() (colossalai.logging.logging.Logger method) (colossalai.logging.logging.LoggerAdapter method) getLevelName() (in module colossalai.logging.logging) getLogger() (in module colossalai.logging.logging) getLoggerClass() (in module colossalai.logging.logging) getLogRecordFactory() (in module colossalai.logging.logging) getMessage() (colossalai.logging.logging.LogRecord 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() (colossalai.logging.logging.Handler method) (colossalai.logging.logging.Logger method) (colossalai.logging.logging.NullHandler method) 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)

handleError() (colossalai.logging.logging.Handler method) Handler (class in colossalai.logging.logging) has() (colossalai.registry.registry.Registry method) hasHandlers() (colossalai.logging.logging.Logger method) (colossalai.logging.logging.LoggerAdapter method)

I

info() (colossalai.logging.DistributedLogger method) (colossalai.logging.logging.Logger method) (colossalai.logging.logging.LoggerAdapter method) (in module colossalai.logging.logging) 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_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_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) isEnabledFor() (colossalai.logging.logging.Logger method) (colossalai.logging.logging.LoggerAdapter 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() (colossalai.logging.logging.Logger method) (colossalai.logging.logging.LoggerAdapter method) (in module colossalai.logging.logging) log_to_file() (colossalai.logging.DistributedLogger method) Logger (class in colossalai.logging.logging) LoggerAdapter (class in colossalai.logging.logging) LogMemoryByEpochHook (class in colossalai.trainer.hooks) LogMetricByEpochHook (class in colossalai.trainer.hooks) LogMetricByStepHook (class in colossalai.trainer.hooks) LogRecord (class in colossalai.logging.logging) LogTimingByEpochHook (class in colossalai.trainer.hooks) LossHook (class in colossalai.trainer.hooks) LRSchedulerHook (class in colossalai.trainer.hooks)

M

makeLogRecord() (in module colossalai.logging.logging) makeRecord() (colossalai.logging.logging.Logger method) MetricHook (class in colossalai.trainer.hooks) model (colossalai.engine.Engine property) ModelDataTracer (class in colossalai.utils.memory_tracer.model_data_memtracer) 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.grad_scaler colossalai.amp.naive_amp.grad_scaler.base_grad_scaler colossalai.amp.naive_amp.grad_scaler.constant_grad_scaler colossalai.amp.naive_amp.grad_scaler.dynamic_grad_scaler 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.constants 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.core colossalai.engine colossalai.engine.gradient_handler colossalai.engine.gradient_handler.utils colossalai.engine.ophooks colossalai.engine.ophooks.zero_hook colossalai.engine.paramhooks colossalai.engine.schedule colossalai.global_variables 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.experts colossalai.nn.layer.moe.layers colossalai.nn.layer.moe.utils 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_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 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.cpu_adam 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.testing colossalai.testing.comparison colossalai.testing.utils 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_tracer colossalai.utils.memory_tracer.async_memtracer colossalai.utils.memory_tracer.memstats_collector colossalai.utils.memory_tracer.model_data_memtracer colossalai.utils.memory_tracer.test_memstats_collector colossalai.utils.moe colossalai.utils.multi_tensor_apply colossalai.utils.multi_tensor_apply.multi_tensor_apply colossalai.utils.profiler colossalai.utils.profiler.comm_profiler colossalai.utils.profiler.pcie_profiler colossalai.utils.profiler.prof_utils colossalai.utils.tensor_detector colossalai.utils.tensor_detector.tensor_detector colossalai.utils.timer colossalai.zero colossalai.zero.init_ctx colossalai.zero.init_ctx.init_context colossalai.zero.shard_utils colossalai.zero.shard_utils.base_shard_strategy colossalai.zero.shard_utils.bucket_tensor_shard_strategy colossalai.zero.shard_utils.tensor_shard_strategy colossalai.zero.sharded_model colossalai.zero.sharded_model.param_manager colossalai.zero.sharded_model.reduce_scatter colossalai.zero.sharded_model.sharded_grad colossalai.zero.sharded_model.sharded_model colossalai.zero.sharded_model.sharded_model_v2 colossalai.zero.sharded_model.utils colossalai.zero.sharded_optim colossalai.zero.sharded_optim.bookkeeping colossalai.zero.sharded_optim.bookkeeping.base_store colossalai.zero.sharded_optim.bookkeeping.bucket_store colossalai.zero.sharded_optim.bookkeeping.gradient_store colossalai.zero.sharded_optim.bookkeeping.parameter_store colossalai.zero.sharded_optim.bookkeeping.tensor_bucket colossalai.zero.sharded_optim.sharded_optim colossalai.zero.sharded_optim.sharded_optim_v2 colossalai.zero.sharded_param colossalai.zero.sharded_param.sharded_param colossalai.zero.sharded_param.sharded_tensor

MoeContext (class in colossalai.context.moe_context) MoeCrossEntropyLoss (class in colossalai.nn.loss.loss_moe) MoeExperts (class in colossalai.nn.layer.moe.experts) 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) MoeModule (class in colossalai.nn.layer.moe) (class in colossalai.nn.layer.moe.layers) 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) (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) (class in colossalai.nn.layer.moe.utils) NullHandler (class in colossalai.logging.logging)

O

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

optimizer (colossalai.engine.Engine property) OptimState (class in colossalai.zero.sharded_optim.sharded_optim_v2)

P

ParallelContext (class in colossalai.context.parallel_context) ParallelMode (class in colossalai.context.parallel_mode) parameterize() (in module colossalai.testing) (in module colossalai.testing.utils) 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) PcieEvent (class in colossalai.utils.profiler.pcie_profiler) PcieProfiler (class in colossalai.utils.profiler.pcie_profiler)

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) process() (colossalai.logging.logging.LoggerAdapter method) ProcessGroupInitializer (class in colossalai.context.process_group_initializer) (class in colossalai.context.process_group_initializer.process_group_initializer) ProfilerContext (class in colossalai.utils.profiler.prof_utils)

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) 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) release() (colossalai.logging.logging.Handler method) removeHandler() (colossalai.logging.logging.Logger method) report_memory_usage() (in module colossalai.utils) rerun_on_exception() (in module colossalai.testing) (in module colossalai.testing.utils) 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) 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) setFormatter() (colossalai.logging.logging.Handler method) setLevel() (colossalai.logging.logging.Handler method) (colossalai.logging.logging.Logger method) (colossalai.logging.logging.LoggerAdapter method)

setLoggerClass() (in module colossalai.logging.logging) setLogRecordFactory() (in module colossalai.logging.logging) setStream() (colossalai.logging.logging.StreamHandler method) ShardedModelV2 (class in colossalai.zero.sharded_model) (class in colossalai.zero.sharded_model.sharded_model_v2) ShardedOptimizerV2 (class in colossalai.zero.sharded_optim) (class in colossalai.zero.sharded_optim.sharded_optim_v2) shutdown() (in module colossalai.logging.logging) 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) StreamHandler (class in colossalai.logging.logging) sync_model_param() (in module colossalai.utils) (in module colossalai.utils.common) sync_moe_model_param() (in module colossalai.utils.moe) synchronize() (in module colossalai.utils) (in module colossalai.utils.cuda)

T

TensorboardHook (class in colossalai.trainer.hooks) TensorShardStrategy (class in colossalai.zero.shard_utils) (class in colossalai.zero.shard_utils.tensor_shard_strategy) 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) TPExperts (class in colossalai.nn.layer.moe) (class in colossalai.nn.layer.moe.experts) train() (colossalai.engine.Engine method) Trainer (class in colossalai.trainer) TrainingState (class in colossalai.zero.sharded_model.sharded_model) TransformerSelfAttentionRing (class in colossalai.nn.layer.parallel_sequence) (class in colossalai.nn.layer.parallel_sequence.layers)

U

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

usesTime() (colossalai.logging.logging.Formatter method)

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)

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) (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.Logger method) (colossalai.logging.logging.LoggerAdapter method) (in module colossalai.logging.logging)

weight (colossalai.nn.layer.moe.layers.FP32LinearGate attribute) 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)

ZeroHook (class in colossalai.engine.ophooks.zero_hook) ZeroInitContext (class in colossalai.zero.init_ctx) (class in colossalai.zero.init_ctx.init_context)