[Discussion] Unnecessary hack code for FSDP

[wanxinzzz]

Hello, thanks for your great work!

I ran this repo successfully and i was confused by the hack code for combine CLIPVisionModel with FSDP in https://github.com/haotian-liu/LLaVA/blob/main/llava/model/llava.py#L48

class LlavaLlamaModel(LlamaModel):
    config_class = LlavaConfig

    def __init__(self, config: LlamaConfig, mm_vision_tower=None, mm_hidden_size=None):
        super(LlavaLlamaModel, self).__init__(config)

        if hasattr(config, "mm_vision_tower"):
            # HACK: for FSDP
            self.vision_tower = [CLIPVisionModel.from_pretrained(config.mm_vision_tower)]
            # self.vision_tower = CLIPVisionModel.from_pretrained(config.mm_vision_tower)

I don't think this is necessary code for FSDP. And i tried to assign model directly to vision_tower instead of loading it with list.

class LlavaLlamaModel(LlamaModel):
    config_class = LlavaConfig

    def __init__(self, config: LlamaConfig, mm_vision_tower=None, mm_hidden_size=None):
        super(LlavaLlamaModel, self).__init__(config)

        if hasattr(config, "mm_vision_tower"):
            # HACK: for FSDP
            # self.vision_tower = [CLIPVisionModel.from_pretrained(config.mm_vision_tower)]
            self.vision_tower = CLIPVisionModel.from_pretrained(config.mm_vision_tower)

        if hasattr(config, "use_mm_proj"):
            self.mm_projector = nn.Linear(config.mm_hidden_size, config.hidden_size)

    def initialize_vision_modules(self, vision_tower, mm_vision_select_layer,
                                  pretrain_mm_mlp_adapter=None, tune_mm_mlp_adapter=False):
        self.config.mm_vision_tower = vision_tower

        image_processor = CLIPImageProcessor.from_pretrained(vision_tower)

        if not hasattr(self, 'vision_tower'):
            vision_tower = CLIPVisionModel.from_pretrained(vision_tower)
        else:
            vision_tower = self.vision_tower
        vision_tower.requires_grad_(False)
        # vision_tower = vision_tower.to(torch.float16)
        self.vision_tower = vision_tower

After that, i commented out all the code that adjusted the vision_tower device or dtype and it works.

The codes include

1. https://github.com/haotian-liu/LLaVA/blob/main/llava/model/llava.py#L65

vision_tower = vision_tower.to(torch.float16)

2. https://github.com/haotian-liu/LLaVA/blob/main/llava/train/train.py#L611

model.get_model().vision_tower[0].to(dtype=dtype, device=training_args.device)

Is that right, or does it cause some other problem?

[zsxm1998]

Also curious about this.