fix(moe): Kimi-K3 correctness fixes (routing, NCCL warmup, PP, per-GPU tps) (#3731)

* fix(kimi-k3): honor backend.fake_balanced_gate in KimiK3MoE

KimiK3MoE.__init__ unconditionally built KimiK3Gate, silently ignoring
backend.fake_balanced_gate, which the base MoE class honors. With
random- or zero-initialized weights the learned gate produces near-equal
scores, so topk returns experts [0..topk) for every token and all
traffic collapses onto each EP group's first rank; benchmark configs
that requested forced load balance were degenerate. In a 64xGB200
Kimi-K3 run this routing collapse alone cost ~6x throughput
(34.4 -> 248.5 tok/s/GPU once fixed, wall-clock verified).

Build FakeBalancedGate(moe_config, noise=backend.fake_gate_noise) when
the flag is set, exactly as the base MoE does.

Signed-off-by: Yisong Li <yisongbetter@gmail.com>

* fix(moe): tolerate never-forwarded FSDP modules in MoE PP post-backward

With FusedLinearCrossEntropy under pipeline parallelism the lm_head
weight is consumed inside the loss function, so its FSDP module never
runs forward and its FSDPCommContext is never lazily initialized.
_run_post_backward_hooks then crashes the last pipeline stage with
AttributeError: 'FSDPCommContext' object has no attribute
'post_forward_order'.

Guard the per-state post_backward() calls and the root final callback
against this AttributeError. A never-forwarded FSDP group has no
gradients to reduce, so skipping it is safe; verified on a 2-stage
Kimi-K3 pipeline probe by comparing per-step losses against a non-fused
control run (no gradient desynchronization).

Signed-off-by: Yisong Li <yisongbetter@gmail.com>

* fix(train_ft): report per-GPU tokens/sec using world size under pipeline parallelism

tps is the global tokens/sec of the step (num_tokens_in_batch is summed
over the DP group), but tps_per_gpu divided it by dp_size * cp_size
only. Ranks belonging to the same DP replica across the pp (and tp)
dimensions were not counted, so the printed per-GPU throughput -- and
any MFU derived from it -- was inflated by the pp*tp factor.

Measured on a 256-GPU Kimi-K3 run with pp8 (dp=32): the log line
reported 275.2 tok/s/GPU while wall-clock-verified global throughput
divided by 256 GPUs was 34.4 tok/s/GPU -- exactly 8x, the pp factor.

Divide by the recipe's world size instead, matching the divisor the MFU
computation already uses.

Signed-off-by: Yisong Li <yisongbetter@gmail.com>

* fix(moe): narrow the never-forwarded FSDP guard to the lazy-init signature

Address review feedback: a broad 'except AttributeError' could silently
swallow unrelated failures inside post_backward() and drop a real group's
gradient reduction. Match the never-forwarded signature explicitly
(missing 'post_forward_order'), log the skip at debug level, and re-raise
anything else. Adds a regression test asserting unrelated AttributeErrors
propagate.

Signed-off-by: Yisong Li <yisongbetter@gmail.com>

* fix(kimi-k3): enable comm-group prewarm in the K3 recipe

prewarm.comm_groups exists exactly to prevent the NCCL 'cuda failure 2
(out of memory)' that fires when the grad-norm communicator initializes
lazily at the step-0 memory peak (its docstring quotes this failure), but
it is opt-in and the K3 recipe never enables it -- so recipe users still
hit the crash at memory-tight settings (reproduced at 2.78T, pp8,
pp_microbatch_size=2: all last-stage ranks die identically at the end of
step 0). Enable it where it matters.

Signed-off-by: Yisong Li <yisongbetter@gmail.com>

* fix(recipes): keep the ci section last in the K3 SFT config

The nightly-recipe validator requires the top-level ci section to be
the last section; the prewarm section added earlier landed after it.
Move prewarm above ci.

Signed-off-by: Yisong Li <yisongbetter@gmail.com>

---------

Signed-off-by: Yisong Li <yisongbetter@gmail.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>