bagua.torch_api.algorithms.onebit_adam

Module Contents

class bagua.torch_api.algorithms.onebit_adam.OnebitAdamAlgorithm(onebit_optimizer, warmup_steps, hierarchical_reduce=True)

Bases: bagua.torch_api.algorithms.Algorithm

need_reset(self)
init_tensors(self, bagua_module)
Parameters

bagua_module (bagua.torch_api.distributed.BaguaModule) –

tensors_to_buckets(self, tensors)

Given the bucketing suggestion from Bagua, return the actual Bagua buckets. The default implementation follows the suggestion to do the bucketing.

Parameters

tensors (List[List[bagua.torch_api.tensor.BaguaTensor]]) – Bagua tensors grouped in different lists, representing Bagua’s suggestion on how to bucketing the tensors.

Returns

A list of Bagua buckets.

Return type

List[bagua.torch_api.bucket.BaguaBucket]

init_operations(self, bagua_module, bucket)
Parameters
init_backward_hook(self, bagua_module)
Parameters

bagua_module (bagua.torch_api.distributed.BaguaModule) –

class bagua.torch_api.algorithms.onebit_adam.OnebitAdamOptimizer(params, lr=0.001, warmup_steps=100, is_bert=False, freeze_test_step=- 1, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)

Bases: torch.optim.optimizer.Optimizer

step(self, closure=None)