nemo_automodel.components.loss.dist_utils

Module Contents

Name	Description
`_AllGatherWithGrad`	All-gather on dim-0 with autograd support across distributed ranks.

Name	Description
`all_gather_no_grad`	Gather `x` across ranks on dim-0 without autograd.
`all_gather_with_grad`	Gather `x` across ranks on dim-0 while preserving autograd.

class nemo_automodel.components.loss.dist_utils._AllGatherWithGrad()

Bases: Function

All-gather on dim-0 with autograd support across distributed ranks.

nemo_automodel.components.loss.dist_utils._AllGatherWithGrad.backward(
    ctx,
    grads: torch.Tensor = ()
) -> torch.Tensor

staticmethod

nemo_automodel.components.loss.dist_utils._AllGatherWithGrad.forward(
    ctx,
    x: torch.Tensor
) -> tuple[torch.Tensor, ...]

staticmethod

nemo_automodel.components.loss.dist_utils.all_gather_no_grad(
    x: torch.Tensor
) -> torch.Tensor

Gather x across ranks on dim-0 without autograd.

nemo_automodel.components.loss.dist_utils.all_gather_with_grad(
    x: torch.Tensor
) -> torch.Tensor

Gather x across ranks on dim-0 while preserving autograd.