Shorten comments from the review-fix rounds

Trim the added comments and test docstrings to two lines each. No code change.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

Author: sugunav14

examples/llm_qat/export.py

@@ -50,9 +50,8 @@ def get_model(
 
     # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this.
     # For QLoRA the base checkpoint is quantized, so from_pretrained already restored the state.
-    # Skipping the restore below is only safe because QATTrainer writes modelopt_state_train.pth at
-    # trainer init from that same base state, so it carries no quantizer values from_pretrained did
-    # not already load. Revisit this if the trainer ever snapshots state later in training.
+    # Skipping is safe only because QATTrainer writes modelopt_state_train.pth at trainer init,
+    # from that same base state.
     if hasattr(model, "peft_config") and not ModeloptStateManager.is_converted(model):
         modelopt_state = mto.load_modelopt_state(f"{ckpt_path}/modelopt_state_train.pth")
         restore_from_modelopt_state(model, modelopt_state)

modelopt/torch/opt/plugins/transformers.py

@@ -110,12 +110,8 @@ def _restore_qtensor_wrappers(model, model_path):
         q_tensor_state = mode_config.get("metadata", {}).get("q_tensor_state", {})
         if not q_tensor_state:
             continue
-        # PEFT nests the quantized linear as `<name>.base_layer`, and either side may carry that
-        # suffix: the state is saved without it when the base model was compressed before adapters
-        # were attached (`quantize.py --compress`) and with it when compressed after
-        # (`QATTrainer._quantize_model`). Normalize both so the lookup works in either direction.
-        # This assumes transformers injects adapters in place; moving the trainer to
-        # `get_peft_model` would prefix module names with `base_model.model.` and break it.
+        # PEFT nests the quantized linear as `<name>.base_layer`, and either the saved keys or the
+        # live names may carry that suffix. Normalize both so the lookup works in either direction.
         q_tensor_state = {k.removesuffix(".base_layer"): v for k, v in q_tensor_state.items()}
 
         pending = [
@@ -134,8 +130,8 @@ def _restore_qtensor_wrappers(model, model_path):
             )
             matched += 1
 
-        # A total miss means the module names were remapped by a wrapper we do not know about.
-        # Warn here rather than letting it surface as an opaque shape error during dequantization.
+        # A total miss means some wrapper renamed the modules. Warn instead of letting it surface
+        # as an opaque shape error at dequantization.
         if pending and not matched:
             warnings.warn(
                 f"Found {len(q_tensor_state)} compressed weight(s) in {modelopt_state_path} but "

tests/examples/llm_qat/test_llm_qat.py

@@ -181,8 +181,8 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path):
         cache_dir=cache_dir,
     )
 
-    # Step 3: Export. Unlike QLoRA this checkpoint is fake-quantized, so the calibrated amaxes
-    # rather than packed weights are what must survive the load.
+    # Step 3: Export. This checkpoint is fake-quantized, so the calibrated amaxes rather than
+    # packed weights are what must survive the load.
     export_dir = tmp_path / "lora_qat_export"
     _run_export(str(lora_qat_output_dir), str(export_dir))
 
@@ -193,9 +193,8 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path):
     base_weights = load_file(base_model_dir / "model.safetensors")
     assert not any("base_layer" in k or k.endswith("_amax") for k in base_weights)
 
-    # LoRA freezes the base model, so exporting the PTQ checkpoint directly is a trusted oracle
-    # for every calibrated value. Comparing against it catches scales that survive as keys but
-    # were silently reset to defaults, which key-presence assertions alone would miss.
+    # LoRA freezes the base model, so a direct PTQ export is a trusted oracle for every calibrated
+    # value. This catches scales that keep their key but were reset to defaults.
     ptq_export_dir = tmp_path / "ptq_export"
     _run_export(str(ptq_output_dir), str(ptq_export_dir))
     reference = load_file(ptq_export_dir / "model.safetensors")

tests/gpu/torch/export/test_export.py

@@ -262,8 +262,7 @@ def test_postprocess_state_dict(state_dict, quantization, maxbound, expected_sta
 def test_postprocess_state_dict_qlora_strips_base_layer():
     """Every QLoRA `base_layer.*` tensor needed for deployment must survive the rename.
 
-    `weight_scale_2` is the NVFP4 global scale and `bias` matters for architectures such as
-    Qwen2; dropping either yields an undeployable checkpoint.
+    Dropping the NVFP4 global scale or a bias yields an undeployable checkpoint.
     """
     state_dict = {
         "layer1.base_layer.weight": torch.ones(4, 2, dtype=torch.uint8),

tests/unit/torch/opt/plugins/test_hf_patching.py

@@ -78,16 +78,14 @@ def _compressed_model_and_state_dir(tmp_path, state_keyed_with_base_layer=False)
 
     state = mto.modelopt_state(model)
     if state_keyed_with_base_layer:
-        # Compressing after the adapters are attached (QATTrainer._quantize_model) saves the
-        # keys with the peft suffix already in them.
+        # Compressing after the adapters are attached saves the keys with the peft suffix.
         for _, mode_config in state["modelopt_state_dict"]:
             q_tensor_state = mode_config.get("metadata", {}).get("q_tensor_state", {})
             for key in list(q_tensor_state):
                 q_tensor_state[f"{key}.base_layer"] = q_tensor_state.pop(key)
     torch.save(state, tmp_path / "modelopt_state.pth")
 
-    # transformers>=5 loads weights by assigning a plain Parameter holding the already-packed
-    # data, which leaves the module without its QTensorWrapper.
+    # transformers>=5 assigns a plain Parameter holding the packed data, dropping the wrapper.
     packed = model.fc.weight.data.clone()
     del model.fc._parameters["weight"]
     model.fc._parameters["weight"] = nn.Parameter(packed, requires_grad=False)