noether.core.distributed.gather¶

Classes¶

AllGatherGradAutograd

Gathers tensors from all process and supports backward propagation for the gradients across processes.

Functions¶

`get_device_and_bfloat16supported`()
`get_bool_gather_supported`()
`all_gather_grad`(x[, batch_dim])
`all_gather_grad_autograd`(x)
`all_gather_nograd`(x)
`all_gather_nograd_clipped`(x, max_length)
`all_reduce_sum_nograd`(x)
`all_reduce_sum_grad`(x)
`reduce_mean_grad`(x[, dest_rank])
`reduce_mean_nograd`(x[, dest_rank])
`reduce_max_grad`(x[, dest_rank])
`reduce_max_nograd`(x[, dest_rank])
`all_reduce_mean_grad`(x)
`all_reduce_mean_nograd`(x)

Module Contents¶

class noether.core.distributed.gather.AllGatherGradAutograd¶

Bases: torch.autograd.Function

Gathers tensors from all process and supports backward propagation for the gradients across processes.

static forward(ctx, x)¶

static backward(ctx, *grads)¶

noether.core.distributed.gather.get_device_and_bfloat16supported()¶

noether.core.distributed.gather.get_bool_gather_supported()¶

noether.core.distributed.gather.all_gather_grad(x, batch_dim=0)¶

noether.core.distributed.gather.all_gather_grad_autograd(x)¶

noether.core.distributed.gather.all_gather_nograd(x)¶

noether.core.distributed.gather.all_gather_nograd_clipped(x, max_length)¶

noether.core.distributed.gather.all_reduce_sum_nograd(x)¶

noether.core.distributed.gather.all_reduce_sum_grad(x)¶

noether.core.distributed.gather.reduce_mean_grad(x, dest_rank=0)¶

noether.core.distributed.gather.reduce_mean_nograd(x, dest_rank=0)¶

noether.core.distributed.gather.reduce_max_grad(x, dest_rank=0)¶

noether.core.distributed.gather.reduce_max_nograd(x, dest_rank=0)¶

noether.core.distributed.gather.all_reduce_mean_grad(x)¶

noether.core.distributed.gather.all_reduce_mean_nograd(x)¶