Skip to content

feat(refit): support BF16 to MXFP8 NCCL reshard - #3477

Open
seonjinn wants to merge 4 commits into
NVIDIA-NeMo:mainfrom
seonjinn:sna/nccl-reshard-bf16-mxfp8
Open

feat(refit): support BF16 to MXFP8 NCCL reshard#3477
seonjinn wants to merge 4 commits into
NVIDIA-NeMo:mainfrom
seonjinn:sna/nccl-reshard-bf16-mxfp8

Conversation

@seonjinn

@seonjinn seonjinn commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

This PR extends NCCL Reshard refit to support BF16 trainer storage with an MXFP8 vLLM rollout engine.

  • Receive each resharded FFN/MoE shard into a BF16 destination buffer.
  • Quantize the local shard on the vLLM worker and populate its E4M3 value and E8M0 scale tensors.
  • Restore grouped-expert tensor shapes after FlashInfer CuTeDSL quantization, which internally flattens inputs with more than two dimensions.
  • Reuse the existing vLLM load path for non-bulk parameters.
  • Preserve the existing matching BF16 and matching blockwise-FP8 paths.
  • Reject blockwise-FP8 trainer storage with an MXFP8 destination because their scale formats are not interchangeable.

This PR does not add trainer-side prequantization or alter the MoE layout transform. The independent batched-shuffle optimization is in #3478.

Why?

NCCL Reshard previously required matching trainer and rollout storage. A BF16-training/MXFP8-rollout recipe therefore failed validation even though the destination worker can quantize its already-resharded local shard without changing the collective protocol.

Correctness

  • Adds precision-matrix tests for supported and rejected storage pairs.
  • Adds direct and merged-MoE mapping tests for BF16 receive buffers, E4M3 values, and E8M0 scales.
  • Tests grouped-expert shape restoration against the flattening behavior of the FlashInfer CuTeDSL quantizer.
  • Validates MXFP8 scale shape and dtype before refit starts.
  • Keeps the quantization helper shared with the existing vLLM refit load path, including zero-scale sanitization.

@seonjinn
seonjinn requested review from a team as code owners August 3, 2026 19:49
@copy-pr-bot

copy-pr-bot Bot commented Aug 3, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@seonjinn seonjinn added the CI:L1 Run doctests, unit tests, and functional tests label Aug 3, 2026
@seonjinn

seonjinn commented Aug 3, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 83010d8

Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
seonjinn force-pushed the sna/nccl-reshard-bf16-mxfp8 branch from 83010d8 to f3dcd2a Compare August 3, 2026 19:56
@seonjinn

seonjinn commented Aug 3, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test f3dcd2a

Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn

seonjinn commented Aug 3, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 4d1522b

@NVIDIA-NeMo NVIDIA-NeMo deleted a comment from copy-pr-bot Bot Aug 4, 2026
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn

seonjinn commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test d63ed16

@seonjinn

seonjinn commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 6f57c1b799bf980e6a99bc9dc530ce190029e7df

@copy-pr-bot

copy-pr-bot Bot commented Aug 5, 2026

Copy link
Copy Markdown

/ok to test 6f57c1b799bf980e6a99bc9dc530ce190029e7df

@seonjinn, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

@seonjinn

seonjinn commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 6f57c1b

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:L1 Run doctests, unit tests, and functional tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant