# ZeroRedundancyOptimizer

Source: https://www.tensorplay.cn/docs/generated/tensorplay.distributed.optim.ZeroRedundancyOptimizer.html

```python
class tensorplay.distributed.optim.ZeroRedundancyOptimizer(params, optimizer_class: type[Optimizer], process_group=None, parameters_as_bucket_view: bool = False, overlap_with_ddp: bool = False, **defaults: Any)
```

Wrap an arbitrary optim.Optimizer and shards its states across ranks in the group.

The sharing is done as described by [ZeRO](https://arxiv.org/abs/1910.02054).

The local optimizer instance in each rank is only
responsible for updating approximately 1 / world_size parameters and
hence only needs to keep 1 / world_size optimizer states. After
parameters are updated locally, each rank will broadcast its parameters to
all other peers to keep all model replicas in the same state.
ZeroRedundancyOptimizer uses a sorted-greedy algorithm to pack a number
of parameters at each rank. Each parameter belongs to a single rank and is
not divided among ranks.

Parameters:

params – an Iterable of tensors or dicts giving all parameters.

Keyword Arguments:

- optimizer_class – the class of the local optimizer.

- process_group – ProcessGroup (default: dist.GroupMember.WORLD).

- parameters_as_bucket_view ([bool](https://docs.python.org/3/builtins/functions.html#bool), optional) – if True, parameters are packed into buckets to speed up communication.

- overlap_with_ddp ([bool](https://docs.python.org/3/builtins/functions.html#bool), optional) – if True, requires a functional optimizer and registering one of the DDP communication hooks from ddp_zero_hook.py.

- **defaults – forwarded to the local optimizer.

> **Warning**
>
> Currently, ZeroRedundancyOptimizer requires that all of the passed-in parameters are the same dense type.

> **Warning**
>
> ZeroRedundancyOptimizer is experimental and subject to change.

```python
add_param_group(param_group: dict[str, Any]) → None
```

Add a parameter group to the Optimizer’s param_groups.

> **Warning**
>
> This method handles updating the shards on all partitions but needs to be called on all ranks.

```python
consolidate_state_dict(to: int = 0) → None
```

Consolidate a list of state_dict s (one per rank) on the target rank.

Parameters:

to ([int](https://docs.python.org/3/builtins/functions.html#int)) – the rank that receives the optimizer states (default: 0).

> **Warning**
>
> This needs to be called on all ranks.

```python
property join_device
```

Return default device.

```python
join_hook(**_kwargs: Any) → JoinHook
```

Return the ZeRO join hook.

It enables training on uneven inputs by
shadowing the collective communications in the optimizer step.

Gradients must be properly set before this hook is called.

```python
property join_process_group: Any
```

Return process group.

```python
load_state_dict(state_dict: dict[str, Any]) → None
```

Load the state pertaining to the given rank from the input state_dict, updating the local optimizer as needed.

```python
state_dict() → dict[str, Any]
```

Return the last global optimizer state known to this rank.

Raises:

[RuntimeError](https://docs.python.org/3/builtins/exceptions.html#RuntimeError) – if this method is called without a preceding call
    to [consolidate_state_dict()](#tensorplay.distributed.optim.ZeroRedundancyOptimizer.consolidate_state_dict).

```python
step(closure=None, **kwargs: Any) → float | None
```

Perform a single optimizer step and syncs parameters across all ranks.

Parameters:

closure (Callable) – a closure that re-evaluates the model and
returns the loss; optional for most optimizers.

Returns:

Optional loss depending on the underlying local optimizer.

> **Note**
>
> Any extra parameters are passed to the base optimizer as-is.