Automatic model checkpointing for pytorch-lightning training

[WeichenXu123]

🛠 DevTools 🛠

Install mlflow from this PR

pip install git+https://github.com/mlflow/mlflow.git@refs/pull/10935/merge

Checkout with GitHub CLI

gh pr checkout 10935

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Automatic model checkpointing for pytorch-lightning training,

Design doc: https://docs.google.com/document/d/1Ke7-8og_KzV3WE5xOS4XKSLZsISdTsVLlol14vJ3S70/edit

Demo notebook: https://e2-dogfood.staging.cloud.databricks.com/?o=6051921418418893#notebook/2173893049403456

How is this PR tested?

• Existing unit/integration tests
• New unit/integration tests
• Manual tests

Does this PR require documentation update?

• No. You can skip the rest of this section.
• Yes. I've updated:
  • Examples
  • API references
  • Instructions

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/deployments: MLflow Deployments client APIs, server, and third-party Deployments integrations
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
• area/projects: MLproject format, project running backends
• area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

Language

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

• rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
• rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
• rn/feature - A new user-facing feature worth mentioning in the release notes
• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
• rn/documentation - A user-facing documentation change worth mentioning in the release notes

[github-actions]

Documentation preview for 7a8527c will be available when this CircleCI job completes successfully.

• Top page
• Changed pages

More info

• Ignore this comment if this PR does not change the documentation.
• It takes a few minutes for the preview to be available.
• The preview is updated when a new commit is pushed to this PR.
• This comment was created by https://github.com/mlflow/mlflow/actions/runs/7897383534.

[chenmoneygithub]

Thanks Weichen for the PR!

I took a high-level pass, and the logic looks good to me. Left one comment for argument name choice. Can you share a reproducible notebook or github gist with how to use this new feature? Would be easier to review and spot errors from there, thank you!

[chenmoneygithub]

this variable name is a bit strange. can we combine model_checkpoint_every_n_epochs and model_checkpoint_train_time_interval_S tosave_freq? I find that's pretty clean and easy to use:
https://github.com/keras-team/keras/blob/fe2f54aa5bc42fb23a96449cf90434ab9bb6a2cd/keras/callbacks/model_checkpoint.py#L112

[WeichenXu123]

One question:

The save_freq in keras checkpointing callback supports "do checkpoint after N batches (steps) execution", do we want to support this ? Or we only want to support epoch based checkpointing ?

My suggestion is we only support epoch based checkpointing, for batch based checkpointing, per-batch metric validation result is less accurate , and in pytorch-lightning, per-batch validation is not available except you log the metric with on_step=True i.e. LightningModule.log(metric_name, value, on_epoch=True, on_step=True)

[WeichenXu123]

Demo notebook is attached in the PR description.

[chenmoneygithub]

save_freq supports both per-epoch saving and per-N-steps saving. Both scenarios are commonly used based on my experience with the production team.

in pytorch-lightning, per-batch validation is not available

I think we can just use the callback hook? It should have the training stats available: https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.Callback.html#lightning.pytorch.callbacks.Callback.on_after_backward

[WeichenXu123]

summary:

In ModelCheckpointCallback, if we use per-N-steps saving way, it needs to check "monitor" metric and the metric must be updated on per step execution, we need like log(metric_name, value, on_step=True) to enable per-step updated metric.

We have 2 options:
(1) In Mlflow ModelCheckpointCallback, we update the monitor metric in on_after_backward.
(2) Document this to tell user to add code in their module class train_step method: log(metric_name, value, on_step=True) for the metric used as "monitor"

Option (1) has one issue: In on_after_backward we can't get the data for metric computation So we might have to choose option (2)

Option (2) is the way that current built-in pytorch-lightning ModelCheckpointCallback uses.

@chenmoneygithub

[chenmoneygithub]

Yea if 2) is lightning's behavior, we can proceed with that one.

[WeichenXu123]

Note:

save_checkpoint can save both model weights and trainer states.

