colossalai

• colossalai.amp
  • colossalai.amp.apex_amp
  • colossalai.amp.naive_amp
  • colossalai.amp.torch_amp
  • colossalai.amp.amp_type
• 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.process_group_initializer
  • colossalai.context.random
  • colossalai.context.config
  • colossalai.context.parallel_context
  • colossalai.context.parallel_mode
• colossalai.engine
  • colossalai.engine.gradient_handler
  • colossalai.engine.schedule
• colossalai.kernel
  • colossalai.kernel.cuda_native
  • colossalai.kernel.jit
• colossalai.logging
  • colossalai.logging.logging
• colossalai.nn
  • colossalai.nn.layer
  • colossalai.nn.loss
  • colossalai.nn.lr_scheduler
  • colossalai.nn.metric
  • colossalai.nn.model
  • colossalai.nn.optimizer
  • colossalai.nn.init
• colossalai.registry
  • colossalai.registry.registry
• colossalai.trainer
  • colossalai.trainer.hooks
• colossalai.utils
  • colossalai.utils.data_sampler
  • colossalai.utils.gradient_accumulation
  • colossalai.utils.multi_tensor_apply
  • colossalai.utils.activation_checkpoint
  • colossalai.utils.checkpointing
  • colossalai.utils.common
  • colossalai.utils.cuda
  • colossalai.utils.memory
  • colossalai.utils.timer
• colossalai.zero
  • colossalai.zero.loss_scaler
  • colossalai.zero.zero_redundancy_optimizer_level_2
  • colossalai.zero.zero_redundancy_optimizer_level_3

• colossalai.initialize