Add MoE expert parallel training runtime

[alay2shah]

Summary

This PR adds first-class MoE training support with an expert-parallel runtime, manual sharded checkpoint export, and a cleaned-up package structure rebased onto the current GRPO mainline.

Major changes:

• Add MoE SFT/DPO training paths with EP token dispatch, grouped expert compute, FSDP2 wrapping over the DP mesh, activation checkpointing, MoE metrics, and memory tracing.
• Add manual sharded checkpoint save/export support that can reconstruct HF-compatible checkpoints from FSDP2/EP runs.
• Reorganize runtime modules into clearer packages: data_loading, training, checkpointing, distribution, config, cli, and rl.
• Keep GRPO helpers in the reorganized layout: training loops under training/, rollout/reward/env helpers under rl/, and user reward examples under repo-root rewards/.
• Add Ray/Slurm distribution helpers, multinode Ray startup/teardown, and rank/data sharding protections.
• Include one standard MoE EP SFT example config alongside the other committed examples.
• Expand regression coverage for config parsing, data loading, MoE losses/runtime, rank groups, checkpoint metadata/export, GRPO integration, and trainer lifecycle behavior.

Validation

• uv run ruff check src tests
• uv run ruff format --check src tests
• uv run pre-commit run --all-files
• uv run pytest tests/test_config_parsing.py tests/test_grpo_config.py tests/test_grpo_data.py tests/test_rewards_loader.py tests/test_vllm_server.py tests/test_rl_env_adapter.py tests/test_load_models.py tests/test_ray_cluster_support.py tests/test_training_configs.py tests/test_moe_ep_runtime.py tests/test_moe_losses.py tests/test_moe_trainers.py tests/test_sft_loss_masks.py tests/test_trainer_mixins.py -> 323 passed, 5 warnings in 52.33s

Notes

The GPU e2e tests remain marked/skipped in local CPU-style test runs. Recent Slurm validation runs covered the DP4/EP2 and DP4/EP4 save paths before this final cleanup pass.

[alay2shah]

Tested on DP4, EP2 packed 32k sequences on 1 node and multinode DP4, EP4 packed 32k sequences

[alay2shah]

Currently rebasing on new GRPO main

[EdoardoMosca]

Apologies for the delay, partially my fault that now you have merging conflict with main

The branch feels very clean and nicely structured

🟢 formatting and linting are good
🟢 ran non-GPU tests

given the PR size I tackled it with codex but everything it found are things that don't worry me

nit: manual_sharded_checkpoint_format was being removed from the Trainer/DPOConfig kwargs before validation ran, so invalid values were not caught. I applied a small fix. You can revert if you feel I'm being overconcerned