Lightning Lite Examples

[awaelchli]

What does this PR do?

This is the V1 for the new Lightning Lite package. It bundles all major changes together, but individual PRs will be done to get it merged (e.g., #10175, #10176)
Planned to be released as part of 1.5.

Demo

TODOs

• Precision support

  • NativeMixedPrecision
  • Delay support ApexMixedPrecision
  • TPUBfloat (verify it works)
  • Double

• Plugin support

  • DeepSpeed
  • DeepSpeed with multiple models (@tchaton, @awaelchli)
  • TPUSpawn: DDPSpawn changes need to be carried down into this plugin. A minor refactor could be merged to PL master. (@kaushikb11)
  • Sharded (optimizer sharding)
  • Delay support FullySharded (model sharding)

• Move data to device automatically

• Move model to device automatically

• Allow only one model per setup() call

• DataLoader setup: Currently, there is no distributed sampler.

• Resolve miscellaneous TODOs in the code base

• Fix changes that broke Lightning tests

• Make self.setup() take model and optimizers positionally.

• Unit testing, parity tests

• Typing (mypy)

Discussions

• LightningLite constructor arguments: We are currently changing the Trainer constructor arguments to support a new pattern: Trainer(accelerator="cpu/tpu/gpu", strategy="ddp/deepspeed/...", devices=X) Should we start promoting this directly in the LightningLite API?

• Deepspeed API for backward: The user can't call loss.backward(), it needs to be called on the model. Which API do we want to offer:

  A) self.backward(loss, model) # model is optional for plugins other than deepspeed
  B) model.backward(loss)

  In both cases, user needs to change their code if they switch from one strategy to the next.

• Deepspeed API for optimization step: Plain Deepspeed requires a call to model.step() as opposed to the usual optimizer.step(). Since we wrap the optimizers of the user anyway, we could still offer optimizer.step() and redirect to model.step(). It would mean less code changes for the user switching between plugins, but for deepspeed users it might be confusing!

Related work:

• We are currently also redesigning the Precision/TrainingType plugin interactions in collaboration with FB.

Part of #1 (it's a lie, this is just here to avoid noisy GitHub bot)

[ananthsub]

The design document isn't visible to people outside of grid.aI so it's hard to know the context for this

[tchaton]

The design document isn't visible to people outside of grid.aI so it's hard to know the context for this

Hey @ananthsub. Find the document here: https://docs.google.com/document/d/1b10LMNqnv1ellVTAEIlJFV5KvBuxIlFCTnNB3SYIFok/edit#heading=h.jl44rslqge7e.

Best,
T.C

[awaelchli]

This does not belong in this PR. Why did we add this?
Needs to be addressed in #10297