Skip to content

Commit

Permalink
Update distributed_inference.md to include transformer.device_map (#…
Browse files Browse the repository at this point in the history
…9553)

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

* Update docs/source/en/training/distributed_inference.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

---------

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
  • Loading branch information
sayakpaul and stevhliu authored Oct 8, 2024
1 parent 1287822 commit 63a5c87
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion docs/source/en/training/distributed_inference.md
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,7 @@ transformer = FluxTransformer2DModel.from_pretrained(
```

> [!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.

Expand Down

0 comments on commit 63a5c87

Please sign in to comment.