deepspeedai / deepspeedai/DeepSpeed

For Flux inpaint distribute inference (Multi-GPU)

Open
#7,005 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
43.1k
Forks
5k
Avg merge
4d 15h
Merged PRs (30d)
112

Description

import torch
from diffusers.utils import load_image, check_min_version
from src.controlnet_flux import FluxControlNetModel
from src.transformer_flux import FluxTransformer2DModel
from src.pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline
from PIL import Image
import deepspeed
import torchvision.transforms as transforms

check_min_version("0.30.2")


# Set image path , mask path and prompt
image_path='/data2/workspace/bydeng/DATASETS/VAE/train/image/alf_main_1136580672_1636813_1_00_1_artist_0.png'
mask_path='/data2/workspace/bydeng/DATASETS/VAE/train/mask/alf_main_1136580672_1636813_1_00_1_artist_0.png'
prompt='nothing'

# Build pipeline
controlnet = FluxControlNetModel.from_pretrained("/data1/common_models/FLUX.1-dev-Controlnet-Inpainting-Beta/", torch_dtype=torch.bfloat16)
transformer = FluxTransformer2DModel.from_pretrained(
        "/data1/common_models/FLUX.1-dev/", subfolder='transformer', torch_dtype=torch.bfloat16
    )
pipe = FluxControlNetInpaintingPipeline.from_pretrained(
    "/data1/common_models/FLUX.1-dev/",
    controlnet=controlnet,
    transformer=transformer,
    torch_dtype=torch.bfloat16,
).to("cpu")
pipe.transformer.to(torch.bfloat16)
pipe.controlnet.to(torch.bfloat16)



ds_model = deepspeed.init_inference(
    model=pipe,      
    # tensor_parallel={"tp_size": 1},
    mp_size=4,
    dtype='fp16', 
    # replace_method="auto", 
    replace_with_kernel_inject=True, 
)

# Load image and mask
size = (768, 768)
image = load_image(image_path).convert("RGB").resize(size)
mask = load_image(mask_path).convert("RGB").resize(size)
generator = torch.Generator(device="cpu").manual_seed(24)

# Inpaint
result = ds_model.module(
    prompt=prompt,
    height=size[1],
    width=size[0],
    control_image=image,
    control_mask=mask,
    num_inference_steps=28,
    generator=generator,
    controlnet_conditioning_scale=0.9,
    guidance_scale=3.5,
    negative_prompt="",
    true_guidance_scale=1.0 # default: 3.5 for alpha and 1.0 for beta
).images[0]

result.save('flux_inpaint_ddp_1.png')
print("Successfully inpaint image")

can not get the output image when setting tp_size>1, is there any solution?

