Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.rst
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ Changelog
- Add mergeable LoRA QAD for fake-quantized dense Megatron/Transformer Engine linear layers. Enable it with ``--lora_rank`` in Megatron-Bridge PTQ, train through the normal QAD flow, and export a standalone quantized checkpoint without deployment adapters.
- Add optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver.
- Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/>`_ for details.
- Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for NVIDIA Nemotron 3.5 Super VL, covering chat-data blend preparation, assistant-only distillation, Hugging Face export and vLLM serving. See `examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/>`_ for details.
- Add ``--mlflow <tracking-uri>`` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed.
- Add ``--ep_size`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` so large MoE models with grouped-GEMM experts can be exported with their experts sharded across GPUs. Checkpoints built with ``--no_moe_grouped_gemm`` must still be exported at ``--ep_size 1``.
- Add unified HF export of GLM-5 / GLM-5.2 (``glm_moe_dsa``) checkpoints quantized with ``examples/megatron_bridge``. The MTP layer, which Megatron-Bridge does not build for these models, is copied through unquantized from the source checkpoint. DSA sparse-attention models also gain FP8 KV-cache quantization on Megatron.
Expand Down
2 changes: 1 addition & 1 deletion examples/megatron_bridge/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ This directory contains examples of using Model Optimizer with the [NeMo Megatro
>
> Or the [Qwen3.6-35B-A3B W4A4 NVFP4 + QAD tutorial](tutorials/Qwen3.6-35B-A3B/README.md) for an end-to-end quantization-aware distillation workflow.
>
> Or the [NemotronH omni VL W4A4 + QAD tutorial](tutorials/NemotronH-Omni-W4A4-QAD/README.md) for a Mamba-Transformer MoE vision-language model.
> Or the [NVIDIA Nemotron 3.5 Super VL W4A4 + QAD tutorial](tutorials/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/README.md) for a Mamba-Transformer MoE vision-language model.

## Pre-Requisites

Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# NemotronH Omni VL: W4A4 NVFP4 PTQ + QAD + Deployment
# NVIDIA Nemotron 3.5 Super VL: W4A4 NVFP4 PTQ + QAD + Deployment

End-to-end W4A4 quantization of a NemotronH omni vision-language checkpoint (architecture `NemotronH_Omni_Reasoning_V3`), a Mamba-Transformer MoE language model with a vision tower: mixed NVFP4/FP8 post-training quantization → Quantization-Aware Distillation (QAD) → Hugging Face export → vLLM serving. This document covers:
End-to-end W4A4 quantization of [NVIDIA Nemotron 3.5 Super VL](https://huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16) (architecture `NemotronH_Omni_Reasoning_V3`), a Mamba-Transformer MoE language model with a vision tower: mixed NVFP4/FP8 post-training quantization → Quantization-Aware Distillation (QAD) → Hugging Face export → vLLM serving. This document covers:

1. **[Data Preparation](#1-data-preparation)**: the chat-format QAD blend, from `data_blend.yaml`
2. **[Quantization](#2-quantization)**: W4A4 PTQ with `examples/megatron_bridge/quantize.py`
Expand All @@ -10,6 +10,8 @@ End-to-end W4A4 quantization of a NemotronH omni vision-language checkpoint (arc

Only the language model (including its MTP head) is quantized; the vision tower and projector stay BF16 and are copied from the source checkpoint.

The PTQ recipe is the same as the one used for the released [`nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4) checkpoint. The calibration data and the QAD training data are all public Hugging Face datasets. The evaluation results below differ from the official model card scores.

## Results

