Forward the per-head Muon tag at every muon_update call site

[vineethsaivs]

Per-head Muon (#8384) is silently ignored on three of the six paths that call muon_update, so the attention projections it tags are orthogonalized whole.

deepspeed.initialize tags each attention projection with muon_num_heads and the call sites read it back off the parameter. Three do not: the round-robin branch of ZeRO 1/2 get_flat_partition (its unpadded twin does), MuonWithAuxAdam.step, and the new BF16_Optimizer Muon support. Which branch of get_flat_partition runs depends on round_robin_gradients, so whether the feature takes effect turns on an unrelated setting, and the update stays finite either way. Forward num_heads at all three.

Test: TORCHDYNAMO_DISABLE=1 DS_ACCELERATOR=cpu PYTHONPATH=.:tests python -m pytest tests/unit/runtime/zero/test_per_head_muon.py -q. Two new cases fail before and 78 pass after: one steps MuonWithAuxAdam on a tagged projection and asserts the applied update is the per-head one rather than the whole-matrix one (0.32 apart here), and one names any call site that drops the tag. yapf 0.43.0 and flake8 pass on the changed files. Executed on Apple M2 Pro CPU; no GPU or distributed run.