Support reentrant MFSDP activation recomputation

[wujingyue]

Summary

• remember when an FSDP module runs the grad-disabled original forward used by reentrant checkpointing
• recognize the later grad-enabled invocation as activation recomputation
• suppress forward-order prefetch and keep recomputed parameters materialized until backward consumes them

Why

Reentrant checkpointing can recompute a child module before that child's backward-pre hook runs. The context phase alone therefore does not reliably identify every recomputed forward. A sticky module-local marker records the static checkpoint usage observed during the original forward.

This PR is stacked on #6153 and targets its pull-request/6153 diff base. It uses the existing reentrant and non-reentrant activation-recompute test from that PR.

Validation

• python3 -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py::test_fully_shard_activation_recompute_reshards_parameters
• python3 -m ruff check megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py
• python3 -m black --check megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py
• git diff --check

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.