fix: refresh TE fused MLP DDP hooks dynamically

[nvegesna-netizen]

What does this PR do?

Refreshes the Transformer Engine fused MLP's forwarded descendant pre-hooks at
call time so Megatron distributed-data-parallel parameter-gather hooks continue
to run when DDP enables or replaces them after the fused implementation is
constructed.

TEFusedMLP lazily creates its fused implementation during the first forward.
The previous implementation captured descendant hook objects only once at that
point. Distributed optimizer training can temporarily disable its parameter
gather hooks for the first iteration and enable them afterward, leaving the
fused implementation with an empty or stale snapshot. Subsequent forwards then
skip the current DDP hooks.

This change:

• caches only the descendant module identities whose calls are skipped by the
  fused implementation;
• reads their current pre-hook registries on every fused invocation;
• avoids replaying hooks registered on TEFusedMLP itself, since normal
  torch.nn.Module.__call__ already executes those hooks;
• retains the existing rejection of descendant backward hooks and
  input-modifying pre-hooks; and
• adds regression coverage for hooks added, removed, and replaced after fused
  implementation construction.

Validation

• Focused hook regression on 8 GB200 GPUs: both tests passed on every rank.
• Distributed-optimizer parameter-shard ownership test with two independent
  4-rank jobs: both parameterizations passed on both nodes.
• 50-step Llama 3 8B NVFP4 training on 8 GB200 GPUs with the Transformer Engine
  op fuser: completed successfully.
• Repository-pinned Black 24.4.2, isort 5.13.2, and Ruff 0.9.10 checks passed for
  both changed files.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.