Index

A

abort() (bagua.torch_api.algorithms.async_model_average.AsyncModelAverageAlgorithmImpl method) Algorithm (class in bagua.torch_api.algorithms) (class in bagua.torch_api.algorithms.base) AlgorithmImpl (class in bagua.torch_api.algorithms) (class in bagua.torch_api.algorithms.base) all_reduce() (in module bagua.torch_api.data_parallel.functional) allgather() (in module bagua.torch_api) (in module bagua.torch_api.communication) allgather_inplace() (in module bagua.torch_api) (in module bagua.torch_api.communication) allreduce() (in module bagua.torch_api) (in module bagua.torch_api.communication) allreduce_inplace() (in module bagua.torch_api) (in module bagua.torch_api.communication) alltoall() (in module bagua.torch_api) (in module bagua.torch_api.communication)

alltoall_inplace() (in module bagua.torch_api) (in module bagua.torch_api.communication) alltoall_v() (in module bagua.torch_api) alltoall_v_inplace() (in module bagua.torch_api) append_asynchronous_model_average_op() (bagua.torch_api.bucket.BaguaBucket method) append_centralized_synchronous_op() (bagua.torch_api.bucket.BaguaBucket method) append_decentralized_synchronous_op() (bagua.torch_api.bucket.BaguaBucket method) append_low_precision_decentralized_synchronous_op() (bagua.torch_api.bucket.BaguaBucket method) append_python_op() (bagua.torch_api.bucket.BaguaBucket method) AsyncModelAverageAlgorithm (class in bagua.torch_api.algorithms.async_model_average) AsyncModelAverageAlgorithmImpl (class in bagua.torch_api.algorithms.async_model_average) AVG (bagua.torch_api.communication.ReduceOp attribute) (bagua.torch_api.data_parallel.functional.ReduceOp attribute) (bagua.torch_api.ReduceOp attribute)

B

bagua module bagua.torch_api module bagua.torch_api.algorithms module bagua.torch_api.algorithms.async_model_average module bagua.torch_api.algorithms.base module bagua.torch_api.algorithms.bytegrad module bagua.torch_api.algorithms.decentralized module bagua.torch_api.algorithms.gradient_allreduce module bagua.torch_api.algorithms.q_adam module bagua.torch_api.bucket module bagua.torch_api.checkpoint module bagua.torch_api.checkpoint.checkpointing module bagua.torch_api.communication module bagua.torch_api.contrib module bagua.torch_api.contrib.cache_loader module bagua.torch_api.contrib.cached_dataset module bagua.torch_api.contrib.fuse module bagua.torch_api.contrib.fuse.optimizer module bagua.torch_api.contrib.load_balancing_data_loader module bagua.torch_api.contrib.sync_batchnorm module bagua.torch_api.contrib.utils module bagua.torch_api.contrib.utils.redis_store module bagua.torch_api.contrib.utils.store module bagua.torch_api.data_parallel module bagua.torch_api.data_parallel.bagua_distributed module bagua.torch_api.data_parallel.distributed module bagua.torch_api.data_parallel.functional module

bagua.torch_api.distributed module bagua.torch_api.env module bagua.torch_api.model_parallel module bagua.torch_api.model_parallel.moe module bagua.torch_api.model_parallel.moe.experts module bagua.torch_api.model_parallel.moe.layer module bagua.torch_api.model_parallel.moe.sharded_moe module bagua.torch_api.model_parallel.moe.utils module bagua.torch_api.tensor module bagua_backend_tensor() (bagua.torch_api.BaguaTensor method) (bagua.torch_api.tensor.BaguaTensor method) bagua_ensure_grad() (bagua.torch_api.BaguaTensor method) (bagua.torch_api.tensor.BaguaTensor method) bagua_getter_closure() (bagua.torch_api.BaguaTensor method) (bagua.torch_api.tensor.BaguaTensor method) bagua_mark_communication_ready() (bagua.torch_api.BaguaTensor method) (bagua.torch_api.tensor.BaguaTensor method) bagua_mark_communication_ready_without_synchronization() (bagua.torch_api.BaguaTensor method) (bagua.torch_api.tensor.BaguaTensor method) bagua_set_storage() (bagua.torch_api.BaguaTensor method) (bagua.torch_api.tensor.BaguaTensor method) bagua_setter_closure() (bagua.torch_api.BaguaTensor method) (bagua.torch_api.tensor.BaguaTensor method) BaguaBucket (class in bagua.torch_api.bucket) BaguaModule (class in bagua.torch_api) (class in bagua.torch_api.distributed) BaguaProcessGroup (class in bagua.torch_api.communication) BaguaTensor (class in bagua.torch_api) (class in bagua.torch_api.tensor) BAND (bagua.torch_api.communication.ReduceOp attribute) (bagua.torch_api.data_parallel.functional.ReduceOp attribute) (bagua.torch_api.ReduceOp attribute) barrier() (in module bagua.torch_api.communication) Base (in module bagua.torch_api.model_parallel.moe.sharded_moe) BOR (bagua.torch_api.communication.ReduceOp attribute) (bagua.torch_api.data_parallel.functional.ReduceOp attribute) (bagua.torch_api.ReduceOp attribute) broadcast() (in module bagua.torch_api) (in module bagua.torch_api.communication) broadcast_object() (in module bagua.torch_api.communication) BXOR (bagua.torch_api.communication.ReduceOp attribute) (bagua.torch_api.data_parallel.functional.ReduceOp attribute) (bagua.torch_api.ReduceOp attribute) ByteGradAlgorithm (class in bagua.torch_api.algorithms.bytegrad) ByteGradAlgorithmImpl (class in bagua.torch_api.algorithms.bytegrad) bytes() (bagua.torch_api.bucket.BaguaBucket method)

