Feature: nemotron_omni decode is bound by 2944 conditional expert subgraphs per token (about 50 ms/token for a 3B-active model)

Đang mở
  • Có gì không ổn?
#2,060 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 1 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
5/5
Thời gian dự kiến
Hơn một tuần
Mức phù hợp với người mới
35/100
Loại issue
Tính năng
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, pytorch
Lĩnh vực
machine-learning

Hướng nghiên cứu

Start with families/nemotron_omni/moe.py and inspect how each routed expert is built as an IIfConditional subgraph. Compare this with another family’s gathered MoE implementation, then investigate whether the selected experts’ weights can be run as a grouped or batched GEMM. Done means reducing per-token decode time and decoder plan and build size; the issue gives B200/B300 decode measurements for comparison.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

enhancement

Request type: Improvement to existing behavior

Problem and use case

families/nemotron_omni builds each MoE layer of the Nemotron-H language model as 128 IIfConditional expert subgraphs (moe.py, one per routed expert; top-6 run). With 23 MoE layers that is 2944 conditional expert subgraphs per decode step, and the decoder plan has about 25k to 64k layers.

On the official NVFP4 checkpoint (the manifests' nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-NVFP4), decode takes about 48 to 53 ms per token on B200 / B300 (about 20 tokens/s), for a model with about 3B active parameters (reading 3B parameters at about 8 TB/s would take well under 1 ms). Prefill is token by token through the same step, so it has the same cost per token.

The step time does not depend on the weight precision, which shows the step is bound by the conditional structure and kernel count, not by weight reads (decoder only, TensorRT 11.4, batch 1, ms per token at cache positions 100 / 2000):

GPU fp16 (from the BF16 checkpoint, built with the family decoder code) FP8 (official FP8 checkpoint) NVFP4 (as on main)
B200 47.7 / 47.8 48.3 / 48.4 49.7 / 49.8
B300 48.9 / 48.8 48.9 / 48.7 49.7 / 49.6

Decoder build time is also large: 24 to 52 minutes on B200 / B300 and about 101 minutes on H100 for NVFP4.

Proposed outcome

Build the routed experts of each MoE layer without one conditional subgraph per expert, for example a gathered/batched expert formulation (gather the selected experts' weights and run them as one grouped or batched GEMM, as other families do with a gathered MoE), so the decode step time scales with the active weights and the plan and build time shrink.

Model and target details

nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-NVFP4 @ 16993199e436da4ba75ddc410855f87e0d996ee6 (the nemotron-omni-nvfp4-native manifest). Language model: 52 layers (23 Mamba2, 23 MoE with 128 routed experts top-6 plus 1 shared, 6 attention), hidden 2688. TensorRT 11.4; B200 / B300 (also seen on H100).

Alternatives considered

FP8 weights (official FP8 checkpoint) and fp16 weights give the same step time as NVFP4 (table above), so weight precision does not help while the MoE stays conditional.

Additional context

Found during an FP8 study of this family (not merged). Separate note for anyone building a Hugging Face reference for Nemotron-H: without mamba-ssm, the checkpoint's own PyTorch fallback in modeling_nemotron_h.py diverges after the first 128-token chunk (KL about 1e-8 on positions 0 to 127, 0.07 to 0.39 after); a plain sequential scan matches the fp16 engine (KL 0.0032).

Ngôn ngữ chính
Python
Star
288
Fork
73
Merge trung bình
1 ngày 3 giờ
Pull request đã merge (30 ngày)
273

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của NVIDIA/TensorRT-Model-Connect

Tất cả issue của NVIDIA/TensorRT-Model-Connect

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.