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
- init_operations(self, bagua_module, bucket)¶
- Parameters
bagua_module (bagua.torch_api.distributed.BaguaModule) –
bucket (bagua.torch_api.bucket.BaguaBucket) –
- init_backward_hook(self, bagua_module)¶
- Parameters
bagua_module (bagua.torch_api.distributed.BaguaModule) –