C

cache_loader (bagua.torch_api.contrib.cached_dataset.CachedDataset attribute) (bagua.torch_api.contrib.CachedDataset attribute) CachedDataset (class in bagua.torch_api.contrib) (class in bagua.torch_api.contrib.cached_dataset) CacheLoader (class in bagua.torch_api.contrib) (class in bagua.torch_api.contrib.cache_loader)

check_flatten() (bagua.torch_api.bucket.BaguaBucket method) clear() (bagua.torch_api.contrib.utils.store.Store method) clear_ops() (bagua.torch_api.bucket.BaguaBucket method) ClusterStore (class in bagua.torch_api.contrib.utils.store) convert_sync_batchnorm() (bagua.torch_api.contrib.sync_batchnorm.SyncBatchNorm class method)

D

DecentralizedAlgorithm (class in bagua.torch_api.algorithms.decentralized) DecentralizedAlgorithmImpl (class in bagua.torch_api.algorithms.decentralized)

DistributedDataParallel() (in module bagua.torch_api.data_parallel) (in module bagua.torch_api.data_parallel.distributed)

E

ensure_bagua_tensor() (bagua.torch_api.BaguaTensor method) (bagua.torch_api.tensor.BaguaTensor method)

exp_selection_uniform_map (in module bagua.torch_api.model_parallel.moe.sharded_moe) Experts (class in bagua.torch_api.model_parallel.moe.experts)

F

find_free_network_port() (in module bagua.torch_api.env) flattened_tensor() (bagua.torch_api.bucket.BaguaBucket method) forward() (bagua.torch_api.contrib.sync_batchnorm.SyncBatchNorm method) (bagua.torch_api.model_parallel.moe.experts.Experts method) (bagua.torch_api.model_parallel.moe.layer.MoE method) (bagua.torch_api.model_parallel.moe.MoE method) (bagua.torch_api.model_parallel.moe.sharded_moe.MOELayer method) (bagua.torch_api.model_parallel.moe.sharded_moe.TopKGate method)

from_torch_group() (in module bagua.torch_api.communication) fuse_optimizer() (in module bagua.torch_api.contrib) (in module bagua.torch_api.contrib.fuse.optimizer) fuse_step() (in module bagua.torch_api.contrib.fuse.optimizer)

G

gather() (in module bagua.torch_api) (in module bagua.torch_api.communication) gather_inplace() (in module bagua.torch_api) (in module bagua.torch_api.communication) get() (bagua.torch_api.contrib.cache_loader.CacheLoader method) (bagua.torch_api.contrib.CacheLoader method) (bagua.torch_api.contrib.utils.store.Store method) get_autotune_server_wait_time() (in module bagua.torch_api.env) get_backend() (in module bagua.torch_api) get_default_bucket_size() (in module bagua.torch_api.env) get_global_communicator() (bagua.torch_api.communication.BaguaProcessGroup method) get_inter_node_communicator() (bagua.torch_api.communication.BaguaProcessGroup method) get_intra_node_communicator() (bagua.torch_api.communication.BaguaProcessGroup method)

