Add FSDP-aware hybrid Muon optimizer

[maocheng23]

Motivation

SpecForge's unified trainer currently exposes only AdamW. DSpark ablations also
need an opt-in Muon setup without moving embedding/head parameters to Muon or
replicating optimizer state under FSDP.

Modifications

• Add typed training.optimizer: adamw | muon configuration plus independent
  Muon/AdamW learning-rate, weight-decay, momentum, Nesterov, and Newton--Schulz
  controls. AdamW remains the default.
• Partition hidden nn.Linear.weight matrices into Muon while keeping embedding
  projections, Markov/confidence/output heads, norms, and biases on auxiliary
  AdamW. Capture the classification and logical matrix shapes before FSDP wraps
  the draft.
• Use torch.optim.Muon for replicated parameters. For FSDP local shards, keep
  FP32 masters and momentum sharded, gather one BF16 update at a time for the
  same PyTorch 2.11 Newton--Schulz transform, and write back only the local
  slice.
• Preserve the existing AdamW checkpoint schema. Add a versioned hybrid state
  format with strict parameter-partition validation and independent scheduler
  state for both optimizer groups.
• Log lr_muon and lr_adamw, validate incompatible CPU offload settings, and
  document the new training configuration.

Related Issues

None.

Accuracy Test

The unit suite compares a three-rank sharded update against native full-matrix
torch.optim.Muon using uneven (7, 5, 0) shards, including an empty
rank-local parameter. It also checks the concrete DSpark backbone/head
partition and hybrid checkpoint continuation.

Validation with the repository's pinned PyTorch 2.11 environment:

• pytest -q tests/test_optimizer: 17 passed, 1 skipped
• pytest -q tests/test_config: 55 passed, 431 subtests passed
• CPU runtime suite excluding the two modules that require a local SGLang
  installation: 542 passed, 20 skipped, 1 expected failure, 68 subtests passed

No end-to-end model accuracy run has been performed yet.

Benchmark & Profiling

Not yet measured. The implementation keeps persistent optimizer state sharded,
but Newton--Schulz adds one transient BF16 logical-matrix gather per Muon
parameter. An H200/NCCL smoke run and step-time/memory profile are still needed
before marking this PR ready for review.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.