tensorplay.distributed.optim API

DistributedOptimizer takes remote references to parameters and runs the

This class requires the RPC framework.

Wraps an arbitrary tensorplay.optim.Optimizer and runs post-local SGD, 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()

load_state_dict(state_dict)[source]

This is the same as tensorplay.optim.Optimizer 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.

state_dict()[source]

This is the same as tensorplay.optim.Optimizer 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.

step()[source]

Performs a single optimization step (parameter update).

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

The sharing is done as described by ZeRO.

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, optional) – if True, parameters are packed into buckets to speed up communication.

• overlap_with_ddp (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.

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

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.

consolidate_state_dict(to: int = 0) → None[source]

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

Parameters:

to (int) – the rank that receives the optimizer states (default: 0).

Warning

This needs to be called on all ranks.

property join_device

Return default device.

join_hook(**_kwargs: Any) → JoinHook[source]

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.

property join_process_group: Any

Return process group.

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

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

state_dict() → dict[str, Any][source]

Return the last global optimizer state known to this rank.

Raises:

RuntimeError – if this method is called without a preceding call to consolidate_state_dict().

step(closure=None, **kwargs: Any) → float | None[source]

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.