Feature: Add validation for loaded modelopt state files

[akh64bit]

What needs to be done?

Currently, when a modelopt state file is loaded using modelopt.torch.opt.conversion.load_modelopt_state(), there is no validation performed to ensure that the loaded file actually contains a valid modelopt state object (or dictionary) with the expected structure.

Proposed Solution

Add validation logic inside load_modelopt_state (located in modelopt/torch/opt/conversion.py) to verify the schema of the loaded dictionary before returning it.

Specifically:

1. Ensure the loaded object is a dictionary.
2. Verify that it contains the expected keys (such as "modelopt_state" or other expected schemas).
3. Raise a clear ValueError or TypeError if the file doesn't match the expected structure, rather than allowing downstream errors to occur.

Context

This addresses the existing TODO in modelopt/torch/opt/conversion.py:

# TODO: Add some validation to ensure the file is a valid modelopt state file.
modelopt_state = torch.load(modelopt_state_path, **kwargs)
return modelopt_state

[kevalmorabia97]

Thanks @akh64bit for the suggestion. Would you like to contribute a PR with this change? Shold ne a simple few line code change

[ChenhanYu]

Triage Analysis

Field	Value
Classification	Feature Request
Severity	Medium
Complexity	Simple
Auto-fixable	Yes

Summary

Add validation to the load_modelopt_state() function to ensure loaded files contain valid modelopt state objects with expected schema. Currently there is a TODO comment indicating this validation is needed, and missing validation can lead to unclear downstream errors when invalid files are loaded.

Root Cause / Approach

The load_modelopt_state() function in modelopt/torch/opt/conversion.py currently contains a TODO comment indicating that validation logic needs to be added to verify the schema of loaded modelopt state files before returning them. Without this validation, invalid or corrupted files can be loaded, causing confusing errors downstream.

Suggested Fix

Add validation logic inside the load_modelopt_state() function to: (1) verify the loaded object is a dictionary, (2) check for expected keys like 'modelopt_state' or validate the overall structure, (3) raise a clear ValueError or TypeError with descriptive messages if validation fails. The validation should check the state structure matches the expected format from ModeloptStateManager (e.g., a list of tuples containing mode names and ModeState objects).

Relevant Files

• modelopt/torch/opt/conversion.py
• tests/unit/torch/opt/test_chaining.py
• tests/unit/torch/opt/plugins/test_transformers_save_load.py

---

Auto-triaged by pensieve /magic-triage