Flexibility to set device number < total device count

[yhl48]

Before submitting

• Was this discussed/agreed via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Fixes #152. However, this implementation is not fully correct, it will break if there is more than one node.

[yhl48]

@tchaton is there a way to get the devices variable?

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 4 lines in your changes missing coverage. Please review.

Please upload report for BASE (main@935eef5). Learn more about missing BASE report.

Additional details and impacted files

@@          Coverage Diff          @@
##             main   #155   +/-   ##
=====================================
  Coverage        ?    77%           
=====================================
  Files           ?     30           
  Lines           ?   4126           
  Branches        ?      0           
=====================================
  Hits            ?   3186           
  Misses          ?    940           
  Partials        ?      0           

[yhl48]

the problem might be more complicated than I thought, usually devices is set in the lightning Trainer class, but _DistributedEnv is called before that