Skip to content

Commit

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

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

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

Co-authored-by: Steven Liu <[email protected]>

---------

Co-authored-by: Steven Liu <[email protected]>
  • Loading branch information
2 people authored and 蒋硕 committed Oct 11, 2024
1 parent b70d147 commit 7b8ef40
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 7b8ef40

Please sign in to comment.