Model-Optimizer

Root cause: In ModelOpt's modelopt.torch.utils.plugins.megatron_generate.megatron_generate (affects 0.46.0rc1 and main), the vision inputs pixel_values, image_grid_thw, and image_sizes are gated only by step == 0. With enable_kv_cache=False, every decoding step recomputes the full growing token prefix, but steps 1+ run without vision embeddings, so image-placeholder tokens are never replaced and generation silently becomes visually ungrounded (no error is raised). This also triggers when model.config.sequence_parallel=True, since megatron_generate then auto-falls back to no-cache decoding.

Fix: Change the gating condition for all three visual tensors from step == 0 to:

replay_vision_inputs = inference_context is None or step == 0

Apply this condition to pixel_values, image_grid_thw, and image_sizes so vision inputs are passed on every full-prefix recomputation when no inference context (KV cache) exists, and only during prefill when KV cache is active.

Verified result: with the patch, prefix logits are step-invariant (max|diff| = 0.0, argmax agreement 1.000 vs. 0.529 stock), and Qwen3-VL-8B-Instruct correctly describes the image instead of hallucinating unrelated content.

Workarounds if you cannot patch: (1) enable KV cache (enable_kv_cache=True) so visual tensors are only needed at prefill — not possible with sequence_parallel; or (2) monkey-patch the generator function in memory with the corrected gating condition before calling it.

This crash occurs on Megatron-Core models with share_embeddings_and_output_weights=True (e.g., Qwen3-0.6B GPTModel), because the output_layer stores weight = None (borrowing the embedding weight at forward time), while mtq.fold_weight dereferences the stored weight without a None guard. **Fix:** Upgrade TensorRT-Model-Optimizer to a version including PR #2140, which makes fold_weight skip pairs whose weight is not a tensor. Skipping is correct since the borrowed tensor belongs to the embedding module and is folded there. **Workaround (before the fix):** Reimplement the per-pair weight fold manually with a None guard instead of calling mtq.fold_weight, i.e., skip when getattr(module, weight_name) is not a torch.Tensor — this is the approach used in NeMo-RL PR #3441. Note that normal QAT forward passes don't hit this bug because the Megatron plugin quantizes the weight argument passed to the linear functional, not the stored module.weight; only fold_weight reads the stored attribute.

Puzzletron's calculate_kv_dim in modelopt/torch/puzzletron/utils/misc.py computes head_size = n_embd // n_head instead of reading head_dim from the model config. For models where head_dim != hidden_size / num_attention_heads, the KV-cache cost of every attention layer is wrong: Mistral-Small-24B (head_dim 128 vs computed 160) is overestimated by 25%, gpt-oss-20b (head_dim 64 vs computed 45) is underestimated by 30%, and Gemma-3-12B (256 vs 240) is also affected. Under a target_memory constraint this biases the solver toward wrongly pruning attention blocks. Fix (PR #2484): resolve the head dimension with the precedence qk_head_dim (if set) → model config head_dim → fallback n_embd // n_head, e.g. head_dim = getattr(config, 'head_dim', n_embd // n_head) then kv_dim = 2 * num_kv_heads * head_dim. Apply the same resolved head dimension to the attention FLOP estimate in calculate_additive_metrics (which previously used qk_head_dim or (n_embd // n_head)). Llama-3.x, Qwen2.5-7B and Qwen3-8B are unaffected because their head_dim equals hidden_size/num_attention_heads.

The most common cause of poor QAD accuracy recovery is too-small global batch size and too-short sequence length (e.g., gbs=16, seq_len=8192 will show little improvement). Per NVIDIA maintainers, the paper configs (e.g., AceReason Nemotron 1.1 7B) use:

- seq_len = 32768 with sequence packing enabled
- global_batch_size = 100 (equivalent to ~1024 packed samples)
- warmup_steps = 10
- Train only 200–600 steps (~1B tokens total); do not reuse samples if you have enough data
- Total tokens: ~1B
- Loss: forward KL, temperature T=1

Critical points:
1. Use the model's ORIGINAL SFT dataset as the distillation data — this is emphasized as essential.
2. Select the best checkpoint by running actual downstream benchmarks, not just validation loss.
3. The paper used the NeMo/Megatron-LM path (examples/llm_qad), which is more actively maintained than the Hugging Face/FSDP2 path (examples/llm_qat); use Megatron-LM for reproduction.
4. Standard optimizer settings (AdamW, lr ~1e-5, cosine schedule, grad clip 1.0, BF16) are fine — the data and batch/sequence scale are what matter.

In modelopt/onnx/quantization/fp8.py, INT8 scales are converted with np_fp8_scale = (np_scale * 448.0) / 127.0. Since int8_scale = amax / 127 and fp8_scale = amax / 448, the mathematically consistent conversion is fp8_scale = int8_scale * (127.0 / 448.0). The current multiplier inflates scales by ~3.5×, so quantized tensors occupy only about ±36 of the FP8 ±448 range and can trigger activation blow-up, saturation/clipping, and layer imbalance. NVIDIA's recommended workaround: do not convert INT8 scales — run direct amax calibration for FP8 instead. This avoids both inflated values (current code) and values underflowing to zero (naive reverse ratio). Note: the source still contains scale * 448 / 127, so prefer native FP8 amax calibration until an official fix lands.