Lightning-AI / Lightning-AI/pytorch-lightning

"Migrations" for lightning models and trainers

Open
#14,646 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

discussion feature won't fix
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## 🚀 Feature

There should be a simple way to "migrate" (mutate) a lightning model and trainer / optimizer.

### Motivation

It's common when training large scale neural networks, or building off existing large-scale nets, to mess with it in unholy ways. This allows one to leverage previous long-term training times, but still customize the network (or optimizer) as need be.

* For example, take a vision net and add a few layers at the end or chop off a few layers.
* Shrink a linear layer by randomly dropping some of the output dimensions.

This is an old-school strategy that Nico Pinto taught me over ten years ago (I was shocked by how gross but effective it was). My guess is that this is used in the 128-node clusters with super advanced ML practitioners.

### Pitch

This would allow lightning to use some FAANG-type black-magic stuff that isn't talked about much.

Large scale networks tend to diverge for unknown reasons very late in the training process. Sometimes just adjusting the LR is fine, but also modifying the network architecture might be in order.

Additionally, modifying an existing network is a great way to take some huge model you trained a week ago, and then run much simpler faster experiments on a truncated version. (You can fine-tune it quickly and do ad-hoc ablation studies to see if your net really needs that large size.)

### Alternatives

Futz around manually and maybe footgun on some lightning safeguard I'm not aware of.

### Additional context

This is an advanced version of the features requested in #14531 and #5339

I think just futzing with the optimizer (in those PRs) would be a great start. The model migration stuff is a little trickier to design in a clean way.

cc @borda

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading the related requests in #14531 and #5339, then examine how Lightning models, trainers, and optimizers are currently handled. The issue identifies optimizer migration as a possible first step, while model migration remains a design problem; done requires a clearly defined migration API and behavior for customized architectures or optimizer state.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.