[tests] make TestDeepSpeedModelZoo device-agnostic

[faaany]

What does this PR do?

This PR is to enable the TestDeepSpeedModelZoo tests to work on a non-CUDA accelerator.

@amyeroberts and @ydshieh

[amyeroberts]

Thanks!

[amyeroberts]

Let's get a second 👍 from @ydshieh the test master to confirm this is OK on our CI runs

[faaany]

Let's get a second 👍 from @ydshieh the test master to confirm this is OK on our CI runs

Sure!

[faaany]

I further updated my code to get the real device count. But the CI seems to not use the latest transformers version. @ydshieh is this something intended?

[ydshieh]

@faaany . It uses the latest dev. transformers version (i.e. this PR).

The problem is that, some variables like BACKEND_DEVICE_COUNT is only defined when torch is available. But we have other CI environment where only Tensorflow or only Flax is installed without torch.

Let me check what to do with it.

if is_torch_available():
    # Mappings from device names to callable functions to support device agnostic
    # testing.
    BACKEND_MANUAL_SEED = {"cuda": torch.cuda.manual_seed, "cpu": torch.manual_seed, "default": torch.manual_seed}
    BACKEND_EMPTY_CACHE = {"cuda": torch.cuda.empty_cache, "cpu": None, "default": None}
    BACKEND_DEVICE_COUNT = {"cuda": torch.cuda.device_count, "cpu": lambda: 0, "default": lambda: 1}

[ydshieh]

Could you add

if is_torch_available():
    # Mappings from device names to callable functions to support device agnostic
    # testing.
    BACKEND_MANUAL_SEED = {"cuda": torch.cuda.manual_seed, "cpu": torch.manual_seed, "default": torch.manual_seed}
    BACKEND_EMPTY_CACHE = {"cuda": torch.cuda.empty_cache, "cpu": None, "default": None}
    BACKEND_DEVICE_COUNT = {"cuda": torch.cuda.device_count, "cpu": lambda: 0, "default": lambda: 1}

to

if is_torch_available():
    # Mappings from device names to callable functions to support device agnostic
    # testing.
    BACKEND_MANUAL_SEED = {"cuda": torch.cuda.manual_seed, "cpu": torch.manual_seed, "default": torch.manual_seed}
    BACKEND_EMPTY_CACHE = {"cuda": torch.cuda.empty_cache, "cpu": None, "default": None}
    BACKEND_DEVICE_COUNT = {"cuda": torch.cuda.device_count, "cpu": lambda: 0, "default": lambda: 1}
else:
    BACKEND_MANUAL_SEED = {"default": None}
    BACKEND_EMPTY_CACHE = {"default": None}
    BACKEND_DEVICE_COUNT = {"default": lambda: 0}

and see how it goes? 🙏 thank you in advance

[faaany]

@ydshieh yes, it works, thx for the suggestion! The CI still doesn't run through, but I think it is not caused by this PR.

[ydshieh]

Yeah, rebase on main to include #31407 would make CI green.

[ydshieh]

Thank you !

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.