[MoE] DeepEP refactor and fix memory leak during training and inference

[shuhuayu]

1. Simplified the token permutation logic.
2. Updated the handle management so there will be no memory leak during training and inference. Related issue: [DeepEP] How is _handle_cache handled during inference? #2273

On training on 16b deepseek v3 model, before the fix there was a growing memory usage.

After the fix, the memory usage stabilizes.

[tianyu-l]

@elfiegg please take a look

[elfiegg]

Looks good, is my understanding correct that this PR mainly implements:

1. use torch.argsort for sorting the tokens for performance reasons
2. save handle in dispatch's setup_context, and clear handle cache in combine forward if inference mode else clear it in setup_context

[elfiegg]

FYI @goldhuang to unblock your work

[shuhuayu]

Looks good, is my understanding correct that this PR mainly implements:

1. use torch.argsort for sorting the tokens for performance reasons
2. save handle in dispatch's setup_context, and clear handle cache in combine forward if inference mode else clear it in setup_context

Yes. During training, handles are saved in both dispatch_ctx and combine_ctx, and the handle in _handle_cache is cleared in combine_setup_context. During inference, no handles saved in op_ctx, and _handle_cache is cleared during combine_forward.

[elfiegg]

Sounds good, logic makes totally sense to me. Good to know inference mode setup_context won't be called

[goldhuang]

@shuhuayu With the changes in the PR, can I train without selective activation checkpointing on deepep? The extra memory cost from SAC is huge. It may not be a good idea to enable SAC in some cases.

[elfiegg]

@goldhuang You can configure --activation_checkpoint.mode=full and --parallelism.expert_parallel_comm_backend=deepep to turn on full checkpointing with deepep being all-to-all backend

[goldhuang]

@shuhuayu @elfiegg My point is that deepep integration in current main branch also has memory leak when you use it without SAC (meaning it's running in the pattern of fwd-fwd-bwd). You may want to make sure your changes also cover the fwd-fwd-bwd case.
The current main branch can only do fwd-bwd without a leak. Both fwd and fwd-fwd-bwd will cause a leak. I only reported the fwd case earlier.

[shuhuayu]

@shuhuayu With the changes in the PR, can I train without selective activation checkpointing on deepep? The extra memory cost from SAC is huge. It may not be a good idea to enable SAC in some cases.

@goldhuang Thanks for the question.

1. You can use full AC to save memory.
2. Currently, during training, dispatch_op.ctx and combine_op.ctx save layout medata, i.e., the handle regardless AC configuration, which should be not big.
3. If we use selective and op in AC, deepep ops are included in the op_sac_save_list, so their activations get saved. We can remove them from op_sac_save_list by commenting out these two lines:

   torchtitan/torchtitan/models/llama4/infra/parallelize.py

   Lines 132 to 133 in cee9482

   	_op_sac_save_list.add(torch.ops.deepep.dispatch.default)
   	_op_sac_save_list.add(torch.ops.deepep.combine.default)

I think 1 saves most memory and 3 saves memory specifically from deepep communications.

In my test of a deepseek 16b model on 16 h100s (seqlen=4096, bsz=4, fsdp=ep=8, pp=2, attention=sdpa, compile=loss, moe_communication=deepep, no moe force load balance), the results are:

1. selective ac with op: mfu 11.46%, memory 27.5%.
2. selective ac with op, but recompute deepep ops (bullet 3 above): mfu 11.5%, memory 27.38%
3. full ac: mfu 12.8%, memory 24.5%.

So in these tests the savings from excluding deepep ops from sac save list are not significant in small scale.

[shuhuayu]

@shuhuayu @elfiegg My point is that deepep integration in current main branch also has memory leak when you use it without SAC (meaning it's running in the pattern of fwd-fwd-bwd). You may want to make sure your changes also cover the fwd-fwd-bwd case. The current main branch can only do fwd-bwd without a leak. Both fwd and fwd-fwd-bwd will cause a leak. I only reported the fwd case earlier.

@goldhuang Thanks for pointing this out.

1. fwd case should be fixed if you use torch.inference_mode().
2. fwd-fwd-bwd case like in full ac is now covered in this pr and tested by running full ac test.