fix "Expected all tensors to be on the same device, but found at least two devices" error #11690
Conversation
|
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. |
|
@SunMarc , could you pls help review? Thx very much. |
| if hidden_states.device != res_hidden_states.device: | ||
| res_hidden_states = res_hidden_states.to(hidden_states.device) |
There was a problem hiding this comment.
we shouldn't need that since both hidden_states and res_hidden_states should be on the same device no ? The pre-forward hook added by accelerate should be move all the inputs to the same device.
There was a problem hiding this comment.
@SunMarc , i suppose this is a corner case? torch.cat is a weight-less function, so seems cannot covered by the pre-forward hook set by accelerate...
There was a problem hiding this comment.
I mean since hidden_states and res_hidden_states_tuple are in the forward definition, they should be moved to the same device by the pre-forward hook added by accelerate
There was a problem hiding this comment.
@SunMarc We run into a corner case here. Since we have 8 cards here, so the determined device_map(by https://github.com/huggingface/diffusers/blob/1bc6f3dc0f21779480db70a4928d14282c0198ed/src/diffusers/models/model_loading_utils.py#L64C5-L64C26) is
device_map: OrderedDict([('conv_in', 0), ('time_proj', 0), ('time_embedding', 0), ('down_blocks.0', 0), ('down_blocks.1.resnets.0', 1), ('up_blocks.0.resnets.0', 1), ('up_blocks.0.resnets.1', 2), ('up_blocks.0.upsamplers', 2), ('up_blocks.1', 3), ('mid_block.attentions', 3), ('conv_norm_out', 4), ('conv_act', 4), ('conv_out', 4), ('mid_block.resnets', 4)])
We can see UpBlock is not the atomic module, its submodules are assigned to different devices(up_blocks.0.resnets.0, up_blocks.0.resnets.1), so pre-hook for UpBlock will not help in this case. And since torch.cat is not pre-hooked(and cannot since it's a function rather than a module?), so the issue happens.
If there is no a torch.cat btw the sub-blocks in UpBlock, things will be all fine.
There was a problem hiding this comment.
@SunMarc, need your inputs in how to proceed for this corner case, thx.
pytest -rA tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_load_sharded_checkpoint_device_map_from_hub_localon 8 devices(CUDA, XPU), there will be a RuntimeError "RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:4 and cuda:1!", fix it by moving to same device@sayakpaul , pls help review, thx.