We can restore the "trainer" by https://pytorch-lightning.readthedocs.io/en/0.8.5/weights_loading.html#restoring-training-state

So shall we provide a helper function to restore the trainer ? (similar to the load_latest_checkpoint)

[chenmoneygithub]

yes, resuming training from a saved checkpoint is a common use case. Also users are not necessarily loading the latest checkpoint, so I thnk we can provide a public API mlflow.pytorch.load_checkpoint, which can load pytorch checkpoint, which includes model weights and optimizer states. This can also be used for vanilla pytorch workflow, not necessarily lightning workflows.

[WeichenXu123]

oh, the new version API becomes:

trainer.fit(model, ckpt_path="some/path/to/my_checkpoint.ckpt")

So we can't return a trainer directly without fit invocation.

https://lightning.ai/docs/pytorch/stable/common/checkpointing_basic.html#resume-training-state

[chenmoneygithub]

Got it. My previous comment was like we can also load the model weights specifically without relying on the trainer instance, which means:

model = mlflow.pytorch.load_checkpoint(model_class, mlflow_uri)
trainer = trainer(model)

For example, if we are sharing this model, and someone using vanilla pytorch wants to finetune it with a custom training loop, then they can load the checkpoints.

[WeichenXu123]

Note:

No matter "save_weights_only" true or false, when loading back, it always requires the "model_class", and it does not need the model object.

This is different with Keras side API:
https://github.com/mlflow/mlflow/pull/10955/files#r1471374884

[chenmoneygithub]

makes sense to me!

[chenmoneygithub]

Thanks Weichen! I played with the notebook, the experience is pretty smooth. Two things about the file we save with checkpoints:

• Can we put checkpoints along with the metrics into one directory? Now it's flat files in a parent directory. With a large number of checkpoints logged, it's hard to navigate.
• We don't necessarily have eval metrics for every checkpoint we save, e.g., if we save checkpoints per 5000 steps, we may not have the eval metrics. We can make "eval at checkpoint saving" optional when saving per N steps.

[chenmoneygithub]

Thanks Weichen for the PR, looks pretty good! Left some comments on nit and argument choice.

[chenmoneygithub]

is train: "pl.Trainer" a legal type annotation? Shall we import the class?

[WeichenXu123]

It is legal if pl is available.

[chenmoneygithub]

ohh I mean shall we remove the double quotes?

[WeichenXu123]

double quotes is legal syntax, and other mlflow code also uses it. :)

[chenmoneygithub]

gotcha, I learned something new, thank you!

[chenmoneygithub]

It might be a better UX to raise an explicit error than silently failing.

[WeichenXu123]

I prefer to use:

                _logger.error(
                    "If MLflowModelCheckpoint 'save_best_only' config is True, it requires to "
                    "compare the monitored metric value, but the provided monitored metric value "
                    "is not available."
                )

instead of raise exception,

because the default checkpoint configs are:

    checkpoint=True,
    checkpoint_monitor="val_loss",
    checkpoint_mode="min",
    checkpoint_save_best_only=True,

but sometimes the model might not log "val_loss" metric, so that if it raise exception here, it breaks the whole autologging , but we don't need to break other part autologging (e.g. log params / metrics)

@chenmoneygithub

[chenmoneygithub]

makes sense to me!

On top of that, may we call out "checkpoint logging is skipped" in the error message?

[harupy]

why do we need to set this?

[WeichenXu123]

For testing "save_best_only" related logic. e.g. "mode" is "min" , then only loss decreases, the new checkpoint is a better one

[harupy]

I feel like this will be broken in the future once the pytorch-lighting team removes _callback_metrics or renames it because they don't know this. Relying on private attributes in third-party packages is too risky and would increase maintenance burden.

[WeichenXu123]

Updated! :)

[WeichenXu123]

Updated again.

[WeichenXu123]

@mlflow-automation autoformat

[WeichenXu123]

@mlflow-automation autoformat