transformers
90a491d6 - Fix Trainer checkpoint resume crashing on CPU with multiple processes (#49123)

Commit
11 days ago
Fix Trainer checkpoint resume crashing on CPU with multiple processes (#49123) * fix(trainer): resume from checkpoint on CPU with multiple processes Trainer._load_optimizer_and_scheduler passes self.args.device as map_location when world_size > 1, to load the optimizer state directly on each GPU. On CPU (e.g. torchrun --nproc_per_node=2 with the gloo backend), self.args.device is an indexed device like "cpu:0" rather than plain "cpu" - but torch's CPU deserializer only understands the plain "cpu" location tag, so resuming crashes on every rank with: RuntimeError: don't know how to restore data location of torch.storage.UntypedStorage (tagged with cpu:0) The single-process case already avoids this by using "cpu" directly. Extend that same fallback to the multi-process case whenever the device is CPU, and only use the indexed device object when it's actually a GPU device. Fixes #49121 * Apply review: simplify map_location, drop the added test Per SunMarc's review: the branch fits in one line, and the test isn't needed. Co-authored-by: Neev Modh <141932641+neevmodh@users.noreply.github.com>
Author
Parents
Loading