Fixed spmd typecheck for qwen35 and kimi27

[jinsooihm]

There were some changes in #4121 that made the spmd typecheck fail.
The integration tests do not test for spmd_typecheck with TP > 1 for qwen35 and kimi27, so this was not caught.

Test plan

• Qwen3.5 multimodal, DP=4/TP=2, spmd_types strict typechecking: 1 training step passed.
• Kimi K2.5 multimodal, DP=2/TP=2, spmd_types strict typechecking: 1 training step passed using a test-only AdamW configuration with the TP-incompatible QK-clip hook disabled.
• scripts/loss_compare.py: compared parent 4ab7596c0 against c302e9fb1 on Qwen3.5 TP=2 for 10 deterministic steps. Full-precision loss and grad_norm were bitwise identical at every step.
• Focused ufmt, flake8, pydoclint, and git diff --check passed.

[tianyu-l]

The integration tests do not test for spmd_typecheck with TP > 1

should add TP