get_local_rank() (in module bagua.torch_api) (in module bagua.torch_api.env) get_local_size() (in module bagua.torch_api) (in module bagua.torch_api.env) get_node_rank() (in module bagua.torch_api.env) get_rank() (in module bagua.torch_api) (in module bagua.torch_api.env) get_world_size() (in module bagua.torch_api) (in module bagua.torch_api.env) GlobalAlgorithmRegistry (in module bagua.torch_api.algorithms) GradientAllReduceAlgorithm (class in bagua.torch_api.algorithms.gradient_allreduce) GradientAllReduceAlgorithmImpl (class in bagua.torch_api.algorithms.gradient_allreduce) gumbel_map (in module bagua.torch_api.model_parallel.moe.sharded_moe) gumbel_rsample() (in module bagua.torch_api.model_parallel.moe.sharded_moe)

I

init() (bagua.torch_api.algorithms.Algorithm class method) (bagua.torch_api.algorithms.base.Algorithm class method) init_backward_hook() (bagua.torch_api.algorithms.AlgorithmImpl method) (bagua.torch_api.algorithms.base.AlgorithmImpl method) init_forward_pre_hook() (bagua.torch_api.algorithms.AlgorithmImpl method) (bagua.torch_api.algorithms.base.AlgorithmImpl method) init_operations() (bagua.torch_api.algorithms.AlgorithmImpl method) (bagua.torch_api.algorithms.base.AlgorithmImpl method) init_post_backward_hook() (bagua.torch_api.algorithms.AlgorithmImpl method) (bagua.torch_api.algorithms.base.AlgorithmImpl method) init_post_optimizer_step_hook() (bagua.torch_api.algorithms.AlgorithmImpl method) (bagua.torch_api.algorithms.base.AlgorithmImpl method)

init_process_group() (in module bagua.torch_api) (in module bagua.torch_api.communication) init_tensors() (bagua.torch_api.algorithms.AlgorithmImpl method) (bagua.torch_api.algorithms.base.AlgorithmImpl method) is_bagua_tensor() (bagua.torch_api.BaguaTensor method) (bagua.torch_api.tensor.BaguaTensor method) is_fused_optimizer() (in module bagua.torch_api.contrib.fuse.optimizer) is_initialized() (in module bagua.torch_api.communication) is_moe_param() (in module bagua.torch_api.model_parallel.moe) (in module bagua.torch_api.model_parallel.moe.utils)

L

load_checkpoint() (in module bagua.torch_api.checkpoint) (in module bagua.torch_api.checkpoint.checkpointing) LoadBalancingDistributedBatchSampler (class in bagua.torch_api.contrib) (class in bagua.torch_api.contrib.load_balancing_data_loader)

LoadBalancingDistributedSampler (class in bagua.torch_api.contrib) (class in bagua.torch_api.contrib.load_balancing_data_loader) LowPrecisionDecentralizedAlgorithm (class in bagua.torch_api.algorithms.decentralized) LowPrecisionDecentralizedAlgorithmImpl (class in bagua.torch_api.algorithms.decentralized)

M

MAX (bagua.torch_api.communication.ReduceOp attribute) (bagua.torch_api.data_parallel.functional.ReduceOp attribute) (bagua.torch_api.ReduceOp attribute) mget() (bagua.torch_api.contrib.utils.store.Store method) MIN (bagua.torch_api.communication.ReduceOp attribute) (bagua.torch_api.data_parallel.functional.ReduceOp attribute) (bagua.torch_api.ReduceOp attribute) module bagua bagua.torch_api bagua.torch_api.algorithms bagua.torch_api.algorithms.async_model_average bagua.torch_api.algorithms.base bagua.torch_api.algorithms.bytegrad bagua.torch_api.algorithms.decentralized bagua.torch_api.algorithms.gradient_allreduce bagua.torch_api.algorithms.q_adam bagua.torch_api.bucket bagua.torch_api.checkpoint bagua.torch_api.checkpoint.checkpointing bagua.torch_api.communication bagua.torch_api.contrib bagua.torch_api.contrib.cache_loader bagua.torch_api.contrib.cached_dataset bagua.torch_api.contrib.fuse bagua.torch_api.contrib.fuse.optimizer bagua.torch_api.contrib.load_balancing_data_loader bagua.torch_api.contrib.sync_batchnorm bagua.torch_api.contrib.utils bagua.torch_api.contrib.utils.redis_store bagua.torch_api.contrib.utils.store bagua.torch_api.data_parallel bagua.torch_api.data_parallel.bagua_distributed bagua.torch_api.data_parallel.distributed bagua.torch_api.data_parallel.functional bagua.torch_api.distributed bagua.torch_api.env bagua.torch_api.model_parallel bagua.torch_api.model_parallel.moe bagua.torch_api.model_parallel.moe.experts bagua.torch_api.model_parallel.moe.layer bagua.torch_api.model_parallel.moe.sharded_moe bagua.torch_api.model_parallel.moe.utils bagua.torch_api.tensor

