Hit a bug trying to run Step 3.7 Flash NVFP4 with MTP on a 2-node Spark cluster (TP=2).
The MTP draft loader crashes when it tries to copy the full 4096-dim vocab weight into the TP-sharded 2048-dim slot. Basically the loader doesn’t know the embedding was split across ranks.
File ".../models/step3p5_mtp.py", line 273, in load_weights
RuntimeError: The size of tensor a (2048) must match the size of tensor b (4096) at non-singleton dimension 1
Everything else loads fine — NVFP4 MoE on CUTLASS, FlashInfer attention, NCCL multi-rail RoCE all good. No-MTP serving works great SMH on the same setup (14 tok/s SMH, correct answers). It’s just the MTP weight loading that breaks at TP=2.
StepFun tested at TP=8 so they probably never saw this. For those of us on 2 Sparks, TP=2 is the only option and MTP is dead until this gets patched.
Anyone run into this or know a workaround?