(bydeng) (base) bydeng@iv-ydgjd36ghsk36d113bdn:/data2/workspace/bydeng/Projects/FLUX-Controlnet-Inpainting$ deepspeed --include localhost:6 deepspeed_test.py 
[2025-02-05 15:42:28,251] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-02-05 15:42:30,222] [WARNING] [runner.py:215:fetch_hostfile] Unable to find hostfile, will proceed with training with local resources only.
[2025-02-05 15:42:30,222] [INFO] [runner.py:607:main] cmd = /data2/users/bydeng/.conda/envs/bydeng/bin/python -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbNl19 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None deepspeed_test.py
[2025-02-05 15:42:31,506] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-02-05 15:42:33,289] [INFO] [launch.py:146:main] WORLD INFO DICT: {'localhost': [6]}
[2025-02-05 15:42:33,289] [INFO] [launch.py:152:main] nnodes=1, num_local_procs=1, node_rank=0
[2025-02-05 15:42:33,289] [INFO] [launch.py:163:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0]})
[2025-02-05 15:42:33,289] [INFO] [launch.py:164:main] dist_world_size=1
[2025-02-05 15:42:33,289] [INFO] [launch.py:168:main] Setting CUDA_VISIBLE_DEVICES=6
[2025-02-05 15:42:33,289] [INFO] [launch.py:256:main] process 400760 spawned with command: ['/data2/users/bydeng/.conda/envs/bydeng/bin/python', '-u', 'deepspeed_test.py', '--local_rank=0']
[2025-02-05 15:42:36,151] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  8.45it/s]
Loading pipeline components...:  29%|█████████████████████████████████████████▏                                                                                                      | 2/7 [00:00<00:00,  6.81it/s]You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
Loading pipeline components...: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  9.78it/s]
[2025-02-05 15:42:38,297] [INFO] [logging.py:128:log_dist] [Rank -1] DeepSpeed info: version=0.16.3, git-hash=unknown, git-branch=unknown
[2025-02-05 15:42:38,298] [WARNING] [config_utils.py:70:_process_deprecated_field] Config parameter mp_size is deprecated use tensor_parallel.tp_size instead
[2025-02-05 15:42:38,298] [INFO] [logging.py:128:log_dist] [Rank -1] quantize_bits = 8 mlp_extra_grouping = False, quantize_groups = 1
**** found and replaced vae w. <class 'deepspeed.model_implementations.diffusers.vae.DSVAE'>
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 28/28 [00:15<00:00,  1.78it/s]
Successfully inpaint image
[2025-02-05 15:43:04,293] [INFO] [launch.py:351:main] Process 400760 exits successfully.
(bydeng) (base) bydeng@iv-ydgjd36ghsk36d113bdn:/data2/workspace/bydeng/Projects/FLUX-Controlnet-Inpainting$ deepspeed --include localhost:6,7 deepspeed_test.py 
[2025-02-05 15:43:35,537] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-02-05 15:43:37,288] [WARNING] [runner.py:215:fetch_hostfile] Unable to find hostfile, will proceed with training with local resources only.
[2025-02-05 15:43:37,289] [INFO] [runner.py:607:main] cmd = /data2/users/bydeng/.conda/envs/bydeng/bin/python -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbNiwgN119 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None deepspeed_test.py
[2025-02-05 15:43:38,617] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-02-05 15:43:40,529] [INFO] [launch.py:146:main] WORLD INFO DICT: {'localhost': [6, 7]}
[2025-02-05 15:43:40,529] [INFO] [launch.py:152:main] nnodes=1, num_local_procs=2, node_rank=0
[2025-02-05 15:43:40,529] [INFO] [launch.py:163:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0, 1]})
[2025-02-05 15:43:40,529] [INFO] [launch.py:164:main] dist_world_size=2
[2025-02-05 15:43:40,529] [INFO] [launch.py:168:main] Setting CUDA_VISIBLE_DEVICES=6,7
[2025-02-05 15:43:40,530] [INFO] [launch.py:256:main] process 407930 spawned with command: ['/data2/users/bydeng/.conda/envs/bydeng/bin/python', '-u', 'deepspeed_test.py', '--local_rank=0']
[2025-02-05 15:43:40,530] [INFO] [launch.py:256:main] process 407931 spawned with command: ['/data2/users/bydeng/.conda/envs/bydeng/bin/python', '-u', 'deepspeed_test.py', '--local_rank=1']
[2025-02-05 15:43:43,482] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-02-05 15:43:43,536] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
Loading pipeline components...:   0%|                                                                                                                                                        | 0/7 [00:00<?, ?it/s]You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 10.62it/s]
Loading pipeline components...: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 11.75it/s]
                                                                                                                                                                                                                  [2025-02-05 15:43:45,197] [INFO] [logging.py:128:log_dist] [Rank -1] DeepSpeed info: version=0.16.3, git-hash=unknown, git-branch=unknown                                                     | 0/2 [00:00<?, ?it/s]
[2025-02-05 15:43:45,197] [WARNING] [config_utils.py:70:_process_deprecated_field] Config parameter mp_size is deprecated use tensor_parallel.tp_size instead
[2025-02-05 15:43:45,197] [INFO] [logging.py:128:log_dist] [Rank -1] quantize_bits = 8 mlp_extra_grouping = False, quantize_groups = 1
[2025-02-05 15:43:45,197] [INFO] [comm.py:652:init_distributed] cdb=None
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  8.07it/s]
Loading pipeline components...:  43%|█████████████████████████████████████████████████████████████▋                                                                                  | 3/7 [00:00<00:00,  6.22it/s]You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
Loading pipeline components...: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  9.88it/s]
[2025-02-05 15:43:45,679] [INFO] [logging.py:128:log_dist] [Rank -1] DeepSpeed info: version=0.16.3, git-hash=unknown, git-branch=unknown
[2025-02-05 15:43:45,680] [WARNING] [config_utils.py:70:_process_deprecated_field] Config parameter mp_size is deprecated use tensor_parallel.tp_size instead
[2025-02-05 15:43:45,680] [INFO] [logging.py:128:log_dist] [Rank -1] quantize_bits = 8 mlp_extra_grouping = False, quantize_groups = 1
[2025-02-05 15:43:45,680] [INFO] [comm.py:652:init_distributed] cdb=None
[2025-02-05 15:43:45,680] [INFO] [comm.py:683:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
**** found and replaced vae w. <class 'deepspeed.model_implementations.diffusers.vae.DSVAE'>
**** found and replaced vae w. <class 'deepspeed.model_implementations.diffusers.vae.DSVAE'>
[2025-02-05 15:43:55,532] [INFO] [launch.py:319:sigkill_handler] Killing subprocess 407930
[2025-02-05 15:43:55,532] [INFO] [launch.py:319:sigkill_handler] Killing subprocess 407931
[2025-02-05 15:43:55,566] [ERROR] [launch.py:325:sigkill_handler] ['/data2/users/bydeng/.conda/envs/bydeng/bin/python', '-u', 'deepspeed_test.py', '--local_rank=1'] exits with return code = -11

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with deepspeed_test.py and the deepspeed.init_inference entry point, focusing on tensor_parallel.tp_size and the two-GPU launch command using localhost:6,7. Reproduce the successful single-GPU run and the multi-GPU segmentation fault, then inspect DeepSpeed inference initialization around the diffusers pipeline. Done means multi-GPU inference completes and saves the output image without a crash.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.