Skip to content

test_distributed_resume[ddp/fsdp] fails on any 2+ GPU machine: resumed scheduler lacks initial_lr in param_groups #1048

Description

@zihanlin-ai

Symptom

pytest "tests/unit/train/test_setup_model.py::test_distributed_resume"
...
KeyError: "param 'initial_lr' is not specified in param_groups[0] when resuming
scheduler with last_epoch >= 0."

Both parametrizations (ddp, fsdp) fail identically on unmodified main (reproduced at c05b95b). CI never sees this: the test is @requires_multi_gpu and the single-GPU runners skip it.

Where

Trainer.setup_optimizer (trainer.py#L342-L371): on resume it loads the optimizer state, then builds a fresh scheduler with last_epoch=self.checkpointer.previous_epoch >= 0. PyTorch's LRScheduler.__init__ then requires initial_lr in every param group — which the distributed optimizer-state load does not restore, so the constructor raises.

Environment

torch 2.11.0+cu130, 8x H20, python 3.12; found while validating #1029 (unrelated to that change — fails without it).

Happy to send a fix if useful — e.g. stamping initial_lr into the restored param groups before scheduler construction, or resuming the scheduler via its own state_dict (related to the exact-resume work in #1026/#1027).

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions