QLoRA DDP export #353
QLoRA DDP export #353
Conversation
WalkthroughAdds QLoRA-aware export tooling and flows: new export script, export APIs extended with an is_modelopt_qlora flag, quant utils updated to remap/skip QLoRA keys and conditionally emit per-layer quant configs based on weight-quantizer enablement, LoRA adapter cleanup, trainer hooks for LoRA best-model loading, and an enabled QLoRA NVFP4 test. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant User
participant ExportScript as export.py
participant Model as QLoRA Model
participant Calib as modelopt_state_train.pth
participant Exporter as _export_hf_checkpoint
participant Utils as postprocess_state_dict
User->>ExportScript: run main(args)
ExportScript->>Model: load base model + LoRA adapters
ExportScript->>Calib: load modelopt_state_train.pth (if present)
Calib-->>Model: restore calibration (+ optional quantizer weights)
ExportScript->>Exporter: _export_hf_checkpoint(model, is_modelopt_qlora=True)
Exporter->>Utils: postprocess_state_dict(state_dict, maxbound, quant, is_modelopt_qlora=True)
Utils-->>Exporter: processed_state_dict + hf_quant_config
Exporter-->>ExportScript: return artifacts
ExportScript->>User: write model, adapters, quant config, tokenizer
sequenceDiagram
autonumber
participant Trainer as QAT/QAD Trainer
participant Model
participant FS as Filesystem
participant Loader as _load_best_model
Trainer->>Model: calibrate quantizers
Trainer->>FS: save modelopt_state_calib.pth (immediately)
alt load best model (no FSDP, LoRA present)
Trainer->>Loader: _load_best_model(...)
Loader-->>Trainer: custom load handling adapters
else
Trainer->>Loader: delegate to superclass
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
sugunav14
force-pushed
the
svelury/qlora-ddp-export
branch
from
035117f to
6254cad
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #353 +/- ##
=======================================
Coverage 74.37% 74.37%
=======================================
Files 182 182
Lines 18219 18219
=======================================
Hits 13550 13550
Misses 4669 4669 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
modelopt/torch/export/quant_utils.py (1)
743-745: Verify dtype casting behavior and potential precision loss.The dtype casting is applied unconditionally when a dtype is provided, which could lead to precision loss or unexpected behavior if the original weight dtype is more precise than the target dtype.
Consider adding a warning or validation:
if dtype: + if weight.dtype != dtype: + logger.info(f"Converting weight from {weight.dtype} to {dtype}") weight = weight.to(dtype)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
examples/llm_qat/README.md(1 hunks)examples/llm_qat/main.py(1 hunks)modelopt/torch/export/quant_utils.py(5 hunks)modelopt/torch/export/unified_export_hf.py(4 hunks)modelopt/torch/quantization/plugins/transformers_trainer.py(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (4)
modelopt/torch/quantization/plugins/transformers_trainer.py (1)
modelopt/torch/export/unified_export_hf.py (1)
export_hf_checkpoint(503-556)
modelopt/torch/export/unified_export_hf.py (3)
modelopt/torch/export/quant_utils.py (2)
maybe_transpose_expert_weight_dimensions(91-120)to_quantized_weight(724-790)modelopt/torch/export/layer_utils.py (1)
is_quantlinear(346-348)modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)
is_enabled(389-391)
modelopt/torch/export/quant_utils.py (1)
modelopt/torch/quantization/qtensor/base_qtensor.py (1)
to(115-123)
examples/llm_qat/main.py (1)
modelopt/torch/quantization/plugins/transformers_trainer.py (1)
export_base_model(291-295)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: linux
🔇 Additional comments (14)
modelopt/torch/export/quant_utils.py (4)
730-730: LGTM! Optional dtype parameter added correctly.The addition of the optional
dtypeparameter follows good API design practices with a sensible default ofNone.
833-836: LGTM! Base layer key mappings added appropriately.The new base layer mappings support QLoRA export workflows by handling the transformation from LoRA-specific keys to standard model keys.
847-847: LGTM! Proper exclusion of base_layer keys.The condition correctly filters out base_layer keys from the main processing loop, which is consistent with the new mappings approach.
899-903: LGTM! LoRA adapter cleanup implemented correctly.The LoRA adapter removal logic is properly implemented to clean up adapter-specific parameters from the exported state dict, ensuring a clean base model export.
examples/llm_qat/main.py (1)
276-278: LGTM! QLoRA export integration implemented correctly.The conditional call to
trainer.export_base_model()is properly guarded by both the LoRA and compression flags, ensuring the base model is only exported when appropriate for QLoRA workflows.modelopt/torch/export/unified_export_hf.py (5)
88-91: LGTM! Early return for LoRA models prevents unnecessary processing.The early return correctly skips processing for LoRA-finetuned models by detecting the presence of a
base_modelattribute, avoiding potential issues with the requantize/resmooth operations.
329-336: Consistent dtype parameter usage.The non-NVFP4 quantization path correctly passes the dtype parameter to
to_quantized_weight, maintaining consistency with the NVFP4 path above.
465-470: Enhanced guard conditions for quantized weight export.The additional check for
hasattr(sub_module, "weight_quantizer")andsub_module.weight_quantizer.is_enabledprovides better safety by ensuring quantizers exist and are active before attempting export.
531-536: LGTM! Proper base model export for QLoRA models.The logic correctly identifies QLoRA models by checking for the
base_modelattribute and exports the underlying base model instead of the wrapper, which is essential for proper deployment.
317-323: Resolved — internal dtype cast remains; no action needed.
quant_utils.py performsif dtype: weight = weight.to(dtype)(modelopt/torch/export/quant_utils.py lines 743–744), so removing the pre-cast in unified_export_hf.py does not change quantization behavior.modelopt/torch/quantization/plugins/transformers_trainer.py (3)
31-31: LGTM! Required import added.The import of
export_hf_checkpointis correctly added to support the new export functionality.
279-290: LoRA-specific best model loading logic implemented correctly.The implementation correctly handles the difference between LoRA and non-LoRA models. For LoRA models, it properly removes and re-loads the adapter from the best checkpoint path.
Note: The TODO comment indicates this is temporary until
get_peft_model()is used, which aligns with the PR description mentioning temporary fixes.
291-296: Simple and effective base model export.The implementation correctly calls
export_hf_checkpointwith the appropriate output directory structure, and the main process check ensures only one process performs the export.examples/llm_qat/README.md (1)
357-362: Fix vLLM serve example: point --tokenizer to the base/merged model (keep --lora-modules syntax)
- Location: examples/llm_qat/README.md (lines 357–362).
- Replace the example so the served model is the merged or original base model and --tokenizer points to the base-model (or HF name). The current --lora-modules adapter=llama3-fp4-qlora is correct.
- Suggested command (use actual paths/names from your repo):
vllm serve --enable-lora --lora-modules adapter= --tokenizer --tokenizer-mode auto --trust-remote-code --port 8000- Verify tokenizer.pad_token_id == model.config.pad_token_id (or set pad_token = eos_token) to avoid generation/padding issues. Repo check found no base_model directory—confirm exact paths before committing this change.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelopt/torch/export/quant_utils.py (1)
732-747: Guard the new dtype pre-cast to prevent fp8/integer pitfallsCasting to an arbitrary dtype before quantization can raise “Promotion for Float8 Types is not supported” or produce invalid math if an integer/bool dtype is passed. Restrict to safe float dtypes and reject fp8/ints.
Apply this diff:
def to_quantized_weight( weight: torch.Tensor, weights_scaling_factor: torch.Tensor, quantization: str, weights_scaling_factor2: torch.Tensor | None = None, block_size: int | None = None, dtype: torch.dtype | None = None, ): @@ - if dtype: - weight = weight.to(dtype) + if dtype is not None: + # Only allow >=16-bit float dtypes here; fp8 and non-floats break downstream ops. + allowed_dtypes = {torch.float32, torch.float16, torch.bfloat16} + disallowed = {getattr(torch, "float8_e4m3fn", None)} + if dtype in disallowed or dtype not in allowed_dtypes: + raise ValueError(f"Unsupported pre-quant cast dtype: {dtype}. Use float32/float16/bfloat16.") + weight = weight.to(dtype)
🧹 Nitpick comments (1)
modelopt/torch/export/quant_utils.py (1)
611-612: Replace commented-out debug print with logger.debug or removeKeep logs consistent and avoid dead code.
Apply this diff:
- # print(f"DEBUG LOG: Processing layer {k} with quantization {v}, block size {block_size_value}") + logger.debug("Processing layer %s with quantization=%s, block_size=%s", k, v, block_size_value)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
modelopt/torch/export/quant_utils.py(7 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
modelopt/torch/export/quant_utils.py (1)
modelopt/torch/quantization/qtensor/base_qtensor.py (1)
to(115-123)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: linux
- GitHub Check: code-quality
- GitHub Check: build-docs
🔇 Additional comments (1)
modelopt/torch/export/quant_utils.py (1)
835-853: Base-layer key remap: avoid silent drops and verify collisionsUnmapped base_layer.* keys are dropped by design. If other base attributes exist (e.g., bias), they’ll vanish silently, and remaps may overwrite existing top-level keys.
Please confirm:
- All required base_layer.* fields are covered by replacements.
- Remapped targets (e.g., “weight”, “input_scale”, “weight_scale”) won’t collide with already-present keys.
If needed, I can add a warning when a base_layer key is encountered but not remapped. Want a patch?
realAsma
left a comment
realAsma
left a comment
There was a problem hiding this comment.
Overall implementation looks good. However we dont have to combine the export for QLoRA with transformer_trainer. We should do the the export via hf_ptq.py
Okay, I will try refactoring the PR to do that! Thank you! |
|
@sugunav14 sounds good, for example here is how we support regular qat deployment - #353 (comment) I am thinking we should have something like: For deployment of QAT and QLoRA checkpoint, we still need to specify For QAT/QLoRA, can we support the following usage: |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Could you test with Phi4-multimodal-instruct export (FP8 and NVFP4) and make sure the before and after the change the safetensors are the same?
To quant Phi4-multimodal-instruct, you need to :
- Download https://huggingface.co/microsoft/Phi-4-multimodal-instruct
- modify https://huggingface.co/microsoft/Phi-4-multimodal-instruct/blob/main/modeling_phi4mm.py#L2101 and enforce using InputMode.LANGUAGE.
- use transformers 4.48
- Run FP8 and NVFP4 PTQ
- Compare the generated safetensors with https://huggingface.co/nvidia/Phi-4-multimodal-instruct-FP4 and https://huggingface.co/nvidia/Phi-4-multimodal-instruct-FP8. Make sure the tensor keys are the same.
modelopt/torch/export/quant_utils.py
Outdated
| # Get the corresponding AWQ block size | ||
| block_size_value = layer_config_dict.get(awq_key, 0) | ||
|
|
||
| # print(f"DEBUG LOG: Processing layer {k} with quantization {v}, block size {block_size_value}") |
modelopt/torch/export/quant_utils.py
Outdated
| keys_to_delete.append(key) | ||
|
|
||
| # remove LoRA adapters from state dict | ||
| for key, value in post_state_dict.items(): |
There was a problem hiding this comment.
what if the original model has lora adapters like phi4-multimodal?
There was a problem hiding this comment.
I did not think of that case! Thanks for the catch!
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (9)
examples/llm_qat/export.py (7)
32-56: Add error handling for missing modelopt state file.Line 47 loads
modelopt_state_calibration.pthwithout checking if it exists. If the file is missing (e.g., model was not trained with QLoRA or was trained with an older version),torch.loadwill raise aFileNotFoundErrorwith an unclear error message.Consider adding explicit validation:
+from pathlib import Path + def get_lora_model( ckpt_path: str, device="cuda", ): """ Loads a QLoRA model that has been trained using modelopt trainer. """ + # Validate modelopt state file exists + modelopt_state_path = Path(ckpt_path) / "modelopt_state_calibration.pth" + if not modelopt_state_path.exists(): + raise FileNotFoundError( + f"modelopt_state_calibration.pth not found in {ckpt_path}. " + "Ensure the model was trained with QLoRA using the modelopt trainer." + ) + device_map = "auto" if device == "cpu": device_map = "cpu" # Load model with adapters model = AutoModelForCausalLM.from_pretrained(ckpt_path, device_map=device_map) # Restore modelopt state - modelopt_state = torch.load(f"{ckpt_path}/modelopt_state_calibration.pth", weights_only=False) + modelopt_state = torch.load(modelopt_state_path, weights_only=False) restore_from_modelopt_state(model, modelopt_state)
53-53: Consider using print_rank_0 for distributed consistency.The print statement on line 53 will execute on all ranks in a distributed setting, potentially causing log clutter. If this script may be used in a distributed context, consider using
print_rank_0frommodelopt.torch.utils.logging(already used in transformers_trainer.py).
87-87: Remove or conditionalize debug print of config data.Line 87 prints the entire config_data dictionary, which can be very verbose. This appears to be debug code left in the script.
Consider removing or making it conditional:
- print(config_data) + # Optionally log config for debugging + # print(config_data)
98-102: Improve error message clarity.The error message "Cannot export model to the model_config" is unclear. Consider clarifying what failed and providing actionable guidance.
warnings.warn( - "Cannot export model to the model_config. The modelopt-optimized model state_dict" - " can be saved with torch.save for further inspection." + f"Failed to export model to {export_dir}. The modelopt-optimized model state_dict " + "can be saved with torch.save for further inspection." )
29-29: Remove unused RAND_SEED constant.The
RAND_SEEDconstant is defined but never used in the script.-RAND_SEED = 1234 -
32-35: Add type hints to function signatures.The function lacks type hints, which would improve code clarity and enable static type checking. Consider adding return type annotation.
-def get_lora_model( - ckpt_path: str, - device="cuda", -): +def get_lora_model( + ckpt_path: str, + device: str = "cuda", +) -> torch.nn.Module:
59-59: Add type hints to main function.Consider adding type hints for better code clarity.
-def main(args): +def main(args: argparse.Namespace) -> None:modelopt/torch/export/quant_utils.py (2)
903-907: Improve LoRA key detection to avoid false positives.The substring check
"lora" in keyon line 906 may inadvertently match keys that contain "lora" as part of a larger word (e.g., "flora", "explorer") or match model-native LoRA adapters that should be preserved.Based on learnings
Apply this diff to use more precise path-segment matching:
# remove LoRA adapters from state dict if is_modelopt_qlora: - for key in post_state_dict: - if "lora" in key and key not in keys_to_delete: + for key in list(post_state_dict.keys()): + parts = key.split(".") + # Check if "lora" appears as a complete path segment or as a prefix (e.g., lora_A, lora_B) + if (("lora" in parts or any(p.startswith("lora_") for p in parts)) + and key not in keys_to_delete): keys_to_delete.append(key)Note: Also changed to iterate over
list(post_state_dict.keys())to avoid issues with dictionary modification during iteration.
1086-1096: Improvement: NVFP4-specific block_size handling partially addresses past concerns.The new logic correctly skips layer config entries for NVFP4 formats when
block_size == 0, which indicates the weight_quantizer is not enabled. This is an improvement over the previous approach that skipped ALL formats withblock_size == 0.However,
awq_block_sizeis still written for all formats on line 1096, even those that don't use block quantization (e.g.,QUANTIZATION_INT8_SQ,QUANTIZATION_FP8,QUANTIZATION_FP8_PC_PT).Consider only writing
awq_block_sizefor formats that actually use it:# Construct per layer config dictionary layer_config_dict[name + ".quantization"] = quantization_format - layer_config_dict[name + ".awq_block_size"] = block_size + # Only write block_size for block-quantized formats + if block_size > 0: + layer_config_dict[name + ".awq_block_size"] = block_sizeThis avoids polluting the config with unnecessary zero values for per-tensor formats, though the
process_layer_quant_configfunction (line 601) already filters these out withif "awq_block_size" in k: continue.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
examples/llm_qat/README.md(1 hunks)examples/llm_qat/export.py(1 hunks)modelopt/torch/export/quant_utils.py(4 hunks)modelopt/torch/export/unified_export_hf.py(3 hunks)modelopt/torch/quantization/plugins/transformers_trainer.py(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- modelopt/torch/export/unified_export_hf.py
🧰 Additional context used
🧬 Code graph analysis (3)
examples/llm_qat/export.py (4)
modelopt/torch/export/convert_hf_config.py (1)
convert_hf_quant_config_format(21-117)modelopt/torch/export/unified_export_hf.py (1)
_export_hf_checkpoint(336-495)modelopt/torch/opt/conversion.py (2)
restore_from_modelopt_state(510-567)modelopt_state(444-486)modelopt/torch/quantization/utils.py (1)
set_quantizer_state_dict(459-466)
modelopt/torch/quantization/plugins/transformers_trainer.py (3)
modelopt/torch/opt/conversion.py (2)
modelopt_state(444-486)save(489-507)modelopt/torch/quantization/utils.py (1)
get_quantizer_state_dict(446-456)modelopt/torch/utils/logging.py (1)
print_rank_0(92-95)
modelopt/torch/export/quant_utils.py (2)
modelopt/torch/export/unified_export_megatron.py (1)
state_dict(465-469)modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)
maxbound(188-194)
🔇 Additional comments (5)
modelopt/torch/quantization/plugins/transformers_trainer.py (3)
212-219: LGTM! Pre-compression state capture enables correct export.Saving the modelopt state and quantizer weights before compression is essential for the export workflow. The export.py script (line 47) loads this exact file to restore the quantizer configuration needed for proper model export. The timing (post-calibration, pre-compression) ensures the quantization metadata is preserved for deployment.
287-297: Ignore attribute name inconsistency suggestion. args.lora is the boolean toggle and args.lora_config holds the adapter config; both are used correctly.Likely an incorrect or invalid review comment.
149-149: Verify default adapter naming inadd_adaptercalls.
- Confirm that
self.model.add_adapter(self.args.lora_config)(lines 149 & 361) assigns a defaultadapter_namethatself.model.active_adapter()(line 293) returns, sodelete_adapterandload_adapteroperate correctly.- If not, pass an explicit
adapter_nametoadd_adapterto guarantee consistent lifecycle management.examples/llm_qat/export.py (1)
71-76: Clarify which quantization config format is saved to hf_quant_config.json.Line 73-74 saves the original
hf_quant_config(modelopt format), while line 76 converts it to llm-compressor format and embeds it in config.json (line 89). This means two different formats are persisted:
base_model/hf_quant_config.json: modelopt formatbase_model/config.json(quantization_config field): llm-compressor formatIf both formats are needed for different consumers, consider adding a comment explaining why both are saved. Otherwise, consider saving only the converted format.
modelopt/torch/export/quant_utils.py (1)
835-846: Potential logic issue: skip_keys may prevent replacements from being applied.On line 846,
"base_layer"is appended toskip_keys. Then on line 852, the code skips any key whereall(skip_key not in key for skip_key in skip_keys)is true. This means keys containing"base_layer"will be skipped entirely and never reach the replacement logic (lines 857-882).However, the replacements dictionary (lines 839-845) includes patterns like
"base_layer.weight"which should be transformed to"weight". These replacements won't be applied because the keys are filtered out first.Consider refactoring to apply replacements first, then skip keys that should be removed:
- skip_keys = ["output_quantizer", "_amax", "_bias_value", "input_quantizer._pre_quant_scale"] + # Keys to skip entirely (not related to quantizers or base transformations) + skip_patterns = ["output_quantizer"] # For modelopt-trained LoRA models, we need to remove the base_layer prefix from the keys for deployment if is_modelopt_qlora: replacements.update( { "base_layer.weight": "weight", "base_layer.input_scale": "input_scale", "base_layer.weight_scale": "weight_scale", } ) - skip_keys.append("base_layer") post_state_dict = {} for key, value in state_dict.items(): - # Skip keys not related to quantizers - if all(skip_key not in key for skip_key in skip_keys): - post_state_dict[key] = value - continue + # Skip keys that should be entirely filtered out + if any(skip_pattern in key for skip_pattern in skip_patterns): + continue - # Apply replacements if the key matches any suffix in the replacements dict + # Try to apply replacements first + replaced = False for old_suffix, new_suffix in replacements.items(): if key.endswith(old_suffix): prefix = key[: -len(old_suffix)] if "_amax" in key: # ... existing _amax handling ... post_state_dict[prefix + new_suffix] = value + replaced = True break + + # If no replacement was applied and key doesn't contain quantizer-specific suffixes, keep it + if not replaced: + quantizer_suffixes = ["_amax", "_bias_value", "input_quantizer._pre_quant_scale"] + if not any(key.endswith(suffix) for suffix in quantizer_suffixes): + post_state_dict[key] = valueLikely an incorrect or invalid review comment.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
examples/llm_qat/export.py (4)
29-29: Remove unused constant.
RAND_SEEDis defined but never used in the script.-RAND_SEED = 1234 -
32-56: Consider improving robustness and observability.The function lacks error handling for missing checkpoint files and uses print instead of logging. Consider:
- Error handling: Wrap file loading in try-except to provide clear error messages if
modelopt_state_calibration.pthis missing or corrupted.- Logging: Replace
print("Restoring modelopt weights")with proper logging (e.g.,logging.info(...)).- Documentation: Expand the docstring to document parameters, return value, and expected checkpoint structure.
Example improvements:
+import logging + def get_lora_model( ckpt_path: str, device="cuda", ): """ - Loads a QLoRA model that has been trained using modelopt trainer. + Loads a QLoRA model that has been trained using modelopt trainer. + + Args: + ckpt_path: Path to the checkpoint directory containing the model and modelopt state. + device: Device to load the model on ("cuda" or "cpu"). + + Returns: + The loaded model with restored modelopt and quantizer state. """ device_map = "auto" if device == "cpu": device_map = "cpu" # Load model with adapters model = AutoModelForCausalLM.from_pretrained(ckpt_path, device_map=device_map) # Restore modelopt state + try: - modelopt_state = torch.load(f"{ckpt_path}/modelopt_state_calibration.pth", weights_only=False) + modelopt_state = torch.load(f"{ckpt_path}/modelopt_state_calibration.pth", weights_only=False) + except FileNotFoundError as e: + raise FileNotFoundError( + f"modelopt_state_calibration.pth not found in {ckpt_path}. " + "Ensure the checkpoint was saved correctly during training." + ) from e restore_from_modelopt_state(model, modelopt_state) # Restore modelopt quantizer state dict modelopt_weights = modelopt_state.pop("modelopt_state_weights", None) if modelopt_weights is not None: - print("Restoring modelopt weights") + logging.info("Restoring modelopt quantizer weights") set_quantizer_state_dict(model, modelopt_weights) return model
59-101: Use Path objects consistently and improve error messages.The function mixes Path objects with f-string concatenation. For consistency and robustness, use Path objects throughout. Also, the exception handler's warning message is generic and doesn't help users diagnose the issue.
Apply this diff:
def main(args): # Load model model = get_lora_model(args.pyt_ckpt_path, args.device) tokenizer = AutoTokenizer.from_pretrained(args.pyt_ckpt_path) # Export HF checkpoint export_dir = Path(args.export_path) export_dir.mkdir(parents=True, exist_ok=True) base_model_dir = export_dir / "base_model" base_model_dir.mkdir(parents=True, exist_ok=True) try: post_state_dict, hf_quant_config = _export_hf_checkpoint(model, is_modelopt_qlora=True) - with open(f"{export_dir}/base_model/hf_quant_config.json", "w") as file: + with open(base_model_dir / "hf_quant_config.json", "w") as file: json.dump(hf_quant_config, file, indent=4) hf_quant_config = convert_hf_quant_config_format(hf_quant_config) # Save base model - model.base_model.save_pretrained(f"{export_dir}/base_model", state_dict=post_state_dict) + model.base_model.save_pretrained(base_model_dir, state_dict=post_state_dict) # Save adapters model.save_pretrained(export_dir) - config_path = f"{export_dir}/base_model/config.json" + config_path = base_model_dir / "config.json" # In the case of LoRA model.save_pretrained does not save the correct config.json config_data = model.config.to_dict() config_data["quantization_config"] = hf_quant_config with open(config_path, "w") as file: json.dump(config_data, file, indent=4) # Save tokenizer tokenizer.save_pretrained(export_dir) except Exception as e: warnings.warn( - "Cannot export model to the model_config. The modelopt-optimized model state_dict" - " can be saved with torch.save for further inspection." + f"Export failed: {e}. The modelopt-optimized model state_dict " + "can be saved with torch.save for further inspection." ) raise e
85-86: Clarify the comment about config.json handling.The comment is grammatically incomplete. Consider rephrasing for clarity.
- # In the case of LoRA model.save_pretrained does not save the correct config.json + # For LoRA models, save_pretrained does not include quantization_config in config.json, + # so we manually reconstruct and write it here.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
examples/llm_qat/export.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
examples/llm_qat/export.py (4)
modelopt/torch/export/convert_hf_config.py (1)
convert_hf_quant_config_format(21-117)modelopt/torch/export/unified_export_hf.py (1)
_export_hf_checkpoint(336-495)modelopt/torch/opt/conversion.py (2)
restore_from_modelopt_state(510-567)modelopt_state(444-486)modelopt/torch/quantization/utils.py (1)
set_quantizer_state_dict(459-466)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: linux
- GitHub Check: code-quality
- GitHub Check: build-docs
🔇 Additional comments (2)
examples/llm_qat/export.py (2)
104-118: LGTM!The argument parser is well-structured with appropriate defaults and a required checkpoint path.
1-118: Acknowledge planned refactor per PR discussion.Based on the PR objectives, this script is temporary. The author has agreed to refactor the export logic to integrate with
hf_ptq.pyper reviewer feedback. This standalone script serves the immediate QLoRA export use case but should be consolidated with the existing export flow in a follow-up.Based on PR objectives summary.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
examples/llm_qat/README.md (1)
368-370: Consider clarifying the adapter name convention.The vLLM command uses
--lora-modules adapter=llama3-fp4-qlora-hfwhere "adapter" is a generic name. It would be helpful to clarify whether this name is:
- A fixed convention expected by vLLM
- User-configurable and should match a specific naming pattern
- Related to the export output structure
Adding a brief note about the adapter naming would improve usability, especially given the PR discussion about simplifying the CLI for QLoRA checkpoints.
modelopt/torch/export/quant_utils.py (1)
904-908: Substring matching may catch unintended keys.The LoRA adapter removal logic uses substring matching (
"lora" in key), which could inadvertently match keys like "flora", "exploration", or "coloration". While theis_modelopt_qloragate reduces this risk, consider more precise matching.Apply this diff for more precise LoRA key detection:
# remove LoRA adapters from state dict if is_modelopt_qlora: for key in post_state_dict: - if "lora" in key and key not in keys_to_delete: + # Match LoRA keys more precisely: check for .lora. or lora_ patterns + parts = key.split(".") + if (any(p.startswith("lora_") or p == "lora" for p in parts)) and key not in keys_to_delete: keys_to_delete.append(key)This ensures matching only keys with "lora" as a complete path component or prefix (e.g., "model.lora_A.weight", "adapter.lora.bias"), not arbitrary substrings.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
examples/llm_qat/README.md(1 hunks)examples/llm_qat/export.py(1 hunks)modelopt/torch/export/quant_utils.py(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- examples/llm_qat/export.py
🧰 Additional context used
🧬 Code graph analysis (1)
modelopt/torch/export/quant_utils.py (4)
modelopt/torch/export/unified_export_megatron.py (1)
state_dict(465-469)modelopt/torch/opt/conversion.py (1)
state_dict(130-132)modelopt/torch/distill/distillation_model.py (1)
state_dict(189-192)modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)
maxbound(188-194)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: linux
- GitHub Check: build-docs
- GitHub Check: code-quality
🔇 Additional comments (6)
examples/llm_qat/README.md (1)
357-370: Documentation provides clear QLoRA deployment flow.The new export and deployment instructions successfully replace the previous experimental note with actionable guidance. The command examples are well-structured and the reference to vLLM documentation is helpful.
modelopt/torch/export/quant_utils.py (5)
812-827: LGTM: Parameter documented.The new
is_modelopt_qloraparameter has been properly documented in the docstring at line 824. The documentation is clear and consistent with the existing style.
836-836: LGTM: Centralized skip logic.The
skip_keyslist provides a cleaner, more maintainable approach to filtering non-quantizer keys. This addresses the concern from previous reviews about structured key handling.
838-847: LGTM: LoRA handling properly gated.The LoRA-specific key transformations are now correctly gated behind the
is_modelopt_qloraflag. This addresses the previous concern about unconditional LoRA adapter removal and ensures the logic only applies to ModelOpt QLoRA exports.
1087-1093: LGTM: NVFP4 block_size filtering fixed.The guard correctly skips only NVFP4-related formats when
block_size == 0, indicating the weight quantizer is not enabled. This fixes the previous issue where non-block formats (INT8_SQ, FP8, FP8_PC_PT) were incorrectly dropped. Non-block formats now proceed to lines 1095-1097 regardless of block_size.
851-855: postprocess_state_dict filtering logic is correct
Non-quantizer keys are retained, and all quantizer/base_layer keys are either transformed per the replacements mapping or removed as intended.
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Co-authored-by: realAsma <[email protected]> Signed-off-by: sugunav14 <[email protected]>
Co-authored-by: realAsma <[email protected]> Signed-off-by: sugunav14 <[email protected]>
Co-authored-by: realAsma <[email protected]> Signed-off-by: sugunav14 <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
sugunav14
force-pushed
the
svelury/qlora-ddp-export
branch
from
79ef31b to
7cca3b7
Compare
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
Signed-off-by: Suguna Velury <[email protected]>
What does this PR do?
Type of change: New example
Overview: This PR provides an e2e example for fine-tuning a model using QLoRA with DDP and exporting checkpoint for deployment using vllm.
Usage
Refer to README.md changes
Testing
Trainer
./launch.sh --model meta-llama/Meta-Llama-3-8B --num_epochs 0.01 --lr 1e-3 --do_train True --output_dir test --quant_cfg FP8_DEFAULT_CFG --compress True --lora TrueExport
python export.py --pyt_ckpt_path test --export_dir test-fp8Deployment
vllm serve test-fp8/base_model --enable-lora --lora-modules sql-lora=test-fp8 --port 8090 --tokenizer test-fp8e2e unit test
Sanity check weights, dtypes of generated checkpoint
Test phi4
Before your PR is "Ready for review"
Additional Information
Summary by CodeRabbit
New Features
Improvements
Documentation
Tests