DeepSpeed

ebf1531e - Add keyword argument support to activation checkpointing (#8182)

Add keyword argument support to activation checkpointing (#8182) ## Summary DeepSpeed's activation checkpointing wrapper currently accepts positional arguments only, while `torch.utils.checkpoint.checkpoint` also supports keyword arguments. This change allows callers to pass keyword arguments through `deepspeed.checkpointing.checkpoint`. Keyword names and non-Tensor values are retained for reconstruction during the forward and recompute passes. Tensor keyword values are flattened into the inputs passed to `CheckpointFunction`, so autograd tracks them and returns their gradients correctly. The activation checkpointing documentation now describes keyword argument support, and the regression test covers both Tensor and non-Tensor keyword arguments as well as gradient propagation. ## Validation - Activation checkpointing unit tests: 27 passed - Pre-commit checks for all changed files: passed - 1-GPU CUDA correctness smoke: direct and checkpointed execution matched, with zero maximum gradient error - 2-GPU distributed CUDA correctness smoke: each rank matched direct execution, with zero maximum gradient error - DCO sign-off is included in the commit Fixes #7038 Signed-off-by: Wang Zupeng <zupenwang@gmail.com>