Skip to content

fix: PTQ 1GPU, export PP divisibility, hidden states conversations key #1293

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChenhanYu merged 1 commit into from
Apr 20, 2026
Merged

fix: PTQ 1GPU, export PP divisibility, hidden states conversations key #1293

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion examples/speculative_decoding/collect_hidden_states/compute_hidden_states_trtllm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,7 +256,7 @@ async def submit_generates():
for entry in dataset:
conversation_id = entry.get("conversation_id", entry.get("uuid"))

conversations = entry["conversations"]
conversations = entry.get("messages") or entry.get("conversations")
if not conversations or not isinstance(conversations, list):
num_invalid += 1
continue
Expand Down
17 changes: 14 additions & 3 deletions tools/launcher/common/megatron_lm/quantize/quantize.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,11 +41,22 @@ TP=${TP:-1} PP=${PP:-1} EP=${EP:-1} ETP=${ETP:-1} ${QUANTIZE_EXE} ${MLM_MODEL_CF
export MLM_EXTRA_ARGS="--mmlu-dataset ${MMLU_DATASET:-/hf-local/cais/mmlu} --fraction 0.01 --lower-bound ${MMLU_LOWER_BOUND:-0.38} --disable-tqdm"
TP=${TP:-1} PP=${PP:-1} EP=${EP:-1} ETP=${ETP:-1} MLM_MODEL_CKPT=${MLM_MODEL_SAVE} ${MMLU_EXE} ${MLM_MODEL_CFG}

# Export quantized checkpoint to HF format (PP=all GPUs)
# Export quantized checkpoint to HF format
# Use largest PP <= total GPUs that divides the model's num_hidden_layers
TOTAL_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo ${NUM_GPUS:-1})
echo "=== Exporting ${MLM_MODEL_CFG} ${QUANT_CFG} (PP=${TOTAL_GPUS}) ==="
EXPORT_PP=$(python3 -c "
import json, os
cfg = os.path.join('${HF_MODEL_CKPT}', 'config.json')
n_layers = json.load(open(cfg)).get('num_hidden_layers', 1) if os.path.exists(cfg) else 1
gpus = ${TOTAL_GPUS}
pp = gpus
while pp > 1 and n_layers % pp != 0:
pp -= 1
print(pp)
" 2>/dev/null || echo ${TOTAL_GPUS})
echo "=== Exporting ${MLM_MODEL_CFG} ${QUANT_CFG} (PP=${EXPORT_PP}, ${TOTAL_GPUS} GPUs) ==="
export MLM_EXTRA_ARGS=
TP=1 PP=${TOTAL_GPUS} EP=1 ETP=1 MLM_MODEL_CKPT=${MLM_MODEL_SAVE} ${EXPORT_EXE} ${MLM_MODEL_CFG}
TP=1 PP=${EXPORT_PP} EP=1 ETP=1 MLM_MODEL_CKPT=${MLM_MODEL_SAVE} ${EXPORT_EXE} ${MLM_MODEL_CFG}
Comment on lines 46 to +59

@coderabbitai coderabbitai Bot Apr 20, 2026

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.

⚠️ Potential issue | 🟠 Major

Validate and clamp GPU/PP values before export invocation.

At Line 46, torch.cuda.device_count() can return 0; then Line 55 can emit EXPORT_PP=0, and Line 59 runs export with invalid PP=0. Also, the values are used without numeric validation.

🔧 Suggested hardening patch 
-TOTAL_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo ${NUM_GPUS:-1})
+TOTAL_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo "${NUM_GPUS:-1}")
+if ! [[ "${TOTAL_GPUS}" =~ ^[0-9]+$ ]] || [[ "${TOTAL_GPUS}" -lt 1 ]]; then
+    TOTAL_GPUS=1
+fi
 EXPORT_PP=$(python3 -c "
 import json, os
 cfg = os.path.join('${HF_MODEL_CKPT}', 'config.json')
 n_layers = json.load(open(cfg)).get('num_hidden_layers', 1) if os.path.exists(cfg) else 1
-gpus = ${TOTAL_GPUS}
+gpus = int('${TOTAL_GPUS}')
 pp = gpus
 while pp > 1 and n_layers % pp != 0:
     pp -= 1
 print(pp)
-" 2>/dev/null || echo ${TOTAL_GPUS})
+" 2>/dev/null || echo "${TOTAL_GPUS}")
 echo "=== Exporting ${MLM_MODEL_CFG} ${QUANT_CFG} (PP=${EXPORT_PP}, ${TOTAL_GPUS} GPUs) ==="
 export MLM_EXTRA_ARGS=
-TP=1 PP=${EXPORT_PP} EP=1 ETP=1 MLM_MODEL_CKPT=${MLM_MODEL_SAVE} ${EXPORT_EXE} ${MLM_MODEL_CFG}
+TP=1 PP="${EXPORT_PP}" EP=1 ETP=1 MLM_MODEL_CKPT="${MLM_MODEL_SAVE}" ${EXPORT_EXE} "${MLM_MODEL_CFG}"
🧰 Tools
🪛 Shellcheck (0.11.0)

[info] 46-46: Double quote to prevent globbing and word splitting.

(SC2086)

[info] 56-56: Double quote to prevent globbing and word splitting.

(SC2086)

[info] 59-59: Double quote to prevent globbing and word splitting.

(SC2086)

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@tools/launcher/common/megatron_lm/quantize/quantize.sh` around lines 46 - 59,
TOTAL_GPUS and EXPORT_PP can be zero/unvalidated causing an invalid PP=0 during
the export invocation; before using them (the TOTAL_GPUS assignment and the
EXPORT_PP calculation plus the final export invocation that sets TP/PP/EP/ETP),
validate and clamp both to >=1 (and optionally to NUM_GPUS if provided) and
fallback to 1 on errors; update the python one-liners or add a small shell check
after them to coerce non-numeric or zero values to 1 and log the corrected
values so the final export invocation always receives a valid PP (and
TOTAL_GPUS) >=1.

ls ${EXPORT_DIR}
cat ${EXPORT_DIR}/hf_quant_config.json

Expand Down
16 changes: 8 additions & 8 deletions tools/launcher/examples/Qwen/Qwen3-8B/megatron_lm_ptq.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ pipeline:
config:
model: Qwen/Qwen3-8B
quant_cfg: NVFP4_DEFAULT_CFG
tp: 8
tp: 1
calib_dataset: abisee/cnn_dailymail
calib_size: 32
mmlu_dataset: cais/mmlu
Expand All @@ -33,15 +33,15 @@ pipeline:
slurm_config:
_factory_: "slurm_factory"
nodes: 1
ntasks_per_node: 8
gpus_per_node: 8
ntasks_per_node: 1
gpus_per_node: 1

task_1:
_target_: common.megatron_lm.quantize.task.MegatronLMQuantizeTask
config:
model: Qwen/Qwen3-8B
quant_cfg: FP8_DEFAULT_CFG
tp: 8
tp: 1
calib_dataset: abisee/cnn_dailymail
calib_size: 32
mmlu_dataset: cais/mmlu
Expand All @@ -50,18 +50,18 @@ pipeline:
slurm_config:
_factory_: "slurm_factory"
nodes: 1
ntasks_per_node: 8
gpus_per_node: 8
ntasks_per_node: 1
gpus_per_node: 1

# Step 3: TRT-LLM eval MMLU on all exported checkpoints
task_2:
script: common/tensorrt_llm/eval.sh
environment:
- HF_MODEL_CKPT: /scratchspace/export
- TP: "8"
- TP: "1"
- EP: "1"
slurm_config:
_factory_: "slurm_factory"
nodes: 1
ntasks_per_node: 1
gpus_per_node: 8
gpus_per_node: 1
Loading