Does DeepSpeed support OpenFlamingo (or other VL models)?

Hi there,

thanks for this great repo! I am trying to model parrallel inference of Deepspeed on OpenFlamingo using the following code.

import logging
import huggingface_hub
import os
from open_flamingo import create_model_and_transforms
from huggingface_hub import hf_hub_download
import torch
import deepspeed

local_rank = int(os.getenv('LOCAL_RANK', '0'))
world_size = int(os.getenv('WORLD_SIZE', '1'))

MODEL = {
    "language": "anas-awadalla/mpt-7b",
    "flamingo": "openflamingo/OpenFlamingo-9B-vitl-mpt7b",
    "cross_attn_every_n_layers": 4
}


model, image_processor, tokenizer = create_model_and_transforms(
    clip_vision_encoder_path="ViT-L-14",
    clip_vision_encoder_pretrained="openai",
    lang_encoder_path=MODEL["language"],
    tokenizer_path=MODEL["language"],
    cross_attn_every_n_layers=MODEL["cross_attn_every_n_layers"],
)


# grab model checkpoint from huggingface hub
huggingface_hub.login(
    token=<TOKEN>
)
checkpoint_path = hf_hub_download(MODEL["flamingo"], "checkpoint.pt")
model.load_state_dict(torch.load(checkpoint_path), strict=False)

ds_engine = deepspeed.init_inference(
    model,
    mp_size=world_size,
    dtype=torch.float32,
    replace_with_kernel_inject=True,
)



from PIL import Image
import requests

"""
Step 1: Load images
"""
demo_image_one = Image.open(
    requests.get(
        "http://images.cocodataset.org/val2017/000000039769.jpg", stream=True
    ).raw
)

demo_image_two = Image.open(
    requests.get(
        "http://images.cocodataset.org/test-stuff2017/000000028137.jpg",
        stream=True
    ).raw
)

query_image = Image.open(
    requests.get(
        "http://images.cocodataset.org/test-stuff2017/000000028352.jpg",
        stream=True
    ).raw
)

BS = 64
"""
Step 2: Preprocessing images
Details: For OpenFlamingo, we expect the image to be a torch tensor of shape 
 batch_size x num_media x num_frames x channels x height x width. 
 In this case batch_size = 1, num_media = 3, num_frames = 1,
 channels = 3, height = 224, width = 224.
"""
vision_x = [image_processor(demo_image_one).unsqueeze(0), image_processor(demo_image_two).unsqueeze(0), image_processor(query_image).unsqueeze(0)]
vision_x = torch.cat(vision_x, dim=0)
vision_x = vision_x.unsqueeze(1).unsqueeze(0)
vision_x = vision_x.expand(BS, -1, -1, -1, -1, -1)

"""
Step 3: Preprocessing text
Details: In the text we expect an <image> special token to indicate where an image is.
 We also expect an <|endofchunk|> special token to indicate the end of the text 
 portion associated with an image.
"""
tokenizer.padding_side = "left" # For generation padding tokens should be on the left
lang_x = tokenizer(
    ["<image>An image of two dogs.<|endofchunk|><image>An image of a basketball.<|endofchunk|><image>An image of"],
    return_tensors="pt",
)
# duplicate along the batch dimension
lang_x = {k: v.expand(BS, -1) for k, v in lang_x.items()}



"""
Step 4: Generate text
"""
model = ds_engine.module
device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
vision_x = vision_x.to(device)
lang_x = {k: v.to(device) for k, v in lang_x.items()}

generated_text = model.generate(
    vision_x=vision_x,
    lang_x=lang_x["input_ids"],
    attention_mask=lang_x["attention_mask"],
    max_new_tokens=20,
    num_beams=3,
)

for i in range(BS):
    print("Generated text: ", tokenizer.decode(generated_text[i]))

However, the models are still loaded on each GPU rather than model parallel and the memory is still the same compared to DDP. I have following questions along with some thoughts after reading related issues.

1. Regarding the model loading, is it because this statement model.load_state_dict(torch.load(checkpoint_path), strict=False) before ds_engine = deepspeed.init_inference(...) ?

2. What is the correct way to load the weights of such VL model? It consists of both vision and language encoders as well as some interleaved layers. It seems like that this setting does not belong to the three scenarios in this tutorial, i.e., 1. directly loading from huggingface, 2. deepspeed models, 3. Megatron models.

3. Can I use Deepspeed to easily support such VL models (vision encoder + large language models + some interleaved layers inside the LMs)? From what I have seen in the supported models, most of the supported models are Large Language Models. If yes, any suggestions, such as tutorials or notes?

Thanks in advance and looking forward to your reply!

Start with the DeepSpeed inference entry point used here, deepspeed.init_inference, and inspect the model policies in deepspeed/module_inject/replace_policy.py. Compare the OpenFlamingo loading flow with the inference tutorial's checkpoint-loading scenarios; done would require a documented, tested path for model-parallel vision-language inference.