All checkpoints were served with the same vLLM settings (see [Section 5](#5-evaluation-and-serving)); the quantized ones with an FP8 KV cache. BF16 is the mean of two runs.
Expand Down Expand Up @@ -72,14 +74,14 @@ PTQ runs on **one 4-GPU GB300 node** with 512 calibration samples.

```bash
torchrun --nproc_per_node 4 /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
--hf_model_name_or_path <nemotron_h_omni-checkpoint> \
--hf_model_name_or_path nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16 \
--trust_remote_code \
--recipe models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/ptq/nvfp4-mse \
--tp_size 1 --ep_size 4 --pp_size 1 \
--calib_dataset_name cnn_nemotron_v2_mix \
--calib_num_samples 512 \
--calib_batch_size 1 \
--export_megatron_path /path/to/nemotron_h_omni_w4a4_megatron \
--export_megatron_path /path/to/nemotron_3_5_super_vl_w4a4_megatron \
--skip_generate
```

Expand Down Expand Up @@ -108,9 +110,9 @@ We used **8 nodes × 4 GB300 (32 GPUs)** with context parallelism 4 and expert p
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) MASTER_PORT=29500
srun bash -c 'export RANK=$SLURM_PROCID WORLD_SIZE=$SLURM_NTASKS LOCAL_RANK=$SLURM_LOCALID
python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
--teacher_hf_path <nemotron_h_omni-checkpoint> \
--student_hf_path <nemotron_h_omni-checkpoint> \
--student_megatron_path /path/to/nemotron_h_omni_w4a4_megatron \
--teacher_hf_path nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16 \
--student_hf_path nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16 \
--student_megatron_path /path/to/nemotron_3_5_super_vl_w4a4_megatron \
--trust_remote_code \
--tp_size 1 --pp_size 1 --cp_size 4 --ep_size 32 \
--sft --sft_hf_dataset /path/to/qad_blend/train.jsonl \
Expand Down Expand Up @@ -145,10 +147,10 @@ Export the QAD checkpoint to a Hugging Face checkpoint on one 4-GPU node (pipeli

```bash
torchrun --nproc_per_node 4 /opt/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py \
--hf_model_name_or_path <nemotron_h_omni-checkpoint> \
--hf_model_name_or_path nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16 \
--trust_remote_code \
--megatron_path /path/to/qad_output/checkpoints \
--export_unified_hf_path /path/to/nemotron_h_omni_w4a4_hf \
--export_unified_hf_path /path/to/nemotron_3_5_super_vl_w4a4_hf \
--pp_size 4
```

Expand All @@ -167,9 +169,10 @@ Serve with vLLM on one 4-GPU GB300 node, TP=4 with expert parallelism. Quantized
| Setting | Value |
| --- | --- |
| `--kv-cache-dtype` | `fp8` |
| `--max-model-len` | 262144 |
| `--max-model-len` | 262144 (32768 for MMLU-ProX) |

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] This --max-model-len exception contradicts the "same vLLM settings" claim the results table rests on.

Line 17 says "All checkpoints were served with the same vLLM settings (see Section 5)", and the whole point of the new note above the table is to be precise about what the numbers do and don't represent. But MMLU-ProX — the one benchmark with a large, surprising QAD gain (74.41 → 81.15, the row that gets its own paragraph of explanation) — was served at 32768 rather than 262144.

A reader comparing the MMLU-ProX column will want to know whether that 32768 applied to all three checkpoints or only some. If it applied to all three uniformly the comparison is sound and one clause fixes it; if not, the row isn't apples-to-apples and that matters a lot given it's the headline movement. Either way it's worth stating rather than leaving in a parenthetical two tables away from the claim it qualifies.

Suggest saying it explicitly at line 17, e.g. "All checkpoints were served with the same vLLM settings (see Section 5), MMLU-ProX at --max-model-len 32768 for all three", and then the parenthetical here is just a pointer rather than the only mention.

| `--mamba-ssm-cache-dtype` | `float16` |
| `--mamba-conv-cache-dtype` | `float32` |
| Mamba cache rounding | stochastic rounding, 5 Philox rounds |
| `--mamba-cache-dtype` | `float32` |
| `--enable-mamba-cache-stochastic-rounding` | on |
| `--mamba-cache-philox-rounds` | 5 |
| Reasoning parser | `nemotron_v3` |
| Chat template | the checkpoint's `chat_template.jinja` |
2 changes: 1 addition & 1 deletion examples/megatron_bridge/tutorials/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,5 +8,5 @@ Each one walks through a complete workflow using the scripts in [examples/megatr
| Tutorial | What it covers |
| --- | --- |
| [NVIDIA-Nemotron-3-Nano-30B-A3B-BF16](NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md) | End-to-end optimization of the Nemotron-3-Nano-30B-A3B-BF16 (MoE + Mamba-Transformer hybrid) model: Minitron structured pruning (31.6B/A3.6B → 22B/A3.0B) → two-phase knowledge distillation (100B tokens, 8K then 32K seq length) → quantization → vLLM deployment. Includes data-blend preparation, evaluation setup, and detailed pruning / data-blend / long-context ablations. |
| [NemotronH-Omni-W4A4-QAD](NemotronH-Omni-W4A4-QAD/README.md) | **W4A4** quantization of the language model of a NemotronH omni vision-language checkpoint (`NemotronH_Omni_Reasoning_V3`, Mamba-Transformer MoE): mixed NVFP4/FP8 PTQ → KD-only QAD on chat data → Hugging Face export → vLLM serving. |
| [NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16](NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/README.md) | **W4A4** quantization of the language model of NVIDIA Nemotron 3.5 Super VL (`NemotronH_Omni_Reasoning_V3`, Mamba-Transformer MoE): mixed NVFP4/FP8 PTQ → KD-only QAD on chat data → Hugging Face export → vLLM serving. |
| [Qwen3.6-35B-A3B](Qwen3.6-35B-A3B/README.md) | End-to-end **W4A4 NVFP4** optimization of the Qwen3.6-35B-A3B (MoE + hybrid linear/full attention VLM) model: NVFP4 W4A4 post-training quantization → quantization-aware distillation (QAD) to recover the accuracy W4A4 costs → evaluation → vLLM deployment. |
20 changes: 14 additions & 6 deletions modelopt/torch/utils/plugins/prepare_megatron_data_blend.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,8 +199,9 @@ def _to_chat(
if row.get(key):
chat[key] = row[key]
else:
text = row.get("text") or ""
if not text:
text = row.get("text")
# A non-string value (e.g. a list) would otherwise pass as list content.
if not isinstance(text, str) or not text:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] This type check is the right fix for text; the messages branch above it is left relying on the new downstream except for the same class of input.

For text the PR validates at the source — the root cause. For messages (line 191) the equivalent bad input ({"messages": "hello"}, this PR's own new test case) is only caught after it has been fed through prepare_messages_for_chat_template, which iterates the string character by character and appends each char (chars aren't dicts, so they take the not isinstance(msg, dict) pass-through at dataset_utils.py:86), producing {"messages": ["h", "e", "l", "l", "o"]}. The AttributeError only surfaces later, at line 215's turn.get("role").

That works today, but it's fragile in a way specific to this function: if _to_chat ever gains a step between the conversion and the assistant-turn check that tolerates a list of strings, a string messages value silently becomes a per-character conversation instead of being rejected. Validating in the same place as text closes it and makes the two branches read alike:

    if content_field == "messages":
        messages = row.get("messages")
        if not messages or not isinstance(messages, list):
            return None

Per CONTRIBUTING's "fix the bug cause, not the side effect", this is the cause; the broad except in _take_chat_rows then stays a backstop for the unforeseen rather than the mechanism a case you already know about depends on. Worth keeping the except either way — this just stops it from being load-bearing here.

return None
if "<extra_id_" not in text:
return {"messages": [{"role": "assistant", "content": text}]}
Expand Down Expand Up @@ -277,9 +278,16 @@ def _take_chat_rows(
max_rejected = 10 * want + 100
try:
for row in _iter_rows(source, output_dir, shuffle, seed):
chat = _to_chat(row, source["content_field"], reasoning_content)
# Judge the source row's top-level media keys too, which the conversion does not keep.
if chat is None or not is_text_only_chat_example({**row, **chat}):
try:
chat = _to_chat(row, source["content_field"], reasoning_content)
# Judge the source row's top-level media keys too, which the conversion drops.
if chat is not None and not is_text_only_chat_example({**row, **chat}):
chat = None
except (AttributeError, KeyError, TypeError, ValueError) as e:
# As in the tokenizing reader: one malformed record must not lose the source.
print(f"[WARN] {name}: skipping a row that cannot be converted: {e!r}", flush=True)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] This warning is unbounded and prints once per rejected row.

The existing malformed-JSONL warning in _iter_rows has the same shape, but it fires only on genuinely corrupt lines — rare. This one fires on every row that raises, and the give-up budget is 10 * want + 100. A source whose schema systematically mismatches (every row raising, e.g. messages stored as a string throughout) emits up to that many identical lines before the run gives up — for the tutorial's 20-weight source, ~7,000 lines of the same repr, which buries the actual RuntimeError at the end.

Consider capping it, e.g. warn on the first few and then let the existing ({rejected} rejected) / give-up message carry the count:

            except (AttributeError, KeyError, TypeError, ValueError) as e:
                # As in the tokenizing reader: one malformed record must not lose the source.
                if unconvertible < 5:
                    print(f"[WARN] {name}: skipping a row that cannot be converted: {e!r}", flush=True)
                unconvertible += 1
                chat = None

Reporting unconvertible alongside rejected in the final summary line would also distinguish "this source is the wrong schema" from "this source is mostly multimodal", which the single rejected counter currently conflates.

chat = None
Comment on lines +281 to +289

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Compatibility] The per-row try also swallows the config-level KeyError that used to fail fast, turning a typo into a silent 100% rejection.

source["content_field"] is evaluated inside the new try. prepare_chat_data_blend validates the values of content_field (line 377: {source["content_field"] for ...} - {"messages", "text"}) — but that validation itself raises KeyError only on the messages/text path, and crucially _take_chat_rows is also reachable with a source dict whose content_field key is absent in the files: shape used by this PR's own test fixture path. Any such KeyError now gets caught, printed once per row as "skipping a row that cannot be converted", and counted as a rejection.

The consequence: instead of one clear KeyError: 'content_field', the user gets 10 * want + 100 warning lines (for the tutorial's largest source, ~7,000 lines) and then only 0/6890 usable rows (gave up after 6891 rejected rows) — which reads as "this dataset is bad", not "my YAML is wrong". Same for reasoning_content being an invalid mode: prepare_messages_for_chat_template raises ValueError, which is now per-row-swallowed rather than surfacing.

Hoist the lookup out of the loop so a config error can't be mistaken for data:

    content_field = source["content_field"]
    ...
    for row in _iter_rows(source, output_dir, shuffle, seed):
        try:
            chat = _to_chat(row, content_field, reasoning_content)

That keeps the intended fix (a malformed row is skipped) while a malformed source still fails on the first access, before any rows are read.

if chat is None:
rejected += 1
if rejected > max_rejected:
break
Expand All @@ -288,7 +296,7 @@ def _take_chat_rows(
if len(rows) >= want:
break
except Exception as e: # e.g. a gated split, or a schema the streaming reader cannot cast
raise RuntimeError(f"{name}: failed to read or convert") from e
raise RuntimeError(f"{name}: failed to read") from e

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] The outer except Exception still catches conversion failures, so the narrowed message can now be wrong.

The inner handler only catches (AttributeError, KeyError, TypeError, ValueError). Anything else raised by _to_chat or by Megatron-Bridge's is_text_only_chat_example — an IndexError on an unexpected message shape, a RecursionError on a deeply nested content, a KeyError subclass that isn't one — still propagates to line 298 and is reported as "{name}: failed to read". The old wording (failed to read or convert) covered exactly that case; the new one sends someone debugging a 36k-row, 11-source blend looking at dataset access and gating when the read actually succeeded.

Either keep the outer message honest, or make the inner catch the only conversion boundary. The cheapest fix is the former:

    except Exception as e:  # e.g. a gated split, or a schema the streaming reader cannot cast
        raise RuntimeError(f"{name}: failed to read or convert") from e

and revert the matching test change. If the intent really is "the outer handler is read-only now", then the inner except has to be except Exception so no conversion error can reach it — but that is the broader catch, and it makes the masking concern on line 288 worse.

if len(rows) < want:
gave_up = rejected > max_rejected
reason = f"gave up after {rejected} rejected rows" if gave_up else "source exhausted"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,7 @@ def test_to_chat_handles_all_row_shapes():
"messages": [{"role": "assistant", "content": "def f(): pass"}]
}
assert blend_module._to_chat({"text": ""}, "text") is None
assert blend_module._to_chat({"text": ["x"]}, "text") is None
# Markers that parse to no turns: dropped, never trained as assistant text.
assert blend_module._to_chat({"text": "<extra_id_0>\nraw body\n"}, "text") is None
# A role without a mapping would cut its turn out of the middle: the whole row is dropped.
Expand Down Expand Up @@ -342,7 +343,7 @@ def iter_rows(source, output_dir, shuffle, seed):
raise TypeError("cannot cast")

monkeypatch.setattr(blend_module, "_iter_rows", iter_rows)
with pytest.raises(RuntimeError, match="org/a train: failed to read or convert"):
with pytest.raises(RuntimeError, match="org/a train: failed to read"):
blend_module.prepare_chat_data_blend(write_config())
assert not (tmp_path / "blend" / "train.jsonl").exists()

Expand All @@ -357,8 +358,10 @@ def test_write_jsonl_is_atomic(tmp_path: Path):
def test_chat_blend_reads_files_sources(tmp_path: Path, monkeypatch):
plain, packed = tmp_path / "a.jsonl", tmp_path / "b.jsonl.gz"
lines = [json.dumps(_chat(i)) for i in range(20)]
# A malformed record is skipped with a warning, as in the tokenizing reader.
plain.write_text("\n".join([*lines[:10], '{"messages": [', *lines[10:]]) + "\n", "utf-8")
# Malformed records, unparseable or unconvertible, are skipped with a warning as in the
# tokenizing reader.
bad = ['{"messages": [', '{"messages": "hello"}']
plain.write_text("\n".join([*lines[:10], *bad, *lines[10:]]) + "\n", "utf-8")
with gzip.open(packed, "wt", encoding="utf-8") as f:
f.writelines(json.dumps(_chat(i)) + "\n" for i in range(20, 40))
paths = {"data/a.jsonl": str(plain), "data/b.jsonl.gz": str(packed)}
Expand Down
Loading