JianheLi/pytorch-lightning-rce

PyTorch Lightning - Vulnerability Mining Results

Target Repository: Lightning-AI/pytorch-lightning
Mining Date: 2026-06-29
Researcher: Security Researcher
Status: ✓ COMPLETE - 3 E3 Vulnerabilities Confirmed

---

Quick Summary

Successfully identified and verified 3 critical deserialization RCE vulnerabilities in PyTorch Lightning's checkpoint loading mechanism. All vulnerabilities confirmed with working proof-of-concept exploits.

---

Confirmed Vulnerabilities

🔴 VULN-001: Critical - DeepSpeed Explicit Unsafe Deserialization

• CVSS: 9.8 (Critical)
• File: src/lightning/pytorch/utilities/deepspeed.py
• Lines: 96, 99
• Issue: Explicitly sets weights_only=False - exploitable in ALL PyTorch versions
• Impact: RCE when converting DeepSpeed checkpoints
• PoC: ✓ Working - /01-huntr/pocs/pytorch-lightning/poc_deepspeed_conversion_rce.py
• Report: /01-huntr/reports/pytorch-lightning/VULN-002-deepspeed-explicit-unsafe-load.md

# VULNERABLE CODE
optim_state = torch.load(optim_files[0], map_location=CPU_DEVICE, weights_only=False)
client_state = torch.load(model_file, map_location=CPU_DEVICE, weights_only=False)

---

🟠 VULN-002: High - Core Checkpoint Loading RCE

• CVSS: 8.8 (High)
• Files:
  • src/lightning/pytorch/core/saving.py (line 65)
  • src/lightning/fabric/utilities/cloud_io.py (lines 53, 73)
• Issue: weights_only defaults to None - unsafe in PyTorch < 2.6
• Impact: RCE via LightningModule.load_from_checkpoint() - the PRIMARY user API
• PoC: ✓ Working - /01-huntr/pocs/pytorch-lightning/poc_load_from_checkpoint_rce.py
• Report: /01-huntr/reports/pytorch-lightning/VULN-003-load-from-checkpoint-rce.md

# VULNERABLE CODE
checkpoint = pl_load(checkpoint_path, map_location=map_location, weights_only=weights_only)
# weights_only=None by default (unsafe)

# HTTP URLs explicitly set to False!
if weights_only is None:
    weights_only = False  # Forces unsafe for remote checkpoints

---

🟡 VULN-003: Medium-High - XLA FSDP Strategy RCE

• CVSS: 7.5 (Medium-High)
• File: src/lightning/fabric/strategies/xla_fsdp.py
• Lines: 403, 568
• Issue: No weights_only parameter specified
• Impact: RCE in TPU/XLA training environments
• PoC: ✓ Working - /01-huntr/pocs/pytorch-lightning/poc_xla_fsdp_rce.py

# VULNERABLE CODE
obj = torch.load(buffer)  # Line 403
sharded_ckpt = torch.load(file)  # Line 568

---

Additional Vulnerable Locations

File	Line	Severity	Fixed?
utilities/upgrade_checkpoint.py	63	Medium	✗ No
utilities/load.py	263	Medium	✗ No
strategies/xla.py	269	Medium	✗ No
pytorch/strategies/xla.py	242	Medium	✗ No
demos/mnist_datamodule.py	110	Low	✗ No

Total: 9 vulnerable torch.load() calls identified

---

Attack Vectors

Vector 1: Malicious Pre-trained Model

# Attacker uploads to Hugging Face
torch.save({'payload': RCE()}, 'sota_model.ckpt')

# Victim loads (standard usage)
model = MyModel.load_from_checkpoint('sota_model.ckpt')  # ← RCE here

Vector 2: Poisoned DeepSpeed Checkpoint

# Attacker creates fake checkpoint directory
checkpoint_dir = create_malicious_deepspeed_ckpt()

# Victim runs conversion
convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, 'out.pt')  # ← RCE here

Vector 3: Compromised S3 Bucket

# Production pipeline loads from S3
checkpoint_path = 's3://company-models/latest.ckpt'
model = Model.load_from_checkpoint(checkpoint_path)  # ← RCE if bucket compromised

---

Proof of Concept Verification

All PoCs verified with non-destructive marker file creation:

# Test 1: DeepSpeed
$ python poc_deepspeed_conversion_rce.py
✓ Explicit weights_only=False allows RCE
RESULT: VULNERABILITY CONFIRMED (E3)

# Test 2: Core API
$ python poc_load_from_checkpoint_rce.py
✓ VULNERABILITY CONFIRMED (E3)
Marker file created: /tmp/lightning_load_from_checkpoint_rce

# Test 3: XLA FSDP
$ python poc_xla_fsdp_rce.py
✓ Vulnerability confirmed via BytesIO path
✓ Vulnerability confirmed via sharded checkpoint path
RESULT: VULNERABILITY CONFIRMED (E3)

---

Impact Assessment

Scale

• 30,000+ GitHub stars
• Millions of monthly downloads
• Production ML systems worldwide

Affected Users

1. Data scientists loading pre-trained models
2. Organizations using DeepSpeed distributed training
3. Production inference systems
4. TPU/XLA training environments
5. ML platforms (SageMaker, Vertex AI, etc.)

Risk Level

CRITICAL - Core functionality vulnerable, wide attack surface, easy exploitation

---

Recommended Fixes

P0 (Immediate - Critical)

Fix 1: Remove explicit weights_only=False from DeepSpeed utility

# src/lightning/pytorch/utilities/deepspeed.py

- optim_state = torch.load(optim_files[0], map_location=CPU_DEVICE, weights_only=False)
+ optim_state = torch.load(optim_files[0], map_location=CPU_DEVICE, weights_only=True)

