Question about loading checkpoint

I was able to resume training by only ‘load_state_dict’ in monitor peer before using hivemind 1.0.0 version. The code looks like this:

# monitor peer
if load_from_pretrained:
  self.model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"), strict=False)
  ...
  self.collaborative_optimizer.load_state_dict(torch.load("optimizer.pt", map_location="cpu"))

The peers would load from monitor’s state after start up.

However, in ver 1.0.0 or master code, ‘load_state_dict’ in monitor seems not work. My question is am I using the wrong method or should I load the checkpoint on the worker peer?

[finger92]

I solved this! In example/albert The trainer peer used scheduler while monitor peer not. Which results in some differences in the “optimizer state_dict” of the two peers(scheduler will add a ‘initial_lr’ to optimizer’s state). After adding a non-functional “scheduler” in monitor peer it works fine.

By the way, I changed “prefix” in state_averager in monitor peer’s code to let trainer could download state from monitor

self.state_averager = TrainingStateAverager(
            dht=dht,
            optimizer=opt,
            prefix=f"{experiment_prefix}_state_averager",
            state_compression=hivemind.Float16Compression(),
            bandwidth=optimizer_args.bandwidth,
            client_mode=optimizer_args.client_mode,
            start=True,
            **asdict(averager_args),
        )

[justheuristic]

Hi! Awesome work! Feel free to ping us if you encounter any more oddities 😃

We’ll incorporate your fixes into the example in the coming days (within a week or two at most) and write back to you with an update.