Optimizer state-to-parameter association
ID: optimizer-state-to-parameter-association
Optimizer state-to-parameter association identifies which parameter each saved piece of optimizer state belongs to. An optimization checkpoint can contain every stored number and still resume incorrectly if those numbers are attached to different parameters. Equal shapes and data types do not establish mathematical identity.
The PyTorch optimizer dictionary separates per-parameter state from parameter-group settings. Its loader pairs saved parameter identifiers with the current parameters in group order, without independently verifying the intended association.[ref]
torch.optim in the PyTorch documentation. PyTorch2.15, Optimizer.state_dict, state/param_groups and loading note
For an original two-parameter example, use and . Let the named parameters be and , their current gradients both be 1, and their saved momentum buffers be and . Correct restoration givesIf a reversed parameter ordering attaches the same saved buffers to the wrong scalar parameters, the result is insteadBoth buffers had valid scalar shapes, so a shape check would not expose this error.
Consistently permuting the parameters, gradients and state is merely a change of ordering; permuting only the state changes the algorithm's continuation. Record a stable parameter schema and group order with the checkpoint, and check any intentional remapping against that schema. Parameter values and optimizer dictionaries are distinct saved objects, so restoring one does not automatically restore the other's meaning.[ref]
Saving and Loading Models in the PyTorch tutorials. What is a state_dict?; Saving & Loading a General Checkpoint
New to topics? Read the docs here!