[Question] How to deal with multiple loss backward and optimize?

[WhoisZihan]

TorchRec will fuse the backward&optimize procedure for performance, but in some cases when some parts are not differentiable, I have to manually calculate part of the gradients and manually apply the gradients.

For example, in the code snippet below, there are two losses, one can be directly calculated using nn.MSELoss, another has non-differentiable part and must be calculated manually.

• Without fusion, there are 2 gradients accumulation and 1 optimization step
• With fusion, there are 2 (gradient+optimize) fused step.

My question is: now I perform 2 optimize, will it affect the convergence?
What is the best practice for such situation?

# linear layer and loss function
mse_loss = nn.MSELoss()

# input data and labels
x = torch.rand(3 ,5)
labels = torch.randint(low=0, high=2, size=(3, 1))

# logits
logits = model(x)

# loss
loss1 = mse_loss(logits, labels)
loss1.backward() # this will perform a backward&optimize fusion, and parameters will be updated?

# another loss with non-differentiable part, must be calculated manually
custom_gradient = NonDifferentiableLogic(...)
torch.autograd.backward(logits, grad_tensors=custom_gradient) # this will perform another backward&optimize fusion, and parameters will be updated again?

[dstaay-fb]

I think it probably depends a bit on model architecture; I think it might help to describe what 'fused' means. The assumption is your model is using TorchRec components for embedding tables, and by default we 'fuse' the optimizer step with backward pass for the embedding modules (it saves the requirement to materialize the full gradient tensor in memory).

However, this only applies to embedding modules, and not your other MLP/dense layers. Presumably for those layers you would have typical gradient accumulation (although i think you need to back propagate with retains_grad=True flag on first backward call).

So, Really the different is you would be calling optim.step() on the gradient propagated in each loss case; as opposed to aggregate of two losses.

FWIW, you can disable optimizer fusion by constraining your Sharder to only support EmbeddingComputeKernel.DENSE.

torchrec/torchrec/distributed/embedding_types.py

Line 456 in 1afbf08

EmbeddingComputeKernel.DENSE.value,

[WhoisZihan]

Thanks for the suggestion.
By "fused", I mean calling backward twice would result in two updates if the embedding, even if the retain_graph parameter is set to True.
Will setting kernel to DENSE make the embedding data-parallel? If I still want it to get sharded across ranks without backward fusion, is it possible?

[dstaay-fb]

No, Dense implies the optimizer is not fused, nothing about sharding choice. Choice of sharding configuration is independent. One caveat - Data Parallel embeddings do not support Dense mode (this hasn't been prioritized to date, given majority of models we work with have have hash sizes that require Model Parallel).