> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.loss.dist_utils

## Module Contents

### Classes

| Name                                                                                  | Description                                                         |
| ------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| [`_AllGatherWithGrad`](#nemo_automodel-components-loss-dist_utils-_AllGatherWithGrad) | All-gather on dim-0 with autograd support across distributed ranks. |

### Functions

| Name                                                                                      | Description                                                 |
| ----------------------------------------------------------------------------------------- | ----------------------------------------------------------- |
| [`all_gather_no_grad`](#nemo_automodel-components-loss-dist_utils-all_gather_no_grad)     | Gather `x` across ranks on dim-0 without autograd.          |
| [`all_gather_with_grad`](#nemo_automodel-components-loss-dist_utils-all_gather_with_grad) | Gather `x` across ranks on dim-0 while preserving autograd. |

### API

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

**Bases:** `Function`

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

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

staticmethod

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

staticmethod

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

Gather `x` across ranks on dim-0 without autograd.

```python
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.