[Device map] nn.Parameter don't have children

[patrickvonplaten]

This is needed to allow to use the "device_map" for models that have nn.Parameter at the top level, such as the safety checker of diffusers: https://github.com/huggingface/diffusers/pull/772/files#r990269773

[HuggingFaceDocBuilderDev]

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

[sgugger]

What's a reproducer? I'm not sure how you get to a parameter here since we only apply the children method.

[patrickvonplaten]

What's a reproducer? I'm not sure how you get to a parameter here since we only apply the children method.

Haha that was fast. Here a repro:

from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker

StableDiffusionSafetyChecker._no_split_modules = ["CLIPEncoderLayer"]
pipe = StableDiffusionSafetyChecker.from_pretrained('CompVis/stable-diffusion-safety-checker', device_map='auto')

(sorry, it's 1GB of download)

[sgugger]

Which branch should I check? On main I just get an error telling me device_map auto is not supported yet

[patrickvonplaten]

Ah sorry just updated it above, can you try:

from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker

StableDiffusionSafetyChecker._no_split_modules = ["CLIPEncoderLayer"]
pipe = StableDiffusionSafetyChecker.from_pretrained('CompVis/stable-diffusion-safety-checker', device_map='auto')

[sgugger]

Thanks for the reproducer! The function is indeed used with some parameters inside accelerate, so this is needed.

Just suggested another test.

[patrickvonplaten]

Awesome! Can I merge after one ✔️ here or wait for another reviewer?

[sgugger]

Nope, you can go ahead and merge!