Add keyword argument support to activation checkpointing

[zupengwang]

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