Fix expert parallelism through Trainer
Loading a model with DistributedConfig(tp_size=N, enable_expert_parallel=True)
and handing it to Trainer fails on main in three places:
1. maybe_distribute_model never assigns model._tp_size, so the Trainer builds
no ParallelismConfig and accelerate wraps the DTensor model in DDP:
ValueError: Your model contains DTensor parameters, which is incompatible
with DDP.
2. _get_grad_norm calls clip_grad_norm_ over the full parameter set, and
_foreach_norm cannot span a mix of DTensor (experts) and plain parameters.
3. The fused/foreach AdamW kernels cannot span that mix either.
Set _tp_size where the mesh is recorded, compute the gradient norm (and clip)
per-gradient with replication-aware discounting when parameters live on
different meshes, and fall back to per-parameter AdamW stepping for mixed
parameter sets.
With this and #48205, expert-parallel full fine-tuning through Trainer runs
end-to-end and tracks a single-GPU control step by step (OLMoE-1B-7B, tp=4:
12.13 -> 12.45 -> 11.52 -> ... -> 11.04 vs 12.13 -> 12.45 -> 11.51 -> ... ->
11.09).
