Add back support for stage.remap_qualname()

[kwen2501]

Description

stage.remap_qualname(key) maps a stage's parameter name (key) back to the original model's parameter name.

This now works:

    # Stage module's state dict
    sd = stage.submod.state_dict()
    remapped_keys = [stage.remap_qualname(k) for k in sd.keys()]

    # Original model's state dict
    old_keys = mod.state_dict().keys()

    # Confirm they match
    assert all(rk in old_keys for rk in remapped_keys)

[H-Huang]

nit: wouldn't this evaluate to be True if remapped_keys was a subset of old_keys.

Maybe a better check would be sorted(remapped_keys) == sorted(list(old_keys))

[kwen2501]

remapped_keys is always a subset of old_keys, because remapped_keys are only keys of a stage, while old_keys contains keys of the full original model.