PushToHub#

class PushToHub(repo_id, checkpoint='model_best', weights_only=True, private=True, token=None, revision=None, name=None, on_train_end=True, on_best=False, every_n_epochs=None, every_n_steps=None, background=True, squash_history=True, squash_every_n_pushes=None)[source]#

Bases: HookBase

Push model checkpoints to the Hugging Face Hub during and/or after training.

Entirely opt-in. By default it uploads the raw model_best checkpoint once at the end of training (weights_only=True), so it loads back byte-identically via weight=hf://<repo>/model_best.pth; set weights_only=False to push a consolidated pimm export artifact instead. For cross-cluster sync it can also push periodically: with on_best=True it pushes model_best in after_step whenever the best metric improves, and with every_n_epochs set it pushes the rolling checkpoint in after_epoch every N epochs. The final end-of-training push happens in after_train, which also drains any in-flight background uploads. Registered as PushToHub.

Parameters:

  • repo_id (str) – Target Hub repo, e.g. "youngsm/sonata-pilarnet-L".

  • checkpoint (str) – Which checkpoint to push at end of training, e.g. "model_best" or "last"/"model_last" (the latter resolve to the newest rolling checkpoint). Defaults to "model_best".

  • weights_only (bool) – If True, upload the raw checkpoint file; if False, build and push a consolidated pimm export artifact. Defaults to True.

  • private (bool) – Create the repo as private if it does not exist. Defaults to True.

  • token (str, optional) – Hugging Face token; falls back to the ambient credential when None. Defaults to None.

  • revision (str, optional) – Target branch/revision for uploads. Defaults to None.

  • name (str, optional) – Destination filename in the repo for raw uploads; defaults to the source file’s basename when None. Defaults to None.

  • on_train_end (bool) – Push checkpoint in after_train. Defaults to True.

  • on_best (bool) – Push model_best in after_step whenever best_metric_value improves. Defaults to False.

  • every_n_epochs (int, optional) – Push the rolling checkpoint in after_epoch every N epochs. None disables. Defaults to None.

  • every_n_steps (int, optional) – Push the rolling checkpoint in after_step every N global steps – use this to match an iteration-oriented saver like CheckpointSaverIteration (set it to the saver’s save_freq). The step counter is seeded from resumed progress in before_train. None disables. Defaults to None.

  • background (bool) – Run periodic raw uploads in a background thread so they never block the training step (uploads to the same path are serialized). The final after_train push is always blocking. Defaults to True.

  • squash_history (bool) – Squash the repo’s git/LFS history once in after_train (after the final push) to reclaim orphaned LFS blobs. Cheap server-side metadata op; irreversible (drops history). Defaults to True.

  • squash_every_n_pushes (int, optional) – Also squash after every N successful periodic pushes (on_best/every_n_epochs) so a long run does not accumulate unbounded LFS storage before it ends. None disables mid-run squashing. Defaults to None.

Note

Uploads run on rank 0 only, and failures are logged but never crash the run. Place PushToHub after the checkpoint saver in the hooks list so the files it reads are already written. Repeatedly uploading large checkpoints to the same path accumulates orphaned LFS blobs in the repo’s history; squash_history (end of run) and squash_every_n_pushes (mid run) reclaim that storage via super_squash_history.

Example

Add to cfg.hooks after the checkpoint saver; by default it uploads model_best to the Hub once at the end of training:

hooks = [
    dict(type="CheckpointSaver"),
    dict(type="PushToHub", repo_id="youngsm/sonata-pilarnet-L",
         private=True),
]
# → in after_train (rank 0, blocking) uploads
#   <save_path>/model/model_best.pth to
#   hf://youngsm/sonata-pilarnet-L/model_best.pth; with on_best=True it
#   also pushes model_best in after_step whenever best_metric_value
#   improves, and every_n_epochs=N pushes the rolling ckpt in after_epoch
after_epoch()[source]#

Push the rolling checkpoint every N epochs when configured.

after_step()[source]#

Push model_best on improvement and/or the rolling ckpt on a step cadence.

after_train()[source]#

Push the configured checkpoint when training finishes (blocking).

before_train()[source]#

Seed the step counter from resumed progress (mirrors the iteration saver).