feat: add Megatron support for on-policy distillation

[zpqiu]

This pull request introduces support for Megatron-based training in the on-policy distillation pipeline. It adds new configuration files to enable Megatron parallelism and sequence packing, and removes previous restrictions against Megatron.

Issues

List issues that this PR closes (syntax):

Closes #1151

Usage

• You can potentially add a usage example below

uv run examples/run_distillation_math.py --config examples/configs/distillation_math_megatron.yaml

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 run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

Alignment Experiments between DTensor and Megatron Pathes

• Blue: DTenor path, baseline
• Green: Megatron, CP + Sequence Packing + TP
• Red: Megatron, CP + Sequence Packing + TP + PP

Summary by CodeRabbit

• New Features

  • Enabled Megatron backend for distillation (student and teacher) with scheduler steps auto-aligned to run length.
  • Implemented top‑k logits for Megatron policies with support for sequence packing, tensor/pipeline/context parallelism.

• Documentation

  • Removed outdated limitation indicating Megatron was unsupported.

• Examples

  • Added comprehensive Megatron distillation configs and a Qwen3 32B→1.7B recipe.

• Tests

  • Added functional and unit tests for Megatron distillation and top‑k; included a new nightly test entry and slightly increased the nightly GPU‑hour threshold.