Add support for using custom Environments and Strategies ⚡

[amorehead]

What does this PR do?

Adds the ability for one to employ custom (e.g., Lightning Fabric) Environments (e.g., SlurmEnvironment) and Strategies (e.g., DeepSpeedStrategy) during model training or evaluation.

Before submitting

• Did you make sure title is self-explanatory and the description concisely explains the PR?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you list all the breaking changes introduced by this pull request?
• Did you test your PR locally with pytest command?
• Did you run pre-commit hooks with pre-commit run -a command?

Did you have fun?

Make sure you had fun coding 🙃 ⚡

[codecov-commenter]

Codecov Report

Attention: 29 lines in your changes are missing coverage. Please review.

Comparison is base (bddbc24) 83.24% compared to head (27fd562) 78.60%.

Files	Patch %	Lines
src/__init__.py	37.50%	10 Missing ⚠️
src/train.py	60.86%	9 Missing ⚠️
src/eval.py	50.00%	8 Missing ⚠️
src/models/mnist_module.py	50.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #608      +/-   ##
==========================================
- Coverage   83.24%   78.60%   -4.65%     
==========================================
  Files          11       12       +1     
  Lines         376      430      +54     
==========================================
+ Hits          313      338      +25     
- Misses         63       92      +29     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[giladturok]

@amorehead Can you explain your rationale behind adding these additional configs? Super curious to know your thought process.

Is it related to PyTorch Lighting's use of a "strategy" defined here: https://lightning.ai/docs/pytorch/stable/extensions/strategy.html?

And some example strategies:

[amorehead]

Hey, @gil2rok. Yes, Lightning's strategy class is what I had in mind here. I wanted to make it possible with this template to use any arbitrary (e.g., advanced) training strategy such as DeepSpeed with only a few lines of code changes.

[giladturok]

Hey, @gil2rok. Yes, Lightning's strategy class is what I had in mind here. I wanted to make it possible with this template to use any arbitrary (e.g., advanced) training strategy such as DeepSpeed with only a few lines of code changes.

Because this pull request has not been accepted, how did you use these advanced strategies? Did you just manually add them the PyTorch Lightning trainer arguments?

[amorehead]

As you can see here, now one can specify an optional strategy for Lightning to use e.g., via OmegaConf YAML config files.

[amorehead]

@gil2rok

[amorehead]

@ashleve, may I ask for you to review these changes when you get a chance? I've been using them in my own forked projects for a couple years now, and I think the community would benefit from having them be available by default.

[giladturok]

+1 to add this feature. Would be very helpful IMO. Cool work @amorehead!!

…

On Sat, Apr 19, 2025 at 7:41 PM Alex Morehead ***@***.***> wrote: @ashleve <https://github.com/ashleve>, may I ask for you to review these changes when you get a chance? I've been using them in my own forked projects for a couple years now, and I think the community would benefit from having them be available by default. — Reply to this email directly, view it on GitHub <#608 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AIZ4NS6UTLZW76PPRZMNE5L22J4FLAVCNFSM6AAAAAB3OXKDR6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQMJWG43TOMJSGQ> . You are receiving this because you were mentioned.Message ID: ***@***.***> *amorehead* left a comment (ashleve/lightning-hydra-template#608) <#608 (comment)> @ashleve <https://github.com/ashleve>, may I ask for you to review these changes when you get a chance? I've been using them in my own forked projects for a couple years now, and I think the community would benefit from having them be available by default. — Reply to this email directly, view it on GitHub <#608 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AIZ4NS6UTLZW76PPRZMNE5L22J4FLAVCNFSM6AAAAAB3OXKDR6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQMJWG43TOMJSGQ> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[amorehead]

Hi, @ashleve. Are you open to making me a collaborator of this GitHub repository so I can help manage some of its pending pull requests (such as this one)?

[giladturok]

@ashleve +1 for the above. Would love to have another maintainer (if they're willing) and to get this amazing feature into this template!