apply_optimizer_in_backward (TorchRec)

This API is an open-source API of the TorchRec and is not an external API of the Rec SDK Torch. This section describes the parameter ranges supported by the TorchRec APIs called when the Rec SDK Torch is used.

Function

Specifies the optimizer used by a table.

Prototype

def apply_optimizer_in_backward(
    optimizer_class: Type[torch.optim.Optimizer],
    params: Iterable[torch.nn.Parameter],
    optimizer_kwargs: Dict[str, Any],
) -> None:

Parameters

Parameter	Type	Mandatory/Optional	Description
optimizer_class	Type[torch.optim.Optimizer]	Mandatory	Optimizer type. The options are as follows: torch.optim.Adagrad: Adagrad optimizer. torch.optim.Adam: Adam optimizer. torch.optim.SGD: SGD optimizer. torchrec.optim.AccumulateAdagrad: Adagrad optimizer with the gradient accumulation function. torchrec.optim.AccumulateAdam: Adam optimizer with the gradient accumulation function. torchrec.optim.AccumulateSGD: SGD optimizer with the gradient accumulation function. NOTE: The optimizer type must be the same as that configured during table creation. Currently, the gradient accumulation function supports only the table creation interface of EmbCacheEmbeddingCollection. When torchrec.optim.AccumulateSGD is used, it is recommended that the number of aggregated embeddings be less than or equal to 10,000. If the limit is exceeded, the computation precision may not meet the requirement of 0.01%.
params	Iterable[torch.nn.Parameter]	Mandatory	Sets the torch.nn.Parameter object of the optimizer. Pass the parameters of the HashEmbeddingBagCollection object by referring to 5. NOTICE: For performance purposes, the params parameter cannot be verified. Ensure that the type is correct.
optimizer_kwargs	Dict[str, Any]	Mandatory	Set this parameter based on the parameters and range of optimizer_class. Ensure that the range of this parameter meets the restrictions of the corresponding optimizer. If the gradient accumulation function is used, optimizer_class must be set to an optimizer with the gradient accumulation function, and use_accumulate (bool type) whose value is True (default value: False) and accumulate_step (int type) whose value is 1 (default value: 1) must be transferred.

Parameter

Type

Mandatory/Optional

Description

optimizer_class

Type[torch.optim.Optimizer]

Mandatory

Optimizer type.

The options are as follows:

torch.optim.Adagrad: Adagrad optimizer.
torch.optim.Adam: Adam optimizer.
torch.optim.SGD: SGD optimizer.
torchrec.optim.AccumulateAdagrad: Adagrad optimizer with the gradient accumulation function.
torchrec.optim.AccumulateAdam: Adam optimizer with the gradient accumulation function.
torchrec.optim.AccumulateSGD: SGD optimizer with the gradient accumulation function.

NOTE:

The optimizer type must be the same as that configured during table creation. Currently, the gradient accumulation function supports only the table creation interface of EmbCacheEmbeddingCollection.
When torchrec.optim.AccumulateSGD is used, it is recommended that the number of aggregated embeddings be less than or equal to 10,000. If the limit is exceeded, the computation precision may not meet the requirement of 0.01%.

params

Iterable[torch.nn.Parameter]

Mandatory

Sets the torch.nn.Parameter object of the optimizer. Pass the parameters of the HashEmbeddingBagCollection object by referring to 5.

NOTICE:

For performance purposes, the params parameter cannot be verified. Ensure that the type is correct.

optimizer_kwargs

Dict[str, Any]

Mandatory

Set this parameter based on the parameters and range of optimizer_class. Ensure that the range of this parameter meets the restrictions of the corresponding optimizer.

If the gradient accumulation function is used, optimizer_class must be set to an optimizer with the gradient accumulation function, and use_accumulate (bool type) whose value is True (default value: False) and accumulate_step (int type) whose value is 1 (default value: 1) must be transferred.

Return Value

Success: None
Failure: An exception is thrown.

Example

from torchrec.optim.apply_optimizer_in_backward import apply_optimizer_in_backward
import torch
embedding_optimizer = torch.optim.Adagrad
optimizer_kwargs = {"lr": 0.001, "eps": 0.1}
apply_optimizer_in_backward(
	embedding_optimizer,
	ebc.parameters(),
	optimizer_kwargs=optimizer_kwargs)

apply_optimizer_in_backward (TorchRec)

Function

Prototype

Parameters

Return Value

Example

See Also