fix(dflash): apply activation checkpointing to the draft

[kashif]

What does this PR do ?

DFlash never wired up distributed.activation_checkpointing for the draft model, so the flag was silently ignored. DSpark already had this. Moved DSpark's helper into the shared _spec_train_utils and call it from DFlash too.

Changelog

• Add apply_draft_activation_checkpointing to _spec_train_utils.py, shared by DSpark and DFlash.
• Wire it into TrainDFlashRecipe.setup() before the draft's DDP wrap (covers DFlash2, Domino, and JetSpec via inheritance).
• Move DSpark's activation checkpointing tests next to the shared helper in test_spec_train_utils.py.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Additional Information

[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.

[HuiyingLi]

/claude review

[claude]

LGTM