[QST] Control over gradient communication (AllReduce) and optimizer steps for decoupled DynamicEmb

[jlinswenmogu]

What is your question?

Hi team,

inside the readme.md future plan section, i saw that there is a plan to separate gradient computation (the backward step) and the update step (the optimizer step). currently inside dynamicemb, they are fused together.

i would love to have this separation, because this will make gradient accumulation easier. currently when these two steps are fused at loss.backward, the emb get updates for every mini-batch whereas the nn module weight gradients accumulate for a few mini-batches before update at the optimizer.step().

i have the following question:

i guess the gradient computation would be naturally triggered by loss.backward. however, under the distributed training, regarding emb gradient allreduce, will the team make it in a way that the user would have control over when to trigger emb gradient allreduce and the update step? for example, under megatron DDP mode, user has the autonomy over when to call the megatron.finalize_model_grad, and this is separate from the optimizer.step().

thanks!

---

By submitting this issue, you agree to follow our code of conduct and our contributing guidelines.

[shijieliu]

hi @jlinswenmogu thanks for your interests for the recsys-examples!

First, I’d like to clarify that in DynamicEmb’s backward pass we use all-to-all, not all-reduce, to communicate gradients, because the embedding weights are model-parallel. After each GPU receives its own gradients, it aggregates them by key and then updates its local shard of the table.

If I understand correctly, what you want is to use micro-batches, accumulate gradients over multiple batches, and then perform a single optimizer step. Is it? In that case, are you looking for a model-parallel embedding table, or do you actually just need a data-parallel embedding table?

[jlinswenmogu]

hi @shijieliu , thanks for the reply!

on the all-to-all: yup thats what i meant, sorry for the abuse of terms

on the "use micro-batches, accumulate gradients over multiple batches, then perform a single optimizer step": yes, thats exactly what i meant.

on "are you looking for a model-parallel embedding table, or do you actually just need a data-parallel embedding table": since dynamicemb is model-parallel (which is exactly the point that makes it stand out), i am looking at this direction.

just out of interest: when you mentioned "data-parallel embedding table": do you mean the "each gpu gets the complete copy of the embedding table" approach or the industrial norm parameter server approach?

[jiashuy]

Hi @jlinswenmogu , we think we can support this feature; we estimate it will take a week to develop, but the start date hasn't been determined yet due to prioritization.