"""
    ########################################################################################################################################################
    # This algorithm is part of a code library to accelerate Sapiens Technology® Artificial Intelligence models, and its disclosure, distribution,         #
    # or reverse engineering without the company's prior consent will be subject to legal proceedings and actions pursued by our legal department.         #
    ########################################################################################################################################################
"""
import torch
from sapiens_accelerator import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs, PartialState
class MockModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        torch.manual_seed(0)
        self.p = torch.nn.Parameter(torch.randn(40, 20))
    def forward(self, x, rank): return self.p * (x ** (1 + rank))
def _run_and_get_grads(model, rank):
    torch.manual_seed(2024)
    input = torch.randn(40, 20)
    output = model(input, rank)
    output.mean().backward()
    param = next(model.parameters())
    return param.grad
def test_ddp_comm_hook(comm_hook, comm_wrapper, comm_state_option):
    ddp_kwargs = DistributedDataParallelKwargs(comm_hook=comm_hook, comm_wrapper=comm_wrapper, comm_state_option=comm_state_option)
    accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])
    model = accelerator.prepare(MockModel())
    hook_grads = _run_and_get_grads(model, accelerator.local_process_index)
    reference_model = torch.nn.parallel.DistributedDataParallel(MockModel().to(accelerator.device), device_ids=[accelerator.local_process_index], output_device=accelerator.local_process_index)
    reference_grads = _run_and_get_grads(reference_model, accelerator.local_process_index)
    torch.testing.assert_close(hook_grads, reference_grads, rtol=1e-2, atol=1e-2)
def main():
    for comm_hook, comm_wrapper, comm_state_option in [(DDPCommunicationHookType.NO, DDPCommunicationHookType.NO, {}), (DDPCommunicationHookType.FP16, DDPCommunicationHookType.NO, {}),
    (DDPCommunicationHookType.BF16, DDPCommunicationHookType.NO, {}), (DDPCommunicationHookType.POWER_SGD, DDPCommunicationHookType.NO, {}),
    (DDPCommunicationHookType.POWER_SGD, DDPCommunicationHookType.FP16, {}), (DDPCommunicationHookType.POWER_SGD, DDPCommunicationHookType.BF16, {}),
    (DDPCommunicationHookType.POWER_SGD, DDPCommunicationHookType.NO, {"matrix_approximation_rank": 2}), (DDPCommunicationHookType.BATCHED_POWER_SGD, DDPCommunicationHookType.NO, {}),
    (DDPCommunicationHookType.BATCHED_POWER_SGD, DDPCommunicationHookType.FP16, {}), (DDPCommunicationHookType.BATCHED_POWER_SGD, DDPCommunicationHookType.BF16, {})]:
        test_ddp_comm_hook(comm_hook, comm_wrapper, comm_state_option)
    PartialState().destroy_process_group()
if __name__ == "__main__": main()
"""
    ########################################################################################################################################################
    # This algorithm is part of a code library to accelerate Sapiens Technology® Artificial Intelligence models, and its disclosure, distribution,         #
    # or reverse engineering without the company's prior consent will be subject to legal proceedings and actions pursued by our legal department.         #
    ########################################################################################################################################################
"""