MoE (class in bagua.torch_api.model_parallel.moe) (class in bagua.torch_api.model_parallel.moe.layer) MOELayer (class in bagua.torch_api.model_parallel.moe.sharded_moe) mset() (bagua.torch_api.contrib.utils.store.Store method) multiplicative_jitter() (in module bagua.torch_api.model_parallel.moe.sharded_moe)

N

name (bagua.torch_api.bucket.BaguaBucket attribute) need_reset() (bagua.torch_api.algorithms.AlgorithmImpl method) (bagua.torch_api.algorithms.base.AlgorithmImpl method)

new_group() (in module bagua.torch_api.communication) num_keys() (bagua.torch_api.contrib.cache_loader.CacheLoader method) (bagua.torch_api.contrib.CacheLoader method) (bagua.torch_api.contrib.utils.store.Store method)

P

PRODUCT (bagua.torch_api.communication.ReduceOp attribute) (bagua.torch_api.data_parallel.functional.ReduceOp attribute) (bagua.torch_api.ReduceOp attribute)

Q

QAdamAlgorithm (class in bagua.torch_api.algorithms.q_adam)

QAdamAlgorithmImpl (class in bagua.torch_api.algorithms.q_adam) QAdamOptimizer (class in bagua.torch_api.algorithms.q_adam)

R

recv() (in module bagua.torch_api) (in module bagua.torch_api.communication) RedisStore (class in bagua.torch_api.contrib.utils.redis_store) reduce() (in module bagua.torch_api) (in module bagua.torch_api.communication) reduce_inplace() (in module bagua.torch_api) (in module bagua.torch_api.communication) reduce_scatter() (in module bagua.torch_api) (in module bagua.torch_api.communication)

reduce_scatter_inplace() (in module bagua.torch_api) (in module bagua.torch_api.communication) ReduceOp (class in bagua.torch_api) (class in bagua.torch_api.communication) (class in bagua.torch_api.data_parallel.functional) reify() (bagua.torch_api.algorithms.Algorithm method) (bagua.torch_api.algorithms.base.Algorithm method) resume() (bagua.torch_api.algorithms.async_model_average.AsyncModelAverageAlgorithmImpl method)

S

save_checkpoint() (in module bagua.torch_api.checkpoint) (in module bagua.torch_api.checkpoint.checkpointing) scatter() (in module bagua.torch_api) (in module bagua.torch_api.communication) scatter_inplace() (in module bagua.torch_api) (in module bagua.torch_api.communication) send() (in module bagua.torch_api) (in module bagua.torch_api.communication) set() (bagua.torch_api.contrib.utils.store.Store method) set_epoch() (bagua.torch_api.contrib.load_balancing_data_loader.LoadBalancingDistributedBatchSampler method) (bagua.torch_api.contrib.load_balancing_data_loader.LoadBalancingDistributedSampler method) (bagua.torch_api.contrib.LoadBalancingDistributedBatchSampler method) (bagua.torch_api.contrib.LoadBalancingDistributedSampler method)

shutdown() (bagua.torch_api.contrib.utils.store.Store method) status() (bagua.torch_api.contrib.utils.store.Store method) step() (bagua.torch_api.algorithms.q_adam.QAdamOptimizer method) Store (class in bagua.torch_api.contrib.utils.store) SUM (bagua.torch_api.communication.ReduceOp attribute) (bagua.torch_api.data_parallel.functional.ReduceOp attribute) (bagua.torch_api.ReduceOp attribute) SyncBatchNorm (class in bagua.torch_api.contrib.sync_batchnorm)

T

tensors (bagua.torch_api.bucket.BaguaBucket attribute) tensors_to_buckets() (bagua.torch_api.algorithms.AlgorithmImpl method) (bagua.torch_api.algorithms.base.AlgorithmImpl method) to_bagua_process_group() (in module bagua.torch_api.data_parallel.distributed)

to_bagua_tensor() (bagua.torch_api.BaguaTensor method) (bagua.torch_api.tensor.BaguaTensor method) top1gating() (in module bagua.torch_api.model_parallel.moe.sharded_moe) top2gating() (in module bagua.torch_api.model_parallel.moe.sharded_moe) TopKGate (class in bagua.torch_api.model_parallel.moe.sharded_moe)

U

uniform_map (in module bagua.torch_api.model_parallel.moe.sharded_moe)

V

version (in module bagua.torch_api)

W

wg (bagua.torch_api.model_parallel.moe.sharded_moe.TopKGate attribute)

with_bagua() (bagua.torch_api.BaguaModule method) (bagua.torch_api.distributed.BaguaModule method)