Repository navigation
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
svcnvidia-nemo-ci
left a comment
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
Reviewed all 10 changed files and their directly affected contracts. Two test-coverage gaps remain: the new MoELayer execution branch is only tested at construction, and the changed external GPU dispatch lacks the required bit-exact replay coverage.
No concrete runtime logic defect was established in the inspected source. Tests were not executed; the companion MoK implementation is not available in the mounted repository.
Note: formal approval is not available for this request, so this review is posted as a comment only.
Review generated by codex. This was a full review. Request a full review with /review reset=true.
| return apply_module(self.megakernel_experts)(hidden_states, probs, routing_map) | ||
| if self.megakernel_fuses_shared_expert: | ||
| return apply_module(self.megakernel_experts)(hidden_states, probs, routing_map) | ||
| hidden_states, probs = self.preprocess(hidden_states, probs, routing_map) |
There was a problem hiding this comment.
Exercise the native-shared/latent branch through forward and backward. The added MoELayer test stops after construction. The adapter tests invoke MoKMegakernel separately. Consequently, omitting a latent projection or the shared-output addition from this new branch would still pass the added tests; existing latent/recompute tests do not enable MoK.
Correction: add a regression around MoELayer.forward. A differentiable three-argument routed-backend stub is sufficient to test integration without the companion package. For model width 128 and latent width 64, verify that routing/shared computation receive width 128, the backend receives width 64, and the result equals fc2_latent_proj(backend(fc1_latent_proj(x), ...)) + shared_experts(x). Fail if any native dispatcher method is called. Backward must reach both projections, the shared expert and the input.
Preservation / acceptance (proposed, not executed): cover no shared expert, forced-native shared without latent projections, and compatible fused shared execution, which must not compute shared experts twice. Include TP2/SP with ETP1 and rank-distinct token shards: outputs retain each rank’s local input shape, and finalized replicated-projection gradients match a reference over the combined shards. Include recompute on/off with matching outputs and gradients. These checks establish the new composition and TP behavior rather than merely successful construction.
| gate_kwargs["shared_output_gate_weight"] = shared_output_gate_weight | ||
| forward_kwargs["shared_output_gate_weight"] = shared_output_gate_weight | ||
| if module.situ_beta is not None: | ||
| forward_kwargs["situ_beta"] = module.situ_beta |
There was a problem hiding this comment.
Add real-kernel bit-exact replay coverage for the new MoK modes. Passing situ_beta and allowing absent shared weights changes external GPU dispatch, but the new contract cases replace forward with an input clone and backward with synthetic gradients. Numerical or replay failures in the new modes therefore remain undetectable. No kernel replay test or manifest is present under the head tree’s determinism test directory. The selected light-review rubric explicitly requires a bit-exact test and manifest registration for changed external-library dispatch.
Correction: add a real-MoK replay test under the required kernel-test directory and register this dispatch in its manifest. The invariant is byte-identical output and gradients for identical executions—not just correct keyword forwarding. Exact test implementation depends on the companion package and supported hardware, which are unavailable here.
Preservation / acceptance (proposed, not executed): use fixed nonzero inputs, weights, probabilities and routes; reset gradient buffers before at least three forward/backward replays; compare raw bytes of output, input/probability gradients and every routed parameter’s main_grad. Cover routed-only BF16 and MXFP8 SiTU with nondefault betas, plus a fused-shared SwiGLU control retaining shared FC1/FC2 and optional output-gate gradients. Retain the existing mocked assertions for absent shared weights/gradients, gate/up aliases and DDP accumulation. Make dependency/hardware skips explicit rather than replacing numerical execution with stand-ins.
Signed-off-by: qizhang <qizhang@nvidia.com>
Keep shared experts and latent projections in MCore when fusion is structurally incompatible or explicitly disabled. Support TP/SP with ETP=1, centralize the backend shared-execution option, reject unsupported expert bias, and preserve native dispatcher behavior. Include the corresponding configuration and adapter regressions. Signed-off-by: qizhang <qizhang@nvidia.com>
e975047 to
48eb7db
Compare
48eb7db to
367b1a6
Compare
Add an independent MoELayer composition oracle and finalized TP/SP projection-gradient reference. Cover real BF16/MXFP8 routed-only SiTU and legacy fused shared execution, including router-probability and main-gradient byte replay. Register external MOK dispatch and preserve frozen cuDNN flags during fixture restoration. Signed-off-by: qizhang <qizhang@nvidia.com>
367b1a6 to
8dc4a6e
Compare
…#8497 Bring the Mixture-of-Kittens (MoK) megakernel integration for Kimi-K3 onto the dev branch, as submitted upstream: - support the SiTU-GLU activation in MOK independently of the shared-expert activation (situ_beta is passed to the backend from the config); - keep the native shared-expert and MoE latent-projection paths around a routed-only MOK: the shared execution policy is the ``shared_expert_mode`` entry of ``moe_megakernel_backend_config`` (``auto`` fuses compatible shared MLPs, ``fused`` requires compatibility, ``native`` runs shared experts in MCore); MCore splits it off before the remaining options reach the backend; - enable routed-only MOK with tensor-parallel sequence parallelism (expert TP stays 1; TP>1 requires sequence_parallel; routed weights use the latent size); - reject expert bias with MOK (add_bias_linear must be False). MoK stays opt-in through moe_megakernel_backend; recipes are unchanged. Upstream PR: NVIDIA#8497 Squashed-from: 0e3f344 (PR head; commits 5247d0b, 07b9fe6, bd8e991, 3a3befe, 0e3f344) Rebase note (dev ff4fbc9): the PR is based on an older dev where NVIDIA#6615 had not yet routed the padding mask into MoEFlexTokenDispatcher.dispatch_preprocess; dev's padding-mask branch is kept and nested inside the ``if self.config.moe_megakernel_backend is None:`` guard of MoELayer.preprocess, and dev's moe_use_norm_before_up_proj handling stays in postprocess. All other files are taken from the PR verbatim. This commit replaces the earlier dev_dev squash of the same work, which exposed the shared-expert policy as a separate moe_megakernel_shared_expert_mode config field; recipes now pass ``--moe-megakernel-backend-config '{"shared_expert_mode": "native"}'`` instead. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>
Summary
Validation
Requires the companion MoK routed-only and SiTU-GLU changes.