Adopt composable activation checkpointing once compile + DDP ready

Summary

Track migration from module-replacing CheckpointWrapper activation checkpointing to PyTorch's type-preserving composable activation-checkpointing API:

from torch.distributed._composable import checkpoint

checkpoint(module)

The composable API installs checkpointing through forward hooks, preserving module identity, concrete type, model structure, and fully-qualified parameter names.

Motivation

PR #2955 fixes AM-370, where replacing block.mlp with CheckpointWrapper broke forward-time isinstance(self.mlp, MLP | MoE) dispatch in GPT-OSS and Qwen3-MoE.

Using the composable API centrally would avoid changing the apparent module type and eliminate this class of wrapper-aware model logic. PyTorch already exercises composable activation checkpointing together with FSDP2, but its compiled-FSDP parity test currently skips the composable variant.

Readiness gates

Adopt the composable API after all of the following are true:

• PyTorch exposes the API as stable/public, or AutoModel explicitly accepts depending on the private torch.distributed._composable namespace.
• Composable activation checkpointing is supported and covered with torch.compile + FSDP2; the current upstream compile skip is removed.
• The AutoModel DDP activation-checkpointing path is validated with the composable API.
• Full and selective activation checkpointing remain correct with context_fn, RNG preservation, and per-layer compile.
• Context-parallel attention hooks fire correctly during both forward and recomputation.
• KV-sharing fallback behavior remains correct.
• State-dict save/load, DCP, SafeTensors export, and FQN normalization remain compatible.
• Pipeline-parallel and expert-parallel configurations are covered.

Proposed migration

1. Migrate the ordinary non-compiled submodule path in apply_submodule_checkpointing().
2. Migrate the DDP submodule-checkpointing path after parity coverage exists.
3. Migrate compile/selective-AC paths only after upstream compile support is established.
4. Remove model-specific _checkpoint_wrapped_module inspection once no supported path replaces those submodules with CheckpointWrapper.

References

• PyTorch composable API: https://github.com/pytorch/pytorch/blob/23b1588b32ea2d46410793347d635017fdc49c75/torch/distributed/_composable/checkpoint_activation.py#L38-L46
• PyTorch FSDP2 parity coverage and current compiled-FSDP skip: https://github.com/pytorch/pytorch/blob/23b1588b32ea2d46410793347d635017fdc49c75/test/distributed/_composable/fsdp/test_fully_shard_training.py#L790-L855
• AutoModel PR #2955: https://github.com/NVIDIA-NeMo/Automodel/pull/2955
• Linear AM-370: https://linear.app/nvidia/issue/AM-370/nemo-bench-gpt-oss-qwen3-moe-pretrain-fsdp2-wrapped-selfmlp-fails

Start with apply_submodule_checkpointing() and AutoModel PR #2955, then inspect the referenced PyTorch FSDP2 parity test and its compiled-FSDP skip. Trace the ordinary, DDP, compile, and selective-checkpointing paths against the listed readiness gates. Done means the supported paths use the composable API, compatibility coverage passes, and obsolete CheckpointWrapper inspection can be removed.