[Discussion] Evaluate safer serialization alternatives to pickle/torch.load

[RinZ27]

Several core modules (e.g., modelopt/torch/opt/plugins/megatron.py, distribute.py) currently rely on torch.load(weights_only=False) and pickle.loads for state management. While these are documented as internal use cases, they present RCE risks if checkpoints are sourced from untrusted environments (e.g., public hubs).

I suggest exploring safetensors or implementing strict safe_globals allowlists for future releases to harden the library against malicious model files.

[kevalmorabia97]

@RinZ27 yes this is very important item on our TODO. Would you like to contribute a PR with some suggestions?

[RinZ27]

@kevalmorabia97 Glad to hear this is on the roadmap! I'd be happy to help.

As a first step, I suggest implementing weights_only=True in torch.load calls combined with a curated allow-list of internal Modelopt classes (using torch.serialization.add_safe_globals). This approach provides an immediate security boost while maintaining compatibility for existing state files.

I'll prepare a draft PR for modelopt/torch/opt/conversion.py to demonstrate this pattern. Looking forward to your thoughts on this direction.

[ChenhanYu]

Triage Analysis

Field	Value
Classification	Enhancement
Severity	High
Complexity	Moderate
Auto-fixable	No

Summary

The library currently uses unsafe deserialization methods (torch.load with weights_only=False and pickle.loads) for checkpoint handling, creating potential RCE vulnerabilities when loading from untrusted sources. The issue proposes implementing safer alternatives using weights_only=True with safe_globals allowlists or exploring safetensors format.

Root Cause / Approach

The codebase relies on torch.load(weights_only=False) and pickle.loads for state management across multiple modules. While documented as internal use, this creates security risks when checkpoints originate from untrusted environments. The current approach lacks input validation and safe deserialization constraints.

Suggested Fix

1. Implement torch.load with weights_only=True across all checkpoint loading operations. 2. Create a curated allowlist of internal Modelopt classes using torch.serialization.add_safe_globals for any classes that must be unpickled. 3. Add helper functions to centralize safe loading logic. 4. Consider optional safetensors support for future releases. 5. Add security documentation about checkpoint source validation. 6. Update all torch.load calls in megatron.py, distribute.py, mcore_dist_checkpointing.py with safe parameters.

Relevant Files

• modelopt/torch/opt/conversion.py
• modelopt/torch/opt/plugins/megatron.py
• modelopt/torch/export/distribute.py
• modelopt/torch/quantization/conversion.py
• modelopt/torch/opt/plugins/mcore_dist_checkpointing.py
• modelopt/torch/export/plugins/megatron_importer.py
• modelopt/torch/nas/conversion.py

---

Auto-triaged by pensieve /magic-triage