New models

• All models inherit from a VLM class. See documentation in src/vlm.py

    """
    Generic vision-language model (VLM) class.
    A VLM consists of four components:
        1. A vision encoder that extracts features from pixels, e.g. CLIP
            input: (B, T_img, F, C, H, W)
            output: (B, T_img, F, v, d)
        2. A vision tokenizer that converts these features to visual token-like embeddings, e.g. Perceiver, or a linear projection head
            input: (B, T_img, F, v, d)
            output: (B, T_img, n, d)
        3. A fusion method that allows the language model to attend to these tokens, e.g. cross-attention, or placing the tokens directly in the language model's input sequence
        4. A language model
    """

• Models are further split into those that inherit from VLMWithCrossAttention (dense xattn to fuse vision + language, Flamingo-style) vs. VLMWithLanguageStream (insert vision tokens into the language stream, Kosmos-style)
• BLIP, Kosmos, and Flamingo implemented

FSDP Updates

• FSDP rewritten to reflect recent updates to Pytorch nightly, which allow for mixed frozen/unfrozen params within an FSDP wrapper, and for frozen params to reshard during the backward pass.
• GPU power util for Flamingo model with facebook/opt-6.7B backbone is > 70%, compared to ~30% previously (this is with gradient ckpting).
• This also removes the need to untie LM embeddings, and seems to address the issues w/ backward resharding shapes.

Training code refactor

• train_one_epoch now accepts a list of datasets and executes the same loss function on all of them. This permits us to decide the datasets to train on at runtime (e.g. just LAION) and makes adding in datasets more flexible. To train on a dataset, set the --{dataset_name}_shards arg (e.g. --laion_shards). If this is None, then we will not train on that dataset (i.e., skip LAION)
• train_one_epoch also now accepts a loss function decided at runtime. Losses are found in train/losses.py. Currently, only next token prediction is implemented, but this allows us to work on adding contrastive-generative losses.
• Most of the FSDP / Deepspeed code has been moved to train/distributed.py in an attempt to streamline train/train.py

Steps before merging

• Replicate OF-3B as a sanity check
• Test training with deepspeed
• Check embedding logic (cc @anas-awadalla)
• Check whether anything in eval code needs to change; I think the image caching for classification needs to be completely rewritten now.
• Merge in deepspeed eval code from the deepspeed branch
• Update released weights or modify this code to be compatible with released weights (e.g., I have switched to using lang_model instead of lang_encoder; this will not play well with the released weights; we need to decide what to do about the embeddings).
• Merge in updated / modular eval code from wilds branch

Steps after merging

• Add support for image generation components, e.g. Emu-style
• Add support for additional losses, e.g. CoCA-style
• Add back in additional datasets that were being explored previously on this branch (only mmc4 and laion are in the code for now)