[deepspeed] fix load_best_model test

[stas00]

Fixes #17151 to run from_pretrained with emulated dist env.

While it was working on my setup on CI it failed with:

tests/deepspeed/test_deepspeed.py:756: in test_load_best_model
    model = T5ForConditionalGeneration.from_pretrained(T5_TINY)
src/transformers/modeling_utils.py:2116: in from_pretrained
    init_contexts = [deepspeed.zero.Init(config_dict_or_path=deepspeed_config())] + init_contexts
/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py:693: in __init__
    self.local_device = torch.device('cuda:{}'.format(os.environ["LOCAL_RANK"]))
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
self = environ({'NPP_VERSION': '11.3.2.139', 'NVIDIA_VISIBLE_DEVICES': 'all', 'DALI_BUILD': '2054952', 'GITHUB_WORKSPACE': '/...RRENT_TEST': 'tests/deepspeed/test_deepspeed.py::TrainerIntegrationDeepSpeed::test_load_best_model_zero3_fp16 (call)'})
key = 'LOCAL_RANK'
    def __getitem__(self, key):
        try:
            value = self._data[self.encodekey(key)]
        except KeyError:
            # raise KeyError with the original key value
>           raise KeyError(key) from None
E           KeyError: 'LOCAL_RANK'

the test is exactly the same, just moved a big chunk of it into the with mockenv_context - no code changes

@sgugger

[sgugger]

Thanks for fixing!

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.