[Unified TorchTrainer] Add PyTorch Lightning Trainer Utilities

[woshiyyya]

Why are these changes needed?

Add Lightning related utilities for Unified TorchTrainer.

• 1 basic example using new API
• User guides for Lightning + TorchTrainer

POC example: Text Classification with Ray Data + Lightning

Rendered Doc: https://anyscale-ray--37989.com.readthedocs.build/en/37989/

New utilities

• Strategy Class: RayDDPStrategy, RayFSDPStrategy, RayDeepSpeedStrategy
• Environment Class: RayLightningEnvironment
• Metrics and checkpoint reporting: RayTrainReportCallback
• prepare_trainer(): To check if the users correctly configured the pl.Trainer

Users can inject these utilities to run Lightning code in TorchTrainer.

Metrics and checkpoint reporting

We propose to let the users define the logs in a lightning callback themselves.

from pytorch_lightning.callbacks import Callback
class RayTrainReportCallback(Callback):
    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        with TemporaryDirectory() as tmpdir:
            # Fetch metrics
            metrics = trainer.callback_metrics
            metrics = {k: v.item() for k, v in metrics.items()}
           
            # Save checkpoint to local
            ckpt_path = os.path.join(tmpdir, f"ckpt_epoch_{trainer.current_epoch}.pth")
            trainer.save_checkpoint(ckpt_path, weights_only=False)

            # Report to train session
            checkpoint = Checkpoint.from_directory(tmpdir)
            ray.train.report(metrics=metrics, checkpoint=checkpoint)

The users can choose whether to report the metrics to Ray Train or not.

from pytorch_lightning.callbacks import ModelCheckpoint

def train_loop_per_worker():
    # Report to ray train
    trainer = pl.Trainer(
         callbacks=[RayTrainReportCallback()]
    )

   # Keep using the original lightning's logics
   trainer = pl.Trainer(
         callbacks=[ModelCheckpoint(**kwags)]
    )

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[matthewdeng]

Looks really clean. I think we should showcase one example with the previous API compare it with the new API.

[scottsun94]

If the user chooses not to report to train, we won't be able to print any related metrics in the train output right?

What will happen to this training result table?

Training completed after 0 iterations at 2023-06-26 16:06:41. Total running time: 24min 41s
╭──────────────────────────────────────────────────────╮
│ Training result                                      │
├──────────────────────────────────────────────────────┤
│ config/train_loop_config/args   ...weight_decay=0.0) │
╰──────────────────────────────────────────────────────╯

[krfricke]

Read half of it, will continue later today. Looks great so far! Would be great to get a version that renders in RTD for preview.

[krfricke]

Nit: let's use comments here so that the python code at least would parse correctly (even if it isn't executed)

Suggested change

	...
	# ...

[matthewdeng]

Technically this actually works in python!

[woshiyyya]

@krfricke Updated rendered doc: https://anyscale-ray--37989.com.readthedocs.build/en/37989/
cc @matthewdeng

[matthewdeng]

Generally looks great to me, I think that we can polish and shift around some of the documentation content in a follow up PR.

[krfricke]

This is great, minor nits