[Ray SGD] LightningModule integration + MNIST Example

[amogkam]

Docs to come later

Why are these changes needed?

Related issue number

Checks

• 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 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 :(

[richardliaw]

looks pretty good; good work!

can you add some unit tests and address comments?

[sumanthratna]

amazing!! I haven't had a chance to pull and try this out but @amogkam do you know if the following work with PTL x RaySGD?

• mixed precision
• batch size and lr finding (trainer.tune(model))

we might also want to add a sentence describing whether any functionality is lost to the docs PR

[amogkam]

@sumanthratna Thanks for the comment- let me know how it goes if you plan to try it out!

This will have all the same functionality as standard Ray SGD so there is mixed precision support with Nvidia Apex. But Pytorch native mixed precision, auto batch size tuning, and auto lr rate tuning are not supported yet. These are probably good next steps to work on.

And yep, I have a list of all features that we currently don't support and will add it to the docs :)

[richardliaw]

looks good; ping when tests pass?

[amogkam]

@richardliaw travis passing, should be ready to merge