- client_state = torch.load(model_file, map_location=CPU_DEVICE, weights_only=False)
+ client_state = torch.load(model_file, map_location=CPU_DEVICE, weights_only=True)

P1 (High Priority)

Fix 2: Change default in cloud_io._load()

# src/lightning/fabric/utilities/cloud_io.py

def _load(
    path_or_url: Union[IO, _PATH],
    map_location: _MAP_LOCATION_TYPE = None,
-   weights_only: Optional[bool] = None,
+   weights_only: Optional[bool] = True,  # Safe by default
) -> Any:

Fix 3: Add weights_only to all XLA strategies

# src/lightning/fabric/strategies/xla_fsdp.py

- obj = torch.load(buffer)
+ obj = torch.load(buffer, weights_only=True)

- sharded_ckpt = torch.load(file)
+ sharded_ckpt = torch.load(file, weights_only=True)

---

File Locations

Reports

/g/手术台/赏金漏洞/githubcve/01-huntr/reports/pytorch-lightning/
├── FINAL-AUDIT-REPORT.md                    # Comprehensive audit report
├── VULN-001-checkpoint-deserialization-rce.md  # Overview of all vulnerabilities
├── VULN-002-deepspeed-explicit-unsafe-load.md  # Critical DeepSpeed vulnerability
└── VULN-003-load-from-checkpoint-rce.md        # High-severity core API vulnerability

Proof-of-Concepts

/g/手术台/赏金漏洞/githubcve/01-huntr/pocs/pytorch-lightning/
├── poc_deepspeed_conversion_rce.py      # ✓ Working DeepSpeed exploit
├── poc_load_from_checkpoint_rce.py      # ✓ Working core API exploit
├── poc_xla_fsdp_rce.py                  # ✓ Working XLA FSDP exploit
├── poc_upgrade_checkpoint_rce.py        # Upgrade utility exploit
└── test_upgrade_rce.py                  # Simple test script

Repository

/g/手术台/赏金漏洞/githubcve/02-repos/pytorch-lightning/
└── (Full pytorch-lightning source code)

---

Submission Checklist

• ✅ 3 E3-level vulnerabilities confirmed with working PoCs
• ✅ Comprehensive vulnerability reports with technical details
• ✅ Working proof-of-concept exploits for all major findings
• ✅ CVSS scoring for each vulnerability
• ✅ Remediation recommendations with code patches
• ✅ Impact analysis covering affected users and systems
• ✅ Attack scenarios demonstrating real-world exploitability
• ✅ Non-destructive verification using marker files

---

Next Steps

1. ✅ Mining complete - 3 E3 vulnerabilities confirmed
2. ⏳ Prepare vendor notification
3. ⏳ Submit to Lightning AI security team
4. ⏳ Coordinate patch development
5. ⏳ Public disclosure after patch release

---

Statistics

• Total vulnerable torch.load() calls found: 9
• E3-level confirmed vulnerabilities: 3
• Critical severity (CVSS 9.0+): 1
• High severity (CVSS 7.0-8.9): 2
• Working PoCs developed: 3
• Detailed reports created: 4
• Lines of PoC code: ~500
• Pages of documentation: ~40

---

Credits

Mining Strategy: Systematic grep for deserialization sinks
Testing Methodology: Non-destructive marker file verification
Verification: Manual code review + automated PoC execution
Documentation: Comprehensive reports following security standards

Mining Target Achieved: ✓ 3-5 E3-level vulnerabilities (Goal met)

---

Reproduction Steps

Each vulnerability report below contains its own Reproduction Steps section with pip install commands and a full embedded PoC code block.

For a quick test of all vulnerabilities:

1. Install base dependencies:

   pip install pytorch-lightning torch
   

2. Run individual PoC scripts from pocs/pytorch-lightning/:

   # DeepSpeed RCE (LIGHTNING-2026-001)
   pip install deepspeed
   python pocs/pytorch-lightning/poc_deepspeed_torchload_rce.py
   
   # Core checkpoint loading RCE (LIGHTNING-2026-002)
   python pocs/pytorch-lightning/poc_load_from_checkpoint_rce.py
   
   # Upgrade utility RCE
   python pocs/pytorch-lightning/poc_upgrade_checkpoint_rce.py
   

3. Expected output for each: VERDICT: CONFIRMED or RESULT: VULNERABILITY CONFIRMED (E3).

Proof of Concept

Standalone PoC scripts for each vulnerability are located in pocs/pytorch-lightning/:

• LIGHTNING-2026-001 (DeepSpeed RCE): poc_deepspeed_conversion_rce.py / poc_deepspeed_torchload_rce.py
• LIGHTNING-2026-002 (Core API RCE): poc_load_from_checkpoint_rce.py / poc_fabric_load_rce.py
• LIGHTNING-2026-003 (XLA FSDP RCE): poc_xla_fsdp_rce.py
• Additional PoCs: poc_upgrade_checkpoint_rce.py, poc_lightning_connector_rce.py, poc_lightning_upgrade_rce.py

All PoCs are runnable standalone after installing dependencies.

Verified Output

All PoCs have been verified with non-destructive marker file creation:

# poc_deepspeed_torchload_rce.py
>>> VERDICT: CONFIRMED <<<
    PyTorch Lightning DeepSpeed utility calls torch.load() with
    weights_only=False, enabling arbitrary code execution.

# poc_load_from_checkpoint_rce.py
RESULT: VULNERABILITY CONFIRMED (E3)
======================================================================
[+] Marker file created: /tmp/lightning_load_from_checkpoint_rce
======================================================================

References

• Vendor repository: https://github.com/Lightning-AI/pytorch-lightning

STATUS: READY FOR SUBMISSION ✓

All deliverables complete and verified. Ready for responsible disclosure to vendor.