Unverified Commit 63a5c874 authored by Sayak Paul's avatar Sayak Paul Committed by GitHub
Browse files

Update distributed_inference.md to include `transformer.device_map` (#9553)



* Update distributed_inference.md to include `transformer.device_map`

* Update docs/source/en/training/distributed_inference.md
Co-authored-by: default avatarSteven Liu <59462357+stevhliu@users.noreply.github.com>

---------
Co-authored-by: default avatarSteven Liu <59462357+stevhliu@users.noreply.github.com>
parent 12878229
...@@ -177,7 +177,7 @@ transformer = FluxTransformer2DModel.from_pretrained( ...@@ -177,7 +177,7 @@ transformer = FluxTransformer2DModel.from_pretrained(
``` ```
> [!TIP] > [!TIP]
> At any point, you can try `print(pipeline.hf_device_map)` to see how the various models are distributed across devices. This is useful for tracking the device placement of the models. > At any point, you can try `print(pipeline.hf_device_map)` to see how the various models are distributed across devices. This is useful for tracking the device placement of the models. You can also try `print(transformer.hf_device_map)` to see how the transformer model is sharded across devices.
Add the transformer model to the pipeline for denoising, but set the other model-level components like the text encoders and VAE to `None` because you don't need them yet. Add the transformer model to the pipeline for denoising, but set the other model-level components like the text encoders and VAE to `None` because you don't need them yet.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment