# tensorplay.distributed.optim API

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

## Functions 1

[#](#api-tensorplay.distributed.optim.as_functional_optim)

### as_functional_optim

function[Full reference ↗](/docs/generated/tensorplay.distributed.optim.as_functional_optim.html)

```python
tensorplay.distributed.optim.as_functional_optim(optim_cls: type, *args, **kwargs)
```

## Classes 3

[#](#api-tensorplay.distributed.optim.DistributedOptimizer)

### DistributedOptimizer

class[Full reference ↗](/docs/generated/tensorplay.distributed.optim.DistributedOptimizer.html)

```python
class tensorplay.distributed.optim.DistributedOptimizer(optimizer_class, params_rref, *args, **kwargs)
```

DistributedOptimizer takes remote references to parameters and runs the

This class requires the RPC framework.

[#](#api-tensorplay.distributed.optim.PostLocalSGDOptimizer)

### PostLocalSGDOptimizer

class[Full reference ↗](/docs/generated/tensorplay.distributed.optim.PostLocalSGDOptimizer.html)

```python
class tensorplay.distributed.optim.PostLocalSGDOptimizer(optim: Optimizer, averager: ModelAverager)
```

Wraps an arbitrary tensorplay.optim.Optimizer and runs [post-local SGD](https://arxiv.org/abs/1808.07217),
This optimizer runs local optimizer at every step.
After the warm-up stage, it averages parameters periodically after the local optimizer is applied.

Parameters:

- optim – The local optimizer.

- averager – A model averager instance to run post-localSGD algorithm.

Example:

```
>>> # xdoctest: +SKIP("undefined variables")
>>> local_optim = tp.optim.SGD(params=model.parameters(), lr=0.01)
>>> opt = PostLocalSGDOptimizer(
>>>     optim=local_optim,
>>>     averager=averagers.PeriodicModelAverager(period=4, warmup_steps=100)
>>> )
>>> for step in range(0, 200):
>>>    opt.zero_grad()
>>>    loss = loss_fn(output, labels)
>>>    loss.backward()
>>>    opt.step()
```

```python
load_state_dict(state_dict)
```

This is the same as tensorplay.optim.Optimizer [load_state_dict()](#tensorplay.distributed.optim.PostLocalSGDOptimizer.load_state_dict),
but also restores model averager’s step value to the one
saved in the provided state_dict.

If there is no "step" entry in state_dict,
it will raise a warning and initialize the model averager’s step to 0.

```python
state_dict()
```

This is the same as tensorplay.optim.Optimizer [state_dict()](#tensorplay.distributed.optim.PostLocalSGDOptimizer.state_dict),
but adds an extra entry to record model averager’s step to the checkpoint
to ensure reload does not cause unnecessary warm up again.

```python
step()
```

Performs a single optimization step (parameter update).

[#](#api-tensorplay.distributed.optim.ZeroRedundancyOptimizer)

### ZeroRedundancyOptimizer

class[Full reference ↗](/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.